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,264 @@
+package llm
+
+import (
+	"context"
+	"fmt"
+	"knowledge-graph-backend/internal/model"
+	"testing"
+
+	"github.com/tmc/langchaingo/llms"
+)
+
+type mockModel struct {
+	generateContent func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error)
+}
+
+func (m *mockModel) Call(ctx context.Context, prompt string, opts ...llms.CallOption) (string, error) {
+	return "", nil
+}
+
+func (m *mockModel) GenerateContent(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+	if m.generateContent != nil {
+		return m.generateContent(ctx, messages, opts...)
+	}
+	return nil, nil
+}
+
+func TestAgent_ChatWithTools_DirectResponse(t *testing.T) {
+	mock := &mockModel{
+		generateContent: func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+			return &llms.ContentResponse{
+				Choices: []*llms.ContentChoice{
+					{Content: "hello from agent"},
+				},
+			}, nil
+		},
+	}
+
+	client := &Client{model: mock}
+	executor := NewToolExecutor(&mockGraphService{})
+	agent := NewAgent(client, executor)
+
+	result, err := agent.ChatWithTools(context.Background(), "system", []llms.MessageContent{
+		llms.TextParts(llms.ChatMessageTypeHuman, "hi"),
+	})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result != "hello from agent" {
+		t.Errorf("expected 'hello from agent', got %q", result)
+	}
+}
+
+func TestAgent_ChatWithTools_ToolCallThenResponse(t *testing.T) {
+	callCount := 0
+	mock := &mockModel{
+		generateContent: func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+			callCount++
+			if callCount == 1 {
+				return &llms.ContentResponse{
+					Choices: []*llms.ContentChoice{
+						{
+							ToolCalls: []llms.ToolCall{
+								{
+									ID: "call_1",
+									FunctionCall: &llms.FunctionCall{
+										Name:      "search_nodes",
+										Arguments: `{"query":"AI"}`,
+									},
+								},
+							},
+						},
+					},
+				}, nil
+			}
+			return &llms.ContentResponse{
+				Choices: []*llms.ContentChoice{
+					{Content: "found AI nodes"},
+				},
+			}, nil
+		},
+	}
+
+	graphSvc := &mockGraphService{
+		searchNodes: func(query string) []model.Node {
+			return []model.Node{{ID: "1", Label: "AI"}}
+		},
+	}
+
+	client := &Client{model: mock}
+	executor := NewToolExecutor(graphSvc)
+	agent := NewAgent(client, executor)
+
+	result, err := agent.ChatWithTools(context.Background(), "system", []llms.MessageContent{
+		llms.TextParts(llms.ChatMessageTypeHuman, "search for AI"),
+	})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result != "found AI nodes" {
+		t.Errorf("expected 'found AI nodes', got %q", result)
+	}
+	if callCount != 2 {
+		t.Errorf("expected 2 GenerateContent calls, got %d", callCount)
+	}
+}
+
+func TestAgent_ChatWithTools_EmptyResponse(t *testing.T) {
+	mock := &mockModel{
+		generateContent: func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+			return &llms.ContentResponse{Choices: []*llms.ContentChoice{}}, nil
+		},
+	}
+
+	client := &Client{model: mock}
+	executor := NewToolExecutor(&mockGraphService{})
+	agent := NewAgent(client, executor)
+
+	_, err := agent.ChatWithTools(context.Background(), "system", nil)
+	if err == nil {
+		t.Fatal("expected error for empty response, got nil")
+	}
+}
+
+func TestAgent_ChatWithTools_NilResponse(t *testing.T) {
+	mock := &mockModel{
+		generateContent: func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+			return nil, nil
+		},
+	}
+
+	client := &Client{model: mock}
+	executor := NewToolExecutor(&mockGraphService{})
+	agent := NewAgent(client, executor)
+
+	_, err := agent.ChatWithTools(context.Background(), "system", nil)
+	if err == nil {
+		t.Fatal("expected error for nil response, got nil")
+	}
+}
+
+func TestAgent_ChatWithTools_GenerateError(t *testing.T) {
+	mock := &mockModel{
+		generateContent: func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+			return nil, fmt.Errorf("network error")
+		},
+	}
+
+	client := &Client{model: mock}
+	executor := NewToolExecutor(&mockGraphService{})
+	agent := NewAgent(client, executor)
+
+	_, err := agent.ChatWithTools(context.Background(), "system", nil)
+	if err == nil {
+		t.Fatal("expected error, got nil")
+	}
+}
+
+func TestAgent_ChatWithTools_ToolExecutionError_ContinuesLoop(t *testing.T) {
+	callCount := 0
+	mock := &mockModel{
+		generateContent: func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+			callCount++
+			if callCount == 1 {
+				return &llms.ContentResponse{
+					Choices: []*llms.ContentChoice{
+						{
+							ToolCalls: []llms.ToolCall{
+								{
+									ID: "call_1",
+									FunctionCall: &llms.FunctionCall{
+										Name:      "delete_node",
+										Arguments: `{"id":"nonexistent"}`,
+									},
+								},
+							},
+						},
+					},
+				}, nil
+			}
+			return &llms.ContentResponse{
+				Choices: []*llms.ContentChoice{
+					{Content: "handled error"},
+				},
+			}, nil
+		},
+	}
+
+	graphSvc := &mockGraphService{
+		deleteNode: func(id string) error {
+			return fmt.Errorf("node not found")
+		},
+	}
+
+	client := &Client{model: mock}
+	executor := NewToolExecutor(graphSvc)
+	agent := NewAgent(client, executor)
+
+	result, err := agent.ChatWithTools(context.Background(), "system", []llms.MessageContent{
+		llms.TextParts(llms.ChatMessageTypeHuman, "delete nonexistent"),
+	})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result != "handled error" {
+		t.Errorf("expected 'handled error', got %q", result)
+	}
+}
+
+func TestAgent_ChatWithTools_MaxIterations(t *testing.T) {
+	mock := &mockModel{
+		generateContent: func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+			return &llms.ContentResponse{
+				Choices: []*llms.ContentChoice{
+					{
+						ToolCalls: []llms.ToolCall{
+							{
+								ID: "call_loop",
+								FunctionCall: &llms.FunctionCall{
+									Name:      "search_nodes",
+									Arguments: `{"query":"loop"}`,
+								},
+							},
+						},
+					},
+				},
+			}, nil
+		},
+	}
+
+	graphSvc := &mockGraphService{
+		searchNodes: func(query string) []model.Node {
+			return []model.Node{}
+		},
+	}
+
+	client := &Client{model: mock}
+	executor := NewToolExecutor(graphSvc)
+	agent := NewAgent(client, executor)
+
+	_, err := agent.ChatWithTools(context.Background(), "system", nil)
+	if err == nil {
+		t.Fatal("expected max iterations error, got nil")
+	}
+}
+
+func TestAgent_ChatWithTools_ContextCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	mock := &mockModel{
+		generateContent: func(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
+			return nil, ctx.Err()
+		},
+	}
+
+	client := &Client{model: mock}
+	executor := NewToolExecutor(&mockGraphService{})
+	agent := NewAgent(client, executor)
+
+	_, err := agent.ChatWithTools(ctx, "system", nil)
+	if err == nil {
+		t.Fatal("expected error due to cancelled context, got nil")
+	}
+}