knowledge-graph-agent/backend/internal/llm/context_test.go

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")
	}
}