gitdataai/libs/room/src/service/ai_nonstreaming.rs
ZhenYi f5e3da35b0 feat(room): store ordered streaming chunks + billing integration
- Save thinking_content as {"__chunks__": [{type, content}]} for replay
- Tool call sanitization — don't expose raw results to frontend
- Billing record_ai_usage integration
- Room service module refactoring into service/ directory
2026-04-26 13:10:42 +08:00

95 lines
3.1 KiB
Rust

use std::sync::Arc;
use chrono::Utc;
use db::cache::AppCache;
use db::database::AppDatabase;
use models::rooms::room_ai;
use queue::MessageProducer;
use sea_orm::{sea_query::Expr, ColumnTrait, EntityTrait, ExprTrait, QueryFilter};
use uuid::Uuid;
use super::ai_common::create_and_publish_ai_message;
use crate::connection::RoomConnectionManager;
use agent::chat::{AiRequest, ChatService};
pub async fn process_message_ai_nonstreaming(
chat_service: Arc<ChatService>,
request: AiRequest,
room_id: Uuid,
project_id: Uuid,
model_id: Uuid,
lock_guard: crate::room_ai_queue::RoomAiLockGuard,
db: AppDatabase,
cache: AppCache,
queue: MessageProducer,
room_manager: Arc<RoomConnectionManager>,
) {
let chat_service = chat_service.clone();
tokio::spawn(async move {
let _lock_guard = lock_guard;
let model_display_name = request.model.name.clone();
match chat_service.process(request).await {
Ok(result) => {
if let Err(e) = create_and_publish_ai_message(
&db,
&cache,
&queue,
&room_manager,
room_id,
project_id,
Uuid::now_v7(),
result.content,
model_id,
Some(model_display_name),
)
.await
{
tracing::error!(error = %e, "Failed to create AI message");
} else {
let now = Utc::now();
if let Err(e) = room_ai::Entity::update_many()
.col_expr(
room_ai::Column::CallCount,
Expr::col(room_ai::Column::CallCount).add(1),
)
.col_expr(room_ai::Column::LastCallAt, Expr::value(Some(now)))
.filter(room_ai::Column::Room.eq(room_id))
.filter(room_ai::Column::Model.eq(model_id))
.exec(&db)
.await
{
tracing::warn!(error = %e, "Failed to update room_ai call stats");
}
// Record billing (non-fatal)
let _ = super::billing::record_ai_usage(
&db,
project_id,
model_id,
result.input_tokens,
result.output_tokens,
)
.await;
}
}
Err(e) => {
tracing::error!(error = %e, "AI processing failed");
let _ = create_and_publish_ai_message(
&db,
&cache,
&queue,
&room_manager,
room_id,
project_id,
Uuid::now_v7(),
format!("[AI error: {}]", e),
model_id,
Some(model_display_name),
)
.await;
}
}
});
}