From a2879c1faa61583c6ed0a77bc15c11d994d27cf8 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Mon, 4 May 2026 17:32:02 +0200 Subject: [PATCH 1/8] feat(gemini): add native API translator --- .env.template | 5 + CLAUDE.md | 2 +- README.md | 1 + config/config.example.yaml | 2 + helm/README.md | 1 + helm/templates/_helpers.tpl | 4 + helm/values.schema.json | 1 + helm/values.yaml | 2 + internal/providers/gemini/gemini.go | 98 ++- internal/providers/gemini/gemini_test.go | 251 ++++++++ internal/providers/gemini/native.go | 697 +++++++++++++++++++++ internal/providers/gemini/native_stream.go | 227 +++++++ 12 files changed, 1281 insertions(+), 10 deletions(-) create mode 100644 internal/providers/gemini/native.go create mode 100644 internal/providers/gemini/native_stream.go diff --git a/.env.template b/.env.template index 97281700f..64010a94d 100644 --- a/.env.template +++ b/.env.template @@ -263,6 +263,11 @@ # Google Gemini # GEMINI_API_KEY=... +# Use Gemini's native generateContent API for chat/responses (default: true). +# Set to false to use Gemini's OpenAI-compatible API for chat/responses. +# USE_GOOGLE_GEMINI_NATIVE_API=true +# OpenAI-compatible Gemini base URL used when native chat is disabled, and by +# Gemini embeddings/files/batches which still rely on that compatibility surface. # GEMINI_BASE_URL=https://generativelanguage.googleapis.com/v1beta/openai # xAI (Grok) diff --git a/CLAUDE.md b/CLAUDE.md index bc9aeb06e..5307fe745 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -122,5 +122,5 @@ Full reference: `.env.template` and `config/config.yaml` - **Resilience:** Configured via `config/config.yaml` - global `resilience.retry.*` and `resilience.circuit_breaker.*` defaults with optional per-provider overrides under `providers..resilience.retry.*` and `providers..resilience.circuit_breaker.*`. Retry defaults: `max_retries` (3), `initial_backoff` (1s), `max_backoff` (30s), `backoff_factor` (2.0), `jitter_factor` (0.1). Circuit breaker defaults: `failure_threshold` (5), `success_threshold` (2), `timeout` (30s) - **Metrics:** `METRICS_ENABLED` (false), `METRICS_ENDPOINT` (/metrics) - **Guardrails:** Configured via `config/config.yaml` only (except `GUARDRAILS_ENABLED` env var) -- **Providers:** `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `GEMINI_API_KEY`, `XAI_API_KEY`, `GROQ_API_KEY`, `OPENROUTER_API_KEY`, `ZAI_API_KEY`, `ZAI_BASE_URL` (optional Z.ai endpoint override), `MINIMAX_API_KEY`, `MINIMAX_BASE_URL` (optional MiniMax endpoint override), `AZURE_API_KEY`, `AZURE_BASE_URL` (Azure OpenAI deployment base URL), `AZURE_API_VERSION` (optional Azure API version), `ORACLE_API_KEY` (Oracle API key), `ORACLE_BASE_URL` (Oracle OpenAI-compatible base URL), `[_SUFFIX]_MODELS` (comma-separated configured model list for any provider type), `OLLAMA_BASE_URL`, `VLLM_BASE_URL`, `VLLM_API_KEY` (optional upstream vLLM bearer token) +- **Providers:** `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `GEMINI_API_KEY`, `USE_GOOGLE_GEMINI_NATIVE_API` (true by default; false uses Gemini's OpenAI-compatible chat API), `XAI_API_KEY`, `GROQ_API_KEY`, `OPENROUTER_API_KEY`, `ZAI_API_KEY`, `ZAI_BASE_URL` (optional Z.ai endpoint override), `MINIMAX_API_KEY`, `MINIMAX_BASE_URL` (optional MiniMax endpoint override), `AZURE_API_KEY`, `AZURE_BASE_URL` (Azure OpenAI deployment base URL), `AZURE_API_VERSION` (optional Azure API version), `ORACLE_API_KEY` (Oracle API key), `ORACLE_BASE_URL` (Oracle OpenAI-compatible base URL), `[_SUFFIX]_MODELS` (comma-separated configured model list for any provider type), `OLLAMA_BASE_URL`, `VLLM_BASE_URL`, `VLLM_API_KEY` (optional upstream vLLM bearer token) - **Provider model metadata:** `providers..models` accepts either model IDs (strings) or `{id, metadata}` objects. When `metadata` is supplied (`display_name`, `context_window`, `max_output_tokens`, `modes`, `capabilities`, `pricing`, …) it is merged onto the remote ai-model-list entry during enrichment, with operator values winning per-field. Primary use case: advertising context windows, capabilities, and pricing for local models (Ollama) and other custom endpoints whose IDs are not in the upstream registry. diff --git a/README.md b/README.md index db467c04c..5ae3c4399 100644 --- a/README.md +++ b/README.md @@ -250,6 +250,7 @@ Key settings: | `ENABLE_PASSTHROUGH_ROUTES` | `true` | Enable provider-native passthrough routes under `/p/{provider}/...` | | `ALLOW_PASSTHROUGH_V1_ALIAS` | `true` | Allow `/p/{provider}/v1/...` aliases while keeping `/p/{provider}/...` canonical | | `ENABLED_PASSTHROUGH_PROVIDERS` | `openai,anthropic,openrouter,zai,vllm` | Comma-separated list of enabled passthrough providers | +| `USE_GOOGLE_GEMINI_NATIVE_API` | `true` | Use Gemini native `generateContent` for chat/responses; set `false` for Gemini's OpenAI-compatible API | | `STORAGE_TYPE` | `sqlite` | Storage backend (`sqlite`, `postgresql`, `mongodb`) | | `METRICS_ENABLED` | `false` | Enable Prometheus metrics (experimental) | | `LOGGING_ENABLED` | `false` | Enable audit logging | diff --git a/config/config.example.yaml b/config/config.example.yaml index eea45a2df..269e6c2b8 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -201,6 +201,8 @@ providers: gemini: type: gemini api_key: "..." + # Chat/responses use Gemini's native generateContent API by default. + # Set USE_GOOGLE_GEMINI_NATIVE_API=false to use Gemini's OpenAI-compatible API instead. xai: type: xai diff --git a/helm/README.md b/helm/README.md index 7f2e07528..d8e5e6e38 100644 --- a/helm/README.md +++ b/helm/README.md @@ -58,6 +58,7 @@ helm install gomodel ./helm \ | `providers.openai.enabled` | Enable OpenAI | `false` | | `providers.anthropic.enabled` | Enable Anthropic | `false` | | `providers.gemini.enabled` | Enable Gemini | `false` | +| `providers.gemini.useNativeApi` | Use Gemini native generateContent for chat/responses; set false for Gemini OpenAI compatibility | `true` | | `providers.groq.enabled` | Enable Groq | `false` | | `providers.xai.enabled` | Enable xAI | `false` | | `providers.zai.enabled` | Enable Z.ai | `false` | diff --git a/helm/templates/_helpers.tpl b/helm/templates/_helpers.tpl index ebae494f0..1ac6472f5 100644 --- a/helm/templates/_helpers.tpl +++ b/helm/templates/_helpers.tpl @@ -173,6 +173,10 @@ Generate provider environment variables for the Deployment. value: {{ $config.baseUrl | quote }} {{- end }} {{- end }} +{{- if and (eq $name "gemini") (hasKey $config "useNativeApi") }} +- name: USE_GOOGLE_GEMINI_NATIVE_API + value: {{ $config.useNativeApi | quote }} +{{- end }} {{- end }} {{- end }} {{- end }} diff --git a/helm/values.schema.json b/helm/values.schema.json index 23d9441a1..80230252b 100644 --- a/helm/values.schema.json +++ b/helm/values.schema.json @@ -255,6 +255,7 @@ "properties": { "enabled": { "type": "boolean" }, "apiKey": { "type": "string" }, + "useNativeApi": { "type": "boolean" }, "baseUrl": { "type": "string" } } }, diff --git a/helm/values.yaml b/helm/values.yaml index d9caf5b12..0851e8a14 100644 --- a/helm/values.yaml +++ b/helm/values.yaml @@ -64,6 +64,8 @@ providers: enabled: false # -- Gemini API key (ignored if providers.existingSecret is set) apiKey: "" + # -- Use Gemini native generateContent API for chat/responses. Set false to use Gemini's OpenAI-compatible API. + useNativeApi: true # -- Optional: Override Gemini base URL baseUrl: "" diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 38f03faa1..c88fc3e37 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "net/url" + "os" "slices" "strconv" "strings" @@ -30,16 +31,19 @@ var Registration = providers.Registration{ const ( // Gemini provides an OpenAI-compatible endpoint defaultOpenAICompatibleBaseURL = "https://generativelanguage.googleapis.com/v1beta/openai" - // Native Gemini API endpoint for models listing + // Native Gemini API endpoint for generateContent and models listing defaultModelsBaseURL = "https://generativelanguage.googleapis.com/v1beta" + useNativeAPIEnvVar = "USE_GOOGLE_GEMINI_NATIVE_API" ) // Provider implements the core.Provider interface for Google Gemini type Provider struct { client *llmclient.Client + nativeClient *llmclient.Client httpClient *http.Client hooks llmclient.Hooks apiKey string + useNativeAPI bool modelsURL string modelsClientConf llmclient.Config } @@ -48,10 +52,11 @@ type Provider struct { func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { baseURL := providers.ResolveBaseURL(providerCfg.BaseURL, defaultOpenAICompatibleBaseURL) p := &Provider{ - httpClient: nil, - apiKey: providerCfg.APIKey, - hooks: opts.Hooks, - modelsURL: defaultModelsBaseURL, + httpClient: nil, + apiKey: providerCfg.APIKey, + hooks: opts.Hooks, + useNativeAPI: useNativeAPIFromEnv(), + modelsURL: defaultModelsBaseURL, modelsClientConf: llmclient.Config{ ProviderName: "gemini", BaseURL: defaultModelsBaseURL, @@ -68,6 +73,7 @@ func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) c CircuitBreaker: opts.Resilience.CircuitBreaker, } p.client = llmclient.New(clientCfg, p.setHeaders) + p.nativeClient = llmclient.New(p.modelsClientConf, p.setNativeHeaders) return p } @@ -78,10 +84,11 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H httpClient = http.DefaultClient } p := &Provider{ - httpClient: httpClient, - apiKey: apiKey, - hooks: hooks, - modelsURL: defaultModelsBaseURL, + httpClient: httpClient, + apiKey: apiKey, + hooks: hooks, + useNativeAPI: useNativeAPIFromEnv(), + modelsURL: defaultModelsBaseURL, } modelsCfg := llmclient.DefaultConfig("gemini", defaultModelsBaseURL) modelsCfg.Hooks = hooks @@ -89,12 +96,16 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H cfg := llmclient.DefaultConfig("gemini", defaultOpenAICompatibleBaseURL) cfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) + p.nativeClient = llmclient.NewWithHTTPClient(httpClient, modelsCfg, p.setNativeHeaders) return p } // SetBaseURL allows configuring a custom base URL for the provider func (p *Provider) SetBaseURL(url string) { p.client.SetBaseURL(url) + if p.nativeClient != nil { + p.nativeClient.SetBaseURL(url) + } } // SetModelsURL allows configuring a custom models API base URL. @@ -114,6 +125,28 @@ func (p *Provider) setHeaders(req *http.Request) { } } +// setNativeHeaders sets the required headers for Gemini native API requests. +func (p *Provider) setNativeHeaders(req *http.Request) { + req.Header.Set("x-goog-api-key", p.apiKey) + + if requestID := core.GetRequestID(req.Context()); requestID != "" { + req.Header.Set("X-Request-Id", requestID) + } +} + +func useNativeAPIFromEnv() bool { + value, ok := os.LookupEnv(useNativeAPIEnvVar) + if !ok || strings.TrimSpace(value) == "" { + return true + } + switch strings.ToLower(strings.TrimSpace(value)) { + case "0", "false", "no", "off": + return false + default: + return true + } +} + // adaptChatRequest rewrites a ChatRequest for Gemini's OpenAI-compatible endpoint. // Gemini uses "reasoning_effort" as a top-level string (e.g. "low", "medium", "high"), // not the nested "reasoning": {"effort": "..."} format. @@ -140,6 +173,12 @@ func adaptChatRequest(req *core.ChatRequest) (any, error) { // ChatCompletion sends a chat completion request to Gemini func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (*core.ChatResponse, error) { + if req == nil { + return nil, core.NewInvalidRequestError("chat request is required", nil) + } + if p.useNativeAPI { + return p.nativeChatCompletion(ctx, req) + } body, err := adaptChatRequest(req) if err != nil { return nil, err @@ -159,8 +198,31 @@ func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (* return &resp, nil } +func (p *Provider) nativeChatCompletion(ctx context.Context, req *core.ChatRequest) (*core.ChatResponse, error) { + body, err := convertChatRequestToGemini(req) + if err != nil { + return nil, err + } + var geminiResp geminiGenerateContentResponse + err = p.nativeClient.Do(ctx, llmclient.Request{ + Method: http.MethodPost, + Endpoint: nativeGenerateEndpoint(req.Model), + Body: body, + }, &geminiResp) + if err != nil { + return nil, err + } + return nativeChatResponse(req, &geminiResp), nil +} + // StreamChatCompletion returns a raw response body for streaming (caller must close) func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { + if req == nil { + return nil, core.NewInvalidRequestError("chat request is required", nil) + } + if p.useNativeAPI { + return p.nativeStreamChatCompletion(ctx, req) + } streamReq := req.WithStreaming() body, err := adaptChatRequest(streamReq) if err != nil { @@ -179,6 +241,24 @@ func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatReque return stream, nil } +func (p *Provider) nativeStreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { + streamReq := req.WithStreaming() + body, err := convertChatRequestToGemini(streamReq) + if err != nil { + return nil, err + } + stream, err := p.nativeClient.DoStream(ctx, llmclient.Request{ + Method: http.MethodPost, + Endpoint: nativeStreamEndpoint(req.Model), + Body: body, + }) + if err != nil { + return nil, err + } + includeUsage := streamReq.StreamOptions != nil && streamReq.StreamOptions.IncludeUsage + return newGeminiNativeStream(stream, req.Model, includeUsage), nil +} + // geminiModel represents a model in Gemini's native API response type geminiModel struct { Name string `json:"name"` diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index ede7b4c59..9d573050b 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -39,6 +39,8 @@ func TestNew_ReturnsProvider(t *testing.T) { } func TestChatCompletion(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "false") + tests := []struct { name string statusCode int @@ -166,7 +168,183 @@ func TestChatCompletion(t *testing.T) { } } +func TestChatCompletion_UsesNativeGenerateContentByDefault(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("Method = %q, want %q", r.Method, http.MethodPost) + } + if r.URL.Path != "/models/gemini-2.5-flash:generateContent" { + t.Errorf("Path = %q, want native generateContent endpoint", r.URL.Path) + } + if got := r.Header.Get("x-goog-api-key"); got != "test-api-key" { + t.Errorf("x-goog-api-key = %q, want test-api-key", got) + } + if got := r.Header.Get("Authorization"); got != "" { + t.Errorf("Authorization = %q, want empty for native Gemini API", got) + } + + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + if _, ok := payload["messages"]; ok { + t.Fatal("native request should not contain OpenAI messages") + } + if _, ok := payload["contents"]; !ok { + t.Fatal("native request should contain contents") + } + generationConfig, ok := payload["generationConfig"].(map[string]any) + if !ok { + t.Fatalf("generationConfig = %#v, want object", payload["generationConfig"]) + } + if got := generationConfig["maxOutputTokens"]; got != float64(128) { + t.Fatalf("maxOutputTokens = %#v, want 128", got) + } + if _, ok := payload["system_instruction"]; !ok { + t.Fatal("system_instruction missing") + } + + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "gemini-native-123", + "candidates": [{ + "index": 0, + "content": {"role": "model", "parts": [{"text": "Hello from native Gemini"}]}, + "finishReason": "STOP" + }], + "usageMetadata": { + "promptTokenCount": 7, + "candidatesTokenCount": 5, + "totalTokenCount": 12 + } + }`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetBaseURL(server.URL) + + maxTokens := 128 + resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + MaxTokens: &maxTokens, + Messages: []core.Message{ + {Role: "system", Content: "Be concise."}, + {Role: "user", Content: "Hello"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.ID != "gemini-native-123" { + t.Fatalf("ID = %q, want gemini-native-123", resp.ID) + } + if resp.Provider != "gemini" { + t.Fatalf("Provider = %q, want gemini", resp.Provider) + } + if got := resp.Choices[0].Message.Content; got != "Hello from native Gemini" { + t.Fatalf("content = %q, want native text", got) + } + if got := resp.Choices[0].FinishReason; got != "stop" { + t.Fatalf("finish_reason = %q, want stop", got) + } + if resp.Usage.TotalTokens != 12 { + t.Fatalf("TotalTokens = %d, want 12", resp.Usage.TotalTokens) + } +} + +func TestChatCompletion_NativeFunctionCallTranslation(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + tools, ok := payload["tools"].([]any) + if !ok || len(tools) != 1 { + t.Fatalf("tools = %#v, want one Gemini tool", payload["tools"]) + } + toolConfig, ok := payload["toolConfig"].(map[string]any) + if !ok { + t.Fatalf("toolConfig = %#v, want object", payload["toolConfig"]) + } + functionConfig := toolConfig["functionCallingConfig"].(map[string]any) + if functionConfig["mode"] != "ANY" { + t.Fatalf("tool mode = %#v, want ANY", functionConfig["mode"]) + } + + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "gemini-tools-123", + "candidates": [{ + "content": {"role": "model", "parts": [{ + "functionCall": { + "id": "call_native", + "name": "lookup_weather", + "args": {"city": "Warsaw"} + } + }]}, + "finishReason": "STOP" + }] + }`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetBaseURL(server.URL) + + resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Weather?"}, + }, + Tools: []map[string]any{{ + "type": "function", + "function": map[string]any{ + "name": "lookup_weather", + "description": "Look up weather", + "parameters": map[string]any{ + "type": "object", + "properties": map[string]any{ + "city": map[string]any{"type": "string"}, + }, + "required": []any{"city"}, + }, + }, + }}, + ToolChoice: map[string]any{ + "type": "function", + "function": map[string]any{"name": "lookup_weather"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got := resp.Choices[0].FinishReason; got != "tool_calls" { + t.Fatalf("finish_reason = %q, want tool_calls", got) + } + if len(resp.Choices[0].Message.ToolCalls) != 1 { + t.Fatalf("tool calls = %d, want 1", len(resp.Choices[0].Message.ToolCalls)) + } + call := resp.Choices[0].Message.ToolCalls[0] + if call.ID != "call_native" || call.Function.Name != "lookup_weather" { + t.Fatalf("tool call = %+v, want native function call", call) + } + if call.Function.Arguments != `{"city":"Warsaw"}` { + t.Fatalf("arguments = %q, want JSON object", call.Function.Arguments) + } +} + func TestStreamChatCompletion(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "false") + tests := []struct { name string statusCode int @@ -257,6 +435,77 @@ data: [DONE] } } +func TestStreamChatCompletion_UsesNativeStreamByDefault(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/models/gemini-2.5-flash:streamGenerateContent" { + t.Errorf("Path = %q, want native streamGenerateContent endpoint", r.URL.Path) + } + if got := r.URL.Query().Get("alt"); got != "sse" { + t.Errorf("alt = %q, want sse", got) + } + if got := r.Header.Get("x-goog-api-key"); got != "test-api-key" { + t.Errorf("x-goog-api-key = %q, want test-api-key", got) + } + + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + if _, ok := payload["stream"]; ok { + t.Fatal("native stream request should not contain OpenAI stream flag") + } + + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`data: {"responseId":"gemini-stream-123","candidates":[{"content":{"role":"model","parts":[{"text":"Hello"}]}}]} + +data: {"responseId":"gemini-stream-123","candidates":[{"content":{"role":"model","parts":[{"text":"!"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}} + +`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetBaseURL(server.URL) + + body, err := provider.StreamChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + StreamOptions: &core.StreamOptions{IncludeUsage: true}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer func() { _ = body.Close() }() + + raw, err := io.ReadAll(body) + if err != nil { + t.Fatalf("failed to read stream: %v", err) + } + stream := string(raw) + if !strings.Contains(stream, `"id":"gemini-stream-123"`) { + t.Fatalf("stream = %q, want native response id", stream) + } + if !strings.Contains(stream, `"content":"Hello"`) || !strings.Contains(stream, `"content":"!"`) { + t.Fatalf("stream = %q, want converted content chunks", stream) + } + if !strings.Contains(stream, `"finish_reason":"stop"`) { + t.Fatalf("stream = %q, want stop finish reason", stream) + } + if !strings.Contains(stream, `"usage"`) || !strings.Contains(stream, `"total_tokens":6`) { + t.Fatalf("stream = %q, want usage chunk", stream) + } + if !strings.Contains(stream, "data: [DONE]") { + t.Fatalf("stream = %q, want [DONE]", stream) + } +} + func TestListModels(t *testing.T) { tests := []struct { name string @@ -391,6 +640,8 @@ func TestChatCompletionWithContext(t *testing.T) { } func TestResponses(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "false") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{ diff --git a/internal/providers/gemini/native.go b/internal/providers/gemini/native.go new file mode 100644 index 000000000..b39b71d8d --- /dev/null +++ b/internal/providers/gemini/native.go @@ -0,0 +1,697 @@ +package gemini + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "mime" + "net/http" + "net/url" + "strconv" + "strings" + "time" + + "gomodel/internal/core" +) + +type geminiGenerateContentRequest struct { + SystemInstruction *geminiContent `json:"system_instruction,omitempty"` + Contents []geminiContent `json:"contents"` + Tools []geminiTool `json:"tools,omitempty"` + ToolConfig *geminiToolConfig `json:"toolConfig,omitempty"` + GenerationConfig map[string]any `json:"generationConfig,omitempty"` + SafetySettings []map[string]any `json:"safetySettings,omitempty"` + CachedContent string `json:"cachedContent,omitempty"` +} + +type geminiContent struct { + Role string `json:"role,omitempty"` + Parts []geminiPart `json:"parts"` +} + +type geminiPart struct { + Text string `json:"text,omitempty"` + InlineData *geminiBlob `json:"inline_data,omitempty"` + FileData *geminiFileData `json:"file_data,omitempty"` + FunctionCall *geminiFunctionCall `json:"functionCall,omitempty"` + FunctionCallAlt *geminiFunctionCall `json:"function_call,omitempty"` + FunctionResponse *geminiFunctionResponse `json:"functionResponse,omitempty"` + FunctionResponseAlt *geminiFunctionResponse `json:"function_response,omitempty"` + Thought bool `json:"thought,omitempty"` + ThoughtSignature string `json:"thoughtSignature,omitempty"` +} + +func (p geminiPart) functionCall() *geminiFunctionCall { + if p.FunctionCall != nil { + return p.FunctionCall + } + return p.FunctionCallAlt +} + +type geminiBlob struct { + MimeType string `json:"mime_type,omitempty"` + Data string `json:"data,omitempty"` +} + +type geminiFileData struct { + MimeType string `json:"mime_type,omitempty"` + FileURI string `json:"file_uri,omitempty"` +} + +type geminiFunctionCall struct { + ID string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + Args json.RawMessage `json:"args,omitempty"` +} + +type geminiFunctionResponse struct { + ID string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + Response json.RawMessage `json:"response,omitempty"` +} + +type geminiTool struct { + FunctionDeclarations []geminiFunctionDeclaration `json:"functionDeclarations,omitempty"` +} + +type geminiFunctionDeclaration struct { + Name string `json:"name,omitempty"` + Description string `json:"description,omitempty"` + Parameters json.RawMessage `json:"parameters,omitempty"` +} + +type geminiToolConfig struct { + FunctionCallingConfig geminiFunctionCallingConfig `json:"functionCallingConfig"` +} + +type geminiFunctionCallingConfig struct { + Mode string `json:"mode,omitempty"` + AllowedFunctionNames []string `json:"allowedFunctionNames,omitempty"` +} + +type geminiGenerateContentResponse struct { + Candidates []geminiCandidate `json:"candidates,omitempty"` + PromptFeedback json.RawMessage `json:"promptFeedback,omitempty"` + UsageMetadata geminiUsageMetadata `json:"usageMetadata,omitempty"` + ModelVersion string `json:"modelVersion,omitempty"` + ResponseID string `json:"responseId,omitempty"` + ModelStatus json.RawMessage `json:"modelStatus,omitempty"` +} + +type geminiCandidate struct { + Content geminiContent `json:"content,omitempty"` + FinishReason string `json:"finishReason,omitempty"` + Index int `json:"index,omitempty"` + SafetyRatings json.RawMessage `json:"safetyRatings,omitempty"` +} + +type geminiUsageMetadata struct { + PromptTokenCount int `json:"promptTokenCount,omitempty"` + CachedContentTokenCount int `json:"cachedContentTokenCount,omitempty"` + CandidatesTokenCount int `json:"candidatesTokenCount,omitempty"` + ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount,omitempty"` + ThoughtsTokenCount int `json:"thoughtsTokenCount,omitempty"` + TotalTokenCount int `json:"totalTokenCount,omitempty"` + PromptTokensDetails json.RawMessage `json:"promptTokensDetails,omitempty"` + CacheTokensDetails json.RawMessage `json:"cacheTokensDetails,omitempty"` + CandidatesTokensDetails json.RawMessage `json:"candidatesTokensDetails,omitempty"` + ToolUsePromptTokensDetails json.RawMessage `json:"toolUsePromptTokensDetails,omitempty"` +} + +func convertChatRequestToGemini(req *core.ChatRequest) (*geminiGenerateContentRequest, error) { + if req == nil { + return nil, core.NewInvalidRequestError("chat request is required", nil) + } + + out := &geminiGenerateContentRequest{ + Contents: make([]geminiContent, 0, len(req.Messages)), + } + + systemParts := make([]geminiPart, 0) + toolCallNames := make(map[string]string) + for _, msg := range req.Messages { + var ( + parts []geminiPart + err error + ) + if strings.TrimSpace(msg.Role) == "tool" { + parts, err = geminiPartsFromToolMessage(msg, toolCallNames[msg.ToolCallID]) + } else { + parts, err = geminiPartsFromMessage(msg) + } + if err != nil { + return nil, err + } + if len(parts) == 0 { + continue + } + + switch strings.TrimSpace(msg.Role) { + case "system", "developer": + systemParts = append(systemParts, parts...) + case "assistant": + out.Contents = append(out.Contents, geminiContent{Role: "model", Parts: parts}) + for _, call := range msg.ToolCalls { + if call.ID != "" && call.Function.Name != "" { + toolCallNames[call.ID] = call.Function.Name + } + } + case "tool": + out.Contents = append(out.Contents, geminiContent{Role: "user", Parts: parts}) + default: + out.Contents = append(out.Contents, geminiContent{Role: "user", Parts: parts}) + } + } + if len(systemParts) > 0 { + out.SystemInstruction = &geminiContent{Parts: systemParts} + } + + tools, err := geminiToolsFromOpenAI(req.Tools) + if err != nil { + return nil, err + } + out.Tools = tools + out.ToolConfig = geminiToolConfigFromOpenAI(req.ToolChoice) + out.GenerationConfig = geminiGenerationConfig(req) + out.SafetySettings = geminiSafetySettings(req) + out.CachedContent = geminiCachedContent(req) + return out, nil +} + +func geminiPartsFromMessage(msg core.Message) ([]geminiPart, error) { + if len(msg.ToolCalls) > 0 { + parts := make([]geminiPart, 0, len(msg.ToolCalls)+1) + if text := strings.TrimSpace(core.ExtractTextContent(msg.Content)); text != "" { + parts = append(parts, geminiPart{Text: text}) + } + for _, call := range msg.ToolCalls { + args := json.RawMessage(strings.TrimSpace(call.Function.Arguments)) + if len(args) == 0 { + args = json.RawMessage(`{}`) + } + parts = append(parts, geminiPart{FunctionCall: &geminiFunctionCall{ + ID: call.ID, + Name: call.Function.Name, + Args: args, + }}) + } + return parts, nil + } + + switch content := msg.Content.(type) { + case nil: + return nil, nil + case string: + if content == "" { + return nil, nil + } + return []geminiPart{{Text: content}}, nil + default: + parts, ok := core.NormalizeContentParts(content) + if !ok { + text := core.ExtractTextContent(content) + if text == "" { + return nil, nil + } + return []geminiPart{{Text: text}}, nil + } + return geminiPartsFromContentParts(parts) + } +} + +func geminiPartsFromToolMessage(msg core.Message, functionName string) ([]geminiPart, error) { + if functionName == "" { + functionName = msg.ToolCallID + } + response, err := geminiToolResponsePayload(core.ExtractTextContent(msg.Content)) + if err != nil { + return nil, err + } + return []geminiPart{{ + FunctionResponse: &geminiFunctionResponse{ + ID: msg.ToolCallID, + Name: functionName, + Response: response, + }, + }}, nil +} + +func geminiToolResponsePayload(content string) (json.RawMessage, error) { + trimmed := strings.TrimSpace(content) + if trimmed == "" { + return json.RawMessage(`{}`), nil + } + if json.Valid([]byte(trimmed)) { + if strings.HasPrefix(trimmed, "{") { + return json.RawMessage(trimmed), nil + } + encoded, err := json.Marshal(map[string]json.RawMessage{"result": json.RawMessage(trimmed)}) + return json.RawMessage(encoded), err + } + encoded, err := json.Marshal(map[string]string{"result": content}) + if err != nil { + return nil, core.NewInvalidRequestError("failed to marshal Gemini tool response", err) + } + return json.RawMessage(encoded), nil +} + +func geminiPartsFromContentParts(parts []core.ContentPart) ([]geminiPart, error) { + out := make([]geminiPart, 0, len(parts)) + for _, part := range parts { + switch part.Type { + case "text", "input_text": + if part.Text != "" { + out = append(out, geminiPart{Text: part.Text}) + } + case "image_url", "input_image": + if part.ImageURL == nil || part.ImageURL.URL == "" { + continue + } + geminiPart, err := geminiPartFromImageURL(part.ImageURL) + if err != nil { + return nil, err + } + out = append(out, geminiPart) + case "input_audio": + if part.InputAudio == nil { + continue + } + out = append(out, geminiPart{InlineData: &geminiBlob{ + MimeType: mimeTypeForAudioFormat(part.InputAudio.Format), + Data: part.InputAudio.Data, + }}) + } + } + return out, nil +} + +func geminiPartFromImageURL(image *core.ImageURLContent) (geminiPart, error) { + rawURL := strings.TrimSpace(image.URL) + if strings.HasPrefix(rawURL, "data:") { + mimeType, data, err := parseDataURL(rawURL) + if err != nil { + return geminiPart{}, err + } + if mimeType == "" { + mimeType = image.MediaType + } + if mimeType == "" { + mimeType = "image/jpeg" + } + return geminiPart{InlineData: &geminiBlob{MimeType: mimeType, Data: data}}, nil + } + + mimeType := image.MediaType + if mimeType == "" { + mimeType = "image/jpeg" + } + return geminiPart{FileData: &geminiFileData{MimeType: mimeType, FileURI: rawURL}}, nil +} + +func parseDataURL(rawURL string) (string, string, error) { + header, data, ok := strings.Cut(rawURL, ",") + if !ok { + return "", "", core.NewInvalidRequestError("invalid data URL in image_url", nil) + } + mediaType := strings.TrimPrefix(header, "data:") + mediaType = strings.TrimSuffix(mediaType, ";base64") + if parsed, _, err := mime.ParseMediaType(mediaType); err == nil { + mediaType = parsed + } + if _, err := base64.StdEncoding.DecodeString(data); err != nil { + return "", "", core.NewInvalidRequestError("invalid base64 data in image_url", err) + } + return mediaType, data, nil +} + +func mimeTypeForAudioFormat(format string) string { + format = strings.Trim(strings.ToLower(strings.TrimSpace(format)), ".") + if format == "" { + return "audio/mpeg" + } + if strings.Contains(format, "/") { + return format + } + return "audio/" + format +} + +func geminiToolsFromOpenAI(tools []map[string]any) ([]geminiTool, error) { + if len(tools) == 0 { + return nil, nil + } + declarations := make([]geminiFunctionDeclaration, 0, len(tools)) + for _, tool := range tools { + if strings.TrimSpace(fmt.Sprint(tool["type"])) != "function" { + continue + } + fn, ok := tool["function"].(map[string]any) + if !ok { + continue + } + name, _ := fn["name"].(string) + if strings.TrimSpace(name) == "" { + continue + } + description, _ := fn["description"].(string) + var parameters json.RawMessage + if raw, ok := fn["parameters"]; ok { + encoded, err := json.Marshal(raw) + if err != nil { + return nil, core.NewInvalidRequestError("failed to marshal Gemini tool parameters", err) + } + parameters = encoded + } + declarations = append(declarations, geminiFunctionDeclaration{ + Name: name, + Description: description, + Parameters: parameters, + }) + } + if len(declarations) == 0 { + return nil, nil + } + return []geminiTool{{FunctionDeclarations: declarations}}, nil +} + +func geminiToolConfigFromOpenAI(choice any) *geminiToolConfig { + mode := "" + var allowed []string + + switch value := choice.(type) { + case string: + switch strings.ToLower(strings.TrimSpace(value)) { + case "none": + mode = "NONE" + case "required": + mode = "ANY" + case "auto": + mode = "AUTO" + } + case map[string]any: + choiceType, _ := value["type"].(string) + if strings.TrimSpace(choiceType) == "function" { + mode = "ANY" + if fn, ok := value["function"].(map[string]any); ok { + if name, _ := fn["name"].(string); name != "" { + allowed = []string{name} + } + } + } + } + + if mode == "" { + return nil + } + return &geminiToolConfig{FunctionCallingConfig: geminiFunctionCallingConfig{ + Mode: mode, + AllowedFunctionNames: allowed, + }} +} + +func geminiGenerationConfig(req *core.ChatRequest) map[string]any { + cfg := make(map[string]any) + if req.MaxTokens != nil { + cfg["maxOutputTokens"] = *req.MaxTokens + } else if raw := req.ExtraFields.Lookup("max_completion_tokens"); len(raw) > 0 { + var maxTokens int + if err := json.Unmarshal(raw, &maxTokens); err == nil && maxTokens > 0 { + cfg["maxOutputTokens"] = maxTokens + } + } + if req.Temperature != nil { + cfg["temperature"] = *req.Temperature + } + copyJSONNumber(req.ExtraFields.Lookup("top_p"), cfg, "topP") + copyJSONNumber(req.ExtraFields.Lookup("top_k"), cfg, "topK") + copyJSONNumber(req.ExtraFields.Lookup("candidate_count"), cfg, "candidateCount") + copyJSONNumber(req.ExtraFields.Lookup("presence_penalty"), cfg, "presencePenalty") + copyJSONNumber(req.ExtraFields.Lookup("frequency_penalty"), cfg, "frequencyPenalty") + copyStopSequences(req.ExtraFields.Lookup("stop"), cfg) + copyResponseFormat(req.ExtraFields.Lookup("response_format"), cfg) + copyGoogleThinkingConfig(req.ExtraFields.Lookup("extra_body"), cfg) + if req.Reasoning != nil && strings.TrimSpace(req.Reasoning.Effort) != "" { + if _, exists := cfg["thinkingConfig"]; !exists { + if thinkingConfig := thinkingConfigForEffort(req.Model, req.Reasoning.Effort); len(thinkingConfig) > 0 { + cfg["thinkingConfig"] = thinkingConfig + } + } + } + if len(cfg) == 0 { + return nil + } + return cfg +} + +func copyJSONNumber(raw json.RawMessage, cfg map[string]any, key string) { + if len(raw) == 0 { + return + } + var value any + if err := json.Unmarshal(raw, &value); err == nil { + cfg[key] = value + } +} + +func copyStopSequences(raw json.RawMessage, cfg map[string]any) { + if len(raw) == 0 { + return + } + var one string + if err := json.Unmarshal(raw, &one); err == nil && one != "" { + cfg["stopSequences"] = []string{one} + return + } + var many []string + if err := json.Unmarshal(raw, &many); err == nil && len(many) > 0 { + cfg["stopSequences"] = many + } +} + +func copyResponseFormat(raw json.RawMessage, cfg map[string]any) { + if len(raw) == 0 { + return + } + var responseFormat map[string]any + if err := json.Unmarshal(raw, &responseFormat); err != nil { + return + } + formatType, _ := responseFormat["type"].(string) + switch formatType { + case "json_object": + cfg["responseMimeType"] = "application/json" + case "json_schema": + cfg["responseMimeType"] = "application/json" + if schemaObj, ok := responseFormat["json_schema"].(map[string]any); ok { + if schema, ok := schemaObj["schema"]; ok { + cfg["responseSchema"] = schema + } + } + } +} + +func copyGoogleThinkingConfig(raw json.RawMessage, cfg map[string]any) { + if len(raw) == 0 { + return + } + var extra struct { + Google struct { + ThinkingConfig map[string]any `json:"thinking_config"` + } `json:"google"` + } + if err := json.Unmarshal(raw, &extra); err == nil && len(extra.Google.ThinkingConfig) > 0 { + cfg["thinkingConfig"] = normalizeSnakeMapKeys(extra.Google.ThinkingConfig) + } +} + +func normalizeSnakeMapKeys(src map[string]any) map[string]any { + out := make(map[string]any, len(src)) + for key, value := range src { + switch key { + case "thinking_budget": + out["thinkingBudget"] = value + case "thinking_level": + out["thinkingLevel"] = value + case "include_thoughts": + out["includeThoughts"] = value + default: + out[key] = value + } + } + return out +} + +func thinkingConfigForEffort(model, effort string) map[string]any { + effort = strings.ToLower(strings.TrimSpace(effort)) + if strings.Contains(strings.ToLower(model), "gemini-2.5") { + switch effort { + case "none": + return map[string]any{"thinkingBudget": 0} + case "minimal", "low": + return map[string]any{"thinkingBudget": 1024} + case "medium": + return map[string]any{"thinkingBudget": 8192} + case "high": + return map[string]any{"thinkingBudget": 24576} + default: + return nil + } + } + if effort == "none" { + effort = "minimal" + } + return map[string]any{"thinkingLevel": effort} +} + +func geminiSafetySettings(req *core.ChatRequest) []map[string]any { + raw := req.ExtraFields.Lookup("safety_settings") + if len(raw) == 0 { + return nil + } + var settings []map[string]any + if err := json.Unmarshal(raw, &settings); err != nil { + return nil + } + return settings +} + +func geminiCachedContent(req *core.ChatRequest) string { + raw := req.ExtraFields.Lookup("cached_content") + if len(raw) == 0 { + return "" + } + var cached string + _ = json.Unmarshal(raw, &cached) + return cached +} + +func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContentResponse) *core.ChatResponse { + created := time.Now().Unix() + respID := geminiResp.ResponseID + if respID == "" { + respID = "chatcmpl-gemini-" + strconv.FormatInt(created, 10) + } + resp := &core.ChatResponse{ + ID: respID, + Object: "chat.completion", + Created: created, + Model: req.Model, + Provider: "gemini", + Choices: make([]core.Choice, 0, len(geminiResp.Candidates)), + Usage: usageFromGemini(geminiResp.UsageMetadata), + } + for i, candidate := range geminiResp.Candidates { + index := candidate.Index + if index == 0 && i > 0 { + index = i + } + content, toolCalls := openAIMessageFromGeminiParts(candidate.Content.Parts) + resp.Choices = append(resp.Choices, core.Choice{ + Index: index, + Message: core.ResponseMessage{ + Role: "assistant", + Content: content, + ToolCalls: toolCalls, + }, + FinishReason: finishReasonFromGemini(candidate.FinishReason, len(toolCalls) > 0), + }) + } + return resp +} + +func openAIMessageFromGeminiParts(parts []geminiPart) (string, []core.ToolCall) { + var text strings.Builder + toolCalls := make([]core.ToolCall, 0) + for i, part := range parts { + if part.Text != "" && !part.Thought { + text.WriteString(part.Text) + } + if call := part.functionCall(); call != nil { + id := call.ID + if id == "" { + id = "call_" + strconv.Itoa(i) + } + args := strings.TrimSpace(string(call.Args)) + if args == "" { + args = "{}" + } else { + var compact bytes.Buffer + if err := json.Compact(&compact, []byte(args)); err == nil { + args = compact.String() + } + } + toolCalls = append(toolCalls, core.ToolCall{ + ID: id, + Type: "function", + Function: core.FunctionCall{ + Name: call.Name, + Arguments: args, + }, + }) + } + } + return text.String(), toolCalls +} + +func usageFromGemini(usage geminiUsageMetadata) core.Usage { + out := core.Usage{ + PromptTokens: usage.PromptTokenCount, + CompletionTokens: usage.CandidatesTokenCount, + TotalTokens: usage.TotalTokenCount, + } + if out.TotalTokens == 0 { + out.TotalTokens = out.PromptTokens + out.CompletionTokens + } + raw := make(map[string]any) + if usage.CachedContentTokenCount > 0 { + raw["cached_content_token_count"] = usage.CachedContentTokenCount + out.PromptTokensDetails = &core.PromptTokensDetails{CachedTokens: usage.CachedContentTokenCount} + } + if usage.ToolUsePromptTokenCount > 0 { + raw["tool_use_prompt_token_count"] = usage.ToolUsePromptTokenCount + } + if usage.ThoughtsTokenCount > 0 { + raw["thoughts_token_count"] = usage.ThoughtsTokenCount + out.CompletionTokensDetails = &core.CompletionTokensDetails{ReasoningTokens: usage.ThoughtsTokenCount} + } + if len(raw) > 0 { + out.RawUsage = raw + } + return out +} + +func finishReasonFromGemini(reason string, hasToolCalls bool) string { + if hasToolCalls { + return "tool_calls" + } + switch strings.ToUpper(strings.TrimSpace(reason)) { + case "", "FINISH_REASON_UNSPECIFIED": + return "" + case "STOP": + return "stop" + case "MAX_TOKENS": + return "length" + case "SAFETY", "RECITATION", "LANGUAGE", "BLOCKLIST", "PROHIBITED_CONTENT", "SPII", "IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT", "IMAGE_RECITATION": + return "content_filter" + default: + return strings.ToLower(reason) + } +} + +func nativeGenerateEndpoint(model string) string { + return "/models/" + url.PathEscape(normalizeGeminiModelID(model)) + ":generateContent" +} + +func nativeStreamEndpoint(model string) string { + return "/models/" + url.PathEscape(normalizeGeminiModelID(model)) + ":streamGenerateContent?alt=sse" +} + +func normalizeGeminiModelID(model string) string { + model = strings.TrimSpace(model) + model = strings.TrimPrefix(model, "models/") + return model +} + +func nativeProviderError(message string, err error) *core.GatewayError { + return core.NewProviderError("gemini", http.StatusBadGateway, message, err) +} diff --git a/internal/providers/gemini/native_stream.go b/internal/providers/gemini/native_stream.go new file mode 100644 index 000000000..32168e01f --- /dev/null +++ b/internal/providers/gemini/native_stream.go @@ -0,0 +1,227 @@ +package gemini + +import ( + "bufio" + "encoding/json" + "io" + "strconv" + "strings" + "time" + + "gomodel/internal/core" +) + +type geminiNativeStream struct { + reader *io.PipeReader + body io.ReadCloser +} + +func newGeminiNativeStream(body io.ReadCloser, model string, includeUsage bool) io.ReadCloser { + pr, pw := io.Pipe() + stream := &geminiNativeStream{reader: pr, body: body} + go convertGeminiNativeStream(body, pw, model, includeUsage) + return stream +} + +func (s *geminiNativeStream) Read(p []byte) (int, error) { + return s.reader.Read(p) +} + +func (s *geminiNativeStream) Close() error { + _ = s.reader.Close() + return s.body.Close() +} + +func convertGeminiNativeStream(body io.ReadCloser, out *io.PipeWriter, model string, includeUsage bool) { + defer func() { _ = body.Close() }() + + scanner := bufio.NewScanner(body) + scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + + state := geminiStreamState{ + model: model, + includeUsage: includeUsage, + created: time.Now().Unix(), + } + var data strings.Builder + for scanner.Scan() { + line := strings.TrimRight(scanner.Text(), "\r") + if line == "" { + if err := state.consumeEvent(out, data.String()); err != nil { + _ = out.CloseWithError(err) + return + } + data.Reset() + continue + } + if strings.HasPrefix(line, "data:") { + if data.Len() > 0 { + data.WriteByte('\n') + } + data.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:"))) + } + } + if data.Len() > 0 { + if err := state.consumeEvent(out, data.String()); err != nil { + _ = out.CloseWithError(err) + return + } + } + if err := scanner.Err(); err != nil { + _ = out.CloseWithError(err) + return + } + _, _ = io.WriteString(out, "data: [DONE]\n\n") + _ = out.Close() +} + +type geminiStreamState struct { + model string + includeUsage bool + created int64 + responseID string + roleSent bool + sawToolCalls bool +} + +func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { + raw = strings.TrimSpace(raw) + if raw == "" || raw == "[DONE]" { + return nil + } + + var event geminiGenerateContentResponse + if err := json.Unmarshal([]byte(raw), &event); err != nil { + return nativeProviderError("failed to parse native Gemini stream event", err) + } + if s.responseID == "" { + s.responseID = event.ResponseID + if s.responseID == "" { + s.responseID = "chatcmpl-gemini-" + strconv.FormatInt(s.created, 10) + } + } + + for i, candidate := range event.Candidates { + choice, ok := s.chatChunkChoice(candidate, i) + if !ok { + continue + } + chunk := map[string]any{ + "id": s.responseID, + "object": "chat.completion.chunk", + "created": s.created, + "model": s.model, + "provider": "gemini", + "choices": []map[string]any{choice}, + } + if s.includeUsage { + if usage := geminiUsageMap(event.UsageMetadata); usage != nil { + chunk["usage"] = usage + } + } + if err := writeOpenAIStreamChunk(out, chunk); err != nil { + return err + } + } + + if len(event.Candidates) == 0 && s.includeUsage { + if usage := geminiUsageMap(event.UsageMetadata); usage != nil { + chunk := map[string]any{ + "id": s.responseID, + "object": "chat.completion.chunk", + "created": s.created, + "model": s.model, + "provider": "gemini", + "choices": []map[string]any{}, + "usage": usage, + } + if err := writeOpenAIStreamChunk(out, chunk); err != nil { + return err + } + } + } + return nil +} + +func (s *geminiStreamState) chatChunkChoice(candidate geminiCandidate, fallbackIndex int) (map[string]any, bool) { + index := candidate.Index + if index == 0 && fallbackIndex > 0 { + index = fallbackIndex + } + + delta := make(map[string]any) + if !s.roleSent { + delta["role"] = "assistant" + s.roleSent = true + } + + content, toolCalls := openAIMessageFromGeminiParts(candidate.Content.Parts) + if content != "" { + delta["content"] = content + } + if len(toolCalls) > 0 { + s.sawToolCalls = true + delta["tool_calls"] = streamToolCalls(toolCalls) + } + + finish := finishReasonFromGemini(candidate.FinishReason, s.sawToolCalls) + if len(delta) == 0 && finish == "" { + return nil, false + } + choice := map[string]any{ + "index": index, + "delta": delta, + "finish_reason": nil, + } + if finish != "" { + choice["finish_reason"] = finish + } + return choice, true +} + +func streamToolCalls(toolCalls []core.ToolCall) []map[string]any { + out := make([]map[string]any, 0, len(toolCalls)) + for i, call := range toolCalls { + out = append(out, map[string]any{ + "index": i, + "id": call.ID, + "type": "function", + "function": map[string]any{ + "name": call.Function.Name, + "arguments": call.Function.Arguments, + }, + }) + } + return out +} + +func geminiUsageMap(usage geminiUsageMetadata) map[string]any { + coreUsage := usageFromGemini(usage) + if coreUsage.PromptTokens == 0 && coreUsage.CompletionTokens == 0 && coreUsage.TotalTokens == 0 && len(coreUsage.RawUsage) == 0 { + return nil + } + out := map[string]any{ + "prompt_tokens": coreUsage.PromptTokens, + "completion_tokens": coreUsage.CompletionTokens, + "total_tokens": coreUsage.TotalTokens, + } + if coreUsage.PromptTokensDetails != nil { + out["prompt_tokens_details"] = coreUsage.PromptTokensDetails + } + if coreUsage.CompletionTokensDetails != nil { + out["completion_tokens_details"] = coreUsage.CompletionTokensDetails + } + if len(coreUsage.RawUsage) > 0 { + out["raw_usage"] = coreUsage.RawUsage + } + return out +} + +func writeOpenAIStreamChunk(out io.Writer, chunk map[string]any) error { + body, err := json.Marshal(chunk) + if err != nil { + return err + } + _, err = out.Write([]byte("data: " + string(body) + "\n\n")) + return err +} From b2527e31c1a04d6a9aaa06755d638fc3d1322f52 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Mon, 4 May 2026 17:53:38 +0200 Subject: [PATCH 2/8] fix(gemini): account for native usage pricing --- internal/core/types.go | 2 + internal/core/types_test.go | 18 ++++++ internal/providers/gemini/gemini_test.go | 64 +++++++++++++++++++ internal/providers/gemini/native.go | 78 ++++++++++++++++++++++-- internal/usage/cost.go | 69 +++++++++++++++++++++ internal/usage/cost_test.go | 61 ++++++++++++++++++ internal/usage/extractor.go | 6 ++ internal/usage/extractor_test.go | 33 ++++++++++ 8 files changed, 327 insertions(+), 4 deletions(-) diff --git a/internal/core/types.go b/internal/core/types.go index 17c6a18ff..831fdb879 100644 --- a/internal/core/types.go +++ b/internal/core/types.go @@ -260,6 +260,7 @@ type ModelPricing struct { // ModelPricingTier represents a volume-based pricing tier. type ModelPricingTier struct { + UpToTokens *float64 `json:"up_to_tokens,omitempty" yaml:"up_to_tokens,omitempty"` UpToMtok *float64 `json:"up_to_mtok,omitempty" yaml:"up_to_mtok,omitempty"` InputPerMtok *float64 `json:"input_per_mtok,omitempty" yaml:"input_per_mtok,omitempty"` OutputPerMtok *float64 `json:"output_per_mtok,omitempty" yaml:"output_per_mtok,omitempty"` @@ -318,6 +319,7 @@ func (p *ModelPricing) Clone() *ModelPricing { tiers := make([]ModelPricingTier, len(p.Tiers)) for i, t := range p.Tiers { tiers[i] = ModelPricingTier{ + UpToTokens: cloneFloatPtr(t.UpToTokens), UpToMtok: cloneFloatPtr(t.UpToMtok), InputPerMtok: cloneFloatPtr(t.InputPerMtok), OutputPerMtok: cloneFloatPtr(t.OutputPerMtok), diff --git a/internal/core/types_test.go b/internal/core/types_test.go index 776a0e9e4..c81f0d95c 100644 --- a/internal/core/types_test.go +++ b/internal/core/types_test.go @@ -5,6 +5,24 @@ import ( "testing" ) +func TestModelPricingTierUnmarshalUpToTokens(t *testing.T) { + var pricing ModelPricing + if err := json.Unmarshal([]byte(`{ + "currency": "USD", + "tiers": [ + {"up_to_tokens": 200000, "input_per_mtok": 1.25, "output_per_mtok": 10.0} + ] + }`), &pricing); err != nil { + t.Fatalf("json.Unmarshal() error = %v, want nil", err) + } + if len(pricing.Tiers) != 1 { + t.Fatalf("len(Tiers) = %d, want 1", len(pricing.Tiers)) + } + if pricing.Tiers[0].UpToTokens == nil || *pricing.Tiers[0].UpToTokens != 200000 { + t.Fatalf("UpToTokens = %#v, want 200000", pricing.Tiers[0].UpToTokens) + } +} + func TestMessageUnmarshalJSON_AllowsNullContent(t *testing.T) { payload := []byte(`{ "role":"assistant", diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 9d573050b..3c138b7be 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -257,6 +257,70 @@ func TestChatCompletion_UsesNativeGenerateContentByDefault(t *testing.T) { } } +func TestChatCompletion_NativeUsageMetadata(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "gemini-native-usage", + "candidates": [{ + "index": 0, + "content": {"role": "model", "parts": [{"text": "Done"}]}, + "finishReason": "STOP" + }], + "usageMetadata": { + "promptTokenCount": 100, + "cachedContentTokenCount": 40, + "candidatesTokenCount": 20, + "thoughtsTokenCount": 7, + "totalTokenCount": 127, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 60}, + {"modality": "AUDIO", "tokenCount": 40} + ], + "candidatesTokensDetails": [ + {"modality": "AUDIO", "tokenCount": 5} + ] + } + }`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetBaseURL(server.URL) + + resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if resp.Usage.PromptTokens != 100 { + t.Fatalf("PromptTokens = %d, want 100", resp.Usage.PromptTokens) + } + if resp.Usage.CompletionTokens != 27 { + t.Fatalf("CompletionTokens = %d, want candidates + thoughts = 27", resp.Usage.CompletionTokens) + } + if resp.Usage.TotalTokens != 127 { + t.Fatalf("TotalTokens = %d, want 127", resp.Usage.TotalTokens) + } + if resp.Usage.PromptTokensDetails == nil || resp.Usage.PromptTokensDetails.CachedTokens != 40 || resp.Usage.PromptTokensDetails.AudioTokens != 40 { + t.Fatalf("PromptTokensDetails = %+v, want cached=40 audio=40", resp.Usage.PromptTokensDetails) + } + if resp.Usage.CompletionTokensDetails == nil || resp.Usage.CompletionTokensDetails.ReasoningTokens != 7 || resp.Usage.CompletionTokensDetails.AudioTokens != 5 { + t.Fatalf("CompletionTokensDetails = %+v, want reasoning=7 audio=5", resp.Usage.CompletionTokensDetails) + } + if resp.Usage.RawUsage["prompt_cached_tokens"] != 40 { + t.Fatalf("RawUsage[prompt_cached_tokens] = %#v, want 40", resp.Usage.RawUsage["prompt_cached_tokens"]) + } + if resp.Usage.RawUsage["completion_reasoning_tokens"] != 7 { + t.Fatalf("RawUsage[completion_reasoning_tokens] = %#v, want 7", resp.Usage.RawUsage["completion_reasoning_tokens"]) + } +} + func TestChatCompletion_NativeFunctionCallTranslation(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, err := io.ReadAll(r.Body) diff --git a/internal/providers/gemini/native.go b/internal/providers/gemini/native.go index b39b71d8d..33bd1d2f2 100644 --- a/internal/providers/gemini/native.go +++ b/internal/providers/gemini/native.go @@ -634,25 +634,45 @@ func openAIMessageFromGeminiParts(parts []geminiPart) (string, []core.ToolCall) } func usageFromGemini(usage geminiUsageMetadata) core.Usage { + completionTokens := usage.CandidatesTokenCount + usage.ThoughtsTokenCount out := core.Usage{ PromptTokens: usage.PromptTokenCount, - CompletionTokens: usage.CandidatesTokenCount, + CompletionTokens: completionTokens, TotalTokens: usage.TotalTokenCount, } - if out.TotalTokens == 0 { + minimumTotal := out.PromptTokens + out.CompletionTokens + if out.TotalTokens < minimumTotal { out.TotalTokens = out.PromptTokens + out.CompletionTokens } + raw := make(map[string]any) + promptDetails := promptTokenDetailsFromGemini(usage.PromptTokensDetails) if usage.CachedContentTokenCount > 0 { raw["cached_content_token_count"] = usage.CachedContentTokenCount - out.PromptTokensDetails = &core.PromptTokensDetails{CachedTokens: usage.CachedContentTokenCount} + raw["prompt_cached_tokens"] = usage.CachedContentTokenCount + if promptDetails == nil { + promptDetails = &core.PromptTokensDetails{} + } + promptDetails.CachedTokens = usage.CachedContentTokenCount } + if promptDetails != nil { + out.PromptTokensDetails = promptDetails + } + if usage.ToolUsePromptTokenCount > 0 { raw["tool_use_prompt_token_count"] = usage.ToolUsePromptTokenCount } + completionDetails := completionTokenDetailsFromGemini(usage.CandidatesTokensDetails) if usage.ThoughtsTokenCount > 0 { raw["thoughts_token_count"] = usage.ThoughtsTokenCount - out.CompletionTokensDetails = &core.CompletionTokensDetails{ReasoningTokens: usage.ThoughtsTokenCount} + raw["completion_reasoning_tokens"] = usage.ThoughtsTokenCount + if completionDetails == nil { + completionDetails = &core.CompletionTokensDetails{} + } + completionDetails.ReasoningTokens = usage.ThoughtsTokenCount + } + if completionDetails != nil { + out.CompletionTokensDetails = completionDetails } if len(raw) > 0 { out.RawUsage = raw @@ -660,6 +680,56 @@ func usageFromGemini(usage geminiUsageMetadata) core.Usage { return out } +func promptTokenDetailsFromGemini(raw json.RawMessage) *core.PromptTokensDetails { + if len(raw) == 0 { + return nil + } + var counts []geminiModalityTokenCount + if err := json.Unmarshal(raw, &counts); err != nil { + return nil + } + var out core.PromptTokensDetails + for _, count := range counts { + switch strings.ToUpper(strings.TrimSpace(count.Modality)) { + case "AUDIO": + out.AudioTokens += count.TokenCount + case "IMAGE": + out.ImageTokens += count.TokenCount + case "TEXT": + out.TextTokens += count.TokenCount + } + } + if out == (core.PromptTokensDetails{}) { + return nil + } + return &out +} + +func completionTokenDetailsFromGemini(raw json.RawMessage) *core.CompletionTokensDetails { + if len(raw) == 0 { + return nil + } + var counts []geminiModalityTokenCount + if err := json.Unmarshal(raw, &counts); err != nil { + return nil + } + var out core.CompletionTokensDetails + for _, count := range counts { + if strings.EqualFold(strings.TrimSpace(count.Modality), "AUDIO") { + out.AudioTokens += count.TokenCount + } + } + if out == (core.CompletionTokensDetails{}) { + return nil + } + return &out +} + +type geminiModalityTokenCount struct { + Modality string `json:"modality"` + TokenCount int `json:"tokenCount"` +} + func finishReasonFromGemini(reason string, hasToolCalls bool) string { if hasToolCalls { return "tool_calls" diff --git a/internal/usage/cost.go b/internal/usage/cost.go index eb3f571b6..567e452c9 100644 --- a/internal/usage/cost.go +++ b/internal/usage/cost.go @@ -71,7 +71,13 @@ var providerMappings = map[string][]tokenCostMapping{ "gemini": { {rawDataKey: "cached_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok, includedInBase: true}, {rawDataKey: "prompt_cached_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok, includedInBase: true}, + {rawDataKey: "cached_content_token_count", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok, includedInBase: true}, {rawDataKey: "thought_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok, includedInBase: true}, + {rawDataKey: "thoughts_token_count", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok, includedInBase: true}, + {rawDataKey: "reasoning_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok, includedInBase: true}, + {rawDataKey: "completion_reasoning_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok, includedInBase: true}, + {rawDataKey: "prompt_audio_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.AudioInputPerMtok }, side: sideInput, unit: unitPerMtok, includedInBase: true}, + {rawDataKey: "completion_audio_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.AudioOutputPerMtok }, side: sideOutput, unit: unitPerMtok, includedInBase: true}, }, "groq": { {rawDataKey: "cached_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok, includedInBase: true}, @@ -95,6 +101,7 @@ var providerMappings = map[string][]tokenCostMapping{ var informationalFields = map[string]struct{}{ "prompt_text_tokens": {}, "prompt_image_tokens": {}, + "tool_use_prompt_token_count": {}, "completion_accepted_prediction_tokens": {}, "completion_rejected_prediction_tokens": {}, } @@ -109,6 +116,7 @@ func CalculateGranularCost(inputTokens, outputTokens int, rawData map[string]any if pricing == nil { return CostResult{} } + pricing = pricingForTokenCount(pricing, inputTokens) var inputCost, outputCost float64 var hasInput, hasOutput bool @@ -221,6 +229,67 @@ func CalculateGranularCost(inputTokens, outputTokens int, rawData map[string]any return result } +func pricingForTokenCount(pricing *core.ModelPricing, inputTokens int) *core.ModelPricing { + if pricing == nil || inputTokens <= 0 || len(pricing.Tiers) == 0 { + return pricing + } + + tier, ok := selectPricingTier(pricing.Tiers, inputTokens) + if !ok { + return pricing + } + + effective := *pricing + if tier.InputPerMtok != nil { + effective.InputPerMtok = tier.InputPerMtok + } + if tier.OutputPerMtok != nil { + effective.OutputPerMtok = tier.OutputPerMtok + } + return &effective +} + +func selectPricingTier(tiers []core.ModelPricingTier, inputTokens int) (core.ModelPricingTier, bool) { + type tierWithLimit struct { + tier core.ModelPricingTier + limit float64 + } + + limited := make([]tierWithLimit, 0, len(tiers)) + for _, tier := range tiers { + limit, ok := tierLimitTokens(tier) + if !ok || limit <= 0 { + continue + } + limited = append(limited, tierWithLimit{tier: tier, limit: limit}) + } + if len(limited) == 0 { + return core.ModelPricingTier{}, false + } + + sort.Slice(limited, func(i, j int) bool { + return limited[i].limit < limited[j].limit + }) + + tokenCount := float64(inputTokens) + for _, candidate := range limited { + if tokenCount <= candidate.limit { + return candidate.tier, true + } + } + return limited[len(limited)-1].tier, true +} + +func tierLimitTokens(tier core.ModelPricingTier) (float64, bool) { + if tier.UpToTokens != nil { + return *tier.UpToTokens, true + } + if tier.UpToMtok != nil { + return *tier.UpToMtok * 1_000_000, true + } + return 0, false +} + func baseRateForSide(pricing *core.ModelPricing, side costSide) *float64 { if pricing == nil { return nil diff --git a/internal/usage/cost_test.go b/internal/usage/cost_test.go index 2decd38f1..716cc7e98 100644 --- a/internal/usage/cost_test.go +++ b/internal/usage/cost_test.go @@ -130,6 +130,67 @@ func TestCalculateGranularCost_Gemini_PromptCachedTokens(t *testing.T) { } } +func TestCalculateGranularCost_Gemini_NativeUsageAliases(t *testing.T) { + pricing := &core.ModelPricing{ + InputPerMtok: new(0.30), + OutputPerMtok: new(2.50), + CachedInputPerMtok: new(0.03), + ReasoningOutputPerMtok: new(5.00), + } + rawData := map[string]any{ + "cached_content_token_count": 50_000, + "prompt_cached_tokens": 50_000, + "thoughts_token_count": 20_000, + "completion_reasoning_tokens": 20_000, + "tool_use_prompt_token_count": 1000, + } + result := CalculateGranularCost(100_000, 120_000, rawData, "gemini", pricing) + + // Input: 100k * 0.30/1M + 50k * (0.03-0.30)/1M + assertCostNear(t, "InputCost", result.InputCost, 0.0165) + // Output: 120k * 2.50/1M + 20k * (5.00-2.50)/1M + assertCostNear(t, "OutputCost", result.OutputCost, 0.35) + assertCostNear(t, "TotalCost", result.TotalCost, 0.3665) + if result.Caveat != "" { + t.Fatalf("expected no caveat for native Gemini usage aliases, got %q", result.Caveat) + } +} + +func TestCalculateGranularCost_Gemini_AudioTokens(t *testing.T) { + pricing := &core.ModelPricing{ + InputPerMtok: new(0.30), + OutputPerMtok: new(2.50), + AudioInputPerMtok: new(1.00), + AudioOutputPerMtok: new(12.00), + } + rawData := map[string]any{ + "prompt_audio_tokens": 30_000, + "completion_audio_tokens": 10_000, + } + result := CalculateGranularCost(100_000, 40_000, rawData, "gemini", pricing) + + // Input: 100k * 0.30/1M + 30k * (1.00-0.30)/1M + assertCostNear(t, "InputCost", result.InputCost, 0.051) + // Output: 40k * 2.50/1M + 10k * (12.00-2.50)/1M + assertCostNear(t, "OutputCost", result.OutputCost, 0.195) +} + +func TestCalculateGranularCost_TieredPricingUsesPromptTokenThreshold(t *testing.T) { + pricing := &core.ModelPricing{ + InputPerMtok: new(1.25), + OutputPerMtok: new(10.0), + Tiers: []core.ModelPricingTier{ + {UpToTokens: new(200_000.0), InputPerMtok: new(1.25), OutputPerMtok: new(10.0)}, + {UpToTokens: new(1_048_576.0), InputPerMtok: new(2.50), OutputPerMtok: new(15.0)}, + }, + } + result := CalculateGranularCost(250_000, 10_000, nil, "gemini", pricing) + + assertCostNear(t, "InputCost", result.InputCost, 0.625) + assertCostNear(t, "OutputCost", result.OutputCost, 0.15) + assertCostNear(t, "TotalCost", result.TotalCost, 0.775) +} + func TestCalculateGranularCost_XAI_ImageTokens(t *testing.T) { pricing := &core.ModelPricing{ InputPerMtok: new(2.0), diff --git a/internal/usage/extractor.go b/internal/usage/extractor.go index 746515265..7c91a2883 100644 --- a/internal/usage/extractor.go +++ b/internal/usage/extractor.go @@ -361,11 +361,17 @@ func pricingForEndpoint(pricing *core.ModelPricing, endpoint string) *core.Model } effective := *pricing + usesBatchRate := false if pricing.BatchInputPerMtok != nil { effective.InputPerMtok = pricing.BatchInputPerMtok + usesBatchRate = true } if pricing.BatchOutputPerMtok != nil { effective.OutputPerMtok = pricing.BatchOutputPerMtok + usesBatchRate = true + } + if usesBatchRate { + effective.Tiers = nil } return &effective } diff --git a/internal/usage/extractor_test.go b/internal/usage/extractor_test.go index d2c6399a3..da87983c2 100644 --- a/internal/usage/extractor_test.go +++ b/internal/usage/extractor_test.go @@ -707,6 +707,39 @@ func TestExtractFromChatResponse_WithBatchPricingEndpoint(t *testing.T) { } } +func TestExtractFromChatResponse_BatchPricingIgnoresStandardTiers(t *testing.T) { + pricing := &core.ModelPricing{ + InputPerMtok: new(4.0), + OutputPerMtok: new(8.0), + BatchInputPerMtok: new(1.0), + BatchOutputPerMtok: new(2.0), + Tiers: []core.ModelPricingTier{ + {UpToTokens: new(200_000.0), InputPerMtok: new(4.0), OutputPerMtok: new(8.0)}, + {UpToTokens: new(1_048_576.0), InputPerMtok: new(40.0), OutputPerMtok: new(80.0)}, + }, + } + resp := &core.ChatResponse{ + ID: "chatcmpl-batch-tiered", + Model: "gpt-4o", + Usage: core.Usage{ + PromptTokens: 250_000, + CompletionTokens: 10_000, + TotalTokens: 260_000, + }, + } + + entry := ExtractFromChatResponse(resp, "req-batch-tiered", "openai", "/v1/batches", pricing) + if entry == nil { + t.Fatal("expected non-nil entry") + } + if math.Abs(*entry.InputCost-0.25) > 1e-9 { + t.Errorf("InputCost = %f, want 0.25", *entry.InputCost) + } + if math.Abs(*entry.OutputCost-0.02) > 1e-9 { + t.Errorf("OutputCost = %f, want 0.02", *entry.OutputCost) + } +} + func TestExtractFromChatResponse_WithBatchPricingSubpathEndpoint(t *testing.T) { pricing := &core.ModelPricing{ InputPerMtok: new(4.0), From ec2e6664dbe926b8bc1ca0859a3329ce5eb8c490 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Mon, 4 May 2026 18:13:42 +0200 Subject: [PATCH 3/8] fix(gemini): handle native stream edge cases --- internal/providers/gemini/gemini.go | 14 +- internal/providers/gemini/gemini_test.go | 224 +++++++++++++++++- internal/providers/gemini/native.go | 53 ++++- internal/providers/gemini/native_stream.go | 53 ++++- .../contract/replay_provider_helpers_test.go | 1 + 5 files changed, 319 insertions(+), 26 deletions(-) diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index c88fc3e37..8e16104c3 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -51,11 +51,15 @@ type Provider struct { // New creates a new Gemini provider. func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { baseURL := providers.ResolveBaseURL(providerCfg.BaseURL, defaultOpenAICompatibleBaseURL) + useNativeAPI := useNativeAPIFromEnv() + if baseURL != defaultOpenAICompatibleBaseURL { + useNativeAPI = false + } p := &Provider{ httpClient: nil, apiKey: providerCfg.APIKey, hooks: opts.Hooks, - useNativeAPI: useNativeAPIFromEnv(), + useNativeAPI: useNativeAPI, modelsURL: defaultModelsBaseURL, modelsClientConf: llmclient.Config{ ProviderName: "gemini", @@ -103,9 +107,6 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H // SetBaseURL allows configuring a custom base URL for the provider func (p *Provider) SetBaseURL(url string) { p.client.SetBaseURL(url) - if p.nativeClient != nil { - p.nativeClient.SetBaseURL(url) - } } // SetModelsURL allows configuring a custom models API base URL. @@ -113,6 +114,9 @@ func (p *Provider) SetBaseURL(url string) { func (p *Provider) SetModelsURL(url string) { p.modelsURL = url p.modelsClientConf.BaseURL = url + if p.nativeClient != nil { + p.nativeClient.SetBaseURL(url) + } } // setHeaders sets the required headers for Gemini API requests @@ -212,7 +216,7 @@ func (p *Provider) nativeChatCompletion(ctx context.Context, req *core.ChatReque if err != nil { return nil, err } - return nativeChatResponse(req, &geminiResp), nil + return nativeChatResponse(req, &geminiResp) } // StreamChatCompletion returns a raw response body for streaming (caller must close) diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 3c138b7be..6c93ebdf2 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -3,6 +3,7 @@ package gemini import ( "context" "encoding/json" + "errors" "io" "net/http" "net/http/httptest" @@ -38,6 +39,75 @@ func TestNew_ReturnsProvider(t *testing.T) { } } +func TestNew_CustomBaseURLDisablesNativeMode(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + provider := New(providers.ProviderConfig{ + APIKey: "test-api-key", + BaseURL: "https://proxy.example.com/v1beta/openai", + }, providers.ProviderOptions{}) + + geminiProvider, ok := provider.(*Provider) + if !ok { + t.Fatalf("provider type = %T, want *Provider", provider) + } + if geminiProvider.useNativeAPI { + t.Fatal("useNativeAPI = true, want false for custom OpenAI-compatible base URL") + } +} + +func TestSetBaseURLDoesNotMutateNativeClient(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + nativeHit := false + nativeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + nativeHit = true + if r.URL.Path != "/models/gemini-2.5-flash:generateContent" { + t.Errorf("native path = %q, want generateContent path", r.URL.Path) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "gemini-native-baseurl", + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP" + }] + }`)) + })) + defer nativeServer.Close() + + openAICompatHit := false + openAICompatServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + openAICompatHit = true + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error":{"message":"native client used OpenAI-compatible base URL"}}`)) + })) + defer openAICompatServer.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetModelsURL(nativeServer.URL) + provider.SetBaseURL(openAICompatServer.URL + "/v1beta/openai") + + resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp == nil || resp.ID != "gemini-native-baseurl" { + t.Fatalf("response = %+v, want native response", resp) + } + if !nativeHit { + t.Fatal("native server was not called") + } + if openAICompatHit { + t.Fatal("OpenAI-compatible server was called by native client") + } +} + func TestChatCompletion(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "false") @@ -226,7 +296,7 @@ func TestChatCompletion_UsesNativeGenerateContentByDefault(t *testing.T) { defer server.Close() provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) - provider.SetBaseURL(server.URL) + provider.SetModelsURL(server.URL) maxTokens := 128 resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ @@ -286,7 +356,7 @@ func TestChatCompletion_NativeUsageMetadata(t *testing.T) { defer server.Close() provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) - provider.SetBaseURL(server.URL) + provider.SetModelsURL(server.URL) resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ Model: "gemini-2.5-flash", @@ -321,6 +391,71 @@ func TestChatCompletion_NativeUsageMetadata(t *testing.T) { } } +func TestChatCompletion_NativeBlockedPromptReturnsError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "gemini-blocked", + "promptFeedback": { + "blockReason": "SAFETY", + "blockReasonMessage": "unsafe prompt" + } + }`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetModelsURL(server.URL) + + _, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "blocked"}, + }, + }) + if err == nil { + t.Fatal("expected blocked prompt error, got nil") + } + var gatewayErr *core.GatewayError + if !errors.As(err, &gatewayErr) { + t.Fatalf("error = %T %[1]v, want *core.GatewayError", err) + } + if gatewayErr.Type != core.ErrorTypeProvider { + t.Fatalf("error type = %q, want provider_error", gatewayErr.Type) + } + if !strings.Contains(gatewayErr.Message, "SAFETY: unsafe prompt") { + t.Fatalf("message = %q, want block reason", gatewayErr.Message) + } +} + +func TestChatCompletion_NativeRejectsRemoteImageURL(t *testing.T) { + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + + _, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{{ + Role: "user", + Content: []core.ContentPart{ + {Type: "text", Text: "Describe the image."}, + {Type: "image_url", ImageURL: &core.ImageURLContent{URL: "https://example.com/image.png"}}, + }, + }}, + }) + if err == nil { + t.Fatal("expected remote image_url error, got nil") + } + var gatewayErr *core.GatewayError + if !errors.As(err, &gatewayErr) { + t.Fatalf("error = %T %[1]v, want *core.GatewayError", err) + } + if gatewayErr.Type != core.ErrorTypeInvalidRequest { + t.Fatalf("error type = %q, want invalid_request_error", gatewayErr.Type) + } + if !strings.Contains(gatewayErr.Message, "supports only data: URLs") { + t.Fatalf("message = %q, want data URL guidance", gatewayErr.Message) + } +} + func TestChatCompletion_NativeFunctionCallTranslation(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, err := io.ReadAll(r.Body) @@ -362,7 +497,7 @@ func TestChatCompletion_NativeFunctionCallTranslation(t *testing.T) { defer server.Close() provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) - provider.SetBaseURL(server.URL) + provider.SetModelsURL(server.URL) resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ Model: "gemini-2.5-flash", @@ -534,7 +669,7 @@ data: {"responseId":"gemini-stream-123","candidates":[{"content":{"role":"model" defer server.Close() provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) - provider.SetBaseURL(server.URL) + provider.SetModelsURL(server.URL) body, err := provider.StreamChatCompletion(context.Background(), &core.ChatRequest{ Model: "gemini-2.5-flash", @@ -570,6 +705,83 @@ data: {"responseId":"gemini-stream-123","candidates":[{"content":{"role":"model" } } +func TestStreamChatCompletion_NativePerChoiceState(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`data: {"responseId":"gemini-stream-choice-state","candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"id":"call_0","name":"lookup_weather","args":{"city":"Warsaw"}}}]},"finishReason":"STOP"},{"index":1,"content":{"role":"model","parts":[{"text":"plain text"}]},"finishReason":"STOP"}]} + +`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetModelsURL(server.URL) + + body, err := provider.StreamChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer func() { _ = body.Close() }() + + raw, err := io.ReadAll(body) + if err != nil { + t.Fatalf("failed to read stream: %v", err) + } + stream := string(raw) + if got := strings.Count(stream, `"role":"assistant"`); got != 2 { + t.Fatalf("assistant role count = %d, want 2 in stream %q", got, stream) + } + if got := strings.Count(stream, `"finish_reason":"tool_calls"`); got != 1 { + t.Fatalf("tool_calls finish count = %d, want 1 in stream %q", got, stream) + } + if got := strings.Count(stream, `"finish_reason":"stop"`); got != 1 { + t.Fatalf("stop finish count = %d, want 1 in stream %q", got, stream) + } +} + +func TestStreamChatCompletion_NativeBlockedPromptEmitsError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`data: {"responseId":"gemini-stream-blocked","promptFeedback":{"blockReason":"SAFETY","blockReasonMessage":"unsafe prompt"}} + +`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetModelsURL(server.URL) + + body, err := provider.StreamChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "blocked"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer func() { _ = body.Close() }() + + raw, err := io.ReadAll(body) + if err != nil { + t.Fatalf("failed to read stream: %v", err) + } + stream := string(raw) + if !strings.Contains(stream, `"type":"provider_error"`) || !strings.Contains(stream, "Gemini blocked prompt: SAFETY: unsafe prompt") { + t.Fatalf("stream = %q, want normalized provider error", stream) + } + if !strings.Contains(stream, "data: [DONE]") { + t.Fatalf("stream = %q, want [DONE]", stream) + } +} + func TestListModels(t *testing.T) { tests := []struct { name string @@ -657,7 +869,7 @@ func TestListModels(t *testing.T) { defer server.Close() provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) - provider.modelsURL = server.URL + provider.SetModelsURL(server.URL) resp, err := provider.ListModels(context.Background()) @@ -755,6 +967,8 @@ func TestResponses(t *testing.T) { } func TestStreamResponses(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "false") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`data: {"id":"gemini-123","object":"chat.completion.chunk","created":1677652288,"model":"gemini-2.0-flash","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]} diff --git a/internal/providers/gemini/native.go b/internal/providers/gemini/native.go index 33bd1d2f2..4fcc8a026 100644 --- a/internal/providers/gemini/native.go +++ b/internal/providers/gemini/native.go @@ -99,6 +99,11 @@ type geminiGenerateContentResponse struct { ModelStatus json.RawMessage `json:"modelStatus,omitempty"` } +type geminiPromptFeedback struct { + BlockReason string `json:"blockReason,omitempty"` + BlockReasonMessage string `json:"blockReasonMessage,omitempty"` +} + type geminiCandidate struct { Content geminiContent `json:"content,omitempty"` FinishReason string `json:"finishReason,omitempty"` @@ -302,11 +307,10 @@ func geminiPartFromImageURL(image *core.ImageURLContent) (geminiPart, error) { return geminiPart{InlineData: &geminiBlob{MimeType: mimeType, Data: data}}, nil } - mimeType := image.MediaType - if mimeType == "" { - mimeType = "image/jpeg" - } - return geminiPart{FileData: &geminiFileData{MimeType: mimeType, FileURI: rawURL}}, nil + return geminiPart{}, core.NewInvalidRequestError( + "gemini native image_url supports only data: URLs; remote URLs must be uploaded via the Gemini Files API or fetched by a future adapter path", + nil, + ) } func parseDataURL(rawURL string) (string, string, error) { @@ -565,7 +569,11 @@ func geminiCachedContent(req *core.ChatRequest) string { return cached } -func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContentResponse) *core.ChatResponse { +func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContentResponse) (*core.ChatResponse, error) { + if err := geminiBlockedPromptError(geminiResp); err != nil { + return nil, err + } + created := time.Now().Unix() respID := geminiResp.ResponseID if respID == "" { @@ -596,7 +604,38 @@ func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContent FinishReason: finishReasonFromGemini(candidate.FinishReason, len(toolCalls) > 0), }) } - return resp + return resp, nil +} + +func geminiBlockedPromptError(resp *geminiGenerateContentResponse) *core.GatewayError { + if resp == nil || len(resp.Candidates) > 0 { + return nil + } + reason := geminiPromptBlockReason(resp.PromptFeedback) + if reason == "" { + return nil + } + return nativeProviderError("Gemini blocked prompt: "+reason, nil) +} + +func geminiPromptBlockReason(raw json.RawMessage) string { + if len(raw) == 0 { + return "" + } + var feedback geminiPromptFeedback + if err := json.Unmarshal(raw, &feedback); err != nil { + return "" + } + reason := strings.TrimSpace(feedback.BlockReason) + message := strings.TrimSpace(feedback.BlockReasonMessage) + switch { + case reason != "" && message != "": + return fmt.Sprintf("%s: %s", reason, message) + case reason != "": + return reason + default: + return message + } } func openAIMessageFromGeminiParts(parts []geminiPart) (string, []core.ToolCall) { diff --git a/internal/providers/gemini/native_stream.go b/internal/providers/gemini/native_stream.go index 32168e01f..c888f4f89 100644 --- a/internal/providers/gemini/native_stream.go +++ b/internal/providers/gemini/native_stream.go @@ -80,13 +80,18 @@ type geminiStreamState struct { includeUsage bool created int64 responseID string + choices map[int]*geminiChoiceStreamState + stopped bool +} + +type geminiChoiceStreamState struct { roleSent bool sawToolCalls bool } func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { raw = strings.TrimSpace(raw) - if raw == "" || raw == "[DONE]" { + if raw == "" || raw == "[DONE]" || s.stopped { return nil } @@ -101,6 +106,11 @@ func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { } } + if err := geminiBlockedPromptError(&event); err != nil { + s.stopped = true + return writeOpenAIStreamError(out, err) + } + for i, candidate := range event.Candidates { choice, ok := s.chatChunkChoice(candidate, i) if !ok { @@ -144,15 +154,13 @@ func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { } func (s *geminiStreamState) chatChunkChoice(candidate geminiCandidate, fallbackIndex int) (map[string]any, bool) { - index := candidate.Index - if index == 0 && fallbackIndex > 0 { - index = fallbackIndex - } + index := streamChoiceIndex(candidate, fallbackIndex) + state := s.choiceState(index) delta := make(map[string]any) - if !s.roleSent { + if !state.roleSent { delta["role"] = "assistant" - s.roleSent = true + state.roleSent = true } content, toolCalls := openAIMessageFromGeminiParts(candidate.Content.Parts) @@ -160,11 +168,11 @@ func (s *geminiStreamState) chatChunkChoice(candidate geminiCandidate, fallbackI delta["content"] = content } if len(toolCalls) > 0 { - s.sawToolCalls = true + state.sawToolCalls = true delta["tool_calls"] = streamToolCalls(toolCalls) } - finish := finishReasonFromGemini(candidate.FinishReason, s.sawToolCalls) + finish := finishReasonFromGemini(candidate.FinishReason, state.sawToolCalls) if len(delta) == 0 && finish == "" { return nil, false } @@ -179,6 +187,26 @@ func (s *geminiStreamState) chatChunkChoice(candidate geminiCandidate, fallbackI return choice, true } +func streamChoiceIndex(candidate geminiCandidate, fallbackIndex int) int { + index := candidate.Index + if index == 0 && fallbackIndex > 0 { + index = fallbackIndex + } + return index +} + +func (s *geminiStreamState) choiceState(index int) *geminiChoiceStreamState { + if s.choices == nil { + s.choices = make(map[int]*geminiChoiceStreamState) + } + state := s.choices[index] + if state == nil { + state = &geminiChoiceStreamState{} + s.choices[index] = state + } + return state +} + func streamToolCalls(toolCalls []core.ToolCall) []map[string]any { out := make([]map[string]any, 0, len(toolCalls)) for i, call := range toolCalls { @@ -225,3 +253,10 @@ func writeOpenAIStreamChunk(out io.Writer, chunk map[string]any) error { _, err = out.Write([]byte("data: " + string(body) + "\n\n")) return err } + +func writeOpenAIStreamError(out io.Writer, err *core.GatewayError) error { + if err == nil { + return nil + } + return writeOpenAIStreamChunk(out, err.ToJSON()) +} diff --git a/tests/contract/replay_provider_helpers_test.go b/tests/contract/replay_provider_helpers_test.go index 0cfabef9d..7f5f3452b 100644 --- a/tests/contract/replay_provider_helpers_test.go +++ b/tests/contract/replay_provider_helpers_test.go @@ -13,6 +13,7 @@ import ( func newGeminiReplayProvider(t *testing.T, routes map[string]replayRoute) core.Provider { t.Helper() + t.Setenv("USE_GOOGLE_GEMINI_NATIVE_API", "false") client := newReplayHTTPClient(t, routes) provider := gemini.NewWithHTTPClient("test-api-key", client, llmclient.Hooks{}) provider.SetBaseURL("https://replay.local") From bfd9ee593a7e54e7d8b9ac9cd6df93d4b9610d0b Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Mon, 4 May 2026 18:34:19 +0200 Subject: [PATCH 4/8] fix(gemini): address native api review findings --- internal/core/types_test.go | 18 +++ internal/providers/gemini/gemini.go | 30 +++-- internal/providers/gemini/gemini_test.go | 143 ++++++++++++++++++--- internal/providers/gemini/native.go | 18 ++- internal/providers/gemini/native_stream.go | 7 +- internal/usage/extractor.go | 20 ++- internal/usage/extractor_test.go | 66 ++++++++++ 7 files changed, 259 insertions(+), 43 deletions(-) diff --git a/internal/core/types_test.go b/internal/core/types_test.go index c81f0d95c..b858cd3f8 100644 --- a/internal/core/types_test.go +++ b/internal/core/types_test.go @@ -21,6 +21,24 @@ func TestModelPricingTierUnmarshalUpToTokens(t *testing.T) { if pricing.Tiers[0].UpToTokens == nil || *pricing.Tiers[0].UpToTokens != 200000 { t.Fatalf("UpToTokens = %#v, want 200000", pricing.Tiers[0].UpToTokens) } + + cloned := pricing.Clone() + if cloned == nil { + t.Fatal("Clone() = nil, want pricing copy") + } + if len(cloned.Tiers) != 1 { + t.Fatalf("len(cloned.Tiers) = %d, want 1", len(cloned.Tiers)) + } + if cloned.Tiers[0].UpToTokens == nil || *cloned.Tiers[0].UpToTokens != 200000 { + t.Fatalf("cloned UpToTokens = %#v, want 200000", cloned.Tiers[0].UpToTokens) + } + if cloned.Tiers[0].UpToTokens == pricing.Tiers[0].UpToTokens { + t.Fatal("Clone() reused UpToTokens pointer, want deep copy") + } + *pricing.Tiers[0].UpToTokens = 123 + if *cloned.Tiers[0].UpToTokens != 200000 { + t.Fatalf("cloned UpToTokens after source mutation = %v, want 200000", *cloned.Tiers[0].UpToTokens) + } } func TestMessageUnmarshalJSON_AllowsNullContent(t *testing.T) { diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 8e16104c3..91796de25 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -51,19 +51,16 @@ type Provider struct { // New creates a new Gemini provider. func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { baseURL := providers.ResolveBaseURL(providerCfg.BaseURL, defaultOpenAICompatibleBaseURL) - useNativeAPI := useNativeAPIFromEnv() - if baseURL != defaultOpenAICompatibleBaseURL { - useNativeAPI = false - } + modelsURL := defaultModelsBaseURL p := &Provider{ httpClient: nil, apiKey: providerCfg.APIKey, hooks: opts.Hooks, - useNativeAPI: useNativeAPI, - modelsURL: defaultModelsBaseURL, + useNativeAPI: useNativeAPIForBaseURLs(baseURL, modelsURL), + modelsURL: modelsURL, modelsClientConf: llmclient.Config{ ProviderName: "gemini", - BaseURL: defaultModelsBaseURL, + BaseURL: modelsURL, Retry: opts.Resilience.Retry, Hooks: opts.Hooks, CircuitBreaker: opts.Resilience.CircuitBreaker, @@ -87,17 +84,19 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H if httpClient == nil { httpClient = http.DefaultClient } + baseURL := defaultOpenAICompatibleBaseURL + modelsURL := defaultModelsBaseURL p := &Provider{ httpClient: httpClient, apiKey: apiKey, hooks: hooks, - useNativeAPI: useNativeAPIFromEnv(), - modelsURL: defaultModelsBaseURL, + useNativeAPI: useNativeAPIForBaseURLs(baseURL, modelsURL), + modelsURL: modelsURL, } - modelsCfg := llmclient.DefaultConfig("gemini", defaultModelsBaseURL) + modelsCfg := llmclient.DefaultConfig("gemini", modelsURL) modelsCfg.Hooks = hooks p.modelsClientConf = modelsCfg - cfg := llmclient.DefaultConfig("gemini", defaultOpenAICompatibleBaseURL) + cfg := llmclient.DefaultConfig("gemini", baseURL) cfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) p.nativeClient = llmclient.NewWithHTTPClient(httpClient, modelsCfg, p.setNativeHeaders) @@ -107,6 +106,9 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H // SetBaseURL allows configuring a custom base URL for the provider func (p *Provider) SetBaseURL(url string) { p.client.SetBaseURL(url) + p.modelsURL = url + p.modelsClientConf.BaseURL = url + p.useNativeAPI = useNativeAPIForBaseURLs(url, p.modelsURL) } // SetModelsURL allows configuring a custom models API base URL. @@ -151,6 +153,12 @@ func useNativeAPIFromEnv() bool { } } +func useNativeAPIForBaseURLs(baseURL, modelsURL string) bool { + return useNativeAPIFromEnv() && + baseURL == defaultOpenAICompatibleBaseURL && + modelsURL == defaultModelsBaseURL +} + // adaptChatRequest rewrites a ChatRequest for Gemini's OpenAI-compatible endpoint. // Gemini uses "reasoning_effort" as a top-level string (e.g. "low", "medium", "high"), // not the nested "reasoning": {"effort": "..."} format. diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 6c93ebdf2..eb9de0177 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -56,31 +56,41 @@ func TestNew_CustomBaseURLDisablesNativeMode(t *testing.T) { } } -func TestSetBaseURLDoesNotMutateNativeClient(t *testing.T) { +func TestSetBaseURLDisablesNativeRouting(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "true") nativeHit := false nativeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { nativeHit = true - if r.URL.Path != "/models/gemini-2.5-flash:generateContent" { - t.Errorf("native path = %q, want generateContent path", r.URL.Path) - } - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte(`{ - "responseId": "gemini-native-baseurl", - "candidates": [{ - "content": {"role": "model", "parts": [{"text": "ok"}]}, - "finishReason": "STOP" - }] - }`)) + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error":{"message":"native client should not be used after SetBaseURL"}}`)) })) defer nativeServer.Close() openAICompatHit := false openAICompatServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { openAICompatHit = true - w.WriteHeader(http.StatusInternalServerError) - _, _ = w.Write([]byte(`{"error":{"message":"native client used OpenAI-compatible base URL"}}`)) + if r.URL.Path != "/v1beta/openai/chat/completions" { + t.Errorf("OpenAI-compatible path = %q, want /v1beta/openai/chat/completions", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer test-api-key" { + t.Errorf("Authorization = %q, want bearer API key", got) + } + if got := r.Header.Get("x-goog-api-key"); got != "" { + t.Errorf("x-goog-api-key = %q, want empty for OpenAI-compatible API", got) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "id": "gemini-openai-compatible-baseurl", + "object": "chat.completion", + "created": 1677652288, + "model": "gemini-2.5-flash", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }] + }`)) })) defer openAICompatServer.Close() @@ -97,14 +107,14 @@ func TestSetBaseURLDoesNotMutateNativeClient(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } - if resp == nil || resp.ID != "gemini-native-baseurl" { - t.Fatalf("response = %+v, want native response", resp) + if resp == nil || resp.ID != "gemini-openai-compatible-baseurl" { + t.Fatalf("response = %+v, want OpenAI-compatible response", resp) } - if !nativeHit { - t.Fatal("native server was not called") + if nativeHit { + t.Fatal("native server was called after SetBaseURL") } - if openAICompatHit { - t.Fatal("OpenAI-compatible server was called by native client") + if !openAICompatHit { + t.Fatal("OpenAI-compatible server was not called") } } @@ -239,6 +249,8 @@ func TestChatCompletion(t *testing.T) { } func TestChatCompletion_UsesNativeGenerateContentByDefault(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { t.Errorf("Method = %q, want %q", r.Method, http.MethodPost) @@ -328,6 +340,8 @@ func TestChatCompletion_UsesNativeGenerateContentByDefault(t *testing.T) { } func TestChatCompletion_NativeUsageMetadata(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{ @@ -391,7 +405,47 @@ func TestChatCompletion_NativeUsageMetadata(t *testing.T) { } } +func TestCopyJSONNumberAcceptsOnlyNumericValues(t *testing.T) { + tests := []struct { + name string + raw string + wantSet bool + want float64 + }{ + {name: "number", raw: `42`, wantSet: true, want: 42}, + {name: "numeric string", raw: `"42.5"`, wantSet: true, want: 42.5}, + {name: "object", raw: `{"value":42}`, wantSet: false}, + {name: "array", raw: `[42]`, wantSet: false}, + {name: "boolean", raw: `true`, wantSet: false}, + {name: "non numeric string", raw: `"fast"`, wantSet: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := map[string]any{} + copyJSONNumber(json.RawMessage(tt.raw), cfg, "value") + + got, ok := cfg["value"] + if ok != tt.wantSet { + t.Fatalf("cfg[value] set = %v, want %v; cfg = %#v", ok, tt.wantSet, cfg) + } + if !tt.wantSet { + return + } + gotFloat, ok := got.(float64) + if !ok { + t.Fatalf("cfg[value] = %T(%[1]v), want float64", got) + } + if gotFloat != tt.want { + t.Fatalf("cfg[value] = %v, want %v", gotFloat, tt.want) + } + }) + } +} + func TestChatCompletion_NativeBlockedPromptReturnsError(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{ @@ -429,6 +483,8 @@ func TestChatCompletion_NativeBlockedPromptReturnsError(t *testing.T) { } func TestChatCompletion_NativeRejectsRemoteImageURL(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) _, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ @@ -457,6 +513,8 @@ func TestChatCompletion_NativeRejectsRemoteImageURL(t *testing.T) { } func TestChatCompletion_NativeFunctionCallTranslation(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, err := io.ReadAll(r.Body) if err != nil { @@ -635,6 +693,8 @@ data: [DONE] } func TestStreamChatCompletion_UsesNativeStreamByDefault(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/models/gemini-2.5-flash:streamGenerateContent" { t.Errorf("Path = %q, want native streamGenerateContent endpoint", r.URL.Path) @@ -700,12 +760,53 @@ data: {"responseId":"gemini-stream-123","candidates":[{"content":{"role":"model" if !strings.Contains(stream, `"usage"`) || !strings.Contains(stream, `"total_tokens":6`) { t.Fatalf("stream = %q, want usage chunk", stream) } + usageChunks := 0 + for _, chunk := range parseOpenAIStreamChunks(t, stream) { + if _, ok := chunk["usage"]; !ok { + continue + } + usageChunks++ + choices, ok := chunk["choices"].([]any) + if !ok { + t.Fatalf("usage chunk choices = %T(%[1]v), want array", chunk["choices"]) + } + if len(choices) != 0 { + t.Fatalf("usage chunk choices = %#v, want empty choices", choices) + } + } + if usageChunks != 1 { + t.Fatalf("usage chunk count = %d, want 1 in stream %q", usageChunks, stream) + } if !strings.Contains(stream, "data: [DONE]") { t.Fatalf("stream = %q, want [DONE]", stream) } } +func parseOpenAIStreamChunks(t *testing.T, stream string) []map[string]any { + t.Helper() + + var chunks []map[string]any + for _, line := range strings.Split(stream, "\n") { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "data:") { + continue + } + payload := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if payload == "" || payload == "[DONE]" { + continue + } + var chunk map[string]any + if err := json.Unmarshal([]byte(payload), &chunk); err != nil { + t.Fatalf("failed to parse stream chunk %q: %v", payload, err) + } + chunks = append(chunks, chunk) + } + return chunks +} + func TestStreamChatCompletion_NativePerChoiceState(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) @@ -746,6 +847,8 @@ func TestStreamChatCompletion_NativePerChoiceState(t *testing.T) { } func TestStreamChatCompletion_NativeBlockedPromptEmitsError(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) diff --git a/internal/providers/gemini/native.go b/internal/providers/gemini/native.go index 4fcc8a026..0a5a5a6ce 100644 --- a/internal/providers/gemini/native.go +++ b/internal/providers/gemini/native.go @@ -452,8 +452,22 @@ func copyJSONNumber(raw json.RawMessage, cfg map[string]any, key string) { return } var value any - if err := json.Unmarshal(raw, &value); err == nil { - cfg[key] = value + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + if err := decoder.Decode(&value); err != nil { + return + } + switch v := value.(type) { + case float64: + cfg[key] = v + case json.Number: + if parsed, err := v.Float64(); err == nil { + cfg[key] = parsed + } + case string: + if parsed, err := strconv.ParseFloat(v, 64); err == nil { + cfg[key] = parsed + } } } diff --git a/internal/providers/gemini/native_stream.go b/internal/providers/gemini/native_stream.go index c888f4f89..f3d590153 100644 --- a/internal/providers/gemini/native_stream.go +++ b/internal/providers/gemini/native_stream.go @@ -124,17 +124,12 @@ func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { "provider": "gemini", "choices": []map[string]any{choice}, } - if s.includeUsage { - if usage := geminiUsageMap(event.UsageMetadata); usage != nil { - chunk["usage"] = usage - } - } if err := writeOpenAIStreamChunk(out, chunk); err != nil { return err } } - if len(event.Candidates) == 0 && s.includeUsage { + if s.includeUsage { if usage := geminiUsageMap(event.UsageMetadata); usage != nil { chunk := map[string]any{ "id": s.responseID, diff --git a/internal/usage/extractor.go b/internal/usage/extractor.go index 7c91a2883..3ded63c7a 100644 --- a/internal/usage/extractor.go +++ b/internal/usage/extractor.go @@ -361,17 +361,29 @@ func pricingForEndpoint(pricing *core.ModelPricing, endpoint string) *core.Model } effective := *pricing - usesBatchRate := false + usesBatchInput := false + usesBatchOutput := false if pricing.BatchInputPerMtok != nil { effective.InputPerMtok = pricing.BatchInputPerMtok - usesBatchRate = true + usesBatchInput = true } if pricing.BatchOutputPerMtok != nil { effective.OutputPerMtok = pricing.BatchOutputPerMtok - usesBatchRate = true + usesBatchOutput = true } - if usesBatchRate { + if usesBatchInput && usesBatchOutput { effective.Tiers = nil + } else if usesBatchInput || usesBatchOutput { + effective.Tiers = make([]core.ModelPricingTier, len(pricing.Tiers)) + copy(effective.Tiers, pricing.Tiers) + for i := range effective.Tiers { + if usesBatchInput { + effective.Tiers[i].InputPerMtok = pricing.BatchInputPerMtok + } + if usesBatchOutput { + effective.Tiers[i].OutputPerMtok = pricing.BatchOutputPerMtok + } + } } return &effective } diff --git a/internal/usage/extractor_test.go b/internal/usage/extractor_test.go index da87983c2..5a2647df6 100644 --- a/internal/usage/extractor_test.go +++ b/internal/usage/extractor_test.go @@ -740,6 +740,72 @@ func TestExtractFromChatResponse_BatchPricingIgnoresStandardTiers(t *testing.T) } } +func TestExtractFromChatResponse_PartialBatchPricingPreservesOtherSideTiers(t *testing.T) { + tests := []struct { + name string + pricing *core.ModelPricing + wantInput float64 + wantOutput float64 + }{ + { + name: "batch input preserves output tier", + pricing: &core.ModelPricing{ + InputPerMtok: new(4.0), + OutputPerMtok: new(8.0), + BatchInputPerMtok: new(1.0), + Tiers: []core.ModelPricingTier{ + {UpToTokens: new(200_000.0), InputPerMtok: new(4.0), OutputPerMtok: new(8.0)}, + {UpToTokens: new(1_048_576.0), InputPerMtok: new(40.0), OutputPerMtok: new(80.0)}, + }, + }, + wantInput: 0.25, + wantOutput: 0.8, + }, + { + name: "batch output preserves input tier", + pricing: &core.ModelPricing{ + InputPerMtok: new(4.0), + OutputPerMtok: new(8.0), + BatchOutputPerMtok: new(2.0), + Tiers: []core.ModelPricingTier{ + {UpToTokens: new(200_000.0), InputPerMtok: new(4.0), OutputPerMtok: new(8.0)}, + {UpToTokens: new(1_048_576.0), InputPerMtok: new(40.0), OutputPerMtok: new(80.0)}, + }, + }, + wantInput: 10.0, + wantOutput: 0.02, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resp := &core.ChatResponse{ + ID: "chatcmpl-batch-partial-tiered", + Model: "gpt-4o", + Usage: core.Usage{ + PromptTokens: 250_000, + CompletionTokens: 10_000, + TotalTokens: 260_000, + }, + } + + entry := ExtractFromChatResponse(resp, "req-batch-partial-tiered", "openai", "/v1/batches", tt.pricing) + if entry == nil { + t.Fatal("expected non-nil entry") + } + if entry.InputCost == nil || entry.OutputCost == nil { + t.Fatalf("costs = input:%v output:%v, want both populated", entry.InputCost, entry.OutputCost) + } + if math.Abs(*entry.InputCost-tt.wantInput) > 1e-9 { + t.Errorf("InputCost = %f, want %f", *entry.InputCost, tt.wantInput) + } + if math.Abs(*entry.OutputCost-tt.wantOutput) > 1e-9 { + t.Errorf("OutputCost = %f, want %f", *entry.OutputCost, tt.wantOutput) + } + }) + } +} + func TestExtractFromChatResponse_WithBatchPricingSubpathEndpoint(t *testing.T) { pricing := &core.ModelPricing{ InputPerMtok: new(4.0), From 1af3d32c43edad35a7f5ee3527a19cc8a2c5e816 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Mon, 4 May 2026 18:52:38 +0200 Subject: [PATCH 5/8] fix(gemini): preserve models URL and document image limits --- .env.template | 3 + README.md | 2 +- config/config.example.yaml | 1 + docs/docs.json | 1 + docs/guides/gemini.mdx | 104 +++++++++++++++++++++++ internal/providers/gemini/gemini.go | 5 +- internal/providers/gemini/gemini_test.go | 59 +++++++++++++ 7 files changed, 171 insertions(+), 4 deletions(-) create mode 100644 docs/guides/gemini.mdx diff --git a/.env.template b/.env.template index 64010a94d..40b24b9a1 100644 --- a/.env.template +++ b/.env.template @@ -265,6 +265,9 @@ # GEMINI_API_KEY=... # Use Gemini's native generateContent API for chat/responses (default: true). # Set to false to use Gemini's OpenAI-compatible API for chat/responses. +# Native mode supports inline image data via data: URLs, but GoModel does not +# fetch remote image URLs or upload them through Gemini Files API yet. Set this +# to false when you need OpenAI-compatible image_url pass-through behavior. # USE_GOOGLE_GEMINI_NATIVE_API=true # OpenAI-compatible Gemini base URL used when native chat is disabled, and by # Gemini embeddings/files/batches which still rely on that compatibility surface. diff --git a/README.md b/README.md index 5ae3c4399..da6c4fde6 100644 --- a/README.md +++ b/README.md @@ -250,7 +250,7 @@ Key settings: | `ENABLE_PASSTHROUGH_ROUTES` | `true` | Enable provider-native passthrough routes under `/p/{provider}/...` | | `ALLOW_PASSTHROUGH_V1_ALIAS` | `true` | Allow `/p/{provider}/v1/...` aliases while keeping `/p/{provider}/...` canonical | | `ENABLED_PASSTHROUGH_PROVIDERS` | `openai,anthropic,openrouter,zai,vllm` | Comma-separated list of enabled passthrough providers | -| `USE_GOOGLE_GEMINI_NATIVE_API` | `true` | Use Gemini native `generateContent` for chat/responses; set `false` for Gemini's OpenAI-compatible API | +| `USE_GOOGLE_GEMINI_NATIVE_API` | `true` | Use Gemini native `generateContent` for chat/responses; set `false` for Gemini's OpenAI-compatible API and image_url pass-through behavior | | `STORAGE_TYPE` | `sqlite` | Storage backend (`sqlite`, `postgresql`, `mongodb`) | | `METRICS_ENABLED` | `false` | Enable Prometheus metrics (experimental) | | `LOGGING_ENABLED` | `false` | Enable audit logging | diff --git a/config/config.example.yaml b/config/config.example.yaml index 269e6c2b8..425758134 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -203,6 +203,7 @@ providers: api_key: "..." # Chat/responses use Gemini's native generateContent API by default. # Set USE_GOOGLE_GEMINI_NATIVE_API=false to use Gemini's OpenAI-compatible API instead. + # Native mode accepts image data URLs but does not fetch remote image URLs yet. xai: type: xai diff --git a/docs/docs.json b/docs/docs.json index 974fae706..fec382603 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -85,6 +85,7 @@ "icon": "compass", "pages": [ "guides/openclaw", + "guides/gemini", "guides/oracle", "guides/deepseek", "guides/vllm", diff --git a/docs/guides/gemini.mdx b/docs/guides/gemini.mdx new file mode 100644 index 000000000..76e337ef3 --- /dev/null +++ b/docs/guides/gemini.mdx @@ -0,0 +1,104 @@ +--- +title: "GoModel & Google Gemini" +description: "Configure Google Gemini in GoModel, choose native or OpenAI-compatible routing, and understand image_url behavior." +icon: "sparkles" +--- + +GoModel routes Gemini chat and Responses API requests through Gemini's native +`generateContent` API by default. You can switch those requests back to +Gemini's OpenAI-compatible API when you need compatibility behavior that the +native adapter does not implement yet. + +## Configure Gemini + +Env-only configuration is enough: + +```bash +export GEMINI_API_KEY="..." +``` + +Or in `config.yaml`: + +```yaml +providers: + gemini: + type: gemini + api_key: "${GEMINI_API_KEY}" +``` + +## Native versus OpenAI-compatible mode + +Gemini native mode is enabled by default: + +```bash +export USE_GOOGLE_GEMINI_NATIVE_API=true +``` + +Set it to `false` to route chat and Responses API requests through Gemini's +OpenAI-compatible `/chat/completions` endpoint: + +```bash +export USE_GOOGLE_GEMINI_NATIVE_API=false +``` + +`GEMINI_BASE_URL` configures the OpenAI-compatible Gemini base URL. GoModel also +uses that compatibility surface for Gemini embeddings, files, and batches. + +```bash +export GEMINI_BASE_URL="https://generativelanguage.googleapis.com/v1beta/openai" +``` + + + If you configure a custom `GEMINI_BASE_URL`, GoModel disables native Gemini + chat routing for that provider and sends chat requests to the configured + OpenAI-compatible endpoint. + + +## Image URL behavior + +Gemini models support image input, but the two GoModel routing modes handle +OpenAI-style `image_url` values differently. + +In native mode, GoModel converts OpenAI-compatible messages to Gemini +`generateContent` requests. That adapter currently supports inline image data +only: + +```json +{ + "type": "image_url", + "image_url": { + "url": "data:image/jpeg;base64,..." + } +} +``` + +Remote image URLs such as `https://example.com/image.png` are rejected in native +mode. Google's native Gemini API supports inline image data and Files API +references; for URL-hosted images, Google's examples fetch the URL first and +send the bytes to `generateContent`. + +Set `USE_GOOGLE_GEMINI_NATIVE_API=false` when you need GoModel to pass the +OpenAI-compatible `image_url` request shape through to Gemini's +OpenAI-compatible endpoint instead. Google documents image input for that +endpoint using the OpenAI `image_url` field. + +## Current support + +Integrated: + +- chat completions and streaming +- Responses API and streaming +- model listing through Gemini's native `/models` +- usage metadata normalization for native responses +- tool calls and function-call results +- inline image data via `data:` URLs in native mode + +Not integrated in native mode yet: + +- fetching remote `image_url` values +- uploading remote images through the Gemini Files API before a chat request + +References: + +- [Gemini image understanding](https://ai.google.dev/gemini-api/docs/image-understanding) +- [Gemini OpenAI compatibility](https://ai.google.dev/gemini-api/docs/openai) diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 91796de25..eed5cbf98 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -105,10 +105,9 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H // SetBaseURL allows configuring a custom base URL for the provider func (p *Provider) SetBaseURL(url string) { + useNativeAPI := useNativeAPIForBaseURLs(url, p.modelsURL) p.client.SetBaseURL(url) - p.modelsURL = url - p.modelsClientConf.BaseURL = url - p.useNativeAPI = useNativeAPIForBaseURLs(url, p.modelsURL) + p.useNativeAPI = useNativeAPI } // SetModelsURL allows configuring a custom models API base URL. diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index eb9de0177..d76631aaf 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -118,6 +118,65 @@ func TestSetBaseURLDisablesNativeRouting(t *testing.T) { } } +func TestSetBaseURLPreservesModelsURL(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + modelsHit := false + modelsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + modelsHit = true + if r.URL.Path != "/models" { + t.Errorf("models path = %q, want /models", r.URL.Path) + } + if got := r.Header.Get("x-goog-api-key"); got != "test-api-key" { + t.Errorf("x-goog-api-key = %q, want test-api-key", got) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "models": [{ + "name": "models/gemini-2.5-flash", + "displayName": "Gemini 2.5 Flash", + "supportedGenerationMethods": ["generateContent", "streamGenerateContent"], + "inputTokenLimit": 1048576, + "outputTokenLimit": 8192 + }] + }`)) + })) + defer modelsServer.Close() + + openAICompatHit := false + openAICompatServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + openAICompatHit = true + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error":{"message":"ListModels should use models URL"}}`)) + })) + defer openAICompatServer.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetModelsURL(modelsServer.URL) + provider.SetBaseURL(openAICompatServer.URL + "/v1beta/openai") + + if provider.modelsURL != modelsServer.URL { + t.Fatalf("modelsURL = %q, want %q", provider.modelsURL, modelsServer.URL) + } + if provider.modelsClientConf.BaseURL != modelsServer.URL { + t.Fatalf("modelsClientConf.BaseURL = %q, want %q", provider.modelsClientConf.BaseURL, modelsServer.URL) + } + + resp, err := provider.ListModels(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(resp.Data) != 1 || resp.Data[0].ID != "gemini-2.5-flash" { + t.Fatalf("models = %+v, want gemini-2.5-flash", resp.Data) + } + if !modelsHit { + t.Fatal("models server was not called") + } + if openAICompatHit { + t.Fatal("OpenAI-compatible server was called for ListModels") + } +} + func TestChatCompletion(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "false") From 841b0f8768b26083dead22adf2ad14d8a85641cf Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Mon, 4 May 2026 18:57:12 +0200 Subject: [PATCH 6/8] docs: group provider guides in navigation --- docs/docs.json | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/docs/docs.json b/docs/docs.json index fec382603..8b028cf8f 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -85,17 +85,23 @@ "icon": "compass", "pages": [ "guides/openclaw", - "guides/gemini", - "guides/oracle", - "guides/deepseek", - "guides/vllm", - "guides/multiple-ollama", "guides/claude-code", "guides/codex", "guides/opencode-and-other-agents", "guides/prometheus-metrics" ] }, + { + "tab": "Providers", + "icon": "plug", + "pages": [ + "guides/gemini", + "guides/oracle", + "guides/deepseek", + "guides/vllm", + "guides/multiple-ollama" + ] + }, { "tab": "API Reference", "icon": "braces", From 5e3cacb93317cf4380d2a0302e6f1750e84d1cd1 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Mon, 4 May 2026 19:15:21 +0200 Subject: [PATCH 7/8] fix(gemini): derive native base from configured url --- .env.template | 6 +- docs/guides/gemini.mdx | 17 +- internal/providers/gemini/gemini.go | 52 +++- internal/providers/gemini/gemini_test.go | 355 +++++++++++++++++++---- 4 files changed, 351 insertions(+), 79 deletions(-) diff --git a/.env.template b/.env.template index 40b24b9a1..ce66a2f60 100644 --- a/.env.template +++ b/.env.template @@ -269,8 +269,10 @@ # fetch remote image URLs or upload them through Gemini Files API yet. Set this # to false when you need OpenAI-compatible image_url pass-through behavior. # USE_GOOGLE_GEMINI_NATIVE_API=true -# OpenAI-compatible Gemini base URL used when native chat is disabled, and by -# Gemini embeddings/files/batches which still rely on that compatibility surface. +# Gemini base URL. The official defaults are: +# - native chat/models: https://generativelanguage.googleapis.com/v1beta +# - OpenAI-compatible API: https://generativelanguage.googleapis.com/v1beta/openai +# If this ends in /openai, GoModel derives the native base by stripping /openai. # GEMINI_BASE_URL=https://generativelanguage.googleapis.com/v1beta/openai # xAI (Grok) diff --git a/docs/guides/gemini.mdx b/docs/guides/gemini.mdx index 76e337ef3..85e184951 100644 --- a/docs/guides/gemini.mdx +++ b/docs/guides/gemini.mdx @@ -41,17 +41,24 @@ OpenAI-compatible `/chat/completions` endpoint: export USE_GOOGLE_GEMINI_NATIVE_API=false ``` -`GEMINI_BASE_URL` configures the OpenAI-compatible Gemini base URL. GoModel also -uses that compatibility surface for Gemini embeddings, files, and batches. +`GEMINI_BASE_URL` configures the Gemini base. GoModel keeps separate internal +bases for native Gemini and the OpenAI-compatible API: + +- native chat/models default: `https://generativelanguage.googleapis.com/v1beta` +- OpenAI-compatible default: `https://generativelanguage.googleapis.com/v1beta/openai` + +When `GEMINI_BASE_URL` ends in `/openai`, GoModel uses that value for the +OpenAI-compatible client and derives the native base by stripping `/openai`. +Gemini embeddings, files, and batches still use the OpenAI-compatible surface. ```bash export GEMINI_BASE_URL="https://generativelanguage.googleapis.com/v1beta/openai" ``` - If you configure a custom `GEMINI_BASE_URL`, GoModel disables native Gemini - chat routing for that provider and sends chat requests to the configured - OpenAI-compatible endpoint. + `USE_GOOGLE_GEMINI_NATIVE_API` decides whether chat and Responses API calls + use native Gemini or the OpenAI-compatible API. `GEMINI_BASE_URL` only + configures the upstream base URLs. ## Image URL behavior diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index eed5cbf98..0603c1e59 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -50,13 +50,12 @@ type Provider struct { // New creates a new Gemini provider. func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { - baseURL := providers.ResolveBaseURL(providerCfg.BaseURL, defaultOpenAICompatibleBaseURL) - modelsURL := defaultModelsBaseURL + baseURL, modelsURL := geminiBaseURLs(providerCfg.BaseURL) p := &Provider{ httpClient: nil, apiKey: providerCfg.APIKey, hooks: opts.Hooks, - useNativeAPI: useNativeAPIForBaseURLs(baseURL, modelsURL), + useNativeAPI: useNativeAPIFromEnv(), modelsURL: modelsURL, modelsClientConf: llmclient.Config{ ProviderName: "gemini", @@ -84,13 +83,12 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H if httpClient == nil { httpClient = http.DefaultClient } - baseURL := defaultOpenAICompatibleBaseURL - modelsURL := defaultModelsBaseURL + baseURL, modelsURL := geminiBaseURLs("") p := &Provider{ httpClient: httpClient, apiKey: apiKey, hooks: hooks, - useNativeAPI: useNativeAPIForBaseURLs(baseURL, modelsURL), + useNativeAPI: useNativeAPIFromEnv(), modelsURL: modelsURL, } modelsCfg := llmclient.DefaultConfig("gemini", modelsURL) @@ -105,9 +103,14 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H // SetBaseURL allows configuring a custom base URL for the provider func (p *Provider) SetBaseURL(url string) { - useNativeAPI := useNativeAPIForBaseURLs(url, p.modelsURL) - p.client.SetBaseURL(url) - p.useNativeAPI = useNativeAPI + baseURL, modelsURL := geminiBaseURLs(url) + p.client.SetBaseURL(baseURL) + p.modelsURL = modelsURL + p.modelsClientConf.BaseURL = modelsURL + if p.nativeClient != nil { + p.nativeClient.SetBaseURL(modelsURL) + } + p.useNativeAPI = useNativeAPIFromEnv() } // SetModelsURL allows configuring a custom models API base URL. @@ -152,10 +155,33 @@ func useNativeAPIFromEnv() bool { } } -func useNativeAPIForBaseURLs(baseURL, modelsURL string) bool { - return useNativeAPIFromEnv() && - baseURL == defaultOpenAICompatibleBaseURL && - modelsURL == defaultModelsBaseURL +func geminiBaseURLs(configuredBaseURL string) (openAICompatibleBaseURL, nativeBaseURL string) { + baseURL := strings.TrimRight(strings.TrimSpace(configuredBaseURL), "/") + if baseURL == "" { + return defaultOpenAICompatibleBaseURL, defaultModelsBaseURL + } + if baseURL == defaultOpenAICompatibleBaseURL { + return defaultOpenAICompatibleBaseURL, defaultModelsBaseURL + } + if baseURL == defaultModelsBaseURL { + return defaultOpenAICompatibleBaseURL, defaultModelsBaseURL + } + if nativeBaseURL, ok := nativeBaseURLFromOpenAICompatibleBaseURL(baseURL); ok { + return baseURL, nativeBaseURL + } + return baseURL, baseURL +} + +func nativeBaseURLFromOpenAICompatibleBaseURL(baseURL string) (string, bool) { + const suffix = "/openai" + if !strings.HasSuffix(baseURL, suffix) { + return "", false + } + nativeBaseURL := strings.TrimRight(strings.TrimSuffix(baseURL, suffix), "/") + if nativeBaseURL == "" { + return "", false + } + return nativeBaseURL, true } // adaptChatRequest rewrites a ChatRequest for Gemini's OpenAI-compatible endpoint. diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index d76631aaf..6591a39e0 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -39,7 +39,58 @@ func TestNew_ReturnsProvider(t *testing.T) { } } -func TestNew_CustomBaseURLDisablesNativeMode(t *testing.T) { +func TestGeminiBaseURLs(t *testing.T) { + tests := []struct { + name string + configured string + wantCompat string + wantNative string + }{ + { + name: "empty uses official defaults", + wantCompat: defaultOpenAICompatibleBaseURL, + wantNative: defaultModelsBaseURL, + }, + { + name: "official OpenAI-compatible default derives native default", + configured: defaultOpenAICompatibleBaseURL, + wantCompat: defaultOpenAICompatibleBaseURL, + wantNative: defaultModelsBaseURL, + }, + { + name: "official native default keeps OpenAI-compatible default", + configured: defaultModelsBaseURL, + wantCompat: defaultOpenAICompatibleBaseURL, + wantNative: defaultModelsBaseURL, + }, + { + name: "custom OpenAI-compatible URL derives native sibling", + configured: "https://proxy.example.com/v1beta/openai/", + wantCompat: "https://proxy.example.com/v1beta/openai", + wantNative: "https://proxy.example.com/v1beta", + }, + { + name: "custom URL without OpenAI suffix is used for both clients", + configured: "https://proxy.example.com/gemini", + wantCompat: "https://proxy.example.com/gemini", + wantNative: "https://proxy.example.com/gemini", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotCompat, gotNative := geminiBaseURLs(tt.configured) + if gotCompat != tt.wantCompat { + t.Fatalf("OpenAI-compatible base = %q, want %q", gotCompat, tt.wantCompat) + } + if gotNative != tt.wantNative { + t.Fatalf("native base = %q, want %q", gotNative, tt.wantNative) + } + }) + } +} + +func TestNew_CustomBaseURLDerivesNativeBaseURL(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "true") provider := New(providers.ProviderConfig{ @@ -51,52 +102,51 @@ func TestNew_CustomBaseURLDisablesNativeMode(t *testing.T) { if !ok { t.Fatalf("provider type = %T, want *Provider", provider) } - if geminiProvider.useNativeAPI { - t.Fatal("useNativeAPI = true, want false for custom OpenAI-compatible base URL") + if !geminiProvider.useNativeAPI { + t.Fatal("useNativeAPI = false, want true for custom OpenAI-compatible base URL") + } + if got := geminiProvider.client.BaseURL(); got != "https://proxy.example.com/v1beta/openai" { + t.Fatalf("client.BaseURL() = %q, want OpenAI-compatible base URL", got) + } + if geminiProvider.modelsURL != "https://proxy.example.com/v1beta" { + t.Fatalf("modelsURL = %q, want derived native base URL", geminiProvider.modelsURL) + } + if geminiProvider.modelsClientConf.BaseURL != "https://proxy.example.com/v1beta" { + t.Fatalf("modelsClientConf.BaseURL = %q, want derived native base URL", geminiProvider.modelsClientConf.BaseURL) } } -func TestSetBaseURLDisablesNativeRouting(t *testing.T) { +func TestSetBaseURLDerivesNativeRouting(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "true") nativeHit := false - nativeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.Contains(r.URL.Path, "/openai/") { + t.Fatalf("native routing used OpenAI-compatible path %q", r.URL.Path) + } nativeHit = true - w.WriteHeader(http.StatusInternalServerError) - _, _ = w.Write([]byte(`{"error":{"message":"native client should not be used after SetBaseURL"}}`)) - })) - defer nativeServer.Close() - - openAICompatHit := false - openAICompatServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - openAICompatHit = true - if r.URL.Path != "/v1beta/openai/chat/completions" { - t.Errorf("OpenAI-compatible path = %q, want /v1beta/openai/chat/completions", r.URL.Path) + if r.URL.Path != "/v1beta/models/gemini-2.5-flash:generateContent" { + t.Errorf("native path = %q, want /v1beta/models/gemini-2.5-flash:generateContent", r.URL.Path) } - if got := r.Header.Get("Authorization"); got != "Bearer test-api-key" { - t.Errorf("Authorization = %q, want bearer API key", got) + if got := r.Header.Get("x-goog-api-key"); got != "test-api-key" { + t.Errorf("x-goog-api-key = %q, want test-api-key", got) } - if got := r.Header.Get("x-goog-api-key"); got != "" { - t.Errorf("x-goog-api-key = %q, want empty for OpenAI-compatible API", got) + if got := r.Header.Get("Authorization"); got != "" { + t.Errorf("Authorization = %q, want empty for native Gemini API", got) } w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{ - "id": "gemini-openai-compatible-baseurl", - "object": "chat.completion", - "created": 1677652288, - "model": "gemini-2.5-flash", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "ok"}, - "finish_reason": "stop" + "responseId": "gemini-native-baseurl", + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP" }] }`)) })) - defer openAICompatServer.Close() + defer server.Close() provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) - provider.SetModelsURL(nativeServer.URL) - provider.SetBaseURL(openAICompatServer.URL + "/v1beta/openai") + provider.SetBaseURL(server.URL + "/v1beta/openai") resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ Model: "gemini-2.5-flash", @@ -107,25 +157,25 @@ func TestSetBaseURLDisablesNativeRouting(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } - if resp == nil || resp.ID != "gemini-openai-compatible-baseurl" { - t.Fatalf("response = %+v, want OpenAI-compatible response", resp) - } - if nativeHit { - t.Fatal("native server was called after SetBaseURL") + if resp == nil || resp.ID != "gemini-native-baseurl" { + t.Fatalf("response = %+v, want native response", resp) } - if !openAICompatHit { - t.Fatal("OpenAI-compatible server was not called") + if !nativeHit { + t.Fatal("native server was not called") } } -func TestSetBaseURLPreservesModelsURL(t *testing.T) { +func TestSetBaseURLDerivesModelsURL(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "true") modelsHit := false - modelsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { modelsHit = true - if r.URL.Path != "/models" { - t.Errorf("models path = %q, want /models", r.URL.Path) + if strings.Contains(r.URL.Path, "/openai/") { + t.Fatalf("models request used OpenAI-compatible path %q", r.URL.Path) + } + if r.URL.Path != "/v1beta/models" { + t.Errorf("models path = %q, want /v1beta/models", r.URL.Path) } if got := r.Header.Get("x-goog-api-key"); got != "test-api-key" { t.Errorf("x-goog-api-key = %q, want test-api-key", got) @@ -141,25 +191,19 @@ func TestSetBaseURLPreservesModelsURL(t *testing.T) { }] }`)) })) - defer modelsServer.Close() - - openAICompatHit := false - openAICompatServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - openAICompatHit = true - w.WriteHeader(http.StatusInternalServerError) - _, _ = w.Write([]byte(`{"error":{"message":"ListModels should use models URL"}}`)) - })) - defer openAICompatServer.Close() + defer server.Close() provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) - provider.SetModelsURL(modelsServer.URL) - provider.SetBaseURL(openAICompatServer.URL + "/v1beta/openai") + provider.SetBaseURL(server.URL + "/v1beta/openai") - if provider.modelsURL != modelsServer.URL { - t.Fatalf("modelsURL = %q, want %q", provider.modelsURL, modelsServer.URL) + if provider.client.BaseURL() != server.URL+"/v1beta/openai" { + t.Fatalf("client.BaseURL() = %q, want OpenAI-compatible base URL", provider.client.BaseURL()) } - if provider.modelsClientConf.BaseURL != modelsServer.URL { - t.Fatalf("modelsClientConf.BaseURL = %q, want %q", provider.modelsClientConf.BaseURL, modelsServer.URL) + if provider.modelsURL != server.URL+"/v1beta" { + t.Fatalf("modelsURL = %q, want derived native base URL", provider.modelsURL) + } + if provider.modelsClientConf.BaseURL != server.URL+"/v1beta" { + t.Fatalf("modelsClientConf.BaseURL = %q, want derived native base URL", provider.modelsClientConf.BaseURL) } resp, err := provider.ListModels(context.Background()) @@ -172,9 +216,6 @@ func TestSetBaseURLPreservesModelsURL(t *testing.T) { if !modelsHit { t.Fatal("models server was not called") } - if openAICompatHit { - t.Fatal("OpenAI-compatible server was called for ListModels") - } } func TestChatCompletion(t *testing.T) { @@ -1128,6 +1169,109 @@ func TestResponses(t *testing.T) { } } +func TestResponses_Native(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("Method = %q, want %q", r.Method, http.MethodPost) + } + if r.URL.Path != "/models/gemini-2.5-flash:generateContent" { + t.Errorf("Path = %q, want native generateContent endpoint", r.URL.Path) + } + if got := r.Header.Get("x-goog-api-key"); got != "test-api-key" { + t.Errorf("x-goog-api-key = %q, want test-api-key", got) + } + if got := r.Header.Get("Authorization"); got != "" { + t.Errorf("Authorization = %q, want empty for native Gemini API", got) + } + + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + if _, ok := payload["messages"]; ok { + t.Fatal("native request should not contain OpenAI messages") + } + systemInstruction, ok := payload["system_instruction"].(map[string]any) + if !ok { + t.Fatalf("system_instruction = %#v, want object", payload["system_instruction"]) + } + systemParts, ok := systemInstruction["parts"].([]any) + if !ok || len(systemParts) != 1 || systemParts[0].(map[string]any)["text"] != "Be concise." { + t.Fatalf("system_instruction.parts = %#v, want instruction text", systemInstruction["parts"]) + } + contents, ok := payload["contents"].([]any) + if !ok || len(contents) != 1 { + t.Fatalf("contents = %#v, want one native content", payload["contents"]) + } + firstContent := contents[0].(map[string]any) + if firstContent["role"] != "user" { + t.Fatalf("contents[0].role = %#v, want user", firstContent["role"]) + } + parts := firstContent["parts"].([]any) + if len(parts) != 1 || parts[0].(map[string]any)["text"] != "Hello" { + t.Fatalf("contents[0].parts = %#v, want user text", firstContent["parts"]) + } + generationConfig, ok := payload["generationConfig"].(map[string]any) + if !ok { + t.Fatalf("generationConfig = %#v, want object", payload["generationConfig"]) + } + if got := generationConfig["maxOutputTokens"]; got != float64(64) { + t.Fatalf("maxOutputTokens = %#v, want 64", got) + } + + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "gemini-native-response", + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "Native response"}]}, + "finishReason": "STOP" + }], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 3, + "totalTokenCount": 8 + } + }`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetModelsURL(server.URL) + + maxOutputTokens := 64 + resp, err := provider.Responses(context.Background(), &core.ResponsesRequest{ + Model: "gemini-2.5-flash", + Instructions: "Be concise.", + Input: "Hello", + MaxOutputTokens: &maxOutputTokens, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if resp.ID != "gemini-native-response" { + t.Fatalf("ID = %q, want gemini-native-response", resp.ID) + } + if resp.Object != "response" { + t.Fatalf("Object = %q, want response", resp.Object) + } + if resp.Model != "gemini-2.5-flash" { + t.Fatalf("Model = %q, want gemini-2.5-flash", resp.Model) + } + if resp.Provider != "gemini" { + t.Fatalf("Provider = %q, want gemini", resp.Provider) + } + if len(resp.Output) != 1 || len(resp.Output[0].Content) != 1 || resp.Output[0].Content[0].Text != "Native response" { + t.Fatalf("Output = %+v, want native response text", resp.Output) + } +} + func TestStreamResponses(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "false") @@ -1170,3 +1314,96 @@ data: [DONE] t.Error("response should end with [DONE]") } } + +func TestStreamResponses_Native(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("Method = %q, want %q", r.Method, http.MethodPost) + } + if r.URL.Path != "/models/gemini-2.5-flash:streamGenerateContent" { + t.Errorf("Path = %q, want native streamGenerateContent endpoint", r.URL.Path) + } + if got := r.URL.Query().Get("alt"); got != "sse" { + t.Errorf("alt = %q, want sse", got) + } + if got := r.Header.Get("x-goog-api-key"); got != "test-api-key" { + t.Errorf("x-goog-api-key = %q, want test-api-key", got) + } + if got := r.Header.Get("Authorization"); got != "" { + t.Errorf("Authorization = %q, want empty for native Gemini API", got) + } + + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + if _, ok := payload["stream"]; ok { + t.Fatal("native stream request should not contain OpenAI stream flag") + } + systemInstruction, ok := payload["system_instruction"].(map[string]any) + if !ok { + t.Fatalf("system_instruction = %#v, want object", payload["system_instruction"]) + } + systemParts := systemInstruction["parts"].([]any) + if len(systemParts) != 1 || systemParts[0].(map[string]any)["text"] != "Be concise." { + t.Fatalf("system_instruction.parts = %#v, want instruction text", systemInstruction["parts"]) + } + contents, ok := payload["contents"].([]any) + if !ok || len(contents) != 1 { + t.Fatalf("contents = %#v, want one native content", payload["contents"]) + } + parts := contents[0].(map[string]any)["parts"].([]any) + if len(parts) != 1 || parts[0].(map[string]any)["text"] != "Hello" { + t.Fatalf("contents[0].parts = %#v, want user text", contents[0].(map[string]any)["parts"]) + } + + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`data: {"responseId":"gemini-native-stream-response","candidates":[{"content":{"role":"model","parts":[{"text":"Hello"}]}}]} + +data: {"responseId":"gemini-native-stream-response","candidates":[{"content":{"role":"model","parts":[{"text":"!"}]},"finishReason":"STOP"}]} + +`)) + })) + defer server.Close() + + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) + provider.SetModelsURL(server.URL) + + body, err := provider.StreamResponses(context.Background(), &core.ResponsesRequest{ + Model: "gemini-2.5-flash", + Instructions: "Be concise.", + Input: "Hello", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if body == nil { + t.Fatal("body should not be nil") + } + defer func() { _ = body.Close() }() + + raw, err := io.ReadAll(body) + if err != nil { + t.Fatalf("failed to read response stream: %v", err) + } + stream := string(raw) + if !strings.Contains(stream, "response.created") { + t.Fatalf("stream = %q, want response.created event", stream) + } + if !strings.Contains(stream, "response.output_text.delta") { + t.Fatalf("stream = %q, want response.output_text.delta event", stream) + } + if !strings.Contains(stream, `"delta":"Hello"`) || !strings.Contains(stream, `"delta":"!"`) { + t.Fatalf("stream = %q, want normalized text deltas", stream) + } + if !strings.Contains(stream, "data: [DONE]") { + t.Fatalf("stream = %q, want [DONE]", stream) + } +} From 32261b557c89d5661af68cf20a173fb754bc1eee Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Mon, 4 May 2026 20:00:16 +0200 Subject: [PATCH 8/8] fix(gemini): emit one native stream usage chunk --- internal/providers/gemini/gemini_test.go | 2 +- internal/providers/gemini/native_stream.go | 36 +++++++++++++--------- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 6591a39e0..b8f170675 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -820,7 +820,7 @@ func TestStreamChatCompletion_UsesNativeStreamByDefault(t *testing.T) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte(`data: {"responseId":"gemini-stream-123","candidates":[{"content":{"role":"model","parts":[{"text":"Hello"}]}}]} + _, _ = w.Write([]byte(`data: {"responseId":"gemini-stream-123","candidates":[{"content":{"role":"model","parts":[{"text":"Hello"}]}}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":1,"totalTokenCount":5}} data: {"responseId":"gemini-stream-123","candidates":[{"content":{"role":"model","parts":[{"text":"!"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}} diff --git a/internal/providers/gemini/native_stream.go b/internal/providers/gemini/native_stream.go index f3d590153..f0d88ca73 100644 --- a/internal/providers/gemini/native_stream.go +++ b/internal/providers/gemini/native_stream.go @@ -71,6 +71,10 @@ func convertGeminiNativeStream(body io.ReadCloser, out *io.PipeWriter, model str _ = out.CloseWithError(err) return } + if err := state.writeFinalUsage(out); err != nil { + _ = out.CloseWithError(err) + return + } _, _ = io.WriteString(out, "data: [DONE]\n\n") _ = out.Close() } @@ -82,6 +86,7 @@ type geminiStreamState struct { responseID string choices map[int]*geminiChoiceStreamState stopped bool + latestUsage map[string]any } type geminiChoiceStreamState struct { @@ -130,24 +135,27 @@ func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { } if s.includeUsage { - if usage := geminiUsageMap(event.UsageMetadata); usage != nil { - chunk := map[string]any{ - "id": s.responseID, - "object": "chat.completion.chunk", - "created": s.created, - "model": s.model, - "provider": "gemini", - "choices": []map[string]any{}, - "usage": usage, - } - if err := writeOpenAIStreamChunk(out, chunk); err != nil { - return err - } - } + s.latestUsage = geminiUsageMap(event.UsageMetadata) } return nil } +func (s *geminiStreamState) writeFinalUsage(out io.Writer) error { + if !s.includeUsage || s.latestUsage == nil || s.stopped { + return nil + } + chunk := map[string]any{ + "id": s.responseID, + "object": "chat.completion.chunk", + "created": s.created, + "model": s.model, + "provider": "gemini", + "choices": []map[string]any{}, + "usage": s.latestUsage, + } + return writeOpenAIStreamChunk(out, chunk) +} + func (s *geminiStreamState) chatChunkChoice(candidate geminiCandidate, fallbackIndex int) (map[string]any, bool) { index := streamChoiceIndex(candidate, fallbackIndex) state := s.choiceState(index)