package llm import ( "testing" "github.com/tmc/langchaingo/llms" ) func TestNewMessageBuilder(t *testing.T) { b := NewMessageBuilder() if b == nil { t.Fatal("NewMessageBuilder() returned nil") } if len(b.messages) != 0 { t.Fatalf("expected empty messages, got %d", len(b.messages)) } } func TestMessageBuilder_WithSystemPrompt(t *testing.T) { b := NewMessageBuilder().WithSystemPrompt("you are a helper") msgs := b.Build() if len(msgs) != 1 { t.Fatalf("expected 1 message, got %d", len(msgs)) } if msgs[0].Role != llms.ChatMessageTypeSystem { t.Errorf("expected role %v, got %v", llms.ChatMessageTypeSystem, msgs[0].Role) } } func TestMessageBuilder_WithMessages(t *testing.T) { input := []llms.MessageContent{ llms.TextParts(llms.ChatMessageTypeHuman, "hello"), llms.TextParts(llms.ChatMessageTypeAI, "hi"), } b := NewMessageBuilder().WithMessages(input) msgs := b.Build() if len(msgs) != 2 { t.Fatalf("expected 2 messages, got %d", len(msgs)) } if msgs[0].Role != llms.ChatMessageTypeHuman { t.Errorf("expected role %v for msg[0], got %v", llms.ChatMessageTypeHuman, msgs[0].Role) } if msgs[1].Role != llms.ChatMessageTypeAI { t.Errorf("expected role %v for msg[1], got %v", llms.ChatMessageTypeAI, msgs[1].Role) } } func TestMessageBuilder_AppendAIMessage(t *testing.T) { tc := llms.ToolCall{ ID: "call_1", FunctionCall: &llms.FunctionCall{ Name: "search_nodes", Arguments: `{"query":"AI"}`, }, } b := NewMessageBuilder().AppendAIMessage("thinking...", []llms.ToolCall{tc}) msgs := b.Build() if len(msgs) != 1 { t.Fatalf("expected 1 message, got %d", len(msgs)) } if msgs[0].Role != llms.ChatMessageTypeAI { t.Errorf("expected role %v, got %v", llms.ChatMessageTypeAI, msgs[0].Role) } if len(msgs[0].Parts) != 2 { t.Fatalf("expected 2 parts (text + toolCall), got %d", len(msgs[0].Parts)) } } func TestMessageBuilder_AppendAIMessage_EmptyContent(t *testing.T) { tc := llms.ToolCall{ ID: "call_2", FunctionCall: &llms.FunctionCall{ Name: "get_neighbors", Arguments: `{"node_id":"1"}`, }, } b := NewMessageBuilder().AppendAIMessage("", []llms.ToolCall{tc}) msgs := b.Build() if len(msgs) != 1 { t.Fatalf("expected 1 message, got %d", len(msgs)) } if len(msgs[0].Parts) != 1 { t.Fatalf("expected 1 part (only toolCall, no text), got %d", len(msgs[0].Parts)) } } func TestMessageBuilder_AppendToolResult(t *testing.T) { b := NewMessageBuilder().AppendToolResult("call_1", "search_nodes", `{"nodes":[]}`) msgs := b.Build() if len(msgs) != 1 { t.Fatalf("expected 1 message, got %d", len(msgs)) } if msgs[0].Role != llms.ChatMessageTypeTool { t.Errorf("expected role %v, got %v", llms.ChatMessageTypeTool, msgs[0].Role) } } func TestMessageBuilder_Chaining(t *testing.T) { b := NewMessageBuilder(). WithSystemPrompt("system"). WithMessages([]llms.MessageContent{ llms.TextParts(llms.ChatMessageTypeHuman, "hello"), }). AppendAIMessage("hi", nil). AppendToolResult("call_1", "search_nodes", "{}") msgs := b.Build() if len(msgs) != 4 { t.Fatalf("expected 4 messages, got %d", len(msgs)) } if msgs[0].Role != llms.ChatMessageTypeSystem { t.Errorf("msg[0] expected system, got %v", msgs[0].Role) } if msgs[1].Role != llms.ChatMessageTypeHuman { t.Errorf("msg[1] expected human, got %v", msgs[1].Role) } if msgs[2].Role != llms.ChatMessageTypeAI { t.Errorf("msg[2] expected AI, got %v", msgs[2].Role) } if msgs[3].Role != llms.ChatMessageTypeTool { t.Errorf("msg[3] expected tool, got %v", msgs[3].Role) } } func TestMessageBuilder_Build_ReturnsCopy(t *testing.T) { b := NewMessageBuilder().WithSystemPrompt("system") first := b.Build() second := b.Build() if len(first) != len(second) { t.Fatal("Build() should return consistent results") } if &first[0] == &second[0] { t.Error("Build() should return a copy, not the same slice") } } func TestMessageBuilder_Build_DoesNotMutate(t *testing.T) { b := NewMessageBuilder().WithSystemPrompt("original") msgs := b.Build() msgs[0] = llms.TextParts(llms.ChatMessageTypeHuman, "mutated") original := b.Build() if original[0].Role == llms.ChatMessageTypeHuman { t.Error("modifying Build() result should not affect the builder") } }