Add token usage stats and consolidate credentials in SQLite
build / build (push) Successful in 2m34s
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:
@@ -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(),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user