chuan
6a461bf1f3
- 使用 SQLite 保存关联后的请求生命周期和用量记录 - 支持长上下文阶梯价格和可配置的 Fast 计费倍率 - 添加管理接口和可配置列的用量面板 - 保留失败、取消、重试和 compact 请求，便于计费核对
310 lines
12 KiB
Go
310 lines
12 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
|
|
);
|
|
`
|
|
|
|
// 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
|
|
}
|
|
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 ?`, limit)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("查询最近用量记录: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
records := make([]collection.Record, 0, limit)
|
|
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)
|
|
}
|
|
return records, nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) Close() error {
|
|
return r.db.Close()
|
|
}
|