diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 2958937c1..e14aa22f1 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -258,7 +258,12 @@ func (p *Provider) ChatStream( requestBody := p.buildRequestBody(messages, tools, model, options) requestBody["stream"] = true if supportsStreamingUsage(p.apiBase) { - requestBody["stream_options"] = map[string]any{"include_usage": true} + streamOptions := map[string]any{} + if existing, ok := requestBody["stream_options"].(map[string]any); ok { + streamOptions = maps.Clone(existing) + } + streamOptions["include_usage"] = true + requestBody["stream_options"] = streamOptions } jsonData, err := json.Marshal(requestBody) @@ -474,12 +479,7 @@ func (p *Provider) SupportsNativeSearch() bool { } func isNativeSearchHost(apiBase string) bool { - u, err := url.Parse(apiBase) - if err != nil { - return false - } - host := u.Hostname() - return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com") + return isOpenAINativeBaseURL(apiBase) } // isOpenAINativeBaseURL reports whether the given API base is the OpenAI or diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index d4597751f..280d67d51 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -991,7 +991,7 @@ func chatWithCacheKey(t *testing.T, apiBase string) map[string]any { return requestBody } -func chatStreamWithRequestBody(t *testing.T, apiBase string) map[string]any { +func chatStreamWithRequestBody(t *testing.T, apiBase string, opts ...Option) map[string]any { t.Helper() var requestBody map[string]any @@ -1016,7 +1016,7 @@ func chatStreamWithRequestBody(t *testing.T, apiBase string) map[string]any { })) defer server.Close() - p := NewProvider("key", server.URL, "") + p := NewProvider("key", server.URL, "", opts...) p.apiBase = apiBase p.httpClient = &http.Client{ Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) { @@ -1057,6 +1057,26 @@ func TestProviderChatStream_IncludesUsageForOpenAI(t *testing.T) { } } +func TestProviderChatStream_PreservesExistingStreamOptions(t *testing.T) { + body := chatStreamWithRequestBody( + t, + "https://api.openai.com/v1", + WithExtraBody(map[string]any{ + "stream_options": map[string]any{"custom_flag": true}, + }), + ) + streamOptions, ok := body["stream_options"].(map[string]any) + if !ok { + t.Fatalf("stream_options = %T, want map[string]any", body["stream_options"]) + } + if got := streamOptions["custom_flag"]; got != true { + t.Fatalf("stream_options.custom_flag = %v, want true", got) + } + if got := streamOptions["include_usage"]; got != true { + t.Fatalf("stream_options.include_usage = %v, want true", got) + } +} + func TestProviderChatStream_OmitsUsageForNonOpenAI(t *testing.T) { tests := []struct { name string