go-llm/v2/message_test.go

package llm

import (
	"testing"

	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/provider"
)

func TestUserMessage(t *testing.T) {
	msg := UserMessage("hello")
	if msg.Role != RoleUser {
		t.Errorf("expected role=user, got %v", msg.Role)
	}
	if msg.Content.Text != "hello" {
		t.Errorf("expected text='hello', got %q", msg.Content.Text)
	}
	if len(msg.Content.Images) != 0 {
		t.Errorf("expected no images, got %d", len(msg.Content.Images))
	}
}

func TestUserMessageWithImages(t *testing.T) {
	img1 := Image{URL: "https://example.com/1.png"}
	img2 := Image{Base64: "abc123", ContentType: "image/png"}

	msg := UserMessageWithImages("describe", img1, img2)
	if msg.Role != RoleUser {
		t.Errorf("expected role=user, got %v", msg.Role)
	}
	if msg.Content.Text != "describe" {
		t.Errorf("expected text='describe', got %q", msg.Content.Text)
	}
	if len(msg.Content.Images) != 2 {
		t.Fatalf("expected 2 images, got %d", len(msg.Content.Images))
	}
	if msg.Content.Images[0].URL != "https://example.com/1.png" {
		t.Errorf("expected image[0] URL, got %q", msg.Content.Images[0].URL)
	}
	if msg.Content.Images[1].Base64 != "abc123" {
		t.Errorf("expected image[1] base64='abc123', got %q", msg.Content.Images[1].Base64)
	}
	if msg.Content.Images[1].ContentType != "image/png" {
		t.Errorf("expected image[1] contentType='image/png', got %q", msg.Content.Images[1].ContentType)
	}
}

func TestSystemMessage(t *testing.T) {
	msg := SystemMessage("Be helpful")
	if msg.Role != RoleSystem {
		t.Errorf("expected role=system, got %v", msg.Role)
	}
	if msg.Content.Text != "Be helpful" {
		t.Errorf("expected text='Be helpful', got %q", msg.Content.Text)
	}
}

func TestAssistantMessage(t *testing.T) {
	msg := AssistantMessage("Sure thing")
	if msg.Role != RoleAssistant {
		t.Errorf("expected role=assistant, got %v", msg.Role)
	}
	if msg.Content.Text != "Sure thing" {
		t.Errorf("expected text='Sure thing', got %q", msg.Content.Text)
	}
}

func TestToolResultMessage(t *testing.T) {
	msg := ToolResultMessage("tc-123", "result data")
	if msg.Role != RoleTool {
		t.Errorf("expected role=tool, got %v", msg.Role)
	}
	if msg.ToolCallID != "tc-123" {
		t.Errorf("expected toolCallID='tc-123', got %q", msg.ToolCallID)
	}
	if msg.Content.Text != "result data" {
		t.Errorf("expected text='result data', got %q", msg.Content.Text)
	}
}

func TestConvertMessages(t *testing.T) {
	msgs := []Message{
		SystemMessage("system prompt"),
		UserMessageWithImages("look at this", Image{URL: "https://example.com/img.png"}),
		{
			Role:    RoleAssistant,
			Content: Content{Text: "I'll use a tool"},
			ToolCalls: []ToolCall{
				{ID: "tc1", Name: "search", Arguments: `{"q":"test"}`},
			},
		},
		ToolResultMessage("tc1", "found it"),
	}

	converted := convertMessages(msgs)

	if len(converted) != 4 {
		t.Fatalf("expected 4 converted messages, got %d", len(converted))
	}

	// System message
	if converted[0].Role != "system" {
		t.Errorf("msg[0]: expected role='system', got %q", converted[0].Role)
	}
	if converted[0].Content != "system prompt" {
		t.Errorf("msg[0]: expected content='system prompt', got %q", converted[0].Content)
	}

	// User message with images
	if converted[1].Role != "user" {
		t.Errorf("msg[1]: expected role='user', got %q", converted[1].Role)
	}
	if len(converted[1].Images) != 1 {
		t.Fatalf("msg[1]: expected 1 image, got %d", len(converted[1].Images))
	}
	if converted[1].Images[0].URL != "https://example.com/img.png" {
		t.Errorf("msg[1]: expected image URL, got %q", converted[1].Images[0].URL)
	}

	// Assistant message with tool calls
	if converted[2].Role != "assistant" {
		t.Errorf("msg[2]: expected role='assistant', got %q", converted[2].Role)
	}
	if len(converted[2].ToolCalls) != 1 {
		t.Fatalf("msg[2]: expected 1 tool call, got %d", len(converted[2].ToolCalls))
	}
	if converted[2].ToolCalls[0].ID != "tc1" {
		t.Errorf("msg[2]: expected tool call ID='tc1', got %q", converted[2].ToolCalls[0].ID)
	}
	if converted[2].ToolCalls[0].Name != "search" {
		t.Errorf("msg[2]: expected tool call name='search', got %q", converted[2].ToolCalls[0].Name)
	}
	if converted[2].ToolCalls[0].Arguments != `{"q":"test"}` {
		t.Errorf("msg[2]: expected tool call arguments, got %q", converted[2].ToolCalls[0].Arguments)
	}

	// Tool result message
	if converted[3].Role != "tool" {
		t.Errorf("msg[3]: expected role='tool', got %q", converted[3].Role)
	}
	if converted[3].ToolCallID != "tc1" {
		t.Errorf("msg[3]: expected toolCallID='tc1', got %q", converted[3].ToolCallID)
	}
	if converted[3].Content != "found it" {
		t.Errorf("msg[3]: expected content='found it', got %q", converted[3].Content)
	}
}

func TestConvertProviderResponse(t *testing.T) {
	t.Run("text only", func(t *testing.T) {
		resp := convertProviderResponse(provider.Response{
			Text: "hello",
			Usage: &provider.Usage{
				InputTokens:  10,
				OutputTokens: 5,
				TotalTokens:  15,
			},
		})
		if resp.Text != "hello" {
			t.Errorf("expected text='hello', got %q", resp.Text)
		}
		if resp.HasToolCalls() {
			t.Error("expected no tool calls")
		}
		if resp.Usage == nil {
			t.Fatal("expected usage")
		}
		if resp.Usage.InputTokens != 10 {
			t.Errorf("expected 10 input tokens, got %d", resp.Usage.InputTokens)
		}

		msg := resp.Message()
		if msg.Role != RoleAssistant {
			t.Errorf("expected role=assistant, got %v", msg.Role)
		}
		if msg.Content.Text != "hello" {
			t.Errorf("expected message text='hello', got %q", msg.Content.Text)
		}
	})

	t.Run("with tool calls", func(t *testing.T) {
		resp := convertProviderResponse(provider.Response{
			ToolCalls: []provider.ToolCall{
				{ID: "tc1", Name: "search", Arguments: `{"q":"go"}`},
				{ID: "tc2", Name: "calc", Arguments: `{"a":1}`},
			},
		})
		if !resp.HasToolCalls() {
			t.Fatal("expected tool calls")
		}
		if len(resp.ToolCalls) != 2 {
			t.Fatalf("expected 2 tool calls, got %d", len(resp.ToolCalls))
		}
		if resp.ToolCalls[0].ID != "tc1" || resp.ToolCalls[0].Name != "search" {
			t.Errorf("unexpected tool call[0]: %+v", resp.ToolCalls[0])
		}
		if resp.ToolCalls[1].ID != "tc2" || resp.ToolCalls[1].Name != "calc" {
			t.Errorf("unexpected tool call[1]: %+v", resp.ToolCalls[1])
		}

		msg := resp.Message()
		if len(msg.ToolCalls) != 2 {
			t.Errorf("expected 2 tool calls in message, got %d", len(msg.ToolCalls))
		}
	})

	t.Run("nil usage", func(t *testing.T) {
		resp := convertProviderResponse(provider.Response{Text: "ok"})
		if resp.Usage != nil {
			t.Errorf("expected nil usage, got %+v", resp.Usage)
		}
	})
}