diff --git a/src-tauri/src/services/usage_stats.rs b/src-tauri/src/services/usage_stats.rs index 81cc30257..3ed76595c 100644 --- a/src-tauri/src/services/usage_stats.rs +++ b/src-tauri/src/services/usage_stats.rs @@ -829,89 +829,46 @@ impl Database { } } -/// 标准化模型名称：去除供应商前缀并将点号替换为短横线 -/// 例如：anthropic/claude-haiku-4.5 → claude-haiku-4-5 -fn normalize_model_id(model_id: &str) -> String { - // 1. 去除供应商前缀（如 anthropic/、openai/） - let stripped = if let Some(pos) = model_id.find('/') { - &model_id[pos + 1..] - } else { - model_id - }; - // 2. 将点号替换为短横线（如 claude-haiku-4.5 → claude-haiku-4-5） - stripped.replace('.', "-") -} - pub(crate) fn find_model_pricing_row( conn: &Connection, model_id: &str, ) -> Result, AppError> { - // 0. 标准化模型名称（去除前缀 + 点号转短横线） - // 例如：anthropic/claude-haiku-4.5 → claude-haiku-4-5 - let normalized = normalize_model_id(model_id); + // 1) 去除供应商前缀（/ 之前）与冒号后缀（: 之后），例如 moonshotai/kimi-k2-0905:exa → kimi-k2-0905 + let without_prefix = model_id + .rsplit_once('/') + .map(|(_, rest)| rest) + .unwrap_or(model_id); + let cleaned = without_prefix + .split(':') + .next() + .map(str::trim) + .unwrap_or(without_prefix); - // 1. 精确匹配（先尝试原始名称，再尝试标准化后的名称） - for id in [model_id, normalized.as_str()] { - let exact = 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", - [id], - |row| { - Ok(( - row.get::<_, String>(0)?, - row.get::<_, String>(1)?, - row.get::<_, String>(2)?, - row.get::<_, String>(3)?, - )) - }, - ) - .optional() - .map_err(|e| AppError::Database(format!("查询模型定价失败: {e}")))?; + // 2) 精确匹配清洗后的名称 + let exact = 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", + [cleaned], + |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, String>(1)?, + row.get::<_, String>(2)?, + row.get::<_, String>(3)?, + )) + }, + ) + .optional() + .map_err(|e| AppError::Database(format!("查询模型定价失败: {e}")))?; - if exact.is_some() { - if id != model_id { - log::info!("模型 {model_id} 标准化后精确匹配到: {id}"); - } - return Ok(exact); - } + if exact.is_none() { + log::warn!("模型 {model_id}（清洗后: {cleaned}）未找到定价信息，成本将记录为 0"); } - // 2. 逐步删除后缀匹配（claude-haiku-4-5-20250929 → claude-haiku-4-5 → claude-haiku-4 → claude-haiku） - // 使用标准化后的名称进行后缀匹配 - let mut current = normalized; - while let Some(pos) = current.rfind('-') { - current = current[..pos].to_string(); - - 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", - [¤t], - |row| { - Ok(( - row.get::<_, String>(0)?, - row.get::<_, String>(1)?, - row.get::<_, String>(2)?, - row.get::<_, String>(3)?, - )) - }, - ) - .optional() - .map_err(|e| AppError::Database(format!("查询模型定价失败: {e}")))?; - - if result.is_some() { - log::info!("模型 {model_id} 通过删除后缀匹配到: {current}"); - return Ok(result); - } - } - - log::warn!("模型 {model_id} 未找到定价信息，成本将记录为 0"); - Ok(None) + Ok(exact) } #[cfg(test)] @@ -991,54 +948,36 @@ mod tests { let db = Database::memory()?; let conn = lock_conn!(db.conn); - // 测试精确匹配 + // 准备额外定价数据，覆盖前缀/后缀清洗场景 + conn.execute( + "INSERT OR REPLACE INTO model_pricing ( + model_id, input_cost_per_million, output_cost_per_million, + cache_read_cost_per_million, cache_creation_cost_per_million, currency + ) VALUES (?, ?, ?, ?, ?, 'usd')", + params!["claude-haiku-4.5", "1.0", "2.0", "0.0", "0.0"], + )?; + conn.execute( + "INSERT OR REPLACE INTO model_pricing ( + model_id, input_cost_per_million, output_cost_per_million, + cache_read_cost_per_million, cache_creation_cost_per_million, currency + ) VALUES (?, ?, ?, ?, ?, 'usd')", + params!["kimi-k2-0905", "1.0", "1.0", "0.0", "0.0"], + )?; + + // 测试精确匹配（必须存在同名条目） let result = find_model_pricing_row(&conn, "claude-sonnet-4-5")?; assert!(result.is_some(), "应该能精确匹配 claude-sonnet-4-5"); - // 测试带供应商前缀的模型名称（anthropic/claude-haiku-4.5 → claude-haiku-4-5） + // 清洗：去除前缀和冒号后缀 let result = find_model_pricing_row(&conn, "anthropic/claude-haiku-4.5")?; assert!( result.is_some(), - "应该能匹配带前缀的模型 anthropic/claude-haiku-4.5" + "带前缀的模型 anthropic/claude-haiku-4.5 应能匹配到 claude-haiku-4.5" ); - - // 测试带供应商前缀 + 点号的模型名称 - let result = find_model_pricing_row(&conn, "anthropic/claude-sonnet-4.5")?; + let result = find_model_pricing_row(&conn, "moonshotai/kimi-k2-0905:exa")?; assert!( result.is_some(), - "应该能匹配带前缀的模型 anthropic/claude-sonnet-4.5" - ); - - // 测试逐步删除后缀匹配 - 日期后缀 - let result = find_model_pricing_row(&conn, "claude-sonnet-4-5-20241022")?; - assert!( - result.is_some(), - "应该能通过删除后缀匹配 claude-sonnet-4-5-20241022" - ); - - // 测试逐步删除后缀匹配 - 多个后缀 - let result = find_model_pricing_row(&conn, "claude-haiku-4-5-20240229-preview")?; - assert!( - result.is_some(), - "应该能通过删除后缀匹配 claude-haiku-4-5-20240229-preview" - ); - - // 测试 GPT 模型 - let result = find_model_pricing_row(&conn, "gpt-5-2024-11-20")?; - assert!(result.is_some(), "应该能通过删除后缀匹配 gpt-5-2024-11-20"); - - // 测试 Gemini 模型 - let result = find_model_pricing_row(&conn, "gemini-2.5-flash-exp")?; - assert!( - result.is_some(), - "应该能通过删除后缀匹配 gemini-2.5-flash-exp" - ); - - // 测试 claude-sonnet-4-5 命名格式 - let result = find_model_pricing_row(&conn, "claude-sonnet-4-5-20250929")?; - assert!( - result.is_some(), - "应该能通过删除后缀匹配 claude-sonnet-4-5-20250929" + "带前缀+冒号后缀的模型应清洗后匹配到 kimi-k2-0905" ); // 测试不存在的模型