gitdataai/libs/room/src/service/ai_react_nonstreaming.rs
ZhenYi fdca1fbf86
Some checks are pending
CI / Rust Lint & Check (push) Waiting to run
CI / Rust Tests (push) Waiting to run
CI / Frontend Lint & Type Check (push) Waiting to run
CI / Frontend Build (push) Blocked by required conditions
feat(ai): add comprehensive AI streaming and non-streaming processing services
2026-05-01 00:54:24 +08:00

89 lines
3.0 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_react_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>,
) {
tokio::spawn(async move {
let _lock_guard = lock_guard;
let model_display_name = request.model.name.clone();
let final_answer = chat_service
.process_react(&request, |_step| async move {})
.await;
match final_answer {
Ok((response, _input_tokens, _output_tokens)) => {
if let Err(e) = create_and_publish_ai_message(
&db,
&cache,
&queue,
&room_manager,
room_id,
project_id,
Uuid::now_v7(),
response,
model_id,
Some(model_display_name),
)
.await
{
tracing::error!(error = %e, "Failed to create ReAct AI message");
} else {
// Billing handled internally by chat_service.process_react via record_ai_session
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");
}
}
}
Err(e) => {
tracing::error!(error = ?e, "ReAct agent failed");
let _ = create_and_publish_ai_message(
&db,
&cache,
&queue,
&room_manager,
room_id,
project_id,
Uuid::now_v7(),
format!("AI 处理失败: {}", e),
model_id,
Some(model_display_name),
)
.await;
}
}
});
}