88 lines
2.8 KiB
Rust
88 lines
2.8 KiB
Rust
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::error::AiResult;
|
|
|
|
pub struct UsageService {
|
|
pub db: sea_orm::DatabaseConnection,
|
|
}
|
|
|
|
impl UsageService {
|
|
pub fn new(db: sea_orm::DatabaseConnection) -> Self {
|
|
Self { db }
|
|
}
|
|
|
|
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<ai_usage::Model> {
|
|
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<UsageOverview> {
|
|
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::<UsageOverview>()
|
|
.one(&self.db)
|
|
.await?
|
|
.unwrap_or(UsageOverview { total_count: 0 });
|
|
Ok(result)
|
|
}
|
|
|
|
/// 按分析类型统计
|
|
pub async fn get_by_type(&self, tenant_id: Uuid) -> AiResult<Vec<TypeCount>> {
|
|
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::<TypeCount>()
|
|
.all(&self.db)
|
|
.await?;
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, FromQueryResult)]
|
|
pub struct UsageOverview {
|
|
pub total_count: i64,
|
|
}
|
|
|
|
#[derive(Debug, FromQueryResult)]
|
|
pub struct TypeCount {
|
|
pub analysis_type: String,
|
|
pub count: i64,
|
|
}
|