Browse Source

fix: sync schema of prompt message content

Yeuoly 5 months ago
parent
commit
ff64f54c26
2 changed files with 42 additions and 1 deletions
  1. 2 1
      pkg/entities/model_entities/llm.go
  2. 40 0
      pkg/entities/model_entities/llm_test.go

+ 2 - 1
pkg/entities/model_entities/llm.go

@@ -99,7 +99,8 @@ func isPromptMessageContentType(fl validator.FieldLevel) bool {
 
 type PromptMessageContent struct {
 	Type         PromptMessageContentType `json:"type" validate:"required,prompt_message_content_type"`
-	Data         string                   `json:"data" validate:"required"`
+	Base64Data   string                   `json:"base64_data"` // for multi-modal data
+	Data         string                   `json:"data"`        // for text only
 	EncodeFormat string                   `json:"encode_format"`
 	Format       string                   `json:"format"`
 	MimeType     string                   `json:"mime_type"`

+ 40 - 0
pkg/entities/model_entities/llm_test.go

@@ -1,6 +1,7 @@
 package model_entities
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/parser"
@@ -261,3 +262,42 @@ func TestZeroLLMUsage(t *testing.T) {
 		t.Error(err)
 	}
 }
+
+func TestTextPromptMessage(t *testing.T) {
+	const (
+		promptMessage = `
+		{
+			"role": "user",
+			"content": "hello"
+		}
+		`
+	)
+
+	_, err := parser.UnmarshalJsonBytes[PromptMessage]([]byte(promptMessage))
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestImagePromptMessage(t *testing.T) {
+	const (
+		promptMessage = `
+		{
+			"role": "user",
+			"content": [
+				{
+					"type": "image",
+					"data": "base64"
+				}
+			]
+		}
+		`
+	)
+
+	a, err := parser.UnmarshalJsonBytes[PromptMessage]([]byte(promptMessage))
+
+	fmt.Println(a.Content)
+	if err != nil {
+		t.Error(err)
+	}
+}