LovroG05
8bb62a1c20
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
Any model can name another model as its backup. The gateway tries the backup when the model is offline, rate limited, or its upstream fails before the client has received anything. The response reports the requested model, so the client sees no switch. The request log records each failed attempt and the address that served the request, and the cost uses the backup's prices.
355 lines
12 KiB
Go
355 lines
12 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/netip"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
// RequestLog is one row to append to request_logs (billing source of truth).
|
|
type RequestLog struct {
|
|
VirtualKeyID *uuid.UUID
|
|
UserID *uuid.UUID // owner of the key used (set even for synthetic dashboard keys)
|
|
ProviderModelID *uuid.UUID
|
|
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
|
|
FinishReason string
|
|
Error string
|
|
PromptTokens int
|
|
CompletionTokens int
|
|
CachedTokens int
|
|
TotalTokens int
|
|
Cost float64
|
|
LatencyMS int
|
|
TTFTMs *int
|
|
RequestID string
|
|
ClientIP string
|
|
}
|
|
|
|
func (s *Store) InsertRequestLog(ctx context.Context, e RequestLog) (uuid.UUID, error) {
|
|
var id uuid.UUID
|
|
err := s.pool.QueryRow(ctx, `
|
|
INSERT INTO request_logs (
|
|
virtual_key_id, user_id, provider_model_id, provider_id,
|
|
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, 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
|
|
}
|
|
|
|
// InsertLogContent saves the prompt messages and completion text for a log row.
|
|
func (s *Store) InsertLogContent(ctx context.Context, logID uuid.UUID, messages json.RawMessage, completion string) error {
|
|
if len(messages) == 0 {
|
|
messages = json.RawMessage("[]")
|
|
}
|
|
_, err := s.pool.Exec(ctx, `
|
|
INSERT INTO request_log_content (request_log_id, messages, completion)
|
|
VALUES ($1, $2, $3)
|
|
ON CONFLICT (request_log_id) DO UPDATE SET messages = EXCLUDED.messages, completion = EXCLUDED.completion`,
|
|
logID, messages, completion,
|
|
)
|
|
return err
|
|
}
|
|
|
|
// RequestLogListRow is one row returned from the admin logs list query.
|
|
type RequestLogListRow struct {
|
|
ID string
|
|
CreatedAt time.Time
|
|
RequestID string
|
|
KeyPrefix string
|
|
ModelName string
|
|
ServedModel string // set when a backup model served the request
|
|
ProviderName string
|
|
StatusCode int
|
|
Streamed bool
|
|
Error string
|
|
PromptTokens int
|
|
CompletionTokens int
|
|
CachedTokens int
|
|
TotalTokens int
|
|
Cost float64
|
|
LatencyMS int
|
|
TTFTMs *int
|
|
// Content fields — null when content logging is disabled or unavailable.
|
|
Messages *json.RawMessage `json:"Messages,omitempty"`
|
|
Completion *string `json:"Completion,omitempty"`
|
|
}
|
|
|
|
// ListLogsParams filters the request_logs query.
|
|
type ListLogsParams struct {
|
|
Provider string // empty = all
|
|
KeyPrefix string // empty = all
|
|
UserID *uuid.UUID // non-nil = only logs for keys owned by this user
|
|
StatusGte int // 0 = all; e.g. 500 filters ≥500
|
|
StatusLt int // 0 = no upper bound
|
|
Since time.Time
|
|
// Query is a free-text search matched (case-insensitively, substring) against
|
|
// model_name, provider_name, key_prefix, and request_id. Empty = no filter.
|
|
Query string
|
|
Limit int
|
|
// Before/BeforeID page backwards from a cursor: only rows strictly older
|
|
// than (Before, BeforeID) in the (created_at DESC, id DESC) ordering are
|
|
// returned. Pass the last row of the previous page. Zero value = first page.
|
|
// Keyset pagination instead of OFFSET so paging deep into the log stays
|
|
// O(limit) rather than re-scanning and discarding every skipped row.
|
|
Before time.Time
|
|
BeforeID uuid.UUID
|
|
}
|
|
|
|
// ListRequestLogs returns recent request logs for the admin view. Content
|
|
// (prompt/completion) is fetched separately per row on demand — joining it
|
|
// here would pull every logged conversation's full text on every page load.
|
|
func (s *Store) ListRequestLogs(ctx context.Context, p ListLogsParams) ([]RequestLogListRow, error) {
|
|
if p.Limit <= 0 {
|
|
p.Limit = 50
|
|
}
|
|
var since *time.Time
|
|
if !p.Since.IsZero() {
|
|
since = &p.Since
|
|
}
|
|
var before *time.Time
|
|
var beforeID *uuid.UUID
|
|
if !p.Before.IsZero() {
|
|
before = &p.Before
|
|
beforeID = &p.BeforeID
|
|
}
|
|
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.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
|
|
FROM request_logs rl
|
|
WHERE ($1 = '' OR rl.provider_name = $1)
|
|
AND ($2 = '' OR rl.key_prefix = $2)
|
|
AND ($3::uuid IS NULL OR rl.user_id = $3)
|
|
AND ($4 = 0 OR rl.status_code >= $4)
|
|
AND ($5 = 0 OR rl.status_code < $5)
|
|
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 || '%')
|
|
ORDER BY rl.created_at DESC, rl.id DESC
|
|
LIMIT $10`,
|
|
p.Provider, p.KeyPrefix, p.UserID, p.StatusGte, p.StatusLt, since, before, beforeID, p.Query, p.Limit,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []RequestLogListRow
|
|
for rows.Next() {
|
|
var r RequestLogListRow
|
|
if err := rows.Scan(
|
|
&r.ID, &r.CreatedAt, &r.RequestID,
|
|
&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,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, r)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
// GetLogContent returns the prompt messages and completion text for a single
|
|
// request log row, or (nil, nil, nil) if none was recorded or the log is not
|
|
// visible to userID (non-nil userID restricts to that user's own logs, same
|
|
// scoping as ListRequestLogs).
|
|
func (s *Store) GetLogContent(ctx context.Context, logID uuid.UUID, userID *uuid.UUID) (*json.RawMessage, *string, error) {
|
|
var messages *json.RawMessage
|
|
var completion *string
|
|
err := s.pool.QueryRow(ctx, `
|
|
SELECT rlc.messages, rlc.completion
|
|
FROM request_log_content rlc
|
|
JOIN request_logs rl ON rl.id = rlc.request_log_id
|
|
WHERE rlc.request_log_id = $1
|
|
AND ($2::uuid IS NULL OR rl.user_id = $2)`,
|
|
logID, userID,
|
|
).Scan(&messages, &completion)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
return messages, completion, nil
|
|
}
|
|
|
|
// KeyStat is per-virtual-key aggregate stats for a time window.
|
|
type KeyStat struct {
|
|
KeyPrefix string `json:"key_prefix"`
|
|
Requests int `json:"requests"`
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
CompletionTokens int `json:"completion_tokens"`
|
|
TotalTokens int `json:"total_tokens"`
|
|
Cost float64 `json:"cost"`
|
|
Errors int `json:"errors"`
|
|
}
|
|
|
|
// KeyStats returns per-key aggregate stats since the given time. When userID is
|
|
// non-nil only that user's keys are included.
|
|
func (s *Store) KeyStats(ctx context.Context, since time.Time, userID *uuid.UUID) ([]KeyStat, error) {
|
|
rows, err := s.pool.Query(ctx, `
|
|
SELECT COALESCE(key_prefix, '(no key)'),
|
|
COUNT(*)::int,
|
|
SUM(prompt_tokens)::int,
|
|
SUM(completion_tokens)::int,
|
|
SUM(total_tokens)::int,
|
|
SUM(cost),
|
|
SUM(CASE WHEN status_code >= 400 THEN 1 ELSE 0 END)::int
|
|
FROM request_logs
|
|
WHERE created_at >= $1
|
|
AND ($2::uuid IS NULL OR user_id = $2)
|
|
GROUP BY key_prefix
|
|
ORDER BY SUM(total_tokens) DESC`,
|
|
since, userID,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []KeyStat
|
|
for rows.Next() {
|
|
var s KeyStat
|
|
if err := rows.Scan(&s.KeyPrefix, &s.Requests, &s.PromptTokens, &s.CompletionTokens, &s.TotalTokens, &s.Cost, &s.Errors); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, s)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
// LogsSummary is aggregate totals for a time window (dashboard stat cards).
|
|
type LogsSummary struct {
|
|
Requests int `json:"requests"`
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
TotalCost float64 `json:"total_cost"`
|
|
AvgLatencyMS int `json:"avg_latency_ms"`
|
|
Errors int `json:"errors"`
|
|
}
|
|
|
|
// GetLogsSummary computes aggregate totals since the given time in a single
|
|
// query. Every aggregate is coalesced so an empty window scans as zeros.
|
|
func (s *Store) GetLogsSummary(ctx context.Context, since time.Time, userID *uuid.UUID) (LogsSummary, error) {
|
|
var out LogsSummary
|
|
err := s.pool.QueryRow(ctx, `
|
|
SELECT COUNT(*)::int,
|
|
COALESCE(SUM(total_tokens),0)::bigint,
|
|
COALESCE(SUM(cost),0),
|
|
COALESCE(AVG(latency_ms),0)::int,
|
|
COALESCE(SUM(CASE WHEN status_code >= 400 THEN 1 ELSE 0 END),0)::int
|
|
FROM request_logs
|
|
WHERE created_at >= $1
|
|
AND ($2::uuid IS NULL OR user_id = $2)`,
|
|
since, userID,
|
|
).Scan(&out.Requests, &out.TotalTokens, &out.TotalCost, &out.AvgLatencyMS, &out.Errors)
|
|
return out, err
|
|
}
|
|
|
|
// ModelStat is token usage for one model address, for the "top models" chart.
|
|
type ModelStat struct {
|
|
ModelName string `json:"model_name"`
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
}
|
|
|
|
// TopModels returns the top models by token usage since the given time.
|
|
func (s *Store) TopModels(ctx context.Context, since time.Time, userID *uuid.UUID, limit int) ([]ModelStat, error) {
|
|
if limit <= 0 {
|
|
limit = 6
|
|
}
|
|
rows, err := s.pool.Query(ctx, `
|
|
SELECT COALESCE(model_name,''), COALESCE(SUM(total_tokens),0)::bigint
|
|
FROM request_logs
|
|
WHERE created_at >= $1
|
|
AND ($2::uuid IS NULL OR user_id = $2)
|
|
GROUP BY model_name
|
|
ORDER BY SUM(total_tokens) DESC
|
|
LIMIT $3`,
|
|
since, userID, limit,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []ModelStat
|
|
for rows.Next() {
|
|
var m ModelStat
|
|
if err := rows.Scan(&m.ModelName, &m.TotalTokens); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, m)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
// RequestVolumeSeries buckets request counts evenly between since and now
|
|
// into `buckets` buckets (oldest first), for the dashboard sparkline.
|
|
func (s *Store) RequestVolumeSeries(ctx context.Context, since time.Time, userID *uuid.UUID, buckets int) ([]int, error) {
|
|
if buckets <= 0 {
|
|
buckets = 24
|
|
}
|
|
rows, err := s.pool.Query(ctx, `
|
|
SELECT width_bucket(extract(epoch from rl.created_at), extract(epoch from $1::timestamptz), extract(epoch from now()), $3) AS bucket,
|
|
COUNT(*)::int
|
|
FROM request_logs rl
|
|
WHERE rl.created_at >= $1
|
|
AND ($2::uuid IS NULL OR rl.user_id = $2)
|
|
GROUP BY bucket`,
|
|
since, userID, buckets,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
out := make([]int, buckets)
|
|
for rows.Next() {
|
|
var bucket, count int
|
|
if err := rows.Scan(&bucket, &count); err != nil {
|
|
return nil, err
|
|
}
|
|
// width_bucket returns 1..buckets for in-range values; clamp the rare
|
|
// edge case (e.g. a row created exactly at `now()`) into the last bucket.
|
|
if bucket < 1 {
|
|
bucket = 1
|
|
} else if bucket > buckets {
|
|
bucket = buckets
|
|
}
|
|
out[bucket-1] += count
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func nullInet(s string) any {
|
|
if a, err := netip.ParseAddr(s); err == nil {
|
|
return a
|
|
}
|
|
return nil
|
|
}
|