use sea_orm::{ ActiveModelTrait, ColumnTrait, EntityTrait, FromQueryResult, QueryFilter, QuerySelect, Set, }; use uuid::Uuid; use crate::entity::ai_analysis; use crate::entity::ai_usage; use crate::entity::ai_usage_daily; use crate::error::AiResult; pub struct UsageService { pub db: sea_orm::DatabaseConnection, } impl UsageService { pub fn new(db: sea_orm::DatabaseConnection) -> Self { Self { db } } #[allow(clippy::too_many_arguments)] pub async fn log_usage( &self, tenant_id: Uuid, provider: &str, model: &str, analysis_type: &str, input_tokens: u32, output_tokens: u32, duration_ms: u64, cost_cents: i32, is_cache_hit: bool, ) -> AiResult { let id = Uuid::now_v7(); let active = ai_usage::ActiveModel { id: Set(id), tenant_id: Set(tenant_id), provider: Set(provider.into()), model: Set(model.into()), analysis_type: Set(analysis_type.into()), input_tokens: Set(input_tokens as i32), output_tokens: Set(output_tokens as i32), duration_ms: Set(duration_ms as i32), cost_cents: Set(cost_cents), is_cache_hit: Set(is_cache_hit), created_at: Set(chrono::Utc::now()), }; Ok(active.insert(&self.db).await?) } /// 用量概览 pub async fn get_overview(&self, tenant_id: Uuid) -> AiResult { let result = ai_analysis::Entity::find() .filter(ai_analysis::Column::TenantId.eq(tenant_id)) .filter(ai_analysis::Column::Status.eq("completed")) .filter(ai_analysis::Column::DeletedAt.is_null()) .select_only() .column_as(ai_analysis::Column::Id.count(), "total_count") .into_model::() .one(&self.db) .await? .unwrap_or(UsageOverview { total_count: 0 }); Ok(result) } /// 按分析类型统计 pub async fn get_by_type(&self, tenant_id: Uuid) -> AiResult> { let result = ai_analysis::Entity::find() .filter(ai_analysis::Column::TenantId.eq(tenant_id)) .filter(ai_analysis::Column::Status.eq("completed")) .filter(ai_analysis::Column::DeletedAt.is_null()) .select_only() .column(ai_analysis::Column::AnalysisType) .column_as(ai_analysis::Column::Id.count(), "count") .group_by(ai_analysis::Column::AnalysisType) .into_model::() .all(&self.db) .await?; Ok(result) } /// 按日期范围查询日聚合用量 pub async fn get_daily_usage( &self, tenant_id: Uuid, start_date: chrono::NaiveDate, end_date: chrono::NaiveDate, ) -> AiResult> { let rows = ai_usage_daily::Entity::find() .filter(ai_usage_daily::Column::TenantId.eq(tenant_id)) .filter(ai_usage_daily::Column::Date.gte(start_date)) .filter(ai_usage_daily::Column::Date.lte(end_date)) .all(&self.db) .await?; Ok(rows .into_iter() .map(|r| DailyUsageRow { date: r.date, feature: r.feature, provider: r.provider, model: r.model, total_calls: r.total_calls, total_input_tokens: r.total_input_tokens, total_output_tokens: r.total_output_tokens, total_cost_cents: r.total_cost_cents, }) .collect()) } /// 聚合指定日期的用量到日聚合表(由定时任务调用) pub async fn aggregate_daily(&self, tenant_id: Uuid, date: chrono::NaiveDate) -> AiResult<()> { let date_start = date.and_hms_opt(0, 0, 0).unwrap_or_default(); let date_end = date_start + chrono::Duration::days(1); // 从 ai_usage 按分析类型聚合 #[derive(Debug, FromQueryResult)] struct AggRow { analysis_type: String, total_calls: i64, total_input_tokens: i64, total_output_tokens: i64, total_cost_cents: i64, } let rows: Vec = ai_usage::Entity::find() .filter(ai_usage::Column::TenantId.eq(tenant_id)) .filter(ai_usage::Column::CreatedAt.gte(date_start)) .filter(ai_usage::Column::CreatedAt.lt(date_end)) .select_only() .column(ai_usage::Column::AnalysisType) .column_as(ai_usage::Column::Id.count(), "total_calls") .column_as(ai_usage::Column::InputTokens.sum(), "total_input_tokens") .column_as(ai_usage::Column::OutputTokens.sum(), "total_output_tokens") .column_as(ai_usage::Column::CostCents.sum(), "total_cost_cents") .group_by(ai_usage::Column::AnalysisType) .into_model::() .all(&self.db) .await?; for row in &rows { let id = Uuid::now_v7(); let active = ai_usage_daily::ActiveModel { id: Set(id), tenant_id: Set(tenant_id), date: Set(date), feature: Set(row.analysis_type.clone()), provider: Set("aggregated".into()), model: Set("mixed".into()), total_calls: Set(row.total_calls as i32), total_input_tokens: Set(row.total_input_tokens), total_output_tokens: Set(row.total_output_tokens), total_cost_cents: Set(row.total_cost_cents), created_at: Set(chrono::Utc::now()), }; active.insert(&self.db).await?; } Ok(()) } } #[derive(Debug, FromQueryResult)] pub struct UsageOverview { pub total_count: i64, } #[derive(Debug, FromQueryResult)] pub struct TypeCount { pub analysis_type: String, pub count: i64, } #[derive(Debug, Clone, serde::Serialize)] pub struct DailyUsageRow { pub date: chrono::NaiveDate, pub feature: String, pub provider: String, pub model: String, pub total_calls: i32, pub total_input_tokens: i64, pub total_output_tokens: i64, pub total_cost_cents: i64, }