diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index f73d882..6932298 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -89,6 +89,7 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R var allToolCalls []llm.ToolCall var totalUsage llm.TokenUsage iterations := 0 + nudges := 0 completed := false for iterations < l.cfg.MaxIters { @@ -108,6 +109,17 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R totalUsage.TotalTokens += resp.Usage.TotalTokens if len(resp.ToolCalls) == 0 { + // Same unparsed-tool-call recovery as RunStream: a tool call + // written as plain text was never executed, so ending the turn + // here would silently abandon the work mid-task. + if nudges < maxUnparsedToolCallNudges && containsUnparsedToolCall(resp.Content+resp.Reasoning) { + nudges++ + messages = append(messages, + llm.Message{Role: llm.RoleAssistant, Content: resp.Content}, + llm.Message{Role: llm.RoleUser, Content: unparsedToolCallNudge}, + ) + continue + } finalContent = resp.Content completed = true break @@ -169,6 +181,8 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa // Same as Run: the schemas are identical on every iteration. toolSchemas := l.getToolSchemas() iterations := 0 + nudges := 0 + budgetNudges := 0 for iterations < l.cfg.MaxIters { iterations++ @@ -180,12 +194,24 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa }) var hasToolCalls bool + var budgetExceeded bool var responseBuilder strings.Builder + // detectBuf collects this round's raw text (content AND + // reasoning) only to spot tool calls the model wrote as plain + // text — see the unparsed-tool-call recovery below the loop. + var detectBuf strings.Builder for chunk, err := range stream { if err != nil { yield(llm.StreamChunk{}, err) return } + if chunk.FinishReason == llm.FinishThinkingBudget { + budgetExceeded = true + } + if detectBuf.Len() < unparsedDetectBudget { + detectBuf.WriteString(chunk.ReasoningDelta) + detectBuf.WriteString(chunk.Delta) + } if len(chunk.ToolCalls) > 0 { hasToolCalls = true @@ -258,6 +284,41 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa } if !hasToolCalls { + // Recovery for a failure mode common with local models: the + // model writes its tool call as plain text — typically + // inside its reasoning block — so the server never parses + // it into a real tool call. Ending the turn here (the old + // behavior) silently abandons the work mid-task: the + // transcript reads "now I'll update X:" and then... nothing, + // because nothing was ever executed. Instead, tell the model + // what happened and let it re-issue the call properly. + if nudges < maxUnparsedToolCallNudges && containsUnparsedToolCall(detectBuf.String()) { + nudges++ + messages = append(messages, + llm.Message{Role: llm.RoleAssistant, Content: responseBuilder.String()}, + llm.Message{Role: llm.RoleUser, Content: unparsedToolCallNudge}, + ) + continue + } + // The provider cut this round because the model exceeded its + // thinking budget without ever starting an answer or a tool + // call (reasoning spiral). Ending the turn here would abandon + // the task with nothing to show for it — instead tell the + // model its reasoning was cut and demand direct action. Its + // own nudge counter, so a spiral doesn't consume the + // unparsed-tool-call retries (or vice versa). + if budgetNudges < maxThinkingBudgetNudges && budgetExceeded { + budgetNudges++ + content := responseBuilder.String() + if content == "" { + content = "(reasoning cut off: thinking budget exceeded)" + } + messages = append(messages, + llm.Message{Role: llm.RoleAssistant, Content: content}, + llm.Message{Role: llm.RoleUser, Content: thinkingBudgetNudge}, + ) + continue + } return } } @@ -266,6 +327,41 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa } } +// maxUnparsedToolCallNudges bounds how many times per turn the loop re-prompts +// a model that keeps writing tool calls as plain text, so a model that never +// gets it right can't ping-pong forever. +const maxUnparsedToolCallNudges = 2 + +// unparsedDetectBudget caps how much of a round's raw text is buffered for +// unparsed-tool-call detection — markers appear well within this. +const unparsedDetectBudget = 64 * 1024 + +// unparsedToolCallNudge is the corrective message sent when a round produced +// tool-call markup as text but no parsed tool call. +const unparsedToolCallNudge = "Your tool call was written as plain text (inside your reasoning or answer), " + + "so it was NOT executed - nothing has changed. Issue the tool call again now as a real tool call, " + + "outside of any thinking block, without re-explaining your plan." + +// maxThinkingBudgetNudges bounds how many times per turn the loop re-prompts a +// model whose reasoning was cut for exceeding the thinking budget. Separate +// from maxUnparsedToolCallNudges so one failure mode can't consume the other's +// retries. Each spiral still costs a full budget of reasoning tokens, so this +// is kept low. +const maxThinkingBudgetNudges = 2 + +// thinkingBudgetNudge is the corrective message sent when a round was cut by +// the provider's client-side thinking-budget enforcement. +const thinkingBudgetNudge = "Your reasoning exceeded the thinking budget and was cut off before you took any action. " + + "Do not re-analyze from scratch: act now on your best current plan - issue the tool call or give " + + "the final answer directly, with minimal further thinking." + +// containsUnparsedToolCall reports whether s contains tool-call markup that +// should have been parsed by the provider but wasn't (Qwen-style +// / markers are the ones seen in the wild). +func containsUnparsedToolCall(s string) bool { + return strings.Contains(s, " 0 { + if !yield(llm.StreamChunk{ToolCalls: resp.ToolCalls, FinishReason: "tool_calls"}, nil) { + return + } + } + } +} + +func (s *scriptedLLM) Name() string { return "scripted" } +func (s *scriptedLLM) Capabilities() llm.ProviderCapabilities { return llm.ProviderCapabilities{} } + +func editTestRegistry(t *testing.T, executed *int) tools.Registry { + t.Helper() + reg := tools.NewRegistry() + err := reg.Register(tools.Tool{ + Name: "edit", + Description: "edit", + InputSchema: json.RawMessage(`{"type":"object"}`), + Handler: func(_ context.Context, _ json.RawMessage) (tools.ToolResult, error) { + *executed++ + return tools.ToolResult{Content: "ok"}, nil + }, + }) + if err != nil { + t.Fatal(err) + } + return reg +} + +// TestRunStream_RecoversFromUnparsedToolCall reproduces the failure seen +// live with Qwen3.6 + llama.cpp: the model writes its tool call as plain +// text inside its reasoning ("...") so the server +// never parses it, the round has no tool calls, and the old loop simply +// ended the turn — abandoning the task mid-way with "now I'll fix X:" as the +// last words. The loop must instead nudge the model and let it re-issue the +// call for real. +func TestRunStream_RecoversFromUnparsedToolCall(t *testing.T) { + executed := 0 + stub := &scriptedLLM{responses: []llm.CompletionResponse{ + // Round 1: tool call emitted as text inside reasoning — unparsed. + {Reasoning: "I'll fix it now x.py ", Content: "Voy a corregirlo:"}, + // Round 2 (after the nudge): a real, parsed tool call. + {ToolCalls: []llm.ToolCall{{ID: "1", Name: "edit", Arguments: json.RawMessage(`{}`)}}}, + // Round 3: final answer. + {Content: "Listo, corregido."}, + }} + + loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10}) + + var final strings.Builder + for chunk, err := range loop.RunStream(context.Background(), "arregla x.py") { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + final.WriteString(chunk.Delta) + } + + if executed != 1 { + t.Fatalf("expected the re-issued tool call to execute once, got %d", executed) + } + if !strings.Contains(final.String(), "Listo, corregido.") { + t.Fatalf("expected the turn to continue to a final answer, got %q", final.String()) + } + // The corrective nudge must have been sent to the model. + foundNudge := false + for _, m := range stub.lastMessages { + if m.Role == llm.RoleUser && strings.Contains(m.Content, "NOT executed") { + foundNudge = true + } + } + if !foundNudge { + t.Fatal("expected the corrective nudge in the follow-up request messages") + } +} + +// TestRunStream_NudgeGivesUpAfterLimit keeps a model that never emits a real +// tool call from ping-ponging forever: after maxUnparsedToolCallNudges the +// turn ends normally with whatever content there is. +func TestRunStream_NudgeGivesUpAfterLimit(t *testing.T) { + executed := 0 + bad := llm.CompletionResponse{Content: "texto con falso"} + stub := &scriptedLLM{responses: []llm.CompletionResponse{bad, bad, bad, bad}} + + loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10}) + + rounds := 0 + for _, err := range loop.RunStream(context.Background(), "haz algo") { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + _ = rounds + + if stub.calls != maxUnparsedToolCallNudges+1 { + t.Fatalf("expected %d rounds (original + nudges), got %d", maxUnparsedToolCallNudges+1, stub.calls) + } + if executed != 0 { + t.Fatalf("no tool should have executed, got %d", executed) + } +} + +// TestRun_RecoversFromUnparsedToolCall covers the non-streaming path. +func TestRun_RecoversFromUnparsedToolCall(t *testing.T) { + executed := 0 + stub := &scriptedLLM{responses: []llm.CompletionResponse{ + {Content: "ahora lo edito: x.py"}, + {ToolCalls: []llm.ToolCall{{ID: "1", Name: "edit", Arguments: json.RawMessage(`{}`)}}}, + {Content: "Hecho."}, + }} + + loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10}) + resp, err := loop.Run(context.Background(), "arregla x.py") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if executed != 1 { + t.Fatalf("expected the re-issued tool call to execute once, got %d", executed) + } + if resp.Content != "Hecho." { + t.Fatalf("expected the final answer, got %q", resp.Content) + } +} diff --git a/pkg/llm/providers/llamacpp/client.go b/pkg/llm/providers/llamacpp/client.go index d2a661b..aaa6f08 100644 --- a/pkg/llm/providers/llamacpp/client.go +++ b/pkg/llm/providers/llamacpp/client.go @@ -23,6 +23,12 @@ const defaultMaxTokens = 4096 // defaultContextWindow is reported by Capabilities() when Config.ContextWindow is unset. const defaultContextWindow = 32768 +// reasoningCharsPerToken converts MaxThinkingTokens into a character budget +// for client-side enforcement (token counts aren't available per SSE delta). +// ~4 chars/token is deliberately generous for mixed Spanish/English/code, so +// the cut only ever fires later than the configured token budget, not before. +const reasoningCharsPerToken = 4 + // Config holds the settings needed to create a llama.cpp client. type Config struct { BaseURL string // defaults to http://localhost:8080/v1 @@ -36,7 +42,7 @@ type Config struct { MinP float32 // min-p sampling (llama.cpp extension) PresencePenalty float32 RepetitionPenalty float32 // sent as the server's `repeat_penalty` field - MaxThinkingTokens int // best-effort cap on reasoning tokens; ignored by servers that don't support it + MaxThinkingTokens int // cap on reasoning tokens, enforced client-side during Stream (llama.cpp ignores the JSON field, so the stream is cut and the request aborted once the estimate is exceeded); 0 = unlimited } // Client implements llm.LLMClient for llama.cpp. @@ -191,6 +197,21 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq return calls } + // Client-side thinking-budget enforcement: llama.cpp silently drops + // the max_thinking_tokens JSON field, so without this a model in a + // reasoning spiral runs until max_tokens (seen live: 25k+ tokens of + // nonstop thinking). Token counts aren't available per delta, so the + // budget is tracked as an estimate in characters; once exceeded — and + // only while the model is still purely thinking — the stream ends + // with FinishThinkingBudget and the deferred Body.Close() aborts the + // server-side generation, freeing the slot immediately. + reasoningBudget := 0 + if c.maxThinkingTokens > 0 { + reasoningBudget = c.maxThinkingTokens * reasoningCharsPerToken + } + reasoningChars := 0 + answerStarted := false + scanner := bufio.NewScanner(resp.Body) // A single SSE line can exceed bufio.Scanner's 64KB default cap // (e.g. a large tool-call arguments delta or a long reasoning @@ -261,6 +282,20 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq chunk.ToolCalls = flushToolCalls() } + reasoningChars += len(choice.Delta.ReasoningContent) + if choice.Delta.Content != "" { + answerStarted = true + } + // Cut only while the round is pure reasoning: once the answer + // or a tool call has started streaming, the spiral risk is + // over and cutting would destroy real work in flight. + if reasoningBudget > 0 && reasoningChars > reasoningBudget && + !answerStarted && len(toolCallFrags) == 0 && chunk.FinishReason == "" { + chunk.FinishReason = llm.FinishThinkingBudget + yield(chunk, nil) + return + } + // A fragment-only event (a piece of a tool call's streamed // arguments, with nothing else in this delta) has nothing // yet for the agent loop to act on: it was buffered above, diff --git a/pkg/llm/providers/llamacpp/thinking_budget_test.go b/pkg/llm/providers/llamacpp/thinking_budget_test.go new file mode 100644 index 0000000..4150047 --- /dev/null +++ b/pkg/llm/providers/llamacpp/thinking_budget_test.go @@ -0,0 +1,128 @@ +package llamacpp + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +// TestClient_Stream_ThinkingBudgetCutsPureReasoning: with MaxThinkingTokens +// set, a round that is still pure reasoning past the character budget must be +// cut with FinishThinkingBudget — and nothing after the cut may be delivered. +func TestClient_Stream_ThinkingBudgetCutsPureReasoning(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + // 5 chars per delta; budget = 2 tokens * 4 chars = 8 chars, so the + // second delta (total 10) tips it over. + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"aaaaa\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"bbbbb\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"never delivered\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"never delivered\"},\"finish_reason\":\"stop\"}]}\n")) + w.Write([]byte("data: [DONE]\n")) + })) + defer server.Close() + + client, err := New(Config{BaseURL: server.URL + "/v1", MaxThinkingTokens: 2}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var chunks []llm.StreamChunk + for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 2 { + t.Fatalf("expected 2 chunks (reasoning + budget cut), got %d: %+v", len(chunks), chunks) + } + last := chunks[len(chunks)-1] + if last.FinishReason != llm.FinishThinkingBudget { + t.Errorf("expected finish reason %q, got %q", llm.FinishThinkingBudget, last.FinishReason) + } + if last.ReasoningDelta != "bbbbb" { + t.Errorf("expected the tipping reasoning delta on the final chunk, got %q", last.ReasoningDelta) + } + for _, c := range chunks { + if c.Delta != "" { + t.Errorf("no content should have been delivered, got %q", c.Delta) + } + } +} + +// TestClient_Stream_ThinkingBudgetSparesStartedAnswer: once the model has +// begun its actual answer, exceeding the reasoning budget must NOT cut the +// stream — the spiral risk is over and real work is in flight. +func TestClient_Stream_ThinkingBudgetSparesStartedAnswer(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"aaaaa\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hola\"},\"finish_reason\":null}]}\n")) + // Over budget, but the answer already started. + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"bbbbbbbbbb\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" mundo\"},\"finish_reason\":\"stop\"}]}\n")) + w.Write([]byte("data: [DONE]\n")) + })) + defer server.Close() + + client, err := New(Config{BaseURL: server.URL + "/v1", MaxThinkingTokens: 2}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var content string + var finish string + for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + content += chunk.Delta + if chunk.FinishReason != "" { + finish = chunk.FinishReason + } + } + + if content != "Hola mundo" { + t.Errorf("expected the full answer, got %q", content) + } + if finish != "stop" { + t.Errorf("expected a normal stop, got %q", finish) + } +} + +// TestClient_Stream_NoThinkingBudgetMeansUnlimited: MaxThinkingTokens 0 keeps +// today's behavior — reasoning streams without any client-side cap. +func TestClient_Stream_NoThinkingBudgetMeansUnlimited(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n")) + w.Write([]byte("data: [DONE]\n")) + })) + defer server.Close() + + client, err := New(Config{BaseURL: server.URL + "/v1"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var content, finish string + for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + content += chunk.Delta + if chunk.FinishReason != "" { + finish = chunk.FinishReason + } + } + + if content != "ok" || finish != "stop" { + t.Errorf("expected uncut stream (content %q, finish %q), got content %q finish %q", "ok", "stop", content, finish) + } +} diff --git a/pkg/llm/types.go b/pkg/llm/types.go index 7c8df57..9bd6275 100644 --- a/pkg/llm/types.go +++ b/pkg/llm/types.go @@ -103,6 +103,13 @@ const ( StopReasonStopSeq = "stop_sequence" ) +// FinishThinkingBudget is the StreamChunk.FinishReason set by providers that +// enforce a reasoning-token budget client-side: the stream was cut because the +// model exceeded it without ever starting its answer or a tool call. Callers +// (e.g. the agent loop) can treat it as "re-prompt for a direct answer" rather +// than a normal end of turn. +const FinishThinkingBudget = "thinking_budget_exceeded" + // ToolCall represents a function invocation requested by the model. type ToolCall struct { ID string `json:"id"`