refactor: 重构llm包

2026-04-20 15:34:40 +08:00
parent fe87f5ca80
commit 181edf418e
10 changed files with 1181 additions and 299 deletions
@@ -0,0 +1,155 @@
+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")
+	}
+}