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