feat: gemini adapter with bidirectional translation incl streaming and embeddings

2026-10-20 05:56:56 +00:00 · 2026-08-20 04:16:10 +01:00
parent 876d340054
commit 2891766b5f
48 changed files with 2463 additions and 7 deletions
@@ -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
+}