AskaEth
59a1ac4056
通道：新增 model/session_ctx.go 的 WithSessionID / SessionIDFrom（与 tool.WithTurnPersona
同族的 ctx 范式，避开 model 反向 import loop 的环），loop 在 Run 顶部播种；适配器
record(ctx,...) 接会话；usageSessionArg 空串映射为 NULL。
兼容：sink 未注入或会话未播种时行为逐字不变；既有 model 测试零断言改动即全绿。
文件：model/{session_ctx.go(新), usage.go, usage_store.go, openai.go, anthropic.go,
usage_session_test.go(新)} + loop/{run.go, usage_session_test.go(新)}。
变异三处：去 loop 打标 / 写侧不绑会话 / 适配器不读 ctx ⇒ 各自红行，还原复绿。
契约与日志字段集未动；读侧不按 session 过滤，面板汇总口径不变。
267 lines
8.4 KiB
Go
267 lines
8.4 KiB
Go
package model
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"fmt"
|
||
"log/slog"
|
||
"sync"
|
||
"time"
|
||
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||
)
|
||
|
||
// 本文件是**用量落库的写入侧**（切片 U1，主人裁决「2.4 落」）：一个有界队列 +
|
||
// 单写者 + 批量事务的异步写入器。
|
||
//
|
||
// 为什么不是「在 logUsage 里直接 INSERT」：那条语句跑在**每一次模型调用的返回
|
||
// 路径**上，库慢一点、锁一下，主人的每次对话就跟着慢一点 —— 记成本的开销绝不
|
||
// 能加到被记成本的那次调用上。这里因此把两者彻底解耦：
|
||
//
|
||
// 1. `RecordUsage` 只往**有界** channel 里塞一个样本，塞不进就**丢弃并限流
|
||
// WARN**（绝不阻塞、绝不 panic、绝不无界增长）；
|
||
// 2. 一个常驻单写者把队列攒成批次，在一个事务里落库；
|
||
// 3. 整批失败只记一条带来源的 WARN（**不静默**），调用路径毫无感知。
|
||
//
|
||
// 写入样本里没有凭据、没有正文（见 UsageSample），这条红线在类型层面就成立。
|
||
|
||
// usageSource 是本文件每条日志的来源标签（跨模块可追溯定位，项目规范 §10）。
|
||
const usageSource = "model/usage_store"
|
||
|
||
const (
|
||
// usageQueueCapacity 是落库队列的容量。它是**有界**的：满了就丢弃，
|
||
// 而不是把模型调用拖住。
|
||
usageQueueCapacity = 1024
|
||
// usageBatchMax 是单次批量事务最多写入的行数。
|
||
usageBatchMax = 128
|
||
// usageFlushInterval 是队列非空时的最长等待：低频调用不会因为没攒满一批
|
||
// 就一直不落库（面板看得到的是近 24h，不该差这几十秒）。
|
||
usageFlushInterval = 2 * time.Second
|
||
// usageDropWarnInterval 是「队列满、已丢弃样本」告警的限流间隔：真实部署里
|
||
// 库挂了会让每一次调用都丢弃，不限流就会把日志刷爆。
|
||
usageDropWarnInterval = 30 * time.Second
|
||
// usageFinalFlushTimeout 是停止时把队列剩余样本写完的期限。
|
||
usageFinalFlushTimeout = 5 * time.Second
|
||
)
|
||
|
||
// insertUsageSQL 写入一条用量样本。session_id 取 ctx 播种的会话标识
|
||
// （session_ctx.go），**没有会话归属时写 NULL**：管理面板试连、未接循环的
|
||
// 后台流程本就不属于任何会话，写空串会让它们与真实会话在列里混成一个值。
|
||
const insertUsageSQL = `INSERT INTO llm_usage
|
||
(session_id, purpose, adapter, model, prompt_tokens, completion_tokens,
|
||
cache_hit_tokens, cache_miss_tokens, reasoning_tokens, total_tokens, took_ms)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
|
||
|
||
// usageSessionArg 把会话标识转成可绑定的值：空串 ⇒ NULL。
|
||
func usageSessionArg(sessionID string) any {
|
||
if sessionID == "" {
|
||
return nil
|
||
}
|
||
return sessionID
|
||
}
|
||
|
||
// UsageWriter 把用量样本异步落进 `llm_usage`。它实现 UsageSink，因此在模型适配器
|
||
// 里就是一个「可选出口」：注入即有、不注入则只记日志。
|
||
//
|
||
// 构造后必须 Start 才会真的落库；未 Start 时队列照样收、塞满后照样丢弃并告警
|
||
// （不静默失败）。Stop 必须在关库之前调用。
|
||
type UsageWriter struct {
|
||
db *sql.DB
|
||
log *slog.Logger
|
||
|
||
queue chan UsageSample
|
||
done chan struct{}
|
||
|
||
mu sync.Mutex
|
||
cancel context.CancelFunc
|
||
started bool
|
||
dropped int64
|
||
lastWarn time.Time
|
||
}
|
||
|
||
// NewUsageWriter 构造写入器。db 为 nil（未接库）时返回 nil —— 调用方据此把
|
||
// 「无 sink」原样传给适配器，行为等于这一版之前。
|
||
func NewUsageWriter(db *sql.DB, log *slog.Logger) *UsageWriter {
|
||
if db == nil {
|
||
return nil
|
||
}
|
||
if log == nil {
|
||
log = discardLogger(usageSource)
|
||
}
|
||
return &UsageWriter{
|
||
db: db,
|
||
log: logger.WithSource(log, usageSource),
|
||
queue: make(chan UsageSample, usageQueueCapacity),
|
||
done: make(chan struct{}),
|
||
}
|
||
}
|
||
|
||
// RecordUsage 实现 UsageSink：把样本塞进有界队列，塞不进就丢弃并（限流）告警。
|
||
//
|
||
// 它**绝不阻塞**：整个函数是一次非阻塞 channel 发送加一次可能的加锁自增，
|
||
// 因此模型调用路径的额外开销与库的健康状况无关。
|
||
func (w *UsageWriter) RecordUsage(sample UsageSample) {
|
||
if w == nil {
|
||
return
|
||
}
|
||
select {
|
||
case w.queue <- sample:
|
||
default:
|
||
w.noteDrop()
|
||
}
|
||
}
|
||
|
||
// Start 拉起单写者循环。重复调用是空操作；ctx 取消（或 Stop）后循环退出。
|
||
func (w *UsageWriter) Start(ctx context.Context) {
|
||
if w == nil {
|
||
return
|
||
}
|
||
w.mu.Lock()
|
||
defer w.mu.Unlock()
|
||
if w.started {
|
||
return
|
||
}
|
||
w.started = true
|
||
runCtx, cancel := context.WithCancel(ctx)
|
||
w.cancel = cancel
|
||
go w.loop(runCtx)
|
||
}
|
||
|
||
// Stop 停掉循环并等它退出，退出前尽力把队列里剩下的样本写完。未启动时是空操作。
|
||
// 它必须在关库**之前**调用（收尾写入要用那条连接池）。
|
||
func (w *UsageWriter) Stop() {
|
||
if w == nil {
|
||
return
|
||
}
|
||
w.mu.Lock()
|
||
cancel := w.cancel
|
||
started := w.started
|
||
w.mu.Unlock()
|
||
if !started || cancel == nil {
|
||
return
|
||
}
|
||
cancel()
|
||
<-w.done
|
||
}
|
||
|
||
// loop 是单写者：攒批 → 批量事务落库，直到 ctx 取消。
|
||
func (w *UsageWriter) loop(ctx context.Context) {
|
||
defer close(w.done)
|
||
ticker := time.NewTicker(usageFlushInterval)
|
||
defer ticker.Stop()
|
||
batch := make([]UsageSample, 0, usageBatchMax)
|
||
flush := func() {
|
||
if len(batch) == 0 {
|
||
return
|
||
}
|
||
w.writeBatch(ctx, batch)
|
||
batch = batch[:0]
|
||
}
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
// 收尾用**脱离取消**的上下文：ctx 已经取消，沿用它会让最后一批
|
||
// 必然写失败，等于丢掉已经收下的样本。
|
||
w.flushFinal(drainQueue(w.queue, batch))
|
||
return
|
||
case sample := <-w.queue:
|
||
batch = append(batch, sample)
|
||
if len(batch) >= usageBatchMax {
|
||
flush()
|
||
}
|
||
case <-ticker.C:
|
||
flush()
|
||
}
|
||
}
|
||
}
|
||
|
||
// flushFinal 在停机路径上把剩余样本写完（有期限，绝不无限等库）。
|
||
func (w *UsageWriter) flushFinal(remaining []UsageSample) {
|
||
if len(remaining) == 0 {
|
||
return
|
||
}
|
||
ctx, cancel := context.WithTimeout(context.Background(), usageFinalFlushTimeout)
|
||
defer cancel()
|
||
w.writeBatch(ctx, remaining)
|
||
}
|
||
|
||
// writeBatch 在一个事务里落一批样本。失败只记 WARN（带批次大小与原因），
|
||
// 绝不上抛 —— 调用路径早已返回，没有任何人在等这个结果。
|
||
func (w *UsageWriter) writeBatch(ctx context.Context, batch []UsageSample) {
|
||
if err := insertUsageBatch(ctx, w.db, batch); err != nil {
|
||
w.log.Warn("llm usage batch insert failed; batch dropped", "count", len(batch), "err", err)
|
||
}
|
||
}
|
||
|
||
// noteDrop 记一次「队列满、样本被丢弃」，并按 usageDropWarnInterval 限流告警。
|
||
// 丢弃计数是累计值，因此告警里能看到这一段到底丢了多少条。
|
||
func (w *UsageWriter) noteDrop() {
|
||
w.mu.Lock()
|
||
w.dropped++
|
||
dropped := w.dropped
|
||
now := time.Now()
|
||
shouldWarn := w.lastWarn.IsZero() || now.Sub(w.lastWarn) >= usageDropWarnInterval
|
||
if shouldWarn {
|
||
w.lastWarn = now
|
||
}
|
||
w.mu.Unlock()
|
||
if shouldWarn {
|
||
w.log.Warn("llm usage queue full; samples dropped (never blocks the model call)",
|
||
"dropped", dropped, "capacity", usageQueueCapacity)
|
||
}
|
||
}
|
||
|
||
// drainQueue 把队列里当前排着的样本全部取出（不阻塞），追加到 dst 之后返回。
|
||
func drainQueue(queue chan UsageSample, dst []UsageSample) []UsageSample {
|
||
for {
|
||
select {
|
||
case sample := <-queue:
|
||
dst = append(dst, sample)
|
||
default:
|
||
return dst
|
||
}
|
||
}
|
||
}
|
||
|
||
// insertUsageBatch 在一个事务里写入一批样本：单条 prepared statement 循环执行，
|
||
// 一次提交。用 prepared statement 而不是多值 INSERT，是为了避开各家 SQLite 对
|
||
// 绑定参数个数的不同上限（11 列 × 128 行）。
|
||
func insertUsageBatch(ctx context.Context, db *sql.DB, batch []UsageSample) error {
|
||
if db == nil || len(batch) == 0 {
|
||
return nil
|
||
}
|
||
tx, err := db.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return fmt.Errorf("begin usage batch: %w", err)
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
stmt, err := tx.PrepareContext(ctx, insertUsageSQL)
|
||
if err != nil {
|
||
return fmt.Errorf("prepare usage insert: %w", err)
|
||
}
|
||
defer stmt.Close()
|
||
|
||
for i, sample := range batch {
|
||
if _, err := stmt.ExecContext(ctx,
|
||
usageSessionArg(sample.SessionID),
|
||
sample.Purpose,
|
||
sample.Adapter,
|
||
sample.Model,
|
||
sample.PromptTokens,
|
||
sample.CompletionTokens,
|
||
sample.CacheHitTokens,
|
||
sample.CacheMissTokens,
|
||
sample.ReasoningTokens,
|
||
sample.TotalTokens,
|
||
sample.TookMs,
|
||
); err != nil {
|
||
return fmt.Errorf("insert usage sample %d: %w", i, err)
|
||
}
|
||
}
|
||
if err := tx.Commit(); err != nil {
|
||
return fmt.Errorf("commit usage batch: %w", err)
|
||
}
|
||
return nil
|
||
}
|