admin/CC-Switch
1
0
Fork 0
mirror of https://github.com/farion1231/cc-switch.git synced 2026-07-28 08:44:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(db): add usage tracking schema and types

Browse Source 
Add database tables for proxy request logs and model pricing.
Extend Provider and error types to support usage statistics.
This commit is contained in:
YoVinchen
2025-12-03 00:16:40 +08:00
parent d12347d837
commit c0122d717c
22 changed files with 2815 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/src/commands/usage.rs
+152
View File
@@ -0,0 +1,152 @@
//! 使用统计相关命令
use crate::database::Database;
use crate::error::AppError;
use crate::services::usage_stats::*;
use tauri::State;
/// 获取使用量汇总
#[tauri::command]
pub fn get_usage_summary(
db: State<'_, Database>,
start_date: Option<i64>,
end_date: Option<i64>,
) -> Result<UsageSummary, AppError> {
db.get_usage_summary(start_date, end_date)
}
/// 获取每日趋势
#[tauri::command]
pub fn get_usage_trends(
db: State<'_, Database>,
days: u32,
) -> Result<Vec<DailyStats>, AppError> {
db.get_daily_trends(days)
}
/// 获取 Provider 统计
#[tauri::command]
pub fn get_provider_stats(db: State<'_, Database>) -> Result<Vec<ProviderStats>, AppError> {
db.get_provider_stats()
}
/// 获取模型统计
#[tauri::command]
pub fn get_model_stats(db: State<'_, Database>) -> Result<Vec<ModelStats>, AppError> {
db.get_model_stats()
}
/// 获取请求日志列表
#[tauri::command]
pub fn get_request_logs(
db: State<'_, Database>,
provider_id: Option<String>,
model: Option<String>,
status_code: Option<u16>,
start_date: Option<i64>,
end_date: Option<i64>,
limit: u32,
offset: u32,
) -> Result<Vec<RequestLogDetail>, AppError> {
let filters = LogFilters {
provider_id,
model,
status_code,
start_date,
end_date,
};
db.get_request_logs(&filters, limit, offset)
}
/// 获取单个请求详情
#[tauri::command]
pub fn get_request_detail(
db: State<'_, Database>,
request_id: String,
) -> Result<Option<RequestLogDetail>, AppError> {
db.get_request_detail(&request_id)
}
/// 获取模型定价列表
#[tauri::command]
pub fn get_model_pricing(db: State<'_, Database>) -> Result<Vec<ModelPricingInfo>, AppError> {
let conn = crate::database::lock_conn!(db.conn);
let mut stmt = conn.prepare(
"SELECT model_id, display_name, input_cost_per_million, output_cost_per_million,
cache_read_cost_per_million, cache_creation_cost_per_million
FROM model_pricing
ORDER BY display_name"
)?;
let rows = stmt.query_map([], |row| {
Ok(ModelPricingInfo {
model_id: row.get(0)?,
display_name: row.get(1)?,
input_cost_per_million: row.get(2)?,
output_cost_per_million: row.get(3)?,
cache_read_cost_per_million: row.get(4)?,
cache_creation_cost_per_million: row.get(5)?,
})
})?;
let mut pricing = Vec::new();
for row in rows {
pricing.push(row?);
}
Ok(pricing)
}
/// 更新模型定价
#[tauri::command]
pub fn update_model_pricing(
db: State<'_, Database>,
model_id: String,
display_name: String,
input_cost: String,
output_cost: String,
cache_read_cost: String,
cache_creation_cost: String,
) -> Result<(), AppError> {
let conn = crate::database::lock_conn!(db.conn);
conn.execute(
"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)",
rusqlite::params![
model_id,
display_name,
input_cost,
output_cost,
cache_read_cost,
cache_creation_cost
],
)
.map_err(|e| AppError::Database(format!("更新模型定价失败: {e}")))?;
Ok(())
}
/// 检查 Provider 使用限额
#[tauri::command]
pub fn check_provider_limits(
db: State<'_, Database>,
provider_id: String,
app_type: String,
) -> Result<crate::services::usage_stats::ProviderLimitStatus, AppError> {
db.check_provider_limits(&provider_id, &app_type)
}
/// 模型定价信息
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ModelPricingInfo {
pub model_id: String,
pub display_name: String,
pub input_cost_per_million: String,
pub output_cost_per_million: String,
pub cache_read_cost_per_million: String,
pub cache_creation_cost_per_million: String,
}
src-tauri/src/database/dao/providers.rs
+51
View File
@@ -123,6 +123,57 @@ impl Database {
}
}
/// 根据 ID 获取单个供应商
pub fn get_provider_by_id(
&self,
id: &str,
app_type: &str,
) -> Result<Option<Provider>, AppError> {
let conn = lock_conn!(self.conn);
let result = conn.query_row(
"SELECT name, settings_config, website_url, category, created_at, sort_index, notes, icon, icon_color, meta, is_proxy_target
FROM providers WHERE id = ?1 AND app_type = ?2",
params![id, app_type],
|row| {
let name: String = row.get(0)?;
let settings_config_str: String = row.get(1)?;
let website_url: Option<String> = row.get(2)?;
let category: Option<String> = row.get(3)?;
let created_at: Option<i64> = row.get(4)?;
let sort_index: Option<usize> = row.get(5)?;
let notes: Option<String> = row.get(6)?;
let icon: Option<String> = row.get(7)?;
let icon_color: Option<String> = row.get(8)?;
let meta_str: String = row.get(9)?;
let is_proxy_target: bool = row.get(10)?;
let settings_config = serde_json::from_str(&settings_config_str).unwrap_or(serde_json::Value::Null);
let meta: ProviderMeta = serde_json::from_str(&meta_str).unwrap_or_default();
Ok(Provider {
id: id.to_string(),
name,
settings_config,
website_url,
category,
created_at,
sort_index,
notes,
meta: Some(meta),
icon,
icon_color,
is_proxy_target: Some(is_proxy_target),
})
},
);
match result {
Ok(provider) => Ok(Some(provider)),
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
Err(e) => Err(AppError::Database(e.to_string())),
}
}
/// 获取代理目标供应商 ID
pub fn get_proxy_target_provider(&self, app_type: &str) -> Result<Option<String>, AppError> {
let conn = lock_conn!(self.conn);
src-tauri/src/database/mod.rs
+1 -1
View File
@@ -44,7 +44,7 @@ const DB_BACKUP_RETAIN: usize = 10;
/// 当前 Schema 版本号
/// 每次修改表结构时递增，并在 schema.rs 中添加相应的迁移逻辑
pub(crate) const SCHEMA_VERSION: i32 = 1;
pub(crate) const SCHEMA_VERSION: i32 = 2;
/// 安全地序列化 JSON，避免 unwrap panic
pub(crate) fn to_json_string<T: Serialize>(value: &T) -> Result<String, AppError> {
src-tauri/src/database/schema.rs
+287 -1
View File
@@ -204,6 +204,100 @@ impl Database {
)
.map_err(|e| AppError::Database(e.to_string()))?;
// 11. Proxy Request Logs 表 (详细请求日志)
conn.execute(
"CREATE TABLE IF NOT EXISTS proxy_request_logs (
request_id TEXT PRIMARY KEY,
provider_id TEXT NOT NULL,
app_type TEXT NOT NULL,
model TEXT NOT NULL,
input_tokens INTEGER NOT NULL DEFAULT 0,
output_tokens INTEGER NOT NULL DEFAULT 0,
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
input_cost_usd TEXT NOT NULL DEFAULT '0',
output_cost_usd TEXT NOT NULL DEFAULT '0',
cache_read_cost_usd TEXT NOT NULL DEFAULT '0',
cache_creation_cost_usd TEXT NOT NULL DEFAULT '0',
total_cost_usd TEXT NOT NULL DEFAULT '0',
latency_ms INTEGER NOT NULL,
status_code INTEGER NOT NULL,
error_message TEXT,
session_id TEXT,
created_at INTEGER NOT NULL
)",
[],
)
.map_err(|e| AppError::Database(e.to_string()))?;
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_request_logs_provider
ON proxy_request_logs(provider_id, app_type)",
[],
)
.map_err(|e| AppError::Database(e.to_string()))?;
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_request_logs_created_at
ON proxy_request_logs(created_at)",
[],
)
.map_err(|e| AppError::Database(e.to_string()))?;
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_request_logs_model
ON proxy_request_logs(model)",
[],
)
.map_err(|e| AppError::Database(e.to_string()))?;
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_request_logs_session
ON proxy_request_logs(session_id)",
[],
)
.map_err(|e| AppError::Database(e.to_string()))?;
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_request_logs_status
ON proxy_request_logs(status_code)",
[],
)
.map_err(|e| AppError::Database(e.to_string()))?;
// 12. Model Pricing 表 (模型定价)
conn.execute(
"CREATE TABLE IF NOT EXISTS model_pricing (
model_id TEXT PRIMARY KEY,
display_name TEXT NOT NULL,
input_cost_per_million TEXT NOT NULL,
output_cost_per_million TEXT NOT NULL,
cache_read_cost_per_million TEXT NOT NULL DEFAULT '0',
cache_creation_cost_per_million TEXT NOT NULL DEFAULT '0'
)",
[],
)
.map_err(|e| AppError::Database(e.to_string()))?;
// 13. Usage Daily Stats 表 (每日聚合统计)
conn.execute(
"CREATE TABLE IF NOT EXISTS usage_daily_stats (
date TEXT NOT NULL,
provider_id TEXT NOT NULL,
app_type TEXT NOT NULL,
model TEXT NOT NULL,
request_count INTEGER NOT NULL DEFAULT 0,
total_input_tokens INTEGER NOT NULL DEFAULT 0,
total_output_tokens INTEGER NOT NULL DEFAULT 0,
total_cost_usd TEXT NOT NULL DEFAULT '0',
success_count INTEGER NOT NULL DEFAULT 0,
error_count INTEGER NOT NULL DEFAULT 0,
PRIMARY KEY (date, provider_id, app_type, model)
)",
[],
)
.map_err(|e| AppError::Database(e.to_string()))?;
Ok(())
}
@@ -234,7 +328,12 @@ impl Database {
0 => {
log::info!("检测到 user_version=0，迁移到 1（补齐缺失列并设置版本）");
Self::migrate_v0_to_v1(conn)?;
Self::set_user_version(conn, SCHEMA_VERSION)?;
Self::set_user_version(conn, 1)?;
}
1 => {
log::info!("迁移数据库从 v1 到 v2（添加使用统计表和 provider 限制字段）");
Self::migrate_v1_to_v2(conn)?;
Self::set_user_version(conn, 2)?;
}
_ => {
return Err(AppError::Database(format!(
@@ -321,6 +420,193 @@ impl Database {
Ok(())
}
/// v1 -> v2 迁移：添加使用统计表和 provider 限制字段
fn migrate_v1_to_v2(conn: &Connection) -> Result<(), AppError> {
// 为 providers 表添加成本倍数和限制字段
Self::add_column_if_missing(
conn,
"providers",
"cost_multiplier",
"TEXT NOT NULL DEFAULT '1.0'",
)?;
Self::add_column_if_missing(conn, "providers", "limit_daily_usd", "TEXT")?;
Self::add_column_if_missing(conn, "providers", "limit_monthly_usd", "TEXT")?;
// 创建新表(如果不存在)
conn.execute(
"CREATE TABLE IF NOT EXISTS proxy_request_logs (
request_id TEXT PRIMARY KEY,
provider_id TEXT NOT NULL,
app_type TEXT NOT NULL,
model TEXT NOT NULL,
input_tokens INTEGER NOT NULL DEFAULT 0,
output_tokens INTEGER NOT NULL DEFAULT 0,
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
input_cost_usd TEXT NOT NULL DEFAULT '0',
output_cost_usd TEXT NOT NULL DEFAULT '0',
cache_read_cost_usd TEXT NOT NULL DEFAULT '0',
cache_creation_cost_usd TEXT NOT NULL DEFAULT '0',
total_cost_usd TEXT NOT NULL DEFAULT '0',
latency_ms INTEGER NOT NULL,
status_code INTEGER NOT NULL,
error_message TEXT,
session_id TEXT,
created_at INTEGER NOT NULL
)",
[],
)?;
conn.execute(
"CREATE TABLE IF NOT EXISTS model_pricing (
model_id TEXT PRIMARY KEY,
display_name TEXT NOT NULL,
input_cost_per_million TEXT NOT NULL,
output_cost_per_million TEXT NOT NULL,
cache_read_cost_per_million TEXT NOT NULL DEFAULT '0',
cache_creation_cost_per_million TEXT NOT NULL DEFAULT '0'
)",
[],
)?;
conn.execute(
"CREATE TABLE IF NOT EXISTS usage_daily_stats (
date TEXT NOT NULL,
provider_id TEXT NOT NULL,
app_type TEXT NOT NULL,
model TEXT NOT NULL,
request_count INTEGER NOT NULL DEFAULT 0,
total_input_tokens INTEGER NOT NULL DEFAULT 0,
total_output_tokens INTEGER NOT NULL DEFAULT 0,
total_cost_usd TEXT NOT NULL DEFAULT '0',
success_count INTEGER NOT NULL DEFAULT 0,
error_count INTEGER NOT NULL DEFAULT 0,
PRIMARY KEY (date, provider_id, app_type, model)
)",
[],
)?;
// 插入默认模型定价数据
Self::seed_model_pricing(conn)?;
Ok(())
}
/// 插入默认模型定价数据
fn seed_model_pricing(conn: &Connection) -> Result<(), AppError> {
let pricing_data = [
// Claude 4.x series
(
"claude-4.1-sonnet",
"Claude 4.1 Sonnet",
"3.0",
"15.0",
"0.3",
"3.75",
),
(
"claude-4.1-haiku",
"Claude 4.1 Haiku",
"1.0",
"5.0",
"0.1",
"1.25",
),
(
"claude-4.5-sonnet",
"Claude 4.5 Sonnet",
"3.5",
"17.5",
"0.35",
"4.0",
),
(
"claude-4.5-haiku",
"Claude 4.5 Haiku",
"1.2",
"6.0",
"0.12",
"1.5",
),
// 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
(
"gemini-2.5-pro",
"Gemini 2.5 Pro",
"1.5",
"6.0",
"0.35",
"1.5",
),
(
"gemini-2.5-flash",
"Gemini 2.5 Flash",
"0.12",
"0.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",
),
];
for (model_id, display_name, input, output, cache_read, cache_creation) in pricing_data {
conn.execute(
"INSERT OR IGNORE 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)",
rusqlite::params![
model_id,
display_name,
input,
output,
cache_read,
cache_creation
],
)
.map_err(|e| AppError::Database(format!("插入模型定价失败: {e}")))?;
}
log::info!("已插入 {} 条默认模型定价数据", pricing_data.len());
Ok(())
}
/// 确保模型定价表具备默认数据
pub fn ensure_model_pricing_seeded(&self) -> Result<(), AppError> {
let conn = lock_conn!(self.conn);
Self::ensure_model_pricing_seeded_on_conn(&conn)
}
fn ensure_model_pricing_seeded_on_conn(conn: &Connection) -> Result<(), AppError> {
let count: i64 = conn
.query_row("SELECT COUNT(*) FROM model_pricing", [], |row| row.get(0))
.map_err(|e| AppError::Database(format!("统计模型定价数据失败: {e}")))?;
if count == 0 {
Self::seed_model_pricing(conn)?;
}
Ok(())
}
// --- 辅助方法 ---
pub(crate) fn get_user_version(conn: &Connection) -> Result<i32, AppError> {
src-tauri/src/error.rs
+15
View File
@@ -91,12 +91,27 @@ impl<T> From<PoisonError<T>> for AppError {
}
}
impl From<rusqlite::Error> for AppError {
fn from(err: rusqlite::Error) -> Self {
Self::Database(err.to_string())
}
}
impl From<AppError> for String {
fn from(err: AppError) -> Self {
err.to_string()
}
}
impl serde::Serialize for AppError {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.serialize_str(&self.to_string())
}
}
/// 格式化为 JSON 错误字符串，前端可解析为结构化错误
pub fn format_skill_error(
code: &str,
src-tauri/src/provider.rs
+9
View File
@@ -156,6 +156,15 @@ pub struct ProviderMeta {
skip_serializing_if = "Option::is_none"
)]
pub partner_promotion_key: Option<String>,
/// 成本倍数（用于计算实际成本）
#[serde(rename = "costMultiplier", skip_serializing_if = "Option::is_none")]
pub cost_multiplier: Option<String>,
/// 每日消费限额（USD）
#[serde(rename = "limitDailyUsd", skip_serializing_if = "Option::is_none")]
pub limit_daily_usd: Option<String>,
/// 每月消费限额（USD）
#[serde(rename = "limitMonthlyUsd", skip_serializing_if = "Option::is_none")]
pub limit_monthly_usd: Option<String>,
}
impl ProviderManager {
src-tauri/src/proxy/usage/calculator.rs
+183
View File
@@ -0,0 +1,183 @@
//! Cost Calculator - 计算 API 请求成本
//!
//! 使用高精度 Decimal 类型避免浮点数精度问题
use super::parser::TokenUsage;
use rust_decimal::Decimal;
use std::str::FromStr;
/// 成本明细
#[derive(Debug, Clone)]
pub struct CostBreakdown {
pub input_cost: Decimal,
pub output_cost: Decimal,
pub cache_read_cost: Decimal,
pub cache_creation_cost: Decimal,
pub total_cost: Decimal,
}
/// 模型定价信息
#[derive(Debug, Clone)]
pub struct ModelPricing {
pub input_cost_per_million: Decimal,
pub output_cost_per_million: Decimal,
pub cache_read_cost_per_million: Decimal,
pub cache_creation_cost_per_million: Decimal,
}
/// 成本计算器
pub struct CostCalculator;
impl CostCalculator {
/// 计算请求成本
///
/// # 参数
/// - `usage`: Token 使用量
/// - `pricing`: 模型定价
/// - `cost_multiplier`: 成本倍数 (provider 自定义)
pub fn calculate(
usage: &TokenUsage,
pricing: &ModelPricing,
cost_multiplier: Decimal,
) -> CostBreakdown {
let million = Decimal::from(1_000_000);
let input_cost = Decimal::from(usage.input_tokens) * pricing.input_cost_per_million
/ million
* cost_multiplier;
let output_cost = Decimal::from(usage.output_tokens) * pricing.output_cost_per_million
/ million
* cost_multiplier;
let cache_read_cost = Decimal::from(usage.cache_read_tokens)
* pricing.cache_read_cost_per_million
/ million
* cost_multiplier;
let cache_creation_cost = Decimal::from(usage.cache_creation_tokens)
* pricing.cache_creation_cost_per_million
/ million
* cost_multiplier;
let total_cost = input_cost + output_cost + cache_read_cost + cache_creation_cost;
CostBreakdown {
input_cost,
output_cost,
cache_read_cost,
cache_creation_cost,
total_cost,
}
}
/// 尝试计算成本，如果模型未知则返回 None
pub fn try_calculate(
usage: &TokenUsage,
pricing: Option<&ModelPricing>,
cost_multiplier: Decimal,
) -> Option<CostBreakdown> {
pricing.map(|p| Self::calculate(usage, p, cost_multiplier))
}
}
impl ModelPricing {
/// 从字符串创建定价信息
pub fn from_strings(
input: &str,
output: &str,
cache_read: &str,
cache_creation: &str,
) -> Result<Self, rust_decimal::Error> {
Ok(Self {
input_cost_per_million: Decimal::from_str(input)?,
output_cost_per_million: Decimal::from_str(output)?,
cache_read_cost_per_million: Decimal::from_str(cache_read)?,
cache_creation_cost_per_million: Decimal::from_str(cache_creation)?,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cost_calculation() {
let usage = TokenUsage {
input_tokens: 1000,
output_tokens: 500,
cache_read_tokens: 200,
cache_creation_tokens: 100,
};
let pricing = ModelPricing::from_strings("3.0", "15.0", "0.3", "3.75").unwrap();
let multiplier = Decimal::from_str("1.0").unwrap();
let cost = CostCalculator::calculate(&usage, &pricing, multiplier);
// input: 1000 * 3.0 / 1M = 0.003
assert_eq!(cost.input_cost, Decimal::from_str("0.003").unwrap());
// output: 500 * 15.0 / 1M = 0.0075
assert_eq!(cost.output_cost, Decimal::from_str("0.0075").unwrap());
// cache_read: 200 * 0.3 / 1M = 0.00006
assert_eq!(cost.cache_read_cost, Decimal::from_str("0.00006").unwrap());
// cache_creation: 100 * 3.75 / 1M = 0.000375
assert_eq!(
cost.cache_creation_cost,
Decimal::from_str("0.000375").unwrap()
);
// total: 0.003 + 0.0075 + 0.00006 + 0.000375 = 0.010935
assert_eq!(cost.total_cost, Decimal::from_str("0.010935").unwrap());
}
#[test]
fn test_cost_multiplier() {
let usage = TokenUsage {
input_tokens: 1000,
output_tokens: 0,
cache_read_tokens: 0,
cache_creation_tokens: 0,
};
let pricing = ModelPricing::from_strings("3.0", "15.0", "0", "0").unwrap();
let multiplier = Decimal::from_str("1.5").unwrap();
let cost = CostCalculator::calculate(&usage, &pricing, multiplier);
// input: 1000 * 3.0 / 1M * 1.5 = 0.0045
assert_eq!(cost.input_cost, Decimal::from_str("0.0045").unwrap());
assert_eq!(cost.total_cost, Decimal::from_str("0.0045").unwrap());
}
#[test]
fn test_unknown_model_handling() {
let usage = TokenUsage {
input_tokens: 1000,
output_tokens: 500,
cache_read_tokens: 0,
cache_creation_tokens: 0,
};
let multiplier = Decimal::from_str("1.0").unwrap();
let cost = CostCalculator::try_calculate(&usage, None, multiplier);
assert!(cost.is_none());
}
#[test]
fn test_decimal_precision() {
let usage = TokenUsage {
input_tokens: 1,
output_tokens: 1,
cache_read_tokens: 1,
cache_creation_tokens: 1,
};
let pricing = ModelPricing::from_strings("0.075", "0.3", "0.01875", "0.075").unwrap();
let multiplier = Decimal::from_str("1.0").unwrap();
let cost = CostCalculator::calculate(&usage, &pricing, multiplier);
// 验证高精度计算
assert!(cost.total_cost > Decimal::ZERO);
assert!(cost.total_cost.to_string().len() > 2); // 确保保留了小数位
}
}
src-tauri/src/proxy/usage/logger.rs
+271
View File
@@ -0,0 +1,271 @@
//! Usage Logger - 记录 API 请求使用情况
use super::calculator::{CostBreakdown, CostCalculator, ModelPricing};
use super::parser::TokenUsage;
use crate::database::Database;
use crate::error::AppError;
use rust_decimal::Decimal;
use std::time::SystemTime;
/// 请求日志
#[derive(Debug, Clone)]
pub struct RequestLog {
pub request_id: String,
pub provider_id: String,
pub app_type: String,
pub model: String,
pub usage: TokenUsage,
pub cost: Option<CostBreakdown>,
pub latency_ms: u64,
pub status_code: u16,
pub error_message: Option<String>,
pub session_id: Option<String>,
}
/// 使用量记录器
pub struct UsageLogger<'a> {
db: &'a Database,
}
impl<'a> UsageLogger<'a> {
pub fn new(db: &'a Database) -> Self {
Self { db }
}
/// 记录成功的请求
pub fn log_request(&self, log: &RequestLog) -> Result<(), AppError> {
let conn = crate::database::lock_conn!(self.db.conn);
let (input_cost, output_cost, cache_read_cost, cache_creation_cost, total_cost) =
if let Some(cost) = &log.cost {
(
cost.input_cost.to_string(),
cost.output_cost.to_string(),
cost.cache_read_cost.to_string(),
cost.cache_creation_cost.to_string(),
cost.total_cost.to_string(),
)
} else {
(
"0".to_string(),
"0".to_string(),
"0".to_string(),
"0".to_string(),
"0".to_string(),
)
};
let created_at = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap()
.as_secs() as i64;
conn.execute(
"INSERT INTO proxy_request_logs (
request_id, provider_id, app_type, model,
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd,
latency_ms, status_code, error_message, session_id, created_at
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
rusqlite::params![
log.request_id,
log.provider_id,
log.app_type,
log.model,
log.usage.input_tokens,
log.usage.output_tokens,
log.usage.cache_read_tokens,
log.usage.cache_creation_tokens,
input_cost,
output_cost,
cache_read_cost,
cache_creation_cost,
total_cost,
log.latency_ms as i64,
log.status_code as i64,
log.error_message,
log.session_id,
created_at,
],
)
.map_err(|e| AppError::Database(format!("记录请求日志失败: {e}")))?;
Ok(())
}
/// 记录失败的请求
pub fn log_error(
&self,
request_id: String,
provider_id: String,
app_type: String,
model: String,
status_code: u16,
error_message: String,
latency_ms: u64,
) -> Result<(), AppError> {
let log = RequestLog {
request_id,
provider_id,
app_type,
model,
usage: TokenUsage::default(),
cost: None,
latency_ms,
status_code,
error_message: Some(error_message),
session_id: None,
};
self.log_request(&log)
}
/// 获取模型定价
pub fn get_model_pricing(&self, model_id: &str) -> Result<Option<ModelPricing>, AppError> {
let conn = crate::database::lock_conn!(self.db.conn);
let result = conn.query_row(
"SELECT input_cost_per_million, output_cost_per_million,
cache_read_cost_per_million, cache_creation_cost_per_million
FROM model_pricing WHERE model_id = ?1",
[model_id],
|row| {
Ok(ModelPricing::from_strings(
&row.get::<_, String>(0)?,
&row.get::<_, String>(1)?,
&row.get::<_, String>(2)?,
&row.get::<_, String>(3)?,
))
},
);
match result {
Ok(Ok(pricing)) => Ok(Some(pricing)),
Ok(Err(e)) => Err(AppError::Database(format!("解析定价数据失败: {e}"))),
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
Err(e) => Err(AppError::Database(format!("查询定价失败: {e}"))),
}
}
/// 计算并记录请求
pub fn log_with_calculation(
&self,
request_id: String,
provider_id: String,
app_type: String,
model: String,
usage: TokenUsage,
cost_multiplier: Decimal,
latency_ms: u64,
status_code: u16,
session_id: Option<String>,
) -> Result<(), AppError> {
let pricing = self.get_model_pricing(&model)?;
if pricing.is_none() {
log::warn!("模型 {} 的定价信息未找到，成本将记录为 0", model);
}
let cost = CostCalculator::try_calculate(&usage, pricing.as_ref(), cost_multiplier);
let log = RequestLog {
request_id,
provider_id,
app_type,
model,
usage,
cost,
latency_ms,
status_code,
error_message: None,
session_id,
};
self.log_request(&log)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_log_request() -> Result<(), AppError> {
let db = Database::memory()?;
// 插入测试定价
{
let conn = crate::database::lock_conn!(db.conn);
conn.execute(
"INSERT INTO model_pricing (model_id, display_name, input_cost_per_million, output_cost_per_million)
VALUES ('test-model', 'Test Model', '3.0', '15.0')",
[],
)
.unwrap();
}
let logger = UsageLogger::new(&db);
let usage = TokenUsage {
input_tokens: 1000,
output_tokens: 500,
cache_read_tokens: 0,
cache_creation_tokens: 0,
};
logger
.log_with_calculation(
"req-123".to_string(),
"provider-1".to_string(),
"claude".to_string(),
"test-model".to_string(),
usage,
Decimal::from(1),
100,
200,
None,
)?;
// 验证记录已插入
let conn = crate::database::lock_conn!(db.conn);
let count: i64 = conn
.query_row(
"SELECT COUNT(*) FROM proxy_request_logs WHERE request_id = 'req-123'",
[],
|row| row.get(0),
)
.unwrap();
assert_eq!(count, 1);
Ok(())
}
#[test]
fn test_log_error() -> Result<(), AppError> {
let db = Database::memory()?;
let logger = UsageLogger::new(&db);
logger
.log_error(
"req-error".to_string(),
"provider-1".to_string(),
"claude".to_string(),
"unknown-model".to_string(),
500,
"Internal Server Error".to_string(),
50,
)?;
// 验证错误记录已插入
let conn = crate::database::lock_conn!(db.conn);
let (status, error): (i64, Option<String>) = conn
.query_row(
"SELECT status_code, error_message FROM proxy_request_logs WHERE request_id = 'req-error'",
[],
|row| Ok((row.get(0)?, row.get(1)?)),
)
.unwrap();
assert_eq!(status, 500);
assert_eq!(error, Some("Internal Server Error".to_string()));
Ok(())
}
}
src-tauri/src/proxy/usage/mod.rs
+11
View File
@@ -0,0 +1,11 @@
//! Proxy Usage Tracking Module
//!
//! 提供 API 请求的使用量跟踪、成本计算和日志记录功能
pub mod calculator;
pub mod logger;
pub mod parser;
pub use calculator::{CostBreakdown, CostCalculator, ModelPricing};
pub use logger::{RequestLog, UsageLogger};
pub use parser::{ApiType, TokenUsage};
src-tauri/src/proxy/usage/parser.rs
+263
View File
@@ -0,0 +1,263 @@
//! Response Parser - 从 API 响应中提取 token 使用量
//!
//! 支持多种 API 格式：
//! - Claude API (非流式和流式)
//! - OpenRouter (OpenAI 格式)
//! - Codex API (非流式和流式)
//! - Gemini API (非流式和流式)
use serde::{Deserialize, Serialize};
use serde_json::Value;
/// Token 使用量统计
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TokenUsage {
pub input_tokens: u32,
pub output_tokens: u32,
pub cache_read_tokens: u32,
pub cache_creation_tokens: u32,
}
/// API 类型
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApiType {
Claude,
OpenRouter,
Codex,
Gemini,
}
impl TokenUsage {
/// 从 Claude API 非流式响应解析
pub fn from_claude_response(body: &Value) -> Option<Self> {
let usage = body.get("usage")?;
Some(Self {
input_tokens: usage.get("input_tokens")?.as_u64()? as u32,
output_tokens: usage.get("output_tokens")?.as_u64()? as u32,
cache_read_tokens: usage
.get("cache_read_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
cache_creation_tokens: usage
.get("cache_creation_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
})
}
/// 从 Claude API 流式响应解析
pub fn from_claude_stream_events(events: &[Value]) -> Option<Self> {
let mut usage = Self::default();
for event in events {
if let Some(event_type) = event.get("type").and_then(|v| v.as_str()) {
match event_type {
"message_start" => {
if let Some(msg_usage) = event.get("message").and_then(|m| m.get("usage"))
{
usage.input_tokens =
msg_usage.get("input_tokens")?.as_u64()? as u32;
usage.cache_read_tokens = msg_usage
.get("cache_read_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
usage.cache_creation_tokens = msg_usage
.get("cache_creation_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
}
}
"message_delta" => {
if let Some(delta_usage) = event.get("usage") {
usage.output_tokens =
delta_usage.get("output_tokens")?.as_u64()? as u32;
}
}
_ => {}
}
}
}
if usage.input_tokens > 0 || usage.output_tokens > 0 {
Some(usage)
} else {
None
}
}
/// 从 OpenRouter 响应解析 (OpenAI 格式)
pub fn from_openrouter_response(body: &Value) -> Option<Self> {
let usage = body.get("usage")?;
Some(Self {
input_tokens: usage.get("prompt_tokens")?.as_u64()? as u32,
output_tokens: usage.get("completion_tokens")?.as_u64()? as u32,
cache_read_tokens: 0,
cache_creation_tokens: 0,
})
}
/// 从 Codex API 非流式响应解析
pub fn from_codex_response(body: &Value) -> Option<Self> {
let usage = body.get("usage")?;
Some(Self {
input_tokens: usage.get("input_tokens")?.as_u64()? as u32,
output_tokens: usage.get("output_tokens")?.as_u64()? as u32,
cache_read_tokens: usage
.get("cache_read_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
cache_creation_tokens: usage
.get("cache_creation_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
})
}
/// 从 Codex API 流式响应解析
pub fn from_codex_stream_events(events: &[Value]) -> Option<Self> {
for event in events {
if let Some(event_type) = event.get("type").and_then(|v| v.as_str()) {
if event_type == "response.completed" {
if let Some(response) = event.get("response") {
return Self::from_codex_response(response);
}
}
}
}
None
}
/// 从 Gemini API 非流式响应解析
pub fn from_gemini_response(body: &Value) -> Option<Self> {
let usage = body.get("usageMetadata")?;
Some(Self {
input_tokens: usage.get("promptTokenCount")?.as_u64()? as u32,
output_tokens: usage.get("candidatesTokenCount")?.as_u64()? as u32,
cache_read_tokens: usage
.get("cachedContentTokenCount")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
cache_creation_tokens: 0,
})
}
/// 从 Gemini API 流式响应解析
pub fn from_gemini_stream_chunks(chunks: &[Value]) -> Option<Self> {
let mut total_input = 0u32;
let mut total_output = 0u32;
let mut total_cache_read = 0u32;
for chunk in chunks {
if let Some(usage) = chunk.get("usageMetadata") {
total_input = usage
.get("promptTokenCount")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
total_output += usage
.get("candidatesTokenCount")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
total_cache_read = usage
.get("cachedContentTokenCount")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
}
}
if total_input > 0 || total_output > 0 {
Some(Self {
input_tokens: total_input,
output_tokens: total_output,
cache_read_tokens: total_cache_read,
cache_creation_tokens: 0,
})
} else {
None
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_claude_response_parsing() {
let response = json!({
"usage": {
"input_tokens": 100,
"output_tokens": 50,
"cache_read_input_tokens": 20,
"cache_creation_input_tokens": 10
}
});
let usage = TokenUsage::from_claude_response(&response).unwrap();
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.cache_read_tokens, 20);
assert_eq!(usage.cache_creation_tokens, 10);
}
#[test]
fn test_claude_stream_parsing() {
let events = vec![
json!({
"type": "message_start",
"message": {
"usage": {
"input_tokens": 100,
"cache_read_input_tokens": 20,
"cache_creation_input_tokens": 10
}
}
}),
json!({
"type": "message_delta",
"usage": {
"output_tokens": 50
}
}),
];
let usage = TokenUsage::from_claude_stream_events(&events).unwrap();
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.cache_read_tokens, 20);
assert_eq!(usage.cache_creation_tokens, 10);
}
#[test]
fn test_openrouter_response_parsing() {
let response = json!({
"usage": {
"prompt_tokens": 100,
"completion_tokens": 50
}
});
let usage = TokenUsage::from_openrouter_response(&response).unwrap();
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.cache_read_tokens, 0);
assert_eq!(usage.cache_creation_tokens, 0);
}
#[test]
fn test_gemini_response_parsing() {
let response = json!({
"usageMetadata": {
"promptTokenCount": 100,
"candidatesTokenCount": 50,
"cachedContentTokenCount": 20
}
});
let usage = TokenUsage::from_gemini_response(&response).unwrap();
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.cache_read_tokens, 20);
assert_eq!(usage.cache_creation_tokens, 0);
}
}
src-tauri/src/services/usage_stats.rs
+527
View File
@@ -0,0 +1,527 @@
//! 使用统计服务
//!
//! 提供使用量数据的聚合查询功能
use crate::database::{lock_conn, Database};
use crate::error::AppError;
use rusqlite::params;
use serde::{Deserialize, Serialize};
/// 使用量汇总
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UsageSummary {
pub total_requests: u64,
pub total_cost: String,
pub total_input_tokens: u64,
pub total_output_tokens: u64,
pub success_rate: f32,
}
/// 每日统计
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DailyStats {
pub date: String,
pub request_count: u64,
pub total_cost: String,
pub total_tokens: u64,
}
/// Provider 统计
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderStats {
pub provider_id: String,
pub provider_name: String,
pub request_count: u64,
pub total_tokens: u64,
pub total_cost: String,
pub success_rate: f32,
pub avg_latency_ms: u64,
}
/// 模型统计
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelStats {
pub model: String,
pub request_count: u64,
pub total_tokens: u64,
pub total_cost: String,
pub avg_cost_per_request: String,
}
/// 请求日志过滤器
#[derive(Debug, Clone, Default)]
pub struct LogFilters {
pub provider_id: Option<String>,
pub model: Option<String>,
pub status_code: Option<u16>,
pub start_date: Option<i64>,
pub end_date: Option<i64>,
}
/// 请求日志详情
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequestLogDetail {
pub request_id: String,
pub provider_id: String,
pub app_type: String,
pub model: String,
pub input_tokens: u32,
pub output_tokens: u32,
pub cache_read_tokens: u32,
pub cache_creation_tokens: u32,
pub input_cost_usd: String,
pub output_cost_usd: String,
pub cache_read_cost_usd: String,
pub cache_creation_cost_usd: String,
pub total_cost_usd: String,
pub latency_ms: u64,
pub status_code: u16,
pub error_message: Option<String>,
pub created_at: i64,
}
impl Database {
/// 获取使用量汇总
pub fn get_usage_summary(
&self,
start_date: Option<i64>,
end_date: Option<i64>,
) -> Result<UsageSummary, AppError> {
let conn = lock_conn!(self.conn);
let (where_clause, params_vec) = if start_date.is_some() || end_date.is_some() {
let mut conditions = Vec::new();
let mut params = Vec::new();
if let Some(start) = start_date {
conditions.push("created_at >= ?");
params.push(start);
}
if let Some(end) = end_date {
conditions.push("created_at <= ?");
params.push(end);
}
(format!("WHERE {}", conditions.join(" AND ")), params)
} else {
(String::new(), Vec::new())
};
let sql = format!(
"SELECT
COUNT(*) as total_requests,
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) as total_cost,
COALESCE(SUM(input_tokens + output_tokens), 0) as total_tokens,
COALESCE(SUM(CASE WHEN status_code >= 200 AND status_code < 300 THEN 1 ELSE 0 END), 0) as success_count
FROM proxy_request_logs
{where_clause}"
);
let result = conn.query_row(&sql, rusqlite::params_from_iter(params_vec), |row| {
let total_requests: i64 = row.get(0)?;
let total_cost: f64 = row.get(1)?;
let total_tokens: i64 = row.get(2)?;
let success_count: i64 = row.get(3)?;
let success_rate = if total_requests > 0 {
(success_count as f32 / total_requests as f32) * 100.0
} else {
0.0
};
Ok(UsageSummary {
total_requests: total_requests as u64,
total_cost: format!("{:.6}", total_cost),
total_input_tokens: 0,
total_output_tokens: 0,
success_rate,
})
})?;
Ok(result)
}
/// 获取每日趋势
pub fn get_daily_trends(
&self,
days: u32,
) -> Result<Vec<DailyStats>, AppError> {
let conn = lock_conn!(self.conn);
let sql = "SELECT
date(created_at, 'unixepoch') as date,
COUNT(*) as request_count,
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) as total_cost,
COALESCE(SUM(input_tokens + output_tokens), 0) as total_tokens
FROM proxy_request_logs
WHERE created_at >= strftime('%s', 'now', ?)
GROUP BY date
ORDER BY date DESC";
let mut stmt = conn.prepare(sql)?;
let rows = stmt.query_map([format!("-{} days", days)], |row| {
Ok(DailyStats {
date: row.get(0)?,
request_count: row.get::<_, i64>(1)? as u64,
total_cost: format!("{:.6}", row.get::<_, f64>(2)?),
total_tokens: row.get::<_, i64>(3)? as u64,
})
})?;
let mut stats = Vec::new();
for row in rows {
stats.push(row?);
}
Ok(stats)
}
/// 获取 Provider 统计
pub fn get_provider_stats(&self) -> Result<Vec<ProviderStats>, AppError> {
let conn = lock_conn!(self.conn);
let sql = "SELECT
l.provider_id,
p.name as provider_name,
COUNT(*) as request_count,
COALESCE(SUM(l.input_tokens + l.output_tokens), 0) as total_tokens,
COALESCE(SUM(CAST(l.total_cost_usd AS REAL)), 0) as total_cost,
COALESCE(SUM(CASE WHEN l.status_code >= 200 AND l.status_code < 300 THEN 1 ELSE 0 END), 0) as success_count,
COALESCE(AVG(l.latency_ms), 0) as avg_latency
FROM proxy_request_logs l
LEFT JOIN providers p ON l.provider_id = p.id AND l.app_type = p.app_type
GROUP BY l.provider_id, l.app_type
ORDER BY total_cost DESC";
let mut stmt = conn.prepare(sql)?;
let rows = stmt.query_map([], |row| {
let request_count: i64 = row.get(2)?;
let success_count: i64 = row.get(5)?;
let success_rate = if request_count > 0 {
(success_count as f32 / request_count as f32) * 100.0
} else {
0.0
};
Ok(ProviderStats {
provider_id: row.get(0)?,
provider_name: row.get::<_, Option<String>>(1)?.unwrap_or_else(|| "Unknown".to_string()),
request_count: request_count as u64,
total_tokens: row.get::<_, i64>(3)? as u64,
total_cost: format!("{:.6}", row.get::<_, f64>(4)?),
success_rate,
avg_latency_ms: row.get::<_, f64>(6)? as u64,
})
})?;
let mut stats = Vec::new();
for row in rows {
stats.push(row?);
}
Ok(stats)
}
/// 获取模型统计
pub fn get_model_stats(&self) -> Result<Vec<ModelStats>, AppError> {
let conn = lock_conn!(self.conn);
let sql = "SELECT
model,
COUNT(*) as request_count,
COALESCE(SUM(input_tokens + output_tokens), 0) as total_tokens,
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) as total_cost
FROM proxy_request_logs
GROUP BY model
ORDER BY total_cost DESC";
let mut stmt = conn.prepare(sql)?;
let rows = stmt.query_map([], |row| {
let request_count: i64 = row.get(1)?;
let total_cost: f64 = row.get(3)?;
let avg_cost = if request_count > 0 {
total_cost / request_count as f64
} else {
0.0
};
Ok(ModelStats {
model: row.get(0)?,
request_count: request_count as u64,
total_tokens: row.get::<_, i64>(2)? as u64,
total_cost: format!("{:.6}", total_cost),
avg_cost_per_request: format!("{:.6}", avg_cost),
})
})?;
let mut stats = Vec::new();
for row in rows {
stats.push(row?);
}
Ok(stats)
}
/// 获取请求日志列表
pub fn get_request_logs(
&self,
filters: &LogFilters,
limit: u32,
offset: u32,
) -> Result<Vec<RequestLogDetail>, AppError> {
let conn = lock_conn!(self.conn);
let mut conditions = Vec::new();
let mut params: Vec<Box<dyn rusqlite::ToSql>> = Vec::new();
if let Some(ref provider_id) = filters.provider_id {
conditions.push("provider_id = ?");
params.push(Box::new(provider_id.clone()));
}
if let Some(ref model) = filters.model {
conditions.push("model = ?");
params.push(Box::new(model.clone()));
}
if let Some(status) = filters.status_code {
conditions.push("status_code = ?");
params.push(Box::new(status as i64));
}
if let Some(start) = filters.start_date {
conditions.push("created_at >= ?");
params.push(Box::new(start));
}
if let Some(end) = filters.end_date {
conditions.push("created_at <= ?");
params.push(Box::new(end));
}
let where_clause = if conditions.is_empty() {
String::new()
} else {
format!("WHERE {}", conditions.join(" AND "))
};
params.push(Box::new(limit as i64));
params.push(Box::new(offset as i64));
let sql = format!(
"SELECT request_id, provider_id, app_type, model,
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd,
latency_ms, status_code, error_message, created_at
FROM proxy_request_logs
{where_clause}
ORDER BY created_at DESC
LIMIT ? OFFSET ?"
);
let mut stmt = conn.prepare(&sql)?;
let params_refs: Vec<&dyn rusqlite::ToSql> = params.iter().map(|p| p.as_ref()).collect();
let rows = stmt.query_map(params_refs.as_slice(), |row| {
Ok(RequestLogDetail {
request_id: row.get(0)?,
provider_id: row.get(1)?,
app_type: row.get(2)?,
model: row.get(3)?,
input_tokens: row.get::<_, i64>(4)? as u32,
output_tokens: row.get::<_, i64>(5)? as u32,
cache_read_tokens: row.get::<_, i64>(6)? as u32,
cache_creation_tokens: row.get::<_, i64>(7)? as u32,
input_cost_usd: row.get(8)?,
output_cost_usd: row.get(9)?,
cache_read_cost_usd: row.get(10)?,
cache_creation_cost_usd: row.get(11)?,
total_cost_usd: row.get(12)?,
latency_ms: row.get::<_, i64>(13)? as u64,
status_code: row.get::<_, i64>(14)? as u16,
error_message: row.get(15)?,
created_at: row.get(16)?,
})
})?;
let mut logs = Vec::new();
for row in rows {
logs.push(row?);
}
Ok(logs)
}
/// 获取单个请求详情
pub fn get_request_detail(&self, request_id: &str) -> Result<Option<RequestLogDetail>, AppError> {
let conn = lock_conn!(self.conn);
let result = conn.query_row(
"SELECT request_id, provider_id, app_type, model,
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd,
latency_ms, status_code, error_message, created_at
FROM proxy_request_logs
WHERE request_id = ?",
[request_id],
|row| {
Ok(RequestLogDetail {
request_id: row.get(0)?,
provider_id: row.get(1)?,
app_type: row.get(2)?,
model: row.get(3)?,
input_tokens: row.get::<_, i64>(4)? as u32,
output_tokens: row.get::<_, i64>(5)? as u32,
cache_read_tokens: row.get::<_, i64>(6)? as u32,
cache_creation_tokens: row.get::<_, i64>(7)? as u32,
input_cost_usd: row.get(8)?,
output_cost_usd: row.get(9)?,
cache_read_cost_usd: row.get(10)?,
cache_creation_cost_usd: row.get(11)?,
total_cost_usd: row.get(12)?,
latency_ms: row.get::<_, i64>(13)? as u64,
status_code: row.get::<_, i64>(14)? as u16,
error_message: row.get(15)?,
created_at: row.get(16)?,
})
},
);
match result {
Ok(detail) => Ok(Some(detail)),
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
Err(e) => Err(AppError::Database(e.to_string())),
}
}
/// 检查 Provider 使用限额
pub fn check_provider_limits(
&self,
provider_id: &str,
app_type: &str,
) -> Result<ProviderLimitStatus, AppError> {
let conn = lock_conn!(self.conn);
// 获取 provider 的限额设置
let (limit_daily, limit_monthly) = conn.query_row(
"SELECT meta FROM providers WHERE id = ? AND app_type = ?",
params![provider_id, app_type],
|row| {
let meta_str: String = row.get(0)?;
Ok(meta_str)
},
).ok().and_then(|meta_str| {
serde_json::from_str::<serde_json::Value>(&meta_str).ok()
}).and_then(|meta| {
let daily = meta.get("limitDailyUsd")
.and_then(|v| v.as_str())
.and_then(|s| s.parse::<f64>().ok());
let monthly = meta.get("limitMonthlyUsd")
.and_then(|v| v.as_str())
.and_then(|s| s.parse::<f64>().ok());
Some((daily, monthly))
}).unwrap_or((None, None));
// 计算今日使用量
let daily_usage: f64 = conn.query_row(
"SELECT COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0)
FROM proxy_request_logs
WHERE provider_id = ? AND app_type = ?
AND date(created_at, 'unixepoch') = date('now')",
params![provider_id, app_type],
|row| row.get(0),
).unwrap_or(0.0);
// 计算本月使用量
let monthly_usage: f64 = conn.query_row(
"SELECT COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0)
FROM proxy_request_logs
WHERE provider_id = ? AND app_type = ?
AND strftime('%Y-%m', created_at, 'unixepoch') = strftime('%Y-%m', 'now')",
params![provider_id, app_type],
|row| row.get(0),
).unwrap_or(0.0);
let daily_exceeded = limit_daily.map(|limit| daily_usage >= limit).unwrap_or(false);
let monthly_exceeded = limit_monthly.map(|limit| monthly_usage >= limit).unwrap_or(false);
Ok(ProviderLimitStatus {
provider_id: provider_id.to_string(),
daily_usage: format!("{:.6}", daily_usage),
daily_limit: limit_daily.map(|l| format!("{:.2}", l)),
daily_exceeded,
monthly_usage: format!("{:.6}", monthly_usage),
monthly_limit: limit_monthly.map(|l| format!("{:.2}", l)),
monthly_exceeded,
})
}
}
/// Provider 限额状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderLimitStatus {
pub provider_id: String,
pub daily_usage: String,
pub daily_limit: Option<String>,
pub daily_exceeded: bool,
pub monthly_usage: String,
pub monthly_limit: Option<String>,
pub monthly_exceeded: bool,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_get_usage_summary() -> Result<(), AppError> {
let db = Database::memory()?;
// 插入测试数据
{
let conn = lock_conn!(db.conn);
conn.execute(
"INSERT INTO proxy_request_logs (
request_id, provider_id, app_type, model,
input_tokens, output_tokens, total_cost_usd,
latency_ms, status_code, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
params!["req1", "p1", "claude", "claude-3", 100, 50, "0.01", 100, 200, 1000],
)?;
conn.execute(
"INSERT INTO proxy_request_logs (
request_id, provider_id, app_type, model,
input_tokens, output_tokens, total_cost_usd,
latency_ms, status_code, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
params!["req2", "p1", "claude", "claude-3", 200, 100, "0.02", 150, 200, 2000],
)?;
}
let summary = db.get_usage_summary(None, None)?;
assert_eq!(summary.total_requests, 2);
assert_eq!(summary.success_rate, 100.0);
Ok(())
}
#[test]
fn test_get_model_stats() -> Result<(), AppError> {
let db = Database::memory()?;
// 插入测试数据
{
let conn = lock_conn!(db.conn);
conn.execute(
"INSERT INTO proxy_request_logs (
request_id, provider_id, app_type, model,
input_tokens, output_tokens, total_cost_usd,
latency_ms, status_code, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
params!["req1", "p1", "claude", "claude-3-sonnet", 100, 50, "0.01", 100, 200, 1000],
)?;
}
let stats = db.get_model_stats()?;
assert_eq!(stats.len(), 1);
assert_eq!(stats[0].model, "claude-3-sonnet");
assert_eq!(stats[0].request_count, 1);
Ok(())
}
}