feat: add InvokeLLMWithStructuredOutput functionality (#369)

* feat: add InvokeLLMWithStructuredOutput functionality

- Introduced a new method InvokeLLMWithStructuredOutput to the BackwardsInvocation interface for handling structured output requests.
- Added corresponding request and response types to support structured output.
- Implemented the method in both RealBackwardsInvocation and MockedDifyInvocation for testing purposes.
- Updated permission handling and task execution for the new structured output invocation type.

This enhancement allows for more flexible and detailed responses from the LLM, improving the overall functionality of the invocation system.

* refactor: enhance LLMResultChunkWithStructuredOutput structure

- Updated the LLMResultChunkWithStructuredOutput type to include additional fields: Model, SystemFingerprint, and Delta.
- Added comments to clarify the reasoning behind the structure and the use of type embedding for JSON marshaling.

This change improves the clarity and functionality of the LLMResultChunkWithStructuredOutput type, ensuring proper JSON serialization.

* refactor: streamline LLMResultChunk construction in InvokeLLMWithStructuredOutput

- Simplified the construction of LLMResultChunk and LLMResultChunkWithStructuredOutput by removing unnecessary type embedding.
- Enhanced readability and maintainability of the code while preserving functionality.

This change contributes to cleaner code and improved clarity in the handling of structured output responses.

Author: Yeuoly

internal/core/dify_invocation/invcation.go

@@ -9,6 +9,9 @@ import (
 type BackwardsInvocation interface {
 	// InvokeLLM
 	InvokeLLM(payload *InvokeLLMRequest) (*stream.Stream[model_entities.LLMResultChunk], error)
+	// InvokeLLMWithStructuredOutput
+	InvokeLLMWithStructuredOutput(payload *InvokeLLMWithStructuredOutputRequest) (
+		*stream.Stream[model_entities.LLMResultChunkWithStructuredOutput], error)
 	// InvokeTextEmbedding
 	InvokeTextEmbedding(payload *InvokeTextEmbeddingRequest) (*model_entities.TextEmbeddingResult, error)
 	// InvokeRerank

internal/core/dify_invocation/real/http_request.go

@@ -115,6 +115,12 @@ func (i *RealBackwardsInvocation) InvokeLLM(payload *dify_invocation.InvokeLLMRe
 	return StreamResponse[model_entities.LLMResultChunk](i, "POST", "invoke/llm", http_requests.HttpPayloadJson(payload))
 }
 
+func (i *RealBackwardsInvocation) InvokeLLMWithStructuredOutput(payload *dify_invocation.InvokeLLMWithStructuredOutputRequest) (
+	*stream.Stream[model_entities.LLMResultChunkWithStructuredOutput], error,
+) {
+	return StreamResponse[model_entities.LLMResultChunkWithStructuredOutput](i, "POST", "/invoke/llm/structured-output", http_requests.HttpPayloadJson(payload))
+}
+
 func (i *RealBackwardsInvocation) InvokeTextEmbedding(payload *dify_invocation.InvokeTextEmbeddingRequest) (*model_entities.TextEmbeddingResult, error) {
 	return Request[model_entities.TextEmbeddingResult](i, "POST", "invoke/text-embedding", http_requests.HttpPayloadJson(payload))
 }

internal/core/dify_invocation/tester/mock.go

@@ -136,6 +136,62 @@ func (m *MockedDifyInvocation) InvokeLLM(payload *dify_invocation.InvokeLLMReque
 	return stream, nil
 }
 
+func (m *MockedDifyInvocation) InvokeLLMWithStructuredOutput(payload *dify_invocation.InvokeLLMWithStructuredOutputRequest) (
+	*stream.Stream[model_entities.LLMResultChunkWithStructuredOutput], error,
+) {
+	// generate json from payload.StructuredOutputSchema
+	structuredOutput, err := jsonschema.GenerateValidateJson(payload.StructuredOutputSchema)
+	if err != nil {
+		return nil, err
+	}
+
+	// marshal jsonSchema to string
+	structuredOutputString := parser.MarshalJson(structuredOutput)
+
+	// split structuredOutputString into 10 parts and write them to the stream
+	parts := []string{}
+	for i := 0; i < 10; i++ {
+		start := i * len(structuredOutputString) / 10
+		end := (i + 1) * len(structuredOutputString) / 10
+		if i == 9 { // last part
+			end = len(structuredOutputString)
+		}
+		parts = append(parts, structuredOutputString[start:end])
+	}
+
+	stream := stream.NewStream[model_entities.LLMResultChunkWithStructuredOutput](11)
+	routine.Submit(nil, func() {
+		for i, part := range parts {
+			stream.Write(model_entities.LLMResultChunkWithStructuredOutput{
+				Model:             model_entities.LLMModel(payload.Model),
+				SystemFingerprint: "test",
+				Delta: model_entities.LLMResultChunkDelta{
+					Index: &[]int{i}[0],
+					Message: model_entities.PromptMessage{
+						Role:      model_entities.PROMPT_MESSAGE_ROLE_ASSISTANT,
+						Content:   part,
+						Name:      "test",
+						ToolCalls: []model_entities.PromptMessageToolCall{},
+					},
+				},
+			})
+		}
+		// write the last part
+		stream.Write(model_entities.LLMResultChunkWithStructuredOutput{
+			Model:             model_entities.LLMModel(payload.Model),
+			SystemFingerprint: "test",
+			Delta: model_entities.LLMResultChunkDelta{
+				Index: &[]int{10}[0],
+			},
+			LLMStructuredOutput: model_entities.LLMStructuredOutput{
+				StructuredOutput: structuredOutput,
+			},
+		})
+		stream.Close()
+	})
+	return stream, nil
+}
+
 func (m *MockedDifyInvocation) InvokeTextEmbedding(payload *dify_invocation.InvokeTextEmbeddingRequest) (*model_entities.TextEmbeddingResult, error) {
 	result := model_entities.TextEmbeddingResult{
 		Model: payload.Model,

internal/core/dify_invocation/types.go

@@ -18,6 +18,7 @@ type InvokeType string
 
 const (
 	INVOKE_TYPE_LLM                      InvokeType = "llm"
+	INVOKE_TYPE_LLM_STRUCTURED_OUTPUT    InvokeType = "llm_structured_output"
 	INVOKE_TYPE_TEXT_EMBEDDING           InvokeType = "text_embedding"
 	INVOKE_TYPE_RERANK                   InvokeType = "rerank"
 	INVOKE_TYPE_TTS                      InvokeType = "tts"
@@ -51,6 +52,15 @@ type InvokeLLMRequest struct {
 	InvokeLLMSchema
 }
 
+type InvokeLLMWithStructuredOutputRequest struct {
+	BaseInvokeDifyRequest
+	requests.BaseRequestInvokeModel
+	// requests.InvokeLLMSchema
+	// TODO: as completion_params in requests.InvokeLLMSchema is "model_parameters", we declare another one here
+	InvokeLLMSchema
+	StructuredOutputSchema map[string]any `json:"structured_output_schema" validate:"required"`
+}
+
 type InvokeTextEmbeddingRequest struct {
 	BaseInvokeDifyRequest
 	requests.BaseRequestInvokeModel

internal/core/plugin_daemon/backwards_invocation/task.go

@@ -154,6 +154,12 @@ var (
 			},
 			"error": "permission denied, you need to enable app access in plugin manifest",
 		},
+		dify_invocation.INVOKE_TYPE_LLM_STRUCTURED_OUTPUT: {
+			"func": func(declaration *plugin_entities.PluginDeclaration) bool {
+				return declaration.Resource.Permission.AllowInvokeLLM()
+			},
+			"error": "permission denied, you need to enable llm access in plugin manifest",
+		},
 	}
 )
 
@@ -250,6 +256,9 @@ var (
 		dify_invocation.INVOKE_TYPE_FETCH_APP: func(handle *BackwardsInvocation) {
 			genericDispatchTask(handle, executeDifyInvocationFetchAppTask)
 		},
+		dify_invocation.INVOKE_TYPE_LLM_STRUCTURED_OUTPUT: func(handle *BackwardsInvocation) {
+			genericDispatchTask(handle, executeDifyInvocationLLMStructuredOutputTask)
+		},
 	}
 )
 
@@ -337,6 +346,26 @@ func executeDifyInvocationLLMTask(
 	}
 }
 
+func executeDifyInvocationLLMStructuredOutputTask(
+	handle *BackwardsInvocation,
+	request *dify_invocation.InvokeLLMWithStructuredOutputRequest,
+) {
+	response, err := handle.backwardsInvocation.InvokeLLMWithStructuredOutput(request)
+	if err != nil {
+		handle.WriteError(fmt.Errorf("invoke llm with structured output model failed: %s", err.Error()))
+		return
+	}
+
+	for response.Next() {
+		value, err := response.Read()
+		if err != nil {
+			handle.WriteError(fmt.Errorf("read llm with structured output model failed: %s", err.Error()))
+			return
+		}
+		handle.WriteResponse("stream", value)
+	}
+}
+
 func executeDifyInvocationTextEmbeddingTask(
 	handle *BackwardsInvocation,
 	request *dify_invocation.InvokeTextEmbeddingRequest,

pkg/entities/model_entities/llm.go

@@ -188,6 +188,22 @@ type LLMResultChunk struct {
 	Delta             LLMResultChunkDelta `json:"delta" validate:"required"`
 }
 
+type LLMStructuredOutput struct {
+	StructuredOutput map[string]any `json:"structured_output" validate:"omitempty"`
+}
+
+type LLMResultChunkWithStructuredOutput struct {
+	// You might argue that why not embed LLMResultChunk directly?
+	// `LLMResultChunk` has implemented interface `MarshalJSON`, due to Golang's type embedding,
+	// it also effectively implements the `MarshalJSON` method of `LLMResultChunkWithStructuredOutput`,
+	// resulting in a unexpected JSON marshaling of `LLMResultChunkWithStructuredOutput`
+	Model             LLMModel            `json:"model" validate:"required"`
+	SystemFingerprint string              `json:"system_fingerprint" validate:"omitempty"`
+	Delta             LLMResultChunkDelta `json:"delta" validate:"required"`
+
+	LLMStructuredOutput
+}
+
 /*
 This is a compatibility layer for the old LLMResultChunk format.
 The old one has the `PromptMessages` field, we need to ensure the new one is backward compatible.