From 2ea8721c37a3a04be05ec6f9e96415896ca3a60f Mon Sep 17 00:00:00 2001 From: Stefan Walcz Date: Fri, 2 Oct 2026 15:42:17 +0200 Subject: [PATCH] fix(tools): bound the forced reasoning steps by their own output cap The reasoning steps of WithForceReasoning (before the tool pick and before parameter generation) send no output cap of their own, so a model that never closes the reasoning argument runs to the client's global cap (16384 by default), and the length retry then doubles it. Observed with gemma4-26b-a4b behind LocalAI: one forced reasoning call streamed 12,969 tool-call deltas for 478-501 s until the cap, and the turn never answered. The reasoning steps now carry max_tokens from WithReasoningMaxTokens (default DefaultReasoningMaxTokens = 2048; 0 uses the default, a negative value keeps the previous uncapped behaviour). The cap goes on the request, where the bundled clients take it over their own default (see SetMaxTokens). A finish_reason=length on such a capped step is final: no length retry, the call returns an error wrapping errOutputCapReached; a complete tool call that is cut at the cap is still used. The parameter-reasoning step already falls back on any error; the tool-pick step now continues without the reasoning on this error only, other errors stay fatal. A reasoning that finishes normally is used as before. The option is forwarded to plan sub-executions. Signed-off-by: Stefan Walcz Assisted-by: Claude Opus 5.5 --- decision_output_cap_test.go | 84 ++++++++++++ options.go | 29 +++++ plan.go | 3 + reasoning_cap_test.go | 251 ++++++++++++++++++++++++++++++++++++ tools.go | 76 ++++++++++- 5 files changed, 436 insertions(+), 7 deletions(-) create mode 100644 decision_output_cap_test.go create mode 100644 reasoning_cap_test.go diff --git a/decision_output_cap_test.go b/decision_output_cap_test.go new file mode 100644 index 0000000..b2c7b0c --- /dev/null +++ b/decision_output_cap_test.go @@ -0,0 +1,84 @@ +package cogito + +import ( + "context" + "errors" + "testing" + + "github.com/sashabaranov/go-openai" +) + +// lengthCutLLM is a non-streaming fake whose reasoning call is always cut by +// the output cap, with the arguments left open, unless complete is set. +type lengthCutLLM struct { + complete bool + requests []openai.ChatCompletionRequest +} + +func (m *lengthCutLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { return f, nil } + +func (m *lengthCutLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) { + m.requests = append(m.requests, req) + limit := req.MaxTokens + if limit == 0 { + limit = 16384 + } + args := `{"reasoning":"a a a a` + if m.complete { + args = `{"reasoning":"short"}` + } + usage := LLMUsage{PromptTokens: 100, CompletionTokens: limit, TotalTokens: 100 + limit} + return LLMReply{ + ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{{ + FinishReason: openai.FinishReasonLength, + Message: openai.ChatCompletionMessage{Role: "assistant", ToolCalls: []openai.ToolCall{{ + ID: "c1", Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{Name: "reasoning", Arguments: args}, + }}}, + }}}, + MaxTokens: limit, + }, usage, nil +} + +var outputCapConv = []openai.ChatCompletionMessage{{Role: "user", Content: "find x"}} + +func TestDecisionCappedStopsAtItsOwnCap(t *testing.T) { + llm := &lengthCutLLM{} + _, err := decisionCapped(context.Background(), llm, outputCapConv, Tools{reasoningTool()}, "reasoning", 3, 512) + if !errors.Is(err, errOutputCapReached) { + t.Fatalf("expected errOutputCapReached, got %v", err) + } + if len(llm.requests) != 1 { + t.Fatalf("expected one attempt without a length retry, got %d", len(llm.requests)) + } + if got := llm.requests[0].MaxTokens; got != 512 { + t.Fatalf("expected the request to carry max tokens 512, got %d", got) + } +} + +func TestDecisionCappedKeepsACompleteCallCutAtTheCap(t *testing.T) { + llm := &lengthCutLLM{complete: true} + res, err := decisionCapped(context.Background(), llm, outputCapConv, Tools{reasoningTool()}, "reasoning", 3, 512) + if err != nil { + t.Fatalf("expected the complete call to be used, got %v", err) + } + if len(res.toolChoices) != 1 || res.toolChoices[0].Arguments["reasoning"] != "short" { + t.Fatalf("expected the reasoning call, got %+v", res.toolChoices) + } +} + +// Counter-check: without a step cap the decision behaves as before, it sends +// no cap of its own and raises the cap once on a cut tool call. +func TestDecisionWithoutCapKeepsTheLengthRetry(t *testing.T) { + llm := &lengthCutLLM{} + _, err := decision(context.Background(), llm, outputCapConv, Tools{reasoningTool()}, "reasoning", 3) + if err == nil || errors.Is(err, errOutputCapReached) { + t.Fatalf("expected the uncapped truncation error, got %v", err) + } + if len(llm.requests) != 2 { + t.Fatalf("expected the original attempt plus one length retry, got %d", len(llm.requests)) + } + if llm.requests[0].MaxTokens != 0 || llm.requests[1].MaxTokens != 32768 { + t.Fatalf("expected caps 0 then 32768, got %d then %d", llm.requests[0].MaxTokens, llm.requests[1].MaxTokens) + } +} diff --git a/options.go b/options.go index 1c1c50a..ffdd2a7 100644 --- a/options.go +++ b/options.go @@ -48,6 +48,7 @@ type Options struct { loopDetectionSteps int forceReasoning bool forceReasoningTool bool + reasoningMaxTokens int guidedTools bool parallelToolExecution bool toolImageForwarding bool @@ -409,6 +410,34 @@ func WithForceReasoningTool() func(o *Options) { } } +// DefaultReasoningMaxTokens is the output cap of the forced reasoning steps +// (tool and parameter reasoning) when WithReasoningMaxTokens is not set. +const DefaultReasoningMaxTokens = 2048 + +// WithReasoningMaxTokens sets the output-token cap of the forced reasoning +// steps (see WithForceReasoning). These steps are optional: when the cap cuts +// one, the loop continues without its reasoning instead of letting a model +// that never closes the reasoning argument run to the client's global cap. +// 0 (or unset) uses DefaultReasoningMaxTokens; a negative value removes the +// step's own cap, leaving the request to the client's default as before. +func WithReasoningMaxTokens(n int) func(o *Options) { + return func(o *Options) { + o.reasoningMaxTokens = n + } +} + +// reasoningOutputCap returns the cap the reasoning steps send, 0 for none. +func (o *Options) reasoningOutputCap() int { + switch { + case o.reasoningMaxTokens < 0: + return 0 + case o.reasoningMaxTokens == 0: + return DefaultReasoningMaxTokens + default: + return o.reasoningMaxTokens + } +} + // WithToolImageForwarding, when true, forwards image blocks from an MCP tool // result to the model as a follow-up user message (OpenAI tool-role messages // cannot carry image parts). Default false. diff --git a/plan.go b/plan.go index 8cbe7b4..ae16f6c 100644 --- a/plan.go +++ b/plan.go @@ -708,6 +708,9 @@ func convertOptionsToFunctions(o *Options) []Option { if o.forceReasoning { opts = append(opts, WithForceReasoning()) } + if o.reasoningMaxTokens != 0 { + opts = append(opts, WithReasoningMaxTokens(o.reasoningMaxTokens)) + } if o.maxRetries > 0 { opts = append(opts, WithMaxRetries(o.maxRetries)) } diff --git a/reasoning_cap_test.go b/reasoning_cap_test.go new file mode 100644 index 0000000..daa7773 --- /dev/null +++ b/reasoning_cap_test.go @@ -0,0 +1,251 @@ +package cogito_test + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + + . "github.com/mudler/cogito" + "github.com/mudler/cogito/tests/mock" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/sashabaranov/go-openai" +) + +// clientDefaultCap stands in for the client's global output cap (the bundled +// clients send 16384 when the request carries none). +const clientDefaultCap = 16384 + +// reasoningStreamLLM is a streaming fake that answers by the forced tool of a +// request. Its reasoning tool call can be "endless": the argument never +// closes and the stream only ends when the request's output cap (or the +// client default) is spent, one token per delta, with finish_reason=length, +// which is what a model that loops inside the reasoning argument does. +type reasoningStreamLLM struct { + mu sync.Mutex + + endlessToolReasoning bool // the reasoning before the tool pick never ends + endlessParamReasoning bool // the reasoning before parameter generation never ends + shortReasoning string // reasoning text when it does end + + picks int + + reasoningRequests []openai.ChatCompletionRequest + reasoningDeltas []int + intentionRequests []openai.ChatCompletionRequest + paramRequests []openai.ChatCompletionRequest +} + +func forcedName(req openai.ChatCompletionRequest) string { + switch tc := req.ToolChoice.(type) { + case openai.ToolChoice: + return tc.Function.Name + case *openai.ToolChoice: + return tc.Function.Name + } + return "" +} + +func isToolPickReasoning(req openai.ChatCompletionRequest) bool { + for _, m := range req.Messages { + if strings.Contains(m.Content, "Analyze the current situation and available tools") { + return true + } + } + return false +} + +func (m *reasoningStreamLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { + return f.AddMessage(AssistantMessageRole, "final answer"), nil +} + +func (m *reasoningStreamLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) { + return LLMReply{}, LLMUsage{}, fmt.Errorf("unexpected non-streaming call (forced tool %q)", forcedName(req)) +} + +func (m *reasoningStreamLLM) CreateChatCompletionStream(ctx context.Context, req openai.ChatCompletionRequest) (<-chan StreamEvent, error) { + m.mu.Lock() + defer m.mu.Unlock() + + var events []StreamEvent + call := func(name, args string) { + events = append(events, + StreamEvent{Type: StreamEventToolCall, ToolCallIndex: 0, ToolCallID: "c1", ToolName: name, ToolArgs: args}, + StreamEvent{Type: StreamEventDone, FinishReason: "tool_calls", Usage: LLMUsage{CompletionTokens: 20}}) + } + + switch name := forcedName(req); name { + case "reasoning": + m.reasoningRequests = append(m.reasoningRequests, req) + endless := m.endlessParamReasoning + if isToolPickReasoning(req) { + endless = m.endlessToolReasoning + } + if !endless { + args, _ := json.Marshal(map[string]string{"reasoning": m.shortReasoning}) + m.reasoningDeltas = append(m.reasoningDeltas, 1) + call("reasoning", string(args)) + break + } + limit := req.MaxTokens + if limit == 0 { + limit = clientDefaultCap + } + events = append(events, StreamEvent{Type: StreamEventToolCall, ToolCallIndex: 0, ToolCallID: "c1", ToolName: "reasoning", ToolArgs: `{"reasoning":"`}) + for i := 1; i < limit; i++ { + events = append(events, StreamEvent{Type: StreamEventToolCall, ToolCallIndex: 0, ToolArgs: "a "}) + } + events = append(events, StreamEvent{Type: StreamEventDone, FinishReason: "length", MaxTokens: limit, + Usage: LLMUsage{PromptTokens: 100, CompletionTokens: limit, TotalTokens: 100 + limit}}) + m.reasoningDeltas = append(m.reasoningDeltas, limit) + case "pick_tool", "pick_tools", "": + if name == "" && len(req.Tools) == 0 { + // Plain answer turn. + events = append(events, + StreamEvent{Type: StreamEventContent, Content: "final answer"}, + StreamEvent{Type: StreamEventDone, FinishReason: "stop"}) + break + } + m.intentionRequests = append(m.intentionRequests, req) + m.picks++ + if m.picks == 1 { + call("pick_tool", `{"tool":"search","reasoning":"search first"}`) + } else { + call("pick_tool", `{"tool":"reply","reasoning":"done"}`) + } + case "search": + m.paramRequests = append(m.paramRequests, req) + call("search", `{}`) + default: + return nil, fmt.Errorf("unexpected forced tool %q", name) + } + + ch := make(chan StreamEvent, len(events)) + for _, ev := range events { + ch <- ev + } + close(ch) + return ch, nil +} + +func runWithReasoning(llm *reasoningStreamLLM, opts ...Option) (Fragment, error) { + f := NewEmptyFragment().AddMessage(UserMessageRole, "Wie buche ich eine Eingangsrechnung?") + search := mock.NewMockTool("search", "Search the knowledge base") + mock.SetRunResult(search, "search result") + base := []Option{ + WithTools(search), + WithForceReasoning(), + WithIterations(2), + WithMaxRetries(2), + WithStreamCallback(func(StreamEvent) {}), + } + return ExecuteTools(llm, f, append(base, opts...)...) +} + +func lastAssistant(f Fragment) string { + for i := len(f.Messages) - 1; i >= 0; i-- { + if f.Messages[i].Role == AssistantMessageRole.String() && f.Messages[i].Content != "" { + return f.Messages[i].Content + } + } + return "" +} + +var _ = Describe("Reasoning step output cap", func() { + It("caps an endless tool-pick reasoning and still reaches the tool and a final answer", func() { + llm := &reasoningStreamLLM{endlessToolReasoning: true} + + result, err := runWithReasoning(llm) + Expect(err).ToNot(HaveOccurred()) + + Expect(llm.reasoningRequests).ToNot(BeEmpty()) + for _, req := range llm.reasoningRequests { + Expect(req.MaxTokens).To(Equal(DefaultReasoningMaxTokens), "every reasoning request carries the step's own cap") + } + for _, n := range llm.reasoningDeltas { + Expect(n).To(BeNumerically("<=", DefaultReasoningMaxTokens), "the reasoning stream stops at the step cap, not at the client cap") + } + // A cut step is final: no length retry with a raised cap per pick. + Expect(llm.reasoningRequests).To(HaveLen(llm.picks)) + + // The loop continued: the tool was picked, parametrized and run. + Expect(llm.paramRequests).To(HaveLen(1)) + Expect(result.Status.ToolsCalled.Names()).To(ContainElement("search")) + Expect(lastAssistant(result)).ToNot(BeEmpty()) + }) + + It("caps an endless parameter reasoning and falls back to the original reasoning", func() { + llm := &reasoningStreamLLM{endlessParamReasoning: true, shortReasoning: "use search for the booking rule"} + + result, err := runWithReasoning(llm) + Expect(err).ToNot(HaveOccurred()) + + var capped int + for i, req := range llm.reasoningRequests { + Expect(req.MaxTokens).To(Equal(DefaultReasoningMaxTokens)) + if llm.reasoningDeltas[i] > 1 { + capped++ + Expect(llm.reasoningDeltas[i]).To(Equal(DefaultReasoningMaxTokens)) + } + } + Expect(capped).To(Equal(1), "exactly the one parameter reasoning ran into the cap, once") + + Expect(llm.paramRequests).To(HaveLen(1)) + Expect(llm.paramRequests[0].Messages[0].Content).To(ContainSubstring("use search for the booking rule"), + "the fallback keeps the tool-pick reasoning for parameter generation") + Expect(result.Status.ToolsCalled.Names()).To(ContainElement("search")) + }) + + It("uses a short reasoning that ends normally (counter-check)", func() { + llm := &reasoningStreamLLM{shortReasoning: "the user asks about booking, search the knowledge base"} + + result, err := runWithReasoning(llm) + Expect(err).ToNot(HaveOccurred()) + + // The reasoning reaches the tool pick as an assistant turn ... + Expect(llm.intentionRequests).ToNot(BeEmpty()) + var seen bool + for _, m := range llm.intentionRequests[0].Messages { + if strings.Contains(m.Content, "the user asks about booking, search the knowledge base") { + seen = true + } + } + Expect(seen).To(BeTrue(), "the reasoning is passed on to the tool pick") + // ... and the parameter generation gets the combined reasoning. + Expect(llm.paramRequests).To(HaveLen(1)) + Expect(llm.paramRequests[0].Messages[0].Content).To(ContainSubstring("Parameter Analysis")) + Expect(result.Status.ToolsCalled.Names()).To(ContainElement("search")) + }) + + It("sends a configured cap, the default for 0, and no step cap when negative", func() { + for _, tc := range []struct { + opt Option + want int + }{ + {WithReasoningMaxTokens(512), 512}, + {WithReasoningMaxTokens(0), DefaultReasoningMaxTokens}, + {nil, DefaultReasoningMaxTokens}, + {WithReasoningMaxTokens(-1), 0}, + } { + llm := &reasoningStreamLLM{endlessToolReasoning: true} + var opts []Option + if tc.opt != nil { + opts = append(opts, tc.opt) + } + _, err := runWithReasoning(llm, opts...) + Expect(llm.reasoningRequests).ToNot(BeEmpty()) + Expect(llm.reasoningRequests[0].MaxTokens).To(Equal(tc.want)) + if tc.want > 0 { + Expect(err).ToNot(HaveOccurred()) + continue + } + // Without a step cap the old behaviour stays: the endless + // reasoning runs to the client cap, the length retry doubles it, + // and the turn fails. This is the measured production failure. + Expect(err).To(HaveOccurred()) + Expect(llm.reasoningDeltas[0]).To(Equal(clientDefaultCap)) + } + }) +}) diff --git a/tools.go b/tools.go index fd1bb7d..32d5e80 100644 --- a/tools.go +++ b/tools.go @@ -282,16 +282,41 @@ func mergeConsecutiveAssistantMessages(messages []openai.ChatCompletionMessage) // Falls back to decision() when streaming is not possible. func decisionWithStreaming(ctx context.Context, llm LLM, conversation []openai.ChatCompletionMessage, tools Tools, forceTool string, maxRetries int, streamCB StreamCallback) (*decisionResult, error) { + return decisionWithStreamingCapped(ctx, llm, conversation, tools, forceTool, maxRetries, streamCB, 0) +} + +// errOutputCapReached reports that a decision with its own output cap (see +// decisionWithStreamingCapped) was cut by that cap. Callers treat the step as +// optional and continue without its result. +var errOutputCapReached = errors.New("output cap of an optional step reached") + +// outputCapReachedError wraps errOutputCapReached with the figures of the cut. +func outputCapReachedError(what string, outCap int, usage LLMUsage) error { + return fmt.Errorf("%s: %w (max tokens %d, completion tokens %d)", what, errOutputCapReached, outCap, usage.CompletionTokens) +} + +// decisionWithStreamingCapped is decisionWithStreaming with a request-level +// output cap. With outCap > 0 every attempt carries max_tokens=outCap, which +// the bundled clients honour over their own default, and a +// finish_reason=length is final: no length retry with a raised cap, the call +// returns an error wrapping errOutputCapReached instead. Meant for optional +// steps (the reasoning tool) that must not run to the client's global cap. +// With outCap <= 0 it behaves exactly like decisionWithStreaming. +func decisionWithStreamingCapped(ctx context.Context, llm LLM, conversation []openai.ChatCompletionMessage, + tools Tools, forceTool string, maxRetries int, streamCB StreamCallback, outCap int) (*decisionResult, error) { sllm, isStreaming := llm.(StreamingLLM) if !isStreaming || streamCB == nil { - return decision(ctx, llm, conversation, tools, forceTool, maxRetries) + return decisionCapped(ctx, llm, conversation, tools, forceTool, maxRetries, outCap) } req := openai.ChatCompletionRequest{ Messages: mergeConsecutiveAssistantMessages(normalizeSystemMessages(conversation)), Tools: tools.ToOpenAI(), } + if outCap > 0 { + req.MaxTokens = outCap + } if forceTool != "" { req.ToolChoice = openai.ToolChoice{ @@ -389,6 +414,16 @@ func decisionWithStreaming(ctx context.Context, llm LLM, conversation []openai.C xlog.Debug("[decisionWithStreaming] processed", "message", content, "reasoning", reasoning) + if outCap > 0 && finishReason == "length" { + // The step's own cap cut it. A complete tool call is still + // usable; anything else (no call, a call with cut arguments) ends + // the optional step here instead of retrying with a raised cap. + if choices, bad := parseToolCalls(toolCalls, finishReason); len(toolCalls) > 0 && bad == nil { + return &decisionResult{toolChoices: choices, message: content, reasoning: reasoning, usage: usage}, nil + } + return nil, outputCapReachedError("streaming decision", outCap, usage) + } + if len(toolCalls) == 0 { if content == "" { // The model produced no visible content and selected no tool. @@ -469,11 +504,21 @@ func backoffOrCancel(ctx context.Context, attempt int) error { // Similar to agent.go's decision function but adapted for cogito's architecture func decision(ctx context.Context, llm LLM, conversation []openai.ChatCompletionMessage, tools Tools, forceTool string, maxRetries int) (*decisionResult, error) { + return decisionCapped(ctx, llm, conversation, tools, forceTool, maxRetries, 0) +} + +// decisionCapped is decision with a request-level output cap; see +// decisionWithStreamingCapped for the semantics of outCap. +func decisionCapped(ctx context.Context, llm LLM, conversation []openai.ChatCompletionMessage, + tools Tools, forceTool string, maxRetries int, outCap int) (*decisionResult, error) { decision := openai.ChatCompletionRequest{ Messages: mergeConsecutiveAssistantMessages(normalizeSystemMessages(conversation)), Tools: tools.ToOpenAI(), } + if outCap > 0 { + decision.MaxTokens = outCap + } if forceTool != "" { decision.ToolChoice = openai.ToolChoice{ @@ -519,6 +564,14 @@ func decision(ctx context.Context, llm LLM, conversation []openai.ChatCompletion //reasoning := resp.Choices[0].Reasoning xlog.Debug("[decision] processed", "message", msg.Content, "reasoning", reasoning) + if outCap > 0 && resp.ChatCompletionResponse.Choices[0].FinishReason == openai.FinishReasonLength { + // See decisionWithStreamingCapped. + if choices, bad := parseToolCalls(msg.ToolCalls, string(openai.FinishReasonLength)); len(msg.ToolCalls) > 0 && bad == nil { + return &decisionResult{toolChoices: choices, message: msg.Content, reasoning: reasoning, usage: usage}, nil + } + return nil, outputCapReachedError("decision", outCap, usage) + } + if len(msg.ToolCalls) == 0 { // An empty reply cut by length is the same truncation the streaming // path handles; returning it as an empty text answer would end the @@ -610,13 +663,15 @@ func generateToolParameters(o *Options, llm LLM, tool ToolDefinitionInterface, c return nil, err } - // Use decision with reasoning tool to force structured output - paramReasoningResult, err := decisionWithStreaming(o.context, llm, + // Use decision with reasoning tool to force structured output. + // Bounded by the reasoning cap: on truncation the error below takes + // the same fallback as any other failure. + paramReasoningResult, err := decisionWithStreamingCapped(o.context, llm, append(conversation, openai.ChatCompletionMessage{ Role: "system", Content: paramPrompt, }), - Tools{reasoningTool()}, "reasoning", o.maxRetries, o.streamCallback) + Tools{reasoningTool()}, "reasoning", o.maxRetries, o.streamCallback, o.reasoningOutputCap()) if err != nil { xlog.Warn("Failed to get parameter reasoning, using original reasoning", "error", err) // Fall back to original single-step approach @@ -723,14 +778,21 @@ func pickTool(ctx context.Context, llm LLM, fragment Fragment, tools Tools, opts } } - reasoningResult, err := decisionWithStreaming(ctx, llm, + // The reasoning step is optional: bounded by its own output cap, and when + // that cap cuts it the tool pick continues without the reasoning instead + // of letting a model that never closes the argument run to the global cap. + reasoningResult, err := decisionWithStreamingCapped(ctx, llm, append(messages, openai.ChatCompletionMessage{ Role: "user", Content: reasoningPrompt, }), - Tools{reasoningTool()}, "reasoning", o.maxRetries, o.streamCallback) + Tools{reasoningTool()}, "reasoning", o.maxRetries, o.streamCallback, o.reasoningOutputCap()) if err != nil { - return nil, fmt.Errorf("failed to get reasoning: %w", err) + if !errors.Is(err, errOutputCapReached) { + return nil, fmt.Errorf("failed to get reasoning: %w", err) + } + xlog.Warn("[pickTool] reasoning step cut by its output cap, picking the tool without it", "error", err) + reasoningResult = &decisionResult{} } // Extract reasoning from the tool call response