diff --git a/README.md b/README.md index 98a10a0..9a8874a 100644 --- a/README.md +++ b/README.md @@ -629,6 +629,7 @@ result, err := cogito.ExecuteTools(llm, fragment, - When enabled, the sink state tool appears as an option in the tool selection enum - The sink state tool receives a `reasoning` parameter containing the LLM's reasoning about why no tool is needed - Custom sink state tools must accept a `reasoning` parameter in their arguments +- The sink state's arguments are generated under their own output cap, `DefaultSinkStateMaxTokens` (2048) by default. A model that never closes the `reasoning` argument is cut at the cap instead of running to the client's global limit, and a capped step is neither length-retried nor sent through argument correction; the loop then answers as usual. Set it with `cogito.WithSinkStateMaxTokens(n)` (`0` = default, a negative value restores the previous uncapped behaviour). Real tools are not affected by this cap. #### Sub-Agent Spawning 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..8129773 100644 --- a/options.go +++ b/options.go @@ -48,6 +48,8 @@ type Options struct { loopDetectionSteps int forceReasoning bool forceReasoningTool bool + reasoningMaxTokens int + sinkStateMaxTokens int guidedTools bool parallelToolExecution bool toolImageForwarding bool @@ -409,6 +411,64 @@ 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 + } +} + +// DefaultSinkStateMaxTokens is the output cap of the sink state's parameter +// generation when WithSinkStateMaxTokens is not set. +const DefaultSinkStateMaxTokens = 2048 + +// WithSinkStateMaxTokens sets the output-token cap of the step that generates +// the sink state's arguments (see WithSinkState, WithForceReasoning). The sink +// ends the tool loop and its arguments are not executed — the answer comes +// from the closing completion after the loop — so when the cap cuts them the +// loop ends without them instead of letting a model that repeats itself inside +// the argument run to the client's global cap. Real tools are not affected. +// 0 (or unset) uses DefaultSinkStateMaxTokens; a negative value removes the +// step's own cap, leaving the request to the client's default as before. +func WithSinkStateMaxTokens(n int) func(o *Options) { + return func(o *Options) { + o.sinkStateMaxTokens = n + } +} + +// sinkStateOutputCap returns the cap the sink's parameter step sends, 0 for none. +func (o *Options) sinkStateOutputCap() int { + switch { + case o.sinkStateMaxTokens < 0: + return 0 + case o.sinkStateMaxTokens == 0: + return DefaultSinkStateMaxTokens + default: + return o.sinkStateMaxTokens + } +} + // 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..f5b6ed0 100644 --- a/plan.go +++ b/plan.go @@ -708,6 +708,12 @@ func convertOptionsToFunctions(o *Options) []Option { if o.forceReasoning { opts = append(opts, WithForceReasoning()) } + if o.reasoningMaxTokens != 0 { + opts = append(opts, WithReasoningMaxTokens(o.reasoningMaxTokens)) + } + if o.sinkStateMaxTokens != 0 { + opts = append(opts, WithSinkStateMaxTokens(o.sinkStateMaxTokens)) + } 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/sink_cap_test.go b/sink_cap_test.go new file mode 100644 index 0000000..a30a887 --- /dev/null +++ b/sink_cap_test.go @@ -0,0 +1,285 @@ +package cogito_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + + . "github.com/mudler/cogito" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/sashabaranov/go-openai" +) + +// The loop an embedder's sink argument ran into in production: the model +// answered the question inside the sink's "reasoning" argument and then +// repeated the last sentence until the output cap. +const sinkLoopSentence = "17 mal 23 ist 391. " + +// sinkArgs mirrors the sink state an embedder registers (LocalAGI's +// no_tool_to_call): a tool with one free-text argument. +type sinkArgs struct { + Reasoning string `json:"reasoning"` +} + +type sinkRunner struct{} + +func (sinkRunner) Run(args sinkArgs) (string, any, error) { + return "No action needed: " + args.Reasoning, nil, nil +} + +// docArgs is a real tool whose argument is legitimately long (a document). +type docArgs struct { + Content string `json:"content"` +} + +type docRunner struct { + mu sync.Mutex + runs []string +} + +func (d *docRunner) Run(args docArgs) (string, any, error) { + d.mu.Lock() + defer d.mu.Unlock() + d.runs = append(d.runs, args.Content) + return "document written", nil, nil +} + +// sinkTurn is one scripted reply, rendered either as a stream or as a single +// completion, so both decision paths see the same model behaviour. +type sinkTurn struct { + content string + toolName string + toolArgs []string // argument deltas + finishReason string + maxTokens int + completion int +} + +// sinkLLM answers by the forced tool of a request. The parameter generation +// of the sink can be scripted to loop inside its argument until the output +// cap of the request (or the client default) is spent. +type sinkLLM struct { + mu sync.Mutex + + picks []string // tools the intention step picks, in order (last repeats) + + endlessSink bool // the sink argument never closes + maskedLength bool // the backend reports finish_reason=tool_calls instead of length + docDeltas int // length of the real tool's argument, in deltas + sinkText string // sink argument when it does end + + pickIndex int + sinkRequests []openai.ChatCompletionRequest + sinkDeltas []int + docRequests []openai.ChatCompletionRequest +} + +func (m *sinkLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { + return f.AddMessage(AssistantMessageRole, strings.TrimSpace(sinkLoopSentence)), nil +} + +func (m *sinkLLM) turn(req openai.ChatCompletionRequest) (sinkTurn, error) { + m.mu.Lock() + defer m.mu.Unlock() + + call := func(name, args string) sinkTurn { + return sinkTurn{toolName: name, toolArgs: []string{args}, finishReason: "tool_calls", completion: 20} + } + + switch name := forcedName(req); name { + case "reasoning": + return call("reasoning", `{"reasoning":"no tool is needed, answer directly"}`), nil + case "pick_tool": + pick := m.picks[min(m.pickIndex, len(m.picks)-1)] + m.pickIndex++ + return call("pick_tool", fmt.Sprintf(`{"tool":%q,"reasoning":"pick %s"}`, pick, pick)), nil + case "no_tool_to_call": + m.sinkRequests = append(m.sinkRequests, req) + if !m.endlessSink { + args, _ := json.Marshal(sinkArgs{Reasoning: m.sinkText}) + m.sinkDeltas = append(m.sinkDeltas, 1) + return call(name, string(args)), nil + } + limit := req.MaxTokens + if limit == 0 { + limit = clientDefaultCap + } + deltas := make([]string, 0, limit) + deltas = append(deltas, `{"reasoning":"`) + for i := 1; i < limit; i++ { + deltas = append(deltas, sinkLoopSentence) + } + finish := "length" + if m.maskedLength { + // What LocalAI reported in production: the call was cut by the + // cap, the finish reason still said tool_calls. + finish = "tool_calls" + } + m.sinkDeltas = append(m.sinkDeltas, limit) + return sinkTurn{toolName: name, toolArgs: deltas, finishReason: finish, maxTokens: limit, completion: limit}, nil + case "write_document": + m.docRequests = append(m.docRequests, req) + deltas := []string{`{"content":"`} + for i := 0; i < m.docDeltas; i++ { + deltas = append(deltas, "x ") + } + deltas = append(deltas, `"}`) + return sinkTurn{toolName: name, toolArgs: deltas, finishReason: "tool_calls", completion: len(deltas)}, nil + case "": + if len(req.Tools) == 0 { + return sinkTurn{content: strings.TrimSpace(sinkLoopSentence), finishReason: "stop"}, nil + } + } + return sinkTurn{}, fmt.Errorf("unexpected request (forced tool %q)", forcedName(req)) +} + +func (m *sinkLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) { + t, err := m.turn(req) + if err != nil { + return LLMReply{}, LLMUsage{}, err + } + msg := openai.ChatCompletionMessage{Role: openai.ChatMessageRoleAssistant, Content: t.content} + if t.toolName != "" { + msg.ToolCalls = []openai.ToolCall{{ID: "c1", Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{Name: t.toolName, Arguments: strings.Join(t.toolArgs, "")}}} + } + usage := LLMUsage{PromptTokens: 100, CompletionTokens: t.completion, TotalTokens: 100 + t.completion} + return LLMReply{ + ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{ + {Message: msg, FinishReason: openai.FinishReason(t.finishReason)}, + }}, + MaxTokens: t.maxTokens, + }, usage, nil +} + +func (m *sinkLLM) CreateChatCompletionStream(ctx context.Context, req openai.ChatCompletionRequest) (<-chan StreamEvent, error) { + t, err := m.turn(req) + if err != nil { + return nil, err + } + var events []StreamEvent + if t.content != "" { + events = append(events, StreamEvent{Type: StreamEventContent, Content: t.content}) + } + for i, d := range t.toolArgs { + ev := StreamEvent{Type: StreamEventToolCall, ToolCallIndex: 0, ToolArgs: d} + if i == 0 { + ev.ToolCallID, ev.ToolName = "c1", t.toolName + } + events = append(events, ev) + } + events = append(events, StreamEvent{Type: StreamEventDone, FinishReason: t.finishReason, MaxTokens: t.maxTokens, + Usage: LLMUsage{PromptTokens: 100, CompletionTokens: t.completion, TotalTokens: 100 + t.completion}}) + ch := make(chan StreamEvent, len(events)) + for _, ev := range events { + ch <- ev + } + close(ch) + return ch, nil +} + +func runWithSink(llm *sinkLLM, streaming bool, doc *docRunner, opts ...Option) (Fragment, error) { + f := NewEmptyFragment().AddMessage(UserMessageRole, "Was ist 17 mal 23?") + base := []Option{ + WithSinkState(NewToolDefinition(sinkRunner{}, sinkArgs{}, "no_tool_to_call", + "Called when no other tool is needed to respond to the user")), + WithForceReasoning(), + WithIterations(3), + WithMaxRetries(2), + } + if doc != nil { + base = append(base, WithTools(NewToolDefinition[docArgs](doc, docArgs{}, "write_document", "Write a document"))) + } + if streaming { + base = append(base, WithStreamCallback(func(StreamEvent) {})) + } + return ExecuteTools(llm, f, append(base, opts...)...) +} + +// The sink-only turn ends without a tool call recorded; cogito reports that +// as ErrNoToolSelected, which embedders (LocalAGI) treat as a normal answer. +func expectAnswered(result Fragment, err error) { + if err != nil { + Expect(errors.Is(err, ErrNoToolSelected)).To(BeTrue(), "unexpected error: %v", err) + } + Expect(lastAssistant(result)).To(Equal(strings.TrimSpace(sinkLoopSentence))) +} + +var _ = Describe("Sink state output cap", func() { + for _, streaming := range []bool{true, false} { + for _, masked := range []bool{false, true} { + streaming, masked := streaming, masked + It(fmt.Sprintf("caps a sink argument that never closes, once, and still answers (streaming=%v, masked length=%v)", streaming, masked), func() { + llm := &sinkLLM{picks: []string{"no_tool_to_call"}, endlessSink: true, maskedLength: masked} + + result, err := runWithSink(llm, streaming, nil) + expectAnswered(result, err) + + Expect(llm.sinkRequests).To(HaveLen(1), "a capped sink step is final: no length retry, no correction retry") + Expect(llm.sinkRequests[0].MaxTokens).To(Equal(DefaultSinkStateMaxTokens), "the sink step carries its own cap") + Expect(llm.sinkDeltas[0]).To(Equal(DefaultSinkStateMaxTokens), "the loop stops at the sink cap, not the client cap") + }) + } + } + + It("uses a sink argument that ends normally as before (counter-check)", func() { + llm := &sinkLLM{picks: []string{"no_tool_to_call"}, sinkText: "17 mal 23 ist 391."} + + result, err := runWithSink(llm, true, nil) + expectAnswered(result, err) + + Expect(llm.sinkRequests).To(HaveLen(1)) + Expect(llm.sinkRequests[0].MaxTokens).To(Equal(DefaultSinkStateMaxTokens)) + }) + + It("does not cap a real tool with long arguments (counter-check)", func() { + doc := &docRunner{} + // Longer than the sink cap, so a cap leaking onto real tools would cut it. + llm := &sinkLLM{picks: []string{"write_document", "no_tool_to_call"}, sinkText: "done", docDeltas: 3 * DefaultSinkStateMaxTokens} + + result, err := runWithSink(llm, true, doc) + Expect(err).ToNot(HaveOccurred()) + + Expect(llm.docRequests).To(HaveLen(1)) + Expect(llm.docRequests[0].MaxTokens).To(BeZero(), "a real tool's parameter generation carries no sink cap") + Expect(doc.runs).To(HaveLen(1)) + Expect(doc.runs[0]).To(HaveLen(len("x ") * 3 * DefaultSinkStateMaxTokens)) + Expect(result.Status.ToolsCalled.Names()).To(ContainElement("write_document")) + + Expect(llm.sinkRequests).To(HaveLen(1)) + Expect(llm.sinkRequests[0].MaxTokens).To(Equal(DefaultSinkStateMaxTokens)) + }) + + It("sends a configured cap, the default for 0, and no step cap when negative", func() { + for _, tc := range []struct { + opt Option + want int + }{ + {WithSinkStateMaxTokens(512), 512}, + {WithSinkStateMaxTokens(0), DefaultSinkStateMaxTokens}, + {nil, DefaultSinkStateMaxTokens}, + {WithSinkStateMaxTokens(-1), 0}, + } { + llm := &sinkLLM{picks: []string{"no_tool_to_call"}, endlessSink: true} + var opts []Option + if tc.opt != nil { + opts = append(opts, tc.opt) + } + result, err := runWithSink(llm, true, nil, opts...) + expectAnswered(result, err) + Expect(llm.sinkRequests[0].MaxTokens).To(Equal(tc.want)) + if tc.want > 0 { + Expect(llm.sinkRequests).To(HaveLen(1)) + continue + } + // Without a step cap the old behaviour stays: the loop runs to the + // client cap and the length retry doubles it. + Expect(llm.sinkDeltas[0]).To(Equal(clientDefaultCap)) + Expect(llm.sinkRequests).To(HaveLen(2)) + } + }) +}) diff --git a/tools.go b/tools.go index fd1bb7d..d92821c 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,20 @@ func decisionWithStreaming(ctx context.Context, llm LLM, conversation []openai.C xlog.Debug("[decisionWithStreaming] processed", "message", content, "reasoning", reasoning) + if outCap > 0 && (finishReason == "length" || usage.CompletionTokens >= outCap) { + // 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. + // Spending the whole cap counts as cut even when the backend + // reports finish_reason=tool_calls for a call it truncated + // (observed with LocalAI); otherwise the arguments correction + // would re-run the step maxRetries times. + 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 +508,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 +568,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 || usage.CompletionTokens >= outCap) { + // 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 +667,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 @@ -651,8 +710,18 @@ func generateToolParameters(o *Options, llm LLM, tool ToolDefinitionInterface, c } } - // Use decision to force parameter generation - result, err := decisionWithStreaming(o.context, llm, conv, Tools{tool}, toolFunc.Name, o.maxRetries, o.streamCallback) + // Use decision to force parameter generation. Only the sink state's + // arguments are bounded by a cap of their own: the sink ends the loop and + // its arguments are not executed, so a cut leaves the caller's fallback + // (the arguments from the pick) without losing the answer, which the + // closing completion produces. Real tools keep the client's cap — their + // arguments may legitimately be long (a document to write). + paramCap := 0 + if o.sinkState && o.sinkStateTool != nil && o.sinkStateTool.Tool().Function != nil && + toolFunc.Name == o.sinkStateTool.Tool().Function.Name { + paramCap = o.sinkStateOutputCap() + } + result, err := decisionWithStreamingCapped(o.context, llm, conv, Tools{tool}, toolFunc.Name, o.maxRetries, o.streamCallback, paramCap) if err != nil { return nil, fmt.Errorf("failed to generate parameters for tool %s: %w", toolFunc.Name, err) } @@ -723,14 +792,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