diff --git a/internal/llm/client.go b/internal/llm/client.go index 0a31d72..fa72b21 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -622,13 +622,38 @@ func (c *OpenAIClient) CompletionsWithCtx(ctx context.Context, req ChatRequest) // NewStreaming method sets stream=true on the wire itself. When // streaming is NOT enabled, leaving the key in the body would make // the API answer with text/event-stream and the non-streaming path - // fails to decode (see issue #647). - if k == "stream" { + // fails to decode (see issue #647). "stream_options" is owned by the + // streaming branch below for the same reason: providers reject it + // unless stream is true. + if k == "stream" || k == "stream_options" { continue } opts = append(opts, openaiopt.WithJSONSet(k, v)) } if stream, ok := c.cfg.ExtraBody["stream"].(bool); ok && stream { + if streamOptions, ok := c.cfg.ExtraBody["stream_options"]; !ok { + // OpenAI-compatible servers omit token usage from streams unless + // asked, silently losing cost accounting for streamed requests. + // Ask for the final usage chunk by default. + params.StreamOptions = openai.ChatCompletionStreamOptionsParam{IncludeUsage: openai.Bool(true)} + } else if streamOptions != nil { + // An explicit stream_options in extra_body replaces the default, + // but usage stays requested unless include_usage itself is spelled + // out: configuring an unrelated stream option must not silently + // disable cost accounting. An explicit null suppresses the field + // entirely, for gateways that reject stream_options. + if object, ok := streamOptions.(map[string]any); ok { + if _, has := object["include_usage"]; !has { + merged := make(map[string]any, len(object)+1) + for key, value := range object { + merged[key] = value + } + merged["include_usage"] = true + streamOptions = merged + } + } + opts = append(opts, openaiopt.WithJSONSet("stream_options", streamOptions)) + } return c.completionsStreaming(ctx, params, opts...) } diff --git a/internal/llm/client_test.go b/internal/llm/client_test.go index 6966e99..3ad43fa 100644 --- a/internal/llm/client_test.go +++ b/internal/llm/client_test.go @@ -11,6 +11,7 @@ import ( "io" "net/http" "net/http/httptest" + "reflect" "strings" "sync/atomic" "testing" @@ -848,6 +849,139 @@ func TestOpenAIClient_StreamOnlyGateway(t *testing.T) { } } +// TestOpenAIClient_StreamingRequestsUsageByDefault verifies that streaming +// requests ask for the final usage chunk via stream_options.include_usage +// when the provider config does not set stream_options itself. +// OpenAI-compatible servers omit usage from streams unless asked, silently +// losing token accounting for every streamed request. +func TestOpenAIClient_StreamingRequestsUsageByDefault(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("decode request body: %v", err) + http.Error(w, "invalid request body", http.StatusBadRequest) + return + } + + streamOptions, ok := body["stream_options"].(map[string]any) + if !ok || streamOptions["include_usage"] != true { + t.Errorf("stream_options = %#v, want include_usage true", body["stream_options"]) + http.Error(w, "usage request required", http.StatusBadRequest) + return + } + + writeOpenAISSE(t, w, + `{"id":"chatcmpl-usage","object":"chat.completion.chunk","created":1,"model":"gpt-stream","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},"finish_reason":null}]}`, + `{"id":"chatcmpl-usage","object":"chat.completion.chunk","created":1,"model":"gpt-stream","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`, + `{"id":"chatcmpl-usage","object":"chat.completion.chunk","created":1,"model":"gpt-stream","choices":[],"usage":{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15}}`, + ) + })) + defer server.Close() + + client := NewOpenAIClient(ClientConfig{ + URL: server.URL + "/v1", + APIKey: "test-key", + Model: "gpt-stream", + ExtraBody: map[string]any{"stream": true}, + }) + + resp, err := client.CompletionsWithCtx(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if resp.Usage == nil { + t.Fatal("Usage = nil, want usage from the final stream chunk") + } + if resp.Usage.PromptTokens != 12 || resp.Usage.CompletionTokens != 3 || resp.Usage.TotalTokens != 15 { + t.Errorf("Usage = %+v, want prompt 12 / completion 3 / total 15", resp.Usage) + } +} + +// TestOpenAIClient_StreamOptionsConfigOverrides verifies the +// extra_body.stream_options states: an explicit include_usage value is +// respected, an object configuring only unrelated options keeps the usage +// default merged in, an explicit null suppresses the field entirely (for +// gateways that reject stream_options), and non-streaming requests never +// carry the field regardless of config. +func TestOpenAIClient_StreamOptionsConfigOverrides(t *testing.T) { + tests := []struct { + name string + extraBody map[string]any + want any + wantAbsent bool + }{ + { + name: "explicit include_usage false is respected", + extraBody: map[string]any{"stream": true, "stream_options": map[string]any{"include_usage": false}}, + want: map[string]any{"include_usage": false}, + }, + { + name: "unrelated option keeps the usage default", + extraBody: map[string]any{"stream": true, "stream_options": map[string]any{"continuous_usage_stats": true}}, + want: map[string]any{"continuous_usage_stats": true, "include_usage": true}, + }, + { + name: "explicit null suppresses field", + extraBody: map[string]any{"stream": true, "stream_options": nil}, + wantAbsent: true, + }, + { + name: "non-streaming never sends stream_options", + extraBody: map[string]any{"stream_options": map[string]any{"include_usage": true}}, + wantAbsent: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + streaming, _ := tt.extraBody["stream"].(bool) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("decode request body: %v", err) + http.Error(w, "invalid request body", http.StatusBadRequest) + return + } + + value, present := body["stream_options"] + if tt.wantAbsent { + if present { + t.Errorf("stream_options = %#v, want absent", value) + } + } else if !reflect.DeepEqual(value, tt.want) { + t.Errorf("stream_options = %#v, want %#v", value, tt.want) + } + + if streaming { + writeOpenAISSE(t, w, + `{"id":"chatcmpl-s","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}`, + `{"id":"chatcmpl-s","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`, + ) + return + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"id":"chatcmpl-n","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`) + })) + defer server.Close() + + client := NewOpenAIClient(ClientConfig{ + URL: server.URL + "/v1", + APIKey: "test-key", + Model: "m", + ExtraBody: tt.extraBody, + }) + + if _, err := client.CompletionsWithCtx(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + }); err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + }) + } +} + // TestOpenAIClient_NonStreamingRequestDropsStreamField verifies that the // "stream" key in extra_body is NOT forwarded to the non-streaming Chat // Completions request. Forwarding it makes the API answer with