hunjing/shenshi/internal/metering/usage.go

// Package metering 用量计量：每次尝试一条记录（stdout 结构化日志 + 神识自有库落库）。
// 未配置 SHENSHI_DATABASE_URL 时退化为进程内环形缓冲（开发档），/me/usage 与对账抽样同语义。
package metering

import (
	"encoding/json"
	"log"
	"sync"
	"time"

	"hunjing/shenshi/internal/store"
)

// Attempt 一次端点尝试的计量记录（IR 维度齐全，user_id 为归属计量主键）。
type Attempt struct {
	ID                string    `json:"id"` // run_id#endpoint 序号（结算幂等键）
	RunID             string    `json:"run_id"`
	CharacterID       string    `json:"character_id,omitempty"`
	UserID            string    `json:"user_id,omitempty"`
	Model             string    `json:"model"`
	Slot              string    `json:"slot"`
	Endpoint          string    `json:"endpoint"`
	Status            string    `json:"status"` // ok / error / first_byte_timeout
	ErrorCode         string    `json:"error_code,omitempty"`
	FirstTokenMS      int64     `json:"first_token_ms"`
	TotalMS           int64     `json:"total_ms"`
	PromptTokens      int       `json:"prompt_tokens"`
	CompletionTokens  int       `json:"completion_tokens"`
	At                time.Time `json:"at"`
}

// Log 计量器：stdout JSON 始终输出；持久化走 store（nil 时仅进程内有界环形）。
type Log struct {
	mu    sync.Mutex
	store store.Store
	entries []Attempt
	limit int
}

// New 构造计量器（st 为 nil 时进程内保留最近 N 条）。
func New(limit int, st store.Store) *Log {
	return &Log{store: st, entries: make([]Attempt, 0, limit), limit: limit}
}

// Record 记录一次尝试：stdout JSON 日志（对账抽样基准）+ 落库。
func (l *Log) Record(a Attempt) {
	a.At = time.Now().UTC()
	if raw, err := json.Marshal(a); err == nil {
		log.Printf(`{"kind":"shenshi_usage",%s`, string(raw[1:]))
	}
	if l.store != nil {
		_ = l.store.InsertUsage(store.UsageLog{
			ID: a.ID, RunID: a.RunID, CharacterID: a.CharacterID, UserID: a.UserID,
			Model: a.Model, Slot: a.Slot, Endpoint: a.Endpoint,
			Status: a.Status, ErrorCode: a.ErrorCode,
			FirstTokenMS: a.FirstTokenMS, TotalMS: a.TotalMS,
			PromptTokens: a.PromptTokens, CompletionTokens: a.CompletionTokens, At: a.At,
		})
		return
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	if len(l.entries) >= l.limit {
		l.entries = l.entries[1:]
	}
	l.entries = append(l.entries, a)
}

// Aggregate 按 用户×模型×端点 聚合（userID 为空聚合全体；管理面语义）。
func (l *Log) Aggregate(userID string) ([]store.UsageAgg, error) {
	if l.store != nil {
		return l.store.AggregateUsage(userID)
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	agg := map[string]*store.UsageAgg{}
	for _, e := range l.entries {
		if userID != "" && e.UserID != userID {
			continue
		}
		key := e.UserID + "|" + e.Model + "|" + e.Endpoint
		row := agg[key]
		if row == nil {
			row = &store.UsageAgg{UserID: e.UserID, Model: e.Model, Endpoint: e.Endpoint}
			agg[key] = row
		}
		row.Requests++
		if e.Status == "ok" {
			row.OK++
			row.PromptTokens += e.PromptTokens
			row.CompletionTokens += e.CompletionTokens
		}
	}
	out := make([]store.UsageAgg, 0, len(agg))
	for _, row := range agg {
		out = append(out, *row)
	}
	return out, nil
}

// Sample 对账抽样：最近 N 条原始记录（与 Provider 账单比对）。
func (l *Log) Sample(limit int) ([]store.UsageLog, error) {
	if l.store != nil {
		return l.store.SampleUsage(limit)
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	n := limit
	if n > len(l.entries) {
		n = len(l.entries)
	}
	out := make([]store.UsageLog, n)
	for i, e := range l.entries[len(l.entries)-n:] {
		out[i] = store.UsageLog{
			ID: e.ID, RunID: e.RunID, CharacterID: e.CharacterID, UserID: e.UserID,
			Model: e.Model, Slot: e.Slot, Endpoint: e.Endpoint,
			Status: e.Status, ErrorCode: e.ErrorCode,
			FirstTokenMS: e.FirstTokenMS, TotalMS: e.TotalMS,
			PromptTokens: e.PromptTokens, CompletionTokens: e.CompletionTokens, At: e.At,
		}
	}
	return out, nil
}