go-llm/v2/agent/agent_test.go

package agent

import (
	"context"
	"errors"
	"sync"
	"testing"

	llm "gitea.stevedudenhoeffer.com/steve/go-llm/v2"
	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/provider"
)

// mockProvider is a test helper that implements provider.Provider.
type mockProvider struct {
	mu           sync.Mutex
	completeFunc func(ctx context.Context, req provider.Request) (provider.Response, error)
	requests     []provider.Request
}

func (m *mockProvider) Complete(ctx context.Context, req provider.Request) (provider.Response, error) {
	m.mu.Lock()
	m.requests = append(m.requests, req)
	m.mu.Unlock()
	return m.completeFunc(ctx, req)
}

func (m *mockProvider) Stream(ctx context.Context, req provider.Request, events chan<- provider.StreamEvent) error {
	close(events)
	return nil
}

func (m *mockProvider) lastRequest() provider.Request {
	m.mu.Lock()
	defer m.mu.Unlock()
	if len(m.requests) == 0 {
		return provider.Request{}
	}
	return m.requests[len(m.requests)-1]
}

func newMockModel(fn func(ctx context.Context, req provider.Request) (provider.Response, error)) *llm.Model {
	mp := &mockProvider{completeFunc: fn}
	return llm.NewClient(mp).Model("mock-model")
}

func newSimpleMockModel(text string) *llm.Model {
	return newMockModel(func(ctx context.Context, req provider.Request) (provider.Response, error) {
		return provider.Response{Text: text}, nil
	})
}

func TestAgent_Run(t *testing.T) {
	model := newSimpleMockModel("Hello from agent!")
	a := New(model, "You are a helpful assistant.")

	result, err := a.Run(context.Background(), "Say hello")
	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_Run_WithTools(t *testing.T) {
	callCount := 0
	model := newMockModel(func(ctx context.Context, req provider.Request) (provider.Response, error) {
		callCount++
		if callCount == 1 {
			// First call: model requests a tool call
			return provider.Response{
				ToolCalls: []provider.ToolCall{
					{ID: "tc1", Name: "greet", Arguments: `{}`},
				},
			}, nil
		}
		// Second call: model returns text after seeing tool result
		return provider.Response{Text: "Tool said: hello!"}, nil
	})

	tool := llm.DefineSimple("greet", "Says hello", func(ctx context.Context) (string, error) {
		return "hello!", nil
	})

	a := New(model, "You are helpful.", WithTools(llm.NewToolBox(tool)))
	result, err := a.Run(context.Background(), "Use the greet tool")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result != "Tool said: hello!" {
		t.Errorf("expected 'Tool said: hello!', got %q", result)
	}
	if callCount != 2 {
		t.Errorf("expected 2 calls (tool loop), got %d", callCount)
	}
}

func TestAgent_AsTool(t *testing.T) {
	// Create a child agent
	childModel := newSimpleMockModel("child result: 42")
	child := New(childModel, "You compute things.")

	// Create the tool from the child agent
	childTool := child.AsTool("compute", "Delegate computation to child agent")

	// Verify tool metadata
	if childTool.Name != "compute" {
		t.Errorf("expected tool name 'compute', got %q", childTool.Name)
	}
	if childTool.Description != "Delegate computation to child agent" {
		t.Errorf("expected correct description, got %q", childTool.Description)
	}

	// Execute the tool directly (simulating what the parent's Chat.Send loop does)
	result, err := childTool.Execute(context.Background(), `{"input":"what is 6*7?"}`)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result != "child result: 42" {
		t.Errorf("expected 'child result: 42', got %q", result)
	}
}

func TestAgent_AsTool_ParentChild(t *testing.T) {
	// Child agent that always returns a fixed result
	childModel := newSimpleMockModel("researched: Go generics are great")
	child := New(childModel, "You are a researcher.")

	// Parent agent: first call returns tool call, second returns text
	parentCallCount := 0
	parentModel := newMockModel(func(ctx context.Context, req provider.Request) (provider.Response, error) {
		parentCallCount++
		if parentCallCount == 1 {
			return provider.Response{
				ToolCalls: []provider.ToolCall{
					{ID: "tc1", Name: "research", Arguments: `{"input":"Tell me about Go generics"}`},
				},
			}, nil
		}
		// After getting tool result, parent synthesizes final answer
		return provider.Response{Text: "Based on research: Go generics are great"}, nil
	})

	parent := New(parentModel, "You coordinate tasks.",
		WithTools(llm.NewToolBox(
			child.AsTool("research", "Research a topic"),
		)),
	)

	result, err := parent.Run(context.Background(), "Tell me about Go generics")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result != "Based on research: Go generics are great" {
		t.Errorf("expected synthesized result, got %q", result)
	}
	if parentCallCount != 2 {
		t.Errorf("expected 2 parent calls (tool loop), got %d", parentCallCount)
	}
}

func TestAgent_RunMessages(t *testing.T) {
	model := newSimpleMockModel("I see the system and user messages")
	a := New(model, "You are helpful.")

	messages := []llm.Message{
		llm.UserMessage("First question"),
		llm.AssistantMessage("First answer"),
		llm.UserMessage("Follow up"),
	}

	result, err := a.RunMessages(context.Background(), messages)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result != "I see the system and user messages" {
		t.Errorf("unexpected result: %q", result)
	}
}

func TestAgent_ContextCancellation(t *testing.T) {
	model := newMockModel(func(ctx context.Context, req provider.Request) (provider.Response, error) {
		return provider.Response{}, ctx.Err()
	})
	a := New(model, "You are helpful.")

	ctx, cancel := context.WithCancel(context.Background())
	cancel() // Cancel immediately

	_, err := a.Run(ctx, "This should fail")
	if err == nil {
		t.Fatal("expected error from cancelled context")
	}
}

func TestAgent_WithRequestOptions(t *testing.T) {
	var capturedReq provider.Request
	model := newMockModel(func(ctx context.Context, req provider.Request) (provider.Response, error) {
		capturedReq = req
		return provider.Response{Text: "ok"}, nil
	})

	a := New(model, "You are helpful.",
		WithRequestOptions(llm.WithTemperature(0.3), llm.WithMaxTokens(100)),
	)

	_, err := a.Run(context.Background(), "test")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if capturedReq.Temperature == nil || *capturedReq.Temperature != 0.3 {
		t.Errorf("expected temperature 0.3, got %v", capturedReq.Temperature)
	}
	if capturedReq.MaxTokens == nil || *capturedReq.MaxTokens != 100 {
		t.Errorf("expected maxTokens 100, got %v", capturedReq.MaxTokens)
	}
}

func TestAgent_Run_Error(t *testing.T) {
	wantErr := errors.New("model failed")
	model := newMockModel(func(ctx context.Context, req provider.Request) (provider.Response, error) {
		return provider.Response{}, wantErr
	})
	a := New(model, "You are helpful.")

	_, err := a.Run(context.Background(), "test")
	if err == nil {
		t.Fatal("expected error, got nil")
	}
}

func TestAgent_EmptySystem(t *testing.T) {
	model := newSimpleMockModel("no system prompt")
	a := New(model, "") // Empty system prompt

	result, err := a.Run(context.Background(), "test")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result != "no system prompt" {
		t.Errorf("unexpected result: %q", result)
	}
}