diff --git a/clients/localai_client.go b/clients/localai_client.go index 9018f52..54f7afb 100644 --- a/clients/localai_client.go +++ b/clients/localai_client.go @@ -12,6 +12,7 @@ import ( "sync" "github.com/mudler/cogito" + "github.com/mudler/xlog" "github.com/sashabaranov/go-openai" ) @@ -20,6 +21,10 @@ import ( // the model's context window once input grows past a few turns. const defaultMaxTokens = 16384 +// maxSSELineBytes bounds one SSE line. A backend may send a tool call's whole +// arguments in a single chunk (a file body, a patch), so the bound is generous. +const maxSSELineBytes = 16 << 20 + // Ensure LocalAIClient implements cogito.LLM and cogito.StreamingLLM at compile time. var _ cogito.LLM = (*LocalAIClient)(nil) var _ cogito.StreamingLLM = (*LocalAIClient)(nil) @@ -391,6 +396,7 @@ func (llm *LocalAIClient) CreateChatCompletion(ctx context.Context, request open return cogito.LLMReply{ ChatCompletionResponse: response, ReasoningContent: reasoning, + MaxTokens: max(request.MaxTokens, request.MaxCompletionTokens), }, usage, nil } @@ -516,8 +522,34 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques var lastFinishReason string var streamUsage *openai.Usage + maxTokens := request.MaxTokens + if request.MaxCompletionTokens > 0 { + maxTokens = request.MaxCompletionTokens + } + done := func() cogito.StreamEvent { + return cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage), MaxTokens: maxTokens} + } + // dropped counts data lines that did not decode, and sawToolCall + // records whether the stream produced a tool call. A dropped line may + // have carried a slice of a tool call's arguments, which would reach + // the decision incomplete but looking finished. The stream cannot know + // which call the line belonged to, so when both happened it ends with + // an ErrStreamInterrupted error instead of Done, and the decision + // retries the request. Content-only replies keep their Done: a lost + // text slice is not worth a retry. + dropped, sawToolCall := 0, false + finish := func() { + if dropped > 0 && sawToolCall { + ch <- cogito.StreamEvent{Type: cogito.StreamEventError, Error: fmt.Errorf("localai stream: dropped %d malformed chunks of a tool call: %w", dropped, cogito.ErrStreamInterrupted)} + return + } + ch <- done() + } scanner := bufio.NewScanner(resp.Body) + // Some tool parsers send a call's whole arguments in one chunk, well + // past the scanner's 64KB default token. + scanner.Buffer(make([]byte, 0, 64*1024), maxSSELineBytes) for scanner.Scan() { line := scanner.Text() @@ -527,12 +559,18 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques data := strings.TrimPrefix(line, "data: ") if data == "[DONE]" { - ch <- cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage)} + finish() return } var chunk localAIStreamChunk if err := json.Unmarshal([]byte(data), &chunk); err != nil { + dropped++ + head := data + if len(head) > 200 { + head = head[:200] + } + xlog.Warn("localai stream: dropping a chunk that is not valid JSON", "len", len(data), "head", head, "error", err) continue } // End the stream on an error chunk. Without this the chunk has no @@ -566,6 +604,7 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques // Tool call deltas for _, tc := range delta.ToolCalls { + sawToolCall = true idx := 0 if tc.Index != nil { idx = *tc.Index @@ -586,11 +625,18 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques } if err := scanner.Err(); err != nil { - ch <- cogito.StreamEvent{Type: cogito.StreamEventError, Error: err} + ch <- cogito.StreamEvent{Type: cogito.StreamEventError, Error: fmt.Errorf("localai stream: %w", err)} + return + } + // The body ended without [DONE]. Some servers omit it, so a stream + // that reported a finish_reason is complete. Without one the + // connection was cut (or a proxy timed out) mid-reply, and the partial + // reply must not look finished. + if lastFinishReason == "" { + ch <- cogito.StreamEvent{Type: cogito.StreamEventError, Error: fmt.Errorf("localai stream: body ended without [DONE] or a finish_reason: %w", cogito.ErrStreamInterrupted)} return } - // If we reach here without [DONE], still emit done - ch <- cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage)} + finish() }() return ch, nil diff --git a/clients/localai_stream_integrity_test.go b/clients/localai_stream_integrity_test.go new file mode 100644 index 0000000..7ab2be2 --- /dev/null +++ b/clients/localai_stream_integrity_test.go @@ -0,0 +1,126 @@ +package clients + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/mudler/cogito" + "github.com/sashabaranov/go-openai" +) + +// rawSSEServer writes the given body verbatim as an SSE stream and closes the +// connection, so a test controls exactly how the stream ends. +func rawSSEServer(body string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(body)) + })) +} + +// collectStream runs one streaming request against body and returns the +// events the client emitted. +func collectStream(t *testing.T, body string) []cogito.StreamEvent { + t.Helper() + srv := rawSSEServer(body) + defer srv.Close() + llm := NewLocalAILLM("m", "", srv.URL) + ch, err := llm.CreateChatCompletionStream(context.Background(), openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessage{{Role: "user", Content: "hi"}}, + }) + if err != nil { + t.Fatalf("CreateChatCompletionStream: %v", err) + } + var events []cogito.StreamEvent + for ev := range ch { + events = append(events, ev) + } + if len(events) == 0 { + t.Fatal("stream emitted no events") + } + return events +} + +func lastEvent(events []cogito.StreamEvent) cogito.StreamEvent { return events[len(events)-1] } + +func toolArgs(events []cogito.StreamEvent) string { + var b strings.Builder + for _, ev := range events { + if ev.Type == cogito.StreamEventToolCall { + b.WriteString(ev.ToolArgs) + } + } + return b.String() +} + +// A connection cut mid tool call (no finish_reason, no [DONE]) must not look +// like a finished stream: the decision would parse the partial arguments. +func TestStreamEndedWithoutDoneOrFinishReasonIsInterrupted(t *testing.T) { + events := collectStream(t, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"search","arguments":"{\"query\":"}}]}}]}`+"\n\n") + + last := lastEvent(events) + if last.Type != cogito.StreamEventError { + t.Fatalf("want an error event at the end of an interrupted stream, got %q", last.Type) + } + if !errors.Is(last.Error, cogito.ErrStreamInterrupted) { + t.Fatalf("want errors.Is(err, ErrStreamInterrupted), got %v", last.Error) + } +} + +// Some servers omit [DONE]; a stream that reported a finish_reason is complete. +func TestStreamEndedWithoutDoneButWithFinishReasonIsDone(t *testing.T) { + events := collectStream(t, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"search","arguments":"{\"query\":\"x\"}"}}]}}]}`+"\n\n"+ + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}`+"\n\n") + + last := lastEvent(events) + if last.Type != cogito.StreamEventDone { + t.Fatalf("want a done event, got %q (err %v)", last.Type, last.Error) + } + if last.FinishReason != "tool_calls" { + t.Fatalf("want finish_reason tool_calls, got %q", last.FinishReason) + } + if last.MaxTokens != defaultMaxTokens { + t.Fatalf("want the done event to report the cap sent (%d), got %d", defaultMaxTokens, last.MaxTokens) + } +} + +// Some tool parsers send a call's whole arguments in one chunk, which can be +// far larger than bufio.Scanner's 64KB default token. +func TestStreamAcceptsLargeSingleChunk(t *testing.T) { + big := strings.Repeat("a", 200*1024) + events := collectStream(t, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"write","arguments":"{\"text\":\"`+big+`\"}"}}]},"finish_reason":"tool_calls"}]}`+"\n\n"+ + "data: [DONE]\n\n") + + last := lastEvent(events) + if last.Type != cogito.StreamEventDone { + t.Fatalf("want a done event, got %q (err %v)", last.Type, last.Error) + } + if got, want := len(toolArgs(events)), len(big)+len(`{"text":""}`); got != want { + t.Fatalf("want %d bytes of tool arguments, got %d", want, got) + } +} + +// A chunk that does not decode may have carried a slice of the arguments; the +// stream must say it is lossy rather than finish with incomplete arguments. +func TestStreamMalformedToolCallChunkIsNotSilentlyDropped(t *testing.T) { + events := collectStream(t, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"search","arguments":"{\"query\":"}}]}}]}`+"\n\n"+ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"x`+"\n\n"+ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"}"}}]}}]}`+"\n\n"+ + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}`+"\n\n"+ + "data: [DONE]\n\n") + + last := lastEvent(events) + if last.Type != cogito.StreamEventError { + t.Fatalf("want an error event after a dropped tool-call chunk, got %q", last.Type) + } + if !errors.Is(last.Error, cogito.ErrStreamInterrupted) { + t.Fatalf("want errors.Is(err, ErrStreamInterrupted), got %v", last.Error) + } +} diff --git a/clients/openai_client.go b/clients/openai_client.go index 1c5d8fb..7c6ae65 100644 --- a/clients/openai_client.go +++ b/clients/openai_client.go @@ -138,6 +138,7 @@ func (llm *OpenAIClient) CreateChatCompletion(ctx context.Context, request opena return cogito.LLMReply{ ChatCompletionResponse: response, ReasoningContent: response.Choices[0].Message.ReasoningContent, + MaxTokens: max(request.MaxTokens, request.MaxCompletionTokens), }, usage, nil } @@ -181,8 +182,12 @@ func (llm *OpenAIClient) CreateChatCompletionStream(ctx context.Context, request for { resp, err := stream.Recv() + // go-openai returns io.EOF both after [DONE] and when the body + // ends without it, so a cut connection cannot be told from a + // finished stream here, and a server that sends [DONE] without a + // finish_reason is valid. This path is left as it was. if errors.Is(err, io.EOF) { - ch <- cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage)} + ch <- cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage), MaxTokens: max(request.MaxTokens, request.MaxCompletionTokens)} return } if err != nil { diff --git a/decision_length_retry_test.go b/decision_length_retry_test.go index 69befa5..bfdc69f 100644 --- a/decision_length_retry_test.go +++ b/decision_length_retry_test.go @@ -191,9 +191,10 @@ func TestStreamingDecisionLengthRetryCountsFailedAttemptUsage(t *testing.T) { // scriptedChatLLM is the non-streaming counterpart of scriptedStreamLLM. type scriptedChatLLM struct { - replies []openai.ChatCompletionChoice - usages []LLMUsage - requests []openai.ChatCompletionRequest + replies []openai.ChatCompletionChoice + usages []LLMUsage + reasoning string // ReasoningContent of every reply + requests []openai.ChatCompletionRequest } func (m *scriptedChatLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { return f, nil } @@ -204,7 +205,7 @@ func (m *scriptedChatLLM) CreateChatCompletion(ctx context.Context, request open if i >= len(m.replies) { i = len(m.replies) - 1 } - return LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{m.replies[i]}}}, m.usages[i], nil + return LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{m.replies[i]}}, ReasoningContent: m.reasoning}, m.usages[i], nil } func lengthChoice() openai.ChatCompletionChoice { diff --git a/decision_tool_args_test.go b/decision_tool_args_test.go new file mode 100644 index 0000000..790385d --- /dev/null +++ b/decision_tool_args_test.go @@ -0,0 +1,376 @@ +package cogito + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "testing" + + "github.com/sashabaranov/go-openai" +) + +func streamedToolCall(id, args, finish string, completionTokens int) []StreamEvent { + return []StreamEvent{ + {Type: StreamEventToolCall, ToolCallIndex: 0, ToolCallID: id, ToolName: "search", ToolArgs: args}, + {Type: StreamEventDone, FinishReason: finish, Usage: LLMUsage{PromptTokens: 1000, CompletionTokens: completionTokens, TotalTokens: 1000 + completionTokens}}, + } +} + +func chatToolCall(id, args string, finish openai.FinishReason) openai.ChatCompletionChoice { + return openai.ChatCompletionChoice{FinishReason: finish, Message: openai.ChatCompletionMessage{ + Role: "assistant", + ToolCalls: []openai.ToolCall{{ID: id, Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{Name: "search", Arguments: args}}}, + }} +} + +func noopCB(StreamEvent) {} + +// correctionPair returns the assistant tool call and the tool-role reply the +// decision appended after a malformed attempt, failing when either is missing. +func correctionPair(t *testing.T, req openai.ChatCompletionRequest, wantArgs string) (openai.ToolCall, openai.ChatCompletionMessage) { + t.Helper() + msgs := req.Messages + if len(msgs) < 2 { + t.Fatalf("retry request has %d messages, want the correction pair appended", len(msgs)) + } + asst, tool := msgs[len(msgs)-2], msgs[len(msgs)-1] + if asst.Role != openai.ChatMessageRoleAssistant || len(asst.ToolCalls) != 1 { + t.Fatalf("want an assistant message with the malformed tool call, got %+v", asst) + } + tc := asst.ToolCalls[0] + // The history must stay valid JSON: some chat templates (llama.cpp) + // parse earlier tool-call arguments and would reject the retry. + if tc.Function.Arguments != "{}" || tc.Function.Name != "search" { + t.Fatalf("want the echoed call to be search with arguments {}, got %+v", tc) + } + for _, m := range req.Messages { + for _, c := range m.ToolCalls { + if !json.Valid([]byte(c.Function.Arguments)) { + t.Fatalf("retry request carries an assistant tool call with invalid JSON arguments: %q", c.Function.Arguments) + } + } + } + if tc.ID == "" { + t.Fatal("the malformed tool call must carry an ID so the tool reply can pair with it") + } + if tool.Role != openai.ChatMessageRoleTool || tool.ToolCallID != tc.ID { + t.Fatalf("want a tool-role reply for call %q, got %+v", tc.ID, tool) + } + if !strings.Contains(tool.Content, "You sent: "+wantArgs+".") { + t.Fatalf("tool reply must quote the malformed arguments %q, got %q", wantArgs, tool.Content) + } + if !strings.Contains(tool.Content, "not valid JSON") || !strings.Contains(tool.Content, "Call the tool again with valid JSON arguments.") { + t.Fatalf("tool reply does not explain the failure: %q", tool.Content) + } + return tc, tool +} + +func TestStreamingDecisionEmptyArgumentsAreEmptyObject(t *testing.T) { + for _, args := range []string{"", " \n"} { + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{streamedToolCall("c1", args, "tool_calls", 10)}} + res, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 3, noopCB) + if err != nil { + t.Fatalf("args %q: unexpected error: %v", args, err) + } + if len(res.toolChoices) != 1 || res.toolChoices[0].Arguments == nil || len(res.toolChoices[0].Arguments) != 0 { + t.Fatalf("args %q: want one tool choice with an empty map, got %+v", args, res.toolChoices) + } + if len(llm.requests) != 1 { + t.Fatalf("args %q: want no retry, got %d attempts", args, len(llm.requests)) + } + } +} + +func TestDecisionEmptyArgumentsAreEmptyObject(t *testing.T) { + llm := &scriptedChatLLM{ + replies: []openai.ChatCompletionChoice{chatToolCall("c1", "", openai.FinishReasonToolCalls)}, + usages: []LLMUsage{{}}, + } + res, err := decision(context.Background(), llm, lengthRetryConv, Tools{}, "", 3) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(res.toolChoices) != 1 || res.toolChoices[0].Arguments == nil || len(res.toolChoices[0].Arguments) != 0 { + t.Fatalf("want one tool choice with an empty map, got %+v", res.toolChoices) + } + if len(llm.requests) != 1 { + t.Fatalf("want no retry, got %d attempts", len(llm.requests)) + } +} + +func TestStreamingDecisionTruncatedArgumentsRetryWithRaisedCap(t *testing.T) { + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{ + streamedToolCall("c1", `{"query":"x`, "length", 16384), + streamedToolCall("c1", `{"query":"x"}`, "tool_calls", 20000), + }} + res, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 1, noopCB) + if err != nil { + t.Fatalf("want success after one raised-cap retry, got %v", err) + } + if len(res.toolChoices) != 1 || res.toolChoices[0].Arguments["query"] != "x" { + t.Fatalf("want the retried tool call, got %+v", res.toolChoices) + } + if len(llm.requests) != 2 { + t.Fatalf("want exactly 2 attempts, got %d", len(llm.requests)) + } + if first, second := llm.requests[0].MaxTokens, llm.requests[1].MaxTokens; second <= first || second != 32768 { + t.Fatalf("want the retry to raise max tokens (first %d, second %d, want 32768)", first, second) + } + // A truncated call is not a malformed one: no correction is appended. + if len(llm.requests[1].Messages) != len(lengthRetryConv) { + t.Fatalf("the length retry must resend the same messages, got %d", len(llm.requests[1].Messages)) + } +} + +func TestStreamingDecisionTruncatedArgumentsTwiceFails(t *testing.T) { + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{ + streamedToolCall("c1", `{"query":"x`, "length", 16384), + streamedToolCall("c1", `{"query":"xy`, "length", 32768), + }} + _, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 5, noopCB) + if !errors.Is(err, ErrToolArgumentsTruncated) { + t.Fatalf("want ErrToolArgumentsTruncated, got %v", err) + } + if !strings.Contains(err.Error(), "failed to make a streaming decision after") { + t.Fatalf("want the decision error shape, got %q", err) + } + var te *ToolArgumentsTruncatedError + if !errors.As(err, &te) { + t.Fatalf("want errors.As to reach *ToolArgumentsTruncatedError, got %v", err) + } + if te.PromptTokens != 1000 || te.CompletionTokens != 32768 || te.MaxTokens != 32768 { + t.Fatalf("want usage of the last truncated attempt (prompt 1000, completion 32768, max 32768), got %+v", *te) + } + if te.ToolName != "search" || te.ArgumentsBytes != len(`{"query":"xy`) || te.ReasoningBytes != 0 { + t.Fatalf("want the last attempt's tool call figures, got %+v", *te) + } + if len(llm.requests) != 2 { + t.Fatalf("want exactly one length retry (2 attempts), got %d", len(llm.requests)) + } +} + +func TestStreamingDecisionMalformedArgumentsSelfCorrect(t *testing.T) { + conv := []openai.ChatCompletionMessage{{Role: "user", Content: "find x"}} + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{ + streamedToolCall("c1", `{"query": x}`, "tool_calls", 10), + streamedToolCall("c2", `{"query":"x"}`, "tool_calls", 10), + }} + res, err := decisionWithStreaming(context.Background(), llm, conv, Tools{}, "", 3, noopCB) + if err != nil { + t.Fatalf("want success on the corrected attempt, got %v", err) + } + if len(res.toolChoices) != 1 || res.toolChoices[0].Arguments["query"] != "x" { + t.Fatalf("want the corrected tool call, got %+v", res.toolChoices) + } + if len(llm.requests) != 2 { + t.Fatalf("want 2 attempts, got %d", len(llm.requests)) + } + if len(llm.requests[0].Messages) != 1 { + t.Fatalf("the first request must not carry a correction, got %d messages", len(llm.requests[0].Messages)) + } + tc, _ := correctionPair(t, llm.requests[1], `{"query": x}`) + if tc.ID != "c1" { + t.Fatalf("want the backend's call id kept, got %q", tc.ID) + } + if len(conv) != 1 { + t.Fatalf("the caller's conversation must not change, got %d messages", len(conv)) + } +} + +func TestStreamingDecisionMalformedArgumentsWithoutIDGetsOne(t *testing.T) { + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{ + streamedToolCall("", `{"query"`, "tool_calls", 10), + streamedToolCall("", `{"query":"x"}`, "tool_calls", 10), + }} + if _, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 3, noopCB); err != nil { + t.Fatalf("unexpected error: %v", err) + } + correctionPair(t, llm.requests[1], `{"query"`) +} + +func TestStreamingDecisionMalformedArgumentsPersistentFails(t *testing.T) { + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{streamedToolCall("c1", `{"query": x}`, "tool_calls", 10)}} + _, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 3, noopCB) + if !errors.Is(err, ErrToolArgumentsInvalid) { + t.Fatalf("want ErrToolArgumentsInvalid, got %v", err) + } + if !strings.Contains(err.Error(), "failed to make a streaming decision after 3 attempts") { + t.Fatalf("want the decision error shape, got %q", err) + } + if len(llm.requests) != 3 { + t.Fatalf("want 3 attempts, got %d", len(llm.requests)) + } +} + +func TestStreamingDecisionInterruptedStreamIsRetriedAndReported(t *testing.T) { + interrupted := []StreamEvent{ + {Type: StreamEventToolCall, ToolCallIndex: 0, ToolCallID: "c1", ToolName: "search", ToolArgs: `{"query":`}, + {Type: StreamEventError, Error: fmt.Errorf("localai stream: %w", ErrStreamInterrupted)}, + } + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{interrupted}} + _, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 2, noopCB) + if !errors.Is(err, ErrStreamInterrupted) { + t.Fatalf("want ErrStreamInterrupted, got %v", err) + } + if len(llm.requests) != 2 { + t.Fatalf("want the interruption retried (2 attempts), got %d", len(llm.requests)) + } + + llm = &scriptedStreamLLM{scripts: [][]StreamEvent{interrupted, streamedToolCall("c1", `{"query":"x"}`, "tool_calls", 10)}} + if _, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 2, noopCB); err != nil { + t.Fatalf("want success after an interrupted attempt, got %v", err) + } +} + +func TestDecisionTruncatedArguments(t *testing.T) { + llm := &scriptedChatLLM{ + replies: []openai.ChatCompletionChoice{ + chatToolCall("c1", `{"query":"x`, openai.FinishReasonLength), + chatToolCall("c1", `{"query":"x"}`, openai.FinishReasonToolCalls), + }, + usages: []LLMUsage{{CompletionTokens: 16384}, {CompletionTokens: 100}}, + } + if _, err := decision(context.Background(), llm, lengthRetryConv, Tools{}, "", 1); err != nil { + t.Fatalf("want success after one raised-cap retry, got %v", err) + } + if len(llm.requests) != 2 || llm.requests[1].MaxTokens != 32768 { + t.Fatalf("want 2 attempts with the second at 32768, got %d attempts", len(llm.requests)) + } + + llm = &scriptedChatLLM{ + replies: []openai.ChatCompletionChoice{ + chatToolCall("c1", `{"query":"x`, openai.FinishReasonLength), + chatToolCall("c1", `{"query":"xyz`, openai.FinishReasonLength), + }, + usages: []LLMUsage{{PromptTokens: 500, CompletionTokens: 16384}, {PromptTokens: 500, CompletionTokens: 32768}}, + reasoning: "long reasoning", + } + _, err := decision(context.Background(), llm, lengthRetryConv, Tools{}, "", 5) + if !errors.Is(err, ErrToolArgumentsTruncated) { + t.Fatalf("want ErrToolArgumentsTruncated, got %v", err) + } + var te *ToolArgumentsTruncatedError + if !errors.As(err, &te) || te.PromptTokens != 500 || te.CompletionTokens != 32768 || te.MaxTokens != 32768 { + t.Fatalf("want the last attempt's figures through errors.As, got %v (%+v)", err, te) + } + if te.ToolName != "search" || te.ArgumentsBytes != len(`{"query":"xyz`) || te.ReasoningBytes != len("long reasoning") { + t.Fatalf("want the last attempt's tool call and reasoning figures, got %+v", *te) + } + if len(llm.requests) != 2 { + t.Fatalf("want exactly 2 attempts, got %d", len(llm.requests)) + } +} + +func TestDecisionMalformedArguments(t *testing.T) { + llm := &scriptedChatLLM{ + replies: []openai.ChatCompletionChoice{ + chatToolCall("", `{"query": x}`, openai.FinishReasonToolCalls), + chatToolCall("c2", `{"query":"x"}`, openai.FinishReasonToolCalls), + }, + usages: []LLMUsage{{}, {}}, + } + res, err := decision(context.Background(), llm, lengthRetryConv, Tools{}, "", 3) + if err != nil { + t.Fatalf("want success on the corrected attempt, got %v", err) + } + if len(res.toolChoices) != 1 { + t.Fatalf("want the corrected tool call, got %+v", res.toolChoices) + } + correctionPair(t, llm.requests[1], `{"query": x}`) + + llm = &scriptedChatLLM{ + replies: []openai.ChatCompletionChoice{chatToolCall("c1", `{"query": x}`, openai.FinishReasonToolCalls)}, + usages: []LLMUsage{{}}, + } + _, err = decision(context.Background(), llm, lengthRetryConv, Tools{}, "", 3) + if !errors.Is(err, ErrToolArgumentsInvalid) { + t.Fatalf("want ErrToolArgumentsInvalid, got %v", err) + } + if !strings.Contains(err.Error(), "failed to make a decision after 3 attempts") { + t.Fatalf("want the decision error shape, got %q", err) + } + if len(llm.requests) != 3 { + t.Fatalf("want 3 attempts, got %d", len(llm.requests)) + } +} + +// windowStopped is a truncated tool call whose Done reports the cap the +// client sent: completion stopped below it, so the context window ran out. +func streamedTruncatedWithCap(completionTokens, maxTokens int) []StreamEvent { + ev := streamedToolCall("c1", `{"query":"x`, "length", completionTokens) + ev[len(ev)-1].MaxTokens = maxTokens + return append([]StreamEvent{{Type: StreamEventReasoning, Content: "thinking"}, {Type: StreamEventReasoning, Content: " more"}}, ev...) +} + +func TestStreamingDecisionTruncatedArgumentsWindowStoppedFailsFast(t *testing.T) { + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{ + streamedTruncatedWithCap(60000, 90000), + streamedToolCall("c1", `{"query":"x"}`, "tool_calls", 10), + }} + _, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 3, noopCB) + if !errors.Is(err, ErrToolArgumentsTruncated) { + t.Fatalf("want ErrToolArgumentsTruncated, got %v", err) + } + var te *ToolArgumentsTruncatedError + if !errors.As(err, &te) || te.PromptTokens != 1000 || te.CompletionTokens != 60000 || te.MaxTokens != 90000 { + t.Fatalf("want the attempt's figures (1000/60000/90000) through errors.As, got %v (%+v)", err, te) + } + if te.ToolName != "search" || te.ArgumentsBytes != len(`{"query":"x`) || te.ReasoningBytes != len("thinking more") { + t.Fatalf("want tool search, %d argument bytes, %d reasoning bytes, got %+v", len(`{"query":"x`), len("thinking more"), *te) + } + if len(llm.requests) != 1 { + t.Fatalf("a window-stopped truncation must not be retried, got %d attempts", len(llm.requests)) + } +} + +func TestStreamingDecisionTruncatedArgumentsCapReachedRetries(t *testing.T) { + llm := &scriptedStreamLLM{scripts: [][]StreamEvent{ + streamedTruncatedWithCap(16384, 16384), + streamedToolCall("c1", `{"query":"x"}`, "tool_calls", 10), + }} + if _, err := decisionWithStreaming(context.Background(), llm, lengthRetryConv, Tools{}, "", 1, noopCB); err != nil { + t.Fatalf("want success after a raised-cap retry, got %v", err) + } + if len(llm.requests) != 2 || llm.requests[1].MaxTokens != 32768 { + t.Fatalf("want one retry at 32768, got %d attempts", len(llm.requests)) + } +} + +// The raised cap is on the request itself, so a retry that stops below it ran +// out of context window: fail at once instead of retrying again. +func TestDecisionTruncatedArgumentsWindowStoppedAfterRaise(t *testing.T) { + llm := &scriptedChatLLM{ + replies: []openai.ChatCompletionChoice{ + chatToolCall("c1", `{"query":"x`, openai.FinishReasonLength), + chatToolCall("c1", `{"query":"x`, openai.FinishReasonLength), + }, + usages: []LLMUsage{{PromptTokens: 500, CompletionTokens: 16384}, {PromptTokens: 500, CompletionTokens: 20000}}, + } + _, err := decision(context.Background(), llm, lengthRetryConv, Tools{}, "", 5) + var te *ToolArgumentsTruncatedError + if !errors.As(err, &te) || te.CompletionTokens != 20000 || te.MaxTokens != 32768 { + t.Fatalf("want the window-stopped attempt's figures, got %v (%+v)", err, te) + } + if len(llm.requests) != 2 { + t.Fatalf("want 2 attempts, got %d", len(llm.requests)) + } +} + +func TestArgumentsCorrectionQuotesLongArgumentsHeadAndTail(t *testing.T) { + raw := "{" + strings.Repeat("h", 400) + strings.Repeat("t", 400) + msgs := appendArgumentsCorrection(nil, "", &badToolCall{ + call: openai.ToolCall{ID: "c1", Function: openai.FunctionCall{Name: "search", Arguments: raw}}, + err: errors.New("unexpected end of JSON input"), + }) + got := msgs[1].Content + want := "You sent: " + raw[:300] + "…[301 bytes omitted]…" + raw[len(raw)-200:] + "." + if !strings.Contains(got, want) { + t.Fatalf("want the head, an omission marker and the tail, got %q", got) + } + if msgs[0].ToolCalls[0].Function.Arguments != "{}" { + t.Fatalf("want the echoed arguments to be {}, got %q", msgs[0].ToolCalls[0].Function.Arguments) + } +} diff --git a/go.mod b/go.mod index 3031b33..27f89a5 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,7 @@ require ( github.com/docker/go-connections v0.5.0 github.com/google/uuid v1.6.0 github.com/modelcontextprotocol/go-sdk v1.0.0 + github.com/mudler/xlog v0.0.1 github.com/onsi/ginkgo/v2 v2.25.3 github.com/onsi/gomega v1.38.2 github.com/sashabaranov/go-openai v1.41.2 @@ -57,7 +58,6 @@ require ( github.com/moby/sys/userns v0.1.0 // indirect github.com/moby/term v0.5.0 // indirect github.com/morikuni/aec v1.0.0 // indirect - github.com/mudler/xlog v0.0.1 // indirect github.com/opencontainers/go-digest v1.0.0 // indirect github.com/opencontainers/image-spec v1.1.1 // indirect github.com/pkg/errors v0.9.1 // indirect diff --git a/llm.go b/llm.go index b634b9a..6ae61d3 100644 --- a/llm.go +++ b/llm.go @@ -28,4 +28,8 @@ type StreamingLLM interface { type LLMReply struct { ChatCompletionResponse openai.ChatCompletionResponse ReasoningContent string + // MaxTokens is the output cap the client sent with the request, 0 when + // unknown. The decision uses it to tell a context-window stop from a cap + // stop when the request itself carried no cap. + MaxTokens int } diff --git a/stream.go b/stream.go index 20b030e..6453d17 100644 --- a/stream.go +++ b/stream.go @@ -26,6 +26,7 @@ type StreamEvent struct { FinishReason string // "stop", "tool_calls", etc. (populated on done) Error error // populated on error Usage LLMUsage // populated on done + MaxTokens int // on done: the output cap the client sent, 0 when unknown AgentID string // populated for sub-agent events } diff --git a/tool_args.go b/tool_args.go new file mode 100644 index 0000000..f777123 --- /dev/null +++ b/tool_args.go @@ -0,0 +1,171 @@ +package cogito + +import ( + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "slices" + "strings" + + "github.com/mudler/xlog" + "github.com/sashabaranov/go-openai" +) + +var ( + // ErrStreamInterrupted reports a stream that ended before the backend + // finished it: the body closed with no [DONE] and no finish_reason, or a + // chunk was lost, so the accumulated reply is incomplete. + ErrStreamInterrupted = errors.New("stream interrupted before the backend finished") + // ErrToolArgumentsTruncated reports a tool call whose arguments were cut by + // the output limit (finish_reason=length), with no retry left that could + // help. The error that wraps it is a *ToolArgumentsTruncatedError. + ErrToolArgumentsTruncated = errors.New("tool call arguments truncated by the output limit") + // ErrToolArgumentsInvalid reports tool-call arguments that were not valid + // JSON after every attempt. + ErrToolArgumentsInvalid = errors.New("tool call arguments are not valid JSON") +) + +// ToolArgumentsTruncatedError describes the last attempt whose tool call was +// cut by finish_reason=length. The figures let a caller tell the two causes +// apart: CompletionTokens below MaxTokens means the context window ran out, +// not the cap; a large ReasoningBytes against a small ArgumentsBytes means the +// reasoning, not the tool call, used the room. A figure is 0 when unknown. +type ToolArgumentsTruncatedError struct { + PromptTokens int + CompletionTokens int + MaxTokens int // the output cap the attempt carried + ToolName string // the first tool call whose arguments failed to parse + ArgumentsBytes int // length of that call's arguments + ReasoningBytes int // length of the reasoning of that attempt +} + +func (e *ToolArgumentsTruncatedError) Error() string { + cause := "the output cap was reached" + if e.MaxTokens > 0 && e.CompletionTokens > 0 && e.CompletionTokens < e.MaxTokens { + cause = "the context window ran out before the output cap" + } + return fmt.Sprintf("%s: tool %q (finish_reason=length, %s; prompt tokens %d, completion tokens %d, max tokens %d, reasoning %d bytes, arguments %d bytes)", + ErrToolArgumentsTruncated, e.ToolName, cause, e.PromptTokens, e.CompletionTokens, e.MaxTokens, e.ReasoningBytes, e.ArgumentsBytes) +} + +func (e *ToolArgumentsTruncatedError) Is(target error) bool { + return target == ErrToolArgumentsTruncated +} + +// badToolCall is the first tool call of an attempt whose arguments did not parse. +type badToolCall struct { + call openai.ToolCall + err error +} + +// parseToolCalls parses every call's arguments. Empty or blank arguments mean +// no parameters (some backends send "" for a tool without parameters). It +// stops at the first call that fails and returns it. +func parseToolCalls(calls []openai.ToolCall, finishReason string) ([]*ToolChoice, *badToolCall) { + choices := make([]*ToolChoice, 0, len(calls)) + for _, tc := range calls { + arguments := make(map[string]any) + raw := tc.Function.Arguments + if strings.TrimSpace(raw) != "" { + if err := json.Unmarshal([]byte(raw), &arguments); err != nil { + tail := raw + if len(tail) > 80 { + tail = tail[len(tail)-80:] + } + xlog.Warn("Tool call arguments are not valid JSON", "tool", tc.Function.Name, + "finishReason", finishReason, "argumentsLen", len(raw), "tail", tail, "error", err) + return nil, &badToolCall{call: tc, err: err} + } + } + choices = append(choices, &ToolChoice{Name: tc.Function.Name, Arguments: arguments}) + } + return choices, nil +} + +// outputCap returns the output cap an attempt ran with: the one on the +// request, else the one the client reported sending, else 0. +func outputCap(req openai.ChatCompletionRequest, reported int) int { + if req.MaxCompletionTokens > 0 { + return req.MaxCompletionTokens + } + if req.MaxTokens > 0 { + return req.MaxTokens + } + return reported +} + +// truncatedToolCall handles a tool call cut by finish_reason=length. When the +// completion stopped below the cap, the context window ran out and a larger +// cap cannot help, so it fails at once. Otherwise (cap reached, or the figures +// are unknown) it raises the cap once through l and returns nil; when that +// retry is spent or cannot raise the cap, it fails. +func (l *lengthRetry) truncatedToolCall(req *openai.ChatCompletionRequest, usage LLMUsage, reportedCap int, bad *badToolCall, reasoning string) error { + te := &ToolArgumentsTruncatedError{ + PromptTokens: usage.PromptTokens, + CompletionTokens: usage.CompletionTokens, + MaxTokens: outputCap(*req, reportedCap), + ToolName: bad.call.Function.Name, + ArgumentsBytes: len(bad.call.Function.Arguments), + ReasoningBytes: len(reasoning), + } + if te.MaxTokens > 0 && te.CompletionTokens > 0 && te.CompletionTokens < te.MaxTokens { + return te + } + if l.raised != 0 { + return te + } + n := lengthRetryOutputCap(*req, usage) + if n == 0 { + return te + } + xlog.Warn("Tool call arguments truncated by length, retrying with a larger output cap", "tool", te.ToolName, "maxTokens", n) + setOutputCap(req, n) + l.raised = n + return nil +} + +// appendArgumentsCorrection returns messages with the malformed call and a +// tool reply that names the parse error, so the next attempt can correct +// itself instead of repeating an identical request. It never writes into the +// backing array of messages, which may be the caller's. +func appendArgumentsCorrection(messages []openai.ChatCompletionMessage, content string, bad *badToolCall) []openai.ChatCompletionMessage { + call := bad.call + if call.ID == "" { + call.ID = newToolCallID() + } + if call.Type == "" { + call.Type = openai.ToolTypeFunction + } + // Echo the call with empty arguments: chat templates such as llama.cpp's + // parse earlier tool-call arguments as JSON, and the malformed text would + // make the retry itself fail. The tool reply quotes the text instead. + call.Function.Arguments = "{}" + return append(slices.Clip(messages), + openai.ChatCompletionMessage{Role: openai.ChatMessageRoleAssistant, Content: content, ToolCalls: []openai.ToolCall{call}}, + openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleTool, + ToolCallID: call.ID, + Name: call.Function.Name, + Content: fmt.Sprintf("the arguments of this call were not valid JSON (%v). You sent: %s. Call the tool again with valid JSON arguments.", + bad.err, abbreviate(bad.call.Function.Arguments, 300, 200)), + }, + ) +} + +// abbreviate keeps the first head and last tail bytes of s, with a marker +// naming how much it left out between them. +func abbreviate(s string, head, tail int) string { + if len(s) <= head+tail { + return s + } + omitted := len(s) - head - tail + return strings.ToValidUTF8(s[:head], "") + fmt.Sprintf("…[%d bytes omitted]…", omitted) + strings.ToValidUTF8(s[len(s)-tail:], "") +} + +func newToolCallID() string { + var b [12]byte + _, _ = rand.Read(b[:]) + return "call_" + hex.EncodeToString(b[:]) +} diff --git a/tools.go b/tools.go index 09f6ecf..b361454 100644 --- a/tools.go +++ b/tools.go @@ -331,6 +331,7 @@ func decisionWithStreaming(ctx context.Context, llm LLM, conversation []openai.C var streamErr error var usage LLMUsage var finishReason string + var reportedCap int for ev := range ch { streamCB(ev) @@ -359,6 +360,7 @@ func decisionWithStreaming(ctx context.Context, llm LLM, conversation []openai.C case StreamEventDone: usage = ev.Usage finishReason = ev.FinishReason + reportedCap = ev.MaxTokens case StreamEventError: streamErr = ev.Error } @@ -417,27 +419,24 @@ func decisionWithStreaming(ctx context.Context, llm LLM, conversation []openai.C return &decisionResult{message: content, reasoning: reasoning, usage: usage}, nil } - // Process all tool calls - toolChoices := make([]*ToolChoice, 0, len(toolCalls)) - allParsed := true - for _, toolCall := range toolCalls { - arguments := make(map[string]any) - if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &arguments); err != nil { - lastErr = err - xlog.Warn("Attempt to parse streamed tool arguments failed", "attempt", attempts+1, "error", err) - allParsed = false - break - } - toolChoices = append(toolChoices, &ToolChoice{ - Name: toolCall.Function.Name, - Arguments: arguments, - }) - } - - if !allParsed { - if werr := backoffOrCancel(ctx, attempts); werr != nil { - return nil, werr + toolChoices, bad := parseToolCalls(toolCalls, finishReason) + if bad != nil { + if finishReason == "length" { + // The output limit cut the call. An identical request is cut + // identically, so raise the cap once (outside the attempt + // budget), or fail at once when the context window, not the + // cap, stopped it. + if err := lengthRetried.truncatedToolCall(&req, usage, reportedCap, bad, reasoning); err != nil { + return nil, fmt.Errorf("failed to make a streaming decision after %d attempts: %w", attempts+1, err) + } + budget++ + continue } + // Malformed JSON with a normal finish: show the model its call and + // the parse error so the next attempt can correct it. No backoff, + // the failure is not transient. + lastErr = fmt.Errorf("%w: tool %q: %v", ErrToolArgumentsInvalid, bad.call.Function.Name, bad.err) + req.Messages = appendArgumentsCorrection(req.Messages, content, bad) continue } @@ -535,38 +534,29 @@ func decision(ctx context.Context, llm LLM, conversation []openai.ChatCompletion return &decisionResult{message: msg.Content, reasoning: reasoning, usage: usage}, nil } - // Process all tool calls - toolChoices := make([]*ToolChoice, 0, len(msg.ToolCalls)) - for _, toolCall := range msg.ToolCalls { - arguments := make(map[string]any) - - if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &arguments); err != nil { - lastErr = err - xlog.Warn("Attempt to parse tool arguments failed", "attempt", attempts+1, "error", err) - if werr := backoffOrCancel(ctx, attempts); werr != nil { - return nil, werr + finishReason := string(resp.ChatCompletionResponse.Choices[0].FinishReason) + toolChoices, bad := parseToolCalls(msg.ToolCalls, finishReason) + if bad != nil { + // See decisionWithStreaming. + if finishReason == string(openai.FinishReasonLength) { + if err := lengthRetried.truncatedToolCall(&decision, usage, resp.MaxTokens, bad, reasoning); err != nil { + return nil, fmt.Errorf("failed to make a decision after %d attempts: %w", attempts+1, err) } + budget++ continue } - - toolChoices = append(toolChoices, &ToolChoice{ - Name: toolCall.Function.Name, - Arguments: arguments, - }) + lastErr = fmt.Errorf("%w: tool %q: %v", ErrToolArgumentsInvalid, bad.call.Function.Name, bad.err) + decision.Messages = appendArgumentsCorrection(decision.Messages, msg.Content, bad) + continue } xlog.Debug("[decision] tools selected", "message", msg.Content, "toolChoices", len(toolChoices)) - - // If we successfully parsed all tool calls, return the result - if len(toolChoices) == len(msg.ToolCalls) { - result := &decisionResult{ - toolChoices: toolChoices, - message: msg.Content, - reasoning: reasoning, - usage: usage, - } - return result, nil - } + return &decisionResult{ + toolChoices: toolChoices, + message: msg.Content, + reasoning: reasoning, + usage: usage, + }, nil } return nil, fmt.Errorf("failed to make a decision after %d attempts: %w", maxRetries, lastErr)