From 93b289c1c89f7025cfe8ed729569b2487be0d864 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 18:48:02 +0100 Subject: [PATCH] feat: unified image and audio endpoints with gemini image translation --- README.md | 44 ++ llmgateway.go | 18 + provider_anthropic.go | 20 + provider_gemini.go | 66 ++ provider_openai.go | 98 +++ providers.go | 63 +- registry_test.go | 9 + routes_media.go | 244 +++++++ routes_media_test.go | 637 ++++++++++++++++++ testdata/gemini_images/req_basic_in.json | 1 + testdata/gemini_images/req_basic_want.json | 1 + testdata/gemini_images/req_n_zero_in.json | 1 + testdata/gemini_images/req_n_zero_want.json | 1 + .../gemini_images/req_quality_falsy_in.json | 1 + .../gemini_images/req_quality_falsy_want.json | 1 + .../gemini_images/req_quality_truthy_in.json | 1 + .../req_quality_truthy_wanterr.json | 1 + .../req_response_format_b64_in.json | 1 + .../req_response_format_b64_want.json | 1 + .../req_response_format_url_in.json | 1 + .../req_response_format_url_wanterr.json | 1 + testdata/gemini_images/req_size_16x9_in.json | 1 + .../gemini_images/req_size_16x9_want.json | 1 + testdata/gemini_images/req_size_512_in.json | 1 + testdata/gemini_images/req_size_512_want.json | 1 + testdata/gemini_images/req_size_9x16_in.json | 1 + .../gemini_images/req_size_9x16_want.json | 1 + .../gemini_images/req_size_absent_in.json | 1 + .../gemini_images/req_size_absent_want.json | 1 + .../gemini_images/req_size_unknown_in.json | 1 + .../req_size_unknown_wanterr.json | 1 + .../gemini_images/req_style_truthy_in.json | 1 + .../req_style_truthy_wanterr.json | 1 + testdata/gemini_images/resp_basic_in.json | 1 + testdata/gemini_images/resp_basic_want.json | 1 + testdata/gemini_images/resp_empty_in.json | 1 + testdata/gemini_images/resp_empty_want.json | 1 + testdata/gemini_images/resp_multi_in.json | 1 + testdata/gemini_images/resp_multi_want.json | 1 + translate_gemini_images.go | 121 ++++ translate_gemini_images_test.go | 159 +++++ 41 files changed, 1507 insertions(+), 2 deletions(-) create mode 100644 routes_media.go create mode 100644 routes_media_test.go create mode 100644 testdata/gemini_images/req_basic_in.json create mode 100644 testdata/gemini_images/req_basic_want.json create mode 100644 testdata/gemini_images/req_n_zero_in.json create mode 100644 testdata/gemini_images/req_n_zero_want.json create mode 100644 testdata/gemini_images/req_quality_falsy_in.json create mode 100644 testdata/gemini_images/req_quality_falsy_want.json create mode 100644 testdata/gemini_images/req_quality_truthy_in.json create mode 100644 testdata/gemini_images/req_quality_truthy_wanterr.json create mode 100644 testdata/gemini_images/req_response_format_b64_in.json create mode 100644 testdata/gemini_images/req_response_format_b64_want.json create mode 100644 testdata/gemini_images/req_response_format_url_in.json create mode 100644 testdata/gemini_images/req_response_format_url_wanterr.json create mode 100644 testdata/gemini_images/req_size_16x9_in.json create mode 100644 testdata/gemini_images/req_size_16x9_want.json create mode 100644 testdata/gemini_images/req_size_512_in.json create mode 100644 testdata/gemini_images/req_size_512_want.json create mode 100644 testdata/gemini_images/req_size_9x16_in.json create mode 100644 testdata/gemini_images/req_size_9x16_want.json create mode 100644 testdata/gemini_images/req_size_absent_in.json create mode 100644 testdata/gemini_images/req_size_absent_want.json create mode 100644 testdata/gemini_images/req_size_unknown_in.json create mode 100644 testdata/gemini_images/req_size_unknown_wanterr.json create mode 100644 testdata/gemini_images/req_style_truthy_in.json create mode 100644 testdata/gemini_images/req_style_truthy_wanterr.json create mode 100644 testdata/gemini_images/resp_basic_in.json create mode 100644 testdata/gemini_images/resp_basic_want.json create mode 100644 testdata/gemini_images/resp_empty_in.json create mode 100644 testdata/gemini_images/resp_empty_want.json create mode 100644 testdata/gemini_images/resp_multi_in.json create mode 100644 testdata/gemini_images/resp_multi_want.json create mode 100644 translate_gemini_images.go create mode 100644 translate_gemini_images_test.go diff --git a/README.md b/README.md index e158ec4..d39da71 100644 --- a/README.md +++ b/README.md @@ -433,6 +433,50 @@ request. | Unsupported parameters | A semantically meaningful field the target provider cannot express (e.g. `n > 1`, `logit_bias`, `logprobs` on Anthropic) is a **400**, never silently dropped. | Not applicable — the request reaches the provider exactly as sent. | | Non-2xx upstream response | Wrapped in the gateway's own error envelope, with the provider's own body embedded under `error.upstream`. | Forwarded to the client exactly as the upstream sent it — status, headers, and body. | +## Image and audio endpoints + +`POST /v1/images/generations`, `POST /v1/audio/speech`, and +`POST /v1/audio/transcriptions` follow the same pipeline as +`/v1/chat/completions` — auth, model routing, group authorization, limits, +adapter call — minus response caching and token/cost accounting (below). + +| | `openai`-type | `gemini` | `anthropic` | +|---|---|---|---| +| `images/generations` | Native forward. | Translated to Imagen's `:predict` (below). | **501**. | +| `audio/speech` | Native forward; the binary response is streamed to the client as it arrives. | **501** — no OpenAI-compatible text-to-speech endpoint. | **501**. | +| `audio/transcriptions` | Native forward; the client's original multipart body is replayed upstream unchanged. | **501** — no OpenAI-compatible speech-to-text endpoint. | **501**. | + +- **Model routing**: `images/generations` and `audio/speech` read `model` + from the JSON request body. `audio/transcriptions` reads it from the + multipart request's `model` form field instead — the field can appear + in any position among the request's parts. +- **Gemini image translation**: `{model, prompt, n, size, + response_format}` maps to Imagen's `{instances:[{prompt}], parameters: + {sampleCount, aspectRatio}}`. `size` maps to `aspectRatio`: + `1024x1024`/`512x512` → `1:1`, `1792x1024` → `16:9`, `1024x1792` → + `9:16`; an absent `size` defaults to `1:1`; any other value is a + **400**. `response_format: "url"` is a **400** — the gateway stores + nothing to host a URL from. A truthy `quality` or `style` is a **400**, + the same convention chat's translators use for an unsupported field. + Imagen's `predictions[].bytesBase64Encoded` maps back to OpenAI's + `data[].b64_json`. +- **Never cached**: image and audio responses are excluded from the + response cache (`cache.enabled`) regardless of its configuration. +- **Accounting**: request counters (`requestsPerMinute`/`requestsPerDay`) + increment the same as any other route. Token counters never move, and + cost is always recorded as 0 — images and audio are not + cost-accounted. +- **Body limits**: `images/generations` and `audio/speech` share the + unified routes' 10MiB JSON cap. `audio/transcriptions` enforces the + same 10MiB cap explicitly, returning **413** over it — unlike an + oversized JSON body, which is silently truncated into a parse failure, + a truncated multipart body can still parse a valid prefix and would + forward cut-off audio upstream if it were not rejected outright. +- **Retry** (above) applies to all three the same as chat/embeddings; a + streaming `audio/speech` response is retried only before the first + response byte reaches the client, the same zero-bytes rule as chat's + streaming path. + ## MCP and A2A - **Registry** (group-filtered, sorted by name): diff --git a/llmgateway.go b/llmgateway.go index fbc2ae1..eac05a0 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -365,6 +365,24 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } + if r.Method == http.MethodPost && (r.URL.Path == imagesGenerationsPath || r.URL.Path == audioSpeechPath || r.URL.Path == audioTranscriptionsPath) { + u, grp, ok := g.auth.identify(r) + g.logAuthEvent(ok, authEventUserName(u), r) + if !ok { + writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key") + return + } + switch r.URL.Path { + case imagesGenerationsPath: + g.handleImagesGenerations(sw, r, u, grp) + case audioSpeechPath: + g.handleAudioSpeech(sw, r, u, grp) + case audioTranscriptionsPath: + g.handleAudioTranscriptions(sw, r, u, grp) + } + return + } + if r.Method == http.MethodGet && (r.URL.Path == "/v1/mcp/servers" || r.URL.Path == "/v1/agents") { u, grp, ok := g.auth.identify(r) g.logAuthEvent(ok, authEventUserName(u), r) diff --git a/provider_anthropic.go b/provider_anthropic.go index c3571a2..ea403a6 100644 --- a/provider_anthropic.go +++ b/provider_anthropic.go @@ -198,6 +198,26 @@ func (a *anthropicAdapter) embeddings(_ context.Context, _ http.ResponseWriter, return usage{}, &translateError{msg: "embeddings not supported for anthropic models", notSupported: true} } +// imagesGeneration implements providerAdapter: Anthropic has no images API +// (spec §3, v0.2), so this never makes an upstream request. +func (a *anthropicAdapter) imagesGeneration(_ context.Context, _ http.ResponseWriter, _ map[string]any) (usage, error) { + return usage{}, &translateError{msg: "image generation not supported for anthropic models", notSupported: true} +} + +// audioSpeech implements providerAdapter: Anthropic has no +// OpenAI-compatible text-to-speech API (spec §3, v0.2), so this never +// makes an upstream request. +func (a *anthropicAdapter) audioSpeech(_ context.Context, _ http.ResponseWriter, _ []byte, _ string) (usage, error) { + return usage{}, &translateError{msg: "audio speech not supported for anthropic models", notSupported: true} +} + +// audioTranscription implements providerAdapter: Anthropic has no +// OpenAI-compatible speech-to-text API (spec §3, v0.2), so this never +// makes an upstream request. +func (a *anthropicAdapter) audioTranscription(_ context.Context, _ http.ResponseWriter, _ []byte, _ string) (usage, error) { + return usage{}, &translateError{msg: "audio transcription not supported for anthropic models", notSupported: true} +} + // listModels implements providerAdapter: GET {base}/v1/models, parsing // the same "data": [{"id": ...}] shape OpenAI's endpoint uses (Anthropic's // 2026 Models API responds in the same shape). A non-2xx response — diff --git a/provider_gemini.go b/provider_gemini.go index ad8a80d..cd51e45 100644 --- a/provider_gemini.go +++ b/provider_gemini.go @@ -265,6 +265,72 @@ func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, r return u, nil } +// imagesGeneration implements providerAdapter: translates req to an Imagen +// :predict request via geminiImagesRequestFromOpenAI, POSTs it, and +// translates the response back to OpenAI's images.generations shape +// (translate_gemini_images.go). Images never stream. Like +// chatCompletion/embeddings, the model id is read from req (already +// rewritten to the provider's own model id by the caller) and drives the +// URL directly; an OpenAI images.generations response carries no "model" +// field at all, so there is no alias to echo back the way chat/embeddings +// do. +func (a *geminiAdapter) imagesGeneration(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + model := geminiModelFromRequest(req) + delete(req, gatewayAliasKey) + + body, err := geminiImagesRequestFromOpenAI(req) + if err != nil { + return usage{}, err + } + + escapedModel := url.PathEscape(model) + endpoint := a.baseURL + geminiAPIPrefix + escapedModel + ":predict" + + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, a.requestHeaders(true), body, a.retry) + 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) + } + + raw, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) + if err != nil { + return usage{}, fmt.Errorf("%w: read response body: %w", errUpstream, err) + } + + out, err := openAIImagesResponseFromGemini(raw, time.Now().Unix()) + if err != nil { + return usage{}, err + } + + b, err := json.Marshal(out) + if err != nil { + return usage{}, fmt.Errorf("%w: marshal translated images response: %w", errUpstream, err) + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(b) + return usage{}, nil +} + +// audioSpeech implements providerAdapter: Gemini has no OpenAI-compatible +// text-to-speech API (spec §3, v0.2), so this never makes an upstream +// request. +func (a *geminiAdapter) audioSpeech(_ context.Context, _ http.ResponseWriter, _ []byte, _ string) (usage, error) { + return usage{}, &translateError{msg: "audio speech not supported for gemini models", notSupported: true} +} + +// audioTranscription implements providerAdapter: Gemini has no +// OpenAI-compatible speech-to-text API (spec §3, v0.2), so this never +// makes an upstream request. +func (a *geminiAdapter) audioTranscription(_ context.Context, _ http.ResponseWriter, _ []byte, _ string) (usage, error) { + return usage{}, &translateError{msg: "audio transcription not supported for gemini models", notSupported: true} +} + // 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"). diff --git a/provider_openai.go b/provider_openai.go index 7062baa..e5eb7e3 100644 --- a/provider_openai.go +++ b/provider_openai.go @@ -254,6 +254,104 @@ func (a *openaiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, cli return u, nil } +// imagesGeneration implements providerAdapter: a native forward to +// {base}/v1/images/generations. req has already had its "model" field +// rewritten to the provider's own model id by the caller (routes_media.go), +// matching chatCompletion/embeddings' contract. Images are never +// cost-accounted (spec §3, v0.2), so this always returns a zero usage, +// regardless of what the upstream response reports. +func (a *openaiAdapter) imagesGeneration(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/images/generations", a.requestHeaders(true), req, a.retry) + 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) + } + return usage{}, a.forwardMediaBody(w, resp) +} + +// audioSpeech implements providerAdapter: a native forward to +// {base}/v1/audio/speech. body is the caller's already re-marshaled JSON +// request (routes_media.go rewrites "model" to the provider's own model id +// before calling in). The response is binary audio, streamed to w via a +// flushWriter as it arrives rather than buffered in full: once this loop +// starts copying bytes, the response has already committed to the client +// and can no longer be retried — matching spec §1's "retry only while ZERO +// response bytes have reached the client" (the retry itself lives entirely +// inside upstreamRawBytes, above this point). +func (a *openaiAdapter) audioSpeech(ctx context.Context, w http.ResponseWriter, body []byte, contentType string) (usage, error) { + hdr := a.requestHeaders(false) + if contentType != "" { + hdr.Set("Content-Type", contentType) + } + resp, err := upstreamRawBytes(ctx, a.client, http.MethodPost, a.baseURL+"/v1/audio/speech", hdr, body, a.retry) + 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 ct := resp.Header.Get("Content-Type"); ct != "" { + w.Header().Set("Content-Type", ct) + } + w.WriteHeader(resp.StatusCode) + fw := newFlushWriter(w) + if _, err := io.Copy(fw, io.LimitReader(resp.Body, maxResponseBytes)); err != nil { + return usage{}, fmt.Errorf("%w: stream audio response: %w", errUpstream, err) + } + return usage{}, nil +} + +// audioTranscription implements providerAdapter: a native forward to +// {base}/v1/audio/transcriptions. body is the client's original raw +// multipart request, replayed unchanged (routes_media.go never rewrites +// it), and contentType is the client's original Content-Type header, +// forwarded verbatim so the upstream sees the same multipart boundary the +// client used. The response (JSON or plain text, depending on the +// client's requested response_format) is copied through unchanged, the +// same non-streaming path images.generations uses. +func (a *openaiAdapter) audioTranscription(ctx context.Context, w http.ResponseWriter, body []byte, contentType string) (usage, error) { + hdr := a.requestHeaders(false) + if contentType != "" { + hdr.Set("Content-Type", contentType) + } + resp, err := upstreamRawBytes(ctx, a.client, http.MethodPost, a.baseURL+"/v1/audio/transcriptions", hdr, body, a.retry) + 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) + } + return usage{}, a.forwardMediaBody(w, resp) +} + +// forwardMediaBody reads resp's full body (capped at maxResponseBytes) and +// writes resp's Content-Type, status code, and body verbatim to w — the +// copy-through path images.generations and audio-transcriptions share. +// Unlike forwardJSON, it never extracts usage: images/audio are never +// cost-accounted (spec §3, v0.2), so there is nothing to parse the body +// for. +func (a *openaiAdapter) forwardMediaBody(w http.ResponseWriter, resp *http.Response) error { + body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) + if err != nil { + return fmt.Errorf("%w: read response body: %w", errUpstream, err) + } + if ct := resp.Header.Get("Content-Type"); ct != "" { + w.Header().Set("Content-Type", ct) + } + w.WriteHeader(resp.StatusCode) + _, _ = w.Write(body) + return nil +} + // modelsPayload is the "data": [{"id": ...}, ...] shape OpenAI's // /v1/models endpoint returns. type modelsPayload struct { diff --git a/providers.go b/providers.go index 3873fbf..3c5471b 100644 --- a/providers.go +++ b/providers.go @@ -127,6 +127,31 @@ type providerAdapter interface { embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) // listModels returns the provider's available model ids, for discovery. listModels(ctx context.Context) ([]string, error) + // imagesGeneration proxies req to the provider's image-generation + // endpoint and writes the response to w (spec §3, v0.2). req has + // already had its "model" field rewritten to the provider's own model + // id by the caller, matching chatCompletion/embeddings' contract. + // Images are never cached (spec §2) and never cost-accounted (spec + // §3) — every implementation returns a zero usage regardless of what + // its upstream reports. Anthropic has no images API and always + // returns a *translateError with notSupported set, mapped to HTTP 501 + // by the same handleAdapterError chatCompletion/embeddings errors + // already go through. + imagesGeneration(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) + // audioSpeech proxies body — a JSON request the caller has already + // re-marshaled with "model" rewritten to the provider's own model id + // — to the provider's text-to-speech endpoint, streaming the (binary) + // response to w as it arrives rather than buffering it in full (spec + // §3, v0.2). Gemini and Anthropic have no OpenAI-compatible TTS + // endpoint and always return a *translateError with notSupported set. + audioSpeech(ctx context.Context, w http.ResponseWriter, body []byte, contentType string) (usage, error) + // audioTranscription proxies body — the client's original raw + // multipart request, replayed unchanged — to the provider's + // speech-to-text endpoint, forwarding its (JSON or text) response to w + // verbatim (spec §3, v0.2). Gemini and Anthropic have no + // OpenAI-compatible STT endpoint and always return a *translateError + // with notSupported set. + audioTranscription(ctx context.Context, w http.ResponseWriter, body []byte, contentType string) (usage, error) // httpClient returns the adapter's shared *http.Client, so a caller // outside this file (the native passthrough route, routes_passthrough.go) // can issue upstream requests through the same connection-pooled client @@ -190,13 +215,47 @@ func upstreamJSON(ctx context.Context, client *http.Client, method, url string, } bodyBytes = b } + return upstreamBytes(ctx, client, method, url, hdr, bodyBytes, policy) +} +// upstreamRawBytes is upstreamJSON's counterpart for a body the caller has +// already encoded: the audio endpoints' openai-type adapter methods +// (provider_openai.go's audioSpeech and audioTranscription) call this +// directly with a body that must never be marshaled again — audioSpeech's +// is a JSON request the caller (routes_media.go) already re-marshaled +// once after rewriting "model", and audioTranscription's is a raw +// multipart payload upstreamJSON's json.Marshal would corrupt outright. +// Shares upstreamBytes with upstreamJSON, so both get the identical +// retry/zero-bytes-reached invariant documented on upstreamJSON. +func upstreamRawBytes(ctx context.Context, client *http.Client, method, url string, hdr http.Header, body []byte, policy *retryPolicy) (*http.Response, error) { + return upstreamBytes(ctx, client, method, url, hdr, body, policy) +} + +// upstreamBytes is the shared request-build-and-send core behind both +// upstreamJSON (which marshals body to JSON first) and upstreamRawBytes +// (which sends an already-encoded body verbatim): hdr's values are added +// to the request, bodyBytes (nil for a body-less request) becomes the +// request body unchanged, and Content-Type defaults to +// "application/json" only when bodyBytes is non-nil and hdr set none of +// its own. +func upstreamBytes(ctx context.Context, client *http.Client, method, url string, hdr http.Header, bodyBytes []byte, policy *retryPolicy) (*http.Response, error) { call := func() (*http.Response, error) { var r io.Reader if bodyBytes != nil { r = bytes.NewReader(bodyBytes) } - req, err := http.NewRequestWithContext(ctx, method, url, r) + // gosec G704 (SSRF via taint analysis) flags url below: every call + // site builds it from an operator-configured value — + // adapter.baseURL (ProviderConfig.BaseURL, or defaultBaseURLByType + // when omitted) concatenated with a fixed, hardcoded path string — + // never anything request-derived that could change the scheme or + // host. Gemini's endpoints additionally embed a url.PathEscape'd + // model id mid-path (provider_gemini.go), which cannot introduce a + // new host or scheme either. See routes_passthrough.go's + // proxyUpstream for the one adapter call path where a request- + // derived path segment genuinely reaches the outgoing URL, and its + // own comment for why that case is safe. + req, err := http.NewRequestWithContext(ctx, method, url, r) //nolint:gosec // operator-configured host; see comment above if err != nil { return nil, fmt.Errorf("%w: %w: build request: %w", errUpstream, errRequestBuildFailed, err) } @@ -209,7 +268,7 @@ func upstreamJSON(ctx context.Context, client *http.Client, method, url string, req.Header.Set("Content-Type", "application/json") } - resp, err := client.Do(req) //nolint:bodyclose // caller closes resp.Body; upstreamJSON hands the response, not its lifecycle, back + resp, err := client.Do(req) //nolint:bodyclose,gosec // caller closes resp.Body; upstreamBytes hands the response, not its lifecycle, back — same operator-configured URL as above if err != nil { return nil, fmt.Errorf("%w: %w", errUpstream, err) } diff --git a/registry_test.go b/registry_test.go index 5f42f67..a01f936 100644 --- a/registry_test.go +++ b/registry_test.go @@ -34,6 +34,15 @@ func (f *fakeAdapter) chatCompletion(context.Context, http.ResponseWriter, map[s func (f *fakeAdapter) embeddings(context.Context, http.ResponseWriter, map[string]any) (usage, error) { return usage{}, nil } +func (f *fakeAdapter) imagesGeneration(context.Context, http.ResponseWriter, map[string]any) (usage, error) { + return usage{}, nil +} +func (f *fakeAdapter) audioSpeech(context.Context, http.ResponseWriter, []byte, string) (usage, error) { + return usage{}, nil +} +func (f *fakeAdapter) audioTranscription(context.Context, http.ResponseWriter, []byte, string) (usage, error) { + return usage{}, nil +} func (f *fakeAdapter) listModels(ctx context.Context) ([]string, error) { if f.listModelsFn == nil { return nil, nil diff --git a/routes_media.go b/routes_media.go new file mode 100644 index 0000000..ba0e7a6 --- /dev/null +++ b/routes_media.go @@ -0,0 +1,244 @@ +package traefikllmgateway + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "mime" + "mime/multipart" + "net/http" +) + +// Route paths for the three unified media endpoints (spec §3, v0.2): +// image generation and both audio endpoints are openai-type-only native +// forwards, except images.generations, which gemini translates to Imagen +// (translate_gemini_images.go); anthropic answers 501 on all three +// (provider_anthropic.go). See llmgateway.go's ServeHTTP dispatch. +const ( + imagesGenerationsPath = "/v1/images/generations" + audioSpeechPath = "/v1/audio/speech" + audioTranscriptionsPath = "/v1/audio/transcriptions" +) + +// maxMultipartModelFieldBytes caps how much of the "model" form field +// extractMultipartModel reads: a model id is always a short identifier, +// so this is generous while bounding memory against a form field a +// misbehaving or malicious client names "model" but fills with megabytes +// of data. +const maxMultipartModelFieldBytes = 4096 + +// decodeMediaJSONRequest reads r.Body capped at maxRequestBytes (routes_ +// unified.go's cap, shared here) and decodes it as a JSON object, then +// extracts its "model" field. It mirrors runUnified's own decode path +// exactly (spec §3's "reuse the unified decode path") — used by both the +// images and audio-speech routes, whose request bodies are both single +// JSON objects carrying a top-level "model" field, unlike +// audio-transcriptions' multipart body. A read failure or invalid/missing +// JSON is a 400, already written to sw; ok reports whether the caller may +// proceed. +func (g *Gateway) decodeMediaJSONRequest(sw *statusTrackingWriter, r *http.Request) (req map[string]any, model string, ok bool) { + body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes)) + if err != nil { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body") + return nil, "", false + } + if err = json.Unmarshal(body, &req); err != nil { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid JSON body") + return nil, "", false + } + model, _ = req["model"].(string) + if model == "" { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "model is required") + return nil, "", false + } + return req, model, true +} + +// resolveMediaRequest runs the shared model-routing/authz/limits steps +// every media route needs before calling its adapter (spec §3's "auth → +// model routing (registry) → group authz → limits" — group authz is +// folded into modelRegistry.resolve itself, the same as +// routes_unified.go's runUnified). On success it returns the resolved +// adapter and the upstream model id to route to; ok reports whether the +// caller may proceed — a resolve or limit failure has already written its +// response to sw. +func (g *Gateway) resolveMediaRequest(sw *statusTrackingWriter, u *user, grp *group, requestedModel string) (adapter providerAdapter, upstreamModel string, ok bool) { + adapter, upstreamModel, _, err := g.registry.resolve(requestedModel, grp) + if err != nil { + writeModelResolveError(sw, err) + return nil, "", false + } + if violation := g.limiter.checkAndCount(buildLimitScopes(u, grp)); violation != nil { + writeLimitViolation(sw, violation) + return nil, "", false + } + return adapter, upstreamModel, true +} + +// handleImagesGenerations implements POST /v1/images/generations (spec +// §3): an openai-type provider forwards natively; gemini translates to +// Imagen (translate_gemini_images.go); anthropic always answers 501 +// (provider_anthropic.go's imagesGeneration). Images are never cached +// (spec §2) and never cost-accounted (spec §3) — only the request +// counters resolveMediaRequest's checkAndCount call already incremented +// move. +func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request, u *user, grp *group) { + sw := &statusTrackingWriter{ResponseWriter: w} + + req, requestedModel, ok := g.decodeMediaJSONRequest(sw, r) + if !ok { + return + } + adapter, upstreamModel, ok := g.resolveMediaRequest(sw, u, grp, requestedModel) + if !ok { + return + } + req["model"] = upstreamModel + + if _, err := adapter.imagesGeneration(r.Context(), sw, req); err != nil { + g.handleAdapterError(sw, err, adapter.name()) + } +} + +// handleAudioSpeech implements POST /v1/audio/speech (spec §3): an +// openai-type provider forwards natively, streaming the binary audio +// response back via a flushWriter (provider_openai.go's audioSpeech); +// gemini and anthropic both always answer 501. req is re-marshaled to +// JSON after its "model" field is rewritten to the resolved upstream +// model id, mirroring how the unified chat/embeddings routes rewrite +// "model" before handing off to their adapter — audioSpeech's own +// interface method takes a raw []byte body rather than a map +// (providers.go), since its response is binary audio no translator ever +// needs to inspect as JSON again. +func (g *Gateway) handleAudioSpeech(w http.ResponseWriter, r *http.Request, u *user, grp *group) { + sw := &statusTrackingWriter{ResponseWriter: w} + + req, requestedModel, ok := g.decodeMediaJSONRequest(sw, r) + if !ok { + return + } + adapter, upstreamModel, ok := g.resolveMediaRequest(sw, u, grp, requestedModel) + if !ok { + return + } + req["model"] = upstreamModel + + body, err := json.Marshal(req) + if err != nil { + // req decoded from the client's own JSON body plus one string-field + // overwrite ("model") — a value json.Unmarshal has already accepted + // once cannot fail to re-marshal; this is a defensive 500, not a + // path any known input reaches. + g.errorf("audio speech: re-marshal request: %v", err) + writeOAIError(sw, http.StatusInternalServerError, "server_error", "internal error") + return + } + + if _, err := adapter.audioSpeech(r.Context(), sw, body, "application/json"); err != nil { + g.handleAdapterError(sw, err, adapter.name()) + } +} + +// handleAudioTranscriptions implements POST /v1/audio/transcriptions +// (spec §3): an openai-type provider forwards natively; gemini and +// anthropic both always answer 501. Unlike images/speech, the multipart +// request body is replayed upstream byte-for-byte, unchanged — spec §3's +// model routing only reads the "model" form field to pick a provider and +// check authorization; it never rewrites the field a client's own model +// alias occupies inside the multipart body, since editing one part of a +// multipart payload in place while preserving its exact boundary framing +// is unnecessary complexity this endpoint's contract does not require. +func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Request, u *user, grp *group) { + sw := &statusTrackingWriter{ResponseWriter: w} + + body, oversize, err := readCapped(r.Body, maxRequestBytes) + if err != nil { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body") + return + } + if oversize { + writeOAIError(sw, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body too large") + return + } + + contentType := r.Header.Get("Content-Type") + _, params, err := mime.ParseMediaType(contentType) + if err != nil { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid Content-Type") + return + } + boundary := params["boundary"] + if boundary == "" { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "missing multipart boundary") + return + } + + requestedModel, err := extractMultipartModel(body, boundary) + if err != nil || requestedModel == "" { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "model is required") + return + } + + adapter, _, ok := g.resolveMediaRequest(sw, u, grp, requestedModel) + if !ok { + return + } + + if _, err := adapter.audioTranscription(r.Context(), sw, body, contentType); err != nil { + g.handleAdapterError(sw, err, adapter.name()) + } +} + +// readCapped reads up to limit+1 bytes from r, reporting oversize=true +// when more than limit bytes were available. Unlike +// decodeMediaJSONRequest's plain io.LimitReader-and-decode (which +// silently truncates an oversized JSON body into a parse failure), +// audio-transcriptions' multipart body cannot be truncated the same way — +// a cut multipart payload often still parses a valid prefix of it, which +// would forward truncated audio upstream instead of failing loudly — so +// this endpoint alone reports the cap explicitly (spec §3's "bodies above +// maxRequestBytes (10MB) → 413"). +func readCapped(r io.Reader, limit int64) (body []byte, oversize bool, err error) { + body, err = io.ReadAll(io.LimitReader(r, limit+1)) + if err != nil { + return nil, false, err + } + if int64(len(body)) > limit { + return body[:limit], true, nil + } + return body, false, nil +} + +// extractMultipartModel scans body — a full multipart/form-data payload +// already read into memory (bounded by handleAudioTranscriptions' own +// readCapped call before this is ever reached) — for its "model" form +// field, returning the field's value. It reads over a fresh +// bytes.NewReader(body), never consuming or mutating body itself, so the +// caller can replay the exact same bytes upstream afterward (spec §3). +// Parts are read in whatever order the client sent them — a "model" +// field following a large audio-file part is found the same as one sent +// first. Returns "" with a nil error when no part named "model" is +// present; the caller treats an empty model the same as a missing one. +func extractMultipartModel(body []byte, boundary string) (string, error) { + mr := multipart.NewReader(bytes.NewReader(body), boundary) + for { + part, err := mr.NextPart() + if errors.Is(err, io.EOF) { + return "", nil + } + if err != nil { + return "", err + } + if part.FormName() != "model" { + _ = part.Close() //nolint:errcheck // read-side close; nothing actionable on failure + continue + } + val, err := io.ReadAll(io.LimitReader(part, maxMultipartModelFieldBytes)) + _ = part.Close() //nolint:errcheck // read-side close; nothing actionable on failure + if err != nil { + return "", err + } + return string(val), nil + } +} diff --git a/routes_media_test.go b/routes_media_test.go new file mode 100644 index 0000000..e5a363f --- /dev/null +++ b/routes_media_test.go @@ -0,0 +1,637 @@ +package traefikllmgateway + +import ( + "bytes" + "context" + "encoding/json" + "io" + "mime/multipart" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +// newMultipartTranscriptionRequest builds a POST /v1/audio/transcriptions +// request with a "file" part and, unless omitModel is true, a "model" +// part carrying model. modelFirst controls part order — false writes the +// file part first and the model part last, covering spec §3's "model +// field arriving after other parts" case. +func newMultipartTranscriptionRequest(t *testing.T, apiKey, model string, modelFirst, omitModel bool) *http.Request { + t.Helper() + var buf bytes.Buffer + mw := multipart.NewWriter(&buf) + + writeModel := func() { + fw, err := mw.CreateFormField("model") + if err != nil { + t.Fatalf("CreateFormField(model): %v", err) + } + if _, err := fw.Write([]byte(model)); err != nil { + t.Fatalf("write model field: %v", err) + } + } + writeFile := func() { + fw, err := mw.CreateFormFile("file", "audio.wav") + if err != nil { + t.Fatalf("CreateFormFile(file): %v", err) + } + if _, err := fw.Write([]byte("FAKEAUDIOBYTES")); err != nil { + t.Fatalf("write file field: %v", err) + } + } + + if modelFirst && !omitModel { + writeModel() + } + writeFile() + if !modelFirst && !omitModel { + writeModel() + } + if err := mw.Close(); err != nil { + t.Fatalf("close multipart writer: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, audioTranscriptionsPath, &buf) + req.Header.Set("Content-Type", mw.FormDataContentType()) + if apiKey != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + return req +} + +// newMediaTestConfig returns a base Config wired with an "openai" provider +// pointed at openaiURL, a permissive "default" group, and one inline user +// "alice" — the shared fixture every routes_media.go handler test starts +// from, customized further per test as needed. +func newMediaTestConfig(openaiURL string, models ...string) *Config { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: openaiURL, APIKey: "sk-up", Models: models}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + // Limits is a non-nil (if all-zero, i.e. unlimited) *LimitsConfig so + // buildLimitScopes (routes_media.go/routes_unified.go) includes a + // "user" scope for alice — matching routes_unified_test.go's own + // happy-path fixture — otherwise checkAndCount has no scope to + // increment a counter for at all. + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}} + return cfg +} + +// newMediaTestGateway constructs a *Gateway from cfg via New, failing the +// test on any construction error. +func newMediaTestGateway(t *testing.T, cfg *Config) *Gateway { + t.Helper() + next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + return gw +} + +// ---- routing: ServeHTTP recognizes all three media paths ahead of auth ---- + +// TestServeHTTP_MediaRoutes_RequireAuth proves llmgateway.go's dispatch +// recognizes all three media paths (spec §3) before falling through to +// the unknown-route 404: an unauthenticated request to each gets 401, not +// 404. +func TestServeHTTP_MediaRoutes_RequireAuth(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + gw := newMediaTestGateway(t, cfg) + + for _, path := range []string{imagesGenerationsPath, audioSpeechPath, audioTranscriptionsPath} { + t.Run(path, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{}`)) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401 (proves routing reached auth, not a 404 fallthrough), body=%s", rec.Code, rec.Body.String()) + } + }) + } +} + +// ---- images.generations: openai native forward ---- + +func TestHandleImagesGenerations_OpenAI_HappyPath_NativeForward(t *testing.T) { + const respBody = `{"created":1734000000,"data":[{"b64_json":"xyz"}]}` + var gotBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/images/generations" { + t.Errorf("upstream path = %q, want /v1/images/generations", r.URL.Path) + } + _ = json.NewDecoder(r.Body).Decode(&gotBody) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + gw := newMediaTestGateway(t, newMediaTestConfig(srv.URL, "img-test")) + + body := map[string]any{"model": "img-test", "prompt": "a cat", "n": 1} + req := newUnifiedRequest(t, http.MethodPost, imagesGenerationsPath, "sk-alice", body) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.String() != respBody { + t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody) + } + if _, leaked := gotBody[gatewayAliasKey]; leaked { + t.Errorf("upstream request body leaked %q: %v", gatewayAliasKey, gotBody) + } + if gotBody["model"] != "img-test" { + t.Errorf("upstream request model = %v, want %q", gotBody["model"], "img-test") + } + + reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now()) + if !ok || reqCount != 1 { + t.Errorf("user request/min counter = %d (ok=%v), want 1", reqCount, ok) + } + if tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now()); ok && tok != 0 { + t.Errorf("user token/day counter = %d, want 0 (images are never token-accounted)", tok) + } +} + +func TestHandleImagesGenerations_MissingModel_400(t *testing.T) { + gw := newMediaTestGateway(t, newMediaTestConfig("http://127.0.0.1:1", "img-test")) + + req := newUnifiedRequest(t, http.MethodPost, imagesGenerationsPath, "sk-alice", map[string]any{"prompt": "a cat"}) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400, body=%s", rec.Code, rec.Body.String()) + } +} + +func TestHandleImagesGenerations_InvalidJSON_400(t *testing.T) { + gw := newMediaTestGateway(t, newMediaTestConfig("http://127.0.0.1:1", "img-test")) + + req := httptest.NewRequest(http.MethodPost, imagesGenerationsPath, strings.NewReader(`{not json`)) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400, body=%s", rec.Code, rec.Body.String()) + } +} + +// TestHandleImagesGenerations_RetriedOnTransientFailure proves retry +// (spec §1) covers the images endpoint end to end through the full +// Gateway/ServeHTTP pipeline: the upstream fails once with 503, then +// succeeds, and the client sees only the successful response. +func TestHandleImagesGenerations_RetriedOnTransientFailure(t *testing.T) { + const respBody = `{"created":1,"data":[{"b64_json":"xyz"}]}` + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&calls, 1) == 1 { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := newMediaTestConfig(srv.URL, "img-test") + cfg.Retry = RetryConfig{Enabled: true, Attempts: 1, Backoff: "1ms"} + gw := newMediaTestGateway(t, cfg) + + body := map[string]any{"model": "img-test", "prompt": "a cat"} + req := newUnifiedRequest(t, http.MethodPost, imagesGenerationsPath, "sk-alice", body) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.String() != respBody { + t.Errorf("body = %q, want %q", rec.Body.String(), respBody) + } + if got := atomic.LoadInt32(&calls); got != 2 { + t.Errorf("upstream calls = %d, want 2 (one retry)", got) + } +} + +// TestHandleImagesGenerations_Gemini_TranslatesToImagen exercises the +// full gemini translation chain (translate_gemini_images.go) through +// ServeHTTP: request translated to Imagen's :predict shape, response +// translated back to OpenAI's images.generations shape. +func TestHandleImagesGenerations_Gemini_TranslatesToImagen(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(`{"predictions":[{"bytesBase64Encoded":"b64data"}]}`)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "gemini": {Type: "gemini", BaseURL: srv.URL, APIKey: "sk-gem", Models: []string{"imagen-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + gw := newMediaTestGateway(t, cfg) + + body := map[string]any{"model": "imagen-test", "prompt": "a cat", "size": "1024x1024"} + req := newUnifiedRequest(t, http.MethodPost, imagesGenerationsPath, "sk-alice", body) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if gotPath != "/v1beta/models/imagen-test:predict" { + t.Errorf("upstream path = %q, want /v1beta/models/imagen-test:predict", gotPath) + } + var out map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatalf("decode response: %v", err) + } + data, _ := out["data"].([]any) + if len(data) != 1 { + t.Fatalf("data = %v, want 1 entry", out["data"]) + } + entry, _ := data[0].(map[string]any) + if entry["b64_json"] != "b64data" { + t.Errorf("b64_json = %v, want %q", entry["b64_json"], "b64data") + } +} + +// ---- audio.speech: openai native forward, binary streaming ---- + +// TestHandleAudioSpeech_OpenAI_StreamsBinary_WithFlushes proves the binary +// audio response is streamed to the client chunk by chunk, flushing after +// each one, rather than buffered until the response completes — mirroring +// TestHandleChat_Streaming_FlushesIncrementally's assertions +// (routes_unified_test.go) for the SSE case. Unlike that SSE test, a raw +// binary stream carries no application-level event framing for readSSE's +// line-based scanner to dispatch on, so a short sleep between the fake +// upstream's writes is needed here to force the client's io.Copy to see +// each chunk as its own Read — without it, chunks written back-to-back +// with no delay can legitimately coalesce into a single Read on a fast +// loopback connection, which would make this test flaky rather than prove +// anything about audioSpeech's own streaming behavior. +func TestHandleAudioSpeech_OpenAI_StreamsBinary_WithFlushes(t *testing.T) { + chunks := []string{"RIFF-chunk-one-", "RIFF-chunk-two-", "RIFF-chunk-three"} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "audio/mpeg") + w.WriteHeader(http.StatusOK) + fl, ok := w.(http.Flusher) + if !ok { + t.Fatal("fake upstream ResponseWriter does not support Flush") + } + for i, c := range chunks { + _, _ = w.Write([]byte(c)) + fl.Flush() + if i < len(chunks)-1 { + time.Sleep(20 * time.Millisecond) + } + } + })) + defer srv.Close() + + gw := newMediaTestGateway(t, newMediaTestConfig(srv.URL, "tts-test")) + + body := map[string]any{"model": "tts-test", "input": "hello", "voice": "alloy"} + req := newUnifiedRequest(t, http.MethodPost, audioSpeechPath, "sk-alice", body) + rw := newRecordingWriter() + gw.ServeHTTP(rw, req) + + if rw.status != http.StatusOK { + t.Fatalf("status = %d, want 200", rw.status) + } + writes, flushes := 0, 0 + for _, e := range rw.events { + switch e { + case "write": + writes++ + case "flush": + flushes++ + } + } + if writes < len(chunks) { + t.Fatalf("want >=%d writes, got %d (events=%v)", len(chunks), writes, rw.events) + } + if flushes < writes { + t.Fatalf("want a flush per write (incremental delivery), got %d writes and %d flushes", writes, flushes) + } + if want := strings.Join(chunks, ""); rw.buf.String() != want { + t.Errorf("body = %q, want %q", rw.buf.String(), want) + } + if ct := rw.hdr.Get("Content-Type"); ct != "audio/mpeg" { + t.Errorf("Content-Type = %q, want audio/mpeg", ct) + } +} + +func TestHandleAudioSpeech_MissingModel_400(t *testing.T) { + gw := newMediaTestGateway(t, newMediaTestConfig("http://127.0.0.1:1", "tts-test")) + + req := newUnifiedRequest(t, http.MethodPost, audioSpeechPath, "sk-alice", map[string]any{"input": "hello"}) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400, body=%s", rec.Code, rec.Body.String()) + } +} + +// ---- audio.transcriptions: openai native forward, multipart replay ---- + +// buildMultipartTranscription returns the raw bytes of a multipart/ +// form-data body carrying a "model" field (after the "file" part, per +// spec §3's "model field arriving after other parts" case) and its +// Content-Type header value. +func buildMultipartTranscription(t *testing.T, model string) (body []byte, contentType string) { + t.Helper() + var buf bytes.Buffer + mw := multipart.NewWriter(&buf) + fw, err := mw.CreateFormFile("file", "audio.wav") + if err != nil { + t.Fatalf("CreateFormFile: %v", err) + } + if _, err = fw.Write([]byte("FAKEAUDIOBYTES")); err != nil { + t.Fatalf("write file field: %v", err) + } + mfw, err := mw.CreateFormField("model") + if err != nil { + t.Fatalf("CreateFormField(model): %v", err) + } + if _, err = mfw.Write([]byte(model)); err != nil { + t.Fatalf("write model field: %v", err) + } + if err = mw.Close(); err != nil { + t.Fatalf("close multipart writer: %v", err) + } + return buf.Bytes(), mw.FormDataContentType() +} + +func TestHandleAudioTranscriptions_OpenAI_HappyPath_MultipartReplayed(t *testing.T) { + const respBody = `{"text":"hello world"}` + var gotContentType string + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/audio/transcriptions" { + t.Errorf("upstream path = %q, want /v1/audio/transcriptions", r.URL.Path) + } + gotContentType = r.Header.Get("Content-Type") + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + gw := newMediaTestGateway(t, newMediaTestConfig(srv.URL, "whisper-test")) + + // The "model" field is placed after the "file" part in the built body + // (buildMultipartTranscription), covering spec §3's "model field + // arriving after other parts" case end to end. + origBody, contentType := buildMultipartTranscription(t, "whisper-test") + req := httptest.NewRequest(http.MethodPost, audioTranscriptionsPath, bytes.NewReader(origBody)) + req.Header.Set("Content-Type", contentType) + req.Header.Set("Authorization", "Bearer sk-alice") + + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.String() != respBody { + t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody) + } + if gotContentType != contentType { + t.Errorf("upstream Content-Type = %q, want %q (forwarded verbatim)", gotContentType, contentType) + } + if !bytes.Equal(gotBody, origBody) { + t.Errorf("upstream body was rewritten; want the exact original multipart bytes replayed unchanged") + } +} + +func TestHandleAudioTranscriptions_ModelFieldMissing_400(t *testing.T) { + gw := newMediaTestGateway(t, newMediaTestConfig("http://127.0.0.1:1", "whisper-test")) + + req := newMultipartTranscriptionRequest(t, "sk-alice", "", true, true) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400, body=%s", rec.Code, rec.Body.String()) + } +} + +func TestHandleAudioTranscriptions_InvalidContentType_400(t *testing.T) { + gw := newMediaTestGateway(t, newMediaTestConfig("http://127.0.0.1:1", "whisper-test")) + + req := httptest.NewRequest(http.MethodPost, audioTranscriptionsPath, strings.NewReader("garbage")) + req.Header.Set("Content-Type", "application/json") // not multipart at all + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400, body=%s", rec.Code, rec.Body.String()) + } +} + +// TestHandleAudioTranscriptions_OversizeBody_413 proves a multipart body +// above maxRequestBytes is rejected with 413 (spec §3), the one media +// route with an explicit oversize check — unlike images/speech's shared +// JSON decode path, which silently truncates instead (routes_media.go's +// readCapped doc comment explains why). +func TestHandleAudioTranscriptions_OversizeBody_413(t *testing.T) { + gw := newMediaTestGateway(t, newMediaTestConfig("http://127.0.0.1:1", "whisper-test")) + + var buf bytes.Buffer + mw := multipart.NewWriter(&buf) + fw, err := mw.CreateFormField("model") + if err != nil { + t.Fatalf("CreateFormField: %v", err) + } + _, _ = fw.Write([]byte("whisper-test")) + ffw, err := mw.CreateFormFile("file", "big.wav") + if err != nil { + t.Fatalf("CreateFormFile: %v", err) + } + _, _ = ffw.Write(bytes.Repeat([]byte("A"), maxRequestBytes+1024)) + if err := mw.Close(); err != nil { + t.Fatalf("close multipart writer: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, audioTranscriptionsPath, &buf) + req.Header.Set("Content-Type", mw.FormDataContentType()) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("status = %d, want 413, body=%s", rec.Code, rec.Body.String()) + } +} + +// ---- unit coverage for the new pure helpers ---- + +func TestReadCapped(t *testing.T) { + t.Run("under limit", func(t *testing.T) { + body, oversize, err := readCapped(strings.NewReader("hello"), 10) + if err != nil || oversize || string(body) != "hello" { + t.Errorf("readCapped = (%q, %v, %v), want (\"hello\", false, nil)", body, oversize, err) + } + }) + t.Run("exactly at limit", func(t *testing.T) { + body, oversize, err := readCapped(strings.NewReader("0123456789"), 10) + if err != nil || oversize || string(body) != "0123456789" { + t.Errorf("readCapped = (%q, %v, %v), want (\"0123456789\", false, nil)", body, oversize, err) + } + }) + t.Run("over limit", func(t *testing.T) { + body, oversize, err := readCapped(strings.NewReader("01234567890"), 10) + if err != nil || !oversize || len(body) != 10 { + t.Errorf("readCapped = (len=%d, %v, %v), want (len=10, true, nil)", len(body), oversize, err) + } + }) +} + +func TestExtractMultipartModel(t *testing.T) { + build := func(fields []struct{ name, value string }, includeFile bool) (body []byte, boundary string) { + var buf bytes.Buffer + mw := multipart.NewWriter(&buf) + for _, f := range fields { + fw, err := mw.CreateFormField(f.name) + if err != nil { + t.Fatalf("CreateFormField: %v", err) + } + _, _ = fw.Write([]byte(f.value)) + } + if includeFile { + fw, err := mw.CreateFormFile("file", "x.wav") + if err != nil { + t.Fatalf("CreateFormFile: %v", err) + } + _, _ = fw.Write([]byte("bytes")) + } + if err := mw.Close(); err != nil { + t.Fatalf("close: %v", err) + } + return buf.Bytes(), mw.Boundary() + } + + t.Run("model field present", func(t *testing.T) { + body, boundary := build([]struct{ name, value string }{{"model", "whisper-1"}}, true) + got, err := extractMultipartModel(body, boundary) + if err != nil || got != "whisper-1" { + t.Errorf("extractMultipartModel = (%q, %v), want (\"whisper-1\", nil)", got, err) + } + }) + t.Run("model field missing", func(t *testing.T) { + body, boundary := build(nil, true) + got, err := extractMultipartModel(body, boundary) + if err != nil || got != "" { + t.Errorf("extractMultipartModel = (%q, %v), want (\"\", nil)", got, err) + } + }) + t.Run("no parts at all", func(t *testing.T) { + body, boundary := build(nil, false) + got, err := extractMultipartModel(body, boundary) + if err != nil || got != "" { + t.Errorf("extractMultipartModel = (%q, %v), want (\"\", nil)", got, err) + } + }) +} + +// ---- gemini: audio always 501 ---- + +func newGeminiOnlyConfig(baseURL string, models ...string) *Config { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "gemini": {Type: "gemini", BaseURL: baseURL, APIKey: "sk-gem", Models: models}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + return cfg +} + +func TestHandleAudioSpeech_Gemini_Returns501(t *testing.T) { + gw := newMediaTestGateway(t, newGeminiOnlyConfig("http://127.0.0.1:1", "gem-audio")) + + req := newUnifiedRequest(t, http.MethodPost, audioSpeechPath, "sk-alice", map[string]any{"model": "gem-audio", "input": "hi"}) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusNotImplemented { + t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String()) + } +} + +func TestHandleAudioTranscriptions_Gemini_Returns501(t *testing.T) { + gw := newMediaTestGateway(t, newGeminiOnlyConfig("http://127.0.0.1:1", "gem-audio")) + + req := newMultipartTranscriptionRequest(t, "sk-alice", "gem-audio", true, false) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + + if rec.Code != http.StatusNotImplemented { + t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String()) + } +} + +// ---- anthropic: all three endpoints 501, never dial upstream ---- + +// TestAnthropicMediaEndpoints_Return501_NeverCallUpstream proves anthropic's +// three media adapter methods (provider_anthropic.go) return their +// *translateError{notSupported:true} before making any upstream HTTP +// request: BaseURL points at a port nothing listens on, so a 502 +// (upstream connection error) instead of 501 would mean the adapter tried +// to dial anyway. +func TestAnthropicMediaEndpoints_Return501_NeverCallUpstream(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "anthropic": {Type: "anthropic", BaseURL: "http://127.0.0.1:1", APIKey: "sk-ant", Models: []string{"claude-media"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + gw := newMediaTestGateway(t, cfg) + + t.Run("images.generations", func(t *testing.T) { + req := newUnifiedRequest(t, http.MethodPost, imagesGenerationsPath, "sk-alice", map[string]any{"model": "claude-media", "prompt": "a cat"}) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + if rec.Code != http.StatusNotImplemented { + t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String()) + } + }) + t.Run("audio.speech", func(t *testing.T) { + req := newUnifiedRequest(t, http.MethodPost, audioSpeechPath, "sk-alice", map[string]any{"model": "claude-media", "input": "hi"}) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + if rec.Code != http.StatusNotImplemented { + t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String()) + } + }) + t.Run("audio.transcriptions", func(t *testing.T) { + req := newMultipartTranscriptionRequest(t, "sk-alice", "claude-media", true, false) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + if rec.Code != http.StatusNotImplemented { + t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String()) + } + }) +} diff --git a/testdata/gemini_images/req_basic_in.json b/testdata/gemini_images/req_basic_in.json new file mode 100644 index 0000000..c65823a --- /dev/null +++ b/testdata/gemini_images/req_basic_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","n":2,"size":"1024x1024"} diff --git a/testdata/gemini_images/req_basic_want.json b/testdata/gemini_images/req_basic_want.json new file mode 100644 index 0000000..6f9cc49 --- /dev/null +++ b/testdata/gemini_images/req_basic_want.json @@ -0,0 +1 @@ +{"instances":[{"prompt":"a cat"}],"parameters":{"sampleCount":2,"aspectRatio":"1:1"}} diff --git a/testdata/gemini_images/req_n_zero_in.json b/testdata/gemini_images/req_n_zero_in.json new file mode 100644 index 0000000..b7cb697 --- /dev/null +++ b/testdata/gemini_images/req_n_zero_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","n":0} diff --git a/testdata/gemini_images/req_n_zero_want.json b/testdata/gemini_images/req_n_zero_want.json new file mode 100644 index 0000000..fe08716 --- /dev/null +++ b/testdata/gemini_images/req_n_zero_want.json @@ -0,0 +1 @@ +{"instances":[{"prompt":"a cat"}],"parameters":{"sampleCount":1,"aspectRatio":"1:1"}} diff --git a/testdata/gemini_images/req_quality_falsy_in.json b/testdata/gemini_images/req_quality_falsy_in.json new file mode 100644 index 0000000..56bac83 --- /dev/null +++ b/testdata/gemini_images/req_quality_falsy_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","quality":""} diff --git a/testdata/gemini_images/req_quality_falsy_want.json b/testdata/gemini_images/req_quality_falsy_want.json new file mode 100644 index 0000000..fe08716 --- /dev/null +++ b/testdata/gemini_images/req_quality_falsy_want.json @@ -0,0 +1 @@ +{"instances":[{"prompt":"a cat"}],"parameters":{"sampleCount":1,"aspectRatio":"1:1"}} diff --git a/testdata/gemini_images/req_quality_truthy_in.json b/testdata/gemini_images/req_quality_truthy_in.json new file mode 100644 index 0000000..21eae3c --- /dev/null +++ b/testdata/gemini_images/req_quality_truthy_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","quality":"hd"} diff --git a/testdata/gemini_images/req_quality_truthy_wanterr.json b/testdata/gemini_images/req_quality_truthy_wanterr.json new file mode 100644 index 0000000..f43064c --- /dev/null +++ b/testdata/gemini_images/req_quality_truthy_wanterr.json @@ -0,0 +1 @@ +{"errContains":"quality","notSupported":false} diff --git a/testdata/gemini_images/req_response_format_b64_in.json b/testdata/gemini_images/req_response_format_b64_in.json new file mode 100644 index 0000000..ec6785f --- /dev/null +++ b/testdata/gemini_images/req_response_format_b64_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","response_format":"b64_json"} diff --git a/testdata/gemini_images/req_response_format_b64_want.json b/testdata/gemini_images/req_response_format_b64_want.json new file mode 100644 index 0000000..fe08716 --- /dev/null +++ b/testdata/gemini_images/req_response_format_b64_want.json @@ -0,0 +1 @@ +{"instances":[{"prompt":"a cat"}],"parameters":{"sampleCount":1,"aspectRatio":"1:1"}} diff --git a/testdata/gemini_images/req_response_format_url_in.json b/testdata/gemini_images/req_response_format_url_in.json new file mode 100644 index 0000000..8b4a8b2 --- /dev/null +++ b/testdata/gemini_images/req_response_format_url_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","response_format":"url"} diff --git a/testdata/gemini_images/req_response_format_url_wanterr.json b/testdata/gemini_images/req_response_format_url_wanterr.json new file mode 100644 index 0000000..d4ad74d --- /dev/null +++ b/testdata/gemini_images/req_response_format_url_wanterr.json @@ -0,0 +1 @@ +{"errContains":"response_format","notSupported":false} diff --git a/testdata/gemini_images/req_size_16x9_in.json b/testdata/gemini_images/req_size_16x9_in.json new file mode 100644 index 0000000..a5f40ef --- /dev/null +++ b/testdata/gemini_images/req_size_16x9_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","size":"1792x1024"} diff --git a/testdata/gemini_images/req_size_16x9_want.json b/testdata/gemini_images/req_size_16x9_want.json new file mode 100644 index 0000000..ef11b0f --- /dev/null +++ b/testdata/gemini_images/req_size_16x9_want.json @@ -0,0 +1 @@ +{"instances":[{"prompt":"a cat"}],"parameters":{"sampleCount":1,"aspectRatio":"16:9"}} diff --git a/testdata/gemini_images/req_size_512_in.json b/testdata/gemini_images/req_size_512_in.json new file mode 100644 index 0000000..7bce3e9 --- /dev/null +++ b/testdata/gemini_images/req_size_512_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","size":"512x512"} diff --git a/testdata/gemini_images/req_size_512_want.json b/testdata/gemini_images/req_size_512_want.json new file mode 100644 index 0000000..fe08716 --- /dev/null +++ b/testdata/gemini_images/req_size_512_want.json @@ -0,0 +1 @@ +{"instances":[{"prompt":"a cat"}],"parameters":{"sampleCount":1,"aspectRatio":"1:1"}} diff --git a/testdata/gemini_images/req_size_9x16_in.json b/testdata/gemini_images/req_size_9x16_in.json new file mode 100644 index 0000000..6be8d69 --- /dev/null +++ b/testdata/gemini_images/req_size_9x16_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","size":"1024x1792"} diff --git a/testdata/gemini_images/req_size_9x16_want.json b/testdata/gemini_images/req_size_9x16_want.json new file mode 100644 index 0000000..e8c690a --- /dev/null +++ b/testdata/gemini_images/req_size_9x16_want.json @@ -0,0 +1 @@ +{"instances":[{"prompt":"a cat"}],"parameters":{"sampleCount":1,"aspectRatio":"9:16"}} diff --git a/testdata/gemini_images/req_size_absent_in.json b/testdata/gemini_images/req_size_absent_in.json new file mode 100644 index 0000000..80cdb25 --- /dev/null +++ b/testdata/gemini_images/req_size_absent_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat"} diff --git a/testdata/gemini_images/req_size_absent_want.json b/testdata/gemini_images/req_size_absent_want.json new file mode 100644 index 0000000..fe08716 --- /dev/null +++ b/testdata/gemini_images/req_size_absent_want.json @@ -0,0 +1 @@ +{"instances":[{"prompt":"a cat"}],"parameters":{"sampleCount":1,"aspectRatio":"1:1"}} diff --git a/testdata/gemini_images/req_size_unknown_in.json b/testdata/gemini_images/req_size_unknown_in.json new file mode 100644 index 0000000..4832d81 --- /dev/null +++ b/testdata/gemini_images/req_size_unknown_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","size":"999x999"} diff --git a/testdata/gemini_images/req_size_unknown_wanterr.json b/testdata/gemini_images/req_size_unknown_wanterr.json new file mode 100644 index 0000000..48420d0 --- /dev/null +++ b/testdata/gemini_images/req_size_unknown_wanterr.json @@ -0,0 +1 @@ +{"errContains":"size","notSupported":false} diff --git a/testdata/gemini_images/req_style_truthy_in.json b/testdata/gemini_images/req_style_truthy_in.json new file mode 100644 index 0000000..b16e8f9 --- /dev/null +++ b/testdata/gemini_images/req_style_truthy_in.json @@ -0,0 +1 @@ +{"model":"imagen-3","prompt":"a cat","style":"vivid"} diff --git a/testdata/gemini_images/req_style_truthy_wanterr.json b/testdata/gemini_images/req_style_truthy_wanterr.json new file mode 100644 index 0000000..54284df --- /dev/null +++ b/testdata/gemini_images/req_style_truthy_wanterr.json @@ -0,0 +1 @@ +{"errContains":"style","notSupported":false} diff --git a/testdata/gemini_images/resp_basic_in.json b/testdata/gemini_images/resp_basic_in.json new file mode 100644 index 0000000..9b9fe33 --- /dev/null +++ b/testdata/gemini_images/resp_basic_in.json @@ -0,0 +1 @@ +{"predictions":[{"bytesBase64Encoded":"abc123"}]} diff --git a/testdata/gemini_images/resp_basic_want.json b/testdata/gemini_images/resp_basic_want.json new file mode 100644 index 0000000..1568f6c --- /dev/null +++ b/testdata/gemini_images/resp_basic_want.json @@ -0,0 +1 @@ +{"created":1734000000,"data":[{"b64_json":"abc123"}]} diff --git a/testdata/gemini_images/resp_empty_in.json b/testdata/gemini_images/resp_empty_in.json new file mode 100644 index 0000000..8561c74 --- /dev/null +++ b/testdata/gemini_images/resp_empty_in.json @@ -0,0 +1 @@ +{"predictions":[]} diff --git a/testdata/gemini_images/resp_empty_want.json b/testdata/gemini_images/resp_empty_want.json new file mode 100644 index 0000000..c83c53c --- /dev/null +++ b/testdata/gemini_images/resp_empty_want.json @@ -0,0 +1 @@ +{"created":1734000000,"data":[]} diff --git a/testdata/gemini_images/resp_multi_in.json b/testdata/gemini_images/resp_multi_in.json new file mode 100644 index 0000000..8cb20dd --- /dev/null +++ b/testdata/gemini_images/resp_multi_in.json @@ -0,0 +1 @@ +{"predictions":[{"bytesBase64Encoded":"aaa"},{"bytesBase64Encoded":"bbb"}]} diff --git a/testdata/gemini_images/resp_multi_want.json b/testdata/gemini_images/resp_multi_want.json new file mode 100644 index 0000000..f5c58d1 --- /dev/null +++ b/testdata/gemini_images/resp_multi_want.json @@ -0,0 +1 @@ +{"created":1734000000,"data":[{"b64_json":"aaa"},{"b64_json":"bbb"}]} diff --git a/translate_gemini_images.go b/translate_gemini_images.go new file mode 100644 index 0000000..a4d04f7 --- /dev/null +++ b/translate_gemini_images.go @@ -0,0 +1,121 @@ +package traefikllmgateway + +import ( + "encoding/json" + "fmt" +) + +// geminiImageAspectRatioBySize maps OpenAI's images.generations "size" +// values to the Imagen aspectRatio parameter this translator understands, +// per spec §3's mapping table. A size string absent from this map is one +// Imagen has no known aspect ratio for — geminiImageAspectRatio rejects it +// as a *translateError rather than guessing. +var geminiImageAspectRatioBySize = map[string]string{ + "1024x1024": "1:1", + "1792x1024": "16:9", + "1024x1792": "9:16", + "512x512": "1:1", +} + +// geminiImageAspectRatio resolves req's OpenAI "size" field to an Imagen +// aspectRatio parameter: an absent or empty size defaults to "1:1" (spec +// §3), and any other value must be one of geminiImageAspectRatioBySize's +// keys. An unrecognized size is a *translateError, not a silent fallback +// — the client asked for a specific size the gateway cannot honor, so +// letting Imagen apply its own default instead would silently ignore the +// request rather than fail loudly. +func geminiImageAspectRatio(size string) (string, error) { + if size == "" { + return "1:1", nil + } + ratio, ok := geminiImageAspectRatioBySize[size] + if !ok { + return "", &translateError{msg: fmt.Sprintf("size %q is not supported for gemini image generation", size)} + } + return ratio, nil +} + +// geminiImagesRequestFromOpenAI maps an OpenAI images.generations request +// body ({"model","prompt","n","size","response_format",...}) to Imagen's +// :predict request body ({"instances":[{"prompt"}],"parameters": +// {"sampleCount","aspectRatio"}}), per spec §3. response_format:"url" is +// rejected outright — the gateway stores nothing, so it cannot honor a +// client asking for a hosted URL back; the response this translator +// always produces is b64_json (translate_gemini_images.go's +// openAIImagesResponseFromGemini). A truthy "quality" or "style" field is +// rejected the same way chat's translators reject an unsupported field +// (isTruthy, translate_anthropic.go): Imagen has no equivalent for +// either. n defaults to 1 when absent or non-positive, matching OpenAI's +// own default. The request's "model" field is deliberately never read +// here — the adapter reads it separately to build the :predict URL, the +// same convention chatCompletion/embeddings already follow. +func geminiImagesRequestFromOpenAI(req map[string]any) (map[string]any, error) { + if rf, ok := req["response_format"].(string); ok && rf == "url" { + return nil, &translateError{msg: `response_format "url" is not supported for gemini image generation (the gateway stores nothing to host a URL from)`} + } + if v, ok := req["quality"]; ok && isTruthy(v) { + return nil, geminiUnsupportedFieldError("quality") + } + if v, ok := req["style"]; ok && isTruthy(v) { + return nil, geminiUnsupportedFieldError("style") + } + + prompt, _ := req["prompt"].(string) + + n := int64(1) + if v, ok := req["n"]; ok { + if parsed, ok2 := toInt64(v); ok2 && parsed > 0 { + n = parsed + } + } + + size, _ := req["size"].(string) + aspectRatio, err := geminiImageAspectRatio(size) + if err != nil { + return nil, err + } + + return map[string]any{ + "instances": []any{map[string]any{"prompt": prompt}}, + "parameters": map[string]any{ + "sampleCount": n, + "aspectRatio": aspectRatio, + }, + }, nil +} + +// geminiImagePrediction is one entry of Imagen's :predict response +// "predictions" array this translator reads. +type geminiImagePrediction struct { + BytesBase64Encoded string `json:"bytesBase64Encoded"` +} + +// geminiImagePredictResponse is the subset of Imagen's :predict response +// body this translator reads. +type geminiImagePredictResponse struct { + Predictions []geminiImagePrediction `json:"predictions"` +} + +// openAIImagesResponseFromGemini maps an Imagen :predict response body to +// an OpenAI images.generations response ({"created","data":[{"b64_json"}]}), +// per spec §3's response mapping. created is the caller-supplied response +// timestamp — Imagen's response carries none. An empty predictions array +// decodes to an empty (non-nil) data array, not an error: Imagen +// answering with zero images is a valid, if unusual, response to forward +// as-is. +func openAIImagesResponseFromGemini(body []byte, created int64) (map[string]any, error) { + var resp geminiImagePredictResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("%w: decode gemini images response: %w", errUpstream, err) + } + + data := make([]any, len(resp.Predictions)) + for i, p := range resp.Predictions { + data[i] = map[string]any{"b64_json": p.BytesBase64Encoded} + } + + return map[string]any{ + "created": created, + "data": data, + }, nil +} diff --git a/translate_gemini_images_test.go b/translate_gemini_images_test.go new file mode 100644 index 0000000..0154846 --- /dev/null +++ b/translate_gemini_images_test.go @@ -0,0 +1,159 @@ +package traefikllmgateway + +import ( + "encoding/json" + "os" + "path/filepath" + "reflect" + "sort" + "strings" + "testing" +) + +// geminiImagesGoldenTestdataDir is where every fixture this file walks +// lives. +const geminiImagesGoldenTestdataDir = "testdata/gemini_images" + +// geminiImagesGoldenCreated is the fixed "created" timestamp every +// response golden fixture's "_want.json" bakes in — +// openAIImagesResponseFromGemini takes it as a caller-supplied parameter +// (Imagen's response carries none), so the golden test must supply the +// same fixed value the "_want.json" files were written against. +const geminiImagesGoldenCreated = int64(1734000000) + +// geminiImagesReadGoldenJSON reads and json.Unmarshals name (relative to +// testdata/gemini_images) into a fresh map[string]any. +func geminiImagesReadGoldenJSON(t *testing.T, name string) map[string]any { + t.Helper() + b, err := os.ReadFile(filepath.Join(geminiImagesGoldenTestdataDir, 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 +} + +// TestGeminiImagesRequestGolden walks every "req_*_in.json" fixture under +// testdata/gemini_images, translating it via geminiImagesRequestFromOpenAI, +// 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 TestGeminiImagesRequestGolden(t *testing.T) { + entries, err := os.ReadDir(geminiImagesGoldenTestdataDir) + 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, "req_") { + names = append(names, n) + } + } + sort.Strings(names) + if len(names) == 0 { + t.Fatalf("no req_*_in.json fixtures found under %s", geminiImagesGoldenTestdataDir) + } + + for _, name := range names { + t.Run(name, func(t *testing.T) { + in := geminiImagesReadGoldenJSON(t, name+"_in.json") + + got, err := geminiImagesRequestFromOpenAI(in) + + wantErrPath := filepath.Join(geminiImagesGoldenTestdataDir, 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) + } + assertTranslateError(t, err, want) + return + } + + if err != nil { + t.Fatalf("geminiImagesRequestFromOpenAI: %v", err) + } + want := geminiImagesReadGoldenJSON(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("geminiImagesRequestFromOpenAI(%s) mismatch:\ngot:\n%s\nwant:\n%s", name, gotB, wantB) + } + }) + } +} + +// TestGeminiImagesResponseGolden walks every "resp_*_in.json" fixture, +// translating it via openAIImagesResponseFromGemini with the fixed +// geminiImagesGoldenCreated, and asserts the result against the sibling +// "_want.json" fixture. +func TestGeminiImagesResponseGolden(t *testing.T) { + entries, err := os.ReadDir(geminiImagesGoldenTestdataDir) + 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", geminiImagesGoldenTestdataDir) + } + + for _, name := range names { + t.Run(name, func(t *testing.T) { + body, err := os.ReadFile(filepath.Join(geminiImagesGoldenTestdataDir, name+"_in.json")) + if err != nil { + t.Fatalf("read %s_in.json: %v", name, err) + } + + got, err := openAIImagesResponseFromGemini(body, geminiImagesGoldenCreated) + if err != nil { + t.Fatalf("openAIImagesResponseFromGemini: %v", err) + } + + want := geminiImagesReadGoldenJSON(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("openAIImagesResponseFromGemini(%s) mismatch:\ngot:\n%s\nwant:\n%s", name, gotB, wantB) + } + }) + } +} + +// TestGeminiImageAspectRatio_UnknownSizeErrors pins the boundary the +// golden walker's req_size_unknown fixture already exercises end to end, +// directly against geminiImageAspectRatio: an unrecognized, non-empty +// size string is a *translateError, not a silent fallback to "1:1". +func TestGeminiImageAspectRatio_UnknownSizeErrors(t *testing.T) { + _, err := geminiImageAspectRatio("640x480") + var terr *translateError + if err == nil { + t.Fatal("geminiImageAspectRatio(\"640x480\") = nil error, want *translateError") + } + if te, ok := err.(*translateError); ok { + terr = te + } else { + t.Fatalf("err = %v (%T), want *translateError", err, err) + } + if terr.notSupported { + t.Error("notSupported = true, want false (this is a 400, not a 501)") + } +}