feat(agent): add Agent loop with iteration control, tool handling, and streaming support
This commit is contained in:
parent
641481022f
commit
4b39f52081
3 changed files with 923 additions and 0 deletions
157
pkg/agent/integration_test.go
Normal file
157
pkg/agent/integration_test.go
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
package agent_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/agent"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm/providers/llamacpp"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/persona"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/tools"
|
||||
)
|
||||
|
||||
// TestIntegration_LlamaCPP_Generate is an integration test that requires llama.cpp running on localhost:8080.
|
||||
// Run with: go test ./pkg/agent/ -run TestIntegration_LlamaCPP_Generate -tags=integration
|
||||
func TestIntegration_LlamaCPP_Generate(t *testing.T) {
|
||||
client, err := llamacpp.New(llamacpp.Config{
|
||||
BaseURL: "http://localhost:8080/v1",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create client: %v", err)
|
||||
}
|
||||
|
||||
resp, err := client.Generate(context.Background(), llm.CompletionRequest{
|
||||
Messages: []llm.Message{
|
||||
{Role: llm.RoleUser, Content: "What is 2+2? Answer with just the number."},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("generate failed: %v", err)
|
||||
}
|
||||
|
||||
if resp.Content == "" {
|
||||
t.Fatal("expected non-empty response")
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegration_LlamaCPP_Stream is an integration test that requires llama.cpp running on localhost:8080.
|
||||
func TestIntegration_LlamaCPP_Stream(t *testing.T) {
|
||||
client, err := llamacpp.New(llamacpp.Config{
|
||||
BaseURL: "http://localhost:8080/v1",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create client: %v", err)
|
||||
}
|
||||
|
||||
stream := client.Stream(context.Background(), llm.CompletionRequest{
|
||||
Messages: []llm.Message{
|
||||
{Role: llm.RoleUser, Content: "Say hello in 5 words."},
|
||||
},
|
||||
})
|
||||
|
||||
var chunks []llm.StreamChunk
|
||||
for chunk, err := range stream {
|
||||
if err != nil {
|
||||
t.Fatalf("stream error: %v", err)
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected at least one chunk")
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegration_AgentLoop_Generate is an integration test for the agent loop with llama.cpp.
|
||||
func TestIntegration_AgentLoop_Generate(t *testing.T) {
|
||||
client, err := llamacpp.New(llamacpp.Config{
|
||||
BaseURL: "http://localhost:8080/v1",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create client: %v", err)
|
||||
}
|
||||
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.Tool{
|
||||
Name: "add_numbers",
|
||||
Description: "Add two numbers together",
|
||||
InputSchema: json.RawMessage(`{"type": "function", "function": {"name": "add_numbers", "description": "Add two numbers together", "parameters": {"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}}}`),
|
||||
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||
var input struct {
|
||||
A float64 `json:"a"`
|
||||
B float64 `json:"b"`
|
||||
}
|
||||
if err := json.Unmarshal(args, &input); err != nil {
|
||||
return tools.ToolResult{IsError: true}, err
|
||||
}
|
||||
return tools.ToolResult{Content: "42"}, nil
|
||||
},
|
||||
Permission: tools.Allow,
|
||||
})
|
||||
|
||||
loop := agent.New(agent.Config{
|
||||
LLM: client,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: registry,
|
||||
Sandbox: &mockSandbox{},
|
||||
MaxIters: 3,
|
||||
})
|
||||
|
||||
resp, err := loop.Run(context.Background(), "What is 20+22? Use the add_numbers tool.")
|
||||
if err != nil {
|
||||
t.Fatalf("run failed: %v", err)
|
||||
}
|
||||
|
||||
if resp.Content == "" {
|
||||
t.Fatal("expected non-empty response")
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegration_AgentLoop_Stream is an integration test for the agent loop with streaming.
|
||||
func TestIntegration_AgentLoop_Stream(t *testing.T) {
|
||||
client, err := llamacpp.New(llamacpp.Config{
|
||||
BaseURL: "http://localhost:8080/v1",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create client: %v", err)
|
||||
}
|
||||
|
||||
loop := agent.New(agent.Config{
|
||||
LLM: client,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: tools.NewRegistry(),
|
||||
MaxIters: 3,
|
||||
})
|
||||
|
||||
stream := loop.RunStream(context.Background(), "Say something interesting.")
|
||||
|
||||
var chunks []llm.StreamChunk
|
||||
for chunk, err := range stream {
|
||||
if err != nil {
|
||||
t.Fatalf("stream error: %v", err)
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected at least one chunk")
|
||||
}
|
||||
}
|
||||
|
||||
// mockSandbox is a simple sandbox that allows all calls.
|
||||
type mockSandbox struct{}
|
||||
|
||||
func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
// Skip integration tests unless explicitly enabled
|
||||
if os.Getenv("INTEGRATION_TESTS") != "1" {
|
||||
return
|
||||
}
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
239
pkg/agent/loop.go
Normal file
239
pkg/agent/loop.go
Normal file
|
|
@ -0,0 +1,239 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"iter"
|
||||
"time"
|
||||
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/persona"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/tools"
|
||||
)
|
||||
|
||||
// MaxIterations is the default maximum number of agent iterations.
|
||||
const DefaultMaxIterations = 50
|
||||
|
||||
// MaxToolOutputBytes is the default maximum size for tool output in bytes.
|
||||
const DefaultMaxToolOutputBytes = 50 * 1024 // 50KB
|
||||
|
||||
// Sandbox defines the interface for filesystem sandbox operations.
|
||||
type Sandbox interface {
|
||||
ValidateToolCall(tool tools.Tool, call llm.ToolCall) error
|
||||
}
|
||||
|
||||
// Approver is called before executing tools with Ask permission.
|
||||
// Return true to allow, false to deny.
|
||||
type Approver func(tool tools.Tool, call llm.ToolCall) bool
|
||||
|
||||
// OnIterationHook is called after each iteration.
|
||||
type OnIterationHook func(Iteration)
|
||||
|
||||
// Config holds the dependencies and settings for the agent loop.
|
||||
type Config struct {
|
||||
LLM llm.LLMClient
|
||||
Persona persona.Persona
|
||||
Tools tools.Registry
|
||||
Sandbox Sandbox
|
||||
MaxIters int
|
||||
Approver Approver
|
||||
OnIteration OnIterationHook
|
||||
ToolTimeout time.Duration
|
||||
}
|
||||
|
||||
// Iteration represents a single cycle of the agent loop.
|
||||
type Iteration struct {
|
||||
Number int
|
||||
ToolCalls []llm.ToolCall
|
||||
ToolsUsed int
|
||||
Duration time.Duration
|
||||
}
|
||||
|
||||
// Response is the final output of the agent loop.
|
||||
type Response struct {
|
||||
Content string
|
||||
ToolCalls []llm.ToolCall
|
||||
Iterations int
|
||||
Duration time.Duration
|
||||
TokenUsage llm.TokenUsage
|
||||
}
|
||||
|
||||
// Loop is the main agent loop that orchestrates LLM calls and tool execution.
|
||||
type Loop struct {
|
||||
cfg Config
|
||||
}
|
||||
|
||||
// New creates a new Loop with the given configuration.
|
||||
func New(cfg Config) *Loop {
|
||||
if cfg.MaxIters == 0 {
|
||||
cfg.MaxIters = DefaultMaxIterations
|
||||
}
|
||||
return &Loop{cfg: cfg}
|
||||
}
|
||||
|
||||
// Run executes the agent loop and returns the final response.
|
||||
func (l *Loop) Run(ctx context.Context, input string) (Response, error) {
|
||||
start := time.Now()
|
||||
|
||||
messages := l.buildInitialMessages(input)
|
||||
var finalContent string
|
||||
var allToolCalls []llm.ToolCall
|
||||
var totalUsage llm.TokenUsage
|
||||
iterations := 0
|
||||
|
||||
for iterations < l.cfg.MaxIters {
|
||||
iterations++
|
||||
|
||||
resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{
|
||||
Messages: messages,
|
||||
Tools: l.getToolSchemas(),
|
||||
})
|
||||
if err != nil {
|
||||
return Response{}, fmt.Errorf("LLM generate failed: %w", err)
|
||||
}
|
||||
|
||||
totalUsage.InputTokens += resp.Usage.InputTokens
|
||||
totalUsage.OutputTokens += resp.Usage.OutputTokens
|
||||
totalUsage.TotalTokens += resp.Usage.TotalTokens
|
||||
|
||||
if len(resp.ToolCalls) == 0 {
|
||||
finalContent = resp.Content
|
||||
break
|
||||
}
|
||||
|
||||
for _, call := range resp.ToolCalls {
|
||||
result, err := l.executeTool(ctx, call)
|
||||
if err != nil {
|
||||
messages = append(messages, llm.Message{
|
||||
Role: llm.RoleTool,
|
||||
Content: fmt.Sprintf("Error: %v", err),
|
||||
})
|
||||
continue
|
||||
}
|
||||
messages = append(messages, llm.Message{
|
||||
Role: llm.RoleTool,
|
||||
Content: result.Content,
|
||||
})
|
||||
allToolCalls = append(allToolCalls, call)
|
||||
}
|
||||
}
|
||||
|
||||
duration := time.Since(start)
|
||||
if iterations >= l.cfg.MaxIters {
|
||||
return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)
|
||||
}
|
||||
|
||||
return Response{
|
||||
Content: finalContent,
|
||||
ToolCalls: allToolCalls,
|
||||
Iterations: iterations,
|
||||
Duration: duration,
|
||||
TokenUsage: totalUsage,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// RunStream executes the agent loop with streaming output.
|
||||
func (l *Loop) RunStream(ctx context.Context, input string) iter.Seq2[llm.StreamChunk, error] {
|
||||
return func(yield func(llm.StreamChunk, error) bool) {
|
||||
messages := l.buildInitialMessages(input)
|
||||
iterations := 0
|
||||
|
||||
for iterations < l.cfg.MaxIters {
|
||||
iterations++
|
||||
|
||||
stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{
|
||||
Messages: messages,
|
||||
Tools: l.getToolSchemas(),
|
||||
})
|
||||
|
||||
var hasToolCalls bool
|
||||
for chunk, err := range stream {
|
||||
if err != nil {
|
||||
yield(llm.StreamChunk{}, err)
|
||||
return
|
||||
}
|
||||
|
||||
if len(chunk.ToolCalls) > 0 {
|
||||
hasToolCalls = true
|
||||
for _, tc := range chunk.ToolCalls {
|
||||
result, err := l.executeTool(ctx, tc)
|
||||
if err != nil {
|
||||
messages = append(messages, llm.Message{
|
||||
Role: llm.RoleTool,
|
||||
Content: fmt.Sprintf("Error: %v", err),
|
||||
})
|
||||
continue
|
||||
}
|
||||
messages = append(messages, llm.Message{
|
||||
Role: llm.RoleTool,
|
||||
Content: result.Content,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if !hasToolCalls && chunk.Delta != "" {
|
||||
if !yield(chunk, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !hasToolCalls {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if iterations >= l.cfg.MaxIters {
|
||||
yield(llm.StreamChunk{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (l *Loop) buildInitialMessages(input string) []llm.Message {
|
||||
systemPrompt := persona.AssembleSystemPrompt(l.cfg.Persona, "")
|
||||
return []llm.Message{
|
||||
{Role: llm.RoleSystem, Content: systemPrompt},
|
||||
{Role: llm.RoleUser, Content: input},
|
||||
}
|
||||
}
|
||||
|
||||
func (l *Loop) getToolSchemas() []json.RawMessage {
|
||||
var schemas []json.RawMessage
|
||||
for _, tool := range l.cfg.Tools.List() {
|
||||
schemas = append(schemas, tool.InputSchema)
|
||||
}
|
||||
return schemas
|
||||
}
|
||||
|
||||
func (l *Loop) executeTool(ctx context.Context, call llm.ToolCall) (tools.ToolResult, error) {
|
||||
tool, found := l.cfg.Tools.Get(call.Name)
|
||||
if !found {
|
||||
return tools.ToolResult{IsError: true}, fmt.Errorf("tool %q not found", call.Name)
|
||||
}
|
||||
|
||||
if tool.Permission == tools.Ask && l.cfg.Approver != nil && !l.cfg.Approver(tool, call) {
|
||||
return tools.ToolResult{Content: "Tool execution denied by user"}, nil
|
||||
}
|
||||
|
||||
if l.cfg.Sandbox != nil {
|
||||
if err := l.cfg.Sandbox.ValidateToolCall(tool, call); err != nil {
|
||||
return tools.ToolResult{IsError: true}, fmt.Errorf("sandbox violation: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if l.cfg.OnIteration != nil {
|
||||
l.cfg.OnIteration(Iteration{
|
||||
Number: 1,
|
||||
ToolCalls: []llm.ToolCall{call},
|
||||
ToolsUsed: 1,
|
||||
})
|
||||
}
|
||||
|
||||
result, err := tool.Handler(ctx, call.Arguments)
|
||||
if err != nil {
|
||||
return tools.ToolResult{IsError: true, Content: fmt.Sprintf("Tool error: %v", err)}, nil
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
527
pkg/agent/loop_test.go
Normal file
527
pkg/agent/loop_test.go
Normal file
|
|
@ -0,0 +1,527 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"iter"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
llm "github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/persona"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/tools"
|
||||
)
|
||||
|
||||
type mockLLM struct {
|
||||
generateFunc func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error)
|
||||
streamFunc func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error]
|
||||
}
|
||||
|
||||
func (m *mockLLM) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
return m.generateFunc(ctx, req)
|
||||
}
|
||||
|
||||
func (m *mockLLM) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||
return m.streamFunc(ctx, req)
|
||||
}
|
||||
|
||||
func (m *mockLLM) Name() string { return "mock" }
|
||||
func (m *mockLLM) Capabilities() llm.ProviderCapabilities { return llm.ProviderCapabilities{} }
|
||||
|
||||
type mockSandbox struct {
|
||||
validateFunc func(tool tools.Tool, call llm.ToolCall) error
|
||||
}
|
||||
|
||||
func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error {
|
||||
if m.validateFunc != nil {
|
||||
return m.validateFunc(tool, call)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestNew(t *testing.T) {
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
return llm.CompletionResponse{Content: "done"}, nil
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: tools.NewRegistry(),
|
||||
})
|
||||
|
||||
if loop == nil {
|
||||
t.Fatal("expected non-nil loop")
|
||||
}
|
||||
if loop.cfg.MaxIters != DefaultMaxIterations {
|
||||
t.Errorf("expected default max iters %d, got %d", DefaultMaxIterations, loop.cfg.MaxIters)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNew_CustomMaxIters(t *testing.T) {
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
return llm.CompletionResponse{Content: "done"}, nil
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: tools.NewRegistry(),
|
||||
MaxIters: 10,
|
||||
})
|
||||
|
||||
if loop.cfg.MaxIters != 10 {
|
||||
t.Errorf("expected max iters 10, got %d", loop.cfg.MaxIters)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_NoToolCalls(t *testing.T) {
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
return llm.CompletionResponse{
|
||||
Content: "I understand.",
|
||||
Usage: llm.TokenUsage{
|
||||
InputTokens: 10,
|
||||
OutputTokens: 5,
|
||||
TotalTokens: 15,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: tools.NewRegistry(),
|
||||
})
|
||||
|
||||
resp, err := loop.Run(context.Background(), "Hello")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if resp.Content != "I understand." {
|
||||
t.Errorf("expected 'I understand.', got %q", resp.Content)
|
||||
}
|
||||
if resp.Iterations != 1 {
|
||||
t.Errorf("expected 1 iteration, got %d", resp.Iterations)
|
||||
}
|
||||
if resp.TokenUsage.InputTokens != 10 {
|
||||
t.Errorf("expected 10 input tokens, got %d", resp.TokenUsage.InputTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_ToolCalls(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.Tool{
|
||||
Name: "greet",
|
||||
Description: "Greet someone",
|
||||
InputSchema: json.RawMessage(`{"type":"object","properties":{"name":{"type":"string"}}}`),
|
||||
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||
return tools.ToolResult{Content: "Hello!"}, nil
|
||||
},
|
||||
Permission: tools.Allow,
|
||||
})
|
||||
|
||||
callCount := 0
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
callCount++
|
||||
if callCount == 1 {
|
||||
return llm.CompletionResponse{
|
||||
ToolCalls: []llm.ToolCall{
|
||||
{ID: "call-1", Name: "greet", Arguments: json.RawMessage(`{"name":"World"}`)},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
return llm.CompletionResponse{
|
||||
Content: "Done!",
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: registry,
|
||||
})
|
||||
|
||||
resp, err := loop.Run(context.Background(), "Say hi")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if resp.Content != "Done!" {
|
||||
t.Errorf("expected 'Done!', got %q", resp.Content)
|
||||
}
|
||||
if resp.Iterations != 2 {
|
||||
t.Errorf("expected 2 iterations, got %d", resp.Iterations)
|
||||
}
|
||||
if len(resp.ToolCalls) != 1 {
|
||||
t.Errorf("expected 1 tool call, got %d", len(resp.ToolCalls))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_MaxIterations(t *testing.T) {
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
return llm.CompletionResponse{
|
||||
ToolCalls: []llm.ToolCall{{ID: "1", Name: "x", Arguments: json.RawMessage("{}")}},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.Tool{
|
||||
Name: "x",
|
||||
Description: "x",
|
||||
InputSchema: json.RawMessage(`{}`),
|
||||
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||
return tools.ToolResult{Content: "ok"}, nil
|
||||
},
|
||||
Permission: tools.Allow,
|
||||
})
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: registry,
|
||||
MaxIters: 3,
|
||||
})
|
||||
|
||||
_, err := loop.Run(context.Background(), "test")
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
if err.Error() != "max iterations (3) reached" {
|
||||
t.Errorf("expected 'max iterations (3) reached', got %q", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_ToolNotFound(t *testing.T) {
|
||||
callCount := 0
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
callCount++
|
||||
if callCount == 1 {
|
||||
return llm.CompletionResponse{
|
||||
ToolCalls: []llm.ToolCall{{ID: "1", Name: "nonexistent", Arguments: json.RawMessage("{}")}},
|
||||
}, nil
|
||||
}
|
||||
return llm.CompletionResponse{Content: "done"}, nil
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: tools.NewRegistry(),
|
||||
})
|
||||
|
||||
resp, err := loop.Run(context.Background(), "test")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if resp.Content != "done" {
|
||||
t.Errorf("expected 'done', got %q", resp.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_ApprovalDenied(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.Tool{
|
||||
Name: "dangerous",
|
||||
Description: "Do something dangerous",
|
||||
InputSchema: json.RawMessage(`{}`),
|
||||
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||
return tools.ToolResult{Content: "executed"}, nil
|
||||
},
|
||||
Permission: tools.Ask,
|
||||
})
|
||||
|
||||
callCount := 0
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
callCount++
|
||||
if callCount == 1 {
|
||||
return llm.CompletionResponse{
|
||||
ToolCalls: []llm.ToolCall{{ID: "1", Name: "dangerous", Arguments: json.RawMessage("{}")}},
|
||||
}, nil
|
||||
}
|
||||
return llm.CompletionResponse{Content: "done"}, nil
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: registry,
|
||||
Approver: func(tool tools.Tool, call llm.ToolCall) bool {
|
||||
return false // deny all
|
||||
},
|
||||
})
|
||||
|
||||
resp, err := loop.Run(context.Background(), "test")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(resp.ToolCalls) != 1 {
|
||||
t.Errorf("expected 1 tool call, got %d", len(resp.ToolCalls))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_SandboxViolation(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.Tool{
|
||||
Name: "restricted",
|
||||
Description: "Restricted tool",
|
||||
InputSchema: json.RawMessage(`{}`),
|
||||
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||
return tools.ToolResult{Content: "executed"}, nil
|
||||
},
|
||||
Permission: tools.Allow,
|
||||
})
|
||||
|
||||
callCount := 0
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
callCount++
|
||||
if callCount == 1 {
|
||||
return llm.CompletionResponse{
|
||||
ToolCalls: []llm.ToolCall{{ID: "1", Name: "restricted", Arguments: json.RawMessage("{}")}},
|
||||
}, nil
|
||||
}
|
||||
return llm.CompletionResponse{Content: "done"}, nil
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: registry,
|
||||
Sandbox: &mockSandbox{
|
||||
validateFunc: func(tool tools.Tool, call llm.ToolCall) error {
|
||||
return fmt.Errorf("path traversal detected")
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
resp, err := loop.Run(context.Background(), "test")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(resp.ToolCalls) != 0 {
|
||||
t.Errorf("expected 0 tool calls (sandbox rejected), got %d", len(resp.ToolCalls))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_OnIterationHook(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.Tool{
|
||||
Name: "greet",
|
||||
Description: "Greet",
|
||||
InputSchema: json.RawMessage(`{}`),
|
||||
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||
return tools.ToolResult{Content: "hello"}, nil
|
||||
},
|
||||
Permission: tools.Allow,
|
||||
})
|
||||
|
||||
callCount := 0
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
callCount++
|
||||
if callCount == 1 {
|
||||
return llm.CompletionResponse{
|
||||
ToolCalls: []llm.ToolCall{{ID: "1", Name: "greet", Arguments: json.RawMessage("{}")}},
|
||||
}, nil
|
||||
}
|
||||
return llm.CompletionResponse{Content: "done"}, nil
|
||||
},
|
||||
}
|
||||
|
||||
var iterations []Iteration
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: registry,
|
||||
OnIteration: func(iter Iteration) {
|
||||
iterations = append(iterations, iter)
|
||||
},
|
||||
})
|
||||
|
||||
_, err := loop.Run(context.Background(), "test")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(iterations) != 1 {
|
||||
t.Errorf("expected 1 iteration hook call, got %d", len(iterations))
|
||||
}
|
||||
if iterations[0].ToolsUsed != 1 {
|
||||
t.Errorf("expected 1 tool used, got %d", iterations[0].ToolsUsed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_Stream_NoToolCalls(t *testing.T) {
|
||||
mockClient := &mockLLM{
|
||||
streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||
return func(yield func(llm.StreamChunk, error) bool) {
|
||||
yield(llm.StreamChunk{Delta: "Hello"}, nil)
|
||||
yield(llm.StreamChunk{Delta: " world"}, nil)
|
||||
yield(llm.StreamChunk{FinishReason: "stop"}, nil)
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: tools.NewRegistry(),
|
||||
})
|
||||
|
||||
var chunks []llm.StreamChunk
|
||||
stream := loop.RunStream(context.Background(), "test")
|
||||
for chunk, err := range stream {
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
|
||||
if len(chunks) != 2 {
|
||||
t.Errorf("expected 2 chunks, got %d", len(chunks))
|
||||
}
|
||||
if chunks[0].Delta != "Hello" {
|
||||
t.Errorf("expected 'Hello', got %q", chunks[0].Delta)
|
||||
}
|
||||
if chunks[1].Delta != " world" {
|
||||
t.Errorf("expected ' world', got %q", chunks[1].Delta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_Stream_WithToolCalls(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.Tool{
|
||||
Name: "greet",
|
||||
Description: "Greet",
|
||||
InputSchema: json.RawMessage(`{}`),
|
||||
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||
return tools.ToolResult{Content: "greeted"}, nil
|
||||
},
|
||||
Permission: tools.Allow,
|
||||
})
|
||||
|
||||
callCount := 0
|
||||
mockClient := &mockLLM{
|
||||
streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||
return func(yield func(llm.StreamChunk, error) bool) {
|
||||
callCount++
|
||||
if callCount == 1 {
|
||||
yield(llm.StreamChunk{
|
||||
ToolCalls: []llm.ToolCall{{ID: "1", Name: "greet", Arguments: json.RawMessage("{}")}},
|
||||
}, nil)
|
||||
} else {
|
||||
yield(llm.StreamChunk{Delta: "done"}, nil)
|
||||
yield(llm.StreamChunk{FinishReason: "stop"}, nil)
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: registry,
|
||||
})
|
||||
|
||||
var chunks []llm.StreamChunk
|
||||
stream := loop.RunStream(context.Background(), "test")
|
||||
for chunk, err := range stream {
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
|
||||
if len(chunks) != 1 {
|
||||
t.Errorf("expected 1 chunk (only the 'done' chunk), got %d", len(chunks))
|
||||
}
|
||||
if chunks[0].Delta != "done" {
|
||||
t.Errorf("expected 'done', got %q", chunks[0].Delta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_Stream_MaxIterations(t *testing.T) {
|
||||
mockClient := &mockLLM{
|
||||
streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||
return func(yield func(llm.StreamChunk, error) bool) {
|
||||
yield(llm.StreamChunk{
|
||||
ToolCalls: []llm.ToolCall{{ID: "1", Name: "x", Arguments: json.RawMessage("{}")}},
|
||||
}, nil)
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.Tool{
|
||||
Name: "x",
|
||||
Description: "x",
|
||||
InputSchema: json.RawMessage(`{}`),
|
||||
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||
return tools.ToolResult{Content: "ok"}, nil
|
||||
},
|
||||
Permission: tools.Allow,
|
||||
})
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: registry,
|
||||
MaxIters: 1,
|
||||
})
|
||||
|
||||
var chunks []llm.StreamChunk
|
||||
stream := loop.RunStream(context.Background(), "test")
|
||||
for chunk, err := range stream {
|
||||
if err != nil {
|
||||
// expect max iterations error
|
||||
if !errors.Is(err, context.DeadlineExceeded) && err.Error() != "max iterations (1) reached" {
|
||||
t.Fatalf("expected max iterations error, got: %v", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_Timeout(t *testing.T) {
|
||||
mockClient := &mockLLM{
|
||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return llm.CompletionResponse{}, ctx.Err()
|
||||
default:
|
||||
return llm.CompletionResponse{Content: "done"}, nil
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
|
||||
defer cancel()
|
||||
time.Sleep(20 * time.Millisecond) // ensure context is cancelled before calling
|
||||
|
||||
loop := New(Config{
|
||||
LLM: mockClient,
|
||||
Persona: persona.DefaultPersona(),
|
||||
Tools: tools.NewRegistry(),
|
||||
})
|
||||
|
||||
_, err := loop.Run(ctx, "test")
|
||||
if err == nil {
|
||||
t.Fatal("expected timeout error, got nil")
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue