From 6275daa702c4096ea06dd72e4f1f0e61e9aa5f61 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 03:06:08 +0100 Subject: [PATCH] feat: provider adapter interface and openai-type adapter with streaming usage capture --- provider_openai.go | 244 +++++++++++++++++++++ provider_openai_test.go | 462 ++++++++++++++++++++++++++++++++++++++++ providers.go | 196 +++++++++++++++++ 3 files changed, 902 insertions(+) create mode 100644 provider_openai.go create mode 100644 provider_openai_test.go create mode 100644 providers.go diff --git a/provider_openai.go b/provider_openai.go new file mode 100644 index 0000000..9103be1 --- /dev/null +++ b/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 +} diff --git a/provider_openai_test.go b/provider_openai_test.go new file mode 100644 index 0000000..d371cf6 --- /dev/null +++ b/provider_openai_test.go @@ -0,0 +1,462 @@ +package traefikllmgateway + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "reflect" + "testing" + "time" +) + +// captureBodyHandler wraps fn, first draining and recording the request +// body (and, if capturedAuth is non-nil, the Authorization header) so a +// test can assert on what the adapter actually sent upstream. +func captureBodyHandler(capturedBody *[]byte, capturedAuth *string, fn http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + if capturedBody != nil { + *capturedBody = b + } + if capturedAuth != nil { + *capturedAuth = r.Header.Get("Authorization") + } + fn(w, r) + } +} + +// decodeJSONBody unmarshals b into a map for structural assertions, +// avoiding brittleness from json.Marshal's key ordering. +func decodeJSONBody(t *testing.T, b []byte) map[string]any { + t.Helper() + var m map[string]any + if err := json.Unmarshal(b, &m); err != nil { + t.Fatalf("decode captured request body: %v (body=%s)", err, b) + } + return m +} + +// TestOpenAIAdapter_ChatCompletion_NonStreaming proves the non-streaming +// path copies the upstream's Content-Type, status, and body verbatim to +// the client, extracts prompt/completion usage from the body, and sends no +// stream_options since the request never asked to stream. +func TestOpenAIAdapter_ChatCompletion_NonStreaming(t *testing.T) { + const respBody = `{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":10,"completion_tokens":20}}` + + var gotBody []byte + var gotAuth string + srv := httptest.NewServer(captureBodyHandler(&gotBody, &gotAuth, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/chat/completions" { + t.Errorf("path = %q, want /v1/chat/completions", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + rec := httptest.NewRecorder() + + req := map[string]any{"model": "gpt-5", "messages": []any{}} + u, err := a.chatCompletion(context.Background(), rec, req) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + + if u.prompt != 10 || u.completion != 20 { + t.Errorf("usage = %+v, want {prompt:10 completion:20}", u) + } + if rec.Code != http.StatusOK { + t.Errorf("status = %d, want 200", rec.Code) + } + if got := rec.Header().Get("Content-Type"); got != "application/json" { + t.Errorf("Content-Type = %q, want application/json", got) + } + if rec.Body.String() != respBody { + t.Errorf("body = %q, want %q (verbatim passthrough)", rec.Body.String(), respBody) + } + if gotAuth != "Bearer sk-test" { + t.Errorf("Authorization = %q, want %q", gotAuth, "Bearer sk-test") + } + + sent := decodeJSONBody(t, gotBody) + if _, present := sent["stream_options"]; present { + t.Errorf("stream_options sent on a non-streaming request: %v", sent["stream_options"]) + } +} + +// TestOpenAIAdapter_Embeddings proves embeddings uses the same +// verbatim-passthrough forwarding as chatCompletion's non-streaming path, +// against /v1/embeddings, extracting only prompt_tokens. +func TestOpenAIAdapter_Embeddings(t *testing.T) { + const respBody = `{"data":[{"embedding":[0.1,0.2]}],"usage":{"prompt_tokens":5}}` + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/embeddings" { + t.Errorf("path = %q, want /v1/embeddings", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + rec := httptest.NewRecorder() + + u, err := a.embeddings(context.Background(), rec, map[string]any{"model": "text-embedding-3", "input": "hello"}) + if err != nil { + t.Fatalf("embeddings: %v", err) + } + if u.prompt != 5 || u.completion != 0 { + t.Errorf("usage = %+v, want {prompt:5 completion:0}", u) + } + if rec.Body.String() != respBody { + t.Errorf("body = %q, want %q", rec.Body.String(), respBody) + } +} + +// streamFrames are the four SSE data payloads a fake streaming upstream +// emits for the tests below: three content chunks, then a final +// choices-empty usage chunk. +var streamFrames = []string{ + `{"id":"1","choices":[{"index":0,"delta":{"content":"Hel"}}]}`, + `{"id":"1","choices":[{"index":0,"delta":{"content":"lo"}}]}`, + `{"id":"1","choices":[{"index":0,"delta":{"content":"!"}}]}`, + `{"id":"1","choices":[],"usage":{"prompt_tokens":7,"completion_tokens":9}}`, +} + +// newStreamingFakeUpstream returns an httptest.Server whose +// /v1/chat/completions handler records the request body, then emits +// streamFrames as SSE, flushing and sleeping 50ms between each, and +// finally "[DONE]". +func newStreamingFakeUpstream(t *testing.T, capturedBody *[]byte) *httptest.Server { + t.Helper() + return httptest.NewServer(captureBodyHandler(capturedBody, nil, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + f, ok := w.(http.Flusher) + if !ok { + t.Fatalf("httptest ResponseWriter does not implement http.Flusher") + } + for _, frame := range streamFrames { + fmt.Fprintf(w, "data: %s\n\n", frame) + f.Flush() + time.Sleep(50 * time.Millisecond) + } + fmt.Fprint(w, "data: [DONE]\n\n") + f.Flush() + })) +} + +// TestOpenAIAdapter_ChatCompletion_Streaming_ClientDidNotAskUsage proves +// that when the client streams without setting stream_options, +// chatCompletion injects include_usage upstream on its own, forwards only +// the 3 content chunks plus [DONE] to the client (suppressing the +// usage-only chunk), still captures the usage it carried, and flushes at +// least once per forwarded event. +func TestOpenAIAdapter_ChatCompletion_Streaming_ClientDidNotAskUsage(t *testing.T) { + var gotBody []byte + srv := newStreamingFakeUpstream(t, &gotBody) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + rec := newRecordingResponseWriter() + + req := map[string]any{"model": "gpt-5", "stream": true} + u, err := a.chatCompletion(context.Background(), rec, req) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + + if u.prompt != 7 || u.completion != 9 { + t.Errorf("usage = %+v, want {prompt:7 completion:9}", u) + } + + want := "data: " + streamFrames[0] + "\n\n" + + "data: " + streamFrames[1] + "\n\n" + + "data: " + streamFrames[2] + "\n\n" + + "data: [DONE]\n\n" + if got := rec.body.String(); got != want { + t.Errorf("forwarded body = %q, want %q (usage chunk must be suppressed)", got, want) + } + if rec.flushed < 4 { + t.Errorf("flushed = %d, want >= 4", rec.flushed) + } + + sent := decodeJSONBody(t, gotBody) + so, ok := sent["stream_options"].(map[string]any) + if !ok { + t.Fatalf("upstream request missing injected stream_options: %v", sent) + } + if so["include_usage"] != true { + t.Errorf("stream_options.include_usage = %v, want true", so["include_usage"]) + } +} + +// TestOpenAIAdapter_ChatCompletion_Streaming_ClientAskedUsage proves that +// when the client already set stream_options itself, chatCompletion leaves +// it untouched (does not overwrite it) and forwards the usage-only chunk +// through to the client instead of suppressing it. +func TestOpenAIAdapter_ChatCompletion_Streaming_ClientAskedUsage(t *testing.T) { + var gotBody []byte + srv := newStreamingFakeUpstream(t, &gotBody) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + rec := newRecordingResponseWriter() + + req := map[string]any{ + "model": "gpt-5", + "stream": true, + "stream_options": map[string]any{"include_usage": true}, + } + u, err := a.chatCompletion(context.Background(), rec, req) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + + if u.prompt != 7 || u.completion != 9 { + t.Errorf("usage = %+v, want {prompt:7 completion:9}", u) + } + + want := "data: " + streamFrames[0] + "\n\n" + + "data: " + streamFrames[1] + "\n\n" + + "data: " + streamFrames[2] + "\n\n" + + "data: " + streamFrames[3] + "\n\n" + + "data: [DONE]\n\n" + if got := rec.body.String(); got != want { + t.Errorf("forwarded body = %q, want %q (usage chunk must be forwarded)", got, want) + } + if rec.flushed < 5 { + t.Errorf("flushed = %d, want >= 5", rec.flushed) + } + + sent := decodeJSONBody(t, gotBody) + so, ok := sent["stream_options"].(map[string]any) + if !ok { + t.Fatalf("upstream request lost the client's stream_options: %v", sent) + } + if so["include_usage"] != true { + t.Errorf("stream_options.include_usage = %v, want true (must not be overwritten)", so["include_usage"]) + } +} + +// TestOpenAIAdapter_ChatCompletion_Upstream429 proves a non-2xx upstream +// response returns a *providerHTTPError carrying the status and body, +// with nothing written to the client — the caller decides envelope vs +// passthrough. +func TestOpenAIAdapter_ChatCompletion_Upstream429(t *testing.T) { + const errBody = `{"error":{"message":"rate limited","type":"rate_limit_error"}}` + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(errBody)) + })) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + rec := httptest.NewRecorder() + + u, err := a.chatCompletion(context.Background(), rec, map[string]any{"model": "gpt-5"}) + if err == nil { + t.Fatalf("chatCompletion: want error, got nil (usage=%+v)", u) + } + if u != (usage{}) { + t.Errorf("usage = %+v, want zero value on error", u) + } + + var httpErr *providerHTTPError + if !errors.As(err, &httpErr) { + t.Fatalf("err = %v (%T), want *providerHTTPError", err, err) + } + if httpErr.status != http.StatusTooManyRequests { + t.Errorf("status = %d, want %d", httpErr.status, http.StatusTooManyRequests) + } + if string(httpErr.body) != errBody { + t.Errorf("body = %q, want %q", httpErr.body, errBody) + } + if rec.Body.Len() != 0 { + t.Errorf("client response body = %q, want empty (caller decides passthrough)", rec.Body.String()) + } +} + +// TestOpenAIAdapter_ListModels proves listModels parses the OpenAI +// /v1/models "data": [{"id": ...}] fixture into a flat slice of ids. +func TestOpenAIAdapter_ListModels(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Errorf("path = %q, want /v1/models", r.URL.Path) + } + if r.Method != http.MethodGet { + t.Errorf("method = %q, want GET", r.Method) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"id":"gpt-5"},{"id":"gpt-5-mini"}]}`)) + })) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + got, err := a.listModels(context.Background()) + if err != nil { + t.Fatalf("listModels: %v", err) + } + + want := []string{"gpt-5", "gpt-5-mini"} + if !reflect.DeepEqual(got, want) { + t.Errorf("got %#v, want %#v", got, want) + } +} + +// TestOpenAIAdapter_Keyless proves a resolved-empty API key sends no +// Authorization header at all, not one with an empty token — the +// operator's own keyless gateway is a real deployment target. +func TestOpenAIAdapter_Keyless(t *testing.T) { + var gotAuth string + sawHeader := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + _, sawHeader = r.Header["Authorization"] + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"1","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}`)) + })) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "") + rec := httptest.NewRecorder() + + u, err := a.chatCompletion(context.Background(), rec, map[string]any{"model": "gpt-5"}) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + if u.prompt != 1 || u.completion != 2 { + t.Errorf("usage = %+v, want {prompt:1 completion:2}", u) + } + if sawHeader { + t.Errorf("Authorization header present (%q), want absent for a keyless adapter", gotAuth) + } +} + +// TestBuildAdapters covers buildAdapters' config-resolution rules: secret +// resolution and propagation, base URL defaulting and trailing-slash +// trimming, keyless as valid (not an error), and the unknown/not-yet- +// implemented type errors. +func TestBuildAdapters(t *testing.T) { + t.Run("good: openai with literal key defaults base URL", func(t *testing.T) { + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "openai", APIKey: "sk-test"}, + }} + adapters, err := buildAdapters(cfg) + if err != nil { + t.Fatalf("buildAdapters: %v", err) + } + a, ok := adapters["p1"] + if !ok { + t.Fatalf("adapters[%q] missing", "p1") + } + if a.name() != "p1" { + t.Errorf("name() = %q, want %q", a.name(), "p1") + } + if a.typeName() != "openai" { + t.Errorf("typeName() = %q, want %q", a.typeName(), "openai") + } + if a.base() != defaultBaseOpenAI { + t.Errorf("base() = %q, want %q", a.base(), defaultBaseOpenAI) + } + + req := httptest.NewRequest(http.MethodPost, "/", nil) + a.injectAuth(req) + if got := req.Header.Get("Authorization"); got != "Bearer sk-test" { + t.Errorf("Authorization = %q, want %q", got, "Bearer sk-test") + } + }) + + t.Run("good: explicit base URL trailing slash trimmed", func(t *testing.T) { + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "openai", BaseURL: "https://custom.example.com/", APIKey: "sk-test"}, + }} + adapters, err := buildAdapters(cfg) + if err != nil { + t.Fatalf("buildAdapters: %v", err) + } + if got := adapters["p1"].base(); got != "https://custom.example.com" { + t.Errorf("base() = %q, want %q", got, "https://custom.example.com") + } + }) + + t.Run("good: empty apiKey is keyless, not an error", func(t *testing.T) { + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "openai"}, + }} + adapters, err := buildAdapters(cfg) + if err != nil { + t.Fatalf("buildAdapters: %v", err) + } + req := httptest.NewRequest(http.MethodPost, "/", nil) + adapters["p1"].injectAuth(req) + if _, present := req.Header["Authorization"]; present { + t.Errorf("Authorization header present for a keyless provider") + } + }) + + t.Run("good: env: apiKey resolves via resolveSecret", func(t *testing.T) { + t.Setenv("LLMGW_T2", "sk-from-env") + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "openai", APIKey: "env:LLMGW_T2"}, //nolint:gosec // "env:NAME" is a resolveSecret indirection, not a credential + }} + adapters, err := buildAdapters(cfg) + if err != nil { + t.Fatalf("buildAdapters: %v", err) + } + req := httptest.NewRequest(http.MethodPost, "/", nil) + adapters["p1"].injectAuth(req) + if got := req.Header.Get("Authorization"); got != "Bearer sk-from-env" { + t.Errorf("Authorization = %q, want %q", got, "Bearer sk-from-env") + } + }) + + t.Run("bad: env: apiKey unset propagates resolveSecret's error", func(t *testing.T) { + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "openai", APIKey: "env:LLMGW_MISSING2"}, //nolint:gosec // "env:NAME" is a resolveSecret indirection, not a credential + }} + if _, err := buildAdapters(cfg); err == nil { + t.Fatalf("buildAdapters: want error for unset env var, got nil") + } + }) + + t.Run("bad: unknown provider type", func(t *testing.T) { + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "not-a-real-type"}, + }} + _, err := buildAdapters(cfg) + if err == nil { + t.Fatalf("buildAdapters: want error for unknown type, got nil") + } + }) + + t.Run("edge: anthropic type not yet implemented", func(t *testing.T) { + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "anthropic", APIKey: "sk-test"}, + }} + _, err := buildAdapters(cfg) + if err == nil { + t.Fatalf("buildAdapters: want placeholder error for anthropic, got nil") + } + }) + + t.Run("edge: gemini type not yet implemented", func(t *testing.T) { + cfg := &Config{Providers: map[string]*ProviderConfig{ + "p1": {Type: "gemini", APIKey: "sk-test"}, + }} + _, err := buildAdapters(cfg) + if err == nil { + t.Fatalf("buildAdapters: want placeholder error for gemini, got nil") + } + }) +} diff --git a/providers.go b/providers.go new file mode 100644 index 0000000..dbd354c --- /dev/null +++ b/providers.go @@ -0,0 +1,196 @@ +package traefikllmgateway + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" +) + +// Provider type identifiers, matched against ProviderConfig.Type. +const ( + providerTypeOpenAI = "openai" + providerTypeAnthropic = "anthropic" + providerTypeGemini = "gemini" +) + +// Default base URLs applied by buildAdapters when a ProviderConfig omits +// BaseURL, keyed by provider type. +const ( + defaultBaseOpenAI = "https://api.openai.com" + defaultBaseAnthropic = "https://api.anthropic.com" + defaultBaseGemini = "https://generativelanguage.googleapis.com" +) + +// defaultBaseURLByType maps a provider type to its default base URL. It is +// read-only after package initialization, the same pattern pricing.go uses +// for builtinPricing. +var defaultBaseURLByType = map[string]string{ + providerTypeOpenAI: defaultBaseOpenAI, + providerTypeAnthropic: defaultBaseAnthropic, + providerTypeGemini: defaultBaseGemini, +} + +// maxResponseBytes caps how much of a non-streaming upstream response body +// an adapter reads into memory: 32MiB, generous for a chat completion or +// embeddings JSON body while bounding memory against a misbehaving or +// oversized upstream response. +const maxResponseBytes = 32 << 20 + +// adapterIdleConnsPerHost bounds each adapter's shared http.Client's idle +// connection pool for its single upstream host, so repeated requests reuse +// connections instead of paying a fresh TLS handshake per call. +const adapterIdleConnsPerHost = 16 + +// providerHTTPError reports a non-2xx response from an upstream provider. +// Adapters return it instead of writing the error to the client themselves, +// so the caller — the unified route in a later task — decides whether to +// pass the body through verbatim or wrap it in the gateway's own error +// envelope. +type providerHTTPError struct { + body []byte + status int +} + +// Error implements the error interface. It deliberately omits body: the +// caller that inspects providerHTTPError.body decides what, if anything, of +// the upstream response to surface. +func (e *providerHTTPError) Error() string { + return fmt.Sprintf("provider returned status %d", e.status) +} + +// newProviderHTTPError builds a providerHTTPError from a non-2xx upstream +// response, reading its body (capped at maxResponseBytes) before the caller +// closes it. +func newProviderHTTPError(resp *http.Response) *providerHTTPError { + body, _ := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) + return &providerHTTPError{status: resp.StatusCode, body: body} +} + +// errUpstream is a sentinel wrapped into errors an adapter returns when it +// cannot reach, or cannot make sense of a response from, its upstream — +// distinct from providerHTTPError, which means the upstream answered, just +// with a non-2xx status. +var errUpstream = errors.New("llmgateway: upstream request failed") + +// providerAdapter is the gateway's uniform interface over one configured +// upstream LLM provider. Task 9 (anthropic) and Task 10 (gemini) implement +// it alongside this task's openai-type adapter; Tasks 11-12 wire a built +// map of these into the request routing path. +type providerAdapter interface { + // name returns the provider's configured name (the key under + // Config.Providers). + name() string + // typeName returns the provider's configured type, e.g. "openai". + typeName() string + // base returns the provider's resolved base URL, trailing slash + // trimmed. + base() string + // injectAuth sets whatever headers r needs to authenticate against the + // provider. It is a no-op when the provider's resolved API key is + // empty — a keyless upstream is a valid configuration, not an error. + injectAuth(r *http.Request) + // chatCompletion proxies req to the provider's chat completions + // endpoint and writes the (possibly streamed) response to w. req has + // already had its "model" field rewritten to the provider's own model + // id by the caller. On a non-2xx upstream response, chatCompletion + // writes nothing to w and returns a *providerHTTPError instead. + chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) + // embeddings mirrors chatCompletion for the provider's embeddings + // endpoint. It never streams. + 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) +} + +// newAdapterHTTPClient returns an *http.Client for a provider adapter to +// hold and reuse across every request it makes: no client-level timeout, +// since a streaming chat completion can legitimately run for minutes and +// callers bound requests with a context deadline instead, and a Transport +// cloned from http.DefaultTransport with a larger per-host idle connection +// pool, since every request from one adapter targets the same upstream +// host. +func newAdapterHTTPClient() *http.Client { + tr := http.DefaultTransport.(*http.Transport).Clone() //nolint:forcetypeassert // http.DefaultTransport is always *http.Transport + tr.MaxIdleConnsPerHost = adapterIdleConnsPerHost + return &http.Client{Transport: tr} +} + +// upstreamJSON issues an HTTP request to url: body, when non-nil, is +// marshaled as the JSON request body; hdr's values are added to the +// request (a caller builds this from its adapter's auth and content-type +// headers). It returns the raw *http.Response for the caller to inspect +// the status code and read or stream the body from — closing resp.Body is +// the caller's responsibility. A failure to encode the body, build the +// request, or reach the upstream returns an error wrapping errUpstream; +// a non-2xx status is not itself an error here, since some callers (e.g. +// listModels) have nothing to write through and every current caller +// checks resp.StatusCode itself. +func upstreamJSON(ctx context.Context, client *http.Client, method, url string, hdr http.Header, body any) (*http.Response, error) { + var r io.Reader + if body != nil { + b, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("%w: encode request body: %v", errUpstream, err) //nolint:errorlint // errUpstream is a sentinel wrapped for errors.Is, %v is fine for the inner cause + } + r = bytes.NewReader(b) + } + + req, err := http.NewRequestWithContext(ctx, method, url, r) + if err != nil { + return nil, fmt.Errorf("%w: build request: %v", errUpstream, err) //nolint:errorlint + } + for k, vs := range hdr { + for _, v := range vs { + req.Header.Add(k, v) + } + } + if body != nil && req.Header.Get("Content-Type") == "" { + 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 + if err != nil { + return nil, fmt.Errorf("%w: %v", errUpstream, err) //nolint:errorlint + } + return resp, nil +} + +// buildAdapters resolves cfg.Providers into a map of providerAdapter keyed +// by provider name. For each provider it resolves APIKey via resolveSecret +// (propagating that error, but treating a resolved empty key as a valid +// keyless upstream rather than an error) and defaults BaseURL by provider +// type when the config omits one, trimming any trailing slash either way. +// An unknown provider type is a config error. A configured anthropic or +// gemini provider currently errors too — Tasks 9 and 10 fill in those +// switch cases; every other case here is unaffected by that. +func buildAdapters(cfg *Config) (map[string]providerAdapter, error) { + adapters := make(map[string]providerAdapter, len(cfg.Providers)) + for name, pc := range cfg.Providers { + apiKey, err := resolveSecret(pc.APIKey) + if err != nil { + return nil, fmt.Errorf("llmgateway: provider %q: %w", name, err) + } + + base := strings.TrimSuffix(pc.BaseURL, "/") + if base == "" { + base = defaultBaseURLByType[pc.Type] + } + + switch pc.Type { + case providerTypeOpenAI: + adapters[name] = newOpenAIAdapter(name, base, apiKey) + case providerTypeAnthropic: + return nil, fmt.Errorf("llmgateway: provider %q: type %q not yet implemented", name, pc.Type) + case providerTypeGemini: + return nil, fmt.Errorf("llmgateway: provider %q: type %q not yet implemented", name, pc.Type) + default: + return nil, fmt.Errorf("llmgateway: provider %q: unknown type %q", name, pc.Type) + } + } + return adapters, nil +}