CI / test (push) Failing after 6m12s
/v1/openapi.json is the API clients use — OpenAI's endpoints, Anthropic's, and what llmgw adds — in OpenAPI 3.1, generated from internal/clientapi's route table by the generator the admin API's document comes from, now told each document's title, server, token and error. A test reads the router's mux.HandleFunc calls and refuses a route described and not mounted, or mounted and not described. Asked for before there is a key, so it needs none. /docs/api shows it in Swagger UI 5.33.1 (Apache 2.0, served from the gateway as the page's policy requires): Authorize takes a key, Try it out sends to this gateway. Staff are offered the admin API's document too, from /docs/api/admin.json, which nobody else is given. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1412 lines
54 KiB
Go
1412 lines
54 KiB
Go
package router
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
llmgw "gitea.mynodes.xyz/ilker/llmgw"
|
|
"io"
|
|
"net/http"
|
|
"slices"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.opentelemetry.io/otel/attribute"
|
|
"go.opentelemetry.io/otel/codes"
|
|
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/accounting"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/admission"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/scim"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/secrets"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/web"
|
|
)
|
|
|
|
// admit waits for the flow control layer to let the request through. With
|
|
// admission disabled it returns immediately, which is the v0.2.0 behaviour.
|
|
func (r *Router) admit(ctx context.Context, req *parsedRequest) (*admission.Lease, error) {
|
|
if r.adm == nil {
|
|
return nil, nil
|
|
}
|
|
// The key is the tenant when there is one: fairness within a band has to be
|
|
// per caller, and a session id is something the caller picks, so one client
|
|
// could take a whole band by opening many sessions. Without authentication
|
|
// the session is still the best available grouping.
|
|
tenant := req.Tenant
|
|
if tenant == "" {
|
|
tenant = req.SessionID
|
|
}
|
|
return r.adm.Acquire(ctx, admission.Request{
|
|
Model: req.Model,
|
|
Band: req.Band,
|
|
Tenant: tenant,
|
|
})
|
|
}
|
|
|
|
// rejectAdmission turns a flow control refusal into a status the client can act
|
|
// on: 503 when the queue is full (try later), 504 when it timed out waiting.
|
|
func (r *Router) rejectAdmission(w http.ResponseWriter, req *parsedRequest, err error, start time.Time) {
|
|
status := http.StatusServiceUnavailable
|
|
switch {
|
|
case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded):
|
|
// The client gave up while queued; nobody is listening.
|
|
r.m.observeAdmissionRejected(req.Band, "client_gone")
|
|
r.m.observeRequest(r.metricModel(req.Model), "", 499, time.Since(start), nil, 0)
|
|
return
|
|
case errors.Is(err, admission.ErrQueueTimeout):
|
|
// The cluster was busy for as long as the band may wait, not broken:
|
|
// said so a client comes back. Without it llmgw-kb, which retries
|
|
// only what says when to, counted every such passage as failed.
|
|
status = http.StatusGatewayTimeout
|
|
w.Header().Set("Retry-After", "10")
|
|
case errors.Is(err, admission.ErrQueueFull):
|
|
w.Header().Set("Retry-After", "10")
|
|
case errors.Is(err, admission.ErrDisplaced):
|
|
// Turned away for higher-priority work, not for good: said so the
|
|
// client backs off and comes back rather than giving up.
|
|
w.Header().Set("Retry-After", "5")
|
|
}
|
|
reason := "unavailable"
|
|
switch {
|
|
case errors.Is(err, admission.ErrDisplaced):
|
|
reason = "displaced"
|
|
case errors.Is(err, admission.ErrQueueFull):
|
|
reason = "queue_full"
|
|
case errors.Is(err, admission.ErrQueueTimeout):
|
|
reason = "ttl"
|
|
case errors.Is(err, admission.ErrClosed):
|
|
reason = "shutdown"
|
|
}
|
|
r.m.observeAdmissionRejected(req.Band, reason)
|
|
writeError(w, status, "server_error", err.Error())
|
|
r.m.observeRequest(r.metricModel(req.Model), "", status, time.Since(start), nil, 0)
|
|
r.accessLog(req, dispatchResult{err: err}, status, time.Since(start))
|
|
}
|
|
|
|
// prefixBytes is how much of the prompt is kept for prefix-aware routing.
|
|
// Large enough to cover a system prompt plus the first message, small enough
|
|
// that hashing it stays cheap.
|
|
const prefixBytes = 2048
|
|
|
|
// Handler returns the client-facing OpenAI API.
|
|
func (r *Router) Handler() http.Handler {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("POST /v1/chat/completions", r.handleInference("/v1/chat/completions"))
|
|
mux.HandleFunc("POST /v1/completions", r.handleInference("/v1/completions"))
|
|
mux.HandleFunc("POST /v1/embeddings", r.handleInference("/v1/embeddings"))
|
|
// Passed through to engines that serve them: vLLM's rerankers and
|
|
// scorers, and speech to text.
|
|
mux.HandleFunc("POST /v1/rerank", r.handleInference("/v1/rerank"))
|
|
mux.HandleFunc("POST /v1/score", r.handleInference("/v1/score"))
|
|
mux.HandleFunc("POST /v1/audio/transcriptions", r.handleInference("/v1/audio/transcriptions"))
|
|
mux.HandleFunc("POST /v1/audio/translations", r.handleInference("/v1/audio/translations"))
|
|
// Pictures, made by ComfyUI from a workflow the catalog names.
|
|
mux.HandleFunc("POST "+imagesPath, r.handleInference(imagesPath))
|
|
// Answered by the guardrails' judge model.
|
|
mux.HandleFunc("POST /v1/moderations", r.handleModerations)
|
|
// Codex, and whatever else speaks only the Responses API, translated to
|
|
// chat completions on the way in and back on the way out.
|
|
mux.HandleFunc("POST /v1/responses", r.handleResponses)
|
|
// Claude Code and Anthropic's SDKs, translated the same way.
|
|
mux.HandleFunc("POST /v1/messages", r.handleMessages)
|
|
mux.HandleFunc("POST /v1/messages/count_tokens", r.handleCountTokens)
|
|
mux.HandleFunc("GET /v1/models", r.handleModels)
|
|
// The API described, for SDK generators and the portal's reference.
|
|
// Not secret, and asked for before a key is in hand: no key needed.
|
|
mux.HandleFunc("GET /v1/openapi.json", func(w http.ResponseWriter, _ *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write(llmgw.ClientOpenAPI)
|
|
})
|
|
// Files and batches: a file of requests served at the lowest priority.
|
|
mux.HandleFunc("POST /v1/files", r.handleUploadFile)
|
|
mux.HandleFunc("GET /v1/files", r.handleListFiles)
|
|
mux.HandleFunc("GET /v1/files/{id}", r.handleGetFile)
|
|
mux.HandleFunc("GET /v1/files/{id}/content", r.handleFileContent)
|
|
mux.HandleFunc("DELETE /v1/files/{id}", r.handleDeleteFile)
|
|
mux.HandleFunc("POST /v1/batches", r.handleCreateBatch)
|
|
mux.HandleFunc("GET /v1/batches", r.handleListBatches)
|
|
mux.HandleFunc("GET /v1/batches/{id}", r.handleGetBatch)
|
|
mux.HandleFunc("POST /v1/batches/{id}/cancel", r.handleCancelBatch)
|
|
// Jobs: 3D models, which take minutes and are asked after rather than
|
|
// waited for. The result is a file.
|
|
mux.HandleFunc("POST "+meshesPath, r.handleCreateJob)
|
|
mux.HandleFunc("GET "+meshesPath, r.handleListJobs)
|
|
mux.HandleFunc("GET "+meshesPath+"/{id}", r.handleGetJob)
|
|
mux.HandleFunc("POST "+meshesPath+"/{id}/cancel", r.handleCancelJob)
|
|
// Knowledge bases: documents a person's llmgw-kb read and embedded.
|
|
r.registerKB(mux)
|
|
// Liveness: this process is up. It says nothing about whether it can
|
|
// serve, which is why it is not the thing to put a load balancer on.
|
|
mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) {
|
|
w.WriteHeader(http.StatusOK)
|
|
})
|
|
|
|
// Readiness: this process has somewhere to send a request.
|
|
//
|
|
// The distinction is not academic. A standby control plane is up, answers
|
|
// /health, and has no endpoints at all — its reconciler placed nothing and
|
|
// it holds no agent streams to adopt from — so a load balancer checking
|
|
// /health would send it half the traffic and half the traffic would be a
|
|
// 503. A leader whose replicas have all failed is in the same position for
|
|
// a different reason, and a load balancer should treat them the same.
|
|
//
|
|
// Open, like /health: a load balancer does not have a key, and this leaks
|
|
// one bit — whether anything is servable — which the next request would
|
|
// reveal anyway.
|
|
//
|
|
// Only this cluster's own endpoints count, not a federated peer's: a
|
|
// peer reads this to decide whether to send here, and two clusters
|
|
// federating each other would otherwise each be ready because the other
|
|
// is, with nothing of their own to serve with.
|
|
mux.HandleFunc("GET /ready", func(w http.ResponseWriter, _ *http.Request) {
|
|
if r.localEndpoints() == 0 {
|
|
http.Error(w, "no endpoints are available to serve", http.StatusServiceUnavailable)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusOK)
|
|
})
|
|
|
|
// Leadership: 200 on the control plane that places work, 503 on a
|
|
// standby. What a floating address follows (llmgw server add --vip):
|
|
// unlike /ready it does not wait for a replica, so a cluster with no
|
|
// node yet still has its address on one machine. A process that is not
|
|
// a control plane, or one with no rival, is always the leader.
|
|
mux.HandleFunc("GET /leader", func(w http.ResponseWriter, _ *http.Request) {
|
|
if f := r.isLeader.Load(); f != nil && !(*f)() {
|
|
http.Error(w, "this control plane is a standby", http.StatusServiceUnavailable)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusOK)
|
|
})
|
|
|
|
// The landing page says how to connect and nothing about what is here, so
|
|
// it is safe to serve without a key — and it has to be, because the person
|
|
// reading it is the one who does not have one yet.
|
|
mux.HandleFunc("GET /", web.Landing(version.Version, r.auth != nil,
|
|
r.accounts != nil && r.accounts.SignupOpen(),
|
|
r.landing.Contact, r.landing.Notice))
|
|
// What llmgw is and how to install it: about the software, not this
|
|
// installation, so it needs no key either.
|
|
mux.HandleFunc("GET /about", web.About(version.Version))
|
|
mux.Handle("GET /_static/", web.Static())
|
|
// The language switcher at the foot of every page, which has to work
|
|
// before anyone has signed in.
|
|
mux.HandleFunc("GET /_lang/{code}", web.SwitchLanguage)
|
|
mux.HandleFunc("GET /me", r.handleMe)
|
|
// Somebody has to be able to say "forget that", and a retention window is
|
|
// no answer to a message sent by mistake. Authenticated like any other
|
|
// client endpoint, and scoped to the caller's project by the store.
|
|
mux.HandleFunc("DELETE /v1/threads/{id}", r.handleDeleteThread)
|
|
|
|
// Notes an agent keeps for itself, over REST and as MCP tools. Off unless a
|
|
// project or the threads block turned storage on.
|
|
mux.HandleFunc("PUT /v1/memory/{key}", r.handlePutNote)
|
|
mux.HandleFunc("GET /v1/memory/{key}", r.handleGetNote)
|
|
mux.HandleFunc("GET /v1/memory", r.handleSearchNotes)
|
|
mux.HandleFunc("DELETE /v1/memory/{key}", r.handleDeleteNote)
|
|
mux.HandleFunc("POST /mcp", r.handleMCP)
|
|
|
|
if r.store != nil {
|
|
// The one-time link that delivers a key to the person it was issued
|
|
// for. Served without a key, because the person opening it is the one
|
|
// who does not have one yet; what protects it is the token.
|
|
claim := web.Claim(r.store, version.Version, r.log)
|
|
mux.HandleFunc("GET /claim/{token}", claim)
|
|
mux.HandleFunc("POST /claim/{token}", claim)
|
|
}
|
|
if r.accounts != nil {
|
|
// Signup, sign-in and the profile page. These are browser surfaces
|
|
// authenticated by a session cookie, not by an API key, which is why
|
|
// they are exempt below.
|
|
pages := web.NewAccountPages(r.accounts, r.store, r.store, version.Version, r.log)
|
|
// So the example on the profile page names a model that is actually
|
|
// running, rather than a placeholder that gets sent as-is.
|
|
pages.ServedModels(r.exampleModels)
|
|
pages.SetModelFacts(func(name string) web.ModelFacts {
|
|
_, m := r.resolve(name)
|
|
return web.ModelFacts{ContextLength: m.ContextLength, MaxOutput: m.MaxOutput, Tools: m.Capabilities.Tools}
|
|
})
|
|
pages.SetPersonalOff(r.personalOff)
|
|
pages.SetLlamaCppVersion(r.llamaVersion)
|
|
pages.SetPersonalDefaults(r.personalDefaults)
|
|
if r.store != nil {
|
|
pages.SetAudit(r.store)
|
|
// A forgotten password: a link mailed when the cluster's
|
|
// settings say how, otherwise asked of an operator.
|
|
pages.SetPasswordReset(r.store, r.mail, r.mail.BaseURL)
|
|
if r.oidc.Enabled() {
|
|
if err := pages.RegisterOIDC(mux, r.oidc, r.store); err != nil {
|
|
r.log.Error("single sign-on is off", "err", err)
|
|
}
|
|
}
|
|
if r.scim.TokenFile != "" {
|
|
if b, err := secrets.Read(context.Background(), r.scim.TokenFile); err != nil || strings.TrimSpace(string(b)) == "" {
|
|
r.log.Error("provisioning (SCIM) is off: its token cannot be read", "file", r.scim.TokenFile, "err", err)
|
|
} else {
|
|
scim.New(r.store, strings.TrimSpace(string(b)),
|
|
scim.Mapping{Roles: r.oidc.Roles, Memberships: r.oidc.Memberships}, r.log).Register(mux)
|
|
}
|
|
}
|
|
}
|
|
if pol, ok := r.guard.(interface {
|
|
Policy(guard.Subject) (guard.Policy, string)
|
|
}); ok {
|
|
pages.SetGuardPolicies(func(s guard.Subject) (string, bool) {
|
|
p, names := pol.Policy(s)
|
|
return names, p.NoStore && p.Mode == "enforce"
|
|
})
|
|
}
|
|
if r.store != nil && r.guardLibrary != nil {
|
|
pages.RegisterAdminGuard(mux, r.store, r.guardLibrary)
|
|
}
|
|
if r.store != nil && r.settingsCheck != nil {
|
|
pages.RegisterAdminSettings(mux, r.store, r.settingsCheck, r.settingsLocal)
|
|
}
|
|
pages.Register(mux)
|
|
// The docs portal, with the menu of whoever is signed in, and only to
|
|
// them when the site asked for that.
|
|
docs := pages.DocsPage(web.Docs(version.Version, true), r.landing.DocsSignIn)
|
|
mux.HandleFunc("GET /docs", docs)
|
|
mux.HandleFunc("GET /docs/{doc}", docs)
|
|
// The APIs' reference, in Swagger UI: the client API's document to
|
|
// whoever may read the docs, the admin API's to staff.
|
|
mux.HandleFunc("GET /docs/api", pages.DocsPage(web.APIDocs(version.Version, true), r.landing.DocsSignIn))
|
|
mux.HandleFunc("GET /docs/api/admin.json", pages.Guard(func(w http.ResponseWriter, req *http.Request, user store.User) {
|
|
if !user.Staff() {
|
|
http.NotFound(w, req)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write(llmgw.AdminOpenAPI)
|
|
}))
|
|
r.registerChat(mux, pages)
|
|
// An organization's own page, for the people who manage it.
|
|
pages.RegisterOrgPages(mux, r.store, r.acct)
|
|
|
|
if r.clusterView != nil {
|
|
// The cluster dashboard, reachable from a browser without exposing
|
|
// the admin listener — which also carries pprof, and pprof can dump
|
|
// the heap that holds every API key in plaintext.
|
|
mux.HandleFunc("GET /admin", pages.AdminPage(web.Dashboard(
|
|
r.clusterView, r.clusterUsage, version.Version, 5*time.Second, false,
|
|
web.NavExtras{
|
|
Mirror: r.mirror != nil, History: r.history, Window: r.histWindow,
|
|
}, r.log)))
|
|
mux.HandleFunc("GET /admin/nodes/{id}/logs", pages.AdminPage(
|
|
web.NodeLogsPage(r.clusterView, false, version.Version)))
|
|
mux.HandleFunc("GET /admin/nvlink", pages.AdminPage(web.NVLinkPage(r.clusterView, false, version.Version)))
|
|
mux.HandleFunc("GET /admin/fabrics", pages.AdminPage(web.FabricsPage(r.clusterView, false, version.Version)))
|
|
mux.HandleFunc("GET /admin/fabrics/{id}", pages.AdminPage(web.FabricsPage(r.clusterView, false, version.Version)))
|
|
mux.HandleFunc("GET /admin/fabrics/{id}/topology.json", pages.AdminPage(web.FabricTopologyJSON(r.clusterView)))
|
|
mux.HandleFunc("GET /admin/nodes/{id}", pages.AdminPage(
|
|
web.NodePage(r.clusterView, r.history, r.histWindow, false, version.Version, r.log)))
|
|
pages.RegisterAdminUsers(mux, r.store)
|
|
// Quotas are written on the same tree the accounts page builds, so
|
|
// they are mounted together.
|
|
// Capacity quotas are the control plane's, reached through the
|
|
// same admin API the CLI uses.
|
|
if c, ok := r.clusterView.(web.CapacityAdmin); ok {
|
|
pages.SetCapacity(c)
|
|
}
|
|
// The engine build each GPU node installs, chosen on the cluster
|
|
// page.
|
|
if e, ok := r.clusterView.(web.NodeEngineAdmin); ok {
|
|
pages.RegisterNodeEngines(mux, e)
|
|
}
|
|
pages.RegisterAdminQuotas(mux, r.store, r.acct)
|
|
// And the report of what that tree actually used.
|
|
if r.history != nil {
|
|
pages.SetHistory(r.history)
|
|
}
|
|
pages.RegisterAdminUsage(mux, r.store, r.acct)
|
|
}
|
|
|
|
// Independent of the cluster view: the mirror is a separate machine
|
|
// with its own disk, and it is worth administering whether or not this
|
|
// process also has a cluster to show.
|
|
if r.mirror != nil {
|
|
pages.RegisterAdminMirror(mux, web.MirrorDeps{
|
|
Mirror: r.mirror, Registry: r.registry,
|
|
Catalog: r.catalogAdmin, Fleet: r.clusterView,
|
|
})
|
|
}
|
|
} else {
|
|
// No accounts, so nobody to show a menu to or to send to sign in.
|
|
docs := web.Docs(version.Version, false)
|
|
mux.HandleFunc("GET /docs", docs)
|
|
mux.HandleFunc("GET /docs/{doc}", docs)
|
|
mux.HandleFunc("GET /docs/api", web.APIDocs(version.Version, false))
|
|
}
|
|
|
|
if r.auth == nil {
|
|
return r.refuseWhileStopped(r.refuseLoops(mux))
|
|
}
|
|
exempt := []string{"/health", "/ready", "/leader", "/v1/openapi.json", "/", "/about", "/_static/", "/_lang/", "/claim/",
|
|
"/signup", "/login", "/login/", "/logout", "/profile", "/profile/", "/scim/",
|
|
"/admin", "/admin/", "/chat", "/chat/", "/images", "/images/", "/3d", "/3d/", "/org", "/org/", "/setup/"}
|
|
// The docs are public unless the site said otherwise, and then guarded by
|
|
// the session check in front of them — or, on a gateway with keys and no
|
|
// accounts to sign in to, by the key like any other path.
|
|
if !r.landing.DocsSignIn || r.accounts != nil {
|
|
exempt = append(exempt, "/docs", "/docs/")
|
|
}
|
|
// Wrapping the whole mux rather than checking inside each handler: adding a
|
|
// route must not be able to add an unauthenticated one by omission. Health
|
|
// stays open because a load balancer has no key.
|
|
return r.refuseWhileStopped(r.refuseLoops(auth.Middleware(r.auth, r.log, exempt...)(mux)))
|
|
}
|
|
|
|
// refuseWhileStopped answers every request for a model with 503 while the
|
|
// operators have stopped the service (llmgw stop): the models are being taken
|
|
// down, and a client is better told so plainly than left to time out.
|
|
func (r *Router) refuseWhileStopped(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
|
if r.stopped.Load() && req.Method == http.MethodPost && strings.HasPrefix(req.URL.Path, "/v1/") {
|
|
w.Header().Set("Retry-After", "300")
|
|
writeErrorCode(w, http.StatusServiceUnavailable, "server_error", "system_stopped",
|
|
"this gateway's service has been stopped by its operators; it answers again once they start it")
|
|
return
|
|
}
|
|
next.ServeHTTP(w, req)
|
|
})
|
|
}
|
|
|
|
// parsedRequest is what the router needs from a request. Everything else in the
|
|
// body is forwarded untouched — clients send engine-specific fields we have no
|
|
// business interpreting.
|
|
type parsedRequest struct {
|
|
Path string
|
|
// ContentType is the body's when it is a form rather than JSON — audio
|
|
// for transcription — and is sent on as it came.
|
|
ContentType string
|
|
// Tags are the client's metadata and x-llmgw-tags; EndUser is OpenAI's
|
|
// user field. Both are recorded with the usage, never sent on.
|
|
Tags map[string]string
|
|
EndUser string
|
|
// logRetain is how long the request and its answer are kept, when its
|
|
// project keeps them and it fell in the share kept.
|
|
logRetain time.Duration
|
|
// Model is what will be dispatched, which is not always what was asked
|
|
// for: a caller out of budget may be moved to a model they can afford.
|
|
Model string
|
|
// RequestedModel is what the client sent, kept so the response can say
|
|
// when something else answered.
|
|
RequestedModel string
|
|
|
|
Stream bool
|
|
SuppressUsageChunk bool
|
|
UpstreamBody []byte
|
|
Header http.Header
|
|
|
|
// RequestID follows the request through the access log, the usage record
|
|
// and the engine, and goes back to the client in X-Request-Id.
|
|
RequestID string
|
|
// QueueMs is how long admission held it.
|
|
QueueMs int64
|
|
|
|
SessionID string
|
|
Band string
|
|
// KeyID, ProjectID and Tenant are empty when authentication is off.
|
|
KeyID string
|
|
UserID string
|
|
ProjectID string
|
|
TeamID, OrgID string
|
|
Tenant string
|
|
ProfileOverride string
|
|
PromptPrefix string
|
|
PromptTokens int
|
|
// ContextTokens is what has to fit in one context: the prompt, except
|
|
// for embeddings, where each input is embedded on its own and the
|
|
// longest one is what a context must hold.
|
|
ContextTokens int
|
|
MaxTokens int
|
|
|
|
// ThreadID is the conversation this request continues, empty for the
|
|
// stateless path every request took before threads existed.
|
|
ThreadID string
|
|
// ThreadTurns are the turns this request itself carried, kept so the
|
|
// exchange can be stored after a reply without re-reading a body that has
|
|
// since had the history spliced into it.
|
|
ThreadTurns []store.Message
|
|
// ThreadProject is the project the thread is stored under, resolved once by
|
|
// loadThread so saving it afterwards cannot disagree.
|
|
ThreadProject string
|
|
|
|
// reservation is carried from the limit check to the point where the
|
|
// request is settled, so the record counts against the same budget window
|
|
// the check read.
|
|
reservation accounting.Reservation
|
|
// guardIn is what the guard decided about the request.
|
|
guardIn guard.Verdict
|
|
// knowledge is set when passages from knowledge bases were added.
|
|
knowledge bool
|
|
// kbSalt is the collections whose passages were added, for the prefix
|
|
// cache's salt.
|
|
kbSalt string
|
|
// noStore is a policy that keeps nothing of this request.
|
|
noStore bool
|
|
// inCluster is a project whose requests may not leave the cluster.
|
|
inCluster bool
|
|
// internal is a request the gateway makes of itself — a judge's — which
|
|
// no guard filters.
|
|
internal bool
|
|
// cache is where its answer is kept, when its project reuses answers
|
|
// and this one may be.
|
|
cache cacheState
|
|
}
|
|
|
|
// inferenceBody covers the fields the router reads across all inference
|
|
// endpoints. Unknown fields are ignored here and preserved in UpstreamBody.
|
|
type inferenceBody struct {
|
|
Model string `json:"model"`
|
|
Stream bool `json:"stream"`
|
|
User string `json:"user"`
|
|
Metadata json.RawMessage `json:"metadata"`
|
|
Messages []struct {
|
|
Role string `json:"role"`
|
|
Content json.RawMessage `json:"content"`
|
|
} `json:"messages"`
|
|
Prompt json.RawMessage `json:"prompt"`
|
|
Input json.RawMessage `json:"input"`
|
|
MaxTokens int `json:"max_tokens"`
|
|
MaxCompletionTokens int `json:"max_completion_tokens"`
|
|
StreamOptions *struct {
|
|
IncludeUsage bool `json:"include_usage"`
|
|
} `json:"stream_options"`
|
|
}
|
|
|
|
func (r *Router) handleInference(path string) http.HandlerFunc {
|
|
return func(w http.ResponseWriter, req *http.Request) {
|
|
start := time.Now()
|
|
// First, so even a request refused as malformed can be named when
|
|
// somebody asks what happened to it.
|
|
rid := requestID(req.Header.Get("X-Request-Id"))
|
|
w.Header().Set("X-Request-Id", rid)
|
|
ctx, span := startInference(req, path, rid)
|
|
defer span.End()
|
|
req = req.WithContext(ctx)
|
|
|
|
parsed, err := r.parse(req, path)
|
|
if err == nil {
|
|
span.SetAttributes(attribute.String("llmgw.model", parsed.Model))
|
|
parsed.RequestID = rid
|
|
if parsed.Header == nil {
|
|
parsed.Header = http.Header{}
|
|
}
|
|
parsed.Header.Set("X-Request-Id", rid)
|
|
}
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
|
r.m.observeRequest("", "", http.StatusBadRequest, time.Since(start), nil, 0)
|
|
return
|
|
}
|
|
|
|
// A picture asked of a model that writes text, or the other way
|
|
// round, is refused before it waits for either.
|
|
if why := r.imageRouteRefusal(path, parsed.Model, parsed.UpstreamBody); why != "" {
|
|
writeError(w, http.StatusBadRequest, "invalid_request_error", why)
|
|
r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusBadRequest, time.Since(start), nil, 0)
|
|
return
|
|
}
|
|
|
|
// Before authentication work and before admission: a request that
|
|
// cannot fit should not consume a quota reservation or a queue slot on
|
|
// its way to being refused.
|
|
if _, entry := r.resolve(parsed.Model); entry.ContextLength > 0 {
|
|
if why := checkContextFits(entry, parsed.contextTokens(), parsed.MaxTokens); why != "" {
|
|
writeError(w, http.StatusBadRequest, "invalid_request_error", why)
|
|
r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusBadRequest, time.Since(start), nil, 0)
|
|
return
|
|
}
|
|
}
|
|
|
|
p0, authed := auth.FromContext(req.Context())
|
|
if r.hiddenModel(parsed.Model) || r.adapterHidden(p0, authed, parsed.Model) {
|
|
// Said as a model that is not there: it is not, for clients.
|
|
writeErrorCode(w, http.StatusNotFound, "invalid_request_error", "model_not_found",
|
|
fmt.Sprintf("the model %q does not exist", parsed.Model))
|
|
r.m.observeRequest("", "", http.StatusNotFound, time.Since(start), nil, 0)
|
|
return
|
|
}
|
|
if p, ok := auth.FromContext(req.Context()); ok {
|
|
if !r.permits(p, parsed.Model) {
|
|
auth.WriteForbiddenModel(w, parsed.Model)
|
|
r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusForbidden, time.Since(start), nil, 0)
|
|
return
|
|
}
|
|
// Limits are checked before admission, not after: a request that is
|
|
// over its quota should be refused rather than queued, and queueing
|
|
// it would have it occupy a band's capacity to no purpose.
|
|
if r.acct != nil {
|
|
// Priced by the model that will serve, not by the name the
|
|
// caller typed. A price table is written against the catalog,
|
|
// and an alias asked for by its short name would miss every
|
|
// entry and fall to the default — charging a 1.5B at whatever
|
|
// rate the largest model set, or at nothing.
|
|
priced, _ := r.resolve(parsed.Model)
|
|
reservation, err := r.acct.Check(p, priced, parsed.estimatedTokens())
|
|
if err != nil && isBudgetSpent(err) {
|
|
// Out of money, not out of options. A smaller model the
|
|
// caller can still afford beats a refusal, and the catalog
|
|
// already says which those are.
|
|
if fb, ok := r.budgetFallback(p, parsed); ok {
|
|
reservation, err = fb, nil
|
|
// Said now rather than at the end: the proxy writes the
|
|
// response itself, so a header set after dispatch never
|
|
// reaches the client.
|
|
w.Header().Set("x-llmgw-served-model", parsed.Model)
|
|
}
|
|
}
|
|
// On every answer, refused or not, so a client can pace itself
|
|
// before it is refused rather than learn the limit from a 429.
|
|
r.rateLimitHeaders(w, p)
|
|
if err != nil {
|
|
r.rejectLimit(w, parsed, err, start)
|
|
return
|
|
}
|
|
parsed.reservation = reservation
|
|
}
|
|
}
|
|
|
|
// Step 5 of the lifecycle: "now?" before "where?" (plan §7.1).
|
|
_, aspan := tracer().Start(req.Context(), "admission")
|
|
lease, err := r.admit(req.Context(), parsed)
|
|
endSpan(aspan, err)
|
|
if err != nil {
|
|
r.rejectAdmission(w, parsed, err, start)
|
|
return
|
|
}
|
|
defer lease.Release()
|
|
queued := time.Since(start)
|
|
parsed.QueueMs = queued.Milliseconds()
|
|
r.m.observeAdmission(lease.Band(), r.metricModel(parsed.Model), queued)
|
|
if queued > time.Millisecond {
|
|
w.Header().Set("x-llmgw-queue-ms", strconv.FormatInt(queued.Milliseconds(), 10))
|
|
}
|
|
|
|
// A project whose requests stay in the cluster is not sent outside it,
|
|
// as the model asked for or, below, as a fallback.
|
|
if r.inCluster(req.Context(), parsed.ProjectID) {
|
|
parsed.inCluster = true
|
|
if _, entry := r.resolve(parsed.Model); entry.IsExternal() {
|
|
writeErrorCode(w, http.StatusForbidden, "invalid_request_error", "in_cluster",
|
|
"this project's requests stay in the cluster, and "+parsed.Model+" is served outside it")
|
|
r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusForbidden, time.Since(start), nil, 0)
|
|
return
|
|
}
|
|
}
|
|
|
|
// A policy that keeps nothing is known before anything is kept: a
|
|
// conversation asked for is refused rather than quietly not stored,
|
|
// and a model outside the cluster is not sent the prompt.
|
|
if model, entry := r.resolve(parsed.Model); r.noStore(parsed.guardSubject(model)) {
|
|
parsed.noStore = true
|
|
if parsed.ThreadID != "" || entry.IsExternal() {
|
|
writeErrorCode(w, http.StatusBadRequest, "invalid_request_error", "no_store", errNoStoreMsg)
|
|
r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusBadRequest, time.Since(start), nil, 0)
|
|
return
|
|
}
|
|
}
|
|
|
|
// After the model is finally settled — a budget fallback may have
|
|
// changed it — because trimming is measured against the context of the
|
|
// model that will actually serve, not the one that was asked for.
|
|
if dropped, err := r.loadThread(req.Context(), parsed); err != nil {
|
|
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
|
r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusBadRequest, time.Since(start), nil, 0)
|
|
return
|
|
} else if dropped > 0 {
|
|
// Dropping turns is the right answer and must not be a silent one.
|
|
w.Header().Set(trimmedHeader, strconv.Itoa(dropped))
|
|
}
|
|
|
|
// The passages a question asked knowledge bases for, before it is
|
|
// judged: they are text the model reads, and a document can carry
|
|
// instructions aimed at its reader.
|
|
if !r.retrieve(req.Context(), w, parsed, start) {
|
|
return
|
|
}
|
|
|
|
// Judged as it will be sent: authorised, within its limits, and with
|
|
// its stored conversation in it, which is text the model reads too.
|
|
if r.guard != nil && !r.guardInput(req.Context(), w, parsed, start) {
|
|
return
|
|
}
|
|
|
|
r.decideRequestLog(req.Context(), parsed)
|
|
// After the guardrails, so a policy still decides whether a request is
|
|
// answered, and the request looked up is the one the model would be
|
|
// sent.
|
|
if a := r.cacheLookup(req.Context(), w, parsed); a != nil {
|
|
r.serveCached(req.Context(), w, parsed, a, start)
|
|
return
|
|
}
|
|
dw, cw := r.cacheWriterFor(w, parsed)
|
|
res := r.dispatch(req.Context(), dw, parsed)
|
|
status := r.finish(w, parsed, res, start)
|
|
r.cacheKeep(parsed, res, status, cw)
|
|
r.observeCanary(parsed, res, status, time.Since(start))
|
|
if parsed.logRetain > 0 {
|
|
go r.keepRequest(context.WithoutCancel(req.Context()), parsed, res, status)
|
|
}
|
|
r.observeOutput(parsed, res, res.judged)
|
|
span.SetAttributes(attribute.Int("http.response.status_code", status),
|
|
attribute.String("llmgw.served_model", res.model),
|
|
attribute.String("llmgw.endpoint", res.endpointID))
|
|
if status >= 500 {
|
|
span.SetStatus(codes.Error, http.StatusText(status))
|
|
}
|
|
// After the client has been served. A conversation that failed to
|
|
// record is a worse outcome than a slow reply, but only just, and the
|
|
// caller has their answer either way.
|
|
if res.err == nil && res.outcome.Status == http.StatusOK {
|
|
r.saveThread(context.WithoutCancel(req.Context()), parsed, res.outcome.Reply, parsed.Model)
|
|
}
|
|
}
|
|
}
|
|
|
|
// permits reports whether a key's allowlist covers the requested model.
|
|
//
|
|
// The decision is per model, not per string. A model's canonical name and every
|
|
// alias of it name the same thing, so the allowlist is matched against all of
|
|
// them: writing either form works, and no alias can reach a model that was not
|
|
// allowed. Matching only what the client sent would make every alias a way
|
|
// around the allowlist; matching only the canonical name would silently ignore
|
|
// an allowlist written in the names clients actually use.
|
|
func (r *Router) permits(p auth.Principal, requested string) bool {
|
|
// A judge model is the guardrails', not anybody's to ask; a canary is
|
|
// reached through the model it is a canary of.
|
|
if r.hiddenModel(requested) {
|
|
return false
|
|
}
|
|
// An adapter named in the allowlist allows that adapter: not its model,
|
|
// and not the model's other adapters, which the checks below would let
|
|
// through if the adapter's own name were not asked about first.
|
|
if _, ok := r.adapter(requested); ok && p.Allows(requested) {
|
|
return true
|
|
}
|
|
canonical, m := r.resolve(requested)
|
|
if p.Allows(canonical) {
|
|
return true
|
|
}
|
|
for _, alias := range m.Aliases {
|
|
if p.Allows(alias) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// finish records metrics and, when nothing was written, produces the error body.
|
|
// It returns the status the request ended with.
|
|
func (r *Router) finish(w http.ResponseWriter, req *parsedRequest, res dispatchResult, start time.Time) int {
|
|
status := res.outcome.Status
|
|
switch {
|
|
case res.err == nil:
|
|
// The served-model header is set before dispatch, not here: the proxy
|
|
// writes the response, so by this point the headers are long gone.
|
|
case errors.Is(res.err, proxy.ErrClientGone):
|
|
// Nothing to report to a client that is gone; the upstream was cancelled.
|
|
status = 499
|
|
r.log.Debug("client disconnected", "model", req.Model, "endpoint", res.endpointID)
|
|
case res.outcome.Wrote:
|
|
// The response is already partly out; the stream just ends short.
|
|
r.log.Warn("stream ended early", "model", req.Model, "endpoint", res.endpointID, "err", res.err)
|
|
if status == 0 {
|
|
status = http.StatusOK
|
|
}
|
|
case errors.Is(res.err, errLoading):
|
|
// Retry-After, because this one succeeds on a retry and a client that
|
|
// is told so will make it. Without the header a well-behaved client
|
|
// backs off exactly as it would from a permanent failure.
|
|
status = http.StatusServiceUnavailable
|
|
w.Header().Set("Retry-After", "30")
|
|
writeError(w, status, "server_error", res.err.Error())
|
|
case errors.Is(res.err, errNoCapacity):
|
|
status = http.StatusServiceUnavailable
|
|
writeError(w, status, "server_error", res.err.Error())
|
|
default:
|
|
status = http.StatusBadGateway
|
|
writeError(w, status, "server_error", res.err.Error())
|
|
}
|
|
|
|
r.m.observeRequest(r.metricModel(req.Model), res.pool, status, time.Since(start), res.outcome.Usage, res.outcome.TTFT)
|
|
if status < 400 {
|
|
r.observeTokens(req, res, time.Since(start))
|
|
}
|
|
r.account(req, res, status, time.Since(start))
|
|
r.accessLog(req, res, status, time.Since(start))
|
|
return status
|
|
}
|
|
|
|
// account settles a finished request against its key.
|
|
//
|
|
// It runs for every outcome, not only success: a request that used a GPU and
|
|
// then failed still used it, and a key that could spend nothing by failing
|
|
// would be a way around a budget.
|
|
func (r *Router) account(req *parsedRequest, res dispatchResult, status int, took time.Duration) {
|
|
if r.acct == nil || req.KeyID == "" {
|
|
return
|
|
}
|
|
u := store.Usage{
|
|
At: time.Now(), KeyID: req.KeyID, UserID: req.UserID, ProjectID: req.ProjectID,
|
|
Model: req.Model, ServedModel: res.model, Status: status,
|
|
LatencyMs: took.Milliseconds(),
|
|
RequestID: req.RequestID, Replica: res.endpointID,
|
|
QueueMs: req.QueueMs, TTFTMs: res.outcome.TTFT.Milliseconds(),
|
|
EndUser: req.EndUser, Tags: req.Tags,
|
|
}
|
|
guardFields(&u, req.guardIn, res.guard)
|
|
// An answer reused generates nothing: its usage is the zero the hit
|
|
// carries, and it costs nothing against a budget or a token limit.
|
|
u.CacheHit = res.cached
|
|
if u.ServedModel == "" {
|
|
u.ServedModel = req.Model
|
|
}
|
|
if usage := res.outcome.Usage; usage != nil {
|
|
u.PromptTokens = int64(usage.PromptTokens)
|
|
u.CompletionTokens = int64(usage.CompletionTokens)
|
|
} else {
|
|
// No usage came back — a failure before the engine answered, or an
|
|
// engine that does not report it. Charging the estimate is closer to
|
|
// the truth than charging nothing.
|
|
u.PromptTokens = int64(req.PromptTokens)
|
|
}
|
|
u.Images = int64(res.outcome.Images)
|
|
u.Meshes = int64(res.outcome.Meshes)
|
|
r.m.tenant.observe(req, status, u.PromptTokens, u.CompletionTokens, res.outcome.TTFT)
|
|
r.acct.Record(req.reservation, u)
|
|
}
|
|
|
|
// rejectLimit turns a quota refusal into a status the client can act on.
|
|
func (r *Router) rejectLimit(w http.ResponseWriter, req *parsedRequest, err error, start time.Time) {
|
|
var le *accounting.LimitError
|
|
if !errors.As(err, &le) {
|
|
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
|
return
|
|
}
|
|
if le.RetryAfter > 0 {
|
|
w.Header().Set("Retry-After", strconv.Itoa(int(le.RetryAfter.Seconds())))
|
|
}
|
|
// 429 for all of them: a spent budget is a quota, and OpenAI clients
|
|
// already know how to surface that code. What differs is Retry-After,
|
|
// which is absent when waiting does not help.
|
|
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", le.Error())
|
|
r.m.observeRequest(r.metricModel(req.Model), "", http.StatusTooManyRequests, time.Since(start), nil, 0)
|
|
r.m.tenant.observe(req, http.StatusTooManyRequests, 0, 0, 0)
|
|
r.acct.RecordRefused(store.Usage{
|
|
At: time.Now(), KeyID: req.KeyID, UserID: req.UserID, ProjectID: req.ProjectID,
|
|
Model: req.Model, ServedModel: req.Model, Status: http.StatusTooManyRequests,
|
|
LatencyMs: time.Since(start).Milliseconds(), RequestID: req.RequestID,
|
|
EndUser: req.EndUser, Tags: req.Tags,
|
|
})
|
|
r.accessLog(req, dispatchResult{err: err}, http.StatusTooManyRequests, time.Since(start))
|
|
}
|
|
|
|
// parse reads the body and extracts the routing-relevant fields.
|
|
func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) {
|
|
body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if ct := req.Header.Get("Content-Type"); strings.HasPrefix(ct, "multipart/form-data") {
|
|
return r.parseForm(req, path, ct, body)
|
|
}
|
|
|
|
var b inferenceBody
|
|
if err := json.Unmarshal(body, &b); err != nil {
|
|
return nil, err
|
|
}
|
|
if b.Model == "" {
|
|
return nil, errors.New("model is required")
|
|
}
|
|
|
|
tags, err := requestTags(b.Metadata, req.Header.Get(tagsHeader))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
p := &parsedRequest{
|
|
Tags: tags,
|
|
EndUser: endUser(b.User),
|
|
Path: path,
|
|
Model: b.Model,
|
|
RequestedModel: b.Model,
|
|
Stream: b.Stream,
|
|
UpstreamBody: body,
|
|
Header: req.Header,
|
|
SessionID: sessionID(req, r.cfg.SessionHeader, b.User),
|
|
Band: req.Header.Get(r.cfg.Admission.BandHeader),
|
|
ProfileOverride: req.Header.Get("x-llmgw-profile"),
|
|
MaxTokens: max(b.MaxTokens, b.MaxCompletionTokens),
|
|
PromptPrefix: promptPrefix(&b),
|
|
ThreadID: req.Header.Get(threadHeader),
|
|
}
|
|
// With authentication on, the key decides the band and the routing profile,
|
|
// and the headers stop being believed. Trusting the header was only ever
|
|
// acceptable while there was no identity to consult — it let any client
|
|
// promote itself into the top priority band.
|
|
if principal, ok := auth.FromContext(req.Context()); ok {
|
|
p.Band = principal.Band
|
|
p.ProfileOverride = principal.RoutingProfile
|
|
p.KeyID = principal.KeyID
|
|
p.UserID = principal.UserID
|
|
p.ProjectID = principal.ProjectID
|
|
p.TeamID, p.OrgID = principal.TeamID, principal.OrgID
|
|
// Tenant fairness within a band is per key rather than per session, so
|
|
// one key cannot take a whole band by opening many sessions.
|
|
p.Tenant = principal.KeyID
|
|
}
|
|
|
|
// Bytes over four: a crude stand-in for a tokenizer, good enough for
|
|
// scoring and for limits. It becomes exact only with plan §7.5 tier 2 (see
|
|
// plan §14). Measured over the whole prompt, not the routing prefix —
|
|
// see promptBytes.
|
|
p.PromptTokens = promptBytes(&b) / 4
|
|
p.ContextTokens = p.PromptTokens
|
|
if path == "/v1/embeddings" {
|
|
if longest, ok := longestInput(b.Input); ok {
|
|
p.ContextTokens = longest / 4
|
|
}
|
|
}
|
|
|
|
// Tags are the gateway's: an engine has no use for them, and its logs
|
|
// are no place for what a client wrote there.
|
|
if len(b.Metadata) > 0 {
|
|
p.UpstreamBody = withoutField(p.UpstreamBody, "metadata")
|
|
}
|
|
|
|
// Ask the engine for usage the client did not request, so accounting works
|
|
// regardless of client behaviour; the extra chunk is dropped on the way back.
|
|
if r.cfg.InjectUsage && b.Stream && (b.StreamOptions == nil || !b.StreamOptions.IncludeUsage) {
|
|
if injected, ok := injectIncludeUsage(p.UpstreamBody); ok {
|
|
p.UpstreamBody = injected
|
|
p.SuppressUsageChunk = true
|
|
}
|
|
}
|
|
return p, nil
|
|
}
|
|
|
|
// sessionID picks the affinity key: an explicit header first, then OpenAI's
|
|
// user field, which most clients already set per end user.
|
|
func sessionID(req *http.Request, header, user string) string {
|
|
if header != "" {
|
|
if v := req.Header.Get(header); v != "" {
|
|
return v
|
|
}
|
|
}
|
|
return user
|
|
}
|
|
|
|
// promptPrefix reconstructs the leading portion of the prompt for prefix-aware
|
|
// routing. Only the first prefixBytes matter, so long conversations stop early.
|
|
func promptPrefix(b *inferenceBody) string {
|
|
var sb strings.Builder
|
|
appendRaw := func(raw json.RawMessage) {
|
|
if len(raw) == 0 || sb.Len() >= prefixBytes {
|
|
return
|
|
}
|
|
var s string
|
|
if err := json.Unmarshal(raw, &s); err != nil {
|
|
s = string(raw)
|
|
}
|
|
sb.WriteString(s)
|
|
}
|
|
|
|
for _, m := range b.Messages {
|
|
if sb.Len() >= prefixBytes {
|
|
break
|
|
}
|
|
sb.WriteString(m.Role)
|
|
sb.WriteString(":")
|
|
appendRaw(m.Content)
|
|
sb.WriteString("\n")
|
|
}
|
|
appendRaw(b.Prompt)
|
|
appendRaw(b.Input)
|
|
|
|
s := sb.String()
|
|
if len(s) > prefixBytes {
|
|
s = s[:prefixBytes]
|
|
}
|
|
return s
|
|
}
|
|
|
|
// promptBytes measures the whole prompt, which promptPrefix deliberately does
|
|
// not.
|
|
//
|
|
// The two were the same function for too long, and the estimate was taken from
|
|
// the prefix: bytes over four of at most prefixBytes, so every prompt above
|
|
// 2 KB counted as 512 tokens however long it really was. Prefix routing wants a
|
|
// bounded head of the conversation; a token estimate wants all of it, and using
|
|
// one for the other meant TPM limits undercounted exactly the requests that
|
|
// consume the most, and nothing could tell whether a prompt fit in a context.
|
|
// longestInput is the length of an embeddings request's longest input, when
|
|
// input is a list of strings: a batch of texts, each its own context.
|
|
func longestInput(raw json.RawMessage) (int, bool) {
|
|
var list []string
|
|
if len(raw) == 0 || raw[0] != '[' || json.Unmarshal(raw, &list) != nil || len(list) == 0 {
|
|
return 0, false
|
|
}
|
|
n := 0
|
|
for _, s := range list {
|
|
n = max(n, len(s))
|
|
}
|
|
return n, true
|
|
}
|
|
|
|
// contextTokens is what has to fit in one context; requests built inside
|
|
// the router set only PromptTokens.
|
|
func (p *parsedRequest) contextTokens() int {
|
|
if p.ContextTokens > 0 {
|
|
return p.ContextTokens
|
|
}
|
|
return p.PromptTokens
|
|
}
|
|
|
|
func promptBytes(b *inferenceBody) int {
|
|
n := 0
|
|
add := func(raw json.RawMessage) {
|
|
if len(raw) == 0 {
|
|
return
|
|
}
|
|
var str string
|
|
if err := json.Unmarshal(raw, &str); err == nil {
|
|
n += len(str)
|
|
return
|
|
}
|
|
// Not a plain string: structured content, an array of parts, an image.
|
|
// The raw JSON overstates the text slightly, which is the safe
|
|
// direction for a limit.
|
|
n += len(raw)
|
|
}
|
|
for _, m := range b.Messages {
|
|
n += len(m.Role) + 2
|
|
add(m.Content)
|
|
}
|
|
add(b.Prompt)
|
|
add(b.Input)
|
|
return n
|
|
}
|
|
|
|
// injectIncludeUsage sets stream_options.include_usage on the outgoing body.
|
|
//
|
|
// This is the one place the router rewrites a request, so it is done by
|
|
// decoding to a generic map and re-encoding: field order changes, but no field
|
|
// is lost. If anything looks unexpected it declines rather than guessing.
|
|
func injectIncludeUsage(body []byte) ([]byte, bool) {
|
|
var m map[string]json.RawMessage
|
|
if err := json.Unmarshal(body, &m); err != nil {
|
|
return nil, false
|
|
}
|
|
|
|
opts := map[string]any{}
|
|
if raw, ok := m["stream_options"]; ok {
|
|
if err := json.Unmarshal(raw, &opts); err != nil {
|
|
return nil, false
|
|
}
|
|
}
|
|
opts["include_usage"] = true
|
|
|
|
encoded, err := json.Marshal(opts)
|
|
if err != nil {
|
|
return nil, false
|
|
}
|
|
m["stream_options"] = encoded
|
|
|
|
out, err := json.Marshal(m)
|
|
if err != nil {
|
|
return nil, false
|
|
}
|
|
return out, true
|
|
}
|
|
|
|
// noEndpointFor explains that nothing serves a model, and says what does.
|
|
//
|
|
// "no endpoint serves \"MODEL\"" was all a client got, which is true and
|
|
// useless: it reads the same whether the name was a typo, a placeholder left
|
|
// in from the documentation, or a real outage. Someone copying the curl from
|
|
// the account page sent the literal word MODEL and had nothing to go on.
|
|
//
|
|
// The list is what /v1/models already returns to the same caller, so naming it
|
|
// here discloses nothing new. It is capped because a large catalog would turn
|
|
// an error message into a page.
|
|
func noEndpointFor(model string, served []string) string {
|
|
if len(served) == 0 {
|
|
return fmt.Sprintf("no endpoint serves %q, and no model is currently being served", model)
|
|
}
|
|
const showAt = 10
|
|
shown := served
|
|
suffix := ""
|
|
if len(shown) > showAt {
|
|
shown, suffix = shown[:showAt], fmt.Sprintf(" and %d more", len(served)-showAt)
|
|
}
|
|
return fmt.Sprintf("no endpoint serves %q; available models are %s%s",
|
|
model, strings.Join(shown, ", "), suffix)
|
|
}
|
|
|
|
// exampleModels is servedModels with the ones that answer for real first.
|
|
//
|
|
// The profile page takes the first name for its copy-pasteable example, and
|
|
// handing it a simulated model would send someone's first request to an engine
|
|
// that replies in Greek letters.
|
|
func (r *Router) exampleModels() []string {
|
|
names := slices.DeleteFunc(r.servedModels(), r.hiddenModel)
|
|
slices.SortStableFunc(names, func(a, b string) int {
|
|
as, bs := r.modelMap()[a].Simulated(), r.modelMap()[b].Simulated()
|
|
switch {
|
|
case as == bs:
|
|
return 0
|
|
case bs:
|
|
return -1
|
|
default:
|
|
return 1
|
|
}
|
|
})
|
|
return names
|
|
}
|
|
|
|
// isGuardModel reports whether a name is a judge model, by any of its names.
|
|
func (r *Router) isGuardModel(name string) bool {
|
|
_, entry := r.resolve(name)
|
|
return entry.IsGuard()
|
|
}
|
|
|
|
// servedModels lists the canonical models at least one endpoint serves.
|
|
func (r *Router) servedModels() []string {
|
|
seen := map[string]struct{}{}
|
|
for _, s := range r.reg.Snapshots("") {
|
|
for _, m := range s.Models {
|
|
seen[m] = struct{}{}
|
|
}
|
|
}
|
|
out := make([]string, 0, len(seen))
|
|
for m := range seen {
|
|
out = append(out, m)
|
|
}
|
|
slices.Sort(out)
|
|
return out
|
|
}
|
|
|
|
// handleModels lists the models that at least one endpoint currently serves.
|
|
//
|
|
// Only those the caller's key may use: listing the rest tells a key holder
|
|
// what the gateway serves to others, and invites requests that are refused.
|
|
func (r *Router) handleModels(w http.ResponseWriter, req *http.Request) {
|
|
seen := map[string]struct{}{}
|
|
for _, s := range r.reg.Snapshots("") {
|
|
for _, m := range s.Models {
|
|
seen[m] = struct{}{}
|
|
}
|
|
}
|
|
// Catalog aliases are usable names too, so clients should see them.
|
|
for name := range r.modelMap() {
|
|
if _, ok := seen[r.modelMap()[name].Name]; ok {
|
|
seen[name] = struct{}{}
|
|
}
|
|
}
|
|
|
|
principal, authed := auth.FromContext(req.Context())
|
|
names := make([]string, 0, len(seen))
|
|
for n := range seen {
|
|
if authed && !r.permits(principal, n) || r.hiddenModel(n) || r.adapterHidden(principal, authed, n) {
|
|
continue
|
|
}
|
|
names = append(names, n)
|
|
}
|
|
slices.Sort(names)
|
|
|
|
data := make([]map[string]any, 0, len(names))
|
|
for _, n := range names {
|
|
// A simulated model is owned by "llmgw-sim", which is what the sim
|
|
// engine already calls itself on its own /v1/models. Flattening it to
|
|
// "llmgw" here was hiding the one distinction a client needs: this
|
|
// model returns placeholder text, not inference.
|
|
owner := "llmgw"
|
|
if r.modelMap()[n].Simulated() {
|
|
owner = "llmgw-sim"
|
|
}
|
|
entry := map[string]any{
|
|
"id": n, "object": "model", "created": time.Now().Unix(), "owned_by": owner,
|
|
}
|
|
// Not OpenAI's either: what a federated peer reading this list needs
|
|
// to leave out — a second name for a model already listed, and a
|
|
// model this cluster has from a peer of its own.
|
|
if m := r.modelMap()[n]; m.Name != "" && m.Name != n {
|
|
entry["alias_of"] = m.Name
|
|
}
|
|
if m := r.modelMap()[n]; m.IsPeer() {
|
|
entry["peer"] = m.External.Peer
|
|
}
|
|
// Only when known. A context_length of 0 would read as "no context",
|
|
// which is a worse answer than saying nothing.
|
|
m := r.modelMap()[n]
|
|
if m.ContextLength > 0 {
|
|
entry["context_length"] = m.ContextLength
|
|
}
|
|
if m.MaxOutput > 0 {
|
|
entry["max_output_tokens"] = m.MaxOutput
|
|
}
|
|
// Not part of the OpenAI shape, and the reason to send it anyway is
|
|
// that this list is the only view most clients have: nine names and
|
|
// nothing to choose between them.
|
|
if m.Description != "" {
|
|
entry["description"] = m.Description
|
|
}
|
|
if m.License != nil {
|
|
entry["license"] = m.License
|
|
}
|
|
entry["capabilities"] = map[string]bool{
|
|
"tools": m.Capabilities.Tools, "vision": m.Capabilities.Vision,
|
|
"embeddings": m.Capabilities.Embeddings, "images": m.MakesImages(), "meshes": m.MakesMeshes(),
|
|
}
|
|
data = append(data, entry)
|
|
}
|
|
if id := r.clusterID().id; id != "" {
|
|
w.Header().Set(clusterHeader, id)
|
|
}
|
|
writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": data})
|
|
}
|
|
|
|
func writeJSON(w http.ResponseWriter, code int, v any) {
|
|
var buf bytes.Buffer
|
|
if err := json.NewEncoder(&buf).Encode(v); err != nil {
|
|
http.Error(w, `{"error":{"message":"encoding failed"}}`, http.StatusInternalServerError)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(code)
|
|
_, _ = w.Write(buf.Bytes())
|
|
}
|
|
|
|
// writeError emits an OpenAI-shaped error so clients need no special handling
|
|
// for gateway failures.
|
|
func writeError(w http.ResponseWriter, code int, kind, msg string) {
|
|
writeErrorCode(w, code, kind, "", msg)
|
|
}
|
|
|
|
// writeErrorCode is writeError with OpenAI's machine-readable code, which
|
|
// clients branch on: content_policy_violation is the one SDKs know.
|
|
func writeErrorCode(w http.ResponseWriter, status int, kind, code, msg string) {
|
|
var c any
|
|
if code != "" {
|
|
c = code
|
|
}
|
|
writeJSON(w, status, map[string]any{"error": map[string]any{
|
|
"message": msg, "type": kind, "code": c,
|
|
}})
|
|
}
|
|
|
|
// estimatedTokens is what a token-rate limit is checked against, before the
|
|
// real figure exists.
|
|
//
|
|
// The prompt estimate plus whatever the client said it wanted back. A client
|
|
// that names no max_tokens gets a modest assumption rather than zero: assuming
|
|
// zero would let an unbounded request through every token limit.
|
|
func (p *parsedRequest) estimatedTokens() int64 {
|
|
const assumedCompletion = 256
|
|
completion := p.MaxTokens
|
|
if completion <= 0 {
|
|
completion = assumedCompletion
|
|
}
|
|
return int64(p.PromptTokens + completion)
|
|
}
|
|
|
|
// handleMe reports what the presented key is and what it may do.
|
|
//
|
|
// It exists so a user can answer "how much have I used" and "why was I
|
|
// refused" without asking an operator to read a dashboard they cannot reach.
|
|
// Deliberately outside /v1: it is not part of the OpenAI API, and putting it
|
|
// there would imply it is.
|
|
func (r *Router) handleMe(w http.ResponseWriter, req *http.Request) {
|
|
p, ok := auth.FromContext(req.Context())
|
|
if !ok {
|
|
// With authentication off there is no identity to describe, and saying
|
|
// so is more useful than an empty object.
|
|
writeJSON(w, http.StatusOK, map[string]any{
|
|
"authenticated": false,
|
|
"note": "this gateway has no api keys configured",
|
|
})
|
|
return
|
|
}
|
|
|
|
out := map[string]any{
|
|
"authenticated": true,
|
|
"key_id": p.KeyID,
|
|
"key": p.KeyPrefix,
|
|
"project": p.ProjectID,
|
|
"team": p.TeamID,
|
|
"organization": p.OrgID,
|
|
}
|
|
if p.Band != "" {
|
|
out["band"] = p.Band
|
|
}
|
|
limits := map[string]any{}
|
|
if p.Limits.RPM != nil {
|
|
limits["requests_per_minute"] = *p.Limits.RPM
|
|
}
|
|
if p.Limits.TPM != nil {
|
|
limits["tokens_per_minute"] = *p.Limits.TPM
|
|
}
|
|
if p.Limits.AllowedModels != nil {
|
|
limits["allowed_models"] = p.Limits.AllowedModels
|
|
}
|
|
if p.Limits.Budget != nil {
|
|
limits["budget_units"] = float64(*p.Limits.Budget) / 1e6
|
|
limits["budget_period"] = string(p.Limits.BudgetPeriod)
|
|
}
|
|
out["limits"] = limits
|
|
|
|
if r.acct != nil {
|
|
spent, rpm, tpm := r.acct.Status(p)
|
|
usage := map[string]any{
|
|
"requests_per_minute": rpm,
|
|
"tokens_per_minute": tpm,
|
|
}
|
|
if p.Limits.Budget != nil {
|
|
usage["spent_units"] = float64(spent) / 1e6
|
|
usage["remaining_units"] = float64(*p.Limits.Budget-spent) / 1e6
|
|
}
|
|
out["usage"] = usage
|
|
}
|
|
writeJSON(w, http.StatusOK, out)
|
|
}
|
|
|
|
// requestID is the client's X-Request-Id when it is one worth keeping, or a
|
|
// new one. Kept only when short and plain: it is written into logs and a
|
|
// header, and a client-chosen string with newlines or megabytes in it is not
|
|
// an id.
|
|
func requestID(sent string) string {
|
|
if n := len(sent); n > 0 && n <= 128 && strings.IndexFunc(sent, func(c rune) bool {
|
|
return !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' ||
|
|
c == '-' || c == '_' || c == '.' || c == ':')
|
|
}) < 0 {
|
|
return sent
|
|
}
|
|
var b [12]byte
|
|
_, _ = rand.Read(b[:])
|
|
return "req_" + hex.EncodeToString(b[:])
|
|
}
|
|
|
|
// accessLog writes one line per finished request: enough to follow a slow or
|
|
// failed request from the id a client reports to the replica that served it,
|
|
// and where its time went. Before this the only per-request lines were for
|
|
// disconnects, early stream ends and upstream retries.
|
|
func (r *Router) accessLog(req *parsedRequest, res dispatchResult, status int, took time.Duration) {
|
|
attrs := []any{
|
|
"request_id", req.RequestID, "status", status, "model", req.Model,
|
|
"duration_ms", took.Milliseconds(), "queue_ms", req.QueueMs, "stream", req.Stream,
|
|
}
|
|
if req.RequestedModel != "" && req.RequestedModel != req.Model {
|
|
attrs = append(attrs, "requested_model", req.RequestedModel)
|
|
}
|
|
if res.model != "" && res.model != req.Model {
|
|
attrs = append(attrs, "served_model", res.model)
|
|
}
|
|
if res.endpointID != "" {
|
|
attrs = append(attrs, "replica", res.endpointID)
|
|
}
|
|
if res.outcome.TTFT > 0 {
|
|
attrs = append(attrs, "ttft_ms", res.outcome.TTFT.Milliseconds())
|
|
}
|
|
if u := res.outcome.Usage; u != nil {
|
|
attrs = append(attrs, "prompt_tokens", u.PromptTokens, "completion_tokens", u.CompletionTokens)
|
|
}
|
|
if req.KeyID != "" {
|
|
attrs = append(attrs, "key", req.KeyID, "project", req.ProjectID)
|
|
}
|
|
if req.UserID != "" {
|
|
attrs = append(attrs, "user", req.UserID)
|
|
}
|
|
if res.err != nil {
|
|
attrs = append(attrs, "err", res.err.Error())
|
|
}
|
|
// What a guard did, by name: never what it matched.
|
|
if v := req.guardIn; res.guard.Stricter(v) {
|
|
attrs = append(attrs, "guard", string(res.guard.Action), "guard_stage", "output", "guard_policy", res.guard.Policy,
|
|
"guard_category", res.guard.Category, "guard_rule", res.guard.Rule)
|
|
} else if v.Action != "" {
|
|
attrs = append(attrs, "guard", string(v.Action), "guard_stage", "input", "guard_policy", v.Policy,
|
|
"guard_category", v.Category, "guard_rule", v.Rule)
|
|
}
|
|
if res.outcome.Stopped {
|
|
attrs = append(attrs, "stopped_by_guard", true)
|
|
}
|
|
if res.cached {
|
|
attrs = append(attrs, "cache", "hit")
|
|
}
|
|
r.log.Info("request", attrs...)
|
|
}
|
|
|
|
// rateLimitHeaders sets OpenAI's x-ratelimit-* headers from the tightest
|
|
// limit that applies. SDKs and editors read these to slow down before they
|
|
// are refused; only Retry-After, and only on a refusal, was sent before. A
|
|
// limit that is not set is not advertised.
|
|
func (r *Router) rateLimitHeaders(w http.ResponseWriter, p auth.Principal) {
|
|
if r.acct == nil {
|
|
return
|
|
}
|
|
h := r.acct.Headroom(p)
|
|
reset := strconv.FormatFloat(h.Reset.Seconds(), 'f', 0, 64) + "s"
|
|
if h.RequestLimit > 0 {
|
|
w.Header().Set("x-ratelimit-limit-requests", strconv.Itoa(h.RequestLimit))
|
|
w.Header().Set("x-ratelimit-remaining-requests", strconv.Itoa(h.RequestsLeft))
|
|
w.Header().Set("x-ratelimit-reset-requests", reset)
|
|
}
|
|
if h.TokenLimit > 0 {
|
|
w.Header().Set("x-ratelimit-limit-tokens", strconv.Itoa(h.TokenLimit))
|
|
w.Header().Set("x-ratelimit-remaining-tokens", strconv.Itoa(h.TokensLeft))
|
|
w.Header().Set("x-ratelimit-reset-tokens", reset)
|
|
}
|
|
}
|
|
|
|
// withModel is body with its model set to name, for an endpoint that knows the
|
|
// model by another name — an external provider. The body unchanged when it is
|
|
// not a JSON object, which the endpoint then refuses as it would have anyway.
|
|
func withModel(body []byte, name string) []byte {
|
|
var m map[string]json.RawMessage
|
|
if err := json.Unmarshal(body, &m); err != nil {
|
|
return body
|
|
}
|
|
encoded, err := json.Marshal(name)
|
|
if err != nil {
|
|
return body
|
|
}
|
|
m["model"] = encoded
|
|
out, err := json.Marshal(m)
|
|
if err != nil {
|
|
return body
|
|
}
|
|
return out
|
|
}
|
|
|
|
// withAdapter is body asking for model with a LoRA adapter applied, in
|
|
// llama.cpp's form: "lora": [{"id": index, "scale": scale}]. Adapters not
|
|
// listed stay at zero, the scale the engine loads each one at.
|
|
func withAdapter(body []byte, model string, index int, scale float64) []byte {
|
|
var m map[string]json.RawMessage
|
|
if err := json.Unmarshal(body, &m); err != nil {
|
|
return body
|
|
}
|
|
name, err1 := json.Marshal(model)
|
|
lora, err2 := json.Marshal([]map[string]any{{"id": index, "scale": scale}})
|
|
if err1 != nil || err2 != nil {
|
|
return body
|
|
}
|
|
m["model"], m["lora"] = name, lora
|
|
out, err := json.Marshal(m)
|
|
if err != nil {
|
|
return body
|
|
}
|
|
return out
|
|
}
|