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,89 @@
|
||||
// Package store owns the single embedded SQLite database backing the zhanlu
|
||||
// proxy: token-usage records (the requests table) and the persisted login
|
||||
// credentials (the credentials table, single-tenant single row). It implements
|
||||
// stats.Recorder for usage tracking and exposes Load/Save for credentials.
|
||||
package store
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
// Store is the single owner of the proxy's SQLite database handle.
|
||||
type Store struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
// Open opens (or creates) the database at path and ensures both tables exist.
|
||||
// SQLite is opened with WAL journaling and a busy timeout so concurrent reads
|
||||
// (stats queries) and writes (request records, credential saves) do not
|
||||
// collide.
|
||||
func Open(path string) (*Store, error) {
|
||||
db, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)&_pragma=foreign_keys(ON)")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open db %q: %w", path, err)
|
||||
}
|
||||
if err := ensureSchema(db); err != nil {
|
||||
db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &Store{db: db}, nil
|
||||
}
|
||||
|
||||
// Close releases the database handle.
|
||||
func (s *Store) Close() error {
|
||||
return s.db.Close()
|
||||
}
|
||||
|
||||
func ensureSchema(db *sql.DB) error {
|
||||
stmts := []string{
|
||||
`CREATE TABLE IF NOT EXISTS requests (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
ts INTEGER NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
stream INTEGER NOT NULL DEFAULT 0,
|
||||
prompt_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
completion_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cached_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'success',
|
||||
latency_ms INTEGER NOT NULL DEFAULT 0
|
||||
)`,
|
||||
`CREATE TABLE IF NOT EXISTS credentials (
|
||||
id INTEGER PRIMARY KEY CHECK (id = 1),
|
||||
access_key TEXT NOT NULL DEFAULT '',
|
||||
secret_key TEXT NOT NULL DEFAULT '',
|
||||
token TEXT NOT NULL DEFAULT '',
|
||||
api_key TEXT NOT NULL DEFAULT '',
|
||||
model_base_url TEXT NOT NULL DEFAULT '',
|
||||
email TEXT NOT NULL DEFAULT '',
|
||||
organization TEXT NOT NULL DEFAULT '',
|
||||
team TEXT NOT NULL DEFAULT '',
|
||||
base_url TEXT NOT NULL DEFAULT '',
|
||||
saved_at INTEGER NOT NULL DEFAULT 0
|
||||
)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_requests_ts ON requests(ts)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_requests_model ON requests(model)`,
|
||||
// Ensure the single credentials row exists so UPSERTs and SELECTs always
|
||||
// have a target.
|
||||
`INSERT INTO credentials (id) VALUES (1) ON CONFLICT(id) DO NOTHING`,
|
||||
}
|
||||
for _, q := range stmts {
|
||||
if _, err := db.Exec(q); err != nil {
|
||||
return fmt.Errorf("schema: %w", err)
|
||||
}
|
||||
}
|
||||
// Add cached_tokens to databases created before this column existed. SQLite
|
||||
// returns "duplicate column name" when it already exists; that is expected
|
||||
// and ignored.
|
||||
if _, err := db.Exec(`ALTER TABLE requests ADD COLUMN cached_tokens INTEGER NOT NULL DEFAULT 0`); err != nil {
|
||||
if !strings.Contains(err.Error(), "duplicate column") {
|
||||
return fmt.Errorf("migrate cached_tokens: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user