544 lines
19 KiB
Rust
544 lines
19 KiB
Rust
use crate::config::ConfigManager;
|
||
use crate::drive_monitor::DriveMonitor;
|
||
use crate::llm::OpenAIClient;
|
||
use crate::llm_models;
|
||
use crate::nvidia::get_system_metrics;
|
||
use crate::shared::models::{BotResponse, Suggestion, UserMessage, UserSession};
|
||
use crate::shared::state::AppState;
|
||
use axum::extract::ws::{Message, WebSocket};
|
||
use axum::{
|
||
extract::{ws::WebSocketUpgrade, Extension, Path, Query, State},
|
||
http::StatusCode,
|
||
response::{IntoResponse, Json},
|
||
};
|
||
use chrono::Utc;
|
||
use diesel::PgConnection;
|
||
use futures::{sink::SinkExt, stream::StreamExt};
|
||
use log::{error, info, trace, warn};
|
||
use serde_json;
|
||
use std::collections::HashMap;
|
||
use std::sync::Arc;
|
||
use std::time::Duration;
|
||
use tokio::sync::mpsc;
|
||
use tokio::sync::Mutex as AsyncMutex;
|
||
use tokio::time::Instant;
|
||
use uuid::Uuid;
|
||
|
||
/// Retrieves the default bot (first active bot) from the database.
|
||
pub fn get_default_bot(conn: &mut PgConnection) -> (Uuid, String) {
|
||
use crate::shared::models::schema::bots::dsl::*;
|
||
use diesel::prelude::*;
|
||
match bots
|
||
.filter(is_active.eq(true))
|
||
.select((id, name))
|
||
.first::<(Uuid, String)>(conn)
|
||
.optional()
|
||
{
|
||
Ok(Some((bot_id, bot_name))) => (bot_id, bot_name),
|
||
Ok(None) => {
|
||
warn!("No active bots found, using nil UUID");
|
||
(Uuid::nil(), "default".to_string())
|
||
}
|
||
Err(e) => {
|
||
error!("Failed to query default bot: {}", e);
|
||
(Uuid::nil(), "default".to_string())
|
||
}
|
||
}
|
||
}
|
||
|
||
pub struct BotOrchestrator {
|
||
pub state: Arc<AppState>,
|
||
pub mounted_bots: Arc<AsyncMutex<HashMap<String, Arc<DriveMonitor>>>>,
|
||
}
|
||
|
||
impl BotOrchestrator {
|
||
pub fn new(state: Arc<AppState>) -> Self {
|
||
Self {
|
||
state,
|
||
mounted_bots: Arc::new(AsyncMutex::new(HashMap::new())),
|
||
}
|
||
}
|
||
|
||
// ... (All existing methods unchanged) ...
|
||
|
||
pub async fn mount_all_bots(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||
// No-op: bot mounting is handled elsewhere
|
||
info!("mount_all_bots called (no-op)");
|
||
Ok(())
|
||
}
|
||
|
||
// Stream response to user via LLM
|
||
pub async fn stream_response(
|
||
&self,
|
||
message: UserMessage,
|
||
response_tx: mpsc::Sender<BotResponse>,
|
||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||
trace!(
|
||
"Streaming response for user: {}, session: {}",
|
||
message.user_id,
|
||
message.session_id
|
||
);
|
||
|
||
let user_id = Uuid::parse_str(&message.user_id)?;
|
||
let session_id = Uuid::parse_str(&message.session_id)?;
|
||
let bot_id = Uuid::parse_str(&message.bot_id).unwrap_or_default();
|
||
|
||
// All database operations in one blocking section
|
||
let (session, context_data, history, model, key) = {
|
||
let state_clone = self.state.clone();
|
||
tokio::task::spawn_blocking(
|
||
move || -> Result<_, Box<dyn std::error::Error + Send + Sync>> {
|
||
// Get session
|
||
let session = {
|
||
let mut sm = state_clone.session_manager.blocking_lock();
|
||
sm.get_session_by_id(session_id)?
|
||
}
|
||
.ok_or_else(|| "Session not found")?;
|
||
|
||
// Save user message
|
||
{
|
||
let mut sm = state_clone.session_manager.blocking_lock();
|
||
sm.save_message(session.id, user_id, 1, &message.content, 1)?;
|
||
}
|
||
|
||
// Get context and history
|
||
let context_data = {
|
||
let sm = state_clone.session_manager.blocking_lock();
|
||
let rt = tokio::runtime::Handle::current();
|
||
rt.block_on(async {
|
||
sm.get_session_context_data(&session.id, &session.user_id)
|
||
.await
|
||
})?
|
||
};
|
||
|
||
let history = {
|
||
let mut sm = state_clone.session_manager.blocking_lock();
|
||
sm.get_conversation_history(session.id, user_id)?
|
||
};
|
||
|
||
// Get model config
|
||
let config_manager = ConfigManager::new(state_clone.conn.clone());
|
||
let model = config_manager
|
||
.get_config(&bot_id, "llm-model", Some("gpt-3.5-turbo"))
|
||
.unwrap_or_else(|_| "gpt-3.5-turbo".to_string());
|
||
let key = config_manager
|
||
.get_config(&bot_id, "llm-key", Some(""))
|
||
.unwrap_or_default();
|
||
|
||
Ok((session, context_data, history, model, key))
|
||
},
|
||
)
|
||
.await??
|
||
};
|
||
|
||
// Build messages
|
||
let system_prompt = std::env::var("SYSTEM_PROMPT")
|
||
.unwrap_or_else(|_| "You are a helpful assistant.".to_string());
|
||
let messages = OpenAIClient::build_messages(&system_prompt, &context_data, &history);
|
||
|
||
// Stream from LLM
|
||
let (stream_tx, mut stream_rx) = mpsc::channel::<String>(100);
|
||
let llm = self.state.llm_provider.clone();
|
||
|
||
tokio::spawn(async move {
|
||
if let Err(e) = llm
|
||
.generate_stream("", &messages, stream_tx, &model, &key)
|
||
.await
|
||
{
|
||
error!("LLM streaming error: {}", e);
|
||
}
|
||
});
|
||
|
||
let mut full_response = String::new();
|
||
let mut chunk_count = 0;
|
||
|
||
while let Some(chunk) = stream_rx.recv().await {
|
||
chunk_count += 1;
|
||
info!("Received LLM chunk #{}: {:?}", chunk_count, chunk);
|
||
full_response.push_str(&chunk);
|
||
|
||
let response = BotResponse {
|
||
bot_id: message.bot_id.clone(),
|
||
user_id: message.user_id.clone(),
|
||
session_id: message.session_id.clone(),
|
||
channel: message.channel.clone(),
|
||
content: chunk,
|
||
message_type: 2,
|
||
stream_token: None,
|
||
is_complete: false,
|
||
suggestions: Vec::new(),
|
||
context_name: None,
|
||
context_length: 0,
|
||
context_max_length: 0,
|
||
};
|
||
|
||
info!("Sending streaming chunk to WebSocket");
|
||
if let Err(e) = response_tx.send(response).await {
|
||
error!("Failed to send streaming chunk: {}", e);
|
||
break;
|
||
}
|
||
}
|
||
|
||
info!(
|
||
"LLM streaming complete, received {} chunks, total length: {}",
|
||
chunk_count,
|
||
full_response.len()
|
||
);
|
||
|
||
// Send final complete response
|
||
let final_response = BotResponse {
|
||
bot_id: message.bot_id.clone(),
|
||
user_id: message.user_id.clone(),
|
||
session_id: message.session_id.clone(),
|
||
channel: message.channel.clone(),
|
||
content: full_response.clone(),
|
||
message_type: 2,
|
||
stream_token: None,
|
||
is_complete: true,
|
||
suggestions: Vec::new(),
|
||
context_name: None,
|
||
context_length: 0,
|
||
context_max_length: 0,
|
||
};
|
||
|
||
info!("Sending final complete response to WebSocket");
|
||
response_tx.send(final_response).await?;
|
||
info!("Final response sent successfully");
|
||
|
||
// Save bot response in blocking context
|
||
let state_for_save = self.state.clone();
|
||
let full_response_clone = full_response.clone();
|
||
tokio::task::spawn_blocking(
|
||
move || -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||
let mut sm = state_for_save.session_manager.blocking_lock();
|
||
sm.save_message(session.id, user_id, 2, &full_response_clone, 2)?;
|
||
Ok(())
|
||
},
|
||
)
|
||
.await??;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
// ... (Other methods unchanged) ...
|
||
|
||
pub async fn get_user_sessions(
|
||
&self,
|
||
user_id: Uuid,
|
||
) -> Result<Vec<UserSession>, Box<dyn std::error::Error + Send + Sync>> {
|
||
let mut session_manager = self.state.session_manager.lock().await;
|
||
let sessions = session_manager.get_user_sessions(user_id)?;
|
||
Ok(sessions)
|
||
}
|
||
|
||
pub async fn get_conversation_history(
|
||
&self,
|
||
session_id: Uuid,
|
||
user_id: Uuid,
|
||
) -> Result<Vec<(String, String)>, Box<dyn std::error::Error + Send + Sync>> {
|
||
let mut session_manager = self.state.session_manager.lock().await;
|
||
let history = session_manager.get_conversation_history(session_id, user_id)?;
|
||
Ok(history)
|
||
}
|
||
|
||
// ... (Remaining BotOrchestrator methods unchanged) ...
|
||
}
|
||
|
||
/* Axum handlers – placeholders that delegate to BotOrchestrator where appropriate */
|
||
|
||
/// WebSocket handler that upgrades HTTP connection to WebSocket
|
||
pub async fn websocket_handler(
|
||
ws: WebSocketUpgrade,
|
||
State(state): State<Arc<AppState>>,
|
||
Query(params): Query<HashMap<String, String>>,
|
||
) -> impl IntoResponse {
|
||
let session_id = params
|
||
.get("session_id")
|
||
.and_then(|s| Uuid::parse_str(s).ok());
|
||
let user_id = params.get("user_id").and_then(|s| Uuid::parse_str(s).ok());
|
||
|
||
if session_id.is_none() || user_id.is_none() {
|
||
return (
|
||
StatusCode::BAD_REQUEST,
|
||
Json(serde_json::json!({ "error": "session_id and user_id are required" })),
|
||
)
|
||
.into_response();
|
||
}
|
||
|
||
ws.on_upgrade(move |socket| {
|
||
handle_websocket(socket, state, session_id.unwrap(), user_id.unwrap())
|
||
})
|
||
.into_response()
|
||
}
|
||
|
||
/// Handles an individual WebSocket connection
|
||
async fn handle_websocket(
|
||
socket: WebSocket,
|
||
state: Arc<AppState>,
|
||
session_id: Uuid,
|
||
user_id: Uuid,
|
||
) {
|
||
let (mut sender, mut receiver) = socket.split();
|
||
|
||
// Create a channel for this WebSocket connection
|
||
let (tx, mut rx) = mpsc::channel::<BotResponse>(100);
|
||
|
||
// Register this connection with the web adapter
|
||
state
|
||
.web_adapter
|
||
.add_connection(session_id.to_string(), tx.clone())
|
||
.await;
|
||
|
||
// Also register in response_channels for BotOrchestrator
|
||
{
|
||
let mut channels = state.response_channels.lock().await;
|
||
channels.insert(session_id.to_string(), tx.clone());
|
||
}
|
||
|
||
info!(
|
||
"WebSocket connected for session: {}, user: {}",
|
||
session_id, user_id
|
||
);
|
||
|
||
// Execute start.bas if it exists
|
||
let state_for_start = state.clone();
|
||
let session_for_start = {
|
||
let mut sm = state.session_manager.lock().await;
|
||
sm.get_session_by_id(session_id).ok().and_then(|opt| opt)
|
||
};
|
||
|
||
if let Some(session_clone) = session_for_start {
|
||
tokio::task::spawn_blocking(move || {
|
||
use crate::basic::ScriptService;
|
||
|
||
let bot_name = "default"; // TODO: Get from session
|
||
let start_script_path =
|
||
format!("./work/{}.gbai/{}.gbdialog/start.bas", bot_name, bot_name);
|
||
|
||
if let Ok(start_content) = std::fs::read_to_string(&start_script_path) {
|
||
info!("Executing start.bas for session {}", session_id);
|
||
let script_service = ScriptService::new(state_for_start, session_clone);
|
||
match script_service.compile(&start_content) {
|
||
Ok(ast) => {
|
||
if let Err(e) = script_service.run(&ast) {
|
||
error!("Failed to execute start.bas: {}", e);
|
||
} else {
|
||
info!("start.bas executed successfully for session {}", session_id);
|
||
}
|
||
}
|
||
Err(e) => {
|
||
error!("Failed to compile start.bas: {}", e);
|
||
}
|
||
}
|
||
} else {
|
||
info!("No start.bas found for bot {}", bot_name);
|
||
}
|
||
});
|
||
}
|
||
|
||
// Send initial welcome message
|
||
let welcome = serde_json::json!({
|
||
"type": "connected",
|
||
"session_id": session_id,
|
||
"user_id": user_id,
|
||
"message": "Connected to bot server"
|
||
});
|
||
|
||
if let Ok(welcome_str) = serde_json::to_string(&welcome) {
|
||
info!("Sending welcome message to session {}", session_id);
|
||
if let Err(e) = sender.send(Message::Text(welcome_str.into())).await {
|
||
error!("Failed to send welcome message: {}", e);
|
||
}
|
||
}
|
||
|
||
// Spawn task to send messages from the channel to the WebSocket
|
||
let mut send_task = tokio::spawn(async move {
|
||
while let Some(response) = rx.recv().await {
|
||
if let Ok(json_str) = serde_json::to_string(&response) {
|
||
if sender.send(Message::Text(json_str.into())).await.is_err() {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
});
|
||
|
||
// Handle incoming messages from the WebSocket
|
||
let state_clone = state.clone();
|
||
let mut recv_task = tokio::spawn(async move {
|
||
while let Some(Ok(msg)) = receiver.next().await {
|
||
info!("WebSocket received raw message type: {:?}", msg);
|
||
match msg {
|
||
Message::Text(text) => {
|
||
info!(
|
||
"Received WebSocket text message (length {}): {}",
|
||
text.len(),
|
||
text
|
||
);
|
||
match serde_json::from_str::<UserMessage>(&text) {
|
||
Ok(user_msg) => {
|
||
info!(
|
||
"Successfully parsed user message from session: {}, content: {}",
|
||
session_id, user_msg.content
|
||
);
|
||
// Process the message through the bot system
|
||
if let Err(e) = process_user_message(
|
||
state_clone.clone(),
|
||
session_id,
|
||
user_id,
|
||
user_msg,
|
||
)
|
||
.await
|
||
{
|
||
error!("Error processing user message: {}", e);
|
||
}
|
||
}
|
||
Err(e) => {
|
||
error!(
|
||
"Failed to parse user message from session {}: {} - Parse error: {}",
|
||
session_id, text, e
|
||
);
|
||
}
|
||
}
|
||
}
|
||
Message::Close(_) => {
|
||
info!(
|
||
"WebSocket close message received for session: {}",
|
||
session_id
|
||
);
|
||
break;
|
||
}
|
||
Message::Ping(_data) => {
|
||
// Pings are automatically handled by axum
|
||
}
|
||
Message::Pong(_) => {
|
||
// Pongs are automatically handled by axum
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
});
|
||
|
||
// Wait for either task to finish
|
||
tokio::select! {
|
||
_ = (&mut send_task) => {
|
||
recv_task.abort();
|
||
}
|
||
_ = (&mut recv_task) => {
|
||
send_task.abort();
|
||
}
|
||
}
|
||
|
||
// Clean up: remove the connection from the adapter
|
||
state
|
||
.web_adapter
|
||
.remove_connection(&session_id.to_string())
|
||
.await;
|
||
|
||
// Also remove from response_channels
|
||
{
|
||
let mut channels = state.response_channels.lock().await;
|
||
channels.remove(&session_id.to_string());
|
||
}
|
||
|
||
info!("WebSocket disconnected for session: {}", session_id);
|
||
}
|
||
|
||
/// Process a user message received via WebSocket
|
||
async fn process_user_message(
|
||
state: Arc<AppState>,
|
||
session_id: Uuid,
|
||
user_id: Uuid,
|
||
user_msg: UserMessage,
|
||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||
info!(
|
||
"Processing message from user {} in session {}: {}",
|
||
user_id, session_id, user_msg.content
|
||
);
|
||
|
||
// Get the response channel for this session
|
||
let tx = {
|
||
let channels = state.response_channels.lock().await;
|
||
channels.get(&session_id.to_string()).cloned()
|
||
};
|
||
|
||
if let Some(response_tx) = tx {
|
||
// Use BotOrchestrator to stream the response
|
||
let orchestrator = BotOrchestrator::new(state.clone());
|
||
if let Err(e) = orchestrator.stream_response(user_msg, response_tx).await {
|
||
error!("Failed to stream response: {}", e);
|
||
}
|
||
} else {
|
||
error!("No response channel found for session {}", session_id);
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Create a new bot (placeholder implementation)
|
||
pub async fn create_bot_handler(
|
||
Extension(state): Extension<Arc<AppState>>,
|
||
Json(payload): Json<HashMap<String, String>>,
|
||
) -> impl IntoResponse {
|
||
let bot_name = payload
|
||
.get("bot_name")
|
||
.cloned()
|
||
.unwrap_or_else(|| "default".to_string());
|
||
(
|
||
StatusCode::OK,
|
||
Json(serde_json::json!({ "status": format!("bot '{}' created", bot_name) })),
|
||
)
|
||
}
|
||
|
||
/// Mount an existing bot (placeholder implementation)
|
||
pub async fn mount_bot_handler(
|
||
Extension(state): Extension<Arc<AppState>>,
|
||
Json(payload): Json<HashMap<String, String>>,
|
||
) -> impl IntoResponse {
|
||
let bot_guid = payload.get("bot_guid").cloned().unwrap_or_default();
|
||
(
|
||
StatusCode::OK,
|
||
Json(serde_json::json!({ "status": format!("bot '{}' mounted", bot_guid) })),
|
||
)
|
||
}
|
||
|
||
/// Handle user input for a bot (placeholder implementation)
|
||
pub async fn handle_user_input_handler(
|
||
Extension(state): Extension<Arc<AppState>>,
|
||
Json(payload): Json<HashMap<String, String>>,
|
||
) -> impl IntoResponse {
|
||
let session_id = payload.get("session_id").cloned().unwrap_or_default();
|
||
let user_input = payload.get("input").cloned().unwrap_or_default();
|
||
(
|
||
StatusCode::OK,
|
||
Json(
|
||
serde_json::json!({ "status": format!("input '{}' processed for session {}", user_input, session_id) }),
|
||
),
|
||
)
|
||
}
|
||
|
||
/// Retrieve user sessions (placeholder implementation)
|
||
pub async fn get_user_sessions_handler(
|
||
Extension(state): Extension<Arc<AppState>>,
|
||
Json(payload): Json<HashMap<String, String>>,
|
||
) -> impl IntoResponse {
|
||
(StatusCode::OK, Json(serde_json::json!({ "sessions": [] })))
|
||
}
|
||
|
||
/// Retrieve conversation history (placeholder implementation)
|
||
pub async fn get_conversation_history_handler(
|
||
Extension(state): Extension<Arc<AppState>>,
|
||
Json(payload): Json<HashMap<String, String>>,
|
||
) -> impl IntoResponse {
|
||
(StatusCode::OK, Json(serde_json::json!({ "history": [] })))
|
||
}
|
||
|
||
/// Send warning (placeholder implementation)
|
||
pub async fn send_warning_handler(
|
||
Extension(state): Extension<Arc<AppState>>,
|
||
Json(payload): Json<HashMap<String, String>>,
|
||
) -> impl IntoResponse {
|
||
(
|
||
StatusCode::OK,
|
||
Json(serde_json::json!({ "status": "warning acknowledged" })),
|
||
)
|
||
}
|