Add token usage stats and consolidate credentials in SQLite

- 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.
2026-08-19 15:52:00 +08:00
parent dd5c560b64
commit fa9640d919
15 changed files with 1101 additions and 103 deletions
@@ -0,0 +1,187 @@
+package store
+
+import (
+	"fmt"
+	"time"
+
+	"git.misaka.ren/M1saka/zhanlu_proxy/internal/stats"
+)
+
+// Compile-time guard: *Store satisfies stats.Recorder.
+var _ stats.Recorder = (*Store)(nil)
+
+const defaultRecentLimit = 200
+
+// Record appends a single usage observation.
+func (s *Store) Record(rec stats.Record) error {
+	if rec.Status == "" {
+		rec.Status = "success"
+	}
+	stream := 0
+	if rec.Stream {
+		stream = 1
+	}
+	_, err := s.db.Exec(`INSERT INTO requests (ts, model, stream, prompt_tokens, completion_tokens, total_tokens, reasoning_tokens, cached_tokens, status, latency_ms) VALUES (?,?,?,?,?,?,?,?,?,?)`,
+		rec.Ts.Unix(), rec.Model, stream, rec.PromptTokens, rec.CompletionTokens, rec.TotalTokens, rec.ReasoningTokens, rec.CachedTokens, rec.Status, rec.LatencyMs)
+	return err
+}
+
+// Stats computes the aggregate summary for the given query.
+func (s *Store) Stats(q stats.Query) (*stats.Summary, error) {
+	q = normalizeQuery(q)
+	where, args := whereClause(q)
+
+	sum := &stats.Summary{}
+	if err := s.scanTotals(sum, where, args); err != nil {
+		return nil, err
+	}
+	if sum.Totals.PromptTokens > 0 {
+		sum.Totals.CacheRate = float64(sum.Totals.CachedTokens) / float64(sum.Totals.PromptTokens)
+	}
+	if err := s.scanPerModel(sum, where, args); err != nil {
+		return nil, err
+	}
+	if err := s.scanDaily(sum, where, args); err != nil {
+		return nil, err
+	}
+	if err := s.scanRecent(sum, q, where, args); err != nil {
+		return nil, err
+	}
+	return sum, nil
+}
+
+// Reset deletes all recorded usage statistics (the credentials row is kept).
+func (s *Store) Reset() error {
+	_, err := s.db.Exec(`DELETE FROM requests`)
+	return err
+}
+
+func normalizeQuery(q stats.Query) stats.Query {
+	if q.Limit <= 0 {
+		q.Limit = defaultRecentLimit
+	}
+	if q.Limit > 5000 {
+		q.Limit = 5000
+	}
+	return q
+}
+
+func whereClause(q stats.Query) (string, []any) {
+	var conds []string
+	var args []any
+	if !q.Since.IsZero() {
+		conds = append(conds, "ts >= ?")
+		args = append(args, q.Since.Unix())
+	}
+	if !q.Until.IsZero() {
+		conds = append(conds, "ts <= ?")
+		args = append(args, q.Until.Unix())
+	}
+	if q.Model != "" {
+		conds = append(conds, "model = ?")
+		args = append(args, q.Model)
+	}
+	if q.Stream != nil {
+		conds = append(conds, "stream = ?")
+		args = append(args, boolToInt(*q.Stream))
+	}
+	if len(conds) == 0 {
+		return "", args
+	}
+	out := ""
+	for i, p := range conds {
+		if i > 0 {
+			out += " AND "
+		}
+		out += p
+	}
+	return " WHERE " + out, args
+}
+
+func boolToInt(b bool) int {
+	if b {
+		return 1
+	}
+	return 0
+}
+
+func (s *Store) scanTotals(sum *stats.Summary, where string, args []any) error {
+	q := `SELECT COUNT(*),
+  COALESCE(SUM(CASE WHEN status='success' THEN 1 ELSE 0 END),0),
+  COALESCE(SUM(CASE WHEN status!='success' THEN 1 ELSE 0 END),0),
+  COALESCE(SUM(prompt_tokens),0),
+  COALESCE(SUM(completion_tokens),0),
+  COALESCE(SUM(total_tokens),0),
+  COALESCE(SUM(reasoning_tokens),0),
+  COALESCE(SUM(cached_tokens),0) FROM requests` + where
+	row := s.db.QueryRow(q, args...)
+	var success, failures int64
+	err := row.Scan(&sum.Totals.Requests, &success, &failures,
+		&sum.Totals.PromptTokens, &sum.Totals.CompletionTokens,
+		&sum.Totals.TotalTokens, &sum.Totals.ReasoningTokens, &sum.Totals.CachedTokens)
+	if err != nil {
+		return fmt.Errorf("scan totals: %w", err)
+	}
+	sum.Totals.SuccessRequests = int(success)
+	sum.Totals.ErrorRequests = int(failures)
+	return nil
+}
+
+func (s *Store) scanPerModel(sum *stats.Summary, where string, args []any) error {
+	q := `SELECT model, COUNT(*), COALESCE(SUM(prompt_tokens),0), COALESCE(SUM(completion_tokens),0), COALESCE(SUM(total_tokens),0), COALESCE(SUM(cached_tokens),0) FROM requests` + where + " GROUP BY model ORDER BY SUM(total_tokens) DESC"
+	rows, err := s.db.Query(q, args...)
+	if err != nil {
+		return fmt.Errorf("scan per-model: %w", err)
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var m stats.ModelStat
+		if err := rows.Scan(&m.Model, &m.Requests, &m.PromptTokens, &m.CompletionTokens, &m.TotalTokens, &m.CachedTokens); err != nil {
+			return err
+		}
+		sum.PerModel = append(sum.PerModel, m)
+	}
+	return rows.Err()
+}
+
+func (s *Store) scanDaily(sum *stats.Summary, where string, args []any) error {
+	q := `SELECT date(ts,'unixepoch','localtime') AS day, COUNT(*), COALESCE(SUM(prompt_tokens),0), COALESCE(SUM(completion_tokens),0), COALESCE(SUM(total_tokens),0), COALESCE(SUM(cached_tokens),0) FROM requests` + where + " GROUP BY day ORDER BY day ASC"
+	rows, err := s.db.Query(q, args...)
+	if err != nil {
+		return fmt.Errorf("scan daily: %w", err)
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var d stats.DayStat
+		if err := rows.Scan(&d.Day, &d.Requests, &d.PromptTokens, &d.CompletionTokens, &d.TotalTokens, &d.CachedTokens); err != nil {
+			return err
+		}
+		sum.Daily = append(sum.Daily, d)
+	}
+	return rows.Err()
+}
+
+func (s *Store) scanRecent(sum *stats.Summary, q stats.Query, where string, args []any) error {
+	limit := q.Limit
+	if limit <= 0 {
+		limit = defaultRecentLimit
+	}
+	query := `SELECT ts, model, stream, prompt_tokens, completion_tokens, total_tokens, reasoning_tokens, cached_tokens, status, latency_ms FROM requests` + where + " ORDER BY id DESC LIMIT ?"
+	rows, err := s.db.Query(query, append(args, limit)...)
+	if err != nil {
+		return fmt.Errorf("scan recent: %w", err)
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var rec stats.Record
+		var ts int64
+		var stream int
+		if err := rows.Scan(&ts, &rec.Model, &stream, &rec.PromptTokens, &rec.CompletionTokens, &rec.TotalTokens, &rec.ReasoningTokens, &rec.CachedTokens, &rec.Status, &rec.LatencyMs); err != nil {
+			return err
+		}
+		rec.Ts = time.Unix(ts, 0).Local()
+		rec.Stream = stream == 1
+		sum.Recent = append(sum.Recent, rec)
+	}
+	return rows.Err()
+}