mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-15 05:29:28 +00:00
feat: unified chat and embeddings endpoints with limits and accounting
Implements Task 12: POST /v1/chat/completions and POST /v1/embeddings
through a shared pipeline (routes_unified.go: handleChat, handleEmbeddings,
runUnified) wired into ServeHTTP's dispatch in llmgateway.go.
Pipeline: decode (10MiB cap) -> resolve model -> enforce user/group
limits -> call adapter -> account usage, translating every failure mode
into its HTTP envelope per the task brief's carried rulings:
- ALIAS ECHO (ruling a): the client-requested model id is threaded to the
adapter via the internal gatewayAliasKey convention and echoed back in
anthropic/gemini's translated response model field; every adapter
strips the key before building its upstream wire body so it never
leaks to a real provider.
- Limit violations (ruling b): storeDown -> 503 server_error; otherwise
429 rate_limit_error with Retry-After when the violation names a
window.
- Stream-safe error handling (ruling c): statusTrackingWriter.wroteHeader
gates whether an adapter error still gets an envelope; providerHTTPError
wraps into {error:{message,type,code,upstream}}; translateError maps to
400/501; context.Canceled logs quietly; anything else is a 502.
- Usage/accounting (ruling d): zero non-streaming usage gets a
ceil(len(body)/4) prompt-token estimate; zero streaming usage logs once
and accounts the request only. unifiedCostMicros extends costMicros'
single-id lookup to canonical-then-bare without changing its signature.
- Scopes (ruling e): user scope only when u.limits is set, group scope
only when grp.limits is set, user first.
routes_unified_test.go drives every case end-to-end through
Gateway.ServeHTTP against httptest fake upstreams (openai- and
anthropic-shaped): happy path with counters asserted via the limiter's
own store, incremental stream flushing, 401/403/404/429/503, embeddings,
alias echo through a provider-prefixed anthropic request, wrapped
upstream error envelopes, and a mid-stream drop that must not append a
second envelope to an already-committed response.
Gates: gofmt, go vet, fieldalignment, golangci-lint (incl. gosec) all
clean; go test ./... -race -count=1 passes.
This commit is contained in:
@@ -254,6 +254,20 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if r.Method == http.MethodPost && (r.URL.Path == "/v1/chat/completions" || r.URL.Path == "/v1/embeddings") {
|
||||||
|
u, grp, ok := g.auth.identify(r)
|
||||||
|
if !ok {
|
||||||
|
writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r.URL.Path == "/v1/chat/completions" {
|
||||||
|
g.handleChat(sw, r, u, grp)
|
||||||
|
} else {
|
||||||
|
g.handleEmbeddings(sw, r, u, grp)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
if g.cfg.PassthroughUnknown {
|
if g.cfg.PassthroughUnknown {
|
||||||
g.next.ServeHTTP(sw, r)
|
g.next.ServeHTTP(sw, r)
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -81,6 +81,14 @@ func (a *anthropicAdapter) requestHeaders(hasBody bool) http.Header {
|
|||||||
// request is made.
|
// request is made.
|
||||||
func (a *anthropicAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
func (a *anthropicAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
||||||
gatewayModel, _ := req["model"].(string)
|
gatewayModel, _ := req["model"].(string)
|
||||||
|
// gatewayAliasKey (ruling a, ALIAS ECHO), when present, is the exact id
|
||||||
|
// the client requested — echoed into the translated response's "model"
|
||||||
|
// field below instead of the bare upstream id, then deleted so it
|
||||||
|
// cannot leak into the upstream request body.
|
||||||
|
if alias, ok := req[gatewayAliasKey].(string); ok && alias != "" {
|
||||||
|
gatewayModel = alias
|
||||||
|
}
|
||||||
|
delete(req, gatewayAliasKey)
|
||||||
streaming, _ := req["stream"].(bool)
|
streaming, _ := req["stream"].(bool)
|
||||||
|
|
||||||
body, err := anthropicRequestFromOpenAI(req)
|
body, err := anthropicRequestFromOpenAI(req)
|
||||||
|
|||||||
+20
-3
@@ -84,6 +84,16 @@ func geminiModelFromRequest(req map[string]any) string {
|
|||||||
// request is made.
|
// request is made.
|
||||||
func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
||||||
model := geminiModelFromRequest(req)
|
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)
|
streaming, _ := req["stream"].(bool)
|
||||||
|
|
||||||
body, err := geminiRequestFromOpenAI(req)
|
body, err := geminiRequestFromOpenAI(req)
|
||||||
@@ -117,11 +127,11 @@ func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite
|
|||||||
}
|
}
|
||||||
|
|
||||||
if streaming && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
|
if streaming && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
|
||||||
return a.forwardStream(w, resp.Body, model)
|
return a.forwardStream(w, resp.Body, responseModel)
|
||||||
}
|
}
|
||||||
// Either a non-streaming request, or a streaming one whose upstream
|
// Either a non-streaming request, or a streaming one whose upstream
|
||||||
// ignored stream=true and answered with a normal JSON body.
|
// ignored stream=true and answered with a normal JSON body.
|
||||||
return a.forwardJSON(w, resp, model)
|
return a.forwardJSON(w, resp, responseModel)
|
||||||
}
|
}
|
||||||
|
|
||||||
// forwardJSON reads resp's full body (capped at maxResponseBytes),
|
// forwardJSON reads resp's full body (capped at maxResponseBytes),
|
||||||
@@ -189,6 +199,13 @@ func (a *geminiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gat
|
|||||||
// Embeddings never stream.
|
// Embeddings never stream.
|
||||||
func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
||||||
model := geminiModelFromRequest(req)
|
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"])
|
texts, single, err := geminiEmbeddingInputTexts(req["input"])
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -223,7 +240,7 @@ func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, r
|
|||||||
return usage{}, fmt.Errorf("%w: read response body: %w", errUpstream, err)
|
return usage{}, fmt.Errorf("%w: read response body: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
out, u, err := openAIEmbeddingResponseFromGemini(raw, single, model, estimatedEmbeddingTokens(texts))
|
out, u, err := openAIEmbeddingResponseFromGemini(raw, single, responseModel, estimatedEmbeddingTokens(texts))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return usage{}, err
|
return usage{}, err
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -92,6 +92,14 @@ type chatStreamChunk struct {
|
|||||||
// resulting usage-only chunk from a client that never asked for it while
|
// resulting usage-only chunk from a client that never asked for it while
|
||||||
// still capturing its usage for accounting.
|
// still capturing its usage for accounting.
|
||||||
func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
||||||
|
// gatewayAliasKey must never reach the real provider: an openai-type
|
||||||
|
// adapter forwards req verbatim as the upstream wire body, so a
|
||||||
|
// leftover alias entry would arrive as an unrecognized request field.
|
||||||
|
// This adapter's own response is a verbatim passthrough of whatever
|
||||||
|
// the upstream returns, so — unlike anthropic/gemini — there is no
|
||||||
|
// alias to echo back into it (documented in gatewayAliasKey's doc
|
||||||
|
// comment, routes_unified.go).
|
||||||
|
delete(req, gatewayAliasKey)
|
||||||
streaming, _ := req["stream"].(bool)
|
streaming, _ := req["stream"].(bool)
|
||||||
clientAskedUsage := false
|
clientAskedUsage := false
|
||||||
if streaming {
|
if streaming {
|
||||||
@@ -135,6 +143,7 @@ func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite
|
|||||||
// stream, so it is forwardJSON end to end, same as chatCompletion's
|
// stream, so it is forwardJSON end to end, same as chatCompletion's
|
||||||
// non-streaming path.
|
// non-streaming path.
|
||||||
func (a *openaiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
func (a *openaiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
||||||
|
delete(req, gatewayAliasKey) // see chatCompletion's identical delete for why
|
||||||
resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/embeddings", a.requestHeaders(true), req)
|
resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/embeddings", a.requestHeaders(true), req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return usage{}, err
|
return usage{}, err
|
||||||
|
|||||||
@@ -1,10 +1,41 @@
|
|||||||
package traefikllmgateway
|
package traefikllmgateway
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"math"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
"strconv"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// maxRequestBytes caps a client request body the unified routes decode
|
||||||
|
// into a map[string]any: 10MiB, generous for a chat or embeddings request
|
||||||
|
// (including an inline base64 image) while bounding memory against an
|
||||||
|
// oversized or malicious body.
|
||||||
|
const maxRequestBytes = 10 << 20
|
||||||
|
|
||||||
|
// gatewayAliasKey is the request map's internal-convention key carrying
|
||||||
|
// the client-requested model id through to an adapter (ruling a, ALIAS
|
||||||
|
// ECHO), distinct from "model" itself, which runUnified rewrites to the
|
||||||
|
// resolved upstream model id before calling the adapter — an operator's
|
||||||
|
// model alias must never reach the upstream API, but the client that used
|
||||||
|
// it must still see it echoed back in the response.
|
||||||
|
//
|
||||||
|
// Every adapter deletes this key before marshaling its upstream wire
|
||||||
|
// body: an openai-type adapter forwards req verbatim, so a leftover
|
||||||
|
// gatewayAliasKey entry would otherwise reach the real provider as an
|
||||||
|
// unrecognized request field. anthropic and gemini's adapters, which
|
||||||
|
// build a brand new response envelope rather than forwarding one
|
||||||
|
// verbatim, read it (falling back to the upstream model id when absent)
|
||||||
|
// to set that envelope's "model" field to the client's original id. An
|
||||||
|
// openai-type adapter's own response is a verbatim passthrough of
|
||||||
|
// whatever the upstream returned and cannot be rewritten this way — its
|
||||||
|
// response body carries the upstream's model id, not the alias. This is
|
||||||
|
// a documented, accepted asymmetry, not a bug.
|
||||||
|
const gatewayAliasKey = "__alias"
|
||||||
|
|
||||||
// handleModels implements GET /v1/models: an authenticated caller gets back
|
// handleModels implements GET /v1/models: an authenticated caller gets back
|
||||||
// an OpenAI-compatible {"object":"list","data":[...]} envelope of the
|
// an OpenAI-compatible {"object":"list","data":[...]} envelope of the
|
||||||
// models their group can see, via modelRegistry.listFor.
|
// models their group can see, via modelRegistry.listFor.
|
||||||
@@ -21,3 +52,214 @@ func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
|
|||||||
"data": g.registry.listFor(grp),
|
"data": g.registry.listFor(grp),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// adapterCall abstracts providerAdapter.chatCompletion and .embeddings:
|
||||||
|
// handleChat and handleEmbeddings share every other step of the pipeline
|
||||||
|
// below (decode, resolve, enforce limits, account usage, translate
|
||||||
|
// errors) and differ only in which of these two methods they invoke.
|
||||||
|
type adapterCall func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error)
|
||||||
|
|
||||||
|
// handleChat implements POST /v1/chat/completions.
|
||||||
|
func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
|
||||||
|
g.runUnified(w, r, u, grp, func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
||||||
|
return a.chatCompletion(ctx, w, req)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleEmbeddings implements POST /v1/embeddings.
|
||||||
|
func (g *Gateway) handleEmbeddings(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
|
||||||
|
g.runUnified(w, r, u, grp, func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) {
|
||||||
|
return a.embeddings(ctx, w, req)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// runUnified is the shared chat/embeddings pipeline: decode the request,
|
||||||
|
// resolve its model, enforce per-user and per-group limits, invoke the
|
||||||
|
// adapter via call, then account the resulting usage. w is wrapped in its
|
||||||
|
// own statusTrackingWriter so a mid-stream adapter error (headers already
|
||||||
|
// sent) can be told apart from one that failed before any write.
|
||||||
|
func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, grp *group, call adapterCall) {
|
||||||
|
sw := &statusTrackingWriter{ResponseWriter: w}
|
||||||
|
|
||||||
|
body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes))
|
||||||
|
if err != nil {
|
||||||
|
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var req map[string]any
|
||||||
|
if err = json.Unmarshal(body, &req); err != nil {
|
||||||
|
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid JSON body")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
requestedModel, _ := req["model"].(string)
|
||||||
|
if requestedModel == "" {
|
||||||
|
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
streaming, _ := req["stream"].(bool)
|
||||||
|
|
||||||
|
adapter, upstreamModel, canonical, err := g.registry.resolve(requestedModel, grp)
|
||||||
|
if err != nil {
|
||||||
|
writeModelResolveError(sw, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
scopes := buildLimitScopes(u, grp)
|
||||||
|
if violation := g.limiter.checkAndCount(scopes); violation != nil {
|
||||||
|
writeLimitViolation(sw, violation)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
req["model"] = upstreamModel
|
||||||
|
req[gatewayAliasKey] = requestedModel
|
||||||
|
|
||||||
|
result, callErr := call(adapter, r.Context(), sw, req)
|
||||||
|
if callErr != nil {
|
||||||
|
g.handleAdapterError(sw, callErr, adapter.name())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if result.total() == 0 {
|
||||||
|
if streaming {
|
||||||
|
g.logf("unified route: zero usage reported for a streaming response from model %q; accounting the request only", canonical)
|
||||||
|
} else {
|
||||||
|
result.prompt = int64(math.Ceil(float64(len(body)) / 4))
|
||||||
|
result.estimated = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
g.limiter.account(scopes, result, unifiedCostMicros(canonical, upstreamModel, result, g.cfg.Pricing))
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildLimitScopes returns the limitScope slice runUnified passes to the
|
||||||
|
// limiter: a user scope only when u has its own limits configured, then a
|
||||||
|
// group scope only when grp does (ruling e) — either, both, or neither may
|
||||||
|
// apply to a given request. The user scope is listed first, so
|
||||||
|
// checkAndCount reports a user's own violation ahead of their group's when
|
||||||
|
// both are breached by the same request.
|
||||||
|
func buildLimitScopes(u *user, grp *group) []limitScope {
|
||||||
|
scopes := make([]limitScope, 0, 2)
|
||||||
|
if u.limits != nil {
|
||||||
|
scopes = append(scopes, limitScope{limits: u.limits, kind: "user", id: u.name})
|
||||||
|
}
|
||||||
|
if grp.limits != nil {
|
||||||
|
scopes = append(scopes, limitScope{limits: grp.limits, kind: "group", id: grp.name})
|
||||||
|
}
|
||||||
|
return scopes
|
||||||
|
}
|
||||||
|
|
||||||
|
// unifiedCostMicros resolves the price to charge one request's usage
|
||||||
|
// against: canonical ("provider/model") first, falling back to bare — the
|
||||||
|
// upstream model id alone — only when canonical has no configured price at
|
||||||
|
// all, in either overrides or the built-in table. This extends
|
||||||
|
// costMicros' single-id lookup to the two-id order ruling (d) requires,
|
||||||
|
// without changing costMicros' own signature: lookupPricing is used purely
|
||||||
|
// as a side-effect-free existence probe on canonical, and the real
|
||||||
|
// computation (and its unknown-model warning, when neither id has a price)
|
||||||
|
// runs through the one costMicros call that is actually charged.
|
||||||
|
func unifiedCostMicros(canonical, bare string, u usage, overrides map[string]*ModelPricing) int64 {
|
||||||
|
if _, ok := lookupPricing(canonical, overrides); ok {
|
||||||
|
return costMicros(canonical, u, overrides)
|
||||||
|
}
|
||||||
|
return costMicros(bare, u, overrides)
|
||||||
|
}
|
||||||
|
|
||||||
|
// writeModelResolveError maps a modelRegistry.resolve error to its HTTP
|
||||||
|
// envelope: errModelUnknown is a 404 (no configured provider knows this
|
||||||
|
// model), errModelDenied is a 403 (a real model the caller's group cannot
|
||||||
|
// use). Any other error is a defensive 500 — resolve's own contract
|
||||||
|
// promises only these two sentinels, so reaching this branch would be a
|
||||||
|
// programming error, not a client mistake.
|
||||||
|
func writeModelResolveError(w http.ResponseWriter, err error) {
|
||||||
|
switch {
|
||||||
|
case errors.Is(err, errModelUnknown):
|
||||||
|
writeOAIError(w, http.StatusNotFound, "invalid_request_error", "unknown model")
|
||||||
|
case errors.Is(err, errModelDenied):
|
||||||
|
writeOAIError(w, http.StatusForbidden, "invalid_request_error", "model access denied")
|
||||||
|
default:
|
||||||
|
writeOAIError(w, http.StatusInternalServerError, "server_error", "internal error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// writeLimitViolation maps a limiter violation to its HTTP envelope
|
||||||
|
// (ruling b): a storeDown violation — the configured limit store was
|
||||||
|
// unreachable and failOpen is false — is a 503 server_error, since it is
|
||||||
|
// not an actual limit breach; any other violation is a 429 rate_limit_error
|
||||||
|
// carrying v's own message, plus a Retry-After header when the violation
|
||||||
|
// names a meaningful retry window.
|
||||||
|
func writeLimitViolation(w http.ResponseWriter, v *limitViolation) {
|
||||||
|
if v.storeDown {
|
||||||
|
writeOAIError(w, http.StatusServiceUnavailable, "server_error", v.message)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if v.retryAfter > 0 {
|
||||||
|
w.Header().Set("Retry-After", strconv.Itoa(v.retryAfter))
|
||||||
|
}
|
||||||
|
writeOAIError(w, http.StatusTooManyRequests, "rate_limit_error", v.message)
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleAdapterError translates an adapter's returned error into a
|
||||||
|
// response, or into nothing at all, per ruling c. sw.wroteHeader is
|
||||||
|
// checked first, ahead of every other case: an error surfacing after the
|
||||||
|
// adapter already started writing a response (a mid-stream connection
|
||||||
|
// drop) must never get a second, conflicting envelope appended, whatever
|
||||||
|
// kind of error it is.
|
||||||
|
func (g *Gateway) handleAdapterError(sw *statusTrackingWriter, err error, providerName string) {
|
||||||
|
if sw.wroteHeader {
|
||||||
|
g.errorf("unified route: adapter error after response started (provider %q): %v", providerName, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var perr *providerHTTPError
|
||||||
|
if errors.As(err, &perr) {
|
||||||
|
writeProviderUpstreamError(sw, providerName, perr)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var terr *translateError
|
||||||
|
if errors.As(err, &terr) {
|
||||||
|
status := http.StatusBadRequest
|
||||||
|
if terr.notSupported {
|
||||||
|
status = http.StatusNotImplemented
|
||||||
|
}
|
||||||
|
writeOAIError(sw, status, "invalid_request_error", terr.msg)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if errors.Is(err, context.Canceled) {
|
||||||
|
g.logf("unified route: client canceled request to provider %q: %v", providerName, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
g.errorf("unified route: upstream connection error (provider %q): %v", providerName, err)
|
||||||
|
writeOAIError(sw, http.StatusBadGateway, "server_error", "upstream connection error")
|
||||||
|
}
|
||||||
|
|
||||||
|
// writeProviderUpstreamError passes perr's status through to the client,
|
||||||
|
// wrapped in the gateway's own envelope shape rather than perr.body
|
||||||
|
// verbatim (unified-route callers get this wrapped form; the openai-type
|
||||||
|
// adapter's own passthrough forwards a non-2xx body unwrapped, since that
|
||||||
|
// path never reaches this function). perr.body is embedded under
|
||||||
|
// error.upstream, decoded to a JSON value when it parses as one and left
|
||||||
|
// as a raw string otherwise, so a caller sees exactly what the provider
|
||||||
|
// said either way.
|
||||||
|
func writeProviderUpstreamError(w http.ResponseWriter, providerName string, perr *providerHTTPError) {
|
||||||
|
var upstream any = string(perr.body)
|
||||||
|
var parsed any
|
||||||
|
if json.Unmarshal(perr.body, &parsed) == nil {
|
||||||
|
upstream = parsed
|
||||||
|
}
|
||||||
|
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(perr.status)
|
||||||
|
_ = json.NewEncoder(w).Encode(map[string]any{ // headers already committed; nothing useful to do on encode failure
|
||||||
|
"error": map[string]any{
|
||||||
|
"message": providerName + " upstream error",
|
||||||
|
"type": "upstream_error",
|
||||||
|
"code": strconv.Itoa(perr.status),
|
||||||
|
"upstream": upstream,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,558 @@
|
|||||||
|
package traefikllmgateway
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// newUnifiedRequest builds a POST request against the unified routes,
|
||||||
|
// JSON-encoding body and attaching apiKey as a Bearer token when non-empty.
|
||||||
|
func newUnifiedRequest(t *testing.T, method, path, apiKey string, body map[string]any) *http.Request {
|
||||||
|
t.Helper()
|
||||||
|
b, err := json.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("marshal request body: %v", err)
|
||||||
|
}
|
||||||
|
req := httptest.NewRequest(method, path, bytes.NewReader(b))
|
||||||
|
if apiKey != "" {
|
||||||
|
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||||
|
}
|
||||||
|
return req
|
||||||
|
}
|
||||||
|
|
||||||
|
// recordingWriter is a minimal http.ResponseWriter/http.Flusher that logs
|
||||||
|
// the sequence of Write and Flush calls it receives, so a streaming test
|
||||||
|
// can assert flushes happen incrementally rather than once at the end.
|
||||||
|
type recordingWriter struct {
|
||||||
|
hdr http.Header
|
||||||
|
events []string
|
||||||
|
buf bytes.Buffer
|
||||||
|
status int
|
||||||
|
}
|
||||||
|
|
||||||
|
func newRecordingWriter() *recordingWriter {
|
||||||
|
return &recordingWriter{hdr: http.Header{}}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *recordingWriter) Header() http.Header { return r.hdr }
|
||||||
|
|
||||||
|
func (r *recordingWriter) WriteHeader(status int) { r.status = status }
|
||||||
|
|
||||||
|
func (r *recordingWriter) Write(b []byte) (int, error) {
|
||||||
|
if r.status == 0 {
|
||||||
|
r.status = http.StatusOK // mirrors net/http's implicit-200 default when Write precedes any explicit WriteHeader
|
||||||
|
}
|
||||||
|
n, _ := r.buf.Write(b)
|
||||||
|
r.events = append(r.events, "write")
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *recordingWriter) Flush() { r.events = append(r.events, "flush") }
|
||||||
|
|
||||||
|
// alwaysErrStore is a counterStore stub whose every operation fails, used
|
||||||
|
// to drive the limiter's fail-closed (storeDown) path deterministically.
|
||||||
|
type alwaysErrStore struct{}
|
||||||
|
|
||||||
|
var errStoreDownStub = errors.New("stub: store down")
|
||||||
|
|
||||||
|
func (alwaysErrStore) incrBy(string, int64, time.Duration) (int64, error) {
|
||||||
|
return 0, errStoreDownStub
|
||||||
|
}
|
||||||
|
|
||||||
|
func (alwaysErrStore) get(string) (int64, error) {
|
||||||
|
return 0, errStoreDownStub
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_HappyPath_NonStreaming_AccountsUsage drives a full
|
||||||
|
// chat-completion request through Gateway.ServeHTTP against a fake
|
||||||
|
// openai-shaped upstream, and asserts the response is forwarded verbatim
|
||||||
|
// and the reported usage lands in the limiter's own counters.
|
||||||
|
func TestHandleChat_HappyPath_NonStreaming_AccountsUsage(t *testing.T) {
|
||||||
|
const respBody = `{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-real","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`
|
||||||
|
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
_, _ = w.Write([]byte(respBody))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{
|
||||||
|
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}},
|
||||||
|
}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{
|
||||||
|
{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}},
|
||||||
|
}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
gw, ok := h.(*Gateway)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("handler is not *Gateway")
|
||||||
|
}
|
||||||
|
|
||||||
|
body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusOK {
|
||||||
|
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
if rec.Body.String() != respBody {
|
||||||
|
t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody)
|
||||||
|
}
|
||||||
|
|
||||||
|
tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||||
|
if !ok || tok != 15 {
|
||||||
|
t.Errorf("user token/day counter = %d (ok=%v), want 15", tok, ok)
|
||||||
|
}
|
||||||
|
reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now())
|
||||||
|
if !ok || reqCount != 1 {
|
||||||
|
t.Errorf("user request/min counter = %d (ok=%v), want 1", reqCount, ok)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_Streaming_FlushesIncrementally proves a streaming chat
|
||||||
|
// completion is forwarded to the client chunk by chunk, flushing after
|
||||||
|
// each one, rather than buffered until the response completes.
|
||||||
|
func TestHandleChat_Streaming_FlushesIncrementally(t *testing.T) {
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
fl, ok := w.(http.Flusher)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("fake upstream ResponseWriter does not support Flush")
|
||||||
|
}
|
||||||
|
for _, c := range []string{
|
||||||
|
`{"id":"c1","object":"chat.completion.chunk","choices":[{"delta":{"content":"Hel"}}]}`,
|
||||||
|
`{"id":"c1","object":"chat.completion.chunk","choices":[{"delta":{"content":"lo"}}]}`,
|
||||||
|
} {
|
||||||
|
_, _ = w.Write([]byte("data: " + c + "\n\n"))
|
||||||
|
fl.Flush()
|
||||||
|
}
|
||||||
|
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||||
|
fl.Flush()
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{
|
||||||
|
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}},
|
||||||
|
}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rw := newRecordingWriter()
|
||||||
|
h.ServeHTTP(rw, req)
|
||||||
|
|
||||||
|
if rw.status != http.StatusOK {
|
||||||
|
t.Fatalf("status = %d, want 200", rw.status)
|
||||||
|
}
|
||||||
|
writes, flushes := 0, 0
|
||||||
|
for _, e := range rw.events {
|
||||||
|
switch e {
|
||||||
|
case "write":
|
||||||
|
writes++
|
||||||
|
case "flush":
|
||||||
|
flushes++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if writes < 3 {
|
||||||
|
t.Fatalf("want >=3 writes (2 content chunks + [DONE]), got %d (events=%v)", writes, rw.events)
|
||||||
|
}
|
||||||
|
if flushes < writes {
|
||||||
|
t.Fatalf("want a flush per write (incremental delivery), got %d writes and %d flushes", writes, flushes)
|
||||||
|
}
|
||||||
|
if !strings.Contains(rw.buf.String(), "data: [DONE]") {
|
||||||
|
t.Errorf("body missing terminal [DONE], got %q", rw.buf.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_NoAPIKey_Returns401 asserts the unified routes require
|
||||||
|
// authentication before any resolve/limit work happens.
|
||||||
|
func TestHandleChat_NoAPIKey_Returns401(t *testing.T) {
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "", map[string]any{"model": "gpt-test"})
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusUnauthorized {
|
||||||
|
t.Fatalf("status = %d, want 401, body=%s", rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_DeniedModel_Returns403 covers a model that a configured
|
||||||
|
// provider knows, but the caller's group is not authorized to use.
|
||||||
|
func TestHandleChat_DeniedModel_Returns403(t *testing.T) {
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{
|
||||||
|
"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}},
|
||||||
|
}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {Models: []string{"gpt-allowed-only"}}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
body := map[string]any{"model": "gpt-test", "messages": []any{}}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusForbidden {
|
||||||
|
t.Fatalf("status = %d, want 403, body=%s", rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_UnknownModel_Returns404 covers a model no configured
|
||||||
|
// provider knows at all.
|
||||||
|
func TestHandleChat_UnknownModel_Returns404(t *testing.T) {
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
body := map[string]any{"model": "no-such-model", "messages": []any{}}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusNotFound {
|
||||||
|
t.Fatalf("status = %d, want 404, body=%s", rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_RequestLimitExceeded_Returns429WithRetryAfter drives two
|
||||||
|
// requests against a user limited to one per minute, and asserts the
|
||||||
|
// second is refused with a Retry-After header.
|
||||||
|
func TestHandleChat_RequestLimitExceeded_Returns429WithRetryAfter(t *testing.T) {
|
||||||
|
const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
_, _ = w.Write([]byte(respBody))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{
|
||||||
|
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}},
|
||||||
|
}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{
|
||||||
|
{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerMinute: 1}},
|
||||||
|
}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
body := map[string]any{"model": "gpt-test", "messages": []any{}}
|
||||||
|
req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec1 := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec1, req1)
|
||||||
|
if rec1.Code != http.StatusOK {
|
||||||
|
t.Fatalf("first request status = %d, want 200, body=%s", rec1.Code, rec1.Body.String())
|
||||||
|
}
|
||||||
|
|
||||||
|
req2 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec2 := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec2, req2)
|
||||||
|
if rec2.Code != http.StatusTooManyRequests {
|
||||||
|
t.Fatalf("second request status = %d, want 429, body=%s", rec2.Code, rec2.Body.String())
|
||||||
|
}
|
||||||
|
if ra := rec2.Header().Get("Retry-After"); ra == "" {
|
||||||
|
t.Error("want a Retry-After header on 429, got none")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_StoreDown_FailClosed_Returns503 wires a counterStore
|
||||||
|
// stub whose every operation fails, with failOpen=false, and asserts the
|
||||||
|
// request is refused with 503 rather than silently enforcing limits
|
||||||
|
// against an unreachable backend.
|
||||||
|
func TestHandleChat_StoreDown_FailClosed_Returns503(t *testing.T) {
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{
|
||||||
|
{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerMinute: 100}},
|
||||||
|
}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
gw, ok := h.(*Gateway)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("handler is not *Gateway")
|
||||||
|
}
|
||||||
|
gw.limiter = newLimiter(alwaysErrStore{}, false)
|
||||||
|
|
||||||
|
body := map[string]any{"model": "gpt-test", "messages": []any{}}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusServiceUnavailable {
|
||||||
|
t.Fatalf("status = %d, want 503, body=%s", rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleEmbeddings_HappyPath mirrors the chat-completion happy path
|
||||||
|
// against POST /v1/embeddings.
|
||||||
|
func TestHandleEmbeddings_HappyPath(t *testing.T) {
|
||||||
|
const respBody = `{"object":"list","data":[{"embedding":[0.1,0.2]}],"model":"text-embedding-3","usage":{"prompt_tokens":5,"total_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()
|
||||||
|
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{
|
||||||
|
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"text-embedding-3"}},
|
||||||
|
}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
body := map[string]any{"model": "text-embedding-3", "input": "hello"}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/embeddings", "sk-alice", body)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusOK {
|
||||||
|
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
if rec.Body.String() != respBody {
|
||||||
|
t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_AnthropicProviderPrefixedModel_EchoesAliasInResponse
|
||||||
|
// requests a provider-prefixed model id ("anthropic/claude-x"): the
|
||||||
|
// upstream call must use the bare upstream id, but the translated
|
||||||
|
// response's "model" field must echo back the exact id the client
|
||||||
|
// requested (ruling a, ALIAS ECHO), not the bare upstream form.
|
||||||
|
func TestHandleChat_AnthropicProviderPrefixedModel_EchoesAliasInResponse(t *testing.T) {
|
||||||
|
const anthResp = `{"id":"msg_01ABC","type":"message","role":"assistant","content":[{"type":"text","text":"hi there"}],"model":"claude-x","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":4}}`
|
||||||
|
|
||||||
|
var gotBody []byte
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
gotBody, _ = io.ReadAll(r.Body)
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
_, _ = w.Write([]byte(anthResp))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{
|
||||||
|
"anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-ant", Models: []string{"claude-x"}},
|
||||||
|
}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
const requestedID = "anthropic/claude-x"
|
||||||
|
body := map[string]any{"model": requestedID, "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusOK {
|
||||||
|
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
|
||||||
|
var out map[string]any
|
||||||
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
||||||
|
t.Fatalf("decode response: %v", err)
|
||||||
|
}
|
||||||
|
if out["model"] != requestedID {
|
||||||
|
t.Errorf("response model = %v, want the client-requested alias %q echoed back", out["model"], requestedID)
|
||||||
|
}
|
||||||
|
|
||||||
|
sentUpstream := decodeJSONBody(t, gotBody)
|
||||||
|
if sentUpstream["model"] != "claude-x" {
|
||||||
|
t.Errorf("upstream request model = %v, want bare upstream id %q", sentUpstream["model"], "claude-x")
|
||||||
|
}
|
||||||
|
if _, present := sentUpstream[gatewayAliasKey]; present {
|
||||||
|
t.Errorf("gatewayAliasKey %q leaked into the upstream request body: %v", gatewayAliasKey, sentUpstream)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_UpstreamNon2xx_WrapsProviderErrorEnvelope asserts a
|
||||||
|
// non-2xx upstream response is passed through with its status code, but
|
||||||
|
// wrapped in the gateway's own envelope with the raw upstream body
|
||||||
|
// embedded under error.upstream (ruling c).
|
||||||
|
func TestHandleChat_UpstreamNon2xx_WrapsProviderErrorEnvelope(t *testing.T) {
|
||||||
|
const upstreamErrBody = `{"error":{"message":"rate limited upstream","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(upstreamErrBody))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{
|
||||||
|
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}},
|
||||||
|
}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
body := map[string]any{"model": "gpt-test", "messages": []any{}}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
if rec.Code != http.StatusTooManyRequests {
|
||||||
|
t.Fatalf("status = %d, want 429, body=%s", rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
|
||||||
|
var out struct {
|
||||||
|
Error struct {
|
||||||
|
Upstream map[string]any `json:"upstream"`
|
||||||
|
Message string `json:"message"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Code string `json:"code"`
|
||||||
|
} `json:"error"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
||||||
|
t.Fatalf("decode response: %v", err)
|
||||||
|
}
|
||||||
|
if out.Error.Type != "upstream_error" {
|
||||||
|
t.Errorf("error.type = %q, want upstream_error", out.Error.Type)
|
||||||
|
}
|
||||||
|
if out.Error.Code != "429" {
|
||||||
|
t.Errorf("error.code = %q, want %q", out.Error.Code, "429")
|
||||||
|
}
|
||||||
|
if out.Error.Upstream == nil {
|
||||||
|
t.Fatal("error.upstream is nil, want the parsed upstream body")
|
||||||
|
}
|
||||||
|
errObj, _ := out.Error.Upstream["error"].(map[string]any)
|
||||||
|
if errObj["message"] != "rate limited upstream" {
|
||||||
|
t.Errorf("error.upstream.error.message = %v, want %q", errObj["message"], "rate limited upstream")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandleChat_MidStreamUpstreamDrop_NoTrailingEnvelope proves an
|
||||||
|
// adapter error arriving after the response has already started streaming
|
||||||
|
// gets logged, not turned into a second, conflicting envelope appended to
|
||||||
|
// a body the client already started receiving (ruling c).
|
||||||
|
func TestHandleChat_MidStreamUpstreamDrop_NoTrailingEnvelope(t *testing.T) {
|
||||||
|
const firstChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n"
|
||||||
|
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
fl, ok := w.(http.Flusher)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("fake upstream ResponseWriter does not support Flush")
|
||||||
|
}
|
||||||
|
_, _ = w.Write([]byte(firstChunk))
|
||||||
|
fl.Flush()
|
||||||
|
|
||||||
|
hj, ok := w.(http.Hijacker)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("fake upstream ResponseWriter does not support Hijack")
|
||||||
|
}
|
||||||
|
conn, _, hjErr := hj.Hijack()
|
||||||
|
if hjErr == nil {
|
||||||
|
_ = conn.Close()
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
cfg := CreateConfig()
|
||||||
|
cfg.Providers = map[string]*ProviderConfig{
|
||||||
|
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}},
|
||||||
|
}
|
||||||
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||||
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||||
|
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||||
|
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
origStderr := os.Stderr
|
||||||
|
pr, pw, _ := os.Pipe()
|
||||||
|
os.Stderr = pw
|
||||||
|
defer func() { os.Stderr = origStderr }()
|
||||||
|
|
||||||
|
body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{}}
|
||||||
|
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
_ = pw.Close() // closing the pipe write end to unblock the read; error not actionable in a test
|
||||||
|
os.Stderr = origStderr
|
||||||
|
var logBuf bytes.Buffer
|
||||||
|
if _, err := io.Copy(&logBuf, pr); err != nil {
|
||||||
|
t.Fatalf("io.Copy: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if rec.Code != http.StatusOK {
|
||||||
|
t.Fatalf("status = %d, want 200 (headers already committed before the drop)", rec.Code)
|
||||||
|
}
|
||||||
|
if got := rec.Body.String(); got != firstChunk {
|
||||||
|
t.Errorf("body = %q, want exactly the one chunk written before the drop (no trailing error envelope)", got)
|
||||||
|
}
|
||||||
|
if !strings.Contains(logBuf.String(), "llmgw[llmgw] ERROR") {
|
||||||
|
t.Errorf("want the mid-stream error logged, got %q", logBuf.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user