mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-18 05:47:38 +00:00
feat: provider adapter interface and openai-type adapter with streaming usage capture
This commit is contained in:
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user