image_content_test.go

 1package llm
 2
 3import (
 4	"encoding/json"
 5	"testing"
 6)
 7
 8func TestImageContent(t *testing.T) {
 9	// Create a Content structure with an image
10	imageContent := Content{
11		Type:      ContentTypeText, // In the future, we might add a specific ContentTypeImage
12		MediaType: "image/jpeg",
13		Data:      "/9j/4AAQSkZJRg...", // Shortened base64 encoded image
14	}
15
16	// Verify the structure is correct
17	if imageContent.MediaType != "image/jpeg" {
18		t.Errorf("Expected MediaType to be 'image/jpeg', got '%s'", imageContent.MediaType)
19	}
20
21	if imageContent.Data != "/9j/4AAQSkZJRg..." {
22		t.Errorf("Expected Data to contain base64 image data")
23	}
24
25	// Create a tool result that contains text and image content
26	toolResult := Content{
27		Type:      ContentTypeToolResult,
28		ToolUseID: "toolu_01A09q90qw90lq917835lq9",
29		ToolResult: []Content{
30			{
31				Type: ContentTypeText,
32				Text: "15 degrees",
33			},
34			imageContent,
35		},
36	}
37
38	// Check that the tool result contains two content items
39	if len(toolResult.ToolResult) != 2 {
40		t.Errorf("Expected tool result to contain 2 content items, got %d", len(toolResult.ToolResult))
41	}
42
43	// Verify JSON marshaling works as expected
44	bytes, err := json.Marshal(toolResult)
45	if err != nil {
46		t.Errorf("Failed to marshal content to JSON: %v", err)
47	}
48
49	// Unmarshal and verify structure is preserved
50	var unmarshaled Content
51	if err := json.Unmarshal(bytes, &unmarshaled); err != nil {
52		t.Errorf("Failed to unmarshal JSON: %v", err)
53	}
54
55	if len(unmarshaled.ToolResult) != 2 {
56		t.Errorf("Expected unmarshaled tool result to contain 2 content items, got %d", len(unmarshaled.ToolResult))
57	}
58
59	if unmarshaled.ToolResult[1].MediaType != "image/jpeg" {
60		t.Errorf("Expected unmarshaled image MediaType to be 'image/jpeg', got '%s'", unmarshaled.ToolResult[1].MediaType)
61	}
62}