248 lines
9.0 KiB
Go
248 lines
9.0 KiB
Go
package agentrun
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
|
|
"enkly/pkg/vllmkit"
|
|
)
|
|
|
|
func TestDefaultThinkPolicy(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
round int
|
|
retrying bool
|
|
want string
|
|
}{
|
|
{"round 0, fresh attempt: deep thinking", 0, false, "high"},
|
|
{"round 0, retry: fast and simple, not deep", 0, true, "low"},
|
|
{"later round, fresh attempt: defer to client default", 3, false, ""},
|
|
{"later round, retry: still fast and simple", 5, true, "low"},
|
|
}
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
if got := DefaultThinkPolicy(tc.round, tc.retrying); got != tc.want {
|
|
t.Errorf("DefaultThinkPolicy(%d, %v) = %q, want %q", tc.round, tc.retrying, got, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCompactThinkingHistory(t *testing.T) {
|
|
msgs := []vllmkit.Message{
|
|
{Role: vllmkit.RoleSystem, Content: "sys"},
|
|
{Role: vllmkit.RoleUser, Content: "do the thing"},
|
|
{Role: vllmkit.RoleAssistant, Content: "", Thinking: "round0 reasoning", ToolCalls: []vllmkit.ToolCallRequest{{Name: "grep"}}},
|
|
{Role: vllmkit.RoleTool, Content: "grep result", ToolName: "grep"},
|
|
{Role: vllmkit.RoleAssistant, Content: "", Thinking: "round1 reasoning", ToolCalls: []vllmkit.ToolCallRequest{{Name: "read_file"}}},
|
|
{Role: vllmkit.RoleTool, Content: "file contents", ToolName: "read_file"},
|
|
}
|
|
|
|
compactThinkingHistory(msgs)
|
|
|
|
if msgs[2].Thinking != "" {
|
|
t.Errorf("round0's thinking should have been stripped once round1's exists, got %q", msgs[2].Thinking)
|
|
}
|
|
if msgs[4].Thinking != "round1 reasoning" {
|
|
t.Errorf("most recent round's thinking should survive, got %q", msgs[4].Thinking)
|
|
}
|
|
// Non-thinking fields must be untouched.
|
|
if len(msgs[2].ToolCalls) != 1 || msgs[2].ToolCalls[0].Name != "grep" {
|
|
t.Errorf("compaction must not touch ToolCalls, got %+v", msgs[2].ToolCalls)
|
|
}
|
|
if msgs[3].Content != "grep result" {
|
|
t.Errorf("compaction must not touch tool messages, got %q", msgs[3].Content)
|
|
}
|
|
}
|
|
|
|
func TestCompactThinkingHistory_NoThinkingIsNoop(t *testing.T) {
|
|
msgs := []vllmkit.Message{
|
|
{Role: vllmkit.RoleUser, Content: "hi"},
|
|
{Role: vllmkit.RoleAssistant, Content: "hello"},
|
|
}
|
|
compactThinkingHistory(msgs)
|
|
if msgs[1].Content != "hello" {
|
|
t.Errorf("content should be untouched, got %q", msgs[1].Content)
|
|
}
|
|
}
|
|
|
|
// fakeChatStreamer scripts one vllmkit.ChatEvent sequence per call,
|
|
// in order, and records the think override and a deep-ish snapshot of
|
|
// the messages it was called with so tests can assert on both the
|
|
// adaptive-thinking policy and the thinking-history compaction working
|
|
// together across a multi-round Run.
|
|
type fakeChatStreamer struct {
|
|
t *testing.T
|
|
rounds [][]vllmkit.ChatEvent
|
|
call int
|
|
gotThink []string
|
|
gotMsgs [][]vllmkit.Message
|
|
}
|
|
|
|
func (f *fakeChatStreamer) ChatStream(ctx context.Context, messages []vllmkit.Message, tools []vllmkit.ToolDef, think string) (<-chan vllmkit.ChatEvent, error) {
|
|
if f.call >= len(f.rounds) {
|
|
f.t.Fatalf("unexpected extra ChatStream call #%d", f.call)
|
|
}
|
|
f.gotThink = append(f.gotThink, think)
|
|
snapshot := make([]vllmkit.Message, len(messages))
|
|
copy(snapshot, messages)
|
|
f.gotMsgs = append(f.gotMsgs, snapshot)
|
|
|
|
ch := make(chan vllmkit.ChatEvent, len(f.rounds[f.call])+1)
|
|
for _, ev := range f.rounds[f.call] {
|
|
ch <- ev
|
|
}
|
|
close(ch)
|
|
f.call++
|
|
return ch, nil
|
|
}
|
|
|
|
type fakeExecutor struct{}
|
|
|
|
func (fakeExecutor) Execute(ctx context.Context, call vllmkit.ToolCallRequest) vllmkit.ToolResult {
|
|
return vllmkit.ToolResult{Name: call.Name, Content: "ok: " + call.Name}
|
|
}
|
|
|
|
// TestRun_AdaptiveThinkingAndCompaction drives a 3-round turn (two
|
|
// tool-call rounds, then a final answer) and checks that: (1)
|
|
// ThinkPolicy's per-round effort actually reaches each ChatStream call,
|
|
// and (2) by the time round 2's request is built, round 0's thinking
|
|
// has been compacted away while round 1's (the most recent) survives —
|
|
// i.e. the adaptive-effort and history-compaction pieces work together
|
|
// exactly as agent.go wires them.
|
|
func TestRun_AdaptiveThinkingAndCompaction(t *testing.T) {
|
|
fake := &fakeChatStreamer{
|
|
t: t,
|
|
rounds: [][]vllmkit.ChatEvent{
|
|
{ // round 0: thinks, then calls a tool
|
|
{Thinking: "round0 reasoning"},
|
|
{ToolCalls: []vllmkit.ToolCallRequest{{Name: "grep", Arguments: map[string]any{}}}},
|
|
},
|
|
{ // round 1: thinks, then calls another tool
|
|
{Thinking: "round1 reasoning"},
|
|
{ToolCalls: []vllmkit.ToolCallRequest{{Name: "read_file", Arguments: map[string]any{}}}},
|
|
},
|
|
{ // round 2: final answer, no tool calls, no thinking
|
|
{Token: "done"},
|
|
},
|
|
},
|
|
}
|
|
|
|
agent := New(fake, fakeExecutor{}, nil)
|
|
agent.MemoryLoader = nil
|
|
|
|
var thinkingRounds []int
|
|
var thinkingTexts []string
|
|
cb := Callbacks{
|
|
OnThinkingRound: func(round int, thinking string) {
|
|
thinkingRounds = append(thinkingRounds, round)
|
|
thinkingTexts = append(thinkingTexts, thinking)
|
|
},
|
|
}
|
|
|
|
final, _, err := agent.Run(context.Background(), []vllmkit.Message{{Role: vllmkit.RoleUser, Content: "investigate X"}}, cb)
|
|
if err != nil {
|
|
t.Fatalf("Run: %v", err)
|
|
}
|
|
if final != "done" {
|
|
t.Errorf("final = %q, want %q", final, "done")
|
|
}
|
|
|
|
// ThinkPolicy effort: round 0 fresh attempt -> "high"; round 1 and 2
|
|
// (not retries) -> "" (defer to client default), per
|
|
// DefaultThinkPolicy.
|
|
wantThink := []string{"high", "", ""}
|
|
if len(fake.gotThink) != len(wantThink) {
|
|
t.Fatalf("got %d ChatStream calls, want %d (think=%v)", len(fake.gotThink), len(wantThink), fake.gotThink)
|
|
}
|
|
for i, want := range wantThink {
|
|
if fake.gotThink[i] != want {
|
|
t.Errorf("round %d: think override = %q, want %q", i, fake.gotThink[i], want)
|
|
}
|
|
}
|
|
|
|
// OnThinkingRound fired once per round that actually thought.
|
|
if len(thinkingRounds) != 2 || thinkingRounds[0] != 0 || thinkingRounds[1] != 1 {
|
|
t.Errorf("thinkingRounds = %v, want [0 1]", thinkingRounds)
|
|
}
|
|
if thinkingTexts[0] != "round0 reasoning" || thinkingTexts[1] != "round1 reasoning" {
|
|
t.Errorf("thinkingTexts = %v", thinkingTexts)
|
|
}
|
|
|
|
// History compaction: by round 2's request, round 0's assistant
|
|
// message (sent on calls 0 and 1, before round1 even existed) is
|
|
// long gone from the picture — what matters is round 2's own
|
|
// snapshot, which must carry round1's thinking (still the most
|
|
// recent) but not round0's (superseded).
|
|
round2Msgs := fake.gotMsgs[2]
|
|
var sawRound0Thinking, sawRound1Thinking bool
|
|
for _, m := range round2Msgs {
|
|
if m.Thinking == "round0 reasoning" {
|
|
sawRound0Thinking = true
|
|
}
|
|
if m.Thinking == "round1 reasoning" {
|
|
sawRound1Thinking = true
|
|
}
|
|
}
|
|
if sawRound0Thinking {
|
|
t.Errorf("round0's thinking should have been compacted away by round 2's request, got messages: %+v", round2Msgs)
|
|
}
|
|
if !sawRound1Thinking {
|
|
t.Errorf("round1's thinking (most recent) should still be present in round 2's request, got messages: %+v", round2Msgs)
|
|
}
|
|
}
|
|
|
|
// TestRun_NilThinkPolicyMeansNoOverride confirms a nil ThinkPolicy (an
|
|
// Agent built by hand rather than via New/NewSubagent) falls back to
|
|
// always passing "" rather than panicking.
|
|
func TestRun_NilThinkPolicyMeansNoOverride(t *testing.T) {
|
|
fake := &fakeChatStreamer{
|
|
t: t,
|
|
rounds: [][]vllmkit.ChatEvent{
|
|
{{Token: "hi"}},
|
|
},
|
|
}
|
|
agent := &Agent{Client: fake, Executor: fakeExecutor{}, Tools: nil, MaxRounds: 4, SystemPrompt: SystemPrompt}
|
|
|
|
_, _, err := agent.Run(context.Background(), []vllmkit.Message{{Role: vllmkit.RoleUser, Content: "hi"}}, Callbacks{})
|
|
if err != nil {
|
|
t.Fatalf("Run: %v", err)
|
|
}
|
|
if len(fake.gotThink) != 1 || fake.gotThink[0] != "" {
|
|
t.Errorf("gotThink = %v, want [\"\"]", fake.gotThink)
|
|
}
|
|
}
|
|
|
|
// TestBuildSystemPrompt_PreferredLanguage confirms a locale attached
|
|
// via WithPreferredLanguage reaches the model as an appended
|
|
// "## Interface language" section naming that locale's own display
|
|
// name, and that a run with no preferred language set (the common
|
|
// case — see WithPreferredLanguage's doc comment) sends the base
|
|
// SystemPrompt completely unchanged, byte for byte.
|
|
func TestBuildSystemPrompt_PreferredLanguage(t *testing.T) {
|
|
fake := &fakeChatStreamer{t: t, rounds: [][]vllmkit.ChatEvent{{{Token: "hi"}}}}
|
|
agent := &Agent{Client: fake, Executor: fakeExecutor{}, MaxRounds: 4, SystemPrompt: SystemPrompt}
|
|
|
|
ctx := WithPreferredLanguage(context.Background(), "fr")
|
|
if _, _, err := agent.Run(ctx, []vllmkit.Message{{Role: vllmkit.RoleUser, Content: "bonjour"}}, Callbacks{}); err != nil {
|
|
t.Fatalf("Run: %v", err)
|
|
}
|
|
got := fake.gotMsgs[0][0].Content
|
|
if !strings.Contains(got, "## Interface language") || !strings.Contains(got, "Français") {
|
|
t.Errorf("system prompt missing French interface-language directive, got suffix: %q", got[len(SystemPrompt):])
|
|
}
|
|
}
|
|
|
|
func TestBuildSystemPrompt_NoPreferredLanguageLeavesPromptUnchanged(t *testing.T) {
|
|
fake := &fakeChatStreamer{t: t, rounds: [][]vllmkit.ChatEvent{{{Token: "hi"}}}}
|
|
agent := &Agent{Client: fake, Executor: fakeExecutor{}, MaxRounds: 4, SystemPrompt: SystemPrompt}
|
|
|
|
if _, _, err := agent.Run(context.Background(), []vllmkit.Message{{Role: vllmkit.RoleUser, Content: "hi"}}, Callbacks{}); err != nil {
|
|
t.Fatalf("Run: %v", err)
|
|
}
|
|
if got := fake.gotMsgs[0][0].Content; got != SystemPrompt {
|
|
t.Errorf("system prompt changed with no preferred language set")
|
|
}
|
|
}
|