diff --git a/crates/erp-ai/src/service/mod.rs b/crates/erp-ai/src/service/mod.rs index 3319347..b4b4b80 100644 --- a/crates/erp-ai/src/service/mod.rs +++ b/crates/erp-ai/src/service/mod.rs @@ -6,6 +6,7 @@ pub mod local_rules; pub mod output_parser; pub mod post_process; pub mod prompt; +pub mod quota; pub mod reanalysis; pub mod suggestion; pub mod usage; diff --git a/crates/erp-ai/src/service/quota.rs b/crates/erp-ai/src/service/quota.rs new file mode 100644 index 0000000..d2e0505 --- /dev/null +++ b/crates/erp-ai/src/service/quota.rs @@ -0,0 +1,188 @@ +use sea_orm::{ColumnTrait, EntityTrait, FromQueryResult, QueryFilter, Statement}; +use uuid::Uuid; + +use crate::entity::ai_tenant_config; +use crate::error::{AiError, AiResult}; + +pub struct QuotaService { + db: sea_orm::DatabaseConnection, + enabled: bool, +} + +impl QuotaService { + pub fn new(db: sea_orm::DatabaseConnection, enabled: bool) -> Self { + Self { db, enabled } + } + + pub async fn get_tenant_config( + &self, + tenant_id: Uuid, + ) -> AiResult> { + let config = ai_tenant_config::Entity::find() + .filter(ai_tenant_config::Column::TenantId.eq(tenant_id)) + .filter(ai_tenant_config::Column::DeletedAt.is_null()) + .filter(ai_tenant_config::Column::IsEnabled.eq(true)) + .one(&self.db) + .await?; + Ok(config) + } + + pub async fn check_quota( + &self, + tenant_id: Uuid, + patient_id: Option, + ) -> AiResult<()> { + if !self.enabled { + return Ok(()); + } + + let config = match self.get_tenant_config(tenant_id).await? { + Some(c) => c, + None => return Ok(()), + }; + + let month_usage = self.get_monthly_token_usage(tenant_id).await?; + if month_usage >= config.monthly_token_budget { + return Err(AiError::QuotaExhausted { + tenant_id, + reason: format!( + "月度 Token 预算已耗尽 ({}/{})", + month_usage, config.monthly_token_budget + ), + }); + } + + if let Some(pid) = patient_id { + let daily_count = self.get_daily_patient_count(tenant_id, pid).await?; + if daily_count >= config.daily_patient_limit as i64 { + return Err(AiError::QuotaExhausted { + tenant_id, + reason: format!( + "患者 {} 今日分析次数已达上限 ({}/{})", + pid, daily_count, config.daily_patient_limit + ), + }); + } + } + + Ok(()) + } + + async fn get_monthly_token_usage(&self, tenant_id: Uuid) -> AiResult { + #[derive(Debug, FromQueryResult)] + struct TokenSum { + total_tokens: i64, + } + + let sql = r#" + SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total_tokens + FROM ai_usage + WHERE tenant_id = $1 + AND created_at >= date_trunc('month', CURRENT_DATE) + "#; + + let result: Option = TokenSum::find_by_statement( + Statement::from_sql_and_values( + sea_orm::DatabaseBackend::Postgres, + sql, + [tenant_id.into()], + ), + ) + .one(&self.db) + .await?; + + Ok(result.map(|r| r.total_tokens).unwrap_or(0)) + } + + async fn get_daily_patient_count( + &self, + tenant_id: Uuid, + patient_id: Uuid, + ) -> AiResult { + #[derive(Debug, FromQueryResult)] + struct CountResult { + count: i64, + } + + let sql = r#" + SELECT COUNT(*) AS count + FROM ai_analysis_results + WHERE tenant_id = $1 + AND patient_id = $2 + AND deleted_at IS NULL + AND created_at >= CURRENT_DATE + "#; + + let result: Option = CountResult::find_by_statement( + Statement::from_sql_and_values( + sea_orm::DatabaseBackend::Postgres, + sql, + [tenant_id.into(), patient_id.into()], + ), + ) + .one(&self.db) + .await?; + + Ok(result.map(|r| r.count).unwrap_or(0)) + } + + pub async fn get_usage_summary( + &self, + tenant_id: Uuid, + ) -> AiResult { + let config = self.get_tenant_config(tenant_id).await?; + let budget = config + .as_ref() + .map(|c| c.monthly_token_budget) + .unwrap_or(1_000_000); + + let used = self.get_monthly_token_usage(tenant_id).await?; + + Ok(QuotaSummary { + tenant_id, + monthly_budget: budget, + monthly_used: used, + daily_patient_limit: config + .as_ref() + .map(|c| c.daily_patient_limit) + .unwrap_or(50), + }) + } +} + +#[derive(Debug, Clone, serde::Serialize)] +pub struct QuotaSummary { + pub tenant_id: Uuid, + pub monthly_budget: i64, + pub monthly_used: i64, + pub daily_patient_limit: i32, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn quota_summary_serialization() { + let summary = QuotaSummary { + tenant_id: Uuid::now_v7(), + monthly_budget: 1_000_000, + monthly_used: 250_000, + daily_patient_limit: 50, + }; + let json = serde_json::to_value(&summary).unwrap(); + assert_eq!(json["monthly_budget"], 1_000_000); + assert_eq!(json["monthly_used"], 250_000); + } + + #[test] + fn quota_error_contains_reason() { + let tid = Uuid::now_v7(); + let err = AiError::QuotaExhausted { + tenant_id: tid, + reason: "月度预算耗尽".into(), + }; + let msg = err.to_string(); + assert!(msg.contains("月度预算耗尽")); + } +}