Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8bb62a1c20 | ||
|
|
599b5542c6 | ||
|
|
abecc4b888 | ||
|
|
00800d3a3f | ||
|
|
b4426ba359 |
@@ -25,6 +25,14 @@ GATEWAY_DB_URL=postgres://lmr:lmr@localhost:5432/lmr?sslmode=disable
|
||||
# GATEWAY_DB_MIN_CONNS=2
|
||||
# GATEWAY_CATALOG_SYNC_INTERVAL=30m
|
||||
|
||||
# InfluxDB 3 statistics exporter. Off until GATEWAY_INFLUX_URL is set.
|
||||
# GATEWAY_INFLUX_URL=http://influxdb:8181
|
||||
# GATEWAY_INFLUX_TOKEN=
|
||||
# GATEWAY_INFLUX_DATABASE=lmr
|
||||
# GATEWAY_INFLUX_MEASUREMENT=lmr_request
|
||||
# GATEWAY_INFLUX_FLUSH_INTERVAL=10s
|
||||
# GATEWAY_INFLUX_BATCH_SIZE=500
|
||||
|
||||
# Host port to expose (maps to container :8080)
|
||||
# PORT=8080
|
||||
|
||||
|
||||
41
README.md
41
README.md
@@ -117,10 +117,38 @@ is loaded automatically; in production, set real environment variables.
|
||||
| `GATEWAY_OIDC_ADMIN_GROUP` | — | **yes** | group required for admin access |
|
||||
| `GATEWAY_SECURITY_ENCRYPTION_KEY` | — | **yes** | 64 hex chars (32 bytes) for AES-256-GCM |
|
||||
| `GATEWAY_CATALOG_SYNC_INTERVAL` | `30m` | no | how often provider model lists are refreshed |
|
||||
| `GATEWAY_INFLUX_URL` | — | no | InfluxDB 3 base URL, e.g. `http://influxdb:8181`. Empty = exporter off |
|
||||
| `GATEWAY_INFLUX_TOKEN` | — | no | InfluxDB 3 token, sent as a Bearer token |
|
||||
| `GATEWAY_INFLUX_DATABASE` | — | with URL | target database |
|
||||
| `GATEWAY_INFLUX_MEASUREMENT` | `lmr_request` | no | measurement name for request points |
|
||||
| `GATEWAY_INFLUX_FLUSH_INTERVAL` | `10s` | no | max time a point waits in the buffer |
|
||||
| `GATEWAY_INFLUX_BATCH_SIZE` | `500` | no | points per write request |
|
||||
|
||||
Config is validated on startup; missing required values or a wrong-length
|
||||
encryption key cause the server to exit immediately.
|
||||
|
||||
### InfluxDB 3 export
|
||||
|
||||
Set `GATEWAY_INFLUX_URL` and `GATEWAY_INFLUX_DATABASE` to export one point per
|
||||
proxied request. Points are batched in memory and written with the
|
||||
`/api/v3/write_lp` endpoint. A full buffer drops points and logs a warning, so
|
||||
the exporter never delays a response. Each point has these tags and fields:
|
||||
|
||||
| Kind | Name | Notes |
|
||||
|---|---|---|
|
||||
| tag | `provider`, `model`, `key_prefix` | group by provider, model address, or key. `key_prefix` is `(no key)` for the dashboard chat |
|
||||
| tag | `user_id`, `provider_id`, `provider_model_id` | UUIDs, set when known |
|
||||
| tag | `status_code`, `streamed` | HTTP status and `true`/`false` |
|
||||
| field | `requests`, `errors` | always `1` and `0` or `1`, so `sum()` gives counts |
|
||||
| field | `prompt_tokens`, `completion_tokens`, `cached_tokens`, `total_tokens` | integers |
|
||||
| field | `cost` | dollars, float |
|
||||
| field | `latency_ms`, `ttft_ms` | integers. `ttft_ms` is `0` when the response was not streamed |
|
||||
| field | `request_id`, `finish_reason`, `error` | strings, empty when not known |
|
||||
|
||||
Every field is written on every point. InfluxDB 3 creates a column on its first
|
||||
write, and a query that names a missing column fails, so optional values are
|
||||
written as `0` or an empty string.
|
||||
|
||||
Generate an encryption key:
|
||||
|
||||
```bash
|
||||
@@ -213,7 +241,7 @@ migrations/000002_virtual_keys.down.sql
|
||||
|
||||
| Name | Kind | Notes |
|
||||
|---|---|---|
|
||||
| `llamacpp` | local | OpenAI-compatible; routes through llama-swap |
|
||||
| `llamacpp` | local | OpenAI-compatible; works with a single llama-server, the llama.cpp router, or llama-swap; reads the running context size from `/props` during sync |
|
||||
| `ollama` | local | OpenAI-compatible |
|
||||
| `openrouter` | api key | OpenAI-compatible; HTTP-Referer/X-Title attribution; reports pricing; serves `/v1/images/generations` through chat completions |
|
||||
| `openai` | api key | reads `cached_tokens` from the usage response |
|
||||
@@ -232,8 +260,15 @@ timer and cached in memory.
|
||||
serving that model and picks one by the configured **routing strategy**
|
||||
(`round_robin` | `lowest_cost`, set from the dashboard Routing page).
|
||||
- **Admin overrides** — per-model `enabled`/`blocked`, parameter overrides,
|
||||
pricing, and rate limits live alongside the discovered data and survive syncs.
|
||||
Clearing an override in the dashboard falls back to the provider-reported price.
|
||||
pricing, context window, rate limits, and a backup model live alongside the
|
||||
discovered data and survive syncs. Clearing an override in the dashboard
|
||||
falls back to the provider-reported value.
|
||||
- **Backup models** — any model can name another model as its backup. When
|
||||
the model is offline, rate limited, or its upstream fails before the client
|
||||
has received anything, the gateway sends the same request to the backup and
|
||||
returns the answer under the requested model's name. The client sees no
|
||||
difference. The request log records each failed attempt, the address that
|
||||
served the request, and the backup's cost.
|
||||
- **Curated aliases** — the schema has `models` / `model_targets` tables for
|
||||
friendly names with priority/weight failover. No code uses them yet.
|
||||
|
||||
|
||||
@@ -34,3 +34,12 @@ GATEWAY_SECURITY_ENCRYPTION_KEY=
|
||||
# --- Catalog ---
|
||||
# How often to sync model lists from providers. Default: 30m
|
||||
GATEWAY_CATALOG_SYNC_INTERVAL=30m
|
||||
|
||||
# --- InfluxDB 3 statistics exporter (optional) ---
|
||||
# Leave GATEWAY_INFLUX_URL empty to disable the exporter.
|
||||
GATEWAY_INFLUX_URL=
|
||||
GATEWAY_INFLUX_TOKEN=
|
||||
GATEWAY_INFLUX_DATABASE=lmr
|
||||
GATEWAY_INFLUX_MEASUREMENT=lmr_request
|
||||
GATEWAY_INFLUX_FLUSH_INTERVAL=10s
|
||||
GATEWAY_INFLUX_BATCH_SIZE=500
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/config"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/crypto"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/db"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/influx"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/oidcauth"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/providermgr"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/ratelimit"
|
||||
@@ -147,8 +148,29 @@ func run() error {
|
||||
}
|
||||
}()
|
||||
|
||||
// The exporter sends one point per request to InfluxDB 3 when configured.
|
||||
var stats api.RequestStats
|
||||
var exporter *influx.Exporter
|
||||
if cfg.Influx.Enabled() {
|
||||
exporter, err = influx.New(influx.Config{
|
||||
URL: cfg.Influx.URL,
|
||||
Token: cfg.Influx.Token,
|
||||
Database: cfg.Influx.Database,
|
||||
Measurement: cfg.Influx.Measurement,
|
||||
FlushInterval: cfg.Influx.FlushInterval,
|
||||
BatchSize: cfg.Influx.BatchSize,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("init influx exporter: %w", err)
|
||||
}
|
||||
exporter.Start()
|
||||
stats = exporter
|
||||
slog.Info("influx exporter enabled", "url", cfg.Influx.URL, "database", cfg.Influx.Database)
|
||||
}
|
||||
|
||||
deps := api.Deps{
|
||||
Store: st,
|
||||
Stats: stats,
|
||||
Cache: cache,
|
||||
Manager: mgr,
|
||||
Cipher: cipher,
|
||||
@@ -193,6 +215,11 @@ func run() error {
|
||||
if err := srv.Shutdown(shutdownCtx); err != nil {
|
||||
return fmt.Errorf("graceful shutdown: %w", err)
|
||||
}
|
||||
if exporter != nil {
|
||||
if err := exporter.Close(shutdownCtx); err != nil {
|
||||
slog.Warn("influx exporter did not flush in time", "err", err)
|
||||
}
|
||||
}
|
||||
slog.Info("server stopped cleanly")
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -36,6 +36,7 @@ export interface AdminModel {
|
||||
enabled: boolean
|
||||
blocked: boolean
|
||||
context_length: number | null
|
||||
context_length_override: number | null
|
||||
last_seen_at: string | null
|
||||
disc_input_per_mtok: number | null
|
||||
disc_output_per_mtok: number | null
|
||||
@@ -44,6 +45,9 @@ export interface AdminModel {
|
||||
price_cached_per_mtok: number | null
|
||||
rpm_limit: number | null
|
||||
tpm_limit: number | null
|
||||
/** Backup model served when this one is down or fails. */
|
||||
fallback_model_id: string | null
|
||||
fallback_address: string | null
|
||||
}
|
||||
|
||||
export interface ChatMessage { role: string; content: unknown }
|
||||
@@ -54,6 +58,8 @@ export interface RequestLog {
|
||||
RequestID: string
|
||||
KeyPrefix: string
|
||||
ModelName: string
|
||||
/** Set when a backup model served the request instead of ModelName. */
|
||||
ServedModel: string
|
||||
ProviderName: string
|
||||
StatusCode: number
|
||||
Streamed: boolean
|
||||
@@ -144,6 +150,8 @@ export interface ModelOverride {
|
||||
cached_price_per_mtok?: number | null
|
||||
rpm_limit?: number | null
|
||||
tpm_limit?: number | null
|
||||
context_length?: number | null
|
||||
fallback_model_id?: string | null
|
||||
}
|
||||
|
||||
export type Settings = {
|
||||
|
||||
@@ -43,6 +43,7 @@ function splitThink(s: string): { reasoning: string; answer: string } {
|
||||
function Detail({ log, content, loading }: { log: RequestLog; content?: LogContent; loading: boolean }) {
|
||||
const meta = [
|
||||
['Request ID', log.RequestID || '—'],
|
||||
['Served by', log.ServedModel ? `${log.ServedModel} (backup)` : log.ModelName],
|
||||
['Streamed', log.Streamed ? 'yes' : 'no'],
|
||||
['Prompt tokens', log.PromptTokens.toLocaleString()],
|
||||
['Cached tokens', log.CachedTokens.toLocaleString()],
|
||||
@@ -204,7 +205,10 @@ export default function Logs() {
|
||||
onKeyDown={e => { if (e.key === 'Enter' || e.key === ' ') { e.preventDefault(); toggle(l.ID) } }}
|
||||
className="row row-hover cursor-pointer" style={{ gridTemplateColumns: COLS, borderBottom: 'none' }}>
|
||||
<span className="num text-[11px] text-ink-3" title={fmtDateTime(l.CreatedAt)}>{fmtTime(l.CreatedAt)}</span>
|
||||
<span className="mono text-[11.5px] text-accent-3 ellip pr-2" title={l.ModelName}>{l.ModelName}</span>
|
||||
<span className="mono text-[11.5px] text-accent-3 ellip pr-2 flex items-center gap-1.5 min-w-0" title={l.ServedModel ? `${l.ModelName} — served by backup ${l.ServedModel}` : l.ModelName}>
|
||||
<span className="ellip">{l.ModelName}</span>
|
||||
{l.ServedModel && <Badge tone="warn" title={`Served by backup ${l.ServedModel}`}>↪ backup</Badge>}
|
||||
</span>
|
||||
<span className="text-xs text-ink-2 ellip">{l.ProviderName}</span>
|
||||
<span className="mono text-[11px] text-ink-3 ellip">{l.KeyPrefix || '—'}</span>
|
||||
<span className="num text-[11px] text-ink-3" title={`${l.PromptTokens.toLocaleString()} in · ${l.CompletionTokens.toLocaleString()} out`}>{fmtTokens(l.TotalTokens)}</span>
|
||||
|
||||
@@ -8,7 +8,7 @@ import { Badge, Button, Card, Checkbox, CodeLine, CopyIconButton, Empty, ErrorNo
|
||||
type Filter = 'all' | 'live' | 'offline' | 'blocked' | 'overridden'
|
||||
|
||||
function hasOverride(m: AdminModel) {
|
||||
return m.price_input_per_mtok != null || m.price_output_per_mtok != null || m.price_cached_per_mtok != null || m.rpm_limit != null || m.tpm_limit != null
|
||||
return m.price_input_per_mtok != null || m.price_output_per_mtok != null || m.price_cached_per_mtok != null || m.rpm_limit != null || m.tpm_limit != null || m.context_length_override != null || m.fallback_model_id != null
|
||||
}
|
||||
|
||||
function price(v: number | null | undefined) {
|
||||
@@ -24,12 +24,14 @@ function parseField(raw: string, int = false): number | null | 'invalid' {
|
||||
return n
|
||||
}
|
||||
|
||||
function OverrideModal({ model, onClose, onSaved }: { model: AdminModel; onClose: () => void; onSaved: () => void }) {
|
||||
function OverrideModal({ model, models, onClose, onSaved }: { model: AdminModel; models: AdminModel[]; onClose: () => void; onSaved: () => void }) {
|
||||
const [inPrice, setInPrice] = useState(price(model.price_input_per_mtok))
|
||||
const [outPrice, setOutPrice] = useState(price(model.price_output_per_mtok))
|
||||
const [cachedPrice, setCachedPrice] = useState(price(model.price_cached_per_mtok))
|
||||
const [rpm, setRpm] = useState(price(model.rpm_limit))
|
||||
const [tpm, setTpm] = useState(price(model.tpm_limit))
|
||||
const [ctx, setCtx] = useState(price(model.context_length_override))
|
||||
const [fallback, setFallback] = useState(model.fallback_model_id ?? '')
|
||||
const [enabled, setEnabled] = useState(model.enabled)
|
||||
const [blocked, setBlocked] = useState(model.blocked)
|
||||
const [loading, setLoading] = useState(false)
|
||||
@@ -44,14 +46,19 @@ function OverrideModal({ model, onClose, onSaved }: { model: AdminModel; onClose
|
||||
cached_price_per_mtok: parseField(cachedPrice),
|
||||
rpm_limit: parseField(rpm, true),
|
||||
tpm_limit: parseField(tpm, true),
|
||||
context_length: parseField(ctx, true),
|
||||
}
|
||||
if (Object.values(fields).includes('invalid')) {
|
||||
setErr('Prices and limits must be empty or a number of 0 or more.')
|
||||
return
|
||||
}
|
||||
if (fields.context_length === 0) {
|
||||
setErr('The context window must be empty or at least 1 token.')
|
||||
return
|
||||
}
|
||||
setLoading(true)
|
||||
try {
|
||||
await api.models.update(model.id, { enabled, blocked, ...(fields as Partial<ModelOverride>) })
|
||||
await api.models.update(model.id, { enabled, blocked, fallback_model_id: fallback || null, ...(fields as Partial<ModelOverride>) })
|
||||
onSaved()
|
||||
onClose()
|
||||
} catch (e) {
|
||||
@@ -62,6 +69,7 @@ function OverrideModal({ model, onClose, onSaved }: { model: AdminModel; onClose
|
||||
}
|
||||
|
||||
const discIn = model.disc_input_per_mtok, discOut = model.disc_output_per_mtok
|
||||
const backups = models.filter(m => m.id !== model.id && !m.blocked && m.enabled)
|
||||
|
||||
return (
|
||||
<Modal title={model.display_name || model.slug} onClose={onClose}>
|
||||
@@ -79,7 +87,16 @@ function OverrideModal({ model, onClose, onSaved }: { model: AdminModel; onClose
|
||||
<div />
|
||||
<Field label="Requests / minute"><Input mono inputMode="numeric" value={rpm} onChange={e => setRpm(e.target.value)} placeholder="unlimited" /></Field>
|
||||
<Field label="Tokens / minute"><Input mono inputMode="numeric" value={tpm} onChange={e => setTpm(e.target.value)} placeholder="unlimited" /></Field>
|
||||
<Field label="Context window (tokens)" hint={model.context_length != null ? `Reported: ${fmtContext(model.context_length)}. Set this when the running instance accepts less.` : 'The provider reports no context window.'}>
|
||||
<Input mono inputMode="numeric" value={ctx} onChange={e => setCtx(e.target.value)} placeholder={model.context_length != null ? String(model.context_length) : 'unknown'} />
|
||||
</Field>
|
||||
</div>
|
||||
<Field label="Backup model" hint="Serves requests for this model when it is offline, rate limited, or its provider fails. The client sees the response under this model's name. The log shows the backup and its cost.">
|
||||
<Select value={fallback} onChange={e => setFallback(e.target.value)}>
|
||||
<option value="">None</option>
|
||||
{backups.map(m => <option key={m.id} value={m.id}>{m.provider_name}/{m.slug}{m.available ? '' : ' (offline)'}</option>)}
|
||||
</Select>
|
||||
</Field>
|
||||
<div className="flex gap-5">
|
||||
<Checkbox label="Enabled" checked={enabled} onChange={setEnabled} />
|
||||
<Checkbox label="Blocked" tone="bad" checked={blocked} onChange={setBlocked} />
|
||||
@@ -204,8 +221,13 @@ export default function Models() {
|
||||
<span className="mono text-[12px] text-accent-3 ellip" title={address}>{address}</span>
|
||||
<CopyIconButton text={address} title="Copy model address" />
|
||||
</div>
|
||||
<span className="text-xs text-ink-2 ellip" title={m.display_name}>{m.display_name || <span className="text-ink-4">—</span>}</span>
|
||||
<span className="num text-[11px] text-ink-3">{m.context_length ? fmtContext(m.context_length).replace(' ctx', '') : '—'}</span>
|
||||
<span className="text-xs text-ink-2 ellip flex items-center gap-1.5 min-w-0" title={m.display_name}>
|
||||
<span className="ellip">{m.display_name || <span className="text-ink-4">—</span>}</span>
|
||||
{m.fallback_address && <Badge tone="neutral" title={`Backup: ${m.fallback_address}`}>↪ backup</Badge>}
|
||||
</span>
|
||||
<span className="num text-[11px] text-ink-3" title={m.context_length_override != null && m.context_length ? `Override. Provider reports ${fmtContext(m.context_length)}` : undefined}>
|
||||
{(m.context_length_override ?? m.context_length) ? fmtContext(m.context_length_override ?? m.context_length!).replace(' ctx', '') : '—'}{m.context_length_override != null && <span className="text-warn">*</span>}
|
||||
</span>
|
||||
<span className="num text-[11px] text-ink-2 flex items-center gap-2">
|
||||
{fmtPrice(inPrice)} / {fmtPrice(outPrice)}
|
||||
{hasOverride(m) && <Badge tone="warn" title="Admin override in effect">override</Badge>}
|
||||
@@ -221,7 +243,7 @@ export default function Models() {
|
||||
</div>
|
||||
</Card>
|
||||
|
||||
{selected && <OverrideModal model={selected} onClose={() => setSelected(null)} onSaved={() => { toast('Model updated', 'good'); load() }} />}
|
||||
{selected && <OverrideModal model={selected} models={models} onClose={() => setSelected(null)} onSaved={() => { toast('Model updated', 'good'); load() }} />}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
@@ -15,6 +16,7 @@ import (
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/catalog"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/providers"
|
||||
@@ -428,6 +430,22 @@ func (a *API) adminUpdateModel(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
if o.ContextLength, err = patchField[int](fields, "context_length"); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
if o.ContextLength.Value != nil && *o.ContextLength.Value < 1 {
|
||||
writeError(w, http.StatusBadRequest, "context_length must be 1 or more")
|
||||
return
|
||||
}
|
||||
if o.FallbackModelID, err = patchField[uuid.UUID](fields, "fallback_model_id"); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
if o.FallbackModelID.Value != nil && *o.FallbackModelID.Value == id {
|
||||
writeError(w, http.StatusBadRequest, "a model cannot be its own backup")
|
||||
return
|
||||
}
|
||||
for _, p := range []store.Patch[float64]{o.InputPrice, o.OutputPrice, o.CachedPrice} {
|
||||
if p.Value != nil && *p.Value < 0 {
|
||||
writeError(w, http.StatusBadRequest, "prices must be 0 or more")
|
||||
@@ -442,6 +460,11 @@ func (a *API) adminUpdateModel(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
if err := a.Store.UpdateModelOverride(r.Context(), id, o); err != nil {
|
||||
var pgErr *pgconn.PgError
|
||||
if errors.As(err, &pgErr) && pgErr.Code == "23503" /* foreign_key_violation */ {
|
||||
writeError(w, http.StatusBadRequest, "backup model not found")
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
@@ -53,11 +53,19 @@ type RequestLogSink interface {
|
||||
TouchVirtualKey(ctx context.Context, id uuid.UUID) error
|
||||
}
|
||||
|
||||
// RequestStats receives a copy of every request log row for export to a
|
||||
// metrics backend. Record must not block. *influx.Exporter satisfies it.
|
||||
type RequestStats interface {
|
||||
Record(e store.RequestLog)
|
||||
}
|
||||
|
||||
// Deps are the dependencies the API handlers need.
|
||||
type Deps struct {
|
||||
Store *store.Store
|
||||
// Logs receives request log rows. Defaults to Store when nil.
|
||||
Logs RequestLogSink
|
||||
Logs RequestLogSink
|
||||
// Stats receives request log rows for the statistics exporter. May be nil.
|
||||
Stats RequestStats
|
||||
Cache *catalog.Cache
|
||||
Manager *providermgr.Manager
|
||||
Cipher *crypto.Cipher
|
||||
|
||||
@@ -43,101 +43,123 @@ func (a *API) chatCompletions(w http.ResponseWriter, r *http.Request) {
|
||||
key := keyFrom(r.Context())
|
||||
address := req.Model
|
||||
|
||||
entry, err := a.Cache.Resolve(address)
|
||||
route, err := a.Cache.Route(address)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusNotFound, err.Error())
|
||||
return
|
||||
}
|
||||
if !allowed(key, entry.Address) {
|
||||
if !allowed(key, route.Primary) {
|
||||
writeError(w, http.StatusForbidden, fmt.Sprintf("model %q is not allowed for this key", address))
|
||||
return
|
||||
}
|
||||
if !a.Limiter.AllowRequest(entry.Address, ptrInt(entry.RPMLimit)) {
|
||||
writeError(w, http.StatusTooManyRequests, "rate limit exceeded (RPM)")
|
||||
return
|
||||
}
|
||||
if !a.Limiter.AllowTokens(entry.Address, ptrInt(entry.TPMLimit)) {
|
||||
writeError(w, http.StatusTooManyRequests, "rate limit exceeded (TPM)")
|
||||
return
|
||||
}
|
||||
|
||||
prov, ok := a.Manager.Get(entry.ProviderName)
|
||||
if !ok {
|
||||
writeError(w, http.StatusServiceUnavailable, "provider unavailable")
|
||||
return
|
||||
}
|
||||
cc, ok := prov.(providers.ChatCompleter)
|
||||
if !ok {
|
||||
writeError(w, http.StatusNotImplemented, "provider does not support chat completions")
|
||||
return
|
||||
}
|
||||
|
||||
// Snapshot the messages before we overwrite req.Model below.
|
||||
messages := marshalMessages(req.Payload["messages"])
|
||||
|
||||
req.Model = entry.Slug
|
||||
applyOverrides(req.Payload, entry.ParamOverrides)
|
||||
|
||||
start := time.Now()
|
||||
if req.Stream {
|
||||
a.streamChat(w, r, cc, req, entry, key, address, messages, start)
|
||||
return
|
||||
var last *failure
|
||||
for i, entry := range route.Targets {
|
||||
if i > 0 && r.Context().Err() != nil {
|
||||
break // the client is gone; a backup cannot help
|
||||
}
|
||||
prov, fail := a.admit(entry, 1)
|
||||
if fail != nil {
|
||||
last = fail
|
||||
continue
|
||||
}
|
||||
cc, ok := prov.(providers.ChatCompleter)
|
||||
if !ok {
|
||||
last = &failure{status: http.StatusNotImplemented, msg: "provider does not support chat completions"}
|
||||
continue
|
||||
}
|
||||
attempt := &providers.ChatRequest{Model: entry.Slug, Stream: req.Stream, Payload: targetPayload(req.Payload, entry)}
|
||||
next := nextTarget(route, i)
|
||||
start := time.Now()
|
||||
if req.Stream {
|
||||
fail = a.streamChat(w, r, cc, attempt, entry, route, next, key, messages, start)
|
||||
} else {
|
||||
fail = a.completeChat(w, r, cc, attempt, entry, route, next, key, messages, start)
|
||||
}
|
||||
if fail == nil {
|
||||
return
|
||||
}
|
||||
last = fail
|
||||
}
|
||||
last.write(w)
|
||||
}
|
||||
|
||||
// completeChat serves a non-streaming completion from one target. It returns
|
||||
// nil once the client has an answer, or the failure when the request should
|
||||
// move on to next. An upstream error with no backup behind it is relayed.
|
||||
func (a *API) completeChat(w http.ResponseWriter, r *http.Request, cc providers.ChatCompleter, req *providers.ChatRequest, entry *catalog.Entry, route *catalog.Route, next *catalog.Entry, key *store.VirtualKey, messages json.RawMessage, start time.Time) *failure {
|
||||
resp, err := cc.ChatCompletion(r.Context(), req)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, err.Error())
|
||||
a.logRequest(r, key, entry, address, http.StatusBadGateway, false, providers.Usage{}, err.Error(), start, nil, messages, "")
|
||||
return
|
||||
a.logRequest(r, key, entry, route, http.StatusBadGateway, false, providers.Usage{}, err.Error()+switchNote(next), start, nil, messages, "")
|
||||
return &failure{status: http.StatusBadGateway, msg: err.Error()}
|
||||
}
|
||||
a.Limiter.RecordTokens(entry.Address, resp.Usage.TotalTokens)
|
||||
respBody := resp.Body
|
||||
errMsg := ""
|
||||
if resp.StatusCode >= 400 {
|
||||
errMsg = upstreamErrorMessage(resp.Body)
|
||||
} else {
|
||||
respBody = injectCreated(resp.Body)
|
||||
errMsg := upstreamErrorMessage(resp.Body)
|
||||
if next != nil && retryable(resp.StatusCode) {
|
||||
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, errMsg+switchNote(next), start, nil, messages, "")
|
||||
return &failure{status: resp.StatusCode, body: resp.Body}
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
_, _ = w.Write(resp.Body)
|
||||
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, errMsg, start, nil, messages, "")
|
||||
return nil
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
_, _ = w.Write(respBody)
|
||||
_, _ = w.Write(clientBody(resp.Body, time.Now().Unix(), reportedModel(route, entry)))
|
||||
|
||||
completion := extractCompletion(resp.Body)
|
||||
a.logRequest(r, key, entry, address, resp.StatusCode, false, resp.Usage, errMsg, start, nil, messages, completion)
|
||||
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, "", start, nil, messages, extractCompletion(resp.Body))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *API) streamChat(w http.ResponseWriter, r *http.Request, cc providers.ChatCompleter, req *providers.ChatRequest, entry *catalog.Entry, key *store.VirtualKey, address string, messages json.RawMessage, start time.Time) {
|
||||
// streamChat serves a streaming completion from one target. The first chunk is
|
||||
// read before the status line goes out, so an upstream that fails at once can
|
||||
// still hand the request to a backup. After the first byte reaches the
|
||||
// client, a failure can only be recorded in the log.
|
||||
func (a *API) streamChat(w http.ResponseWriter, r *http.Request, cc providers.ChatCompleter, req *providers.ChatRequest, entry *catalog.Entry, route *catalog.Route, next *catalog.Entry, key *store.VirtualKey, messages json.RawMessage, start time.Time) *failure {
|
||||
stream, err := cc.ChatCompletionStream(r.Context(), req)
|
||||
if err != nil {
|
||||
var apiErr *providers.APIError
|
||||
if errors.As(err, &apiErr) {
|
||||
errMsg := upstreamErrorMessage(apiErr.Body)
|
||||
if next != nil && retryable(apiErr.StatusCode) {
|
||||
a.logRequest(r, key, entry, route, apiErr.StatusCode, true, providers.Usage{}, errMsg+switchNote(next), start, nil, messages, "")
|
||||
return &failure{status: apiErr.StatusCode, body: apiErr.Body}
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(apiErr.StatusCode)
|
||||
_, _ = w.Write(apiErr.Body)
|
||||
a.logRequest(r, key, entry, address, apiErr.StatusCode, true, providers.Usage{}, upstreamErrorMessage(apiErr.Body), start, nil, messages, "")
|
||||
return
|
||||
a.logRequest(r, key, entry, route, apiErr.StatusCode, true, providers.Usage{}, errMsg, start, nil, messages, "")
|
||||
return nil
|
||||
}
|
||||
writeError(w, http.StatusBadGateway, err.Error())
|
||||
a.logRequest(r, key, entry, address, http.StatusBadGateway, true, providers.Usage{}, err.Error(), start, nil, messages, "")
|
||||
return
|
||||
a.logRequest(r, key, entry, route, http.StatusBadGateway, true, providers.Usage{}, err.Error()+switchNote(next), start, nil, messages, "")
|
||||
return &failure{status: http.StatusBadGateway, msg: err.Error()}
|
||||
}
|
||||
defer stream.Close()
|
||||
|
||||
ev, err := stream.Recv()
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
errMsg := "stream interrupted: " + err.Error()
|
||||
a.logRequest(r, key, entry, route, http.StatusBadGateway, true, providers.Usage{}, errMsg+switchNote(next), start, nil, messages, "")
|
||||
return &failure{status: http.StatusBadGateway, msg: errMsg}
|
||||
}
|
||||
|
||||
flusher, _ := w.(http.Flusher)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
|
||||
model := reportedModel(route, entry)
|
||||
var ttft *int
|
||||
var reasoning, content strings.Builder
|
||||
created := time.Now().Unix()
|
||||
errMsg := ""
|
||||
for {
|
||||
ev, err := stream.Recv()
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
for !errors.Is(err, io.EOF) {
|
||||
if err != nil {
|
||||
// The status line is already sent, so the error can only be
|
||||
// recorded in the log row and signalled by closing the stream.
|
||||
@@ -154,13 +176,14 @@ func (a *API) streamChat(w http.ResponseWriter, r *http.Request, cc providers.Ch
|
||||
}
|
||||
// Accumulate reasoning and content deltas separately so the recorded
|
||||
// completion wraps reasoning in think tags, matching extractCompletion.
|
||||
r, c := extractDelta(ev.Data)
|
||||
reasoning.WriteString(r)
|
||||
rs, c := extractDelta(ev.Data)
|
||||
reasoning.WriteString(rs)
|
||||
content.WriteString(c)
|
||||
_, _ = fmt.Fprintf(w, "data: %s\n\n", injectCreatedChunk(ev.Data, created))
|
||||
_, _ = fmt.Fprintf(w, "data: %s\n\n", clientBody(ev.Data, created, model))
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
ev, err = stream.Recv()
|
||||
}
|
||||
_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
@@ -173,21 +196,26 @@ func (a *API) streamChat(w http.ResponseWriter, r *http.Request, cc providers.Ch
|
||||
}
|
||||
usage := stream.Usage()
|
||||
a.Limiter.RecordTokens(entry.Address, usage.TotalTokens)
|
||||
a.logRequest(r, key, entry, address, http.StatusOK, true, usage, errMsg, start, ttft, messages, completion)
|
||||
a.logRequest(r, key, entry, route, http.StatusOK, true, usage, errMsg, start, ttft, messages, completion)
|
||||
return nil
|
||||
}
|
||||
|
||||
func applyOverrides(payload, overrides map[string]any) {
|
||||
maps.Copy(payload, overrides)
|
||||
}
|
||||
|
||||
func (a *API) logRequest(r *http.Request, key *store.VirtualKey, entry *catalog.Entry, address string, status int, streamed bool, usage providers.Usage, errMsg string, start time.Time, ttft *int, messages json.RawMessage, completion string) {
|
||||
// logRequest records one attempt. ModelName is the address the client sent;
|
||||
// when a backup served the request, ServedModel names it and Cost uses its
|
||||
// prices, so the switch is visible in the log but not to the client.
|
||||
func (a *API) logRequest(r *http.Request, key *store.VirtualKey, entry *catalog.Entry, route *catalog.Route, status int, streamed bool, usage providers.Usage, errMsg string, start time.Time, ttft *int, messages json.RawMessage, completion string) {
|
||||
entryID := entry.ID
|
||||
providerID := entry.ProviderID
|
||||
served := ""
|
||||
if route.IsFallback(entry) {
|
||||
served = entry.Address
|
||||
}
|
||||
rec := store.RequestLog{
|
||||
ProviderModelID: &entryID,
|
||||
ProviderID: &providerID,
|
||||
ProviderName: entry.ProviderName,
|
||||
ModelName: address,
|
||||
ModelName: route.Address,
|
||||
ServedModel: served,
|
||||
StatusCode: status,
|
||||
Streamed: streamed,
|
||||
Error: errMsg,
|
||||
@@ -210,6 +238,10 @@ func (a *API) logRequest(r *http.Request, key *store.VirtualKey, entry *catalog.
|
||||
}
|
||||
}
|
||||
|
||||
if a.Stats != nil {
|
||||
a.Stats.Record(rec)
|
||||
}
|
||||
|
||||
// The log write must outlive the request, so it gets its own context.
|
||||
go func() { //nolint:gosec // G118: intentional, see above
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
@@ -377,34 +409,3 @@ func clientIP(r *http.Request) string {
|
||||
}
|
||||
return host
|
||||
}
|
||||
|
||||
// injectCreated ensures a non-streaming response has a "created" timestamp.
|
||||
// If the upstream already provided one, it is preserved.
|
||||
func injectCreated(body []byte) []byte {
|
||||
var obj map[string]any
|
||||
if json.Unmarshal(body, &obj) == nil {
|
||||
if _, ok := obj["created"]; ok {
|
||||
return body // upstream already set it
|
||||
}
|
||||
obj["created"] = time.Now().Unix()
|
||||
if out, err := json.Marshal(obj); err == nil {
|
||||
return out
|
||||
}
|
||||
}
|
||||
return body // fall back to raw body on any error
|
||||
}
|
||||
|
||||
// injectCreatedChunk ensures a streaming SSE chunk has a "created" timestamp.
|
||||
func injectCreatedChunk(data []byte, created int64) []byte {
|
||||
var obj map[string]any
|
||||
if json.Unmarshal(data, &obj) == nil {
|
||||
if _, ok := obj["created"]; ok {
|
||||
return data // upstream already set it
|
||||
}
|
||||
obj["created"] = created
|
||||
if out, err := json.Marshal(obj); err == nil {
|
||||
return out
|
||||
}
|
||||
}
|
||||
return data // fall back to raw chunk on any error
|
||||
}
|
||||
|
||||
@@ -28,57 +28,62 @@ func (a *API) embeddings(w http.ResponseWriter, r *http.Request) {
|
||||
key := keyFrom(r.Context())
|
||||
address := req.Model
|
||||
|
||||
entry, err := a.Cache.Resolve(address)
|
||||
route, err := a.Cache.Route(address)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusNotFound, err.Error())
|
||||
return
|
||||
}
|
||||
if !allowed(key, entry.Address) {
|
||||
if !allowed(key, route.Primary) {
|
||||
writeError(w, http.StatusForbidden, fmt.Sprintf("model %q is not allowed for this key", address))
|
||||
return
|
||||
}
|
||||
if !a.Limiter.AllowRequest(entry.Address, ptrInt(entry.RPMLimit)) {
|
||||
writeError(w, http.StatusTooManyRequests, "rate limit exceeded (RPM)")
|
||||
|
||||
var last *failure
|
||||
for i, entry := range route.Targets {
|
||||
if i > 0 && r.Context().Err() != nil {
|
||||
break
|
||||
}
|
||||
prov, fail := a.admit(entry, 1)
|
||||
if fail != nil {
|
||||
last = fail
|
||||
continue
|
||||
}
|
||||
emb, ok := prov.(providers.Embedder)
|
||||
if !ok {
|
||||
last = &failure{status: http.StatusNotImplemented, msg: "provider does not support embeddings"}
|
||||
continue
|
||||
}
|
||||
attempt := &providers.EmbeddingRequest{Model: entry.Slug, Payload: targetPayload(req.Payload, entry)}
|
||||
next := nextTarget(route, i)
|
||||
start := time.Now()
|
||||
resp, err := emb.Embeddings(r.Context(), attempt)
|
||||
if err != nil {
|
||||
a.logRequest(r, key, entry, route, http.StatusBadGateway, false, providers.Usage{}, err.Error()+switchNote(next), start, nil, nil, "")
|
||||
last = &failure{status: http.StatusBadGateway, msg: err.Error()}
|
||||
continue
|
||||
}
|
||||
a.Limiter.RecordTokens(entry.Address, resp.Usage.TotalTokens)
|
||||
|
||||
errMsg := ""
|
||||
body := resp.Body
|
||||
if resp.StatusCode >= 400 {
|
||||
errMsg = upstreamErrorMessage(resp.Body)
|
||||
if next != nil && retryable(resp.StatusCode) {
|
||||
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, errMsg+switchNote(next), start, nil, nil, "")
|
||||
last = &failure{status: resp.StatusCode, body: resp.Body}
|
||||
continue
|
||||
}
|
||||
} else {
|
||||
body = withModel(resp.Body, reportedModel(route, entry))
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
_, _ = w.Write(body)
|
||||
|
||||
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, errMsg, start, nil, nil, "")
|
||||
return
|
||||
}
|
||||
if !a.Limiter.AllowTokens(entry.Address, ptrInt(entry.TPMLimit)) {
|
||||
writeError(w, http.StatusTooManyRequests, "rate limit exceeded (TPM)")
|
||||
return
|
||||
}
|
||||
|
||||
prov, ok := a.Manager.Get(entry.ProviderName)
|
||||
if !ok {
|
||||
writeError(w, http.StatusServiceUnavailable, "provider unavailable")
|
||||
return
|
||||
}
|
||||
emb, ok := prov.(providers.Embedder)
|
||||
if !ok {
|
||||
writeError(w, http.StatusNotImplemented, "provider does not support embeddings")
|
||||
return
|
||||
}
|
||||
|
||||
req.Model = entry.Slug
|
||||
applyOverrides(req.Payload, entry.ParamOverrides)
|
||||
|
||||
start := time.Now()
|
||||
resp, err := emb.Embeddings(r.Context(), req)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, err.Error())
|
||||
a.logRequest(r, key, entry, address, http.StatusBadGateway, false, providers.Usage{}, err.Error(), start, nil, nil, "")
|
||||
return
|
||||
}
|
||||
|
||||
a.Limiter.RecordTokens(entry.Address, resp.Usage.TotalTokens)
|
||||
|
||||
errMsg := ""
|
||||
if resp.StatusCode >= 400 {
|
||||
errMsg = upstreamErrorMessage(resp.Body)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
_, _ = w.Write(resp.Body)
|
||||
|
||||
a.logRequest(r, key, entry, address, resp.StatusCode, false, resp.Usage, errMsg, start, nil, nil, "")
|
||||
last.write(w)
|
||||
}
|
||||
|
||||
// ptrInt dereferences an *int, returning 0 (no limit) when nil.
|
||||
|
||||
@@ -47,69 +47,73 @@ func (a *API) imageGenerations(w http.ResponseWriter, r *http.Request) {
|
||||
key := keyFrom(r.Context())
|
||||
address := req.Model
|
||||
|
||||
entry, err := a.Cache.Resolve(address)
|
||||
route, err := a.Cache.Route(address)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusNotFound, err.Error())
|
||||
return
|
||||
}
|
||||
if !allowed(key, entry.Address) {
|
||||
if !allowed(key, route.Primary) {
|
||||
writeError(w, http.StatusForbidden, fmt.Sprintf("model %q is not allowed for this key", address))
|
||||
return
|
||||
}
|
||||
if !slices.Contains(entry.Capabilities, string(providers.CapabilityImages)) {
|
||||
writeError(w, http.StatusBadRequest, fmt.Sprintf("model %q does not generate images", address))
|
||||
return
|
||||
}
|
||||
// Each requested image is one upstream call, so charge all of them
|
||||
// against the RPM limit up front.
|
||||
if !a.Limiter.AllowRequests(entry.Address, ptrInt(entry.RPMLimit), max(req.N, 1)) {
|
||||
writeError(w, http.StatusTooManyRequests, "rate limit exceeded (RPM)")
|
||||
return
|
||||
}
|
||||
if !a.Limiter.AllowTokens(entry.Address, ptrInt(entry.TPMLimit)) {
|
||||
writeError(w, http.StatusTooManyRequests, "rate limit exceeded (TPM)")
|
||||
return
|
||||
}
|
||||
|
||||
prov, ok := a.Manager.Get(entry.ProviderName)
|
||||
if !ok {
|
||||
writeError(w, http.StatusServiceUnavailable, "provider unavailable")
|
||||
return
|
||||
}
|
||||
gen, ok := prov.(providers.ImageGenerator)
|
||||
if !ok {
|
||||
writeError(w, http.StatusNotImplemented, "provider does not support image generation")
|
||||
return
|
||||
}
|
||||
|
||||
// Log the prompt as a one-message conversation so it shows in the request
|
||||
// log like a chat turn.
|
||||
messages, _ := json.Marshal([]map[string]any{{"role": "user", "content": req.Prompt}})
|
||||
|
||||
req.Model = entry.Slug
|
||||
applyOverrides(req.Payload, entry.ParamOverrides)
|
||||
var last *failure
|
||||
for i, entry := range route.Targets {
|
||||
if i > 0 && r.Context().Err() != nil {
|
||||
break
|
||||
}
|
||||
if !slices.Contains(entry.Capabilities, string(providers.CapabilityImages)) {
|
||||
last = &failure{status: http.StatusBadRequest, msg: fmt.Sprintf("model %q does not generate images", address)}
|
||||
continue
|
||||
}
|
||||
// Each requested image is one upstream call, so charge all of them
|
||||
// against the RPM limit up front.
|
||||
prov, fail := a.admit(entry, max(req.N, 1))
|
||||
if fail != nil {
|
||||
last = fail
|
||||
continue
|
||||
}
|
||||
gen, ok := prov.(providers.ImageGenerator)
|
||||
if !ok {
|
||||
last = &failure{status: http.StatusNotImplemented, msg: "provider does not support image generation"}
|
||||
continue
|
||||
}
|
||||
attempt := *req
|
||||
attempt.Model = entry.Slug
|
||||
attempt.Payload = targetPayload(req.Payload, entry)
|
||||
next := nextTarget(route, i)
|
||||
start := time.Now()
|
||||
resp, err := gen.ImageGeneration(r.Context(), &attempt)
|
||||
if err != nil {
|
||||
a.logRequest(r, key, entry, route, http.StatusBadGateway, false, providers.Usage{}, err.Error()+switchNote(next), start, nil, messages, "")
|
||||
last = &failure{status: http.StatusBadGateway, msg: err.Error()}
|
||||
continue
|
||||
}
|
||||
a.Limiter.RecordTokens(entry.Address, resp.Usage.TotalTokens)
|
||||
|
||||
start := time.Now()
|
||||
resp, err := gen.ImageGeneration(r.Context(), req)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, err.Error())
|
||||
a.logRequest(r, key, entry, address, http.StatusBadGateway, false, providers.Usage{}, err.Error(), start, nil, messages, "")
|
||||
errMsg := ""
|
||||
completion := ""
|
||||
if resp.StatusCode >= 400 {
|
||||
errMsg = upstreamErrorMessage(resp.Body)
|
||||
if next != nil && retryable(resp.StatusCode) {
|
||||
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, errMsg+switchNote(next), start, nil, messages, "")
|
||||
last = &failure{status: resp.StatusCode, body: resp.Body}
|
||||
continue
|
||||
}
|
||||
} else {
|
||||
completion = imageCompletionSummary(resp.Body)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
_, _ = w.Write(resp.Body)
|
||||
|
||||
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, errMsg, start, nil, messages, completion)
|
||||
return
|
||||
}
|
||||
a.Limiter.RecordTokens(entry.Address, resp.Usage.TotalTokens)
|
||||
|
||||
errMsg := ""
|
||||
completion := ""
|
||||
if resp.StatusCode >= 400 {
|
||||
errMsg = upstreamErrorMessage(resp.Body)
|
||||
} else {
|
||||
completion = imageCompletionSummary(resp.Body)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
_, _ = w.Write(resp.Body)
|
||||
|
||||
a.logRequest(r, key, entry, address, resp.StatusCode, false, resp.Usage, errMsg, start, nil, messages, completion)
|
||||
last.write(w)
|
||||
}
|
||||
|
||||
// imageCompletionSummary describes an images response for the content log
|
||||
|
||||
@@ -76,17 +76,27 @@ func (f *fakeLogs) TouchVirtualKey(_ context.Context, id uuid.UUID) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// wait blocks until one request has been fully logged.
|
||||
func (f *fakeLogs) wait(t *testing.T) store.RequestLog {
|
||||
// waitN blocks until n requests have been logged and returns the rows in the
|
||||
// order they were written.
|
||||
func (f *fakeLogs) waitN(t *testing.T, n int) []store.RequestLog {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-f.done:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("request log was not written")
|
||||
for range n {
|
||||
select {
|
||||
case <-f.done:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("request log was not written")
|
||||
}
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return f.rows[len(f.rows)-1]
|
||||
return append([]store.RequestLog(nil), f.rows...)
|
||||
}
|
||||
|
||||
// wait blocks until one request has been fully logged.
|
||||
func (f *fakeLogs) wait(t *testing.T) store.RequestLog {
|
||||
t.Helper()
|
||||
rows := f.waitN(t, 1)
|
||||
return rows[len(rows)-1]
|
||||
}
|
||||
|
||||
// upstream is a fake OpenRouter that records the last body it received.
|
||||
@@ -473,3 +483,187 @@ func TestAllowedPatterns(t *testing.T) {
|
||||
t.Error("allow-all key must pass")
|
||||
}
|
||||
}
|
||||
|
||||
// failPrimary answers gpt-4o with status and lets every other model through
|
||||
// to chatReply.
|
||||
func failPrimary(status int) func(http.ResponseWriter, *http.Request, map[string]any) {
|
||||
return func(w http.ResponseWriter, r *http.Request, body map[string]any) {
|
||||
if body["model"] == "openai/gpt-4o" {
|
||||
w.WriteHeader(status)
|
||||
_, _ = io.WriteString(w, `{"error":{"message":"primary down"}}`)
|
||||
return
|
||||
}
|
||||
chatReply(w, r, body)
|
||||
}
|
||||
}
|
||||
|
||||
// withBackup makes the image model the backup of gpt-4o, at a different price.
|
||||
func withBackup(rows []store.CatalogRow) {
|
||||
id := rows[0].ID
|
||||
rows[1].FallbackModelID = &id
|
||||
rows[0].InputPricePerMTok = 1
|
||||
rows[0].OutputPricePerMTok = 2
|
||||
}
|
||||
|
||||
// rowsByStatus indexes log rows by status code.
|
||||
func rowsByStatus(rows []store.RequestLog) map[int]store.RequestLog {
|
||||
out := map[int]store.RequestLog{}
|
||||
for _, r := range rows {
|
||||
out[r.StatusCode] = r
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestChatCompletionsFallsBackOnUpstreamError(t *testing.T) {
|
||||
env := newTestEnv(t, failPrimary(http.StatusServiceUnavailable), withBackup)
|
||||
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
||||
`{"model":"or/openai/gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rec.Code, rec.Body)
|
||||
}
|
||||
// The client sees the model it asked for.
|
||||
var out map[string]any
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
if out["model"] != "openai/gpt-4o" {
|
||||
t.Errorf("reported model = %v, want the requested one", out["model"])
|
||||
}
|
||||
if calls, body := env.upstream.last(); calls != 2 || body["model"] != "google/gemini-2.5-flash-image" {
|
||||
t.Errorf("upstream calls=%d last model=%v", calls, body["model"])
|
||||
}
|
||||
|
||||
// The log sees both attempts and the backup's cost.
|
||||
rows := rowsByStatus(env.logs.waitN(t, 2))
|
||||
failed, ok := rows[http.StatusServiceUnavailable]
|
||||
if !ok || !strings.Contains(failed.Error, "primary down") || !strings.Contains(failed.Error, "switched to backup or/google/gemini-2.5-flash-image") {
|
||||
t.Errorf("failed attempt row = %+v", failed)
|
||||
}
|
||||
if failed.ServedModel != "" || failed.ModelName != "or/openai/gpt-4o" {
|
||||
t.Errorf("failed attempt must be logged against the requested model: %+v", failed)
|
||||
}
|
||||
served := rows[http.StatusOK]
|
||||
if served.ModelName != "or/openai/gpt-4o" || served.ServedModel != "or/google/gemini-2.5-flash-image" {
|
||||
t.Errorf("served row = %+v", served)
|
||||
}
|
||||
want := 100.0/1e6*1 + 50.0/1e6*2
|
||||
if served.Cost < want-1e-12 || served.Cost > want+1e-12 {
|
||||
t.Errorf("cost = %v, want the backup's price %v", served.Cost, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatCompletionsStreamFallsBack(t *testing.T) {
|
||||
env := newTestEnv(t, failPrimary(http.StatusBadGateway), withBackup)
|
||||
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
||||
`{"model":"or/openai/gpt-4o","messages":[],"stream":true}`)
|
||||
if rec.Code != http.StatusOK || rec.Header().Get("Content-Type") != "text/event-stream" {
|
||||
t.Fatalf("status=%d type=%q body=%s", rec.Code, rec.Header().Get("Content-Type"), rec.Body)
|
||||
}
|
||||
body := rec.Body.String()
|
||||
if strings.Count(body, "data: ") != 4 || strings.Count(body, `"model":"openai/gpt-4o"`) != 3 {
|
||||
t.Errorf("stream body = %q", body)
|
||||
}
|
||||
rows := rowsByStatus(env.logs.waitN(t, 2))
|
||||
if rows[http.StatusOK].ServedModel != "or/google/gemini-2.5-flash-image" || !rows[http.StatusOK].Streamed {
|
||||
t.Errorf("served row = %+v", rows[http.StatusOK])
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatCompletionsUnavailableModelServedByBackup(t *testing.T) {
|
||||
env := newTestEnv(t, chatReply, func(rows []store.CatalogRow) {
|
||||
withBackup(rows)
|
||||
rows[1].Available = false
|
||||
})
|
||||
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
||||
`{"model":"or/openai/gpt-4o","messages":[]}`)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rec.Code, rec.Body)
|
||||
}
|
||||
var out map[string]any
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
if out["model"] != "openai/gpt-4o" {
|
||||
t.Errorf("reported model = %v", out["model"])
|
||||
}
|
||||
if calls, _ := env.upstream.last(); calls != 1 {
|
||||
t.Errorf("upstream calls = %d, want 1", calls)
|
||||
}
|
||||
row := env.logs.wait(t)
|
||||
if row.ModelName != "or/openai/gpt-4o" || row.ServedModel != "or/google/gemini-2.5-flash-image" {
|
||||
t.Errorf("log row = %+v", row)
|
||||
}
|
||||
|
||||
// The down model stays listed, because clients can still call it.
|
||||
list := call(env.api.listModels, allowAllKey(), http.MethodGet, "")
|
||||
if !strings.Contains(list.Body.String(), `"or/openai/gpt-4o"`) {
|
||||
t.Errorf("models list = %s", list.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatCompletionsRateLimitOverflowsToBackup(t *testing.T) {
|
||||
env := newTestEnv(t, chatReply, func(rows []store.CatalogRow) {
|
||||
withBackup(rows)
|
||||
one := 1
|
||||
rows[1].RPMLimit = &one
|
||||
})
|
||||
body := `{"model":"or/openai/gpt-4o","messages":[]}`
|
||||
if rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost, body); rec.Code != http.StatusOK {
|
||||
t.Fatalf("first call status=%d", rec.Code)
|
||||
}
|
||||
if row := env.logs.wait(t); row.ServedModel != "" {
|
||||
t.Errorf("first call must use the requested model: %+v", row)
|
||||
}
|
||||
if rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost, body); rec.Code != http.StatusOK {
|
||||
t.Fatalf("second call status=%d body=%s", rec.Code, rec.Body)
|
||||
}
|
||||
if row := env.logs.wait(t); row.ServedModel != "or/google/gemini-2.5-flash-image" {
|
||||
t.Errorf("second call must overflow to the backup: %+v", row)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatCompletionsClientErrorIsNotRetried(t *testing.T) {
|
||||
env := newTestEnv(t, failPrimary(http.StatusBadRequest), withBackup)
|
||||
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
||||
`{"model":"or/openai/gpt-4o","messages":[]}`)
|
||||
if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), "primary down") {
|
||||
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
||||
}
|
||||
if calls, _ := env.upstream.last(); calls != 1 {
|
||||
t.Errorf("upstream calls = %d, want 1", calls)
|
||||
}
|
||||
if row := env.logs.wait(t); strings.Contains(row.Error, "switched") {
|
||||
t.Errorf("log row = %+v", row)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatCompletionsLastFailureIsRelayed(t *testing.T) {
|
||||
// Both the model and its backup are down: the client gets the last error.
|
||||
env := newTestEnv(t, func(w http.ResponseWriter, _ *http.Request, _ map[string]any) {
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_, _ = io.WriteString(w, `{"error":{"message":"all down"}}`)
|
||||
}, withBackup)
|
||||
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
||||
`{"model":"or/openai/gpt-4o","messages":[]}`)
|
||||
if rec.Code != http.StatusServiceUnavailable || !strings.Contains(rec.Body.String(), "all down") {
|
||||
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
||||
}
|
||||
if calls, _ := env.upstream.last(); calls != 2 {
|
||||
t.Errorf("upstream calls = %d, want 2", calls)
|
||||
}
|
||||
env.logs.waitN(t, 2)
|
||||
}
|
||||
|
||||
func TestEmbeddingsFallBack(t *testing.T) {
|
||||
env := newTestEnv(t, func(w http.ResponseWriter, _ *http.Request, body map[string]any) {
|
||||
if body["model"] == "openai/gpt-4o" {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, `{"object":"list","model":"backup","data":[{"embedding":[0.1]}],"usage":{"prompt_tokens":4,"total_tokens":4}}`)
|
||||
}, withBackup)
|
||||
rec := call(env.api.embeddings, allowAllKey(), http.MethodPost, `{"model":"or/openai/gpt-4o","input":"x"}`)
|
||||
if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), `"model":"openai/gpt-4o"`) {
|
||||
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
||||
}
|
||||
rows := rowsByStatus(env.logs.waitN(t, 2))
|
||||
if rows[http.StatusOK].ServedModel != "or/google/gemini-2.5-flash-image" {
|
||||
t.Errorf("served row = %+v", rows[http.StatusOK])
|
||||
}
|
||||
}
|
||||
|
||||
132
backend/internal/api/route.go
Normal file
132
backend/internal/api/route.go
Normal file
@@ -0,0 +1,132 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"maps"
|
||||
"net/http"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/catalog"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/providers"
|
||||
)
|
||||
|
||||
// The proxy handlers walk a route's targets in order. The requested model is
|
||||
// tried first. When it is rate limited, has no provider, or its upstream
|
||||
// fails before the client has received anything, the next backup is tried
|
||||
// with the same request. A backup answers under the requested model's name,
|
||||
// so the client cannot tell the difference. The request log can: each failed
|
||||
// attempt gets its own row, and the row for the backup carries both the
|
||||
// requested address and the address that served it, with the backup's cost.
|
||||
|
||||
// failure is one target's refusal. When no target answers, the last failure
|
||||
// is sent to the client.
|
||||
type failure struct {
|
||||
status int
|
||||
body []byte // upstream error payload to relay; nil writes msg as a gateway error
|
||||
msg string
|
||||
}
|
||||
|
||||
func (f *failure) write(w http.ResponseWriter) {
|
||||
if f.body != nil {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(f.status)
|
||||
_, _ = w.Write(f.body)
|
||||
return
|
||||
}
|
||||
writeError(w, f.status, f.msg)
|
||||
}
|
||||
|
||||
// retryable reports whether an upstream status means the request should move
|
||||
// on to a backup. Client errors are not retried: a backup would get the same
|
||||
// request and most likely give the same answer.
|
||||
func retryable(status int) bool {
|
||||
return status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= http.StatusInternalServerError
|
||||
}
|
||||
|
||||
// admit charges requests against a target's rate limits and returns its
|
||||
// provider. The failure says why the target cannot serve.
|
||||
func (a *API) admit(e *catalog.Entry, requests int) (providers.Provider, *failure) {
|
||||
if !a.Limiter.AllowRequests(e.Address, ptrInt(e.RPMLimit), requests) {
|
||||
return nil, &failure{status: http.StatusTooManyRequests, msg: "rate limit exceeded (RPM)"}
|
||||
}
|
||||
if !a.Limiter.AllowTokens(e.Address, ptrInt(e.TPMLimit)) {
|
||||
return nil, &failure{status: http.StatusTooManyRequests, msg: "rate limit exceeded (TPM)"}
|
||||
}
|
||||
prov, ok := a.Manager.Get(e.ProviderName)
|
||||
if !ok {
|
||||
return nil, &failure{status: http.StatusServiceUnavailable, msg: "provider unavailable"}
|
||||
}
|
||||
return prov, nil
|
||||
}
|
||||
|
||||
// targetPayload copies the client payload for one target and applies the
|
||||
// target's parameter overrides. Each target gets its own copy, so the
|
||||
// overrides of a failed target do not leak into the request to its backup.
|
||||
func targetPayload(payload map[string]any, e *catalog.Entry) map[string]any {
|
||||
out := maps.Clone(payload)
|
||||
maps.Copy(out, e.ParamOverrides)
|
||||
return out
|
||||
}
|
||||
|
||||
// switchNote is appended to the log row of a failed attempt when the request
|
||||
// moves on to a backup. Empty when there is no backup after the attempt.
|
||||
func switchNote(next *catalog.Entry) string {
|
||||
if next == nil {
|
||||
return ""
|
||||
}
|
||||
return "; switched to backup " + next.Address
|
||||
}
|
||||
|
||||
// nextTarget returns the target after index i, or nil.
|
||||
func nextTarget(route *catalog.Route, i int) *catalog.Entry {
|
||||
if i+1 < len(route.Targets) {
|
||||
return route.Targets[i+1]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// reportedModel is the model id a response from e must carry. A backup
|
||||
// reports the requested model's id; the requested model keeps what its
|
||||
// upstream sent.
|
||||
func reportedModel(route *catalog.Route, e *catalog.Entry) string {
|
||||
if route.IsFallback(e) {
|
||||
return route.Slug
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// editJSON decodes an object body, applies edit, and encodes it again. A body
|
||||
// that is not a JSON object is returned unchanged.
|
||||
func editJSON(body []byte, edit func(map[string]any)) []byte {
|
||||
var obj map[string]any
|
||||
if json.Unmarshal(body, &obj) != nil || obj == nil {
|
||||
return body
|
||||
}
|
||||
edit(obj)
|
||||
out, err := json.Marshal(obj)
|
||||
if err != nil {
|
||||
return body
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// clientBody prepares an upstream chat body or chunk for the client. A missing
|
||||
// "created" timestamp is added, and when model is set the reported model is
|
||||
// replaced with it.
|
||||
func clientBody(body []byte, created int64, model string) []byte {
|
||||
return editJSON(body, func(obj map[string]any) {
|
||||
if _, ok := obj["created"]; !ok {
|
||||
obj["created"] = created
|
||||
}
|
||||
if model != "" {
|
||||
obj["model"] = model
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// withModel replaces the reported model when model is set.
|
||||
func withModel(body []byte, model string) []byte {
|
||||
if model == "" {
|
||||
return body
|
||||
}
|
||||
return editJSON(body, func(obj map[string]any) { obj["model"] = model })
|
||||
}
|
||||
42
backend/internal/api/stats_test.go
Normal file
42
backend/internal/api/stats_test.go
Normal file
@@ -0,0 +1,42 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
||||
)
|
||||
|
||||
type fakeStats struct {
|
||||
mu sync.Mutex
|
||||
rows []store.RequestLog
|
||||
}
|
||||
|
||||
func (f *fakeStats) Record(e store.RequestLog) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.rows = append(f.rows, e)
|
||||
}
|
||||
|
||||
func TestLogRequestForwardsRowToStats(t *testing.T) {
|
||||
env := newTestEnv(t, chatReply, nil)
|
||||
stats := &fakeStats{}
|
||||
env.api.Stats = stats
|
||||
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
||||
`{"model":"or/openai/gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rec.Code, rec.Body)
|
||||
}
|
||||
logged := env.logs.wait(t)
|
||||
|
||||
stats.mu.Lock()
|
||||
defer stats.mu.Unlock()
|
||||
if len(stats.rows) != 1 {
|
||||
t.Fatalf("stats rows = %d, want 1", len(stats.rows))
|
||||
}
|
||||
got := stats.rows[0]
|
||||
if got.ModelName != logged.ModelName || got.TotalTokens != logged.TotalTokens || got.Cost != logged.Cost || got.KeyPrefix != "sk-test" {
|
||||
t.Errorf("stats row = %+v, log row = %+v", got, logged)
|
||||
}
|
||||
}
|
||||
@@ -8,6 +8,8 @@ import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
||||
)
|
||||
|
||||
@@ -23,17 +25,26 @@ type Cache struct {
|
||||
|
||||
mu sync.RWMutex
|
||||
byAddr map[string]*Entry
|
||||
byID map[uuid.UUID]*Entry
|
||||
entries []*Entry
|
||||
strategy Strategy
|
||||
// standby holds enabled, unblocked models the provider does not serve
|
||||
// right now. They are not servable, but a request for one can still be
|
||||
// routed to its backup.
|
||||
standby map[string]*Entry
|
||||
|
||||
rr sync.Map // map[string]*atomic.Uint64, round-robin counters per query
|
||||
}
|
||||
|
||||
// NewCache returns an empty cache; call Reload to populate it.
|
||||
func NewCache(s CatalogStore) *Cache {
|
||||
return &Cache{store: s, byAddr: map[string]*Entry{}, strategy: StrategyRoundRobin}
|
||||
return &Cache{store: s, byAddr: map[string]*Entry{}, byID: map[uuid.UUID]*Entry{}, standby: map[string]*Entry{}, strategy: StrategyRoundRobin}
|
||||
}
|
||||
|
||||
// maxFallbackHops bounds a backup chain. A chain may legitimately loop (two
|
||||
// models that back each other up), so the walk also tracks visited ids.
|
||||
const maxFallbackHops = 8
|
||||
|
||||
// Reload rebuilds the cache (servable entries only) and routing strategy from
|
||||
// the store.
|
||||
func (c *Cache) Reload(ctx context.Context) error {
|
||||
@@ -43,9 +54,11 @@ func (c *Cache) Reload(ctx context.Context) error {
|
||||
}
|
||||
|
||||
byAddr := make(map[string]*Entry, len(rows))
|
||||
byID := make(map[uuid.UUID]*Entry, len(rows))
|
||||
standby := map[string]*Entry{}
|
||||
entries := make([]*Entry, 0, len(rows))
|
||||
for _, r := range rows {
|
||||
if !r.Enabled || r.Blocked || !r.Available {
|
||||
if !r.Enabled || r.Blocked {
|
||||
continue
|
||||
}
|
||||
var overrides map[string]any
|
||||
@@ -69,8 +82,14 @@ func (c *Cache) Reload(ctx context.Context) error {
|
||||
CachedPricePerMTok: r.CachedPricePerMTok,
|
||||
RPMLimit: r.RPMLimit,
|
||||
TPMLimit: r.TPMLimit,
|
||||
FallbackID: r.FallbackModelID,
|
||||
}
|
||||
if !r.Available {
|
||||
standby[e.Address] = e
|
||||
continue
|
||||
}
|
||||
byAddr[e.Address] = e
|
||||
byID[e.ID] = e
|
||||
entries = append(entries, e)
|
||||
}
|
||||
|
||||
@@ -82,15 +101,29 @@ func (c *Cache) Reload(ctx context.Context) error {
|
||||
|
||||
c.mu.Lock()
|
||||
c.byAddr = byAddr
|
||||
c.byID = byID
|
||||
c.standby = standby
|
||||
c.entries = entries
|
||||
c.strategy = strategy
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Resolve maps a model address to a servable entry. "any/<slug>" selects among
|
||||
// matching providers using the configured strategy.
|
||||
// Resolve maps a model address to the servable entry a request would reach
|
||||
// first: the model itself, or its first live backup when the model is down.
|
||||
func (c *Cache) Resolve(address string) (*Entry, error) {
|
||||
r, err := c.Route(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.Targets[0], nil
|
||||
}
|
||||
|
||||
// Route maps a model address to the entries that can serve it, in order.
|
||||
// "any/<slug>" selects among matching providers using the configured strategy;
|
||||
// the selected entry's backups follow it. A model that is enabled but not
|
||||
// available routes straight to its backups.
|
||||
func (c *Cache) Route(address string) (*Route, error) {
|
||||
provider, slug, ok := SplitAddress(address)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid model address %q (want <provider>/<slug>)", address)
|
||||
@@ -104,14 +137,51 @@ func (c *Cache) Resolve(address string) (*Entry, error) {
|
||||
if len(cands) == 0 {
|
||||
return nil, fmt.Errorf("no provider serves model %q", slug)
|
||||
}
|
||||
return c.selectLocked(slug, cands), nil
|
||||
e := c.selectLocked(slug, cands)
|
||||
return &Route{Address: address, Primary: e.Address, Slug: e.Slug, Targets: c.chainLocked(e)}, nil
|
||||
}
|
||||
|
||||
e, ok := c.byAddr[address]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("model %q not found", address)
|
||||
if e, ok := c.byAddr[address]; ok {
|
||||
return &Route{Address: address, Primary: address, Slug: e.Slug, Targets: c.chainLocked(e)}, nil
|
||||
}
|
||||
return e, nil
|
||||
if e, ok := c.standby[address]; ok {
|
||||
if targets := c.chainLocked(e); len(targets) > 0 {
|
||||
return &Route{Address: address, Primary: address, Slug: e.Slug, Targets: targets}, nil
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("model %q not found", address)
|
||||
}
|
||||
|
||||
// chainLocked returns e, when it is servable, followed by its live backups.
|
||||
// Backups that are not servable are skipped and the walk continues through
|
||||
// them, so a down backup does not hide the ones behind it.
|
||||
func (c *Cache) chainLocked(e *Entry) []*Entry {
|
||||
var out []*Entry
|
||||
seen := map[uuid.UUID]bool{}
|
||||
for hops := 0; e != nil && hops <= maxFallbackHops && !seen[e.ID]; hops++ {
|
||||
seen[e.ID] = true
|
||||
if c.byID[e.ID] == e {
|
||||
out = append(out, e)
|
||||
}
|
||||
e = c.lookupLocked(e.FallbackID)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// lookupLocked finds an enabled, unblocked entry by id, servable or standby.
|
||||
func (c *Cache) lookupLocked(id *uuid.UUID) *Entry {
|
||||
if id == nil {
|
||||
return nil
|
||||
}
|
||||
if e, ok := c.byID[*id]; ok {
|
||||
return e
|
||||
}
|
||||
for _, e := range c.standby {
|
||||
if e.ID == *id {
|
||||
return e
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// candidatesLocked returns servable entries whose slug equals slug or ends with
|
||||
@@ -149,12 +219,18 @@ func (c *Cache) selectLocked(key string, cands []*Entry) *Entry {
|
||||
|
||||
func cost(e *Entry) float64 { return e.InputPricePerMTok + e.OutputPricePerMTok }
|
||||
|
||||
// List returns a snapshot of servable entries.
|
||||
// List returns a snapshot of the entries clients can call: the servable ones,
|
||||
// plus unavailable models that have a live backup.
|
||||
func (c *Cache) List() []*Entry {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
out := make([]*Entry, len(c.entries))
|
||||
out := make([]*Entry, len(c.entries), len(c.entries)+len(c.standby))
|
||||
copy(out, c.entries)
|
||||
for _, e := range c.standby {
|
||||
if len(c.chainLocked(e)) > 0 {
|
||||
out = append(out, e)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
|
||||
@@ -1,19 +1,139 @@
|
||||
package catalog
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func newTestCache(strategy Strategy, entries ...*Entry) *Cache {
|
||||
c := &Cache{byAddr: map[string]*Entry{}, strategy: strategy}
|
||||
c := &Cache{byAddr: map[string]*Entry{}, byID: map[uuid.UUID]*Entry{}, standby: map[string]*Entry{}, strategy: strategy}
|
||||
for _, e := range entries {
|
||||
if e.Address == "" {
|
||||
e.Address = e.ProviderName + "/" + e.Slug
|
||||
}
|
||||
if e.ID == uuid.Nil {
|
||||
e.ID = uuid.New()
|
||||
}
|
||||
c.byAddr[e.Address] = e
|
||||
c.byID[e.ID] = e
|
||||
c.entries = append(c.entries, e)
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// addStandby registers an enabled model the provider does not serve right now.
|
||||
func addStandby(c *Cache, e *Entry) *Entry {
|
||||
if e.Address == "" {
|
||||
e.Address = e.ProviderName + "/" + e.Slug
|
||||
}
|
||||
if e.ID == uuid.Nil {
|
||||
e.ID = uuid.New()
|
||||
}
|
||||
c.standby[e.Address] = e
|
||||
return e
|
||||
}
|
||||
|
||||
func backup(e *Entry) *uuid.UUID { id := e.ID; return &id }
|
||||
|
||||
func addresses(entries []*Entry) []string {
|
||||
out := make([]string, len(entries))
|
||||
for i, e := range entries {
|
||||
out[i] = e.Address
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestRouteFollowsBackupChainAndSkipsDownBackups(t *testing.T) {
|
||||
a := &Entry{ProviderName: "local", Slug: "a"}
|
||||
b := &Entry{ProviderName: "local", Slug: "b"}
|
||||
cEntry := &Entry{ProviderName: "cloud", Slug: "c"}
|
||||
c := newTestCache(StrategyRoundRobin, a, cEntry)
|
||||
addStandby(c, b)
|
||||
a.FallbackID = backup(b)
|
||||
b.FallbackID = backup(cEntry)
|
||||
|
||||
r, err := c.Route("local/a")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := addresses(r.Targets); len(got) != 2 || got[0] != "local/a" || got[1] != "cloud/c" {
|
||||
t.Fatalf("targets = %v, want [local/a cloud/c]", got)
|
||||
}
|
||||
if r.Primary != "local/a" || r.Slug != "a" || r.IsFallback(a) || !r.IsFallback(cEntry) {
|
||||
t.Errorf("route = %+v", r)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouteUnavailableModelGoesToBackup(t *testing.T) {
|
||||
b := &Entry{ProviderName: "cloud", Slug: "org/a"}
|
||||
c := newTestCache(StrategyRoundRobin, b)
|
||||
a := addStandby(c, &Entry{ProviderName: "local", Slug: "a"})
|
||||
a.FallbackID = backup(b)
|
||||
|
||||
r, err := c.Route("local/a")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := addresses(r.Targets); len(got) != 1 || got[0] != "cloud/org/a" {
|
||||
t.Fatalf("targets = %v", got)
|
||||
}
|
||||
// The client asked for local/a, so the allowlist and the reported model
|
||||
// use it, not the backup.
|
||||
if r.Address != "local/a" || r.Primary != "local/a" || r.Slug != "a" || !r.IsFallback(b) {
|
||||
t.Errorf("route = %+v", r)
|
||||
}
|
||||
if e, err := c.Resolve("local/a"); err != nil || e != b {
|
||||
t.Errorf("Resolve = %v, %v", e, err)
|
||||
}
|
||||
if got := addresses(c.List()); len(got) != 2 {
|
||||
t.Errorf("List must include the down model with a live backup: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouteUnavailableModelWithoutBackupIsNotFound(t *testing.T) {
|
||||
c := newTestCache(StrategyRoundRobin)
|
||||
addStandby(c, &Entry{ProviderName: "local", Slug: "a"})
|
||||
if _, err := c.Route("local/a"); err == nil {
|
||||
t.Fatal("expected not found")
|
||||
}
|
||||
if got := c.List(); len(got) != 0 {
|
||||
t.Errorf("List = %v", addresses(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouteMutualBackupsDoNotLoop(t *testing.T) {
|
||||
a := &Entry{ProviderName: "x", Slug: "a"}
|
||||
b := &Entry{ProviderName: "x", Slug: "b"}
|
||||
c := newTestCache(StrategyRoundRobin, a, b)
|
||||
a.FallbackID = backup(b)
|
||||
b.FallbackID = backup(a)
|
||||
r, err := c.Route("x/b")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := addresses(r.Targets); len(got) != 2 || got[0] != "x/b" || got[1] != "x/a" {
|
||||
t.Fatalf("targets = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouteAnyCarriesBackups(t *testing.T) {
|
||||
a := &Entry{ProviderName: "openai", Slug: "gpt-4o"}
|
||||
b := &Entry{ProviderName: "cloud", Slug: "other"}
|
||||
c := newTestCache(StrategyRoundRobin, a, b)
|
||||
a.FallbackID = backup(b)
|
||||
r, err := c.Route("any/gpt-4o")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := addresses(r.Targets); len(got) != 2 || got[1] != "cloud/other" {
|
||||
t.Fatalf("targets = %v", got)
|
||||
}
|
||||
if r.Address != "any/gpt-4o" || r.Primary != "openai/gpt-4o" {
|
||||
t.Errorf("route = %+v", r)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitAddress(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
|
||||
@@ -41,4 +41,27 @@ type Entry struct {
|
||||
CachedPricePerMTok float64
|
||||
RPMLimit *int
|
||||
TPMLimit *int
|
||||
// FallbackID is the admin-configured backup model, when set. The cache
|
||||
// resolves it to servable entries in Route.
|
||||
FallbackID *uuid.UUID
|
||||
}
|
||||
|
||||
// Route is the ordered list of servable entries for one requested address:
|
||||
// the model the client asked for first, then its backups. A backup is used
|
||||
// when the model before it is unavailable or its upstream fails, and the
|
||||
// client sees no difference in the response.
|
||||
type Route struct {
|
||||
// Address is the address the client sent, for the request log.
|
||||
Address string
|
||||
// Primary is the address of the model the client asked for, after any
|
||||
// "any/<slug>" selection. Key allowlists are checked against it.
|
||||
Primary string
|
||||
// Slug is the primary model's upstream id. Responses served by a backup
|
||||
// report this id, so the switch stays invisible to the client.
|
||||
Slug string
|
||||
// Targets holds the servable entries to try, in order. It is never empty.
|
||||
Targets []*Entry
|
||||
}
|
||||
|
||||
// IsFallback reports whether e is a backup rather than the requested model.
|
||||
func (r *Route) IsFallback(e *Entry) bool { return e.Address != r.Primary }
|
||||
|
||||
@@ -108,4 +108,32 @@ func TestCatalogPipeline(t *testing.T) {
|
||||
if got := len(cache.List()); got != 3 {
|
||||
t.Errorf("catalog size = %d, want 3", got)
|
||||
}
|
||||
|
||||
// An admin context override wins over the discovered value and survives
|
||||
// the next sync.
|
||||
ctxOverride := 98304
|
||||
if err := st.UpdateModelOverride(ctx, e.ID, store.ModelOverride{ContextLength: store.Patch[int]{Set: true, Value: &ctxOverride}}); err != nil {
|
||||
t.Fatalf("UpdateModelOverride: %v", err)
|
||||
}
|
||||
if err := syncer.SyncAll(ctx, []catalog.ProviderRef{{ID: providerID, Provider: fake}}); err != nil {
|
||||
t.Fatalf("SyncAll after override: %v", err)
|
||||
}
|
||||
e, err = cache.Resolve("p1/openai/gpt-4o")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if e.ContextLength == nil || *e.ContextLength != 98304 {
|
||||
t.Errorf("context_length after override = %v, want 98304", e.ContextLength)
|
||||
}
|
||||
// Clearing the override falls back to the discovered value.
|
||||
if err := st.UpdateModelOverride(ctx, e.ID, store.ModelOverride{ContextLength: store.Patch[int]{Set: true}}); err != nil {
|
||||
t.Fatalf("clear override: %v", err)
|
||||
}
|
||||
if err := cache.Reload(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
e, _ = cache.Resolve("p1/openai/gpt-4o")
|
||||
if e.ContextLength == nil || *e.ContextLength != 128000 {
|
||||
t.Errorf("context_length after clearing = %v, want 128000", e.ContextLength)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,6 +100,9 @@ func normalize(providerType string, m providers.Model) store.DiscoveredModel {
|
||||
case "llamacpp":
|
||||
extractLlamaCpp(&d, m.Raw)
|
||||
}
|
||||
if m.ContextLength != nil && *m.ContextLength > 0 {
|
||||
d.ContextLength = m.ContextLength
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
|
||||
@@ -51,3 +51,15 @@ func TestNormalizeOtherProvidersLeaveCapabilitiesEmpty(t *testing.T) {
|
||||
t.Errorf("capabilities = %v, want nil so the column default applies", d.Capabilities)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizePrefersRuntimeContext(t *testing.T) {
|
||||
n := 98304
|
||||
d := normalize("llamacpp", providers.Model{ID: "m", ContextLength: &n, Raw: json.RawMessage(`{"id":"m","meta":{"n_ctx_train":262144}}`)})
|
||||
if d.ContextLength == nil || *d.ContextLength != 98304 {
|
||||
t.Errorf("context = %v, want the runtime value 98304 over the training value", d.ContextLength)
|
||||
}
|
||||
d = normalize("llamacpp", providers.Model{ID: "m", Raw: json.RawMessage(`{"id":"m","meta":{"n_ctx_train":262144}}`)})
|
||||
if d.ContextLength == nil || *d.ContextLength != 262144 {
|
||||
t.Errorf("without a runtime value the training context stays: %v", d.ContextLength)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,8 +15,22 @@ type Config struct {
|
||||
OIDC OIDCConfig
|
||||
Security SecurityConfig
|
||||
Catalog CatalogConfig
|
||||
Influx InfluxConfig
|
||||
}
|
||||
|
||||
// InfluxConfig enables the InfluxDB 3 statistics exporter when URL is set.
|
||||
type InfluxConfig struct {
|
||||
URL string
|
||||
Token string
|
||||
Database string
|
||||
Measurement string
|
||||
FlushInterval time.Duration
|
||||
BatchSize int
|
||||
}
|
||||
|
||||
// Enabled reports whether the exporter is configured.
|
||||
func (c InfluxConfig) Enabled() bool { return c.URL != "" }
|
||||
|
||||
type CatalogConfig struct {
|
||||
SyncInterval time.Duration // how often to refresh the model catalog
|
||||
}
|
||||
@@ -62,6 +76,9 @@ func Load() (*Config, error) {
|
||||
v.SetDefault("db.min_conns", 2)
|
||||
v.SetDefault("db.max_conn_lifetime", 30)
|
||||
v.SetDefault("catalog.sync_interval", "30m")
|
||||
v.SetDefault("influx.measurement", "lmr_request")
|
||||
v.SetDefault("influx.flush_interval", "10s")
|
||||
v.SetDefault("influx.batch_size", 500)
|
||||
|
||||
cfg := &Config{
|
||||
Server: ServerConfig{
|
||||
@@ -87,6 +104,14 @@ func Load() (*Config, error) {
|
||||
Catalog: CatalogConfig{
|
||||
SyncInterval: v.GetDuration("catalog.sync_interval"),
|
||||
},
|
||||
Influx: InfluxConfig{
|
||||
URL: v.GetString("influx.url"),
|
||||
Token: v.GetString("influx.token"),
|
||||
Database: v.GetString("influx.database"),
|
||||
Measurement: v.GetString("influx.measurement"),
|
||||
FlushInterval: v.GetDuration("influx.flush_interval"),
|
||||
BatchSize: v.GetInt("influx.batch_size"),
|
||||
},
|
||||
}
|
||||
|
||||
return cfg, cfg.validate()
|
||||
@@ -123,6 +148,9 @@ func (c *Config) validate() error {
|
||||
if len(missing) > 0 {
|
||||
return fmt.Errorf("missing required OIDC config: %s", strings.Join(missing, ", "))
|
||||
}
|
||||
if c.Influx.Enabled() && c.Influx.Database == "" {
|
||||
return fmt.Errorf("influx.database is required when influx.url is set")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
290
backend/internal/influx/influx.go
Normal file
290
backend/internal/influx/influx.go
Normal file
@@ -0,0 +1,290 @@
|
||||
// Package influx exports one point per proxied request to InfluxDB 3 over the
|
||||
// /api/v3/write_lp HTTP endpoint. Points are batched in memory and flushed by
|
||||
// size or by interval. The exporter never blocks a request: when the buffer is
|
||||
// full the point is dropped and counted.
|
||||
package influx
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
||||
)
|
||||
|
||||
// Config controls the exporter. URL, Database and Measurement are required.
|
||||
type Config struct {
|
||||
URL string
|
||||
Token string
|
||||
Database string
|
||||
Measurement string
|
||||
FlushInterval time.Duration
|
||||
BatchSize int
|
||||
// BufferSize is the number of points held between flushes. Zero = 10 × BatchSize.
|
||||
BufferSize int
|
||||
// Timeout bounds one write request. Zero = 10s.
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
const (
|
||||
defaultFlushInterval = 10 * time.Second
|
||||
defaultBatchSize = 500
|
||||
defaultTimeout = 10 * time.Second
|
||||
maxErrorLen = 512
|
||||
writeAttempts = 3
|
||||
)
|
||||
|
||||
// Exporter batches points and writes them to InfluxDB 3.
|
||||
type Exporter struct {
|
||||
cfg Config
|
||||
writeURL string
|
||||
client *http.Client
|
||||
queue chan Point
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
stopOnce sync.Once
|
||||
dropped atomic.Int64
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
// New validates cfg and returns a stopped exporter. Call Start to run it.
|
||||
func New(cfg Config) (*Exporter, error) {
|
||||
if cfg.URL == "" {
|
||||
return nil, errors.New("influx: url is required")
|
||||
}
|
||||
if cfg.Database == "" {
|
||||
return nil, errors.New("influx: database is required")
|
||||
}
|
||||
if cfg.Measurement == "" {
|
||||
return nil, errors.New("influx: measurement is required")
|
||||
}
|
||||
base, err := url.Parse(cfg.URL)
|
||||
if err != nil || (base.Scheme != "http" && base.Scheme != "https") || base.Host == "" {
|
||||
return nil, fmt.Errorf("influx: url %q must be an http(s) URL", cfg.URL)
|
||||
}
|
||||
if cfg.FlushInterval <= 0 {
|
||||
cfg.FlushInterval = defaultFlushInterval
|
||||
}
|
||||
if cfg.BatchSize <= 0 {
|
||||
cfg.BatchSize = defaultBatchSize
|
||||
}
|
||||
if cfg.BufferSize <= 0 {
|
||||
cfg.BufferSize = 10 * cfg.BatchSize
|
||||
}
|
||||
if cfg.Timeout <= 0 {
|
||||
cfg.Timeout = defaultTimeout
|
||||
}
|
||||
|
||||
q := url.Values{}
|
||||
q.Set("db", cfg.Database)
|
||||
q.Set("precision", "nanosecond")
|
||||
base.Path = strings.TrimRight(base.Path, "/") + "/api/v3/write_lp"
|
||||
base.RawQuery = q.Encode()
|
||||
|
||||
return &Exporter{
|
||||
cfg: cfg,
|
||||
writeURL: base.String(),
|
||||
client: &http.Client{Timeout: cfg.Timeout},
|
||||
queue: make(chan Point, cfg.BufferSize),
|
||||
stop: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
now: time.Now,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Start runs the flush loop in the background.
|
||||
func (e *Exporter) Start() {
|
||||
go e.loop()
|
||||
}
|
||||
|
||||
// Close stops the loop, writes the points still in the buffer, and waits for
|
||||
// that write or for ctx to end.
|
||||
func (e *Exporter) Close(ctx context.Context) error {
|
||||
e.stopOnce.Do(func() { close(e.stop) })
|
||||
select {
|
||||
case <-e.done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// Record converts a request log row to a point and queues it. It never blocks.
|
||||
func (e *Exporter) Record(rec store.RequestLog) {
|
||||
e.Enqueue(e.requestPoint(rec))
|
||||
}
|
||||
|
||||
// Enqueue queues an arbitrary point. It never blocks: a full buffer drops the point.
|
||||
func (e *Exporter) Enqueue(p Point) {
|
||||
select {
|
||||
case e.queue <- p:
|
||||
default:
|
||||
e.dropped.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
// requestPoint maps a log row to the request measurement. Bounded values are
|
||||
// tags, so dashboards can group by them. Per-request values are fields. Every
|
||||
// field is written on every point, because InfluxDB 3 creates a column on its
|
||||
// first write and a query that names a missing column fails.
|
||||
func (e *Exporter) requestPoint(rec store.RequestLog) Point {
|
||||
keyPrefix := rec.KeyPrefix
|
||||
if keyPrefix == "" {
|
||||
keyPrefix = "(no key)"
|
||||
}
|
||||
tags := map[string]string{
|
||||
"provider": rec.ProviderName,
|
||||
"model": rec.ModelName,
|
||||
"key_prefix": keyPrefix,
|
||||
"status_code": strconv.Itoa(rec.StatusCode),
|
||||
"streamed": strconv.FormatBool(rec.Streamed),
|
||||
// fallback is true when a backup model served the request; then
|
||||
// served_model names it and cost is the backup's cost.
|
||||
"fallback": strconv.FormatBool(rec.ServedModel != ""),
|
||||
}
|
||||
if rec.ServedModel != "" {
|
||||
tags["served_model"] = rec.ServedModel
|
||||
}
|
||||
if rec.UserID != nil {
|
||||
tags["user_id"] = rec.UserID.String()
|
||||
}
|
||||
if rec.ProviderID != nil {
|
||||
tags["provider_id"] = rec.ProviderID.String()
|
||||
}
|
||||
if rec.ProviderModelID != nil {
|
||||
tags["provider_model_id"] = rec.ProviderModelID.String()
|
||||
}
|
||||
|
||||
isErr := 0
|
||||
if rec.StatusCode >= 400 || rec.Error != "" {
|
||||
isErr = 1
|
||||
}
|
||||
// ttft_ms is 0 when the response was not streamed.
|
||||
ttft := 0
|
||||
if rec.TTFTMs != nil {
|
||||
ttft = *rec.TTFTMs
|
||||
}
|
||||
msg := rec.Error
|
||||
if len(msg) > maxErrorLen {
|
||||
msg = msg[:maxErrorLen]
|
||||
}
|
||||
fields := map[string]any{
|
||||
"requests": 1,
|
||||
"prompt_tokens": rec.PromptTokens,
|
||||
"completion_tokens": rec.CompletionTokens,
|
||||
"cached_tokens": rec.CachedTokens,
|
||||
"total_tokens": rec.TotalTokens,
|
||||
"cost": rec.Cost,
|
||||
"latency_ms": rec.LatencyMS,
|
||||
"ttft_ms": ttft,
|
||||
"errors": isErr,
|
||||
"request_id": rec.RequestID,
|
||||
"finish_reason": rec.FinishReason,
|
||||
"error": msg,
|
||||
}
|
||||
return Point{Measurement: e.cfg.Measurement, Tags: tags, Fields: fields, Time: e.now()}
|
||||
}
|
||||
|
||||
func (e *Exporter) loop() {
|
||||
defer close(e.done)
|
||||
batch := make([]Point, 0, e.cfg.BatchSize)
|
||||
tick := time.NewTicker(e.cfg.FlushInterval)
|
||||
defer tick.Stop()
|
||||
|
||||
flush := func() {
|
||||
if n := e.dropped.Swap(0); n > 0 {
|
||||
slog.Warn("influx exporter dropped points: buffer full", "points", n)
|
||||
}
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
e.write(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case p := <-e.queue:
|
||||
batch = append(batch, p)
|
||||
if len(batch) >= e.cfg.BatchSize {
|
||||
flush()
|
||||
}
|
||||
case <-tick.C:
|
||||
flush()
|
||||
case <-e.stop:
|
||||
for {
|
||||
select {
|
||||
case p := <-e.queue:
|
||||
batch = append(batch, p)
|
||||
if len(batch) >= e.cfg.BatchSize {
|
||||
flush()
|
||||
}
|
||||
default:
|
||||
flush()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// write sends one batch. Network errors and 5xx responses are retried with a
|
||||
// short backoff. 4xx responses are not, because the same body would fail again.
|
||||
func (e *Exporter) write(batch []Point) {
|
||||
var b strings.Builder
|
||||
for _, p := range batch {
|
||||
p.encode(&b)
|
||||
}
|
||||
body := b.String()
|
||||
if body == "" {
|
||||
return
|
||||
}
|
||||
|
||||
var err error
|
||||
for attempt := 1; attempt <= writeAttempts; attempt++ {
|
||||
var retry bool
|
||||
retry, err = e.post(body)
|
||||
if err == nil {
|
||||
return
|
||||
}
|
||||
if !retry {
|
||||
break
|
||||
}
|
||||
time.Sleep(time.Duration(attempt) * 500 * time.Millisecond)
|
||||
}
|
||||
slog.Warn("influx write failed, batch dropped", "points", len(batch), "err", err)
|
||||
}
|
||||
|
||||
// post sends one write request. The bool reports whether a retry makes sense.
|
||||
func (e *Exporter) post(body string) (bool, error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), e.cfg.Timeout)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.writeURL, strings.NewReader(body))
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "text/plain; charset=utf-8")
|
||||
if e.cfg.Token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+e.cfg.Token)
|
||||
}
|
||||
resp, err := e.client.Do(req)
|
||||
if err != nil {
|
||||
return true, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
return false, nil
|
||||
}
|
||||
msg, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
|
||||
err = fmt.Errorf("influx returned %d: %s", resp.StatusCode, strings.TrimSpace(string(msg)))
|
||||
return resp.StatusCode >= 500 || resp.StatusCode == http.StatusTooManyRequests, err
|
||||
}
|
||||
292
backend/internal/influx/influx_test.go
Normal file
292
backend/internal/influx/influx_test.go
Normal file
@@ -0,0 +1,292 @@
|
||||
package influx
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
||||
)
|
||||
|
||||
func TestEncodeEscapesAndSorts(t *testing.T) {
|
||||
p := Point{
|
||||
Measurement: "my measure,ment",
|
||||
Tags: map[string]string{"z": "last", "a": "has space", "empty": ""},
|
||||
Fields: map[string]any{
|
||||
"s": `quote " back \ newline` + "\nend",
|
||||
"i": int(7),
|
||||
"i64": int64(-3),
|
||||
"f": 1.5,
|
||||
"b": true,
|
||||
"nan": 0.0 / zero(),
|
||||
"bad": struct{}{},
|
||||
},
|
||||
Time: time.Unix(1, 2),
|
||||
}
|
||||
var b strings.Builder
|
||||
if !p.encode(&b) {
|
||||
t.Fatal("encode returned false")
|
||||
}
|
||||
want := `my\ measure\,ment,a=has\ space,z=last b=true,f=1.5,i=7i,i64=-3i,s="quote \" back \\ newline end" 1000000002` + "\n"
|
||||
if b.String() != want {
|
||||
t.Fatalf("got\n%s\nwant\n%s", b.String(), want)
|
||||
}
|
||||
}
|
||||
|
||||
func zero() float64 { return 0 }
|
||||
|
||||
func TestEncodeNoFields(t *testing.T) {
|
||||
var b strings.Builder
|
||||
if (Point{Measurement: "m", Fields: map[string]any{"x": struct{}{}}}).encode(&b) {
|
||||
t.Fatal("expected false for a point without usable fields")
|
||||
}
|
||||
if b.Len() != 0 {
|
||||
t.Fatalf("expected no output, got %q", b.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewRejectsBadConfig(t *testing.T) {
|
||||
cases := []Config{
|
||||
{Database: "d", Measurement: "m"},
|
||||
{URL: "http://x", Measurement: "m"},
|
||||
{URL: "http://x", Database: "d"},
|
||||
{URL: "not a url", Database: "d", Measurement: "m"},
|
||||
{URL: "ftp://x", Database: "d", Measurement: "m"},
|
||||
}
|
||||
for i, c := range cases {
|
||||
if _, err := New(c); err == nil {
|
||||
t.Errorf("case %d: expected error for %+v", i, c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewWriteURL(t *testing.T) {
|
||||
e, err := New(Config{URL: "http://influx:8181/prefix/", Database: "lmr", Measurement: "m"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "http://influx:8181/prefix/api/v3/write_lp?db=lmr&precision=nanosecond"
|
||||
if e.writeURL != want {
|
||||
t.Fatalf("writeURL = %q, want %q", e.writeURL, want)
|
||||
}
|
||||
}
|
||||
|
||||
// captureServer records every write body and the auth header it saw.
|
||||
type captureServer struct {
|
||||
*httptest.Server
|
||||
mu sync.Mutex
|
||||
bodies []string
|
||||
auth string
|
||||
status atomic.Int32
|
||||
got chan struct{}
|
||||
}
|
||||
|
||||
func newCaptureServer(t *testing.T) *captureServer {
|
||||
t.Helper()
|
||||
cs := &captureServer{got: make(chan struct{}, 64)}
|
||||
cs.status.Store(http.StatusNoContent)
|
||||
cs.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/v3/write_lp" {
|
||||
t.Errorf("unexpected path %s", r.URL.Path)
|
||||
}
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
cs.mu.Lock()
|
||||
cs.bodies = append(cs.bodies, string(b))
|
||||
cs.auth = r.Header.Get("Authorization")
|
||||
cs.mu.Unlock()
|
||||
w.WriteHeader(int(cs.status.Load()))
|
||||
cs.got <- struct{}{}
|
||||
}))
|
||||
t.Cleanup(cs.Close)
|
||||
return cs
|
||||
}
|
||||
|
||||
func (cs *captureServer) wait(t *testing.T) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-cs.got:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("no write arrived")
|
||||
}
|
||||
}
|
||||
|
||||
func newTestExporter(t *testing.T, srv *captureServer, cfg Config) *Exporter {
|
||||
t.Helper()
|
||||
cfg.URL = srv.URL
|
||||
cfg.Database = "lmr"
|
||||
cfg.Measurement = "lmr_request"
|
||||
cfg.Token = "secret"
|
||||
e, err := New(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
e.now = func() time.Time { return time.Unix(100, 0) }
|
||||
return e
|
||||
}
|
||||
|
||||
func TestRecordWritesRequestPoint(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
|
||||
e.Start()
|
||||
|
||||
uid := uuid.MustParse("11111111-1111-1111-1111-111111111111")
|
||||
ttft := 42
|
||||
e.Record(store.RequestLog{
|
||||
UserID: &uid,
|
||||
KeyPrefix: "lmr_abc",
|
||||
ModelName: "openai/gpt-4o",
|
||||
ProviderName: "openai",
|
||||
StatusCode: 200,
|
||||
Streamed: true,
|
||||
PromptTokens: 10,
|
||||
CompletionTokens: 5,
|
||||
CachedTokens: 2,
|
||||
TotalTokens: 15,
|
||||
Cost: 0.00125,
|
||||
LatencyMS: 321,
|
||||
TTFTMs: &ttft,
|
||||
RequestID: "req-1",
|
||||
})
|
||||
srv.wait(t)
|
||||
if err := e.Close(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
srv.mu.Lock()
|
||||
defer srv.mu.Unlock()
|
||||
if srv.auth != "Bearer secret" {
|
||||
t.Fatalf("auth = %q", srv.auth)
|
||||
}
|
||||
want := `lmr_request,fallback=false,key_prefix=lmr_abc,model=openai/gpt-4o,provider=openai,status_code=200,streamed=true,user_id=11111111-1111-1111-1111-111111111111 ` +
|
||||
`cached_tokens=2i,completion_tokens=5i,cost=0.00125,error="",errors=0i,finish_reason="",latency_ms=321i,prompt_tokens=10i,request_id="req-1",requests=1i,total_tokens=15i,ttft_ms=42i 100000000000` + "\n"
|
||||
if len(srv.bodies) != 1 || srv.bodies[0] != want {
|
||||
t.Fatalf("bodies = %q\nwant %q", srv.bodies, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordMarksErrors(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
|
||||
long := strings.Repeat("x", maxErrorLen+10)
|
||||
p := e.requestPoint(store.RequestLog{StatusCode: 502, Error: long})
|
||||
if p.Fields["errors"] != 1 {
|
||||
t.Fatalf("errors = %v", p.Fields["errors"])
|
||||
}
|
||||
if got := p.Fields["error"].(string); len(got) != maxErrorLen {
|
||||
t.Fatalf("error length = %d", len(got))
|
||||
}
|
||||
p = e.requestPoint(store.RequestLog{StatusCode: 200, Error: "stream interrupted"})
|
||||
if p.Fields["errors"] != 1 {
|
||||
t.Fatal("a 200 with an error message must count as an error")
|
||||
}
|
||||
}
|
||||
|
||||
// Every field must be present on every point so the columns exist from the
|
||||
// first write. Optional values get an empty string or zero.
|
||||
func TestRequestPointWritesAllFields(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{})
|
||||
p := e.requestPoint(store.RequestLog{StatusCode: 200})
|
||||
for _, k := range []string{"ttft_ms", "error", "finish_reason", "request_id"} {
|
||||
if _, ok := p.Fields[k]; !ok {
|
||||
t.Errorf("field %s missing", k)
|
||||
}
|
||||
}
|
||||
if p.Fields["ttft_ms"] != 0 || p.Fields["error"] != "" {
|
||||
t.Errorf("fields = %v", p.Fields)
|
||||
}
|
||||
if p.Tags["key_prefix"] != "(no key)" {
|
||||
t.Errorf("key_prefix = %q", p.Tags["key_prefix"])
|
||||
}
|
||||
if p.Tags["fallback"] != "false" {
|
||||
t.Errorf("fallback = %q", p.Tags["fallback"])
|
||||
}
|
||||
if _, ok := p.Tags["served_model"]; ok {
|
||||
t.Errorf("served_model must be absent when the requested model served")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestPointTagsBackupModel(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{})
|
||||
p := e.requestPoint(store.RequestLog{StatusCode: 200, ModelName: "local/a", ServedModel: "cloud/b"})
|
||||
if p.Tags["fallback"] != "true" || p.Tags["served_model"] != "cloud/b" || p.Tags["model"] != "local/a" {
|
||||
t.Errorf("tags = %v", p.Tags)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchAndFlushOnClose(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 2})
|
||||
e.Start()
|
||||
for range 3 {
|
||||
e.Record(store.RequestLog{StatusCode: 200})
|
||||
}
|
||||
srv.wait(t) // first batch of 2
|
||||
if err := e.Close(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
srv.mu.Lock()
|
||||
defer srv.mu.Unlock()
|
||||
if len(srv.bodies) != 2 {
|
||||
t.Fatalf("expected 2 writes, got %d", len(srv.bodies))
|
||||
}
|
||||
if n := strings.Count(srv.bodies[0], "\n"); n != 2 {
|
||||
t.Fatalf("first batch has %d lines, want 2", n)
|
||||
}
|
||||
if n := strings.Count(srv.bodies[1], "\n"); n != 1 {
|
||||
t.Fatalf("second batch has %d lines, want 1", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFlushOnInterval(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: 20 * time.Millisecond, BatchSize: 100})
|
||||
e.Start()
|
||||
defer e.Close(context.Background()) //nolint:errcheck
|
||||
e.Record(store.RequestLog{StatusCode: 200})
|
||||
srv.wait(t)
|
||||
}
|
||||
|
||||
func TestEnqueueDropsWhenFull(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1, BufferSize: 1})
|
||||
// Not started: nothing drains the queue.
|
||||
e.Enqueue(Point{Measurement: "m", Fields: map[string]any{"v": 1}})
|
||||
e.Enqueue(Point{Measurement: "m", Fields: map[string]any{"v": 1}})
|
||||
if e.dropped.Load() != 1 {
|
||||
t.Fatalf("dropped = %d, want 1", e.dropped.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteDoesNotRetryClientErrors(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
srv.status.Store(http.StatusBadRequest)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
|
||||
e.write([]Point{{Measurement: "m", Fields: map[string]any{"v": 1}, Time: time.Unix(1, 0)}})
|
||||
srv.mu.Lock()
|
||||
defer srv.mu.Unlock()
|
||||
if len(srv.bodies) != 1 {
|
||||
t.Fatalf("expected 1 attempt, got %d", len(srv.bodies))
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteRetriesServerErrors(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
srv.status.Store(http.StatusInternalServerError)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
|
||||
e.write([]Point{{Measurement: "m", Fields: map[string]any{"v": 1}, Time: time.Unix(1, 0)}})
|
||||
srv.mu.Lock()
|
||||
defer srv.mu.Unlock()
|
||||
if len(srv.bodies) != writeAttempts {
|
||||
t.Fatalf("expected %d attempts, got %d", writeAttempts, len(srv.bodies))
|
||||
}
|
||||
}
|
||||
91
backend/internal/influx/lineprotocol.go
Normal file
91
backend/internal/influx/lineprotocol.go
Normal file
@@ -0,0 +1,91 @@
|
||||
package influx
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Point is one line-protocol row.
|
||||
type Point struct {
|
||||
Measurement string
|
||||
Tags map[string]string
|
||||
Fields map[string]any
|
||||
Time time.Time
|
||||
}
|
||||
|
||||
// encode appends the point as one line-protocol line. It returns false when
|
||||
// the point has no usable field, because InfluxDB rejects such a line.
|
||||
func (p Point) encode(b *strings.Builder) bool {
|
||||
fields := make([]string, 0, len(p.Fields))
|
||||
for _, k := range sortedKeys(p.Fields) {
|
||||
v, ok := formatField(p.Fields[k])
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fields = append(fields, escapeKey(k)+"="+v)
|
||||
}
|
||||
if len(fields) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
b.WriteString(escapeMeasurement(p.Measurement))
|
||||
for _, k := range sortedKeys(p.Tags) {
|
||||
v := p.Tags[k]
|
||||
if v == "" {
|
||||
continue
|
||||
}
|
||||
b.WriteByte(',')
|
||||
b.WriteString(escapeKey(k))
|
||||
b.WriteByte('=')
|
||||
b.WriteString(escapeKey(v))
|
||||
}
|
||||
b.WriteByte(' ')
|
||||
b.WriteString(strings.Join(fields, ","))
|
||||
b.WriteByte(' ')
|
||||
b.WriteString(strconv.FormatInt(p.Time.UnixNano(), 10))
|
||||
b.WriteByte('\n')
|
||||
return true
|
||||
}
|
||||
|
||||
// formatField renders a field value. Unsupported types and non-finite floats
|
||||
// are dropped so one bad value does not reject the whole batch.
|
||||
func formatField(v any) (string, bool) {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return strconv.FormatInt(int64(x), 10) + "i", true
|
||||
case int64:
|
||||
return strconv.FormatInt(x, 10) + "i", true
|
||||
case float64:
|
||||
if math.IsNaN(x) || math.IsInf(x, 0) {
|
||||
return "", false
|
||||
}
|
||||
return strconv.FormatFloat(x, 'f', -1, 64), true
|
||||
case bool:
|
||||
return strconv.FormatBool(x), true
|
||||
case string:
|
||||
return `"` + escapeString(x) + `"`, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func sortedKeys[V any](m map[string]V) []string {
|
||||
keys := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
|
||||
var (
|
||||
measurementEscaper = strings.NewReplacer(",", `\,`, " ", `\ `, "\n", " ")
|
||||
keyEscaper = strings.NewReplacer(",", `\,`, "=", `\=`, " ", `\ `, "\n", " ")
|
||||
stringEscaper = strings.NewReplacer(`\`, `\\`, `"`, `\"`, "\n", " ", "\r", " ")
|
||||
)
|
||||
|
||||
func escapeMeasurement(s string) string { return measurementEscaper.Replace(s) }
|
||||
func escapeKey(s string) string { return keyEscaper.Replace(s) }
|
||||
func escapeString(s string) string { return stringEscaper.Replace(s) }
|
||||
@@ -63,4 +63,9 @@ type Model struct {
|
||||
OwnedBy string `json:"owned_by"`
|
||||
Created int64 `json:"created"`
|
||||
Raw json.RawMessage `json:"-"`
|
||||
// ContextLength is the context window the running instance accepts, when
|
||||
// the provider can find it out (llama.cpp reports it on /props). It wins
|
||||
// over whatever the listing entry itself claims, which is often the
|
||||
// training context.
|
||||
ContextLength *int `json:"-"`
|
||||
}
|
||||
|
||||
204
backend/internal/providers/llamacpp/context.go
Normal file
204
backend/internal/providers/llamacpp/context.go
Normal file
@@ -0,0 +1,204 @@
|
||||
package llamacpp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/providers"
|
||||
)
|
||||
|
||||
// probeTimeout bounds each extra request made during a model sync.
|
||||
const probeTimeout = 5 * time.Second
|
||||
|
||||
// ListModels lists the server's models and, where the server tells us, the
|
||||
// context window each running instance actually accepts. The OpenAI listing
|
||||
// only carries meta.n_ctx_train, the training context, which is often far
|
||||
// larger than the -c the server was started with.
|
||||
//
|
||||
// Three front ends are recognised, in this order:
|
||||
//
|
||||
// - the llama.cpp router (llama-server --models-dir): GET /models lists
|
||||
// every model with a status and its launch arguments; loaded models
|
||||
// answer GET /props?model=<id> with the live n_ctx, unloaded ones give
|
||||
// up their -c from the arguments.
|
||||
// - llama-swap: GET /running lists the loaded models, and each answers
|
||||
// GET /upstream/<id>/props. Unloaded models are left alone, because
|
||||
// asking llama-swap for them would start them.
|
||||
// - a single llama-server: GET /props describes the one loaded model.
|
||||
//
|
||||
// Probe failures never fail the listing. A model the probes do not cover
|
||||
// keeps the value from the listing.
|
||||
func (p *provider) ListModels(ctx context.Context) ([]providers.Model, error) {
|
||||
models, err := p.Transport.ListModels(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p.attachRunningContext(ctx, models)
|
||||
return models, nil
|
||||
}
|
||||
|
||||
func (p *provider) attachRunningContext(ctx context.Context, models []providers.Model) {
|
||||
if len(models) == 0 {
|
||||
return
|
||||
}
|
||||
byID := make(map[string]*providers.Model, len(models))
|
||||
for i := range models {
|
||||
byID[models[i].ID] = &models[i]
|
||||
}
|
||||
root := p.RootURL()
|
||||
|
||||
if found := p.probeRouter(ctx, root, byID); found {
|
||||
return
|
||||
}
|
||||
if found := p.probeSwap(ctx, root, byID); found {
|
||||
return
|
||||
}
|
||||
p.probeSingle(ctx, root, models)
|
||||
}
|
||||
|
||||
// routerListing is the llama.cpp router's GET /models shape.
|
||||
type routerListing struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
Status *struct {
|
||||
Value string `json:"value"`
|
||||
Args []string `json:"args"`
|
||||
} `json:"status"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
// propsReply is the part of GET /props this package reads.
|
||||
type propsReply struct {
|
||||
DefaultGenerationSettings struct {
|
||||
NCtx int `json:"n_ctx"`
|
||||
} `json:"default_generation_settings"`
|
||||
}
|
||||
|
||||
// probeRouter handles the llama.cpp router. It reports whether the server
|
||||
// looked like one, so the caller can stop probing.
|
||||
func (p *provider) probeRouter(ctx context.Context, root string, byID map[string]*providers.Model) bool {
|
||||
var listing routerListing
|
||||
status, err := p.getJSON(ctx, root+"/models", &listing)
|
||||
if err != nil || status >= 300 {
|
||||
return false
|
||||
}
|
||||
isRouter := false
|
||||
for _, entry := range listing.Data {
|
||||
if entry.Status == nil {
|
||||
continue
|
||||
}
|
||||
isRouter = true
|
||||
m, ok := byID[entry.ID]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if entry.Status.Value == "loaded" {
|
||||
if n := p.propsContext(ctx, root+"/props?model="+url.QueryEscape(entry.ID)); n > 0 {
|
||||
m.ContextLength = &n
|
||||
continue
|
||||
}
|
||||
}
|
||||
if n := contextFromArgs(entry.Status.Args); n > 0 {
|
||||
m.ContextLength = &n
|
||||
}
|
||||
}
|
||||
return isRouter
|
||||
}
|
||||
|
||||
// swapRunning is llama-swap's GET /running shape.
|
||||
type swapRunning struct {
|
||||
Running []struct {
|
||||
Model string `json:"model"`
|
||||
State string `json:"state"`
|
||||
} `json:"running"`
|
||||
}
|
||||
|
||||
// probeSwap handles llama-swap. Only running models are asked, because an
|
||||
// /upstream request for a stopped model would start it.
|
||||
func (p *provider) probeSwap(ctx context.Context, root string, byID map[string]*providers.Model) bool {
|
||||
var running swapRunning
|
||||
status, err := p.getJSON(ctx, root+"/running", &running)
|
||||
if err != nil || status >= 300 {
|
||||
return false
|
||||
}
|
||||
for _, r := range running.Running {
|
||||
if r.State != "" && r.State != "ready" {
|
||||
continue
|
||||
}
|
||||
m, ok := byID[r.Model]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if n := p.propsContext(ctx, root+"/upstream/"+url.PathEscape(r.Model)+"/props"); n > 0 {
|
||||
m.ContextLength = &n
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// probeSingle handles a plain llama-server, whose listing has one model and
|
||||
// whose /props describes it.
|
||||
func (p *provider) probeSingle(ctx context.Context, root string, models []providers.Model) {
|
||||
n := p.propsContext(ctx, root+"/props")
|
||||
if n <= 0 {
|
||||
return
|
||||
}
|
||||
for i := range models {
|
||||
v := n
|
||||
models[i].ContextLength = &v
|
||||
}
|
||||
}
|
||||
|
||||
// propsContext reads default_generation_settings.n_ctx from a /props URL.
|
||||
// That value is the per-slot context, which is what one request may use.
|
||||
func (p *provider) propsContext(ctx context.Context, rawURL string) int {
|
||||
var props propsReply
|
||||
status, err := p.getJSON(ctx, rawURL, &props)
|
||||
if err != nil || status >= 300 {
|
||||
return 0
|
||||
}
|
||||
return props.DefaultGenerationSettings.NCtx
|
||||
}
|
||||
|
||||
func (p *provider) getJSON(ctx context.Context, rawURL string, v any) (int, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, probeTimeout)
|
||||
defer cancel()
|
||||
status, err := p.GetJSON(ctx, rawURL, v)
|
||||
if err != nil {
|
||||
slog.Debug("llamacpp: context probe failed", "provider", p.name, "url", rawURL, "err", err)
|
||||
}
|
||||
return status, err
|
||||
}
|
||||
|
||||
// contextFromArgs reads the context size out of llama-server launch
|
||||
// arguments: -c N, -ctx N, --ctx-size N, or --ctx-size=N. When the server
|
||||
// runs several parallel slots (-np N, --parallel N) the context is shared
|
||||
// between them, so the per-request window is the total divided by the slot
|
||||
// count.
|
||||
func contextFromArgs(args []string) int {
|
||||
ctxSize, slots := 0, 1
|
||||
for i := 0; i < len(args); i++ {
|
||||
flag, value, inline := strings.Cut(args[i], "=")
|
||||
if !inline && i+1 < len(args) {
|
||||
value = args[i+1]
|
||||
}
|
||||
switch flag {
|
||||
case "-c", "-ctx", "--ctx-size", "--ctx_size":
|
||||
if n, err := strconv.Atoi(value); err == nil && n > 0 {
|
||||
ctxSize = n
|
||||
}
|
||||
case "-np", "--parallel":
|
||||
if n, err := strconv.Atoi(value); err == nil && n > 1 {
|
||||
slots = n
|
||||
}
|
||||
}
|
||||
}
|
||||
if ctxSize == 0 {
|
||||
return 0
|
||||
}
|
||||
return ctxSize / slots
|
||||
}
|
||||
185
backend/internal/providers/llamacpp/context_test.go
Normal file
185
backend/internal/providers/llamacpp/context_test.go
Normal file
@@ -0,0 +1,185 @@
|
||||
package llamacpp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/providers"
|
||||
)
|
||||
|
||||
// fakeServer plays one of the llama.cpp front ends. routes maps a path (with
|
||||
// query string) to a JSON body; anything else is 404. It records every path
|
||||
// requested so tests can assert that stopped models were not touched.
|
||||
type fakeServer struct {
|
||||
mu sync.Mutex
|
||||
routes map[string]string
|
||||
seen []string
|
||||
}
|
||||
|
||||
func (f *fakeServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
key := r.URL.Path
|
||||
if r.URL.RawQuery != "" {
|
||||
key += "?" + r.URL.RawQuery
|
||||
}
|
||||
f.mu.Lock()
|
||||
f.seen = append(f.seen, key)
|
||||
body, ok := f.routes[key]
|
||||
f.mu.Unlock()
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(body))
|
||||
}
|
||||
|
||||
func newProvider(t *testing.T, routes map[string]string) (*provider, *fakeServer) {
|
||||
t.Helper()
|
||||
fs := &fakeServer{routes: routes}
|
||||
srv := httptest.NewServer(fs)
|
||||
t.Cleanup(srv.Close)
|
||||
p, err := New(providers.Config{Name: "local", BaseURL: srv.URL + "/v1", HTTPClient: srv.Client()})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return p.(*provider), fs
|
||||
}
|
||||
|
||||
func ctxOf(t *testing.T, models []providers.Model, id string) *int {
|
||||
t.Helper()
|
||||
for _, m := range models {
|
||||
if m.ID == id {
|
||||
return m.ContextLength
|
||||
}
|
||||
}
|
||||
t.Fatalf("model %q not listed", id)
|
||||
return nil
|
||||
}
|
||||
|
||||
const twoModels = `{"object":"list","data":[
|
||||
{"id":"qwen3.8-27b","object":"model","meta":{"n_ctx_train":262144}},
|
||||
{"id":"gemma-3-27b","object":"model","meta":{"n_ctx_train":131072}}]}`
|
||||
|
||||
func TestListModelsRouterLoadedAndUnloaded(t *testing.T) {
|
||||
p, fs := newProvider(t, map[string]string{
|
||||
"/v1/models": twoModels,
|
||||
"/models": `{"data":[
|
||||
{"id":"qwen3.8-27b","status":{"value":"loaded","args":["llama-server","-c","262144"]}},
|
||||
{"id":"gemma-3-27b","status":{"value":"unloaded","args":["llama-server","--ctx-size","32768","-np","2"]}}]}`,
|
||||
"/props?model=qwen3.8-27b": `{"default_generation_settings":{"n_ctx":98304}}`,
|
||||
})
|
||||
models, err := p.ListModels(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Loaded: the live /props value wins over the -c argument.
|
||||
if n := ctxOf(t, models, "qwen3.8-27b"); n == nil || *n != 98304 {
|
||||
t.Errorf("loaded model context = %v, want 98304", n)
|
||||
}
|
||||
// Unloaded: parsed from the launch arguments, divided by the slot count.
|
||||
if n := ctxOf(t, models, "gemma-3-27b"); n == nil || *n != 16384 {
|
||||
t.Errorf("unloaded model context = %v, want 16384", n)
|
||||
}
|
||||
for _, path := range fs.seen {
|
||||
if path == "/props?model=gemma-3-27b" {
|
||||
t.Errorf("unloaded model must not be probed through /props")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestListModelsSwapOnlyAsksRunningModels(t *testing.T) {
|
||||
p, fs := newProvider(t, map[string]string{
|
||||
"/v1/models": twoModels,
|
||||
"/running": `{"running":[{"model":"qwen3.8-27b","state":"ready"}]}`,
|
||||
"/upstream/qwen3.8-27b/props": `{"default_generation_settings":{"n_ctx":98304}}`,
|
||||
})
|
||||
models, err := p.ListModels(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := ctxOf(t, models, "qwen3.8-27b"); n == nil || *n != 98304 {
|
||||
t.Errorf("running model context = %v, want 98304", n)
|
||||
}
|
||||
if n := ctxOf(t, models, "gemma-3-27b"); n != nil {
|
||||
t.Errorf("stopped model must keep the listing value, got %d", *n)
|
||||
}
|
||||
for _, path := range fs.seen {
|
||||
if path == "/upstream/gemma-3-27b/props" {
|
||||
t.Errorf("stopped model must not be started by a probe")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestListModelsSingleServer(t *testing.T) {
|
||||
p, _ := newProvider(t, map[string]string{
|
||||
"/v1/models": `{"object":"list","data":[{"id":"qwen3.8-27b","object":"model","meta":{"n_ctx_train":262144}}]}`,
|
||||
"/props": `{"default_generation_settings":{"n_ctx":98304},"total_slots":1}`,
|
||||
})
|
||||
models, err := p.ListModels(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := ctxOf(t, models, "qwen3.8-27b"); n == nil || *n != 98304 {
|
||||
t.Errorf("context = %v, want 98304", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListModelsProbesAreOptional(t *testing.T) {
|
||||
p, _ := newProvider(t, map[string]string{"/v1/models": twoModels})
|
||||
models, err := p.ListModels(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(models) != 2 {
|
||||
t.Fatalf("models = %d", len(models))
|
||||
}
|
||||
for _, m := range models {
|
||||
if m.ContextLength != nil {
|
||||
t.Errorf("%s: no probe answered, context must stay unset", m.ID)
|
||||
}
|
||||
var raw map[string]any
|
||||
if err := json.Unmarshal(m.Raw, &raw); err != nil || raw["meta"] == nil {
|
||||
t.Errorf("%s: raw listing entry lost", m.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestContextFromArgs(t *testing.T) {
|
||||
cases := []struct {
|
||||
args []string
|
||||
want int
|
||||
}{
|
||||
{[]string{"llama-server", "-c", "98304"}, 98304},
|
||||
{[]string{"llama-server", "--ctx-size", "65536"}, 65536},
|
||||
{[]string{"llama-server", "--ctx-size=65536"}, 65536},
|
||||
{[]string{"llama-server", "-ctx", "4096"}, 4096},
|
||||
{[]string{"llama-server", "-c", "65536", "-np", "4"}, 16384},
|
||||
{[]string{"llama-server", "-c", "65536", "--parallel=2"}, 32768},
|
||||
{[]string{"llama-server", "-m", "x.gguf"}, 0},
|
||||
{[]string{"llama-server", "-c", "nope"}, 0},
|
||||
{nil, 0},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
if got := contextFromArgs(tc.args); got != tc.want {
|
||||
t.Errorf("contextFromArgs(%v) = %d, want %d", tc.args, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRootURL(t *testing.T) {
|
||||
for in, want := range map[string]string{
|
||||
"http://h:8080/v1": "http://h:8080",
|
||||
"http://h:8080/v1/": "http://h:8080",
|
||||
"http://h:8080": "http://h:8080",
|
||||
"http://h/api/v1": "http://h/api",
|
||||
} {
|
||||
p, _ := New(providers.Config{Name: "x", BaseURL: in})
|
||||
if got := p.(*provider).RootURL(); got != want {
|
||||
t.Errorf("RootURL(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -188,3 +188,44 @@ func BearerAuth(token string) func(context.Context, *http.Request) error {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// RootURL returns the server root: the BaseURL without a trailing "/v1".
|
||||
// Servers such as llama.cpp expose non-OpenAI endpoints (/props, /models)
|
||||
// there.
|
||||
func (t *Transport) RootURL() string {
|
||||
base := strings.TrimRight(t.BaseURL, "/")
|
||||
return strings.TrimSuffix(base, "/v1")
|
||||
}
|
||||
|
||||
// GetJSON performs an authenticated GET of an absolute URL and decodes a
|
||||
// JSON body into v. It returns the HTTP status; a non-2xx status is not an
|
||||
// error so callers can treat "not found" as "not supported".
|
||||
func (t *Transport) GetJSON(ctx context.Context, rawURL string, v any) (int, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for k, vs := range t.Headers {
|
||||
for _, hv := range vs {
|
||||
req.Header.Add(k, hv)
|
||||
}
|
||||
}
|
||||
if t.Authorize != nil {
|
||||
if err := t.Authorize(ctx, req); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
resp, err := t.client().Do(req)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("provider %q: %w", t.Provider, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode >= http.StatusMultipleChoices {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 64<<10))
|
||||
return resp.StatusCode, nil
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(v); err != nil {
|
||||
return resp.StatusCode, fmt.Errorf("provider %q: decode %s: %w", t.Provider, rawURL, err)
|
||||
}
|
||||
return resp.StatusCode, nil
|
||||
}
|
||||
|
||||
@@ -95,19 +95,23 @@ type CatalogRow struct {
|
||||
CachedPricePerMTok float64
|
||||
RPMLimit *int
|
||||
TPMLimit *int
|
||||
// FallbackModelID is the admin-configured backup model, when set.
|
||||
FallbackModelID *uuid.UUID
|
||||
}
|
||||
|
||||
// ListCatalog returns every catalog entry belonging to an enabled provider.
|
||||
// The context length is the admin override when set, else the discovered
|
||||
// value.
|
||||
func (s *Store) ListCatalog(ctx context.Context) ([]CatalogRow, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT pm.id, pm.created_at, pm.provider_id, p.name, p.type, pm.slug,
|
||||
COALESCE(pm.display_name, ''),
|
||||
pm.capabilities, pm.context_length,
|
||||
pm.capabilities, COALESCE(pm.context_length_override, pm.context_length),
|
||||
pm.enabled, pm.blocked, pm.available, pm.param_overrides,
|
||||
COALESCE(pm.price_input_per_mtok, pm.disc_input_price_per_mtok, 0)::double precision,
|
||||
COALESCE(pm.price_output_per_mtok, pm.disc_output_price_per_mtok, 0)::double precision,
|
||||
COALESCE(pm.price_cached_per_mtok, 0)::double precision,
|
||||
pm.rpm_limit, pm.tpm_limit
|
||||
pm.rpm_limit, pm.tpm_limit, pm.fallback_model_id
|
||||
FROM provider_models pm
|
||||
JOIN providers p ON p.id = pm.provider_id
|
||||
WHERE p.enabled = true`)
|
||||
@@ -124,7 +128,7 @@ func (s *Store) ListCatalog(ctx context.Context) ([]CatalogRow, error) {
|
||||
&r.DisplayName, &r.Capabilities, &r.ContextLength,
|
||||
&r.Enabled, &r.Blocked, &r.Available, &r.ParamOverrides,
|
||||
&r.InputPricePerMTok, &r.OutputPricePerMTok, &r.CachedPricePerMTok,
|
||||
&r.RPMLimit, &r.TPMLimit,
|
||||
&r.RPMLimit, &r.TPMLimit, &r.FallbackModelID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ type AdminModelRow struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Blocked bool `json:"blocked"`
|
||||
ContextLength *int `json:"context_length"`
|
||||
ContextLengthOvr *int `json:"context_length_override"`
|
||||
LastSeenAt *time.Time `json:"last_seen_at"`
|
||||
DiscInputPerMTok *float64 `json:"disc_input_per_mtok"`
|
||||
DiscOutputPerMTok *float64 `json:"disc_output_per_mtok"`
|
||||
@@ -27,18 +28,24 @@ type AdminModelRow struct {
|
||||
PriceCachedPerMTok *float64 `json:"price_cached_per_mtok"`
|
||||
RPMLimit *int `json:"rpm_limit"`
|
||||
TPMLimit *int `json:"tpm_limit"`
|
||||
FallbackModelID *uuid.UUID `json:"fallback_model_id"`
|
||||
// FallbackAddress is the backup model's "<provider>/<slug>" for display.
|
||||
FallbackAddress *string `json:"fallback_address"`
|
||||
}
|
||||
|
||||
// ListModelsAdmin lists every catalog entry (no servability filter) for admin.
|
||||
func (s *Store) ListModelsAdmin(ctx context.Context) ([]AdminModelRow, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT pm.id, p.name, pm.slug, COALESCE(pm.display_name,''), pm.source,
|
||||
pm.available, pm.enabled, pm.blocked, pm.context_length, pm.last_seen_at,
|
||||
pm.available, pm.enabled, pm.blocked, pm.context_length, pm.context_length_override, pm.last_seen_at,
|
||||
pm.disc_input_price_per_mtok, pm.disc_output_price_per_mtok,
|
||||
pm.price_input_per_mtok, pm.price_output_per_mtok, pm.price_cached_per_mtok,
|
||||
pm.rpm_limit, pm.tpm_limit
|
||||
pm.rpm_limit, pm.tpm_limit,
|
||||
pm.fallback_model_id, fp.name || '/' || fb.slug
|
||||
FROM provider_models pm
|
||||
JOIN providers p ON p.id = pm.provider_id
|
||||
LEFT JOIN provider_models fb ON fb.id = pm.fallback_model_id
|
||||
LEFT JOIN providers fp ON fp.id = fb.provider_id
|
||||
ORDER BY p.name, pm.slug`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -48,10 +55,10 @@ func (s *Store) ListModelsAdmin(ctx context.Context) ([]AdminModelRow, error) {
|
||||
for rows.Next() {
|
||||
var m AdminModelRow
|
||||
if err := rows.Scan(&m.ID, &m.ProviderName, &m.Slug, &m.DisplayName, &m.Source,
|
||||
&m.Available, &m.Enabled, &m.Blocked, &m.ContextLength, &m.LastSeenAt,
|
||||
&m.Available, &m.Enabled, &m.Blocked, &m.ContextLength, &m.ContextLengthOvr, &m.LastSeenAt,
|
||||
&m.DiscInputPerMTok, &m.DiscOutputPerMTok,
|
||||
&m.PriceInputPerMTok, &m.PriceOutputPerMTok, &m.PriceCachedPerMTok,
|
||||
&m.RPMLimit, &m.TPMLimit); err != nil {
|
||||
&m.RPMLimit, &m.TPMLimit, &m.FallbackModelID, &m.FallbackAddress); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, m)
|
||||
@@ -76,6 +83,10 @@ type ModelOverride struct {
|
||||
CachedPrice Patch[float64]
|
||||
RPMLimit Patch[int]
|
||||
TPMLimit Patch[int]
|
||||
// ContextLength overrides the provider-reported context window.
|
||||
ContextLength Patch[int]
|
||||
// FallbackModelID is the backup model served when this one fails.
|
||||
FallbackModelID Patch[uuid.UUID]
|
||||
}
|
||||
|
||||
func (s *Store) UpdateModelOverride(ctx context.Context, id uuid.UUID, o ModelOverride) error {
|
||||
@@ -92,7 +103,9 @@ func (s *Store) UpdateModelOverride(ctx context.Context, id uuid.UUID, o ModelOv
|
||||
price_output_per_mtok = CASE WHEN $7::boolean THEN $8 ELSE price_output_per_mtok END,
|
||||
price_cached_per_mtok = CASE WHEN $9::boolean THEN $10 ELSE price_cached_per_mtok END,
|
||||
rpm_limit = CASE WHEN $11::boolean THEN $12 ELSE rpm_limit END,
|
||||
tpm_limit = CASE WHEN $13::boolean THEN $14 ELSE tpm_limit END
|
||||
tpm_limit = CASE WHEN $13::boolean THEN $14 ELSE tpm_limit END,
|
||||
context_length_override = CASE WHEN $15::boolean THEN $16 ELSE context_length_override END,
|
||||
fallback_model_id = CASE WHEN $17::boolean THEN $18 ELSE fallback_model_id END
|
||||
WHERE id = $1`,
|
||||
id, o.Enabled, o.Blocked, params,
|
||||
o.InputPrice.Set, o.InputPrice.Value,
|
||||
@@ -100,6 +113,8 @@ func (s *Store) UpdateModelOverride(ctx context.Context, id uuid.UUID, o ModelOv
|
||||
o.CachedPrice.Set, o.CachedPrice.Value,
|
||||
o.RPMLimit.Set, o.RPMLimit.Value,
|
||||
o.TPMLimit.Set, o.TPMLimit.Value,
|
||||
o.ContextLength.Set, o.ContextLength.Value,
|
||||
o.FallbackModelID.Set, o.FallbackModelID.Value,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ type RequestLog struct {
|
||||
ProviderID *uuid.UUID
|
||||
KeyPrefix string
|
||||
ModelName string // the address the client requested
|
||||
ServedModel string // the address that served it, when a backup model did
|
||||
ProviderName string
|
||||
StatusCode int
|
||||
Streamed bool
|
||||
@@ -43,14 +44,15 @@ func (s *Store) InsertRequestLog(ctx context.Context, e RequestLog) (uuid.UUID,
|
||||
key_prefix, model_name, provider_name,
|
||||
status_code, streamed, finish_reason, error,
|
||||
prompt_tokens, completion_tokens, cached_tokens, total_tokens,
|
||||
cost, latency_ms, ttft_ms, request_id, client_ip
|
||||
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20)
|
||||
cost, latency_ms, ttft_ms, request_id, client_ip, served_model
|
||||
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21)
|
||||
RETURNING id`,
|
||||
e.VirtualKeyID, e.UserID, e.ProviderModelID, e.ProviderID,
|
||||
nullString(e.KeyPrefix), nullString(e.ModelName), nullString(e.ProviderName),
|
||||
e.StatusCode, e.Streamed, nullString(e.FinishReason), nullString(e.Error),
|
||||
e.PromptTokens, e.CompletionTokens, e.CachedTokens, e.TotalTokens,
|
||||
e.Cost, e.LatencyMS, e.TTFTMs, nullString(e.RequestID), nullInet(e.ClientIP),
|
||||
nullString(e.ServedModel),
|
||||
).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
@@ -76,6 +78,7 @@ type RequestLogListRow struct {
|
||||
RequestID string
|
||||
KeyPrefix string
|
||||
ModelName string
|
||||
ServedModel string // set when a backup model served the request
|
||||
ProviderName string
|
||||
StatusCode int
|
||||
Streamed bool
|
||||
@@ -132,7 +135,8 @@ func (s *Store) ListRequestLogs(ctx context.Context, p ListLogsParams) ([]Reques
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT rl.id::text, rl.created_at, COALESCE(rl.request_id,''),
|
||||
COALESCE(rl.key_prefix,''), COALESCE(rl.model_name,''), COALESCE(rl.provider_name,''),
|
||||
COALESCE(rl.key_prefix,''), COALESCE(rl.model_name,''), COALESCE(rl.served_model,''),
|
||||
COALESCE(rl.provider_name,''),
|
||||
rl.status_code, rl.streamed, COALESCE(rl.error,''),
|
||||
rl.prompt_tokens, rl.completion_tokens, rl.cached_tokens, rl.total_tokens,
|
||||
rl.cost, rl.latency_ms, rl.ttft_ms
|
||||
@@ -145,6 +149,7 @@ func (s *Store) ListRequestLogs(ctx context.Context, p ListLogsParams) ([]Reques
|
||||
AND ($6::timestamptz IS NULL OR rl.created_at >= $6)
|
||||
AND ($7::timestamptz IS NULL OR (rl.created_at, rl.id) < ($7, $8))
|
||||
AND ($9 = '' OR rl.model_name ILIKE '%' || $9 || '%'
|
||||
OR rl.served_model ILIKE '%' || $9 || '%'
|
||||
OR rl.provider_name ILIKE '%' || $9 || '%'
|
||||
OR rl.key_prefix ILIKE '%' || $9 || '%'
|
||||
OR rl.request_id ILIKE '%' || $9 || '%')
|
||||
@@ -161,7 +166,7 @@ func (s *Store) ListRequestLogs(ctx context.Context, p ListLogsParams) ([]Reques
|
||||
var r RequestLogListRow
|
||||
if err := rows.Scan(
|
||||
&r.ID, &r.CreatedAt, &r.RequestID,
|
||||
&r.KeyPrefix, &r.ModelName, &r.ProviderName,
|
||||
&r.KeyPrefix, &r.ModelName, &r.ServedModel, &r.ProviderName,
|
||||
&r.StatusCode, &r.Streamed, &r.Error,
|
||||
&r.PromptTokens, &r.CompletionTokens, &r.CachedTokens, &r.TotalTokens,
|
||||
&r.Cost, &r.LatencyMS, &r.TTFTMs,
|
||||
|
||||
@@ -119,14 +119,17 @@ func TestStoreQueries(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("ListModelsAdmin: %v", err)
|
||||
}
|
||||
var m1 uuid.UUID
|
||||
var m1, m2 uuid.UUID
|
||||
for _, r := range rows {
|
||||
if r.ProviderName == "local" && r.Slug == "m1" {
|
||||
m1 = r.ID
|
||||
}
|
||||
if r.ProviderName == "local" && r.Slug == "m2" {
|
||||
m2 = r.ID
|
||||
}
|
||||
}
|
||||
if m1 == uuid.Nil {
|
||||
t.Fatalf("local/m1 missing from %+v", rows)
|
||||
if m1 == uuid.Nil || m2 == uuid.Nil {
|
||||
t.Fatalf("local/m1 or local/m2 missing from %+v", rows)
|
||||
}
|
||||
price := 2.5
|
||||
rpm := 30
|
||||
@@ -157,10 +160,55 @@ func TestStoreQueries(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Backup model: set, read back through the admin list and the catalog,
|
||||
// reject an unknown target, then clear.
|
||||
if err := st.UpdateModelOverride(ctx, m1, store.ModelOverride{FallbackModelID: store.Patch[uuid.UUID]{Set: true, Value: &m2}}); err != nil {
|
||||
t.Fatalf("UpdateModelOverride fallback: %v", err)
|
||||
}
|
||||
rows, _ = st.ListModelsAdmin(ctx)
|
||||
for _, r := range rows {
|
||||
if r.ID == m1 && (r.FallbackModelID == nil || *r.FallbackModelID != m2 || r.FallbackAddress == nil || *r.FallbackAddress != "local/m2") {
|
||||
t.Fatalf("fallback not applied: %+v", r)
|
||||
}
|
||||
}
|
||||
on := true
|
||||
if err := st.UpdateProvider(ctx, provID, store.UpdateProviderParams{Enabled: &on}); err != nil {
|
||||
t.Fatalf("UpdateProvider enable: %v", err)
|
||||
}
|
||||
cat, err := st.ListCatalog(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("ListCatalog: %v", err)
|
||||
}
|
||||
var sawM1 bool
|
||||
for _, r := range cat {
|
||||
if r.ID == m1 {
|
||||
sawM1 = true
|
||||
if r.FallbackModelID == nil || *r.FallbackModelID != m2 {
|
||||
t.Fatalf("catalog row lacks fallback: %+v", r)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sawM1 {
|
||||
t.Fatalf("local/m1 missing from catalog %+v", cat)
|
||||
}
|
||||
unknown := uuid.New()
|
||||
if err := st.UpdateModelOverride(ctx, m1, store.ModelOverride{FallbackModelID: store.Patch[uuid.UUID]{Set: true, Value: &unknown}}); err == nil {
|
||||
t.Fatal("UpdateModelOverride must reject an unknown backup model")
|
||||
}
|
||||
if err := st.UpdateModelOverride(ctx, m1, store.ModelOverride{FallbackModelID: store.Patch[uuid.UUID]{Set: true}}); err != nil {
|
||||
t.Fatalf("UpdateModelOverride clear fallback: %v", err)
|
||||
}
|
||||
rows, _ = st.ListModelsAdmin(ctx)
|
||||
for _, r := range rows {
|
||||
if r.ID == m1 && (r.FallbackModelID != nil || r.FallbackAddress != nil) {
|
||||
t.Fatalf("fallback not cleared: %+v", r)
|
||||
}
|
||||
}
|
||||
|
||||
// Logs: insert, then list and summarize.
|
||||
logID, err := st.InsertRequestLog(ctx, store.RequestLog{
|
||||
VirtualKeyID: &keyID, UserID: &user.ID, ProviderID: &provID, ProviderModelID: &m1,
|
||||
KeyPrefix: "sk-00000001", ModelName: "local/m1", ProviderName: "local",
|
||||
VirtualKeyID: &keyID, UserID: &user.ID, ProviderID: &provID, ProviderModelID: &m2,
|
||||
KeyPrefix: "sk-00000001", ModelName: "local/m1", ServedModel: "local/m2", ProviderName: "local",
|
||||
StatusCode: 502, Error: "upstream down", LatencyMS: 12, ClientIP: "127.0.0.1",
|
||||
})
|
||||
if err != nil {
|
||||
@@ -170,9 +218,13 @@ func TestStoreQueries(t *testing.T) {
|
||||
t.Fatalf("InsertLogContent: %v", err)
|
||||
}
|
||||
list, err := st.ListRequestLogs(ctx, store.ListLogsParams{UserID: &user.ID, Query: "local/m1", StatusGte: 400})
|
||||
if err != nil || len(list) != 1 || list[0].Error != "upstream down" {
|
||||
if err != nil || len(list) != 1 || list[0].Error != "upstream down" || list[0].ServedModel != "local/m2" {
|
||||
t.Fatalf("ListRequestLogs: %v %+v", err, list)
|
||||
}
|
||||
// The search also matches the address that served the request.
|
||||
if list, err := st.ListRequestLogs(ctx, store.ListLogsParams{UserID: &user.ID, Query: "local/m2"}); err != nil || len(list) != 1 {
|
||||
t.Fatalf("ListRequestLogs by served model: %v %+v", err, list)
|
||||
}
|
||||
sum, err = st.GetLogsSummary(ctx, time.Now().Add(-time.Hour), &user.ID)
|
||||
if err != nil || sum.Requests != 1 || sum.Errors != 1 {
|
||||
t.Fatalf("summary = %+v err=%v", sum, err)
|
||||
|
||||
0
backend/internal/web/dist/.gitkeep
vendored
Normal file
0
backend/internal/web/dist/.gitkeep
vendored
Normal file
@@ -0,0 +1 @@
|
||||
ALTER TABLE provider_models DROP COLUMN IF EXISTS context_length_override;
|
||||
4
backend/migrations/000005_context_length_override.up.sql
Normal file
4
backend/migrations/000005_context_length_override.up.sql
Normal file
@@ -0,0 +1,4 @@
|
||||
-- Admin override for the context window. Providers report the training
|
||||
-- context (llama.cpp meta.n_ctx_train) or a catalog value, which can differ
|
||||
-- from what the running instance accepts. NULL means use the discovered value.
|
||||
ALTER TABLE provider_models ADD COLUMN context_length_override integer;
|
||||
2
backend/migrations/000006_fallback_model.down.sql
Normal file
2
backend/migrations/000006_fallback_model.down.sql
Normal file
@@ -0,0 +1,2 @@
|
||||
ALTER TABLE request_logs DROP COLUMN IF EXISTS served_model;
|
||||
ALTER TABLE provider_models DROP COLUMN IF EXISTS fallback_model_id;
|
||||
9
backend/migrations/000006_fallback_model.up.sql
Normal file
9
backend/migrations/000006_fallback_model.up.sql
Normal file
@@ -0,0 +1,9 @@
|
||||
-- Admin-configured backup model. When the model is unavailable or its upstream
|
||||
-- fails, the gateway serves the request with the backup instead. NULL means
|
||||
-- there is no backup. Deleting the backup clears the reference.
|
||||
ALTER TABLE provider_models
|
||||
ADD COLUMN fallback_model_id uuid REFERENCES provider_models (id) ON DELETE SET NULL;
|
||||
|
||||
-- The address that served a request when it differs from model_name, which is
|
||||
-- the address the client requested. NULL when the requested model served it.
|
||||
ALTER TABLE request_logs ADD COLUMN served_model text;
|
||||
@@ -47,6 +47,7 @@ internal/providers provider abstraction and the per-type packages
|
||||
internal/catalog model library: sync, cache, address resolution
|
||||
internal/ratelimit per-model fixed-window RPM and TPM counters
|
||||
internal/billing cost calculation
|
||||
internal/influx InfluxDB 3 exporter: one point per request, batched
|
||||
internal/api chi routes: /v1 proxy, /api/admin, /auth
|
||||
internal/oidcauth OIDC discovery, ID token verification, group check
|
||||
internal/web go:embed of the built dashboard, served as a SPA
|
||||
@@ -101,13 +102,44 @@ name, context length, pricing, and output modalities. Output modalities become
|
||||
the row's capabilities, which is how the image route knows which models accept
|
||||
an image request.
|
||||
|
||||
The llama.cpp provider goes one step further. The OpenAI listing only carries
|
||||
the training context, so the provider asks the server for the running size:
|
||||
`/props` on a single server, `/props?model=` or the launch arguments on the
|
||||
router, and `/upstream/<model>/props` for models llama-swap has loaded. It
|
||||
never asks llama-swap about a stopped model, because that would start it.
|
||||
|
||||
Admin overrides live in separate columns on the same row and survive syncs:
|
||||
enabled, blocked, parameter overrides, prices, and rate limits. A model that
|
||||
disappears from the upstream list is marked unavailable, not deleted.
|
||||
enabled, blocked, parameter overrides, prices, context window, rate limits,
|
||||
and a backup model. A model that disappears from the upstream list is marked
|
||||
unavailable, not deleted.
|
||||
|
||||
The cache holds every servable row in memory, keyed by address. It is rebuilt
|
||||
after each sync, so the request path never queries the database for routing.
|
||||
|
||||
### Backup models
|
||||
|
||||
A model's `fallback_model_id` names another catalog row. The cache turns an
|
||||
address into a route: the requested model first, then its backups in chain
|
||||
order, with models that are not servable skipped. Two models may back each
|
||||
other up; the walk tracks visited ids and stops after a fixed number of hops.
|
||||
A model that is enabled but unavailable still routes, straight to its backups,
|
||||
and still appears in `/v1/models`.
|
||||
|
||||
The proxy handlers walk the route. A target is skipped when the rate limiter
|
||||
refuses it, its provider is missing, the transport fails, or the upstream
|
||||
answers 408, 429, or 5xx. Client errors (4xx) are relayed, because the backup
|
||||
would most likely answer the same. For a stream, the first chunk is read
|
||||
before the status line goes to the client, so a stream that dies at once can
|
||||
still move to a backup. After the first byte is sent, a failure can only be
|
||||
logged.
|
||||
|
||||
The switch is invisible to the client and visible in the log. A backup's
|
||||
response reports the requested model's id. The key allowlist is checked
|
||||
against the requested address. Every failed attempt gets its own log row, with
|
||||
the error and a note that names the backup. The row for the backup keeps the
|
||||
requested address in `model_name`, puts the backup in `served_model`, and
|
||||
uses the backup's prices for the cost.
|
||||
|
||||
## Authentication
|
||||
|
||||
Two systems that do not mix:
|
||||
|
||||
Reference in New Issue
Block a user