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

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