loveuer
343 lines
9.8 KiB
Go
343 lines
9.8 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"testing"
|
|
|
|
"agentu/internal/llm"
|
|
"agentu/internal/tools"
|
|
)
|
|
|
|
type fakeProvider struct {
|
|
calls int
|
|
t *testing.T
|
|
}
|
|
|
|
func (p *fakeProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
|
p.calls++
|
|
switch p.calls {
|
|
case 1:
|
|
if len(req.Tools) == 0 {
|
|
p.t.Fatal("expected tools")
|
|
}
|
|
if req.ToolChoice != "auto" {
|
|
p.t.Fatalf("tool_choice = %q", req.ToolChoice)
|
|
}
|
|
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 0, ID: "call_1", Type: "function", Name: "test_echo", Arguments: `{"text":"he`},
|
|
{Index: 0, Arguments: `llo"}`},
|
|
}})
|
|
case 2:
|
|
last := req.Messages[len(req.Messages)-1]
|
|
if last.Role != llm.RoleTool || last.Content != "hello" {
|
|
p.t.Fatalf("last message = %#v", last)
|
|
}
|
|
return emit(llm.StreamEvent{Content: "done"})
|
|
default:
|
|
p.t.Fatalf("unexpected call count %d", p.calls)
|
|
return nil
|
|
}
|
|
}
|
|
|
|
type echoTool struct{}
|
|
|
|
func (echoTool) Definition() llm.Tool {
|
|
return llm.Tool{
|
|
Type: "function",
|
|
Function: llm.ToolFunction{
|
|
Name: "test_echo",
|
|
Description: "echo test",
|
|
Parameters: json.RawMessage(`{"type":"object"}`),
|
|
},
|
|
}
|
|
}
|
|
|
|
func (echoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
|
var args struct {
|
|
Text string `json:"text"`
|
|
}
|
|
if err := json.Unmarshal(raw, &args); err != nil {
|
|
return "", err
|
|
}
|
|
return args.Text, nil
|
|
}
|
|
|
|
func TestAgentRunsToolLoop(t *testing.T) {
|
|
provider := &fakeProvider{t: t}
|
|
a := New(Options{
|
|
Provider: provider,
|
|
Model: "test",
|
|
SystemPrompt: "system",
|
|
ToolRegistry: tools.NewRegistry(echoTool{}),
|
|
})
|
|
|
|
var out strings.Builder
|
|
var logs strings.Builder
|
|
if err := a.RunTurn(context.Background(), "use a tool", &out, &logs); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if out.String() != "done" {
|
|
t.Fatalf("out = %q", out.String())
|
|
}
|
|
if !strings.Contains(logs.String(), "[tool] test_echo") {
|
|
t.Fatalf("logs = %q", logs.String())
|
|
}
|
|
}
|
|
|
|
type contentProvider struct {
|
|
called bool
|
|
text string
|
|
}
|
|
|
|
func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
|
p.called = true
|
|
return emit(llm.StreamEvent{Content: p.text})
|
|
}
|
|
|
|
type compactProvider struct {
|
|
summaryCalls int
|
|
answerCalls int
|
|
lastSummary string
|
|
}
|
|
|
|
func (p *compactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
|
if len(req.Messages) == 2 && req.Messages[0].Role == llm.RoleSystem && req.Messages[0].Content == summarizerSystemPrompt {
|
|
p.summaryCalls++
|
|
p.lastSummary = req.Messages[1].Content
|
|
if err := emit(llm.StreamEvent{Content: "old work summarized"}); err != nil {
|
|
return err
|
|
}
|
|
return emit(llm.StreamEvent{Usage: &llm.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}})
|
|
}
|
|
|
|
p.answerCalls++
|
|
return emit(llm.StreamEvent{Content: "done"})
|
|
}
|
|
|
|
func TestAgentCompact(t *testing.T) {
|
|
provider := &compactProvider{}
|
|
a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
|
|
a.SetMessages(longConversation())
|
|
|
|
if err := a.Compact(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
messages := a.Messages()
|
|
if len(messages) != 8 {
|
|
t.Fatalf("messages = %d, want 8", len(messages))
|
|
}
|
|
if messages[0].Role != llm.RoleSystem || messages[0].Content != "system" {
|
|
t.Fatalf("system message = %#v", messages[0])
|
|
}
|
|
if messages[1].Role != llm.RoleAssistant || !strings.HasPrefix(messages[1].Content, contextSummaryPrefix) {
|
|
t.Fatalf("summary message = %#v", messages[1])
|
|
}
|
|
if !strings.Contains(messages[1].Content, "old work summarized") {
|
|
t.Fatalf("summary content = %q", messages[1].Content)
|
|
}
|
|
if !strings.Contains(provider.lastSummary, "old-user-0") {
|
|
t.Fatalf("summary input did not include old history: %q", provider.lastSummary)
|
|
}
|
|
if messages[len(messages)-1].Content != "recent-answer-2" {
|
|
t.Fatalf("last recent message = %#v", messages[len(messages)-1])
|
|
}
|
|
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 12 {
|
|
t.Fatalf("usage = %#v", usage)
|
|
}
|
|
}
|
|
|
|
func TestAgentAutoCompact(t *testing.T) {
|
|
provider := &compactProvider{}
|
|
a := New(Options{
|
|
Provider: provider,
|
|
Model: "test",
|
|
SystemPrompt: "system",
|
|
MaxContextTokens: 1,
|
|
ToolRegistry: tools.NewRegistry(),
|
|
MaxToolIterations: 1,
|
|
})
|
|
a.SetMessages(longConversation())
|
|
|
|
var out strings.Builder
|
|
if err := a.RunTurn(context.Background(), "new question", &out, nil); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if provider.summaryCalls != 1 {
|
|
t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
|
|
}
|
|
if provider.answerCalls != 1 {
|
|
t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
|
|
}
|
|
if out.String() != "done" {
|
|
t.Fatalf("out = %q", out.String())
|
|
}
|
|
if !strings.Contains(joinMessageContents(a.Messages()), contextSummaryPrefix) {
|
|
t.Fatalf("messages missing summary: %#v", a.Messages())
|
|
}
|
|
}
|
|
|
|
func TestAgentCompactSkipsShortHistory(t *testing.T) {
|
|
provider := &contentProvider{text: "unused"}
|
|
a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
|
|
a.SetMessages([]llm.Message{
|
|
{Role: llm.RoleSystem, Content: "system"},
|
|
{Role: llm.RoleUser, Content: "question"},
|
|
{Role: llm.RoleAssistant, Content: "answer"},
|
|
})
|
|
|
|
if err := a.Compact(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if provider.called {
|
|
t.Fatal("provider should not be called for short history")
|
|
}
|
|
if len(a.Messages()) != 3 {
|
|
t.Fatalf("messages = %d, want 3", len(a.Messages()))
|
|
}
|
|
}
|
|
|
|
func longConversation() []llm.Message {
|
|
messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}}
|
|
for i := 0; i < 3; i++ {
|
|
messages = append(messages,
|
|
llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("old-user-%d", i)},
|
|
llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("old-answer-%d", i)},
|
|
)
|
|
}
|
|
for i := 0; i < 3; i++ {
|
|
messages = append(messages,
|
|
llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("recent-user-%d", i)},
|
|
llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("recent-answer-%d", i)},
|
|
)
|
|
}
|
|
return messages
|
|
}
|
|
|
|
func joinMessageContents(messages []llm.Message) string {
|
|
var b strings.Builder
|
|
for _, msg := range messages {
|
|
b.WriteString(msg.Content)
|
|
b.WriteByte('\n')
|
|
}
|
|
return b.String()
|
|
}
|
|
func TestAgentLetsProviderDecideShellRequestWhenShellDisabled(t *testing.T) {
|
|
provider := &contentProvider{text: "shell is disabled"}
|
|
a := New(Options{
|
|
Provider: provider,
|
|
Model: "test",
|
|
SystemPrompt: "system",
|
|
ToolRegistry: tools.ReadOnlyBuiltins(t.TempDir()),
|
|
})
|
|
|
|
var out strings.Builder
|
|
if err := a.RunTurn(context.Background(), "能帮我看看 `ssh home-dev` 的运行情况么", &out, nil); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !provider.called {
|
|
t.Fatal("provider should be called so the model can decide whether to call a tool or answer")
|
|
}
|
|
if out.String() != "shell is disabled" {
|
|
t.Fatalf("out = %q", out.String())
|
|
}
|
|
}
|
|
|
|
func TestAgentDoesNotBlockFileRequest(t *testing.T) {
|
|
provider := &contentProvider{text: "ok"}
|
|
a := New(Options{
|
|
Provider: provider,
|
|
Model: "test",
|
|
SystemPrompt: "system",
|
|
ToolRegistry: tools.ReadOnlyBuiltins(t.TempDir()),
|
|
})
|
|
|
|
var out strings.Builder
|
|
if err := a.RunTurn(context.Background(), "请读取 `README.md`", &out, nil); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !provider.called {
|
|
t.Fatal("provider should be called for file requests")
|
|
}
|
|
if out.String() != "ok" {
|
|
t.Fatalf("out = %q", out.String())
|
|
}
|
|
}
|
|
|
|
// alwaysToolProvider always returns a tool call with the same args on every request.
|
|
type alwaysToolProvider struct {
|
|
calls int
|
|
}
|
|
|
|
func (p *alwaysToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
|
p.calls++
|
|
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo", Arguments: `{"text":"loop"}`},
|
|
}})
|
|
}
|
|
|
|
func TestAgentStopsOnDoomLoop(t *testing.T) {
|
|
provider := &alwaysToolProvider{}
|
|
a := New(Options{
|
|
Provider: provider,
|
|
Model: "test",
|
|
SystemPrompt: "system",
|
|
ToolRegistry: tools.NewRegistry(echoTool{}),
|
|
MaxToolIterations: 50, // High hard limit — doom loop should trigger first
|
|
})
|
|
|
|
var out strings.Builder
|
|
err := a.RunTurn(context.Background(), "loop forever", &out, nil)
|
|
if err == nil {
|
|
t.Fatal("expected error when doom loop is detected")
|
|
}
|
|
if !strings.Contains(err.Error(), "doom loop detected") {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
// Should have stopped at doomLoopThreshold (3) rounds, not reached hard limit.
|
|
if provider.calls != doomLoopThreshold {
|
|
t.Fatalf("calls = %d, want %d (doom loop threshold)", provider.calls, doomLoopThreshold)
|
|
}
|
|
}
|
|
|
|
// varyingToolProvider alternates between different tool calls to avoid doom loop detection.
|
|
type varyingToolProvider struct {
|
|
calls int
|
|
}
|
|
|
|
func (p *varyingToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
|
p.calls++
|
|
// Vary the arguments each round to avoid doom loop detection.
|
|
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo",
|
|
Arguments: fmt.Sprintf(`{"text":"round-%d"}`, p.calls)},
|
|
}})
|
|
}
|
|
|
|
func TestAgentHardLimitStopsVaryingCalls(t *testing.T) {
|
|
provider := &varyingToolProvider{}
|
|
const limit = 5
|
|
a := New(Options{
|
|
Provider: provider,
|
|
Model: "test",
|
|
SystemPrompt: "system",
|
|
ToolRegistry: tools.NewRegistry(echoTool{}),
|
|
MaxToolIterations: limit,
|
|
})
|
|
|
|
var out strings.Builder
|
|
err := a.RunTurn(context.Background(), "vary forever", &out, nil)
|
|
if err == nil {
|
|
t.Fatal("expected error when hard limit is reached")
|
|
}
|
|
if !strings.Contains(err.Error(), "tool round limit reached") {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if provider.calls != limit {
|
|
t.Fatalf("calls = %d, want %d", provider.calls, limit)
|
|
}
|
|
}
|