444 lines
17 KiB
Go
444 lines
17 KiB
Go
// Package repository 实现 cpa-ext 的本地持久化。
|
|
package repository
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"time"
|
|
|
|
"cpa-ext/internal/collection"
|
|
_ "github.com/mattn/go-sqlite3"
|
|
)
|
|
|
|
const usageSchema = `
|
|
CREATE TABLE IF NOT EXISTS usage_records (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
request_id TEXT NOT NULL DEFAULT '',
|
|
execution_id TEXT NOT NULL DEFAULT '',
|
|
trace_id TEXT NOT NULL DEFAULT '',
|
|
requested_at TEXT NOT NULL,
|
|
api_key TEXT NOT NULL DEFAULT '',
|
|
key_alias TEXT NOT NULL DEFAULT '',
|
|
model TEXT NOT NULL DEFAULT '',
|
|
reasoning_effort TEXT NOT NULL DEFAULT '',
|
|
service_tier TEXT NOT NULL DEFAULT '',
|
|
speed TEXT NOT NULL DEFAULT '',
|
|
failed INTEGER NOT NULL,
|
|
executor_type TEXT NOT NULL DEFAULT '',
|
|
request_type TEXT NOT NULL DEFAULT '',
|
|
endpoint TEXT NOT NULL DEFAULT '',
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cached_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_write_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
ttft_ns INTEGER NOT NULL DEFAULT 0,
|
|
latency_ns INTEGER NOT NULL DEFAULT 0,
|
|
cost_micros INTEGER,
|
|
price_tier TEXT NOT NULL DEFAULT '',
|
|
fast_requested INTEGER NOT NULL DEFAULT 0,
|
|
fast_pricing_applied INTEGER NOT NULL DEFAULT 0,
|
|
price_multiplier_numerator INTEGER NOT NULL DEFAULT 1,
|
|
price_multiplier_denominator INTEGER NOT NULL DEFAULT 1,
|
|
client_ip TEXT NOT NULL DEFAULT ''
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_usage_records_requested_at ON usage_records(requested_at DESC, id DESC);
|
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_records_execution_id ON usage_records(execution_id) WHERE execution_id <> '';
|
|
|
|
CREATE TABLE IF NOT EXISTS request_records (
|
|
request_id TEXT PRIMARY KEY,
|
|
trace_id TEXT NOT NULL DEFAULT '',
|
|
requested_at TEXT NOT NULL,
|
|
completed_at TEXT NOT NULL,
|
|
model TEXT NOT NULL DEFAULT '',
|
|
requested_model TEXT NOT NULL DEFAULT '',
|
|
source_format TEXT NOT NULL DEFAULT '',
|
|
stream INTEGER NOT NULL DEFAULT 0,
|
|
outcome TEXT NOT NULL,
|
|
status_code INTEGER NOT NULL DEFAULT 0,
|
|
error TEXT NOT NULL DEFAULT '',
|
|
endpoint TEXT NOT NULL DEFAULT '',
|
|
latency_ns INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_request_records_requested_at ON request_records(requested_at DESC);
|
|
|
|
CREATE TABLE IF NOT EXISTS model_prices (
|
|
model TEXT PRIMARY KEY,
|
|
input_rate_micros INTEGER NOT NULL,
|
|
cache_read_rate_micros INTEGER NOT NULL,
|
|
cache_write_rate_micros INTEGER NOT NULL,
|
|
output_rate_micros INTEGER NOT NULL,
|
|
long_context_enabled INTEGER NOT NULL DEFAULT 0,
|
|
long_context_threshold INTEGER NOT NULL DEFAULT 0,
|
|
long_context_comparison TEXT NOT NULL DEFAULT 'gt',
|
|
long_input_rate_micros INTEGER NOT NULL DEFAULT 0,
|
|
long_cache_read_rate_micros INTEGER NOT NULL DEFAULT 0,
|
|
long_cache_write_rate_micros INTEGER NOT NULL DEFAULT 0,
|
|
long_output_rate_micros INTEGER NOT NULL DEFAULT 0,
|
|
fast_pricing_enabled INTEGER NOT NULL DEFAULT 0,
|
|
fast_multiplier_numerator INTEGER NOT NULL DEFAULT 5,
|
|
fast_multiplier_denominator INTEGER NOT NULL DEFAULT 2,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
`
|
|
|
|
const (
|
|
// CLIProxyAPI UsageRecord currently has no RequestID. The usage and terminal
|
|
// callbacks nevertheless carry request-start timestamps from the same host
|
|
// execution, normally only a few milliseconds apart. This window is used
|
|
// only for the read projection and never changes persisted billing facts.
|
|
orphanLifecycleMatchWindow = 250 * time.Millisecond
|
|
// Near-equal candidates are deliberately left separate instead of risking a
|
|
// cross-request association under concurrent same-model traffic.
|
|
orphanLifecycleTieWindow = time.Millisecond
|
|
)
|
|
|
|
// SQLiteUsageRepository 使用单写连接保存全部用量记录。
|
|
type SQLiteUsageRepository struct {
|
|
db *sql.DB
|
|
}
|
|
|
|
func OpenSQLiteUsage(databasePath string) (*SQLiteUsageRepository, error) {
|
|
dsn, err := sqliteDSN(databasePath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
db, err := sql.Open("sqlite3", dsn)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("打开 SQLite: %w", err)
|
|
}
|
|
// 当前插件只需要一个同步 writer，避免并发写争用。
|
|
db.SetMaxOpenConns(1)
|
|
if err := db.Ping(); err != nil {
|
|
_ = db.Close()
|
|
return nil, fmt.Errorf("连接 SQLite: %w", err)
|
|
}
|
|
if _, err := db.Exec(usageSchema); err != nil {
|
|
_ = db.Close()
|
|
return nil, fmt.Errorf("初始化 SQLite schema: %w", err)
|
|
}
|
|
return &SQLiteUsageRepository{db: db}, nil
|
|
}
|
|
|
|
func sqliteDSN(databasePath string) (string, error) {
|
|
databasePath = strings.TrimSpace(databasePath)
|
|
if databasePath == "" {
|
|
return "", errors.New("database_path 不能为空")
|
|
}
|
|
if databasePath == ":memory:" {
|
|
return "file:cpa-ext-memory?mode=memory&cache=shared&_busy_timeout=5000&_foreign_keys=on", nil
|
|
}
|
|
if strings.HasPrefix(databasePath, "file:") {
|
|
separator := "?"
|
|
if strings.Contains(databasePath, "?") {
|
|
separator = "&"
|
|
}
|
|
return databasePath + separator + "_busy_timeout=5000&_foreign_keys=on&_journal_mode=WAL&_synchronous=NORMAL", nil
|
|
}
|
|
absolutePath, err := filepath.Abs(databasePath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("解析 database_path: %w", err)
|
|
}
|
|
if err := os.MkdirAll(filepath.Dir(absolutePath), 0o755); err != nil {
|
|
return "", fmt.Errorf("创建数据库目录: %w", err)
|
|
}
|
|
fileURL := (&url.URL{Scheme: "file", Path: filepath.ToSlash(absolutePath)}).String()
|
|
return fileURL + "?_busy_timeout=5000&_foreign_keys=on&_journal_mode=WAL&_synchronous=NORMAL", nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) Insert(ctx context.Context, record collection.Record) error {
|
|
_, err := r.db.ExecContext(ctx, `
|
|
INSERT OR IGNORE INTO usage_records (
|
|
request_id, execution_id, trace_id, requested_at, api_key, key_alias, model, reasoning_effort, service_tier, speed,
|
|
failed, executor_type, request_type, endpoint, input_tokens, output_tokens, reasoning_tokens,
|
|
cached_tokens, cache_read_tokens, cache_write_tokens, total_tokens,
|
|
ttft_ns, latency_ns, cost_micros, price_tier, fast_requested, fast_pricing_applied,
|
|
price_multiplier_numerator, price_multiplier_denominator, client_ip
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
|
record.RequestID, record.ExecutionID, record.TraceID,
|
|
record.RequestedAt.UTC().Format(time.RFC3339Nano), record.APIKey, record.KeyAlias,
|
|
record.Model, record.ReasoningEffort, record.ServiceTier, record.Speed, record.Failed,
|
|
record.ExecutorType, record.RequestType, record.Endpoint, record.InputTokens, record.OutputTokens,
|
|
record.ReasoningTokens, record.CachedTokens, record.CacheReadTokens,
|
|
record.CacheWriteTokens, record.TotalTokens, int64(record.TTFT), int64(record.Latency),
|
|
record.CostMicros, record.PriceTier, record.FastRequested, record.FastPricingApplied,
|
|
record.PriceMultiplierNumerator, record.PriceMultiplierDenominator, record.ClientIP,
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("写入用量记录: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) UpsertRequest(ctx context.Context, record collection.RequestRecord) error {
|
|
requestID := strings.TrimSpace(record.RequestID)
|
|
if requestID == "" {
|
|
return errors.New("request_id 不能为空")
|
|
}
|
|
requestedAt := record.RequestedAt
|
|
completedAt := record.CompletedAt
|
|
if requestedAt.IsZero() {
|
|
requestedAt = completedAt
|
|
}
|
|
if completedAt.IsZero() {
|
|
completedAt = requestedAt
|
|
}
|
|
latency := completedAt.Sub(requestedAt)
|
|
if latency < 0 {
|
|
latency = 0
|
|
}
|
|
_, err := r.db.ExecContext(ctx, `
|
|
INSERT INTO request_records (
|
|
request_id, trace_id, requested_at, completed_at, model, requested_model,
|
|
source_format, stream, outcome, status_code, error, endpoint, latency_ns
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(request_id) DO UPDATE SET
|
|
trace_id=excluded.trace_id,
|
|
requested_at=excluded.requested_at,
|
|
completed_at=excluded.completed_at,
|
|
model=excluded.model,
|
|
requested_model=excluded.requested_model,
|
|
source_format=excluded.source_format,
|
|
stream=excluded.stream,
|
|
outcome=excluded.outcome,
|
|
status_code=excluded.status_code,
|
|
error=excluded.error,
|
|
endpoint=excluded.endpoint,
|
|
latency_ns=excluded.latency_ns`,
|
|
requestID, record.TraceID, requestedAt.UTC().Format(time.RFC3339Nano),
|
|
completedAt.UTC().Format(time.RFC3339Nano), record.Model, record.RequestedModel,
|
|
record.SourceFormat, record.Stream, record.Outcome, record.StatusCode,
|
|
record.Error, record.Endpoint, int64(latency),
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("写入请求终态: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) ListRecent(ctx context.Context, limit int) ([]collection.Record, error) {
|
|
if limit < 1 {
|
|
return []collection.Record{}, nil
|
|
}
|
|
// A logical request can occupy two source rows until projection time, so read
|
|
// ahead before applying the caller's logical-request limit.
|
|
scanLimit := limit * 2
|
|
if scanLimit < limit {
|
|
scanLimit = limit
|
|
}
|
|
rows, err := r.db.QueryContext(ctx, `
|
|
WITH combined AS (
|
|
SELECT
|
|
u.request_id, u.execution_id, COALESCE(NULLIF(u.trace_id, ''), r.trace_id, '') AS trace_id,
|
|
u.requested_at, u.api_key, u.key_alias, u.model, u.reasoning_effort, u.service_tier, u.speed,
|
|
CASE
|
|
WHEN u.failed = 1 THEN 1
|
|
WHEN r.outcome IN ('failed', 'rejected', 'canceled') THEN 1
|
|
ELSE 0
|
|
END AS failed,
|
|
u.executor_type,
|
|
CASE WHEN r.request_id IS NOT NULL THEN CASE WHEN r.stream THEN 'SSE' ELSE 'JSON' END ELSE u.request_type END AS request_type,
|
|
COALESCE(NULLIF(u.endpoint, ''), r.endpoint, '') AS endpoint,
|
|
u.input_tokens, u.output_tokens, u.reasoning_tokens, u.cached_tokens,
|
|
u.cache_read_tokens, u.cache_write_tokens, u.total_tokens,
|
|
u.ttft_ns, u.latency_ns, u.cost_micros, u.price_tier, u.fast_requested,
|
|
u.fast_pricing_applied, u.price_multiplier_numerator, u.price_multiplier_denominator,
|
|
u.client_ip, COALESCE(r.outcome, '') AS outcome, COALESCE(r.status_code, 0) AS status_code,
|
|
COALESCE(r.error, '') AS error
|
|
FROM usage_records u
|
|
LEFT JOIN request_records r ON r.request_id = u.request_id
|
|
|
|
UNION ALL
|
|
|
|
SELECT
|
|
r.request_id, '', r.trace_id, r.requested_at, '', '',
|
|
COALESCE(NULLIF(r.model, ''), r.requested_model), '', '', '',
|
|
CASE WHEN r.outcome = 'succeeded' THEN 0 ELSE 1 END,
|
|
'', CASE WHEN r.stream THEN 'SSE' ELSE 'JSON' END, r.endpoint,
|
|
0, 0, 0, 0, 0, 0, 0, 0, r.latency_ns, NULL, '', 0, 0, 1, 1, '',
|
|
r.outcome, r.status_code, r.error
|
|
FROM request_records r
|
|
WHERE NOT EXISTS (SELECT 1 FROM usage_records u WHERE u.request_id = r.request_id)
|
|
)
|
|
SELECT request_id, execution_id, trace_id, requested_at, api_key, key_alias, model,
|
|
reasoning_effort, service_tier, speed, failed, executor_type, request_type, endpoint,
|
|
input_tokens, output_tokens, reasoning_tokens, cached_tokens, cache_read_tokens,
|
|
cache_write_tokens, total_tokens, ttft_ns, latency_ns, cost_micros, price_tier,
|
|
fast_requested, fast_pricing_applied, price_multiplier_numerator,
|
|
price_multiplier_denominator, client_ip, outcome, status_code, error
|
|
FROM combined
|
|
ORDER BY requested_at DESC
|
|
LIMIT ?`, scanLimit)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("查询最近用量记录: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
records := make([]collection.Record, 0, scanLimit)
|
|
for rows.Next() {
|
|
var record collection.Record
|
|
var requestedAt string
|
|
var failed bool
|
|
var ttftNS, latencyNS int64
|
|
var cost sql.NullInt64
|
|
if err := rows.Scan(
|
|
&record.RequestID, &record.ExecutionID, &record.TraceID,
|
|
&requestedAt, &record.APIKey, &record.KeyAlias, &record.Model,
|
|
&record.ReasoningEffort, &record.ServiceTier, &record.Speed, &failed, &record.ExecutorType,
|
|
&record.RequestType, &record.Endpoint, &record.InputTokens, &record.OutputTokens,
|
|
&record.ReasoningTokens, &record.CachedTokens, &record.CacheReadTokens,
|
|
&record.CacheWriteTokens, &record.TotalTokens, &ttftNS, &latencyNS,
|
|
&cost, &record.PriceTier, &record.FastRequested, &record.FastPricingApplied,
|
|
&record.PriceMultiplierNumerator, &record.PriceMultiplierDenominator, &record.ClientIP,
|
|
&record.Outcome, &record.StatusCode, &record.Error,
|
|
); err != nil {
|
|
return nil, fmt.Errorf("读取用量记录: %w", err)
|
|
}
|
|
record.RequestedAt, err = time.Parse(time.RFC3339Nano, requestedAt)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("解析用量时间: %w", err)
|
|
}
|
|
record.Failed = failed
|
|
record.TTFT = time.Duration(ttftNS)
|
|
record.Latency = time.Duration(latencyNS)
|
|
if cost.Valid {
|
|
value := cost.Int64
|
|
record.CostMicros = &value
|
|
}
|
|
records = append(records, record)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("遍历用量记录: %w", err)
|
|
}
|
|
records = mergeOrphanLifecycleRecords(records)
|
|
if len(records) > limit {
|
|
records = records[:limit]
|
|
}
|
|
return records, nil
|
|
}
|
|
|
|
type lifecycleMatch struct {
|
|
index int
|
|
distance time.Duration
|
|
ambiguous bool
|
|
}
|
|
|
|
func mergeOrphanLifecycleRecords(records []collection.Record) []collection.Record {
|
|
usageMatches := make(map[int]lifecycleMatch)
|
|
requestMatches := make(map[int]lifecycleMatch)
|
|
|
|
for usageIndex := range records {
|
|
if !isOrphanUsage(records[usageIndex]) {
|
|
continue
|
|
}
|
|
for requestIndex := range records {
|
|
if !isLifecycleOnly(records[requestIndex]) || !sameProjectionModel(records[usageIndex], records[requestIndex]) {
|
|
continue
|
|
}
|
|
distance := records[usageIndex].RequestedAt.Sub(records[requestIndex].RequestedAt)
|
|
if distance < 0 {
|
|
distance = -distance
|
|
}
|
|
if distance > orphanLifecycleMatchWindow {
|
|
continue
|
|
}
|
|
if current, exists := usageMatches[usageIndex]; exists {
|
|
usageMatches[usageIndex] = betterLifecycleMatch(current, requestIndex, distance)
|
|
} else {
|
|
usageMatches[usageIndex] = lifecycleMatch{index: requestIndex, distance: distance}
|
|
}
|
|
if current, exists := requestMatches[requestIndex]; exists {
|
|
requestMatches[requestIndex] = betterLifecycleMatch(current, usageIndex, distance)
|
|
} else {
|
|
requestMatches[requestIndex] = lifecycleMatch{index: usageIndex, distance: distance}
|
|
}
|
|
}
|
|
}
|
|
|
|
remove := make(map[int]struct{})
|
|
for usageIndex, requestMatch := range usageMatches {
|
|
if requestMatch.ambiguous {
|
|
continue
|
|
}
|
|
usageMatch, ok := requestMatches[requestMatch.index]
|
|
if !ok || usageMatch.ambiguous || usageMatch.index != usageIndex {
|
|
continue
|
|
}
|
|
records[usageIndex] = mergeLifecycleIntoUsage(records[usageIndex], records[requestMatch.index])
|
|
remove[requestMatch.index] = struct{}{}
|
|
}
|
|
|
|
if len(remove) == 0 {
|
|
return records
|
|
}
|
|
merged := make([]collection.Record, 0, len(records)-len(remove))
|
|
for index, record := range records {
|
|
if _, drop := remove[index]; !drop {
|
|
merged = append(merged, record)
|
|
}
|
|
}
|
|
return merged
|
|
}
|
|
|
|
func betterLifecycleMatch(current lifecycleMatch, candidate int, distance time.Duration) lifecycleMatch {
|
|
if distance+orphanLifecycleTieWindow < current.distance {
|
|
return lifecycleMatch{index: candidate, distance: distance}
|
|
}
|
|
if current.distance+orphanLifecycleTieWindow < distance {
|
|
return current
|
|
}
|
|
current.ambiguous = true
|
|
return current
|
|
}
|
|
|
|
func isOrphanUsage(record collection.Record) bool {
|
|
if strings.TrimSpace(record.RequestID) != "" {
|
|
return false
|
|
}
|
|
return record.CostMicros != nil || record.TotalTokens != 0 || record.InputTokens != 0 ||
|
|
record.OutputTokens != 0 || record.CacheReadTokens != 0 || record.CacheWriteTokens != 0 ||
|
|
strings.TrimSpace(record.APIKey) != "" || strings.TrimSpace(record.ExecutorType) != ""
|
|
}
|
|
|
|
func isLifecycleOnly(record collection.Record) bool {
|
|
return strings.TrimSpace(record.RequestID) != "" && strings.TrimSpace(record.ExecutionID) == "" &&
|
|
strings.TrimSpace(record.APIKey) == "" && record.CostMicros == nil && record.InputTokens == 0 &&
|
|
record.OutputTokens == 0 && record.CacheReadTokens == 0 && record.CacheWriteTokens == 0 &&
|
|
record.TotalTokens == 0 && strings.TrimSpace(record.Outcome) != ""
|
|
}
|
|
|
|
func sameProjectionModel(left, right collection.Record) bool {
|
|
leftModel := strings.ToLower(strings.TrimSpace(left.Model))
|
|
rightModel := strings.ToLower(strings.TrimSpace(right.Model))
|
|
return leftModel != "" && leftModel == rightModel
|
|
}
|
|
|
|
func mergeLifecycleIntoUsage(usage, lifecycle collection.Record) collection.Record {
|
|
usage.RequestID = lifecycle.RequestID
|
|
if usage.TraceID == "" {
|
|
usage.TraceID = lifecycle.TraceID
|
|
}
|
|
usage.Failed = usage.Failed || lifecycle.Failed
|
|
usage.RequestType = lifecycle.RequestType
|
|
usage.Endpoint = lifecycle.Endpoint
|
|
usage.Outcome = lifecycle.Outcome
|
|
usage.StatusCode = lifecycle.StatusCode
|
|
usage.Error = lifecycle.Error
|
|
if usage.Latency == 0 {
|
|
usage.Latency = lifecycle.Latency
|
|
}
|
|
return usage
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) Close() error {
|
|
return r.db.Close()
|
|
}
|