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() }