feat(ai): 实现 QuotaService 租户配额检查
月度 Token 预算 + 每日患者分析次数限制,raw SQL 聚合查询 可全局开关 (quota_check_enabled),无配置时默认放行
This commit is contained in:
@@ -6,6 +6,7 @@ pub mod local_rules;
|
|||||||
pub mod output_parser;
|
pub mod output_parser;
|
||||||
pub mod post_process;
|
pub mod post_process;
|
||||||
pub mod prompt;
|
pub mod prompt;
|
||||||
|
pub mod quota;
|
||||||
pub mod reanalysis;
|
pub mod reanalysis;
|
||||||
pub mod suggestion;
|
pub mod suggestion;
|
||||||
pub mod usage;
|
pub mod usage;
|
||||||
|
|||||||
188
crates/erp-ai/src/service/quota.rs
Normal file
188
crates/erp-ai/src/service/quota.rs
Normal file
@@ -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<Option<ai_tenant_config::Model>> {
|
||||||
|
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<Uuid>,
|
||||||
|
) -> 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<i64> {
|
||||||
|
#[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> = 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<i64> {
|
||||||
|
#[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> = 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<QuotaSummary> {
|
||||||
|
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("月度预算耗尽"));
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user