118 lines
4.5 KiB
Rust
118 lines
4.5 KiB
Rust
use db::cache::AppCache;
|
|
use db::database::AppDatabase;
|
|
use sea_orm::{EntityTrait, ActiveModelTrait, QueryFilter, ColumnTrait, QueryOrder};
|
|
use uuid::Uuid;
|
|
|
|
/// Record an AI session with cost calculation.
|
|
pub async fn record_ai_session(
|
|
cache: &AppCache,
|
|
db: &AppDatabase,
|
|
project_id: Uuid,
|
|
user_id: Uuid,
|
|
session_id: Uuid,
|
|
room_id: Uuid,
|
|
model_id: Uuid,
|
|
version_id: Uuid,
|
|
input_tokens: i64,
|
|
output_tokens: i64,
|
|
latency_ms: i64,
|
|
) {
|
|
metrics::histogram!("ai_call_latency_ms", "model" => model_id.to_string()).record(latency_ms as f64);
|
|
|
|
let session = models::ai::ai_session::ActiveModel {
|
|
id: sea_orm::Set(session_id),
|
|
room: sea_orm::Set(room_id),
|
|
model: sea_orm::Set(model_id),
|
|
version: sea_orm::Set(version_id),
|
|
token_input: sea_orm::Set(input_tokens),
|
|
token_output: sea_orm::Set(output_tokens),
|
|
latency_ms: sea_orm::Set(Some(latency_ms)),
|
|
cost: sea_orm::Set(None),
|
|
currency: sea_orm::Set(None),
|
|
error_message: sea_orm::Set(None),
|
|
error_code: sea_orm::Set(None),
|
|
created_at: sea_orm::Set(chrono::Utc::now()),
|
|
};
|
|
|
|
if let Err(e) = session.insert(db).await {
|
|
tracing::error!(error = %e, session_id = %session_id, "failed to insert ai session record");
|
|
return;
|
|
}
|
|
|
|
let (cost, currency, error_msg) = match crate::billing::record_ai_usage(
|
|
db, project_id, user_id, version_id, input_tokens, output_tokens,
|
|
).await {
|
|
Ok(crate::billing::BillingResult::Success(record)) => {
|
|
(Some(record.cost), Some(record.currency), None)
|
|
}
|
|
Ok(crate::billing::BillingResult::InsufficientBalance { message }) => {
|
|
create_billing_error_system_message(cache, db, room_id, &message).await;
|
|
(None, None, Some(message))
|
|
}
|
|
Err(e) => (None, None, Some(e.to_string())),
|
|
};
|
|
|
|
use sea_orm::sea_query::Expr;
|
|
let _ = models::ai::ai_session::Entity::update_many()
|
|
.col_expr(models::ai::ai_session::Column::Cost, Expr::value(cost))
|
|
.col_expr(models::ai::ai_session::Column::Currency, Expr::value(currency))
|
|
.col_expr(models::ai::ai_session::Column::ErrorMessage, Expr::value(error_msg))
|
|
.filter(models::ai::ai_session::Column::Id.eq(session_id))
|
|
.exec(db).await;
|
|
}
|
|
|
|
/// Create a system message in the room for billing errors.
|
|
async fn create_billing_error_system_message(
|
|
cache: &AppCache,
|
|
db: &AppDatabase,
|
|
room_id: Uuid,
|
|
message: &str,
|
|
) {
|
|
use models::rooms::{room_message, MessageContentType, MessageSenderType};
|
|
use sea_orm::Set;
|
|
|
|
let seq_key = format!("seq:room:{}", room_id);
|
|
let seq = match cache.conn().await {
|
|
Ok(mut conn) => {
|
|
match redis::cmd("INCR").arg(&seq_key).query_async::<i64>(&mut conn).await {
|
|
Ok(s) => s,
|
|
Err(e) => {
|
|
tracing::warn!(error = %e, "cache INCR failed for system message seq, falling back to DB");
|
|
let last_seq = match room_message::Entity::find()
|
|
.filter(room_message::Column::Room.eq(room_id))
|
|
.order_by_desc(room_message::Column::Seq)
|
|
.one(db).await
|
|
{
|
|
Ok(Some(m)) => m.seq,
|
|
Ok(None) => 0,
|
|
Err(e) => {
|
|
tracing::warn!(error = %e, "Failed to get last seq for system message");
|
|
return;
|
|
}
|
|
};
|
|
last_seq + 1
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(error = %e, "Failed to get Redis connection for system message seq");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let now = chrono::Utc::now();
|
|
let result = room_message::ActiveModel {
|
|
id: Set(Uuid::now_v7()), seq: Set(seq), room: Set(room_id),
|
|
sender_type: Set(MessageSenderType::System), sender_id: Set(None),
|
|
model_id: Set(None), thread: Set(None), in_reply_to: Set(None),
|
|
content: Set(message.to_string()), content_type: Set(MessageContentType::Text),
|
|
thinking_content: Set(None), edited_at: Set(None),
|
|
send_at: Set(now), revoked: Set(None), revoked_by: Set(None),
|
|
}.insert(db).await;
|
|
|
|
match result {
|
|
Ok(_) => tracing::info!(room_id = %room_id, message = %message, "system_message_created_for_billing_error"),
|
|
Err(e) => tracing::warn!(error = %e, room_id = %room_id, "Failed to create system message for billing error"),
|
|
}
|
|
}
|