feat: provider adapter interface and openai-type adapter with streaming usage capture

provider_openai.go

@@ -0,0 +1,244 @@
+package traefikllmgateway
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+)
+
+// openaiAdapter is the providerAdapter for provider type "openai": OpenAI
+// itself, and any OpenAI-compatible upstream (the operator's own gateway
+// included — a real deployment target for this adapter is keyless).
+type openaiAdapter struct {
+	client      *http.Client
+	adapterName string
+	baseURL     string
+	apiKey      string
+}
+
+// newOpenAIAdapter returns an openaiAdapter for provider name, with base as
+// its already-defaulted, trailing-slash-trimmed base URL and apiKey as
+// already resolved by resolveSecret (empty means keyless). It builds one
+// shared *http.Client, reused for every request this adapter makes.
+func newOpenAIAdapter(name, base, apiKey string) *openaiAdapter {
+	return &openaiAdapter{
+		adapterName: name,
+		baseURL:     base,
+		apiKey:      apiKey,
+		client:      newAdapterHTTPClient(),
+	}
+}
+
+// name implements providerAdapter.
+func (a *openaiAdapter) name() string { return a.adapterName }
+
+// typeName implements providerAdapter.
+func (a *openaiAdapter) typeName() string { return providerTypeOpenAI }
+
+// base implements providerAdapter.
+func (a *openaiAdapter) base() string { return a.baseURL }
+
+// injectAuth implements providerAdapter: OpenAI's bearer-token scheme. A
+// no-op when a.apiKey is empty, per the keyless-upstream ruling — a request
+// against a keyless upstream must carry no Authorization header at all,
+// not one with an empty token.
+func (a *openaiAdapter) injectAuth(r *http.Request) {
+	if a.apiKey == "" {
+		return
+	}
+	r.Header.Set("Authorization", "Bearer "+a.apiKey)
+}
+
+// requestHeaders returns the headers for one upstream request: a
+// Content-Type when the request carries a JSON body, plus whatever
+// injectAuth adds. Built through injectAuth (rather than setting
+// Authorization directly) so this adapter's outgoing requests exercise the
+// same auth path the providerAdapter interface exposes to callers.
+func (a *openaiAdapter) requestHeaders(hasBody bool) http.Header {
+	h := http.Header{}
+	if hasBody {
+		h.Set("Content-Type", "application/json")
+	}
+	a.injectAuth(&http.Request{Header: h})
+	return h
+}
+
+// chatStreamUsage is the "usage" object OpenAI's streaming chat completion
+// chunks carry: null on every content chunk, populated only on the final,
+// choices-empty chunk when the client (or this adapter, see chatCompletion)
+// requested it via stream_options.include_usage.
+type chatStreamUsage struct {
+	PromptTokens     int64 `json:"prompt_tokens"`
+	CompletionTokens int64 `json:"completion_tokens"`
+}
+
+// chatStreamChunk is the subset of one SSE chunk's JSON payload
+// chatCompletion's streaming path needs: whether it carries content
+// (Choices non-empty) and, on the final chunk, its usage totals.
+type chatStreamChunk struct {
+	Usage   *chatStreamUsage  `json:"usage"`
+	Choices []json.RawMessage `json:"choices"`
+}
+
+// chatCompletion implements providerAdapter.
+//
+// When req["stream"] is true, chatCompletion injects
+// stream_options.include_usage=true into the upstream request only when
+// the client did not already set stream_options itself — remembering
+// whether the client asked, so the streaming response can suppress the
+// resulting usage-only chunk from a client that never asked for it while
+// still capturing its usage for accounting.
+func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
+	streaming, _ := req["stream"].(bool)
+	clientAskedUsage := false
+	if streaming {
+		if _, present := req["stream_options"]; present {
+			clientAskedUsage = true
+		} else {
+			req["stream_options"] = map[string]any{"include_usage": true}
+		}
+	}
+
+	resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/chat/completions", a.requestHeaders(true), req)
+	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 {
+		return a.forwardStream(w, resp.Body, clientAskedUsage)
+	}
+	return a.forwardJSON(w, resp)
+}
+
+// embeddings implements providerAdapter. Embeddings responses never
+// stream, so it is forwardJSON end to end, same as chatCompletion's
+// non-streaming path.
+func (a *openaiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
+	resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/embeddings", a.requestHeaders(true), req)
+	if err != nil {
+		return usage{}, err
+	}
+	defer resp.Body.Close() //nolint:errcheck
+
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return usage{}, newProviderHTTPError(resp)
+	}
+	return a.forwardJSON(w, resp)
+}
+
+// chatUsagePayload extracts the fields forwardJSON needs from a
+// non-streaming response body: OpenAI's chat completion and embeddings
+// responses both carry a top-level "usage" object, embeddings' simply
+// omitting completion_tokens (it decodes as its zero value).
+type chatUsagePayload struct {
+	Usage chatStreamUsage `json:"usage"`
+}
+
+// forwardJSON reads resp's full body (capped at maxResponseBytes),
+// extracts its usage for accounting, and writes resp's Content-Type,
+// status code, and body verbatim to w — the non-streaming passthrough both
+// chatCompletion and embeddings use. A body that fails to decode as JSON
+// still passes through unchanged; forwardJSON only fails on a read error,
+// since a malformed-but-readable body is the upstream's problem to answer
+// for, not a reason to swallow the response.
+func (a *openaiAdapter) forwardJSON(w http.ResponseWriter, resp *http.Response) (usage, error) {
+	body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
+	if err != nil {
+		return usage{}, fmt.Errorf("%w: read response body: %v", errUpstream, err) //nolint:errorlint
+	}
+
+	var payload chatUsagePayload
+	_ = json.Unmarshal(body, &payload) // best-effort usage extraction; body still forwards verbatim below
+
+	if ct := resp.Header.Get("Content-Type"); ct != "" {
+		w.Header().Set("Content-Type", ct)
+	}
+	w.WriteHeader(resp.StatusCode)
+	_, _ = w.Write(body)
+
+	return usage{prompt: payload.Usage.PromptTokens, completion: payload.Usage.CompletionTokens}, nil
+}
+
+// sseDoneMarker is the exact data payload readSSE hands forwardStream for
+// the OpenAI streaming convention's terminal event.
+const sseDoneMarker = "[DONE]"
+
+// forwardStream reads body as an upstream text/event-stream and forwards
+// each event to w via a fresh sseWriter, flushing per event. The final
+// usage-only chunk (empty "choices", non-null "usage") is captured into
+// the returned usage either way, but only forwarded to the client when
+// clientAskedUsage is true — otherwise it is dropped, since the client
+// never asked for it and chatCompletion only requested it upstream for
+// accounting. The "[DONE]" sentinel is always forwarded.
+func (a *openaiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, clientAskedUsage bool) (usage, error) {
+	sw := newSSEWriter(w)
+	var u usage
+
+	err := readSSE(body, func(ev sseEvent) error {
+		if string(ev.data) == sseDoneMarker {
+			sw.writeDone()
+			return nil
+		}
+
+		var chunk chatStreamChunk
+		hasUsage := json.Unmarshal(ev.data, &chunk) == nil && chunk.Usage != nil
+		if hasUsage {
+			u.prompt = chunk.Usage.PromptTokens
+			u.completion = chunk.Usage.CompletionTokens
+		}
+
+		isUsageOnlyChunk := hasUsage && len(chunk.Choices) == 0
+		if isUsageOnlyChunk && !clientAskedUsage {
+			return nil
+		}
+		return sw.writeData(ev.data)
+	})
+	if err != nil {
+		return usage{}, fmt.Errorf("%w: stream: %v", errUpstream, err) //nolint:errorlint
+	}
+	return u, nil
+}
+
+// modelsPayload is the "data": [{"id": ...}, ...] shape OpenAI's
+// /v1/models endpoint returns.
+type modelsPayload struct {
+	Data []struct {
+		ID string `json:"id"`
+	} `json:"data"`
+}
+
+// listModels implements providerAdapter.
+func (a *openaiAdapter) listModels(ctx context.Context) ([]string, error) {
+	resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+"/v1/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: %v", errUpstream, err) //nolint:errorlint
+	}
+
+	var payload modelsPayload
+	if err := json.Unmarshal(body, &payload); err != nil {
+		return nil, fmt.Errorf("%w: decode models response: %v", errUpstream, err) //nolint:errorlint
+	}
+
+	ids := make([]string, len(payload.Data))
+	for i, d := range payload.Data {
+		ids[i] = d.ID
+	}
+	return ids, nil
+}