zhanlu_proxy/internal/store/stats_recorder.go

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