89 lines
2.5 KiB
Go
89 lines
2.5 KiB
Go
// Package metering 用量计量：每次尝试一条记录（stdout 结构化日志 + 进程内聚合）。
|
||
// 生产部署接 PG（shenshi-gateway.md §7）；M4a 先以日志与内存聚合支撑 /me/usage 与对账抽样。
|
||
package metering
|
||
|
||
import (
|
||
"encoding/json"
|
||
"log"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// Attempt 一次端点尝试的计量记录。
|
||
type Attempt struct {
|
||
RunID string `json:"run_id"`
|
||
CharacterID string `json:"character_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"`
|
||
}
|
||
|
||
// UsageAgg /me/usage 聚合行。
|
||
type UsageAgg struct {
|
||
Model string `json:"model"`
|
||
Endpoint string `json:"endpoint"`
|
||
Requests int `json:"requests"`
|
||
OK int `json:"ok"`
|
||
PromptTokens int `json:"prompt_tokens"`
|
||
CompletionTokens int `json:"completion_tokens"`
|
||
}
|
||
|
||
// Log 进程内计量器（有界环形，避免无界增长）。
|
||
type Log struct {
|
||
mu sync.Mutex
|
||
entries []Attempt
|
||
limit int
|
||
}
|
||
|
||
// New 构造计量器（保留最近 N 条）。
|
||
func New(limit int) *Log {
|
||
return &Log{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:]))
|
||
}
|
||
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 按 模型×端点 聚合。
|
||
func (l *Log) Aggregate() []UsageAgg {
|
||
l.mu.Lock()
|
||
defer l.mu.Unlock()
|
||
agg := map[string]*UsageAgg{}
|
||
for _, e := range l.entries {
|
||
key := e.Model + "|" + e.Endpoint
|
||
row := agg[key]
|
||
if row == nil {
|
||
row = &UsageAgg{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([]UsageAgg, 0, len(agg))
|
||
for _, row := range agg {
|
||
out = append(out, *row)
|
||
}
|
||
return out
|
||
}
|