Add token usage stats and consolidate credentials in SQLite
build / build (push) Successful in 2m34s

- New internal/stats (types) and internal/store (SQLite owner: requests + credentials tables, WAL); store implements stats.Recorder.
- Stream (SSE tee) and non-stream chat paths parse upstream usage incl. cached_tokens and record per-request; add /api/stats, /api/stats/reset, /admin/stats HTML with cache hit rate.
- Drop credentials.json: remove auth file I/O and ZHANLU_CREDENTIALS_FILE; credential precedence is env vars > db row.
This commit is contained in:
2026-08-19 15:52:00 +08:00
parent dd5c560b64
commit fa9640d919
15 changed files with 1101 additions and 103 deletions
+83
View File
@@ -0,0 +1,83 @@
// Package stats defines the types and helpers for OpenAI-compatible token
// usage statistics recorded by the zhanlu proxy. The concrete SQLite-backed
// recorder lives in internal/store; this package is dependency-free so it can
// be referenced by both store and server without import cycles.
package stats
import "time"
// Record is a single chat-completion usage observation.
type Record struct {
Ts time.Time `json:"ts"`
Model string `json:"model"`
Stream bool `json:"stream"`
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
ReasoningTokens int `json:"reasoning_tokens"`
CachedTokens int `json:"cached_tokens"`
Status string `json:"status"` // "success" | "upstream_error"
LatencyMs int64 `json:"latency_ms"`
}
// Query filters the recorded stats. Zero-value time fields mean unbounded on
// that end; empty Model means all models. Limit caps the recent-records list
// (0 = default). Stream filters by streaming mode (nil = both).
type Query struct {
Since time.Time
Until time.Time
Model string
Limit int
Stream *bool
}
// Totals aggregates request counts and token sums over a filtered set.
type Totals struct {
Requests int `json:"requests"`
SuccessRequests int `json:"success_requests"`
ErrorRequests int `json:"error_requests"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
TotalTokens int64 `json:"total_tokens"`
ReasoningTokens int64 `json:"reasoning_tokens"`
CachedTokens int64 `json:"cached_tokens"`
CacheRate float64 `json:"cache_rate"` // cached_tokens / prompt_tokens, 0..1
}
// ModelStat is a per-model aggregation row.
type ModelStat struct {
Model string `json:"model"`
Requests int `json:"requests"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
TotalTokens int64 `json:"total_tokens"`
CachedTokens int64 `json:"cached_tokens"`
}
// DayStat is a per-day aggregation row (server-local time, YYYY-MM-DD).
type DayStat struct {
Day string `json:"day"`
Requests int `json:"requests"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
TotalTokens int64 `json:"total_tokens"`
CachedTokens int64 `json:"cached_tokens"`
}
// Summary is the full result returned by a Recorder's Stats query.
type Summary struct {
Totals Totals `json:"totals"`
PerModel []ModelStat `json:"per_model"`
Daily []DayStat `json:"daily"`
Recent []Record `json:"recent"`
}
// Recorder persists and queries usage statistics. The concrete implementation
// lives in internal/store; the interface is declared here so server code can
// depend on the contract and tests can inject fakes.
type Recorder interface {
Record(r Record) error
Stats(q Query) (*Summary, error)
Reset() error
Close() error
}
+78
View File
@@ -0,0 +1,78 @@
package stats
import (
"encoding/json"
"time"
)
// Usage is the subset of the OpenAI chat-completion usage object the recorder
// persists. Numbers arrive from JSON unmarshal as float64.
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
ReasoningTokens int `json:"reasoning_tokens"`
CachedTokens int `json:"cached_tokens"`
// PromptTokensDetails.CachedTokens is emitted by providers that support
// prompt caching (OpenAI/DeepSeek/Zhipu litellm gateways). Some upstreams
// put cached_tokens at the top level instead.
PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
// CompletionTokensDetails.ReasoningTokens is emitted by reasoning models;
// some upstreams put reasoning_tokens at top level instead.
CompletionTokensDetails struct {
ReasoningTokens int `json:"reasoning_tokens"`
} `json:"completion_tokens_details"`
}
// ExtractUsage decodes a raw usage value (as produced by encoding/json into an
// any) into token counts. It accepts both full usage maps and raw JSON bytes.
// Missing fields default to 0; a nil v yields zero usage.
func ExtractUsage(v any) Usage {
var u Usage
if v == nil {
return u
}
switch t := v.(type) {
case []byte:
_ = json.Unmarshal(t, &u)
case json.RawMessage:
_ = json.Unmarshal(t, &u)
case map[string]any:
// Re-marshal + unmarshal is the simplest robust path for nested
// *_tokens_details; usage payloads are tiny.
if b, err := json.Marshal(t); err == nil {
_ = json.Unmarshal(b, &u)
}
}
if u.ReasoningTokens == 0 {
u.ReasoningTokens = u.CompletionTokensDetails.ReasoningTokens
}
if u.CachedTokens == 0 {
u.CachedTokens = u.PromptTokensDetails.CachedTokens
}
return u
}
// RecordFromUsage builds a Record from a captured usage value plus context.
func RecordFromUsage(model string, stream bool, usage any, status string, start time.Time) Record {
u := ExtractUsage(usage)
if u.TotalTokens == 0 && (u.PromptTokens != 0 || u.CompletionTokens != 0) {
u.TotalTokens = u.PromptTokens + u.CompletionTokens
}
return Record{
Ts: time.Now(),
Model: model,
Stream: stream,
PromptTokens: u.PromptTokens,
CompletionTokens: u.CompletionTokens,
TotalTokens: u.TotalTokens,
ReasoningTokens: u.ReasoningTokens,
CachedTokens: u.CachedTokens,
Status: status,
LatencyMs: time.Since(start).Milliseconds(),
}
}