traefik-llmgateway/provider_gemini.go

package traefikllmgateway

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// geminiAdapter is the providerAdapter for provider type "gemini".
type geminiAdapter struct {
	client *http.Client
	// retry is nil for every adapter built by a bare newGeminiAdapter call
	// (every v0.1 test, and any caller that never wires one) — nil-safe,
	// per retryPolicy.do's own contract, so leaving it unset is exactly
	// v0.1's one-attempt behavior. buildAdapters (providers.go) sets it
	// from cfg.Retry for production use.
	retry       *retryPolicy
	adapterName string
	baseURL     string
	apiKey      string
	// timeout is this adapter's resolved per-request timeout —
	// resolveProviderTimeout, timeout.go; see openaiAdapter.timeout's own
	// doc comment (provider_openai.go) for why this is set to
	// defaultRequestTimeout at construction rather than left at its zero
	// value the way retry is.
	timeout time.Duration
}

// 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(defaultRequestTimeout),
		timeout:     defaultRequestTimeout,
	}, 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 }

// httpClient implements providerAdapter.
func (a *geminiAdapter) httpClient() *http.Client { return a.client }

// requestTimeout implements providerAdapter.
func (a *geminiAdapter) requestTimeout() time.Duration { return a.timeout }

// 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)
	// responseModel (ruling a, ALIAS ECHO): gatewayAliasKey, when present,
	// is the exact id the client requested — echoed into the translated
	// response's "model" field, distinct from model (the bare upstream id),
	// which must keep driving the URL below unchanged. Deleted from req so
	// it cannot leak into the upstream request body.
	responseModel := model
	if alias, ok := req[gatewayAliasKey].(string); ok && alias != "" {
		responseModel = alias
	}
	delete(req, gatewayAliasKey)
	streaming, _ := req["stream"].(bool)
	// clientAskedUsage (M1 fix): stream_options.include_usage is an
	// OpenAI-wire field with no Gemini equivalent, so
	// geminiRequestFromOpenAI never forwards it upstream — it must be read
	// here, before translation, so forwardStream knows whether to
	// synthesize the final usage-only chunk OpenAI-compatible clients
	// expect.
	clientAskedUsage := false
	if streaming {
		if so, ok := req["stream_options"].(map[string]any); ok {
			clientAskedUsage = isTruthy(so["include_usage"])
		}
	}

	body, err := geminiRequestFromOpenAI(req)
	if err != nil {
		return usage{}, err
	}

	hdr := a.requestHeaders(true)
	// escapedModel guards against a model id smuggling extra path segments
	// or query syntax ("/", "?", "..") into the upstream URL — url.PathEscape
	// leaves ordinary ids (letters, digits, ":", ".", "-") untouched, per
	// RFC 3986's pchar grammar, so a normal model id is byte-identical after
	// escaping.
	escapedModel := url.PathEscape(model)
	endpoint := a.baseURL + geminiAPIPrefix + escapedModel + ":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")
		endpoint = a.baseURL + geminiAPIPrefix + escapedModel + ":streamGenerateContent?alt=sse"
	}

	resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, hdr, body, a.retry, a.timeout, a.adapterName)
	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, responseModel, clientAskedUsage)
	}
	// 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, responseModel)
}

// 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]".
//
// clientAskedUsage (M1 fix), when true, makes forwardStream write one
// extra chunk — st.usageChunk(), empty "choices" plus the accumulated
// usage — right before "[DONE]", the same final shape an OpenAI-type
// upstream sends when stream_options.include_usage is set
// (provider_openai.go's forwardStream). Gemini's own stream never carries
// such a chunk, so this adapter synthesizes it instead of relying on
// anything upstream.
func (a *geminiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gatewayModel string, clientAskedUsage bool) (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)
	}
	if clientAskedUsage {
		if werr := sw.writeData(st.usageChunk()); werr != nil {
			return st.usage(), fmt.Errorf("%w: stream: %w", errUpstream, werr)
		}
	}
	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)
	// responseModel (ruling a, ALIAS ECHO): see chatCompletion's identical
	// pattern above.
	responseModel := model
	if alias, ok := req[gatewayAliasKey].(string); ok && alias != "" {
		responseModel = alias
	}
	delete(req, gatewayAliasKey)

	texts, single, err := geminiEmbeddingInputTexts(req["input"])
	if err != nil {
		return usage{}, err
	}

	var body map[string]any
	var endpoint string
	escapedModel := url.PathEscape(model)
	if single {
		body = geminiEmbedContentRequest(texts[0])
		endpoint = a.baseURL + geminiAPIPrefix + escapedModel + ":embedContent"
	} else {
		// geminiBatchEmbedContentsRequest still takes the unescaped bare
		// model — it writes "models/{model}" into the JSON body, not a URL.
		body = geminiBatchEmbedContentsRequest(model, texts)
		endpoint = a.baseURL + geminiAPIPrefix + escapedModel + ":batchEmbedContents"
	}

	resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, a.requestHeaders(true), body, a.retry, a.timeout, a.adapterName)
	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, responseModel, 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
}

// 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, a.timeout, a.adapterName)
	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"), plus
// the cursor-pagination field (M7 fix) the endpoint adds on top —
// nextPageToken, paired with the pageToken query parameter the next
// page's request carries.
type geminiModelsListPayload struct {
	NextPageToken string `json:"nextPageToken"`
	Models        []struct {
		Name string `json:"name"`
	} `json:"models"`
}

// geminiListModelsPageCap bounds how many pages listModels follows (M7
// fix): a misbehaving upstream that always returns a non-empty
// nextPageToken would otherwise page forever. 20 pages covers a large
// catalog without risking an unbounded discovery loop against a single
// provider.
const geminiListModelsPageCap = 20

// geminiListModelsPageSize is the pageSize sent on every models.list page:
// the documented maximum (default 50, at most 1000 per page). Gemini
// requires the same parameters on every page of one listing, so it is sent
// with the pageToken too. One request at the maximum returns the whole
// catalog within registry.go's warmFillTimeout.
const geminiListModelsPageSize = 1000

// 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, following Gemini's nextPageToken cursor pagination (M7
// fix) via the pageToken query parameter, up to geminiListModelsPageCap
// pages.
func (a *geminiAdapter) listModels(ctx context.Context) ([]string, error) {
	var ids []string
	pageToken := ""
	for page := 0; page < geminiListModelsPageCap; page++ {
		pageIDs, nextPageToken, err := a.listModelsPage(ctx, pageToken)
		if err != nil {
			return nil, err
		}
		ids = append(ids, pageIDs...)
		if nextPageToken == "" {
			break
		}
		pageToken = nextPageToken
	}
	return ids, nil
}

// listModelsPage fetches one page of Gemini's /v1beta/models, starting
// at pageToken (empty for the first page), and returns that page's bare
// model ids alongside nextPageToken for listModels' own pagination loop
// above.
func (a *geminiAdapter) listModelsPage(ctx context.Context, pageToken string) (ids []string, nextPageToken string, err error) {
	query := url.Values{}
	query.Set("pageSize", strconv.Itoa(geminiListModelsPageSize))
	if pageToken != "" {
		query.Set("pageToken", pageToken)
	}
	endpoint := a.baseURL + "/v1beta/models?" + query.Encode()
	resp, err := upstreamJSON(ctx, a.client, http.MethodGet, endpoint, a.requestHeaders(false), nil, a.retry, a.timeout, a.adapterName)
	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, payload.NextPageToken, nil
}