Cyrene_v2/internal/runtime/model/usage_query.go

package model

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// ErrUsageUnavailable 表示用量汇总出口不可用（未接库 / 未构造读取器）。
// 网关据此回 500 usage_unavailable，而不是回一份全 0 的假汇总。
var ErrUsageUnavailable = errors.New("model: llm usage store unavailable")

// 本文件是**用量落库的读取侧**（切片 U1）：把 `llm_usage` 表在给定时间窗与用途
// 过滤下聚合成一份只读汇总，供管理端点 `GET /api/v1/admin/llm/usage` 与面板的
// 用量小卡消费。
//
// 三条纪律：
//
//  1. **只读**：只有 SELECT（外加保留期清理的 DELETE，见 PurgeUsage），不建表、
//     不迁移，复用组合根注入的活连接。
//  2. **不含凭据**：查出来的字段只有 token 计数与模型/用途标签 —— 表里本来就没有
//     凭据与正文，这条红线在写入侧就已经成立。
//  3. **不编造**：读不到就返回错误，绝不用一份全 0 的假汇总假装成功（面板据此
//     显示错误，而不是一个看着正常的假数字）。

const (
	// UsageDefaultWindow 是汇总的缺省时间窗（近 24 小时）：不带 since/until 时
	// 「现在」往前推这么久。
	UsageDefaultWindow = 24 * time.Hour
	// UsageMaxWindow 是一次汇总查询允许的最大时间窗（90 天）。超过它的请求由
	// 网关以 400 明确拒绝，**不静默裁剪** —— 静默裁剪会让客户端以为拿到的是
	// 全量，把「我没查那么久」误读成「就那么少」。
	UsageMaxWindow = 90 * 24 * time.Hour
	// usageTimeLayout 与迁移 0017 的 `strftime('%Y-%m-%dT%H:%M:%fZ')` 逐字同形
	// （定宽 UTC ISO-8601，毫秒）。定宽使字典序比较等价于时间序，时间窗因此能
	// 直接用字符串比较、吃上 idx_llm_usage_created_at。
	usageTimeLayout = "2006-01-02T15:04:05.000Z"
)

// UsageQuery 是一次汇总查询的条件。Since/Until 为零值时由调用方按缺省窗口补齐
// （见 UsageDefaultWindow）；Purpose 为空表示不过滤。
type UsageQuery struct {
	Since   time.Time
	Until   time.Time
	Purpose string
}

// UsageTotals 是一个时间窗（或一个分组）内的用量合计。全部是**非敏感计数**。
type UsageTotals struct {
	// Calls 是窗口内成功调用次数。
	Calls int64
	// Prompt / Completion 是输入侧与输出侧 token 合计。
	Prompt     int64
	Completion int64
	// CacheHit / CacheMiss 是输入侧两档 token 合计。
	CacheHit  int64
	CacheMiss int64
}

// HitRate 返回缓存命中率 = cacheHit / (cacheHit + cacheMiss)。分母为 0（窗口内没有
// 任何输入侧明细，或根本没有调用）时返回 0 —— 口径与 Response.CacheHitRate 同源，
// 不编造 100% 或 0% 以外的假值。
func (t UsageTotals) HitRate() float64 {
	total := t.CacheHit + t.CacheMiss
	if total <= 0 {
		return 0
	}
	return float64(t.CacheHit) / float64(total)
}

// UsageGroup 是按某个键（用途 / 模型）聚合的一行。
type UsageGroup struct {
	// Key 是分组键（purpose 或 model 的取值）。
	Key string
	UsageTotals
}

// UsageSummary 是一次汇总查询的结果。
type UsageSummary struct {
	// Window 是本次实际生效的时间窗与过滤条件（回显给客户端，而不是让它猜
	// 缺省值）。
	Window UsageQuery
	// Totals 是整窗合计。
	Totals UsageTotals
	// ByPurpose / ByModel 是按用途与按模型的分解（按调用次数降序，同次数按键
	// 升序 —— 输出确定，客户端可安全 diff）。
	ByPurpose []UsageGroup
	ByModel   []UsageGroup
}

// UsageReader 是 `llm_usage` 表的只读出口。
type UsageReader struct {
	db *sql.DB
}

// NewUsageReader 构造读取器；db 为 nil（未接库）时返回 nil。
func NewUsageReader(db *sql.DB) *UsageReader {
	if db == nil {
		return nil
	}
	return &UsageReader{db: db}
}

// usageRowScanner 是 *sql.Row 与 *sql.Rows 共有的扫码能力，使合计与分组两条
// 查询路径共用同一个扫码函数。
type usageRowScanner interface {
	Scan(dest ...any) error
}

// usageAggregateCols 是合计列清单（顺序与 scanUsageTotals 的扫码顺序逐字一致）。
const usageAggregateCols = `COUNT(*),
    COALESCE(SUM(prompt_tokens),0),
    COALESCE(SUM(completion_tokens),0),
    COALESCE(SUM(cache_hit_tokens),0),
    COALESCE(SUM(cache_miss_tokens),0)`

// Summary 聚合一个时间窗内的用量。q 的时间窗必须已由调用方校验并补齐（零值窗口
// 会原样查库、什么都查不到）。
func (r *UsageReader) Summary(ctx context.Context, q UsageQuery) (UsageSummary, error) {
	if r == nil || r.db == nil {
		return UsageSummary{}, fmt.Errorf("model: %w", ErrUsageUnavailable)
	}
	totals, err := r.totalsOf(ctx, q)
	if err != nil {
		return UsageSummary{}, err
	}
	byPurpose, err := r.groupsOf(ctx, q, usageGroupByPurpose)
	if err != nil {
		return UsageSummary{}, err
	}
	byModel, err := r.groupsOf(ctx, q, usageGroupByModel)
	if err != nil {
		return UsageSummary{}, err
	}
	return UsageSummary{Window: q, Totals: totals, ByPurpose: byPurpose, ByModel: byModel}, nil
}

// totalsOf 读整窗合计。空窗时 SUM 为 NULL，COALESCE 收敛成 0；表不存在（迁移没跑）
// 则直接报错，不静默报 0。
func (r *UsageReader) totalsOf(ctx context.Context, q UsageQuery) (UsageTotals, error) {
	query := `SELECT ` + usageAggregateCols + ` FROM llm_usage WHERE created_at >= ? AND created_at <= ?`
	args := []any{usageTime(q.Since), usageTime(q.Until)}
	if q.Purpose != "" {
		query += ` AND purpose = ?`
		args = append(args, q.Purpose)
	}
	totals, err := scanUsageTotals(r.db.QueryRowContext(ctx, query, args...))
	if err != nil {
		return UsageTotals{}, fmt.Errorf("model: read usage totals: %w", err)
	}
	return totals, nil
}

// usageGroupBy* 是允许的分组列：它们是代码常量，**不是**用户输入（分组列无法参数化，
// 因此这里以白名单常量而非拼接来保证不注入）。
const (
	usageGroupByPurpose = "purpose"
	usageGroupByModel   = "model"
)

// groupsOf 按 groupBy 列聚合。groupBy 只接受 usageGroupBy* 常量，未知值一律报错
// （绝不把任意字符串拼进 SQL）。
func (r *UsageReader) groupsOf(ctx context.Context, q UsageQuery, groupBy string) ([]UsageGroup, error) {
	if groupBy != usageGroupByPurpose && groupBy != usageGroupByModel {
		return nil, fmt.Errorf("model: %w: unknown usage group %q", ErrBadRequest, groupBy)
	}
	query := `SELECT ` + groupBy + `, ` + usageAggregateCols +
		` FROM llm_usage WHERE created_at >= ? AND created_at <= ?`
	args := []any{usageTime(q.Since), usageTime(q.Until)}
	if q.Purpose != "" {
		query += ` AND purpose = ?`
		args = append(args, q.Purpose)
	}
	query += ` GROUP BY ` + groupBy + ` ORDER BY COUNT(*) DESC, ` + groupBy + ` ASC`

	rows, err := r.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("model: read usage groups by %s: %w", groupBy, err)
	}
	defer rows.Close()

	out := make([]UsageGroup, 0, 8)
	for rows.Next() {
		group, err := scanUsageGroup(rows)
		if err != nil {
			return nil, fmt.Errorf("model: scan usage group by %s: %w", groupBy, err)
		}
		out = append(out, group)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("model: iterate usage groups by %s: %w", groupBy, err)
	}
	return out, nil
}

// scanUsageTotals 按 usageAggregateCols 的顺序扫码。
func scanUsageTotals(sc usageRowScanner) (UsageTotals, error) {
	var t UsageTotals
	if err := sc.Scan(&t.Calls, &t.Prompt, &t.Completion, &t.CacheHit, &t.CacheMiss); err != nil {
		return UsageTotals{}, err
	}
	return t, nil
}

// scanUsageGroup 扫码「分组键 + 合计列」。
func scanUsageGroup(sc usageRowScanner) (UsageGroup, error) {
	var key string
	var t UsageTotals
	if err := sc.Scan(&key, &t.Calls, &t.Prompt, &t.Completion, &t.CacheHit, &t.CacheMiss); err != nil {
		return UsageGroup{}, err
	}
	return UsageGroup{Key: key, UsageTotals: t}, nil
}

// PurgeUsage 删除 `created_at` 早于 before 的用量行，返回删除条数。db 为 nil 时是
// 空操作。它与保留期清理（组合根按 LLM_USAGE_RETENTION 定时调用）是唯一会写这张
// 表、又不是追加样本的语句。
func PurgeUsage(ctx context.Context, db *sql.DB, before time.Time) (int64, error) {
	if db == nil {
		return 0, nil
	}
	res, err := db.ExecContext(ctx, `DELETE FROM llm_usage WHERE created_at < ?`, usageTime(before))
	if err != nil {
		return 0, fmt.Errorf("model: purge llm usage: %w", err)
	}
	removed, err := res.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("model: purge llm usage rows affected: %w", err)
	}
	return removed, nil
}

// usageTime 把时间点转成与迁移 0017 的 DEFAULT 逐字同形的 UTC 字符串（定宽毫秒），
// 使时间窗比较是纯字符串比较、且能走 idx_llm_usage_created_at。
func usageTime(t time.Time) string {
	return t.UTC().Format(usageTimeLayout)
}