refactor: message content use struct

liushuangls · Mar 14, 2024 · a2a97c5 · a2a97c5
1 parent 492fee4
commit a2a97c5
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 43 deletions.
diff --git a/README.md b/README.md
@@ -40,7 +40,7 @@ func main() {
 	resp, err := client.CreateMessages(context.Background(), anthropic.MessagesRequest{
 		Model: anthropic.ModelClaudeInstant1Dot2,
 		Messages: []anthropic.Message{
-			{Role: anthropic.RoleUser, Content: "What is your name?"},
+			anthropic.NewUserTextMessage("What is your name?"),
 		},
 		MaxTokens: 1000,
 	})
@@ -75,7 +75,7 @@ func main() {
 		MessagesRequest: anthropic.MessagesRequest{
 			Model: anthropic.ModelClaudeInstant1Dot2,
 			Messages: []anthropic.Message{
-				{Role: anthropic.RoleUser, Content: "What is your name?"},
+				anthropic.NewUserTextMessage("What is your name?"),
 			},
 			MaxTokens:   1000,
 		},
@@ -121,25 +121,19 @@ func main() {
 	if err != nil {
 		panic(err)
 	}
-    
+
 	resp, err := client.CreateMessages(context.Background(), anthropic.MessagesRequest{
-		Model: anthropic.ModelClaude3Opus20240229, // only claude 3 model can use vision
+		Model: anthropic.ModelClaude3Opus20240229,
 		Messages: []anthropic.Message{
 			{
 				Role: anthropic.RoleUser,
-				Content: []any{
-					anthropic.MessageImageContent{
-						Type: "image",
-						Source: anthropic.MessageImageContentSource{
-							Type:      "base64",
-							MediaType: imageMediaType,
-							Data:      imageData,
-						},
-					},
-					anthropic.MessageTextContent{
-						Type: "text",
-						Text: "Describe this image.",
-					},
+				Content: []anthropic.MessageContent{
+					anthropic.NewImageMessageContent(anthropic.MessageContentImageSource{
+						Type:      "base64",
+						MediaType: imageMediaType,
+						Data:      imageData,
+					}),
+					anthropic.NewTextMessageContent("Describe this image."),
 				},
 			},
 		},

diff --git a/message.go b/message.go
@@ -41,21 +41,67 @@ func (m *MessagesRequest) SetTopK(k int) {
 }
 
 type Message struct {
-	Role    string `json:"role"`
-	Content any    `json:"content"` // Content can be string, MessageTextContent or MessageImageContent or slice
+	Role    string           `json:"role"`
+	Content []MessageContent `json:"content"`
 }
 
-type MessageTextContent struct {
-	Type string `json:"type"`
-	Text string `json:"text"`
+func NewUserTextMessage(text string) Message {
+	return Message{
+		Role:    "user",
+		Content: []MessageContent{NewTextMessageContent(text)},
+	}
+}
+
+func NewAssistantTextMessage(text string) Message {
+	return Message{
+		Role:    "assistant",
+		Content: []MessageContent{NewTextMessageContent(text)},
+	}
+}
+
+func (m Message) GetFirstContent() MessageContent {
+	if len(m.Content) == 0 {
+		return MessageContent{}
+	}
+	return m.Content[0]
+}
+
+type MessageContent struct {
+	Type   string                     `json:"type"`
+	Text   *string                    `json:"text,omitempty"`
+	Source *MessageContentImageSource `json:"source,omitempty"`
 }
 
-type MessageImageContent struct {
-	Type   string                    `json:"type"`
-	Source MessageImageContentSource `json:"source"`
+func NewTextMessageContent(text string) MessageContent {
+	return MessageContent{
+		Type: "text",
+		Text: &text,
+	}
+}
+
+func NewImageMessageContent(source MessageContentImageSource) MessageContent {
+	return MessageContent{
+		Type:   "image",
+		Source: &source,
+	}
+}
+
+func (m MessageContent) IsTextContent() bool {
+	return m.Type == "text"
+}
+
+func (m MessageContent) IsImageContent() bool {
+	return m.Type == "image"
+}
+
+func (m MessageContent) GetText() string {
+	if m.IsTextContent() && m.Text != nil {
+		return *m.Text
+	}
+	return ""
 }
 
-type MessageImageContentSource struct {
+type MessageContentImageSource struct {
 	Type      string `json:"type"`
 	MediaType string `json:"media_type"`
 	Data      any    `json:"data"`

diff --git a/message_stream_test.go b/message_stream_test.go
@@ -33,7 +33,7 @@ func TestMessagesStream(t *testing.T) {
 		MessagesRequest: anthropic.MessagesRequest{
 			Model: anthropic.ModelClaudeInstant1Dot2,
 			Messages: []anthropic.Message{
-				{Role: anthropic.RoleUser, Content: "What is your name?"},
+				anthropic.NewUserTextMessage("What is your name?"),
 			},
 			MaxTokens: 1000,
 		},
@@ -81,7 +81,7 @@ func TestMessagesStreamError(t *testing.T) {
 		MessagesRequest: anthropic.MessagesRequest{
 			Model: anthropic.ModelClaudeInstant1Dot2,
 			Messages: []anthropic.Message{
-				{Role: anthropic.RoleUser, Content: "What is your name?"},
+				anthropic.NewUserTextMessage("What is your name?"),
 			},
 			MaxTokens: 1000,
 		},

diff --git a/message_test.go b/message_test.go
@@ -38,7 +38,7 @@ func TestMessages(t *testing.T) {
 	resp, err := client.CreateMessages(context.Background(), anthropic.MessagesRequest{
 		Model: anthropic.ModelClaudeInstant1Dot2,
 		Messages: []anthropic.Message{
-			{Role: anthropic.RoleUser, Content: "What is your name?"},
+			anthropic.NewUserTextMessage("What is your name?"),
 		},
 		MaxTokens: 1000,
 	})
@@ -65,7 +65,7 @@ func TestMessagesTokenError(t *testing.T) {
 	_, err := client.CreateMessages(context.Background(), anthropic.MessagesRequest{
 		Model: anthropic.ModelClaudeInstant1Dot2,
 		Messages: []anthropic.Message{
-			{Role: anthropic.RoleUser, Content: "What is your name?"},
+			anthropic.NewUserTextMessage("What is your name?"),
 		},
 		MaxTokens: 1000,
 	})
@@ -109,19 +109,13 @@ func TestMessagesVision(t *testing.T) {
 		Messages: []anthropic.Message{
 			{
 				Role: anthropic.RoleUser,
-				Content: []any{
-					anthropic.MessageImageContent{
-						Type: "image",
-						Source: anthropic.MessageImageContentSource{
-							Type:      "base64",
-							MediaType: imageMediaType,
-							Data:      imageData,
-						},
-					},
-					anthropic.MessageTextContent{
-						Type: "text",
-						Text: "Describe this image.",
-					},
+				Content: []anthropic.MessageContent{
+					anthropic.NewImageMessageContent(anthropic.MessageContentImageSource{
+						Type:      "base64",
+						MediaType: imageMediaType,
+						Data:      imageData,
+					}),
+					anthropic.NewTextMessageContent("Describe this image."),
 				},
 			},
 		},