diff --git a/provider_gemini.go b/provider_gemini.go new file mode 100644 index 0000000..f00bfa9 --- /dev/null +++ b/provider_gemini.go @@ -0,0 +1,271 @@ +package traefikllmgateway + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" +) + +// geminiAdapter is the providerAdapter for provider type "gemini". +type geminiAdapter struct { + client *http.Client + adapterName string + baseURL string + apiKey string +} + +// newGeminiAdapter returns a geminiAdapter for provider name, with base as +// its already-defaulted, trailing-slash-trimmed base URL and apiKey as +// already resolved by resolveSecret. Like anthropic and unlike openai-type +// providers, Gemini's API requires an API key on every request — a +// resolved-empty apiKey is a constructor error here, mirroring ruling (c). +func newGeminiAdapter(name, base, apiKey string) (*geminiAdapter, error) { + if apiKey == "" { + return nil, fmt.Errorf("llmgateway: provider %q: gemini requires an API key", name) + } + return &geminiAdapter{ + adapterName: name, + baseURL: base, + apiKey: apiKey, + client: newAdapterHTTPClient(), + }, nil +} + +// name implements providerAdapter. +func (a *geminiAdapter) name() string { return a.adapterName } + +// typeName implements providerAdapter. +func (a *geminiAdapter) typeName() string { return providerTypeGemini } + +// base implements providerAdapter. +func (a *geminiAdapter) base() string { return a.baseURL } + +// injectAuth implements providerAdapter: Gemini's x-goog-api-key header +// scheme. Never a no-op — the constructor already rejected an empty apiKey. +func (a *geminiAdapter) injectAuth(r *http.Request) { + r.Header.Set("x-goog-api-key", a.apiKey) +} + +// requestHeaders returns the headers for one upstream request: a +// Content-Type when the request carries a JSON body, plus whatever +// injectAuth adds. +func (a *geminiAdapter) requestHeaders(hasBody bool) http.Header { + h := http.Header{} + if hasBody { + h.Set("Content-Type", "application/json") + } + a.injectAuth(&http.Request{Header: h}) + return h +} + +// geminiModelFromRequest reads req["model"] — already rewritten to the +// provider's own model id by the caller, per the providerAdapter interface +// doc — and strips a leading "models/" prefix if present, so a client (or +// an upstream config) that already qualified the id doesn't produce a +// doubled "models/models/..." URL. +func geminiModelFromRequest(req map[string]any) string { + raw, _ := req["model"].(string) + return strings.TrimPrefix(raw, "models/") +} + +// chatCompletion implements providerAdapter: translates req to a Gemini +// generateContent/streamGenerateContent request via geminiRequestFromOpenAI, +// then forwards the (possibly streamed) response back translated to +// OpenAI's shape. Unlike anthropic and openai-type, the model id is never +// part of the JSON body — Gemini puts it in the URL path — so it is read +// once here and threaded through to both URL construction and the response +// translation's "model" field. A translation failure — a *translateError +// from geminiRequestFromOpenAI — is returned as-is, before any upstream +// request is made. +func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + model := geminiModelFromRequest(req) + streaming, _ := req["stream"].(bool) + + body, err := geminiRequestFromOpenAI(req) + if err != nil { + return usage{}, err + } + + hdr := a.requestHeaders(true) + url := a.baseURL + geminiAPIPrefix + model + ":generateContent" + if streaming { + // Ask the upstream for an SSE response. The Content-Type check + // below still decides which forwarding path runs. + hdr.Set("Accept", "text/event-stream") + url = a.baseURL + geminiAPIPrefix + model + ":streamGenerateContent?alt=sse" + } + + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, url, hdr, body) + if err != nil { + return usage{}, err + } + defer resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return usage{}, newProviderHTTPError(resp) + } + + if streaming && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") { + return a.forwardStream(w, resp.Body, model) + } + // Either a non-streaming request, or a streaming one whose upstream + // ignored stream=true and answered with a normal JSON body. + return a.forwardJSON(w, resp, model) +} + +// forwardJSON reads resp's full body (capped at maxResponseBytes), +// translates it to an OpenAI chat.completion response via +// openAIResponseFromGemini, and writes the translated JSON to w. +// gatewayModel becomes the translated response's "model" field, and the +// response's "created" timestamp is generated here — Gemini's response +// carries neither. +func (a *geminiAdapter) forwardJSON(w http.ResponseWriter, resp *http.Response, gatewayModel string) (usage, error) { + raw, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) + if err != nil { + return usage{}, fmt.Errorf("%w: read response body: %w", errUpstream, err) + } + + out, u, err := openAIResponseFromGemini(raw, gatewayModel, time.Now().Unix()) + if err != nil { + return usage{}, err + } + + b, err := json.Marshal(out) + if err != nil { + return usage{}, fmt.Errorf("%w: marshal translated response: %w", errUpstream, err) + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(b) + return u, nil +} + +// forwardStream reads body as an upstream Gemini text/event-stream, +// translating each event via a geminiStreamState and forwarding every +// resulting OpenAI-shaped chunk to w through a fresh sseWriter, flushing +// per event. gatewayModel and the stream's start time become every chunk's +// "model"/"created" fields. The terminal "[DONE]" is written once readSSE +// returns without error, mirroring ruling (b) — Gemini's stream has no +// distinct terminal event of its own to key off of; a clean EOF is what +// signals the stream ended. On a translate error (a malformed chunk, or a +// genuine stream read error), forwardStream returns without writing +// "[DONE]". +func (a *geminiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gatewayModel string) (usage, error) { + st := newGeminiStreamState(gatewayModel, time.Now().Unix()) + sw := newSSEWriter(w) + + err := readSSE(body, func(ev sseEvent) error { + chunks, terr := st.translate(ev) + for _, c := range chunks { + if werr := sw.writeData(c); werr != nil { + return werr + } + } + return terr + }) + if err != nil { + return st.usage(), fmt.Errorf("%w: stream: %w", errUpstream, err) + } + sw.writeDone() + return st.usage(), nil +} + +// embeddings implements providerAdapter: translates req's OpenAI "input" +// (a string or an array of strings) to a Gemini :embedContent or +// :batchEmbedContents request, per the embeddings mapping table, and +// translates the response back to an OpenAI embeddings list response. +// Embeddings never stream. +func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + model := geminiModelFromRequest(req) + + texts, single, err := geminiEmbeddingInputTexts(req["input"]) + if err != nil { + return usage{}, err + } + + var body map[string]any + var url string + if single { + body = geminiEmbedContentRequest(texts[0]) + url = a.baseURL + geminiAPIPrefix + model + ":embedContent" + } else { + body = geminiBatchEmbedContentsRequest(model, texts) + url = a.baseURL + geminiAPIPrefix + model + ":batchEmbedContents" + } + + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, url, a.requestHeaders(true), body) + if err != nil { + return usage{}, err + } + defer resp.Body.Close() //nolint:errcheck + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return usage{}, newProviderHTTPError(resp) + } + + raw, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) + if err != nil { + return usage{}, fmt.Errorf("%w: read response body: %w", errUpstream, err) + } + + out, u, err := openAIEmbeddingResponseFromGemini(raw, single, model, estimatedEmbeddingTokens(texts)) + if err != nil { + return usage{}, err + } + + b, err := json.Marshal(out) + if err != nil { + return usage{}, fmt.Errorf("%w: marshal translated embeddings response: %w", errUpstream, err) + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(b) + return u, nil +} + +// geminiModelsListPayload is the "models": [{"name": "models/..."}, ...] +// shape Gemini's models.list endpoint returns (its "name" field is the +// full "models/{model}" resource name, unlike OpenAI's bare "id"). +type geminiModelsListPayload struct { + Models []struct { + Name string `json:"name"` + } `json:"models"` +} + +// listModels implements providerAdapter: GET {base}/v1beta/models, +// stripping the "models/" prefix off each entry's "name" to return bare +// model ids, matching the shape openai-type and anthropic-type adapters +// both return. +func (a *geminiAdapter) listModels(ctx context.Context) ([]string, error) { + resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+"/v1beta/models", a.requestHeaders(false), nil) + if err != nil { + return nil, err + } + defer resp.Body.Close() //nolint:errcheck + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, newProviderHTTPError(resp) + } + + body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) + if err != nil { + return nil, fmt.Errorf("%w: read response body: %w", errUpstream, err) + } + + var payload geminiModelsListPayload + if err := json.Unmarshal(body, &payload); err != nil { + return nil, fmt.Errorf("%w: decode models response: %w", errUpstream, err) + } + + ids := make([]string, len(payload.Models)) + for i, m := range payload.Models { + ids[i] = strings.TrimPrefix(m.Name, "models/") + } + return ids, nil +} diff --git a/provider_gemini_test.go b/provider_gemini_test.go new file mode 100644 index 0000000..599ba6f --- /dev/null +++ b/provider_gemini_test.go @@ -0,0 +1,495 @@ +package traefikllmgateway + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// TestNewGeminiAdapter_KeylessIsAConstructorError proves an empty resolved +// API key is a constructor error for gemini-type providers, mirroring +// ruling (c) and anthropic's equivalent test. +func TestNewGeminiAdapter_KeylessIsAConstructorError(t *testing.T) { + if _, err := newGeminiAdapter("p1", "https://generativelanguage.googleapis.com", ""); err == nil { + t.Fatalf("newGeminiAdapter: want error for empty API key, got nil") + } +} + +// TestGeminiAdapter_InjectAuth proves injectAuth sets the x-goog-api-key +// header unconditionally. +func TestGeminiAdapter_InjectAuth(t *testing.T) { + a, err := newGeminiAdapter("p1", "https://generativelanguage.googleapis.com", "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + req := httptest.NewRequest(http.MethodPost, "/", nil) + a.injectAuth(req) + + if got := req.Header.Get("x-goog-api-key"); got != "AIza-test" { //nolint:gosec // test-only literal, not a real credential + t.Errorf("x-goog-api-key = %q, want %q", got, "AIza-test") + } +} + +// TestGeminiAdapter_ChatCompletion_NonStreaming proves the adapter builds +// the {base}/v1beta/models/{model}:generateContent URL from req["model"], +// sends the x-goog-api-key header, translates the OpenAI request to +// Gemini's wire format, and translates the Gemini response back to an +// OpenAI-shaped chat.completion body on the wire, with usage extracted +// from usageMetadata. +func TestGeminiAdapter_ChatCompletion_NonStreaming(t *testing.T) { + const geminiResp = `{"candidates":[{"content":{"role":"model","parts":[{"text":"hi there"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":4},"responseId":"resp_01ABC"}` + + var gotBody []byte + var gotAPIKey, gotPath, gotMethod string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotMethod = r.Method + gotBody, _ = io.ReadAll(r.Body) + gotAPIKey = r.Header.Get("x-goog-api-key") + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(geminiResp)) + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + rec := httptest.NewRecorder() + + req := map[string]any{ + "model": "gemini-2.5-pro", + "messages": []any{map[string]any{"role": "user", "content": "hello"}}, + } + u, err := a.chatCompletion(context.Background(), rec, req) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + + if gotPath != "/v1beta/models/gemini-2.5-pro:generateContent" { + t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:generateContent", gotPath) + } + if gotMethod != http.MethodPost { + t.Errorf("method = %q, want POST", gotMethod) + } + if u.prompt != 10 || u.completion != 4 { + t.Errorf("usage = %+v, want {prompt:10 completion:4}", u) + } + if gotAPIKey != "AIza-test" { //nolint:gosec // test-only literal, not a real credential + t.Errorf("x-goog-api-key = %q, want %q", gotAPIKey, "AIza-test") + } + + sent := decodeJSONBody(t, gotBody) + if _, hasModel := sent["model"]; hasModel { + t.Errorf("upstream body carries a \"model\" field = %v, want none (Gemini's model is URL-only)", sent["model"]) + } + contents, ok := sent["contents"].([]any) + if !ok || len(contents) != 1 { + t.Fatalf("upstream body contents = %v, want a single-element array", sent["contents"]) + } + + if got := rec.Header().Get("Content-Type"); got != "application/json" { + t.Errorf("Content-Type = %q, want application/json", got) + } + var out map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String()) + } + if out["object"] != "chat.completion" { + t.Errorf("object = %v, want %q", out["object"], "chat.completion") + } + if out["id"] != "chatcmpl-resp_01ABC" { + t.Errorf("id = %v, want %q", out["id"], "chatcmpl-resp_01ABC") + } + choices, ok := out["choices"].([]any) + if !ok || len(choices) != 1 { + t.Fatalf("choices = %v, want a single-element array", out["choices"]) + } + choice := choices[0].(map[string]any) + message := choice["message"].(map[string]any) + if message["content"] != "hi there" { + t.Errorf("message.content = %v, want %q", message["content"], "hi there") + } + if choice["finish_reason"] != "stop" { + t.Errorf("finish_reason = %v, want %q", choice["finish_reason"], "stop") + } +} + +// TestGeminiAdapter_ChatCompletion_ModelsPrefixStripped proves a +// req["model"] that already carries a "models/" prefix is not doubled into +// the URL. +func TestGeminiAdapter_ChatCompletion_ModelsPrefixStripped(t *testing.T) { + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`)) + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + rec := httptest.NewRecorder() + + _, err = a.chatCompletion(context.Background(), rec, map[string]any{ + "model": "models/gemini-2.5-pro", + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + }) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + if gotPath != "/v1beta/models/gemini-2.5-pro:generateContent" { + t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:generateContent (no doubled models/ prefix)", gotPath) + } +} + +// geminiStreamFrames are the raw SSE "data:" lines a fake streaming +// upstream emits for the tests below: a role/text chunk, a text-only +// continuation, then a finish chunk carrying usageMetadata. Gemini's SSE +// stream has no "event:" field to set, unlike Anthropic's. +var geminiStreamFrames = []string{ + `{"candidates":[{"content":{"role":"model","parts":[{"text":"Hi!"}]}}],"responseId":"resp_STREAM"}`, + `{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":6,"candidatesTokenCount":3}}`, +} + +// TestGeminiAdapter_ChatCompletion_Streaming proves the adapter requests +// the :streamGenerateContent?alt=sse URL, forwards a streaming Gemini +// response as OpenAI-shaped SSE chunks on the wire, flushing per event, and +// terminates with "[DONE]" once the upstream stream ends — mirroring +// ruling (b): there is no distinct terminal event of Gemini's own to key +// off of. +func TestGeminiAdapter_ChatCompletion_Streaming(t *testing.T) { + var gotPath, gotQuery, gotAccept string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotQuery = r.URL.RawQuery + gotAccept = r.Header.Get("Accept") + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + f, ok := w.(http.Flusher) + if !ok { + t.Errorf("httptest ResponseWriter does not implement http.Flusher") + return + } + for _, data := range geminiStreamFrames { + fmt.Fprintf(w, "data: %s\n\n", data) + f.Flush() + } + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + rec := newRecordingResponseWriter() + + req := map[string]any{ + "model": "gemini-2.5-pro", + "stream": true, + "messages": []any{map[string]any{"role": "user", "content": "hello"}}, + } + u, err := a.chatCompletion(context.Background(), rec, req) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + + if gotPath != "/v1beta/models/gemini-2.5-pro:streamGenerateContent" { + t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:streamGenerateContent", gotPath) + } + if gotQuery != "alt=sse" { + t.Errorf("query = %q, want %q", gotQuery, "alt=sse") + } + if gotAccept != "text/event-stream" { + t.Errorf("Accept = %q, want text/event-stream", gotAccept) + } + if u.prompt != 6 || u.completion != 3 { + t.Errorf("usage = %+v, want {prompt:6 completion:3}", u) + } + + body := rec.body.String() + if got := rec.header.Get("Content-Type"); got != "text/event-stream" { + t.Errorf("Content-Type = %q, want text/event-stream", got) + } + + wantSubstrings := []string{ + `"delta":{"role":"assistant"}`, + `"delta":{"content":"Hi!"}`, + `"finish_reason":"stop"`, + "data: [DONE]\n\n", + } + for _, want := range wantSubstrings { + if !strings.Contains(body, want) { + t.Errorf("streamed body missing %q; got:\n%s", want, body) + } + } + // 3 chunks (role, content delta, finish) plus one more flush for the + // terminal [DONE]: 4 flushes, proving the adapter flushes per event + // rather than buffering the whole stream. + const wantFlushes = 4 + if rec.flushed < wantFlushes { + t.Errorf("flushed = %d, want >= %d", rec.flushed, wantFlushes) + } +} + +// TestGeminiAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType +// proves the streaming-vs-JSON-fallback gate matches an upstream's +// Content-Type case-insensitively, per lesson (g) — tested with +// "Text/Event-Stream; charset=utf-8". +func TestGeminiAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "Text/Event-Stream; charset=utf-8") + w.WriteHeader(http.StatusOK) + f, ok := w.(http.Flusher) + if !ok { + t.Errorf("httptest ResponseWriter does not implement http.Flusher") + return + } + for _, data := range geminiStreamFrames { + fmt.Fprintf(w, "data: %s\n\n", data) + f.Flush() + } + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + rec := newRecordingResponseWriter() + + req := map[string]any{ + "model": "gemini-2.5-pro", + "stream": true, + "messages": []any{map[string]any{"role": "user", "content": "hello"}}, + } + u, err := a.chatCompletion(context.Background(), rec, req) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + if u.prompt != 6 || u.completion != 3 { + t.Errorf("usage = %+v, want {prompt:6 completion:3} (proves the streaming path, not JSON fallback, was taken)", u) + } + if got := rec.header.Get("Content-Type"); got != "text/event-stream" { + t.Errorf("Content-Type on client response = %q, want %q (sseWriter's own header; JSON fallback would copy the upstream's mixed-case value verbatim)", got, "text/event-stream") + } +} + +// TestGeminiAdapter_ChatCompletion_Upstream429 proves a non-2xx upstream +// response returns a *providerHTTPError with nothing written to the +// client. +func TestGeminiAdapter_ChatCompletion_Upstream429(t *testing.T) { + const errBody = `{"error":{"code":429,"message":"rate limited","status":"RESOURCE_EXHAUSTED"}}` + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(errBody)) + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + rec := httptest.NewRecorder() + + u, err := a.chatCompletion(context.Background(), rec, map[string]any{ + "model": "gemini-2.5-pro", + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + }) + if err == nil { + t.Fatalf("chatCompletion: want error, got nil (usage=%+v)", u) + } + var httpErr *providerHTTPError + if !errors.As(err, &httpErr) { + t.Fatalf("err = %v (%T), want *providerHTTPError", err, err) + } + if httpErr.status != http.StatusTooManyRequests { + t.Errorf("status = %d, want %d", httpErr.status, http.StatusTooManyRequests) + } + if rec.Body.Len() != 0 { + t.Errorf("client response body = %q, want empty", rec.Body.String()) + } +} + +// TestGeminiAdapter_ChatCompletion_TranslateErrorSkipsUpstream proves a +// request field geminiRequestFromOpenAI rejects (n>1, here) never reaches +// the upstream at all. +func TestGeminiAdapter_ChatCompletion_TranslateErrorSkipsUpstream(t *testing.T) { + called := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + rec := httptest.NewRecorder() + + _, err = a.chatCompletion(context.Background(), rec, map[string]any{ + "model": "gemini-2.5-pro", + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + "n": float64(2), + }) + var terr *translateError + if !errors.As(err, &terr) { + t.Fatalf("err = %v (%T), want *translateError", err, err) + } + if called { + t.Errorf("upstream was called, want the translate error to short-circuit before any request") + } +} + +// TestGeminiAdapter_Embeddings_Single proves embeddings builds the +// :embedContent URL and translates the response to an OpenAI embeddings +// list body for a single-string "input". +func TestGeminiAdapter_Embeddings_Single(t *testing.T) { + var gotPath string + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"embedding":{"values":[0.1,0.2,0.3]}}`)) + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + rec := httptest.NewRecorder() + + u, err := a.embeddings(context.Background(), rec, map[string]any{ + "model": "text-embedding-004", + "input": "hello world", + }) + if err != nil { + t.Fatalf("embeddings: %v", err) + } + if gotPath != "/v1beta/models/text-embedding-004:embedContent" { + t.Errorf("path = %q, want /v1beta/models/text-embedding-004:embedContent", gotPath) + } + sent := decodeJSONBody(t, gotBody) + if _, hasRequests := sent["requests"]; hasRequests { + t.Errorf("single-input body carries \"requests\", want the :embedContent shape") + } + if !u.estimated { + t.Errorf("usage.estimated = false, want true (Gemini reports no embeddings usage)") + } + + var out map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String()) + } + if out["object"] != "list" { + t.Errorf("object = %v, want %q", out["object"], "list") + } + data, ok := out["data"].([]any) + if !ok || len(data) != 1 { + t.Fatalf("data = %v, want a single-element array", out["data"]) + } +} + +// TestGeminiAdapter_Embeddings_Batch proves an array "input" selects +// Gemini's :batchEmbedContents endpoint and translates its response to an +// OpenAI embeddings list body with one entry per input string. +func TestGeminiAdapter_Embeddings_Batch(t *testing.T) { + var gotPath string + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"embeddings":[{"values":[0.1,0.2]},{"values":[0.3,0.4]}]}`)) + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + rec := httptest.NewRecorder() + + _, err = a.embeddings(context.Background(), rec, map[string]any{ + "model": "text-embedding-004", + "input": []any{"hello", "world!!"}, + }) + if err != nil { + t.Fatalf("embeddings: %v", err) + } + if gotPath != "/v1beta/models/text-embedding-004:batchEmbedContents" { + t.Errorf("path = %q, want /v1beta/models/text-embedding-004:batchEmbedContents", gotPath) + } + sent := decodeJSONBody(t, gotBody) + requests, ok := sent["requests"].([]any) + if !ok || len(requests) != 2 { + t.Fatalf("requests = %v, want a 2-element array", sent["requests"]) + } + first := requests[0].(map[string]any) + if first["model"] != "models/text-embedding-004" { + t.Errorf("requests[0].model = %v, want %q", first["model"], "models/text-embedding-004") + } + + var out map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String()) + } + data, ok := out["data"].([]any) + if !ok || len(data) != 2 { + t.Fatalf("data = %v, want a 2-element array", out["data"]) + } +} + +// TestGeminiAdapter_ListModels proves listModels parses Gemini's +// "models": [{"name": "models/..."}] Models API response into a flat +// slice of bare ids, stripping the "models/" prefix. +func TestGeminiAdapter_ListModels(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1beta/models" { + t.Errorf("path = %q, want /v1beta/models", r.URL.Path) + return + } + if r.Method != http.MethodGet { + t.Errorf("method = %q, want GET", r.Method) + return + } + if got := r.Header.Get("x-goog-api-key"); got != "AIza-test" { + t.Errorf("x-goog-api-key = %q, want %q", got, "AIza-test") + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}]}`)) + })) + defer srv.Close() + + a, err := newGeminiAdapter("p1", srv.URL, "AIza-test") + if err != nil { + t.Fatalf("newGeminiAdapter: %v", err) + } + + got, err := a.listModels(context.Background()) + if err != nil { + t.Fatalf("listModels: %v", err) + } + want := []string{"gemini-2.5-pro", "gemini-2.5-flash"} + if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] { + t.Errorf("got %#v, want %#v", got, want) + } +} diff --git a/provider_openai_test.go b/provider_openai_test.go index a99b2a0..ec13f8b 100644 --- a/provider_openai_test.go +++ b/provider_openai_test.go @@ -629,13 +629,32 @@ func TestBuildAdapters(t *testing.T) { } }) - t.Run("edge: gemini type not yet implemented", func(t *testing.T) { + t.Run("good: gemini with an API key builds successfully", func(t *testing.T) { cfg := &Config{Providers: map[string]*ProviderConfig{ "p1": {Type: "gemini", APIKey: "sk-test"}, }} - _, err := buildAdapters(cfg) - if err == nil { - t.Fatalf("buildAdapters: want placeholder error for gemini, got nil") + adapters, err := buildAdapters(cfg) + if err != nil { + t.Fatalf("buildAdapters: %v", err) + } + a, ok := adapters["p1"] + if !ok { + t.Fatalf("adapters[%q] missing", "p1") + } + if a.typeName() != "gemini" { + t.Errorf("typeName() = %q, want %q", a.typeName(), "gemini") + } + if a.base() != defaultBaseGemini { + t.Errorf("base() = %q, want %q", a.base(), defaultBaseGemini) + } + }) + + t.Run("bad: gemini without an API key is a constructor error", func(t *testing.T) { + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "gemini"}, + }} + if _, err := buildAdapters(cfg); err == nil { + t.Fatalf("buildAdapters: want error for keyless gemini provider, got nil") } }) } diff --git a/providers.go b/providers.go index e6b24d6..5612b9b 100644 --- a/providers.go +++ b/providers.go @@ -182,8 +182,8 @@ func upstreamJSON(ctx context.Context, client *http.Client, method, url string, // keyless upstream for openai-type providers, but a constructor error for // anthropic-type ones — newAnthropicAdapter enforces that, per ruling (c); // this function just propagates whatever error it returns. A configured -// gemini provider currently errors too — Task 10 fills in that switch -// case. +// gemini provider follows the same constructor-error ruling — +// newGeminiAdapter enforces it too. func buildAdapters(cfg *Config) (map[string]providerAdapter, error) { adapters := make(map[string]providerAdapter, len(cfg.Providers)) for name, pc := range cfg.Providers { @@ -207,7 +207,11 @@ func buildAdapters(cfg *Config) (map[string]providerAdapter, error) { } adapters[name] = a case providerTypeGemini: - return nil, fmt.Errorf("llmgateway: provider %q: type %q not yet implemented", name, pc.Type) + g, err := newGeminiAdapter(name, base, apiKey) + if err != nil { + return nil, err + } + adapters[name] = g default: return nil, fmt.Errorf("llmgateway: provider %q: unknown type %q", name, pc.Type) } diff --git a/testdata/gemini/chat_basic_in.json b/testdata/gemini/chat_basic_in.json new file mode 100644 index 0000000..9123c6d --- /dev/null +++ b/testdata/gemini/chat_basic_in.json @@ -0,0 +1,8 @@ +{ + "model": "gemini-2.5-pro", + "messages": [ + {"role": "user", "content": "Hello there"} + ], + "max_tokens": 256, + "temperature": 0.7 +} diff --git a/testdata/gemini/chat_basic_want.json b/testdata/gemini/chat_basic_want.json new file mode 100644 index 0000000..6e1063d --- /dev/null +++ b/testdata/gemini/chat_basic_want.json @@ -0,0 +1,9 @@ +{ + "contents": [ + {"role": "user", "parts": [{"text": "Hello there"}]} + ], + "generationConfig": { + "maxOutputTokens": 256, + "temperature": 0.7 + } +} diff --git a/testdata/gemini/chat_images_http_error_in.json b/testdata/gemini/chat_images_http_error_in.json new file mode 100644 index 0000000..a1707ac --- /dev/null +++ b/testdata/gemini/chat_images_http_error_in.json @@ -0,0 +1,11 @@ +{ + "model": "gemini-2.5-pro", + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}} + ] + } + ] +} diff --git a/testdata/gemini/chat_images_http_error_wanterr.json b/testdata/gemini/chat_images_http_error_wanterr.json new file mode 100644 index 0000000..0578360 --- /dev/null +++ b/testdata/gemini/chat_images_http_error_wanterr.json @@ -0,0 +1 @@ +{"errContains": "must be a data: URI", "notSupported": false} diff --git a/testdata/gemini/chat_images_in.json b/testdata/gemini/chat_images_in.json new file mode 100644 index 0000000..3ee1202 --- /dev/null +++ b/testdata/gemini/chat_images_in.json @@ -0,0 +1,12 @@ +{ + "model": "gemini-2.5-pro", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image_url", "image_url": {"url": "data:image/jpeg;charset=utf-8;base64,/9j/4AAQSkZJRg=="}} + ] + } + ] +} diff --git a/testdata/gemini/chat_images_want.json b/testdata/gemini/chat_images_want.json new file mode 100644 index 0000000..32f422f --- /dev/null +++ b/testdata/gemini/chat_images_want.json @@ -0,0 +1,11 @@ +{ + "contents": [ + { + "role": "user", + "parts": [ + {"text": "What is in this image?"}, + {"inlineData": {"mimeType": "image/jpeg", "data": "/9j/4AAQSkZJRg=="}} + ] + } + ] +} diff --git a/testdata/gemini/chat_parallel_tool_results_in.json b/testdata/gemini/chat_parallel_tool_results_in.json new file mode 100644 index 0000000..84e87a6 --- /dev/null +++ b/testdata/gemini/chat_parallel_tool_results_in.json @@ -0,0 +1,16 @@ +{ + "model": "gemini-2.5-pro", + "messages": [ + {"role": "user", "content": "What's the weather in Paris and Tokyo?"}, + { + "role": "assistant", + "content": null, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{\"location\":\"Paris\"}"}}, + {"id": "call_2", "type": "function", "function": {"name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}"}} + ] + }, + {"role": "tool", "tool_call_id": "call_1", "content": "{\"temp_c\":18}"}, + {"role": "tool", "tool_call_id": "call_2", "content": "{\"temp_c\":27}"} + ] +} diff --git a/testdata/gemini/chat_parallel_tool_results_want.json b/testdata/gemini/chat_parallel_tool_results_want.json new file mode 100644 index 0000000..90abdbd --- /dev/null +++ b/testdata/gemini/chat_parallel_tool_results_want.json @@ -0,0 +1,19 @@ +{ + "contents": [ + {"role": "user", "parts": [{"text": "What's the weather in Paris and Tokyo?"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"name": "get_weather", "args": {"location": "Paris"}}}, + {"functionCall": {"name": "get_weather", "args": {"location": "Tokyo"}}} + ] + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"name": "get_weather", "response": {"content": "{\"temp_c\":18}"}}}, + {"functionResponse": {"name": "get_weather", "response": {"content": "{\"temp_c\":27}"}}} + ] + } + ] +} diff --git a/testdata/gemini/chat_system_array_error_in.json b/testdata/gemini/chat_system_array_error_in.json new file mode 100644 index 0000000..d6e848f --- /dev/null +++ b/testdata/gemini/chat_system_array_error_in.json @@ -0,0 +1,7 @@ +{ + "model": "gemini-2.5-pro", + "messages": [ + {"role": "system", "content": [{"type": "image_url", "image_url": {"url": "https://x"}}]}, + {"role": "user", "content": "Hi"} + ] +} diff --git a/testdata/gemini/chat_system_array_error_wanterr.json b/testdata/gemini/chat_system_array_error_wanterr.json new file mode 100644 index 0000000..9f97b90 --- /dev/null +++ b/testdata/gemini/chat_system_array_error_wanterr.json @@ -0,0 +1 @@ +{"errContains": "not supported for gemini models", "notSupported": false} diff --git a/testdata/gemini/chat_system_array_in.json b/testdata/gemini/chat_system_array_in.json new file mode 100644 index 0000000..c36a3ab --- /dev/null +++ b/testdata/gemini/chat_system_array_in.json @@ -0,0 +1,7 @@ +{ + "model": "gemini-2.5-pro", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "Line one."}, {"type": "text", "text": "Line two."}]}, + {"role": "user", "content": "Hi"} + ] +} diff --git a/testdata/gemini/chat_system_array_want.json b/testdata/gemini/chat_system_array_want.json new file mode 100644 index 0000000..9ea40ed --- /dev/null +++ b/testdata/gemini/chat_system_array_want.json @@ -0,0 +1,6 @@ +{ + "systemInstruction": {"parts": [{"text": "Line one.\nLine two."}]}, + "contents": [ + {"role": "user", "parts": [{"text": "Hi"}]} + ] +} diff --git a/testdata/gemini/chat_tool_choice_none_in.json b/testdata/gemini/chat_tool_choice_none_in.json new file mode 100644 index 0000000..5408483 --- /dev/null +++ b/testdata/gemini/chat_tool_choice_none_in.json @@ -0,0 +1,11 @@ +{ + "model": "gemini-2.5-pro", + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": {"name": "noop", "description": "does nothing", "parameters": {"type": "object", "properties": {}}} + } + ], + "tool_choice": "none" +} diff --git a/testdata/gemini/chat_tool_choice_none_want.json b/testdata/gemini/chat_tool_choice_none_want.json new file mode 100644 index 0000000..85836ed --- /dev/null +++ b/testdata/gemini/chat_tool_choice_none_want.json @@ -0,0 +1,15 @@ +{ + "contents": [ + {"role": "user", "parts": [{"text": "hi"}]} + ], + "tools": [ + { + "functionDeclarations": [ + {"name": "noop", "description": "does nothing", "parameters": {"type": "object", "properties": {}}} + ] + } + ], + "toolConfig": { + "functionCallingConfig": {"mode": "NONE"} + } +} diff --git a/testdata/gemini/chat_tools_in.json b/testdata/gemini/chat_tools_in.json new file mode 100644 index 0000000..10bee49 --- /dev/null +++ b/testdata/gemini/chat_tools_in.json @@ -0,0 +1,22 @@ +{ + "model": "gemini-2.5-pro", + "messages": [ + {"role": "system", "content": "You are a helpful weather bot."}, + {"role": "user", "content": "What's the weather in Paris?"} + ], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"] + } + } + } + ], + "tool_choice": {"type": "function", "function": {"name": "get_weather"}} +} diff --git a/testdata/gemini/chat_tools_want.json b/testdata/gemini/chat_tools_want.json new file mode 100644 index 0000000..693c32e --- /dev/null +++ b/testdata/gemini/chat_tools_want.json @@ -0,0 +1,27 @@ +{ + "systemInstruction": {"parts": [{"text": "You are a helpful weather bot."}]}, + "contents": [ + {"role": "user", "parts": [{"text": "What's the weather in Paris?"}]} + ], + "tools": [ + { + "functionDeclarations": [ + { + "name": "get_weather", + "description": "Get current weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"] + } + } + ] + } + ], + "toolConfig": { + "functionCallingConfig": { + "mode": "ANY", + "allowedFunctionNames": ["get_weather"] + } + } +} diff --git a/testdata/gemini/chat_unsupported_logit_bias_in.json b/testdata/gemini/chat_unsupported_logit_bias_in.json new file mode 100644 index 0000000..48bccd9 --- /dev/null +++ b/testdata/gemini/chat_unsupported_logit_bias_in.json @@ -0,0 +1,5 @@ +{ + "model": "gemini-2.5-pro", + "messages": [{"role": "user", "content": "hi"}], + "logit_bias": {"1234": 10} +} diff --git a/testdata/gemini/chat_unsupported_logit_bias_wanterr.json b/testdata/gemini/chat_unsupported_logit_bias_wanterr.json new file mode 100644 index 0000000..5d3dfef --- /dev/null +++ b/testdata/gemini/chat_unsupported_logit_bias_wanterr.json @@ -0,0 +1 @@ +{"errContains": "logit_bias is not supported for gemini models", "notSupported": false} diff --git a/testdata/gemini/chat_unsupported_n_in.json b/testdata/gemini/chat_unsupported_n_in.json new file mode 100644 index 0000000..1617f9f --- /dev/null +++ b/testdata/gemini/chat_unsupported_n_in.json @@ -0,0 +1,5 @@ +{ + "model": "gemini-2.5-pro", + "messages": [{"role": "user", "content": "hi"}], + "n": 2 +} diff --git a/testdata/gemini/chat_unsupported_n_wanterr.json b/testdata/gemini/chat_unsupported_n_wanterr.json new file mode 100644 index 0000000..2a14758 --- /dev/null +++ b/testdata/gemini/chat_unsupported_n_wanterr.json @@ -0,0 +1 @@ +{"errContains": "n is not supported for gemini models", "notSupported": false} diff --git a/testdata/gemini/embed_batch_in.json b/testdata/gemini/embed_batch_in.json new file mode 100644 index 0000000..0e63ccf --- /dev/null +++ b/testdata/gemini/embed_batch_in.json @@ -0,0 +1,4 @@ +{ + "model": "text-embedding-004", + "input": ["hello", "world!!"] +} diff --git a/testdata/gemini/embed_batch_resp_in.json b/testdata/gemini/embed_batch_resp_in.json new file mode 100644 index 0000000..bebebb1 --- /dev/null +++ b/testdata/gemini/embed_batch_resp_in.json @@ -0,0 +1,6 @@ +{ + "embeddings": [ + {"values": [0.1, 0.2]}, + {"values": [0.3, 0.4]} + ] +} diff --git a/testdata/gemini/embed_batch_resp_want.json b/testdata/gemini/embed_batch_resp_want.json new file mode 100644 index 0000000..3db08cb --- /dev/null +++ b/testdata/gemini/embed_batch_resp_want.json @@ -0,0 +1,9 @@ +{ + "object": "list", + "data": [ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, + {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]} + ], + "model": "gw-embed-model", + "usage": {"prompt_tokens": 3, "total_tokens": 3} +} diff --git a/testdata/gemini/embed_batch_want.json b/testdata/gemini/embed_batch_want.json new file mode 100644 index 0000000..18116b7 --- /dev/null +++ b/testdata/gemini/embed_batch_want.json @@ -0,0 +1,6 @@ +{ + "requests": [ + {"model": "models/text-embedding-004", "content": {"parts": [{"text": "hello"}]}}, + {"model": "models/text-embedding-004", "content": {"parts": [{"text": "world!!"}]}} + ] +} diff --git a/testdata/gemini/embed_single_in.json b/testdata/gemini/embed_single_in.json new file mode 100644 index 0000000..7e5fe7c --- /dev/null +++ b/testdata/gemini/embed_single_in.json @@ -0,0 +1,4 @@ +{ + "model": "text-embedding-004", + "input": "hello world" +} diff --git a/testdata/gemini/embed_single_resp_in.json b/testdata/gemini/embed_single_resp_in.json new file mode 100644 index 0000000..71b5e0d --- /dev/null +++ b/testdata/gemini/embed_single_resp_in.json @@ -0,0 +1,3 @@ +{ + "embedding": {"values": [0.1, 0.2, 0.3]} +} diff --git a/testdata/gemini/embed_single_resp_want.json b/testdata/gemini/embed_single_resp_want.json new file mode 100644 index 0000000..2d8da6a --- /dev/null +++ b/testdata/gemini/embed_single_resp_want.json @@ -0,0 +1,8 @@ +{ + "object": "list", + "data": [ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + ], + "model": "gw-embed-model", + "usage": {"prompt_tokens": 3, "total_tokens": 3} +} diff --git a/testdata/gemini/embed_single_want.json b/testdata/gemini/embed_single_want.json new file mode 100644 index 0000000..b2cb5a8 --- /dev/null +++ b/testdata/gemini/embed_single_want.json @@ -0,0 +1,3 @@ +{ + "content": {"parts": [{"text": "hello world"}]} +} diff --git a/testdata/gemini/resp_funcall_in.json b/testdata/gemini/resp_funcall_in.json new file mode 100644 index 0000000..8275639 --- /dev/null +++ b/testdata/gemini/resp_funcall_in.json @@ -0,0 +1,18 @@ +{ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"functionCall": {"id": "fc_1", "name": "get_weather", "args": {"location": "Paris"}}} + ] + }, + "finishReason": "STOP" + } + ], + "usageMetadata": { + "promptTokenCount": 12, + "candidatesTokenCount": 6 + }, + "responseId": "resp_FUNC1" +} diff --git a/testdata/gemini/resp_funcall_want.json b/testdata/gemini/resp_funcall_want.json new file mode 100644 index 0000000..9390d4a --- /dev/null +++ b/testdata/gemini/resp_funcall_want.json @@ -0,0 +1,24 @@ +{ + "id": "chatcmpl-resp_FUNC1", + "object": "chat.completion", + "created": 1734000000, + "model": "gw-gemini-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{\"location\":\"Paris\"}"} + } + ] + }, + "finish_reason": "tool_calls" + } + ], + "usage": {"prompt_tokens": 12, "completion_tokens": 6, "total_tokens": 18} +} diff --git a/testdata/gemini/resp_maxtokens_in.json b/testdata/gemini/resp_maxtokens_in.json new file mode 100644 index 0000000..6327a5d --- /dev/null +++ b/testdata/gemini/resp_maxtokens_in.json @@ -0,0 +1,13 @@ +{ + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "This response got cut off because it"}] + }, + "finishReason": "MAX_TOKENS" + } + ], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 50}, + "responseId": "resp_MAXTOK" +} diff --git a/testdata/gemini/resp_maxtokens_want.json b/testdata/gemini/resp_maxtokens_want.json new file mode 100644 index 0000000..36e29d5 --- /dev/null +++ b/testdata/gemini/resp_maxtokens_want.json @@ -0,0 +1,14 @@ +{ + "id": "chatcmpl-resp_MAXTOK", + "object": "chat.completion", + "created": 1734000000, + "model": "gw-gemini-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "This response got cut off because it"}, + "finish_reason": "length" + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 50, "total_tokens": 55} +} diff --git a/testdata/gemini/resp_safety_in.json b/testdata/gemini/resp_safety_in.json new file mode 100644 index 0000000..19fa5a6 --- /dev/null +++ b/testdata/gemini/resp_safety_in.json @@ -0,0 +1,10 @@ +{ + "candidates": [ + { + "content": {"role": "model", "parts": []}, + "finishReason": "SAFETY" + } + ], + "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 0}, + "responseId": "resp_SAFETY1" +} diff --git a/testdata/gemini/resp_safety_want.json b/testdata/gemini/resp_safety_want.json new file mode 100644 index 0000000..2d4c7b0 --- /dev/null +++ b/testdata/gemini/resp_safety_want.json @@ -0,0 +1,14 @@ +{ + "id": "chatcmpl-resp_SAFETY1", + "object": "chat.completion", + "created": 1734000000, + "model": "gw-gemini-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": null}, + "finish_reason": "content_filter" + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 0, "total_tokens": 4} +} diff --git a/testdata/gemini/resp_text_in.json b/testdata/gemini/resp_text_in.json new file mode 100644 index 0000000..439fa6d --- /dev/null +++ b/testdata/gemini/resp_text_in.json @@ -0,0 +1,18 @@ +{ + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Hello! How can I help you today?"}] + }, + "finishReason": "STOP" + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 9, + "totalTokenCount": 17 + }, + "modelVersion": "gemini-2.5-pro-001", + "responseId": "resp_ABC123" +} diff --git a/testdata/gemini/resp_text_want.json b/testdata/gemini/resp_text_want.json new file mode 100644 index 0000000..4ea17be --- /dev/null +++ b/testdata/gemini/resp_text_want.json @@ -0,0 +1,14 @@ +{ + "id": "chatcmpl-resp_ABC123", + "object": "chat.completion", + "created": 1734000000, + "model": "gw-gemini-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello! How can I help you today?"}, + "finish_reason": "stop" + } + ], + "usage": {"prompt_tokens": 8, "completion_tokens": 9, "total_tokens": 17} +} diff --git a/testdata/gemini/stream_text.sse b/testdata/gemini/stream_text.sse new file mode 100644 index 0000000..006fa9e --- /dev/null +++ b/testdata/gemini/stream_text.sse @@ -0,0 +1,6 @@ +data: {"candidates":[{"content":{"role":"model","parts":[{"text":"Hello"}]}}],"responseId":"resp_STREAM1"} + +data: {"candidates":[{"content":{"role":"model","parts":[{"text":", world!"}]}}]} + +data: {"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3}} + diff --git a/testdata/gemini/stream_text_then_funcall.sse b/testdata/gemini/stream_text_then_funcall.sse new file mode 100644 index 0000000..15de92d --- /dev/null +++ b/testdata/gemini/stream_text_then_funcall.sse @@ -0,0 +1,6 @@ +data: {"candidates":[{"content":{"role":"model","parts":[{"text":"Let me check that."}]}}],"responseId":"resp_STREAM2"} + +data: {"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"get_weather","args":{"location":"Berlin"}}}]}}]} + +data: {"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":4}} + diff --git a/testdata/gemini/stream_text_then_funcall_usage.json b/testdata/gemini/stream_text_then_funcall_usage.json new file mode 100644 index 0000000..ad04e9b --- /dev/null +++ b/testdata/gemini/stream_text_then_funcall_usage.json @@ -0,0 +1 @@ +{"prompt": 10, "completion": 4} diff --git a/testdata/gemini/stream_text_then_funcall_want.json b/testdata/gemini/stream_text_then_funcall_want.json new file mode 100644 index 0000000..a219082 --- /dev/null +++ b/testdata/gemini/stream_text_then_funcall_want.json @@ -0,0 +1,6 @@ +[ + {"id":"chatcmpl-resp_STREAM2","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}, + {"id":"chatcmpl-resp_STREAM2","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{"content":"Let me check that."},"finish_reason":null}]}, + {"id":"chatcmpl-resp_STREAM2","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_0","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"Berlin\"}"}}]},"finish_reason":null}]}, + {"id":"chatcmpl-resp_STREAM2","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]} +] diff --git a/testdata/gemini/stream_text_usage.json b/testdata/gemini/stream_text_usage.json new file mode 100644 index 0000000..ee07b50 --- /dev/null +++ b/testdata/gemini/stream_text_usage.json @@ -0,0 +1 @@ +{"prompt": 5, "completion": 3} diff --git a/testdata/gemini/stream_text_want.json b/testdata/gemini/stream_text_want.json new file mode 100644 index 0000000..22d8c84 --- /dev/null +++ b/testdata/gemini/stream_text_want.json @@ -0,0 +1,6 @@ +[ + {"id":"chatcmpl-resp_STREAM1","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}, + {"id":"chatcmpl-resp_STREAM1","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}, + {"id":"chatcmpl-resp_STREAM1","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{"content":", world!"},"finish_reason":null}]}, + {"id":"chatcmpl-resp_STREAM1","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]} +] diff --git a/translate_gemini.go b/translate_gemini.go new file mode 100644 index 0000000..fde58eb --- /dev/null +++ b/translate_gemini.go @@ -0,0 +1,836 @@ +package traefikllmgateway + +import ( + "encoding/json" + "fmt" + "strconv" + "strings" +) + +// geminiAPIPrefix is the path segment every Gemini generateContent-family +// and embedding-family URL is built under, followed by "{model}:{method}". +const geminiAPIPrefix = "/v1beta/models/" + +// isTruthyMap reports whether v decodes to a non-empty JSON object. Per +// lesson (i), an OpenAI field this translator rejects when present ( +// logit_bias) is only rejected when it actually carries content — a +// present-but-empty {} or a JSON null must not 400 a request that would +// otherwise translate cleanly. +func isTruthyMap(v any) bool { + m, ok := v.(map[string]any) + return ok && len(m) > 0 +} + +// geminiUnsupportedFieldError returns a *translateError for an OpenAI +// request field Gemini's API has no equivalent for. Deliberately not +// translate_anthropic.go's unsupportedFieldError: that helper's message +// hardcodes "not supported for anthropic models", which would misreport the +// provider on a Gemini request, and this file may only reuse translateError +// itself, not touch translate_anthropic.go to parameterize it. +func geminiUnsupportedFieldError(field string) *translateError { + return &translateError{msg: fmt.Sprintf("%s is not supported for gemini models", field)} +} + +// geminiSystemTextFromContentParts joins an OpenAI system message's +// array-form content ([{"type":"text","text":...}, ...]) into plain text, +// joined with "\n", per lesson (e). Any part that is not {"type":"text"} +// has no equivalent in Gemini's plain-string systemInstruction text, so it +// is a *translateError rather than being silently dropped. Deliberately +// not translate_anthropic.go's systemTextFromContentParts, for the same +// hardcoded-provider-name reason as geminiUnsupportedFieldError above. +func geminiSystemTextFromContentParts(parts []any) (string, error) { + texts := make([]string, 0, len(parts)) + for _, part := range parts { + pm, ok := part.(map[string]any) + if !ok { + return "", &translateError{msg: "system message content[] entries must be objects"} + } + ptype, _ := pm["type"].(string) + if ptype != "text" { + return "", &translateError{msg: fmt.Sprintf("system message content part type %q is not supported for gemini models", ptype)} + } + text, _ := pm["text"].(string) + texts = append(texts, text) + } + return strings.Join(texts, "\n"), nil +} + +// geminiInlineDataFromDataURI parses an OpenAI image_url.url as a +// "data:," URI and returns Gemini's inlineData +// object. Per lesson (d), the media type is everything before the first +// ";" after "data:" — any parameters after it (";charset=utf-8;base64", for +// instance) are stripped rather than rejected or included. Any URL that is +// not a data: URI (an http(s) image URL, which Gemini's generateContent API +// does not fetch) returns a *translateError, per the request mapping +// table's image row. +func geminiInlineDataFromDataURI(uri string) (map[string]any, error) { + rest, ok := strings.CutPrefix(uri, "data:") + if !ok { + return nil, &translateError{msg: "image_url must be a data: URI (http(s) image URLs are not supported for gemini models)"} + } + header, data, ok := strings.Cut(rest, ",") + if !ok { + return nil, &translateError{msg: "malformed data: URI in image_url"} + } + mimeType, _, _ := strings.Cut(header, ";") + return map[string]any{"mimeType": mimeType, "data": data}, nil +} + +// geminiPartFromOpenAI maps one OpenAI content-array entry ({"type":"text", +// ...} or {"type":"image_url",...}) to a Gemini content part. An entry of +// any other type is dropped (nil, nil) rather than erroring — an +// unrecognized part type is forwarded content this translator does not +// understand, not necessarily an invalid request. +func geminiPartFromOpenAI(pm map[string]any) (map[string]any, error) { + switch ptype, _ := pm["type"].(string); ptype { + case "text": + text, _ := pm["text"].(string) + return map[string]any{"text": text}, nil + case "image_url": + iu, _ := pm["image_url"].(map[string]any) + url, _ := iu["url"].(string) + inline, err := geminiInlineDataFromDataURI(url) + if err != nil { + return nil, err + } + return map[string]any{"inlineData": inline}, nil + default: + return nil, nil + } +} + +// geminiFunctionCallPartFromOpenAI maps one assistant tool_calls[] entry +// ({"id","type":"function","function":{"name","arguments"}}) to a Gemini +// {"functionCall":{"name","args"}} content part, parsing the OpenAI +// arguments JSON string into Gemini's structured "args" object. It also +// returns the entry's id and name, so the caller can record them in the +// pending tool_call id map a later role:"tool" message's functionResponse +// needs (Gemini's functionResponse has no id field to correlate by; this +// translator resolves "name" instead, per the request mapping table's +// role:tool row). +func geminiFunctionCallPartFromOpenAI(tc any) (part map[string]any, id, name string, err error) { + tm, ok := tc.(map[string]any) + if !ok { + return nil, "", "", &translateError{msg: "assistant tool_calls[] entry is malformed"} + } + id, _ = tm["id"].(string) + fn, _ := tm["function"].(map[string]any) + name, _ = fn["name"].(string) + argsStr, _ := fn["arguments"].(string) + + args := any(map[string]any{}) + if argsStr != "" { + if err := json.Unmarshal([]byte(argsStr), &args); err != nil { + return nil, "", "", &translateError{msg: "assistant tool_calls[].function.arguments is not valid JSON"} + } + } + return map[string]any{"functionCall": map[string]any{"name": name, "args": args}}, id, name, nil +} + +// geminiMessageFromOpenAI maps one OpenAI user/assistant message to a +// Gemini {"role","parts"} contents entry. geminiRole is "user" or "model" +// (assistant's Gemini-side role name, per the request mapping table). +// toolNameByID accumulates id -> name for every tool_calls[] entry seen on +// an assistant message, so a later role:"tool" message can resolve its +// functionResponse's "name". +func geminiMessageFromOpenAI(msg map[string]any, geminiRole string, toolNameByID map[string]string) (map[string]any, error) { + var parts []any + switch c := msg["content"].(type) { + case string: + if c != "" { + parts = append(parts, map[string]any{"text": c}) + } + case []any: + for _, p := range c { + pm, ok := p.(map[string]any) + if !ok { + continue + } + part, err := geminiPartFromOpenAI(pm) + if err != nil { + return nil, err + } + if part != nil { + parts = append(parts, part) + } + } + } + + if geminiRole == "model" { + toolCalls, _ := msg["tool_calls"].([]any) + for _, tc := range toolCalls { + part, id, name, err := geminiFunctionCallPartFromOpenAI(tc) + if err != nil { + return nil, err + } + if id != "" { + toolNameByID[id] = name + } + parts = append(parts, part) + } + } + + if parts == nil { + parts = []any{} + } + return map[string]any{"role": geminiRole, "parts": parts}, nil +} + +// geminiFunctionResponsePart maps one OpenAI role:"tool" message to a +// single Gemini {"functionResponse":{"name","response":{"content"}}} +// content part, per the request mapping table's role:tool row. name is +// looked up from toolNameByID — the id -> name map built while converting +// the assistant message that produced this tool_call_id — and an unknown +// id is a *translateError rather than a silently empty name. +func geminiFunctionResponsePart(msg map[string]any, toolNameByID map[string]string) (map[string]any, error) { + toolCallID, _ := msg["tool_call_id"].(string) + name, ok := toolNameByID[toolCallID] + if !ok { + return nil, &translateError{msg: fmt.Sprintf("role:tool message references unknown tool_call_id %q", toolCallID)} + } + return map[string]any{ + "functionResponse": map[string]any{ + "name": name, + "response": map[string]any{"content": msg["content"]}, + }, + }, nil +} + +// geminiToolsFromOpenAI maps OpenAI's tools[].function{name, description, +// parameters} to Gemini's tools:[{functionDeclarations:[{name,description, +// parameters}]}], per the request mapping table's tools row. +func geminiToolsFromOpenAI(toolsRaw []any) []any { + decls := make([]any, 0, len(toolsRaw)) + for _, t := range toolsRaw { + tm, ok := t.(map[string]any) + if !ok { + continue + } + fn, _ := tm["function"].(map[string]any) + entry := map[string]any{"name": fn["name"]} + if desc, ok := fn["description"]; ok { + entry["description"] = desc + } + if params, ok := fn["parameters"]; ok { + entry["parameters"] = params + } + decls = append(decls, entry) + } + return []any{map[string]any{"functionDeclarations": decls}} +} + +// geminiToolConfigFromOpenAI maps OpenAI's "tool_choice" to Gemini's +// toolConfig.functionCallingConfig object, per the request mapping table's +// tool_choice rows. Unlike Anthropic, Gemini has a real NONE mode, so +// "none" needs no special "drop tools" handling. An unrecognized shape +// returns nil — the caller omits "toolConfig" entirely rather than sending +// a malformed one. +func geminiToolConfigFromOpenAI(tc any) map[string]any { + switch t := tc.(type) { + case string: + switch t { + case "none": + return map[string]any{"functionCallingConfig": map[string]any{"mode": "NONE"}} + case "auto": + return map[string]any{"functionCallingConfig": map[string]any{"mode": "AUTO"}} + case "required": + return map[string]any{"functionCallingConfig": map[string]any{"mode": "ANY"}} + } + case map[string]any: + if fn, ok := t["function"].(map[string]any); ok { + if name, ok := fn["name"].(string); ok { + return map[string]any{"functionCallingConfig": map[string]any{ + "mode": "ANY", + "allowedFunctionNames": []any{name}, + }} + } + } + } + return nil +} + +// geminiRequestFromOpenAI maps an OpenAI chat-completion request body to a +// Gemini generateContent/streamGenerateContent request body, per this +// file's request mapping table. It returns a *translateError — never +// wrapped — for a field Gemini has no equivalent for (n>1, a truthy +// logit_bias) or a content part this translator cannot map (a non-data-URI +// image URL, a malformed tool_calls argument string, a role:"tool" message +// whose tool_call_id was never seen on a preceding assistant message). +// Consecutive role:"tool" messages are merged into one Gemini {"role": +// "user"} contents entry carrying all of their functionResponse parts, per +// lesson (f). The request's "model" field is deliberately never read here — +// Gemini's body carries no model field at all; the adapter reads it +// separately to build the URL. +func geminiRequestFromOpenAI(req map[string]any) (map[string]any, error) { + if v, ok := req["n"]; ok { + if n, ok2 := toFloat64(v); ok2 && n > 1 { + return nil, geminiUnsupportedFieldError("n") + } + } + if v, ok := req["logit_bias"]; ok && isTruthyMap(v) { + return nil, geminiUnsupportedFieldError("logit_bias") + } + + out := map[string]any{} + + msgsRaw, _ := req["messages"].([]any) + var systemParts []string + contents := make([]any, 0, len(msgsRaw)) + toolNameByID := map[string]string{} + + // pendingFuncResponses accumulates consecutive role:"tool" messages' + // functionResponse parts. Gemini, like Anthropic, expects the N + // tool-result messages following a parallel-tool-call assistant turn + // to collapse into one contents entry, not N separate ones — see + // geminiFunctionResponsePart's doc comment for the id resolution this + // depends on. + var pendingFuncResponses []any + flushPending := func() { + if len(pendingFuncResponses) == 0 { + return + } + contents = append(contents, map[string]any{"role": "user", "parts": pendingFuncResponses}) + pendingFuncResponses = nil + } + + for _, m := range msgsRaw { + msg, ok := m.(map[string]any) + if !ok { + continue + } + switch role, _ := msg["role"].(string); role { + case "system": + switch c := msg["content"].(type) { + case string: + systemParts = append(systemParts, c) + case []any: + text, err := geminiSystemTextFromContentParts(c) + if err != nil { + return nil, err + } + systemParts = append(systemParts, text) + } + case "user": + flushPending() + gm, err := geminiMessageFromOpenAI(msg, "user", toolNameByID) + if err != nil { + return nil, err + } + contents = append(contents, gm) + case "assistant": + flushPending() + gm, err := geminiMessageFromOpenAI(msg, "model", toolNameByID) + if err != nil { + return nil, err + } + contents = append(contents, gm) + case "tool": + part, err := geminiFunctionResponsePart(msg, toolNameByID) + if err != nil { + return nil, err + } + pendingFuncResponses = append(pendingFuncResponses, part) + } + } + flushPending() + + if len(systemParts) > 0 { + out["systemInstruction"] = map[string]any{"parts": []any{map[string]any{"text": strings.Join(systemParts, "\n\n")}}} + } + out["contents"] = contents + + genConfig := map[string]any{} + var maxTokens int64 + var hasMaxTokens bool + if v, ok := req["max_completion_tokens"]; ok { + if n, ok2 := toInt64(v); ok2 { + maxTokens, hasMaxTokens = n, true + } + } else if v, ok := req["max_tokens"]; ok { + if n, ok2 := toInt64(v); ok2 { + maxTokens, hasMaxTokens = n, true + } + } + if hasMaxTokens { + genConfig["maxOutputTokens"] = maxTokens + } + if v, ok := req["temperature"]; ok { + genConfig["temperature"] = v + } + if v, ok := req["top_p"]; ok { + genConfig["topP"] = v + } + if v, ok := req["stop"]; ok { + seqs, err := stopSequencesFrom(v) + if err != nil { + return nil, err + } + if len(seqs) > 0 { + genConfig["stopSequences"] = seqs + } + } + if len(genConfig) > 0 { + out["generationConfig"] = genConfig + } + + if toolsRaw, ok := req["tools"].([]any); ok && len(toolsRaw) > 0 { + out["tools"] = geminiToolsFromOpenAI(toolsRaw) + } + if tc, ok := req["tool_choice"]; ok { + if toolConfig := geminiToolConfigFromOpenAI(tc); toolConfig != nil { + out["toolConfig"] = toolConfig + } + } + + return out, nil +} + +// geminiFunctionCall is the subset of a Gemini "functionCall" content +// part's JSON this translator reads. ID is optional on Gemini's side (the +// "Unique identifier of the function call" field, populated only when the +// model chooses to set it); openAIResponseFromGemini and geminiStreamState +// both fall back to a synthesized id when it is empty, since OpenAI's +// tool_calls[].id is required. +type geminiFunctionCall struct { + ID string `json:"id,omitempty"` + Name string `json:"name"` + Args json.RawMessage `json:"args"` +} + +// geminiPart is the subset of one Gemini response content part's JSON this +// translator reads: a text part carries Text, a function-call part carries +// FunctionCall. +type geminiPart struct { + FunctionCall *geminiFunctionCall `json:"functionCall,omitempty"` + Text string `json:"text,omitempty"` +} + +// geminiContent is a Gemini response candidate's "content" object. +type geminiContent struct { + Role string `json:"role"` + Parts []geminiPart `json:"parts"` +} + +// geminiCandidate is one entry of a Gemini response's "candidates" array. +type geminiCandidate struct { + FinishReason string `json:"finishReason"` + Content geminiContent `json:"content"` +} + +// geminiUsageMetadata is the "usageMetadata" object on a Gemini +// GenerateContentResponse, non-streaming or streaming alike. +type geminiUsageMetadata struct { + PromptTokenCount int64 `json:"promptTokenCount"` + CandidatesTokenCount int64 `json:"candidatesTokenCount"` +} + +// geminiGenerateContentResponse is the subset of a Gemini +// GenerateContentResponse this translator reads — the shape of both a +// non-streaming response body and every streaming chunk (each chunk is a +// full GenerateContentResponse increment, per the stream mapping table). +// UsageMetadata is a pointer so the translator can tell "this chunk carries +// no usage" (nil) apart from "this chunk reports zero tokens" (non-nil, +// zero fields) — the stream mapping table's "last usageMetadata wins" rule +// depends on that distinction. +type geminiGenerateContentResponse struct { + ResponseID string `json:"responseId"` + UsageMetadata *geminiUsageMetadata `json:"usageMetadata"` + Candidates []geminiCandidate `json:"candidates"` +} + +// geminiFinishReason maps a Gemini finishReason to an OpenAI finish_reason, +// per the response mapping table: MAX_TOKENS -> "length"; SAFETY, +// RECITATION, BLOCKLIST, PROHIBITED_CONTENT, and SPII -> "content_filter"; +// STOP and any unrecognized value -> "stop", unless hasFunctionCall is set, +// in which case that default branch reports "tool_calls" instead — Gemini +// has no distinct finishReason value for a function-calling turn the way +// Anthropic's "tool_use" stop_reason does, so the caller's own knowledge of +// whether the candidate carried a functionCall part is what this +// translator uses instead. +func geminiFinishReason(reason string, hasFunctionCall bool) string { + switch reason { + case "MAX_TOKENS": + return "length" + case "SAFETY", "RECITATION", "BLOCKLIST", "PROHIBITED_CONTENT", "SPII": + return "content_filter" + default: + if hasFunctionCall { + return "tool_calls" + } + return "stop" + } +} + +// geminiToolCallID returns a synthesized OpenAI tool_calls[].id for a +// Gemini functionCall part that did not carry its own optional "id" field: +// "call_" followed by ordinal, the part's 0-based position among tool +// calls seen so far in this response or stream (per lesson (a), a +// contiguous per-response/per-stream counter, never a raw content-part +// index). +func geminiToolCallID(ordinal int) string { + return "call_" + strconv.Itoa(ordinal) +} + +// openAIResponseFromGemini maps a non-streaming Gemini +// GenerateContentResponse body to an OpenAI chat.completion response, per +// this file's response mapping table. Only candidates[0] is read, matching +// the table. model and created are supplied by the caller — the gateway- +// facing model id and a response timestamp, since Gemini's own response +// carries neither. +func openAIResponseFromGemini(body []byte, model string, created int64) (map[string]any, usage, error) { + var resp geminiGenerateContentResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, usage{}, fmt.Errorf("%w: decode gemini response: %w", errUpstream, err) + } + + var textParts []string + var toolCalls []any + var finishReasonRaw string + if len(resp.Candidates) > 0 { + c := resp.Candidates[0] + finishReasonRaw = c.FinishReason + ordinal := 0 + for _, p := range c.Content.Parts { + switch { + case p.FunctionCall != nil: + args, err := reMarshalToolInput(p.FunctionCall.Args) + if err != nil { + return nil, usage{}, err + } + id := p.FunctionCall.ID + if id == "" { + id = geminiToolCallID(ordinal) + } + toolCalls = append(toolCalls, map[string]any{ + "id": id, + "type": "function", + "function": map[string]any{ + "name": p.FunctionCall.Name, + "arguments": args, + }, + }) + ordinal++ + case p.Text != "": + textParts = append(textParts, p.Text) + } + } + } + + message := map[string]any{"role": "assistant"} + if len(textParts) > 0 { + message["content"] = strings.Join(textParts, "") + } else { + message["content"] = nil + } + if len(toolCalls) > 0 { + message["tool_calls"] = toolCalls + } + + u := usage{prompt: resp.UsageMetadata.safePrompt(), completion: resp.UsageMetadata.safeCompletion()} + + out := map[string]any{ + "id": chatCompletionIDPrefix + resp.ResponseID, + "object": "chat.completion", + "created": created, + "model": model, + "choices": []any{map[string]any{ + "index": 0, + "message": message, + "finish_reason": geminiFinishReason(finishReasonRaw, len(toolCalls) > 0), + }}, + "usage": map[string]any{ + "prompt_tokens": u.prompt, + "completion_tokens": u.completion, + "total_tokens": u.total(), + }, + } + return out, u, nil +} + +// safePrompt returns m.PromptTokenCount, or 0 when m is nil (no +// usageMetadata was present on the response at all). +func (m *geminiUsageMetadata) safePrompt() int64 { + if m == nil { + return 0 + } + return m.PromptTokenCount +} + +// safeCompletion returns m.CandidatesTokenCount, or 0 when m is nil. +func (m *geminiUsageMetadata) safeCompletion() int64 { + if m == nil { + return 0 + } + return m.CandidatesTokenCount +} + +// geminiStreamState is one streaming chat completion's translation state: +// the OpenAI-shaped chunk envelope fields (id, captured from the first +// chunk that carries a non-empty "responseId"; model and created, fixed at +// construction), the running usage total, and the per-stream tool-call +// ordinal counter lesson (a) requires. A caller constructs one per stream +// and calls translate once per upstream sseEvent, in order. +type geminiStreamState struct { + id string + model string + u usage + created int64 + // nextToolOrdinal is the per-stream contiguous tool-call counter lesson + // (a) requires. + nextToolOrdinal int + startedRole bool + // sawFunctionCall latches true the first time any chunk in this stream + // carries a functionCall part, and stays true for the rest of the + // stream. Gemini's finishReason is "STOP" whether the turn ended in + // text or a function call — unlike Anthropic's distinct "tool_use" + // stop_reason — and the finishReason usually arrives on a later, + // parts-less chunk of its own. A per-call local flag would forget a + // functionCall part seen on an earlier chunk by the time the finish + // chunk arrives, so this must be state on the struct, not a translate + // local. + sawFunctionCall bool +} + +// newGeminiStreamState returns a geminiStreamState for one streaming chat +// completion. model is the gateway-facing model id and created the +// response timestamp every emitted chunk carries. +func newGeminiStreamState(model string, created int64) *geminiStreamState { + return &geminiStreamState{model: model, created: created} +} + +// usage returns the prompt/completion token counts captured so far, from +// whichever chunk's usageMetadata was translated most recently. +func (st *geminiStreamState) usage() usage { + return st.u +} + +// chunk builds one OpenAI-shaped "chat.completion.chunk" payload carrying +// delta, and finishReason as its choices[0].finish_reason — nil for the +// JSON null every non-terminal chunk carries. +func (st *geminiStreamState) chunk(delta map[string]any, finishReason *string) []byte { + var fr any + if finishReason != nil { + fr = *finishReason + } + b, err := json.Marshal(map[string]any{ + "id": chatCompletionIDPrefix + st.id, + "object": "chat.completion.chunk", + "created": st.created, + "model": st.model, + "choices": []any{map[string]any{ + "index": 0, + "delta": delta, + "finish_reason": fr, + }}, + }) + if err != nil { + // delta is always built from this file's own map[string]any/string/ + // int64 literals — never a value that can fail to marshal. + panic(fmt.Sprintf("llmgateway: gemini stream chunk failed to marshal: %v", err)) + } + return b +} + +// translate maps one upstream Gemini SSE event to zero or more +// OpenAI-shaped "chat.completion.chunk" payloads, per this file's stream +// mapping table. Every event's data is a full GenerateContentResponse +// increment — there is no "event:" type field to switch on the way +// Anthropic's stream has, so translate always attempts to decode ev.data +// as one. A malformed data payload returns an error wrapping errUpstream +// rather than being silently skipped, per lesson (c): +// +// - The first chunk translate ever emits (across the whole stream) is +// always a {role:"assistant"} delta. +// - A text part becomes a {content: text} delta. +// - A functionCall part becomes a delta.tool_calls[0] chunk carrying the +// 0-based per-stream tool ordinal as its index (lesson (a)), the +// part's own "id" when Gemini set one or else a synthesized one, and +// the re-marshaled "args" as function.arguments. +// - A non-empty finishReason becomes a finish_reason chunk, mapped per +// geminiFinishReason — using st.sawFunctionCall, latched true the first +// time any chunk in this stream carried a functionCall part, since the +// finish chunk is typically its own later, parts-less chunk with no +// functionCall of its own to inspect. +// - Any usageMetadata present on the chunk overwrites the running usage +// total — "last usageMetadata wins", per the stream mapping table. +func (st *geminiStreamState) translate(ev sseEvent) ([][]byte, error) { + var resp geminiGenerateContentResponse + if err := json.Unmarshal(ev.data, &resp); err != nil { + return nil, fmt.Errorf("%w: decode gemini stream chunk: %w", errUpstream, err) + } + if resp.ResponseID != "" { + st.id = resp.ResponseID + } + if resp.UsageMetadata != nil { + st.u = usage{prompt: resp.UsageMetadata.PromptTokenCount, completion: resp.UsageMetadata.CandidatesTokenCount} + } + + var chunks [][]byte + if !st.startedRole { + st.startedRole = true + chunks = append(chunks, st.chunk(map[string]any{"role": "assistant"}, nil)) + } + + if len(resp.Candidates) == 0 { + return chunks, nil + } + c := resp.Candidates[0] + + for _, p := range c.Content.Parts { + switch { + case p.FunctionCall != nil: + st.sawFunctionCall = true + args, err := reMarshalToolInput(p.FunctionCall.Args) + if err != nil { + return chunks, err + } + id := p.FunctionCall.ID + if id == "" { + id = geminiToolCallID(st.nextToolOrdinal) + } + delta := map[string]any{"tool_calls": []any{map[string]any{ + "index": st.nextToolOrdinal, + "id": id, + "type": "function", + "function": map[string]any{ + "name": p.FunctionCall.Name, + "arguments": args, + }, + }}} + st.nextToolOrdinal++ + chunks = append(chunks, st.chunk(delta, nil)) + case p.Text != "": + chunks = append(chunks, st.chunk(map[string]any{"content": p.Text}, nil)) + } + } + + if c.FinishReason != "" { + reason := geminiFinishReason(c.FinishReason, st.sawFunctionCall) + chunks = append(chunks, st.chunk(map[string]any{}, &reason)) + } + + return chunks, nil +} + +// geminiEmbeddingValues is the "embedding" or one "embeddings[]" entry's +// JSON shape on a Gemini embedContent/batchEmbedContents response. +type geminiEmbeddingValues struct { + Values []float64 `json:"values"` +} + +// geminiEmbedContentResponse is the response body shape for Gemini's +// :embedContent endpoint (the OpenAI single-string-input case). +type geminiEmbedContentResponse struct { + Embedding geminiEmbeddingValues `json:"embedding"` +} + +// geminiBatchEmbedContentsResponse is the response body shape for Gemini's +// :batchEmbedContents endpoint (the OpenAI array-input case). +type geminiBatchEmbedContentsResponse struct { + Embeddings []geminiEmbeddingValues `json:"embeddings"` +} + +// geminiEmbedContentRequest builds the :embedContent request body for a +// single OpenAI embeddings input string, per the embeddings mapping table. +func geminiEmbedContentRequest(text string) map[string]any { + return map[string]any{ + "content": map[string]any{"parts": []any{map[string]any{"text": text}}}, + } +} + +// geminiBatchEmbedContentsRequest builds the :batchEmbedContents request +// body for an array of OpenAI embeddings input strings, per the embeddings +// mapping table. model is the bare (no "models/" prefix) Gemini model id; +// each request entry's own "model" field needs the prefix restored, since +// Gemini's per-request model reference is always the full resource name. +func geminiBatchEmbedContentsRequest(model string, texts []string) map[string]any { + requests := make([]any, 0, len(texts)) + for _, text := range texts { + requests = append(requests, map[string]any{ + "model": "models/" + model, + "content": map[string]any{"parts": []any{map[string]any{"text": text}}}, + }) + } + return map[string]any{"requests": requests} +} + +// geminiEmbeddingInputTexts extracts the OpenAI "input" field (a string or +// an array of strings) as a slice of one-or-more strings, and reports +// whether the original input was a single string — selecting the +// :embedContent endpoint — rather than an array, which selects +// :batchEmbedContents. +func geminiEmbeddingInputTexts(input any) (texts []string, single bool, err error) { + switch v := input.(type) { + case string: + return []string{v}, true, nil + case []any: + texts = make([]string, 0, len(v)) + for _, e := range v { + s, ok := e.(string) + if !ok { + return nil, false, &translateError{msg: "input[] entries must be strings"} + } + texts = append(texts, s) + } + return texts, false, nil + default: + return nil, false, &translateError{msg: "input must be a string or an array of strings"} + } +} + +// estimatedEmbeddingTokens estimates a token count for texts the way +// Gemini's embedding endpoints do not: neither :embedContent nor +// :batchEmbedContents reports any usage at all, so this translator +// estimates ceil(total character count / 4) summed across every input +// string, per the embeddings mapping table. +func estimatedEmbeddingTokens(texts []string) int64 { + var chars int64 + for _, t := range texts { + chars += int64(len(t)) + } + return (chars + 3) / 4 +} + +// openAIEmbeddingResponseFromGemini maps a Gemini embedContent or +// batchEmbedContents response body to an OpenAI embeddings list response, +// per the embeddings mapping table. single selects which of Gemini's two +// response shapes body decodes as. model is the gateway-facing model id and +// estimatedTokens the caller's already-computed estimatedEmbeddingTokens +// result — both usage.prompt_tokens and usage.total_tokens report it, and +// the returned usage is marked estimated, since Gemini reported none. +func openAIEmbeddingResponseFromGemini(body []byte, single bool, model string, estimatedTokens int64) (map[string]any, usage, error) { + var vectors [][]float64 + if single { + var resp geminiEmbedContentResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, usage{}, fmt.Errorf("%w: decode gemini embedContent response: %w", errUpstream, err) + } + vectors = [][]float64{resp.Embedding.Values} + } else { + var resp geminiBatchEmbedContentsResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, usage{}, fmt.Errorf("%w: decode gemini batchEmbedContents response: %w", errUpstream, err) + } + vectors = make([][]float64, len(resp.Embeddings)) + for i, e := range resp.Embeddings { + vectors[i] = e.Values + } + } + + data := make([]any, len(vectors)) + for i, v := range vectors { + data[i] = map[string]any{"object": "embedding", "index": i, "embedding": v} + } + + u := usage{prompt: estimatedTokens, estimated: true} + out := map[string]any{ + "object": "list", + "data": data, + "model": model, + "usage": map[string]any{ + "prompt_tokens": u.prompt, + "total_tokens": u.prompt, + }, + } + return out, u, nil +} diff --git a/translate_gemini_test.go b/translate_gemini_test.go new file mode 100644 index 0000000..4df654b --- /dev/null +++ b/translate_gemini_test.go @@ -0,0 +1,442 @@ +package traefikllmgateway + +import ( + "encoding/json" + "errors" + "os" + "path/filepath" + "reflect" + "sort" + "strings" + "testing" +) + +// geminiGoldenTestdataDir is where every fixture this file walks lives. +const geminiGoldenTestdataDir = "testdata/gemini" + +// geminiGoldenGatewayModel and geminiGoldenCreated are the fixed model/ +// created values every response- and stream-translation golden fixture's +// "_want.json" bakes in — openAIResponseFromGemini and +// newGeminiStreamState both take these as caller-supplied parameters +// (Gemini's own wire formats carry neither), so the golden tests must +// supply the same fixed values the "_want.json" files were written +// against. +const ( + geminiGoldenGatewayModel = "gw-gemini-model" + geminiGoldenCreated = int64(1734000000) +) + +// geminiReadGoldenJSON reads and json.Unmarshals name (relative to +// testdata/gemini) into a fresh map[string]any. +func geminiReadGoldenJSON(t *testing.T, name string) map[string]any { + t.Helper() + b, err := os.ReadFile(filepath.Join(geminiGoldenTestdataDir, name)) + if err != nil { + t.Fatalf("read %s: %v", name, err) + } + var m map[string]any + if err := json.Unmarshal(b, &m); err != nil { + t.Fatalf("decode %s: %v", name, err) + } + return m +} + +// geminiAssertTranslateError asserts err is a *translateError matching want. +func geminiAssertTranslateError(t *testing.T, err error, want wantErrFixture) { + t.Helper() + if err == nil { + t.Fatalf("want a *translateError containing %q, got nil", want.ErrContains) + } + var terr *translateError + if !errors.As(err, &terr) { + t.Fatalf("err = %v (%T), want *translateError", err, err) + } + if !strings.Contains(terr.Error(), want.ErrContains) { + t.Errorf("error message = %q, want it to contain %q", terr.Error(), want.ErrContains) + } + if terr.notSupported != want.NotSupported { + t.Errorf("notSupported = %v, want %v", terr.notSupported, want.NotSupported) + } +} + +// TestGeminiRequestGolden walks every "chat_*_in.json" fixture under +// testdata/gemini, translating it via geminiRequestFromOpenAI, and asserts +// the result against the sibling "_want.json" fixture (a successful +// translation) or "_wanterr.json" fixture (a *translateError), whichever is +// present — exactly one of the two exists per fixture. +func TestGeminiRequestGolden(t *testing.T) { + entries, err := os.ReadDir(geminiGoldenTestdataDir) + if err != nil { + t.Fatalf("read testdata dir: %v", err) + } + + var names []string + for _, e := range entries { + if n, ok := strings.CutSuffix(e.Name(), "_in.json"); ok && strings.HasPrefix(n, "chat_") { + names = append(names, n) + } + } + sort.Strings(names) + if len(names) == 0 { + t.Fatalf("no chat_*_in.json fixtures found under %s", geminiGoldenTestdataDir) + } + + for _, name := range names { + t.Run(name, func(t *testing.T) { + in := geminiReadGoldenJSON(t, name+"_in.json") + + got, err := geminiRequestFromOpenAI(in) + + wantErrPath := filepath.Join(geminiGoldenTestdataDir, name+"_wanterr.json") + if _, statErr := os.Stat(wantErrPath); statErr == nil { + var want wantErrFixture + b, rerr := os.ReadFile(wantErrPath) + if rerr != nil { + t.Fatalf("read %s: %v", wantErrPath, rerr) + } + if uerr := json.Unmarshal(b, &want); uerr != nil { + t.Fatalf("decode %s: %v", wantErrPath, uerr) + } + geminiAssertTranslateError(t, err, want) + return + } + + if err != nil { + t.Fatalf("geminiRequestFromOpenAI: %v", err) + } + want := geminiReadGoldenJSON(t, name+"_want.json") + gotN := normalizeJSON(t, got) + wantN := normalizeJSON(t, want) + if !reflect.DeepEqual(gotN, wantN) { + gotB, _ := json.MarshalIndent(gotN, "", " ") + wantB, _ := json.MarshalIndent(wantN, "", " ") + t.Errorf("geminiRequestFromOpenAI(%s) mismatch:\ngot:\n%s\nwant:\n%s", name, gotB, wantB) + } + }) + } +} + +// TestGeminiResponseGolden walks every "resp_*_in.json" fixture, +// translating it via openAIResponseFromGemini with the fixed +// geminiGoldenGatewayModel/geminiGoldenCreated, and asserts the result +// against the sibling "_want.json" fixture. +func TestGeminiResponseGolden(t *testing.T) { + entries, err := os.ReadDir(geminiGoldenTestdataDir) + if err != nil { + t.Fatalf("read testdata dir: %v", err) + } + + var names []string + for _, e := range entries { + if n, ok := strings.CutSuffix(e.Name(), "_in.json"); ok && strings.HasPrefix(n, "resp_") { + names = append(names, n) + } + } + sort.Strings(names) + if len(names) == 0 { + t.Fatalf("no resp_*_in.json fixtures found under %s", geminiGoldenTestdataDir) + } + + for _, name := range names { + t.Run(name, func(t *testing.T) { + body, err := os.ReadFile(filepath.Join(geminiGoldenTestdataDir, name+"_in.json")) + if err != nil { + t.Fatalf("read %s_in.json: %v", name, err) + } + + got, _, err := openAIResponseFromGemini(body, geminiGoldenGatewayModel, geminiGoldenCreated) + if err != nil { + t.Fatalf("openAIResponseFromGemini: %v", err) + } + + want := geminiReadGoldenJSON(t, name+"_want.json") + gotN := normalizeJSON(t, got) + wantN := normalizeJSON(t, want) + if !reflect.DeepEqual(gotN, wantN) { + gotB, _ := json.MarshalIndent(gotN, "", " ") + wantB, _ := json.MarshalIndent(wantN, "", " ") + t.Errorf("openAIResponseFromGemini(%s) mismatch:\ngot:\n%s\nwant:\n%s", name, gotB, wantB) + } + }) + } +} + +// geminiRunStreamGolden translates every event in +// testdata/gemini/{name}.sse through a fresh geminiStreamState, in order, +// and asserts the flattened list of emitted chunks against +// "{name}_want.json" and the final usage() against "{name}_usage.json". +func geminiRunStreamGolden(t *testing.T, name string) { + t.Helper() + + f, err := os.Open(filepath.Join(geminiGoldenTestdataDir, name+".sse")) + if err != nil { + t.Fatalf("open %s.sse: %v", name, err) + } + defer f.Close() //nolint:errcheck + + st := newGeminiStreamState(geminiGoldenGatewayModel, geminiGoldenCreated) + var got []any + terminalErr := readSSE(f, func(ev sseEvent) error { + chunks, terr := st.translate(ev) + for _, c := range chunks { + var v any + if uerr := json.Unmarshal(c, &v); uerr != nil { + t.Fatalf("translate(%s) produced invalid JSON: %v (%s)", ev.event, uerr, c) + } + got = append(got, v) + } + return terr + }) + if terminalErr != nil { + t.Fatalf("readSSE/translate(%s): %v", name, terminalErr) + } + + wantBytes, err := os.ReadFile(filepath.Join(geminiGoldenTestdataDir, name+"_want.json")) + if err != nil { + t.Fatalf("read %s_want.json: %v", name, err) + } + var want []any + if uerr := json.Unmarshal(wantBytes, &want); uerr != nil { + t.Fatalf("decode %s_want.json: %v", name, uerr) + } + + if !reflect.DeepEqual(got, want) { + gotB, _ := json.MarshalIndent(got, "", " ") + wantB, _ := json.MarshalIndent(want, "", " ") + t.Errorf("translate(%s) chunk list mismatch:\ngot:\n%s\nwant:\n%s", name, gotB, wantB) + } + + usageBytes, err := os.ReadFile(filepath.Join(geminiGoldenTestdataDir, name+"_usage.json")) + if err != nil { + t.Fatalf("read %s_usage.json: %v", name, err) + } + var wantUsage streamUsageFixture + if err := json.Unmarshal(usageBytes, &wantUsage); err != nil { + t.Fatalf("decode %s_usage.json: %v", name, err) + } + if gotUsage := st.usage(); gotUsage.prompt != wantUsage.Prompt || gotUsage.completion != wantUsage.Completion { + t.Errorf("usage() = %+v, want {prompt:%d completion:%d}", gotUsage, wantUsage.Prompt, wantUsage.Completion) + } +} + +// TestGeminiStreamGolden_Text covers a text-only stream: a role chunk, two +// content deltas across two chunks, and a finish chunk on a third, +// parts-less chunk that also carries the final usageMetadata. +func TestGeminiStreamGolden_Text(t *testing.T) { + geminiRunStreamGolden(t, "stream_text") +} + +// TestGeminiStreamGolden_TextThenFuncall covers a text delta, then a +// functionCall part with no upstream "id" (proving the ordinal-synthesized +// "call_0" id, per lesson (a)), then a finish chunk on its own later, +// parts-less chunk — proving st.sawFunctionCall correctly latches across +// chunks so the finish_reason still maps to "tool_calls" rather than +// "stop" even though the finish chunk itself carries no functionCall part. +func TestGeminiStreamGolden_TextThenFuncall(t *testing.T) { + geminiRunStreamGolden(t, "stream_text_then_funcall") +} + +// TestGeminiStreamState_MalformedEventReturnsError proves a malformed +// (invalid-JSON) SSE data payload returns an error wrapping errUpstream, +// rather than being swallowed as (nil, nil) — lesson (c). +func TestGeminiStreamState_MalformedEventReturnsError(t *testing.T) { + st := newGeminiStreamState(geminiGoldenGatewayModel, geminiGoldenCreated) + chunks, err := st.translate(sseEvent{data: []byte(`{"candidates":[`)}) + if err == nil { + t.Fatalf("translate(corrupt data): want error, got nil (chunks=%v)", chunks) + } + if !errors.Is(err, errUpstream) { + t.Errorf("translate(corrupt data): err = %v, want errors.Is(err, errUpstream)", err) + } +} + +// TestGeminiStreamState_UsageMetadataOnlyOverwritesWhenPresent proves "last +// usageMetadata wins" only among chunks that actually carry the field — a +// later chunk with no usageMetadata at all must not reset a usage total an +// earlier chunk already captured. +func TestGeminiStreamState_UsageMetadataOnlyOverwritesWhenPresent(t *testing.T) { + st := newGeminiStreamState(geminiGoldenGatewayModel, geminiGoldenCreated) + + if _, err := st.translate(sseEvent{data: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"hi"}]}}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":2}}`)}); err != nil { + t.Fatalf("translate: %v", err) + } + if got := st.usage(); got.prompt != 7 || got.completion != 2 { + t.Fatalf("usage() after first chunk = %+v, want {prompt:7 completion:2}", got) + } + + if _, err := st.translate(sseEvent{data: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":" there"}]}}]}`)}); err != nil { + t.Fatalf("translate: %v", err) + } + if got := st.usage(); got.prompt != 7 || got.completion != 2 { + t.Errorf("usage() after a usageMetadata-less chunk = %+v, want unchanged {prompt:7 completion:2}", got) + } +} + +// TestGeminiRequestFromOpenAI_MaxCompletionTokensPrecedence proves +// max_completion_tokens takes precedence over max_tokens when an OpenAI +// request sends both — the newer field wins, mirroring the anthropic +// translator's rule. +func TestGeminiRequestFromOpenAI_MaxCompletionTokensPrecedence(t *testing.T) { + req := map[string]any{ + "model": "gemini-2.5-pro", + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + "max_tokens": float64(100), + "max_completion_tokens": float64(200), + } + got, err := geminiRequestFromOpenAI(req) + if err != nil { + t.Fatalf("geminiRequestFromOpenAI: %v", err) + } + genConfig, ok := got["generationConfig"].(map[string]any) + if !ok { + t.Fatalf("generationConfig = %v, want a map", got["generationConfig"]) + } + if genConfig["maxOutputTokens"] != int64(200) { + t.Errorf("maxOutputTokens = %v, want 200 (max_completion_tokens takes precedence)", genConfig["maxOutputTokens"]) + } +} + +// TestGeminiRequestFromOpenAI_LogitBiasEmptyMapAllowed proves a +// present-but-empty logit_bias does not 400 — lesson (i): reject only when +// truthy/non-empty. +func TestGeminiRequestFromOpenAI_LogitBiasEmptyMapAllowed(t *testing.T) { + req := map[string]any{ + "model": "gemini-2.5-pro", + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + "logit_bias": map[string]any{}, + } + if _, err := geminiRequestFromOpenAI(req); err != nil { + t.Errorf("geminiRequestFromOpenAI with empty logit_bias: %v, want no error", err) + } +} + +// TestGeminiRequestFromOpenAI_UnknownToolCallIDErrors proves a role:"tool" +// message whose tool_call_id was never seen on a preceding assistant +// message's tool_calls[] is a *translateError, per the request mapping +// table's role:tool row — Gemini's functionResponse needs the name that +// map resolves, and an unresolvable id must not silently produce an empty +// name. +func TestGeminiRequestFromOpenAI_UnknownToolCallIDErrors(t *testing.T) { + req := map[string]any{ + "model": "gemini-2.5-pro", + "messages": []any{ + map[string]any{"role": "user", "content": "hi"}, + map[string]any{"role": "tool", "tool_call_id": "call_never_seen", "content": "42"}, + }, + } + _, err := geminiRequestFromOpenAI(req) + var terr *translateError + if !errors.As(err, &terr) { + t.Fatalf("err = %v (%T), want *translateError", err, err) + } + if !strings.Contains(terr.Error(), "call_never_seen") { + t.Errorf("error message = %q, want it to mention the unknown id", terr.Error()) + } +} + +// TestGeminiEmbeddings_Single proves a single-string OpenAI embeddings +// input builds Gemini's :embedContent request body and translates its +// response back to an OpenAI embeddings list response, with usage +// estimated (ceil(chars/4)) and marked estimated — per the embeddings +// mapping table. +func TestGeminiEmbeddings_Single(t *testing.T) { + in := geminiReadGoldenJSON(t, "embed_single_in.json") + + texts, single, err := geminiEmbeddingInputTexts(in["input"]) + if err != nil { + t.Fatalf("geminiEmbeddingInputTexts: %v", err) + } + if !single { + t.Fatalf("single = false, want true for a string input") + } + + gotReq := geminiEmbedContentRequest(texts[0]) + wantReq := geminiReadGoldenJSON(t, "embed_single_want.json") + if !reflect.DeepEqual(normalizeJSON(t, gotReq), normalizeJSON(t, wantReq)) { + t.Errorf("geminiEmbedContentRequest mismatch: got %#v, want %#v", gotReq, wantReq) + } + + respBody, err := os.ReadFile(filepath.Join(geminiGoldenTestdataDir, "embed_single_resp_in.json")) + if err != nil { + t.Fatalf("read embed_single_resp_in.json: %v", err) + } + estTokens := estimatedEmbeddingTokens(texts) + if estTokens != 3 { + t.Fatalf("estimatedEmbeddingTokens(%q) = %d, want 3", texts, estTokens) + } + + gotResp, u, err := openAIEmbeddingResponseFromGemini(respBody, single, "gw-embed-model", estTokens) + if err != nil { + t.Fatalf("openAIEmbeddingResponseFromGemini: %v", err) + } + if !u.estimated { + t.Errorf("usage.estimated = false, want true") + } + if u.prompt != 3 { + t.Errorf("usage.prompt = %d, want 3", u.prompt) + } + wantResp := geminiReadGoldenJSON(t, "embed_single_resp_want.json") + if !reflect.DeepEqual(normalizeJSON(t, gotResp), normalizeJSON(t, wantResp)) { + gotB, _ := json.MarshalIndent(gotResp, "", " ") + wantB, _ := json.MarshalIndent(wantResp, "", " ") + t.Errorf("openAIEmbeddingResponseFromGemini mismatch:\ngot:\n%s\nwant:\n%s", gotB, wantB) + } +} + +// TestGeminiEmbeddings_Batch mirrors TestGeminiEmbeddings_Single for an +// array-of-strings OpenAI embeddings input, which selects Gemini's +// :batchEmbedContents endpoint and restores the "models/" prefix on each +// request entry's own "model" field, per the embeddings mapping table. +func TestGeminiEmbeddings_Batch(t *testing.T) { + in := geminiReadGoldenJSON(t, "embed_batch_in.json") + rawModel, _ := in["model"].(string) + + texts, single, err := geminiEmbeddingInputTexts(in["input"]) + if err != nil { + t.Fatalf("geminiEmbeddingInputTexts: %v", err) + } + if single { + t.Fatalf("single = true, want false for an array input") + } + + gotReq := geminiBatchEmbedContentsRequest(rawModel, texts) + wantReq := geminiReadGoldenJSON(t, "embed_batch_want.json") + if !reflect.DeepEqual(normalizeJSON(t, gotReq), normalizeJSON(t, wantReq)) { + gotB, _ := json.MarshalIndent(gotReq, "", " ") + wantB, _ := json.MarshalIndent(wantReq, "", " ") + t.Errorf("geminiBatchEmbedContentsRequest mismatch:\ngot:\n%s\nwant:\n%s", gotB, wantB) + } + + respBody, err := os.ReadFile(filepath.Join(geminiGoldenTestdataDir, "embed_batch_resp_in.json")) + if err != nil { + t.Fatalf("read embed_batch_resp_in.json: %v", err) + } + estTokens := estimatedEmbeddingTokens(texts) + if estTokens != 3 { + t.Fatalf("estimatedEmbeddingTokens(%q) = %d, want 3", texts, estTokens) + } + + gotResp, u, err := openAIEmbeddingResponseFromGemini(respBody, single, "gw-embed-model", estTokens) + if err != nil { + t.Fatalf("openAIEmbeddingResponseFromGemini: %v", err) + } + if !u.estimated { + t.Errorf("usage.estimated = false, want true") + } + wantResp := geminiReadGoldenJSON(t, "embed_batch_resp_want.json") + if !reflect.DeepEqual(normalizeJSON(t, gotResp), normalizeJSON(t, wantResp)) { + gotB, _ := json.MarshalIndent(gotResp, "", " ") + wantB, _ := json.MarshalIndent(wantResp, "", " ") + t.Errorf("openAIEmbeddingResponseFromGemini mismatch:\ngot:\n%s\nwant:\n%s", gotB, wantB) + } +} + +// TestGeminiEmbeddingInputTexts_RejectsNonStringArrayEntry proves an +// input[] array with a non-string entry is a *translateError. +func TestGeminiEmbeddingInputTexts_RejectsNonStringArrayEntry(t *testing.T) { + _, _, err := geminiEmbeddingInputTexts([]any{"ok", float64(5)}) + var terr *translateError + if !errors.As(err, &terr) { + t.Fatalf("err = %v (%T), want *translateError", err, err) + } +}