diff --git a/internal/agent/agent.go b/internal/agent/agent.go index a57b8df1..87c62753 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -163,10 +163,21 @@ type ErrorHandler func(err error) // RetryHandler is called when the LLM request is retried. type RetryHandler func(attempt int, err error) -// PrepareStepHandler is called between steps to allow message modification. -// It receives the step number and current messages, and returns replacement -// messages (or nil to keep unchanged). -type PrepareStepHandler func(stepNumber int, messages []fantasy.Message) []fantasy.Message +// PrepareStepUpdate carries per-step overrides returned by a +// PrepareStepHandler. Nil fields leave the corresponding aspect of the step +// unchanged. +type PrepareStepUpdate struct { + // Messages, when non-nil, replaces the context window for this step. + Messages []fantasy.Message + // ToolChoice, when non-nil, overrides the tool-choice mode for this step. + ToolChoice *fantasy.ToolChoice +} + +// PrepareStepHandler is called between steps to allow message modification +// and per-step tool-choice control. It receives the step number and current +// messages, and returns a PrepareStepUpdate (or nil to keep everything +// unchanged). +type PrepareStepHandler func(stepNumber int, messages []fantasy.Message) *PrepareStepUpdate // GenerateCallbacks consolidates all callback functions for // GenerateWithCallbacks into a single struct, replacing what was previously @@ -923,8 +934,13 @@ func (a *Agent) GenerateWithCallbacks(ctx context.Context, messages []fantasy.Me // Phase 2: Run OnPrepareStep hook (if registered). if hasPrepareStepHook { - if replacement := cb.OnPrepareStep(opts.StepNumber, result.Messages); replacement != nil { - result.Messages = replacement + if update := cb.OnPrepareStep(opts.StepNumber, result.Messages); update != nil { + if update.Messages != nil { + result.Messages = update.Messages + } + if update.ToolChoice != nil { + result.ToolChoice = update.ToolChoice + } } } diff --git a/internal/agent/agent_toolchoice_test.go b/internal/agent/agent_toolchoice_test.go new file mode 100644 index 00000000..6d503562 --- /dev/null +++ b/internal/agent/agent_toolchoice_test.go @@ -0,0 +1,103 @@ +package agent + +import ( + "context" + "testing" + + "charm.land/fantasy" +) + +// toolChoiceAgent is a fake fantasy.Agent that simulates a two-step agentic +// loop and records the ToolChoice each step's PrepareStep callback yields, so +// tests can assert per-step tool-choice overrides are threaded through. +type toolChoiceAgent struct { + stepChoices []*fantasy.ToolChoice +} + +func (f *toolChoiceAgent) Generate(_ context.Context, _ fantasy.AgentCall) (*fantasy.AgentResult, error) { + return &fantasy.AgentResult{}, nil +} + +func (f *toolChoiceAgent) Stream(ctx context.Context, opts fantasy.AgentStreamCall) (*fantasy.AgentResult, error) { + for step := range 2 { + if opts.PrepareStep != nil { + _, prepared, err := opts.PrepareStep(ctx, fantasy.PrepareStepFunctionOptions{ + StepNumber: step, + Model: nil, + Messages: nil, + }) + if err != nil { + return nil, err + } + f.stepChoices = append(f.stepChoices, prepared.ToolChoice) + } + } + return &fantasy.AgentResult{}, nil +} + +// TestPrepareStepToolChoiceForcing verifies that a PrepareStepHandler can +// force a specific tool choice on one step and release it on the next. This +// is the pattern used to guarantee a capture/structured-output tool is +// invoked: force SpecificToolChoice until the call is observed, then flip +// back to nil so the turn can end. +func TestPrepareStepToolChoiceForcing(t *testing.T) { + t.Parallel() + + fake := &toolChoiceAgent{} + a := &Agent{ + fantasyAgent: fake, + streamingEnabled: true, + } + + forced := fantasy.SpecificToolChoice("record_result") + cb := GenerateCallbacks{ + OnPrepareStep: func(stepNumber int, messages []fantasy.Message) *PrepareStepUpdate { + if stepNumber == 0 { + return &PrepareStepUpdate{ToolChoice: &forced} + } + return nil + }, + } + + msgs := []fantasy.Message{fantasy.NewUserMessage("go")} + if _, err := a.GenerateWithCallbacks(context.Background(), msgs, cb); err != nil { + t.Fatalf("GenerateWithCallbacks returned error: %v", err) + } + + if len(fake.stepChoices) != 2 { + t.Fatalf("expected 2 prepared steps, got %d", len(fake.stepChoices)) + } + if fake.stepChoices[0] == nil { + t.Fatal("step 0: expected forced ToolChoice, got nil") + } + if *fake.stepChoices[0] != forced { + t.Errorf("step 0: expected ToolChoice %q, got %q", forced, *fake.stepChoices[0]) + } + if fake.stepChoices[1] != nil { + t.Errorf("step 1: expected nil ToolChoice after release, got %q", *fake.stepChoices[1]) + } +} + +// TestPrepareStepToolChoiceWithoutHandler verifies that when no +// PrepareStepHandler is registered, no ToolChoice override is sent — fantasy +// falls back to its default (auto). +func TestPrepareStepToolChoiceWithoutHandler(t *testing.T) { + t.Parallel() + + fake := &toolChoiceAgent{} + a := &Agent{ + fantasyAgent: fake, + streamingEnabled: true, + } + + msgs := []fantasy.Message{fantasy.NewUserMessage("go")} + if _, err := a.GenerateWithCallbacks(context.Background(), msgs, GenerateCallbacks{}); err != nil { + t.Fatalf("GenerateWithCallbacks returned error: %v", err) + } + + for i, tc := range fake.stepChoices { + if tc != nil { + t.Errorf("step %d: expected nil ToolChoice without handler, got %q", i, *tc) + } + } +} diff --git a/pkg/kit/hooks.go b/pkg/kit/hooks.go index 36df2422..f1ece602 100644 --- a/pkg/kit/hooks.go +++ b/pkg/kit/hooks.go @@ -138,11 +138,24 @@ type PrepareStepHook struct { Messages []LLMMessage } -// PrepareStepResult can replace the context window between steps. +// PrepareStepResult can replace the context window and override the +// tool-choice mode between steps. type PrepareStepResult struct { // Messages replaces the entire context window for this step. If nil, // the original messages (including any steering) are used unchanged. Messages []LLMMessage + // ToolChoice, when non-nil, overrides the tool-choice mode for this step + // only. Use [LLMToolChoiceRequired] or [LLMSpecificToolChoice] to force a + // tool call, [LLMToolChoiceNone] to forbid tool calls, or + // [LLMToolChoiceAuto] to explicitly restore the default. If nil, the + // provider default (auto) is used. + // + // Note that a turn only ends when the model stops emitting tool calls, so + // forcing a tool choice on every step prevents the turn from finishing. + // Callers forcing a tool call should flip back to nil (or + // LLMToolChoiceAuto) on subsequent steps once the desired call has been + // observed. + ToolChoice *LLMToolChoice } // --------------------------------------------------------------------------- @@ -275,12 +288,14 @@ func (m *Kit) OnBeforeCompact(p HookPriority, h func(BeforeCompactHook) *BeforeC // OnPrepareStep registers a hook that fires between steps within a multi-step // agent turn, after steering messages are injected and before the messages are // sent to the LLM. Return a non-nil PrepareStepResult with Messages to replace -// the entire context window for this step. Hooks execute in priority order; -// the first non-nil result wins. Returns an unregister function. +// the entire context window for this step, and/or with ToolChoice to override +// the tool-choice mode for this step. Hooks execute in priority order; the +// first non-nil result wins. Returns an unregister function. // // This is the most powerful interception point in the agent lifecycle. It // enables patterns like transforming tool results, dynamic tool filtering, -// and mid-turn context injection. +// mid-turn context injection, and forcing a specific tool call on a step +// (e.g. guaranteeing a structured-output capture tool is invoked). func (m *Kit) OnPrepareStep(p HookPriority, h func(PrepareStepHook) *PrepareStepResult) func() { return m.prepareStep.register(p, h) } diff --git a/pkg/kit/hooks_test.go b/pkg/kit/hooks_test.go index 214e41ae..9d892792 100644 --- a/pkg/kit/hooks_test.go +++ b/pkg/kit/hooks_test.go @@ -578,6 +578,45 @@ func TestPrepareStepHookRegistry(t *testing.T) { } } +// TestPrepareStepHookToolChoice verifies that a PrepareStep hook can set a +// per-step tool-choice override without replacing messages. +func TestPrepareStepHookToolChoice(t *testing.T) { + hr := newHookRegistry[PrepareStepHook, PrepareStepResult]() + + forced := LLMSpecificToolChoice("record_result") + hr.register(HookPriorityNormal, func(h PrepareStepHook) *PrepareStepResult { + if h.StepNumber == 0 { + return &PrepareStepResult{ToolChoice: &forced} + } + return nil + }) + + input := PrepareStepHook{ + StepNumber: 0, + Messages: []LLMMessage{NewLLMUserMessage("hello")}, + } + result := hr.run(input) + if result == nil { + t.Fatal("expected non-nil result for step 0") + return + } + if result.Messages != nil { + t.Errorf("expected nil Messages (unchanged), got %d messages", len(result.Messages)) + } + if result.ToolChoice == nil { + t.Fatal("expected ToolChoice to be set") + } + if *result.ToolChoice != forced { + t.Errorf("expected ToolChoice %q, got %q", forced, *result.ToolChoice) + } + + // Step 1 — no override. + input.StepNumber = 1 + if result := hr.run(input); result != nil { + t.Errorf("expected nil result for step 1, got %+v", result) + } +} + // TestPrepareStepHookPriority verifies that PrepareStep hooks respect priority ordering. func TestPrepareStepHookPriority(t *testing.T) { hr := newHookRegistry[PrepareStepHook, PrepareStepResult]() diff --git a/pkg/kit/kit.go b/pkg/kit/kit.go index 94f6ea15..731833e7 100644 --- a/pkg/kit/kit.go +++ b/pkg/kit/kit.go @@ -2695,15 +2695,18 @@ func (m *Kit) generate(ctx context.Context, messages []fantasy.Message) (*agent. if !m.prepareStep.hasHooks() { return nil } - return func(stepNumber int, messages []fantasy.Message) []fantasy.Message { + return func(stepNumber int, messages []fantasy.Message) *agent.PrepareStepUpdate { hookResult := m.prepareStep.run(PrepareStepHook{ StepNumber: stepNumber, Messages: messages, }) - if hookResult != nil && hookResult.Messages != nil { - return hookResult.Messages + if hookResult == nil || (hookResult.Messages == nil && hookResult.ToolChoice == nil) { + return nil + } + return &agent.PrepareStepUpdate{ + Messages: hookResult.Messages, + ToolChoice: hookResult.ToolChoice, } - return nil } }(), }) diff --git a/pkg/kit/types.go b/pkg/kit/types.go index 9ff7c7cd..16064bde 100644 --- a/pkg/kit/types.go +++ b/pkg/kit/types.go @@ -228,6 +228,26 @@ const ( // LLMToolInfo describes a tool's name, description, and JSON-Schema parameters. type LLMToolInfo = fantasy.ToolInfo +// LLMToolChoice controls which tool (if any) the model is allowed or required +// to call on a given LLM step. Use the LLMToolChoice* constants for the +// generic modes, or [LLMSpecificToolChoice] to force a single named tool. +type LLMToolChoice = fantasy.ToolChoice + +// LLMToolChoice mode constants mirror the provider's tool-choice values. +const ( + // LLMToolChoiceNone forbids tool calls for the step. + LLMToolChoiceNone = fantasy.ToolChoiceNone + // LLMToolChoiceAuto lets the model decide whether to call tools. This is + // the default behaviour when no tool choice is set. + LLMToolChoiceAuto = fantasy.ToolChoiceAuto + // LLMToolChoiceRequired forces the model to call some tool on the step. + LLMToolChoiceRequired = fantasy.ToolChoiceRequired +) + +// LLMSpecificToolChoice returns a tool choice that forces the model to call +// the named tool on the step. +var LLMSpecificToolChoice = fantasy.SpecificToolChoice + // LLMProviderOptions carries provider-specific key/value option maps, keyed // by provider name (e.g. "anthropic"). Use this when configuring or // inspecting provider-specific tool behaviour.