5b81e7d774
Simplify ai_streaming by delegating to ai_mode_streaming. Extract sequence coordination into dedicated module. Add worker pool management for concurrent AI task handling. Refine ai_react_streaming for better delta chunk handling.
87 lines
2.9 KiB
Rust
87 lines
2.9 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");
|
|
}
|
|
|
|
// Billing handled internally by chat_service.process via record_ai_session
|
|
}
|
|
}
|
|
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(),
|
|
"[AI 处理发生错误，请稍后再试]".to_string(),
|
|
model_id,
|
|
Some(model_display_name),
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
});
|
|
}
|