// 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 }