mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-28 00:35:32 +08:00
feat(database): extend schema with streaming and timing fields
Add new columns to proxy_request_logs table for enhanced usage tracking: - first_token_ms and duration_ms for performance metrics - provider_type and is_streaming for request classification - cost_multiplier for flexible pricing Update model pricing with accurate rates for Claude/GPT/Gemini models. Add ensure_model_pricing_seeded() call on database initialization. Add test for model pricing auto-seeding verification.
This commit is contained in:
@@ -95,6 +95,7 @@ impl Database {
|
||||
};
|
||||
db.create_tables()?;
|
||||
db.apply_schema_migrations()?;
|
||||
db.ensure_model_pricing_seeded()?;
|
||||
|
||||
Ok(db)
|
||||
}
|
||||
@@ -111,6 +112,7 @@ impl Database {
|
||||
conn: Mutex::new(conn),
|
||||
};
|
||||
db.create_tables()?;
|
||||
db.ensure_model_pricing_seeded()?;
|
||||
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
@@ -221,9 +221,14 @@ impl Database {
|
||||
cache_creation_cost_usd TEXT NOT NULL DEFAULT '0',
|
||||
total_cost_usd TEXT NOT NULL DEFAULT '0',
|
||||
latency_ms INTEGER NOT NULL,
|
||||
first_token_ms INTEGER,
|
||||
duration_ms INTEGER,
|
||||
status_code INTEGER NOT NULL,
|
||||
error_message TEXT,
|
||||
session_id TEXT,
|
||||
provider_type TEXT,
|
||||
is_streaming INTEGER NOT NULL DEFAULT 0,
|
||||
cost_multiplier TEXT NOT NULL DEFAULT '1.0',
|
||||
created_at INTEGER NOT NULL
|
||||
)",
|
||||
[],
|
||||
@@ -331,7 +336,7 @@ impl Database {
|
||||
Self::set_user_version(conn, 1)?;
|
||||
}
|
||||
1 => {
|
||||
log::info!("迁移数据库从 v1 到 v2(添加使用统计表和 provider 限制字段)");
|
||||
log::info!("迁移数据库从 v1 到 v2(添加使用统计表和完整字段)");
|
||||
Self::migrate_v1_to_v2(conn)?;
|
||||
Self::set_user_version(conn, 2)?;
|
||||
}
|
||||
@@ -420,9 +425,9 @@ impl Database {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// v1 -> v2 迁移:添加使用统计表和 provider 限制字段
|
||||
/// v1 -> v2 迁移:添加使用统计表和完整字段
|
||||
fn migrate_v1_to_v2(conn: &Connection) -> Result<(), AppError> {
|
||||
// 为 providers 表添加成本倍数和限制字段
|
||||
// providers 表字段
|
||||
Self::add_column_if_missing(
|
||||
conn,
|
||||
"providers",
|
||||
@@ -431,8 +436,9 @@ impl Database {
|
||||
)?;
|
||||
Self::add_column_if_missing(conn, "providers", "limit_daily_usd", "TEXT")?;
|
||||
Self::add_column_if_missing(conn, "providers", "limit_monthly_usd", "TEXT")?;
|
||||
Self::add_column_if_missing(conn, "providers", "provider_type", "TEXT")?;
|
||||
|
||||
// 创建新表(如果不存在)
|
||||
// proxy_request_logs 表(包含所有字段)
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS proxy_request_logs (
|
||||
request_id TEXT PRIMARY KEY,
|
||||
@@ -449,14 +455,37 @@ impl Database {
|
||||
cache_creation_cost_usd TEXT NOT NULL DEFAULT '0',
|
||||
total_cost_usd TEXT NOT NULL DEFAULT '0',
|
||||
latency_ms INTEGER NOT NULL,
|
||||
first_token_ms INTEGER,
|
||||
duration_ms INTEGER,
|
||||
status_code INTEGER NOT NULL,
|
||||
error_message TEXT,
|
||||
session_id TEXT,
|
||||
provider_type TEXT,
|
||||
is_streaming INTEGER NOT NULL DEFAULT 0,
|
||||
cost_multiplier TEXT NOT NULL DEFAULT '1.0',
|
||||
created_at INTEGER NOT NULL
|
||||
)",
|
||||
[],
|
||||
)?;
|
||||
|
||||
// 为已存在的表添加新字段
|
||||
Self::add_column_if_missing(conn, "proxy_request_logs", "provider_type", "TEXT")?;
|
||||
Self::add_column_if_missing(
|
||||
conn,
|
||||
"proxy_request_logs",
|
||||
"is_streaming",
|
||||
"INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
Self::add_column_if_missing(
|
||||
conn,
|
||||
"proxy_request_logs",
|
||||
"cost_multiplier",
|
||||
"TEXT NOT NULL DEFAULT '1.0'",
|
||||
)?;
|
||||
Self::add_column_if_missing(conn, "proxy_request_logs", "first_token_ms", "INTEGER")?;
|
||||
Self::add_column_if_missing(conn, "proxy_request_logs", "duration_ms", "INTEGER")?;
|
||||
|
||||
// model_pricing 表
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS model_pricing (
|
||||
model_id TEXT PRIMARY KEY,
|
||||
@@ -469,6 +498,7 @@ impl Database {
|
||||
[],
|
||||
)?;
|
||||
|
||||
// usage_daily_stats 表
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS usage_daily_stats (
|
||||
date TEXT NOT NULL,
|
||||
@@ -486,91 +516,120 @@ impl Database {
|
||||
[],
|
||||
)?;
|
||||
|
||||
// 插入默认模型定价数据
|
||||
// 清空并重新插入模型定价
|
||||
conn.execute("DELETE FROM model_pricing", [])
|
||||
.map_err(|e| AppError::Database(format!("清空模型定价失败: {e}")))?;
|
||||
Self::seed_model_pricing(conn)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 插入默认模型定价数据
|
||||
/// 格式: (model_id, display_name, input, output, cache_read, cache_creation)
|
||||
fn seed_model_pricing(conn: &Connection) -> Result<(), AppError> {
|
||||
let pricing_data = [
|
||||
// Claude 4.x series
|
||||
// Claude 4.5 系列
|
||||
(
|
||||
"claude-4.1-sonnet",
|
||||
"Claude 4.1 Sonnet",
|
||||
"3.0",
|
||||
"15.0",
|
||||
"0.3",
|
||||
"claude-opus-4-5",
|
||||
"Claude Opus 4.5",
|
||||
"5",
|
||||
"25",
|
||||
"0.50",
|
||||
"6.25",
|
||||
),
|
||||
(
|
||||
"claude-sonnet-4-5",
|
||||
"Claude Sonnet 4.5",
|
||||
"3",
|
||||
"15",
|
||||
"0.30",
|
||||
"3.75",
|
||||
),
|
||||
(
|
||||
"claude-4.1-haiku",
|
||||
"Claude 4.1 Haiku",
|
||||
"1.0",
|
||||
"5.0",
|
||||
"0.1",
|
||||
"claude-haiku-4-5",
|
||||
"Claude Haiku 4.5",
|
||||
"1",
|
||||
"5",
|
||||
"0.10",
|
||||
"1.25",
|
||||
),
|
||||
// Claude 4.1 系列
|
||||
(
|
||||
"claude-4.5-sonnet",
|
||||
"Claude 4.5 Sonnet",
|
||||
"3.5",
|
||||
"17.5",
|
||||
"0.35",
|
||||
"4.0",
|
||||
"claude-opus-4-1",
|
||||
"Claude Opus 4.1",
|
||||
"15",
|
||||
"75",
|
||||
"1.50",
|
||||
"18.75",
|
||||
),
|
||||
(
|
||||
"claude-4.5-haiku",
|
||||
"Claude 4.5 Haiku",
|
||||
"1.2",
|
||||
"6.0",
|
||||
"0.12",
|
||||
"1.5",
|
||||
"claude-sonnet-4-1",
|
||||
"Claude Sonnet 4.1",
|
||||
"3",
|
||||
"15",
|
||||
"0.30",
|
||||
"3.75",
|
||||
),
|
||||
// GPT 5.x series
|
||||
("gpt-5.0", "GPT-5.0", "4.0", "16.0", "0", "0"),
|
||||
("gpt-5.0-mini", "GPT-5.0 Mini", "0.3", "1.2", "0", "0"),
|
||||
("gpt-5.1", "GPT-5.1", "4.5", "18.0", "0", "0"),
|
||||
("gpt-5.1-mini", "GPT-5.1 Mini", "0.4", "1.5", "0", "0"),
|
||||
// Gemini 2.5 / 3.x
|
||||
// Claude 3.7 系列
|
||||
(
|
||||
"claude-sonnet-3-7",
|
||||
"Claude Sonnet 3.7",
|
||||
"3",
|
||||
"15",
|
||||
"0.30",
|
||||
"3.75",
|
||||
),
|
||||
// Claude 3.5 系列
|
||||
(
|
||||
"claude-sonnet-3-5",
|
||||
"Claude Sonnet 3.5",
|
||||
"3",
|
||||
"15",
|
||||
"0.30",
|
||||
"3.75",
|
||||
),
|
||||
(
|
||||
"claude-haiku-3-5",
|
||||
"Claude Haiku 3.5",
|
||||
"0.80",
|
||||
"4",
|
||||
"0.08",
|
||||
"1",
|
||||
),
|
||||
// GPT-5 系列
|
||||
("gpt-5", "GPT-5", "1.25", "10", "0.125", "0"),
|
||||
("gpt-5.1", "GPT-5.1", "1.25", "10", "0.125", "0"),
|
||||
// Gemini 3 系列
|
||||
(
|
||||
"gemini-3-pro-preview",
|
||||
"Gemini 3 Pro Preview",
|
||||
"2",
|
||||
"12",
|
||||
"0",
|
||||
"0",
|
||||
),
|
||||
// Gemini 2.5 系列
|
||||
(
|
||||
"gemini-2.5-pro",
|
||||
"Gemini 2.5 Pro",
|
||||
"1.5",
|
||||
"6.0",
|
||||
"0.35",
|
||||
"1.5",
|
||||
"1.25",
|
||||
"10",
|
||||
"0.125",
|
||||
"0",
|
||||
),
|
||||
(
|
||||
"gemini-2.5-flash",
|
||||
"Gemini 2.5 Flash",
|
||||
"0.12",
|
||||
"0.5",
|
||||
"0.3",
|
||||
"2.5",
|
||||
"0.03",
|
||||
"0.12",
|
||||
),
|
||||
(
|
||||
"gemini-3.0-pro",
|
||||
"Gemini 3.0 Pro",
|
||||
"2.0",
|
||||
"8.0",
|
||||
"0.4",
|
||||
"2.0",
|
||||
),
|
||||
(
|
||||
"gemini-3.0-flash",
|
||||
"Gemini 3.0 Flash",
|
||||
"0.2",
|
||||
"0.8",
|
||||
"0.05",
|
||||
"0.2",
|
||||
"0",
|
||||
),
|
||||
];
|
||||
|
||||
for (model_id, display_name, input, output, cache_read, cache_creation) in pricing_data {
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO model_pricing (
|
||||
"INSERT OR REPLACE INTO model_pricing (
|
||||
model_id, display_name, input_cost_per_million, output_cost_per_million,
|
||||
cache_read_cost_per_million, cache_creation_cost_per_million
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
||||
|
||||
@@ -273,3 +273,62 @@ fn dry_run_validates_schema_compatibility() {
|
||||
"Dry-run should succeed with provider data: {result:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_pricing_is_seeded_on_init() {
|
||||
let db = Database::memory().expect("create memory db");
|
||||
|
||||
let conn = db.conn.lock().expect("lock conn");
|
||||
|
||||
let count: i64 = conn
|
||||
.query_row("SELECT COUNT(*) FROM model_pricing", [], |row| row.get(0))
|
||||
.expect("count pricing");
|
||||
|
||||
assert!(
|
||||
count > 0,
|
||||
"模型定价数据应该在初始化时自动填充,实际数量: {}",
|
||||
count
|
||||
);
|
||||
|
||||
// 验证包含 Claude 模型
|
||||
let claude_count: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM model_pricing WHERE model_id LIKE 'claude-%'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.expect("check claude");
|
||||
assert!(
|
||||
claude_count > 0,
|
||||
"应该包含 Claude 模型定价,实际数量: {}",
|
||||
claude_count
|
||||
);
|
||||
|
||||
// 验证包含 GPT 模型
|
||||
let gpt_count: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM model_pricing WHERE model_id LIKE 'gpt-%'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.expect("check gpt");
|
||||
assert!(
|
||||
gpt_count > 0,
|
||||
"应该包含 GPT 模型定价,实际数量: {}",
|
||||
gpt_count
|
||||
);
|
||||
|
||||
// 验证包含 Gemini 模型
|
||||
let gemini_count: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM model_pricing WHERE model_id LIKE 'gemini-%'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.expect("check gemini");
|
||||
assert!(
|
||||
gemini_count > 0,
|
||||
"应该包含 Gemini 模型定价,实际数量: {}",
|
||||
gemini_count
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user