- Added LLM server readiness check in AutomationService before starting tasks - Renamed `user` parameter to `user_session` in execute_talk for clarity - Updated BotResponse fields to use user_session data instead of hardcoded values - Improved Redis key generation in execute_talk to use user_session fields - Removed commented Redis code in set_current_context_keyword The changes ensure proper initialization of automation tasks by checking LLM server availability first, and improve code clarity by using more descriptive variable names for user session data.
489 lines
16 KiB
Rust
489 lines
16 KiB
Rust
use crate::bot::BotOrchestrator;
|
|
use crate::shared::models::UserSession;
|
|
use crate::shared::state::AppState;
|
|
use actix_web::{web, HttpResponse, Result};
|
|
use chrono::Utc;
|
|
use diesel::prelude::*;
|
|
use diesel::PgConnection;
|
|
use log::{debug, error, info, warn};
|
|
use redis::Client;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::error::Error;
|
|
use std::sync::Arc;
|
|
use uuid::Uuid;
|
|
|
|
#[derive(Clone, Serialize, Deserialize)]
|
|
pub struct SessionData {
|
|
pub id: Uuid,
|
|
pub user_id: Option<Uuid>,
|
|
pub data: String,
|
|
}
|
|
|
|
pub struct SessionManager {
|
|
conn: PgConnection,
|
|
sessions: HashMap<Uuid, SessionData>,
|
|
waiting_for_input: HashSet<Uuid>,
|
|
redis: Option<Arc<Client>>,
|
|
}
|
|
|
|
impl SessionManager {
|
|
pub fn new(conn: PgConnection, redis_client: Option<Arc<Client>>) -> Self {
|
|
SessionManager {
|
|
conn,
|
|
sessions: HashMap::new(),
|
|
waiting_for_input: HashSet::new(),
|
|
redis: redis_client,
|
|
}
|
|
}
|
|
|
|
pub fn provide_input(
|
|
&mut self,
|
|
session_id: Uuid,
|
|
input: String,
|
|
) -> Result<Option<String>, Box<dyn Error + Send + Sync>> {
|
|
info!(
|
|
"SessionManager.provide_input called for session {}",
|
|
session_id
|
|
);
|
|
if let Some(sess) = self.sessions.get_mut(&session_id) {
|
|
sess.data = input;
|
|
self.waiting_for_input.remove(&session_id);
|
|
Ok(Some("user_input".to_string()))
|
|
} else {
|
|
let sess = SessionData {
|
|
id: session_id,
|
|
user_id: None,
|
|
data: input,
|
|
};
|
|
self.sessions.insert(session_id, sess);
|
|
self.waiting_for_input.remove(&session_id);
|
|
Ok(Some("user_input".to_string()))
|
|
}
|
|
}
|
|
|
|
pub fn is_waiting_for_input(&self, session_id: &Uuid) -> bool {
|
|
self.waiting_for_input.contains(session_id)
|
|
}
|
|
|
|
pub fn mark_waiting(&mut self, session_id: Uuid) {
|
|
self.waiting_for_input.insert(session_id);
|
|
}
|
|
|
|
pub fn get_session_by_id(
|
|
&mut self,
|
|
session_id: Uuid,
|
|
) -> Result<Option<UserSession>, Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::user_sessions::dsl::*;
|
|
let result = user_sessions
|
|
.filter(id.eq(session_id))
|
|
.first::<UserSession>(&mut self.conn)
|
|
.optional()?;
|
|
Ok(result)
|
|
}
|
|
|
|
pub fn get_user_session(
|
|
&mut self,
|
|
uid: Uuid,
|
|
bid: Uuid,
|
|
) -> Result<Option<UserSession>, Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::user_sessions::dsl::*;
|
|
let result = user_sessions
|
|
.filter(user_id.eq(uid))
|
|
.filter(bot_id.eq(bid))
|
|
.order(created_at.desc())
|
|
.first::<UserSession>(&mut self.conn)
|
|
.optional()?;
|
|
Ok(result)
|
|
}
|
|
|
|
pub fn get_or_create_user_session(
|
|
&mut self,
|
|
uid: Uuid,
|
|
bid: Uuid,
|
|
session_title: &str,
|
|
) -> Result<Option<UserSession>, Box<dyn Error + Send + Sync>> {
|
|
if let Some(existing) = self.get_user_session(uid, bid)? {
|
|
return Ok(Some(existing));
|
|
}
|
|
self.create_session(uid, bid, session_title).map(Some)
|
|
}
|
|
|
|
pub fn get_or_create_anonymous_user(
|
|
&mut self,
|
|
uid: Option<Uuid>,
|
|
) -> Result<Uuid, Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::users::dsl as users_dsl;
|
|
|
|
let user_id = uid.unwrap_or_else(Uuid::new_v4);
|
|
|
|
let user_exists: Option<Uuid> = users_dsl::users
|
|
.filter(users_dsl::id.eq(user_id))
|
|
.select(users_dsl::id)
|
|
.first(&mut self.conn)
|
|
.optional()?;
|
|
|
|
if user_exists.is_none() {
|
|
let now = Utc::now();
|
|
info!("Creating anonymous user with ID {}", user_id);
|
|
diesel::insert_into(users_dsl::users)
|
|
.values((
|
|
users_dsl::id.eq(user_id),
|
|
users_dsl::username.eq(format!("guest_{}", &user_id.to_string()[..8])),
|
|
users_dsl::email.eq(format!(
|
|
"guest_{}@anonymous.local",
|
|
&user_id.to_string()[..8]
|
|
)),
|
|
users_dsl::password_hash.eq(""),
|
|
users_dsl::is_active.eq(true),
|
|
users_dsl::created_at.eq(now),
|
|
users_dsl::updated_at.eq(now),
|
|
))
|
|
.execute(&mut self.conn)?;
|
|
}
|
|
|
|
Ok(user_id)
|
|
}
|
|
|
|
pub fn create_session(
|
|
&mut self,
|
|
uid: Uuid,
|
|
bid: Uuid,
|
|
session_title: &str,
|
|
) -> Result<UserSession, Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::user_sessions::dsl::*;
|
|
|
|
// Ensure user exists (create anonymous if needed)
|
|
let verified_uid = self.get_or_create_anonymous_user(Some(uid))?;
|
|
let now = Utc::now();
|
|
|
|
let inserted: UserSession = diesel::insert_into(user_sessions)
|
|
.values((
|
|
id.eq(Uuid::new_v4()),
|
|
user_id.eq(verified_uid),
|
|
bot_id.eq(bid),
|
|
title.eq(session_title),
|
|
context_data.eq(serde_json::json!({})),
|
|
answer_mode.eq(0),
|
|
current_tool.eq(None::<String>),
|
|
created_at.eq(now),
|
|
updated_at.eq(now),
|
|
))
|
|
.returning(UserSession::as_returning())
|
|
.get_result(&mut self.conn)
|
|
.map_err(|e| {
|
|
error!("Failed to create session in database: {}", e);
|
|
e
|
|
})?;
|
|
|
|
Ok(inserted)
|
|
}
|
|
|
|
pub fn save_message(
|
|
&mut self,
|
|
sess_id: Uuid,
|
|
uid: Uuid,
|
|
ro: i32,
|
|
content: &str,
|
|
msg_type: i32,
|
|
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::message_history::dsl::*;
|
|
|
|
let next_index = message_history
|
|
.filter(session_id.eq(sess_id))
|
|
.count()
|
|
.get_result::<i64>(&mut self.conn)
|
|
.unwrap_or(0);
|
|
|
|
diesel::insert_into(message_history)
|
|
.values((
|
|
id.eq(Uuid::new_v4()),
|
|
session_id.eq(sess_id),
|
|
user_id.eq(uid),
|
|
role.eq(ro),
|
|
content_encrypted.eq(content),
|
|
message_type.eq(msg_type),
|
|
message_index.eq(next_index),
|
|
created_at.eq(chrono::Utc::now()),
|
|
))
|
|
.execute(&mut self.conn)?;
|
|
|
|
debug!(
|
|
"Message saved for session {} with index {}",
|
|
sess_id, next_index
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn update_session_context(
|
|
&mut self,
|
|
session_id: &Uuid,
|
|
user_id: &Uuid,
|
|
context_data: String,
|
|
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
|
use redis::Commands;
|
|
|
|
let redis_key = format!("context:{}:{}", user_id, session_id);
|
|
if let Some(redis_client) = &self.redis {
|
|
let mut conn = redis_client.get_connection()?;
|
|
conn.set(&redis_key, &context_data)?;
|
|
info!("Updated context in Redis for key {}", redis_key);
|
|
} else {
|
|
warn!("No Redis client configured, context not persisted");
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn get_session_context_data(
|
|
&self,
|
|
session_id: &Uuid,
|
|
user_id: &Uuid,
|
|
) -> Result<String, Box<dyn Error + Send + Sync>> {
|
|
// Bring the Redis command trait into scope so we can call `get`.
|
|
use redis::Commands;
|
|
|
|
let redis_key = format!("context:{}:{}", user_id, session_id);
|
|
if let Some(redis_client) = &self.redis {
|
|
// Attempt to obtain a Redis connection; log and ignore errors
|
|
let conn_option = redis_client
|
|
.get_connection()
|
|
.map_err(|e| {
|
|
warn!("Failed to get Cache connection: {}", e);
|
|
e
|
|
})
|
|
.ok();
|
|
|
|
if let Some(mut connection) = conn_option {
|
|
match connection.get::<_, Option<String>>(&redis_key) {
|
|
Ok(Some(context)) => {
|
|
debug!(
|
|
"Retrieved context from Cache for key {}: {} chars",
|
|
redis_key,
|
|
context.len()
|
|
);
|
|
return Ok(context);
|
|
}
|
|
Ok(None) => {
|
|
debug!("No context found in Cache for key {}", redis_key);
|
|
}
|
|
Err(e) => {
|
|
warn!("Failed to retrieve context from Cache: {}", e);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// If no context found, return empty string
|
|
Ok(String::new())
|
|
}
|
|
|
|
pub fn get_conversation_history(
|
|
&mut self,
|
|
sess_id: Uuid,
|
|
_uid: Uuid,
|
|
) -> Result<Vec<(String, String)>, Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::message_history::dsl::*;
|
|
|
|
let messages = message_history
|
|
.filter(session_id.eq(sess_id))
|
|
.order(message_index.asc())
|
|
.select((role, content_encrypted))
|
|
.load::<(i32, String)>(&mut self.conn)?;
|
|
|
|
let mut history: Vec<(String, String)> = Vec::new();
|
|
for (other_role, content) in messages {
|
|
let role_str = match other_role {
|
|
0 => "user".to_string(),
|
|
1 => "assistant".to_string(),
|
|
2 => "system".to_string(),
|
|
_ => "unknown".to_string(),
|
|
};
|
|
history.push((role_str, content));
|
|
}
|
|
Ok(history)
|
|
}
|
|
|
|
pub fn get_user_sessions(
|
|
&mut self,
|
|
uid: Uuid,
|
|
) -> Result<Vec<UserSession>, Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::user_sessions;
|
|
let sessions = user_sessions::table
|
|
.filter(user_sessions::user_id.eq(uid))
|
|
.order(user_sessions::created_at.desc())
|
|
.load::<UserSession>(&mut self.conn)?;
|
|
Ok(sessions)
|
|
}
|
|
|
|
pub fn update_answer_mode(
|
|
&mut self,
|
|
uid: &str,
|
|
bid: &str,
|
|
mode: i32,
|
|
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::user_sessions::dsl::*;
|
|
|
|
let user_uuid = Uuid::parse_str(uid).map_err(|e| {
|
|
warn!("Invalid user ID format: {}", uid);
|
|
e
|
|
})?;
|
|
let bot_uuid = Uuid::parse_str(bid).map_err(|e| {
|
|
warn!("Invalid bot ID format: {}", bid);
|
|
e
|
|
})?;
|
|
|
|
let updated_count = diesel::update(
|
|
user_sessions
|
|
.filter(user_id.eq(user_uuid))
|
|
.filter(bot_id.eq(bot_uuid)),
|
|
)
|
|
.set((answer_mode.eq(mode), updated_at.eq(chrono::Utc::now())))
|
|
.execute(&mut self.conn)?;
|
|
|
|
if updated_count == 0 {
|
|
warn!("No session found for user {} and bot {}", uid, bid);
|
|
} else {
|
|
info!(
|
|
"Answer mode updated to {} for user {} and bot {}",
|
|
mode, uid, bid
|
|
);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub fn update_user_id(
|
|
&mut self,
|
|
session_id: Uuid,
|
|
new_user_id: Uuid,
|
|
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
|
use crate::shared::models::user_sessions::dsl::*;
|
|
|
|
let updated_count = diesel::update(user_sessions.filter(id.eq(session_id)))
|
|
.set((user_id.eq(new_user_id), updated_at.eq(chrono::Utc::now())))
|
|
.execute(&mut self.conn)?;
|
|
|
|
if updated_count == 0 {
|
|
warn!("No session found with ID: {}", session_id);
|
|
} else {
|
|
debug!("Updated user ID for session {}", session_id);
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[actix_web::post("/api/sessions")]
|
|
async fn create_session(data: web::Data<AppState>) -> Result<HttpResponse> {
|
|
let user_id = Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
|
|
let bot_id = Uuid::nil();
|
|
|
|
let session = {
|
|
let mut session_manager = data.session_manager.lock().await;
|
|
match session_manager.get_or_create_user_session(user_id, bot_id, "New Conversation") {
|
|
Ok(Some(s)) => s,
|
|
Ok(None) => {
|
|
error!("Failed to create session");
|
|
return Ok(HttpResponse::InternalServerError()
|
|
.json(serde_json::json!({"error": "Failed to create session"})));
|
|
}
|
|
Err(e) => {
|
|
error!("Failed to create session: {}", e);
|
|
return Ok(HttpResponse::InternalServerError()
|
|
.json(serde_json::json!({"error": e.to_string()})));
|
|
}
|
|
}
|
|
};
|
|
|
|
Ok(HttpResponse::Ok().json(serde_json::json!({
|
|
"session_id": session.id,
|
|
"title": "New Conversation",
|
|
"created_at": Utc::now()
|
|
})))
|
|
}
|
|
|
|
#[actix_web::get("/api/sessions")]
|
|
async fn get_sessions(data: web::Data<AppState>) -> Result<HttpResponse> {
|
|
let user_id = Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
|
|
let orchestrator = BotOrchestrator::new(Arc::new(data.get_ref().clone()));
|
|
match orchestrator.get_user_sessions(user_id).await {
|
|
Ok(sessions) => Ok(HttpResponse::Ok().json(sessions)),
|
|
Err(e) => {
|
|
error!("Failed to get sessions: {}", e);
|
|
Ok(HttpResponse::InternalServerError()
|
|
.json(serde_json::json!({"error": e.to_string()})))
|
|
}
|
|
}
|
|
}
|
|
|
|
#[actix_web::post("/api/sessions/{session_id}/start")]
|
|
async fn start_session(
|
|
data: web::Data<AppState>,
|
|
path: web::Path<String>,
|
|
) -> Result<HttpResponse> {
|
|
let session_id = path.into_inner();
|
|
|
|
match Uuid::parse_str(&session_id) {
|
|
Ok(session_uuid) => {
|
|
let mut session_manager = data.session_manager.lock().await;
|
|
match session_manager.get_session_by_id(session_uuid) {
|
|
Ok(Some(_session)) => {
|
|
session_manager.mark_waiting(session_uuid);
|
|
Ok(HttpResponse::Ok().json(serde_json::json!({
|
|
"status": "started",
|
|
"session_id": session_id
|
|
})))
|
|
}
|
|
Ok(None) => {
|
|
Ok(HttpResponse::NotFound().json(serde_json::json!({
|
|
"error": "Session not found"
|
|
})))
|
|
}
|
|
Err(e) => {
|
|
error!("Failed to start session {}: {}", session_id, e);
|
|
Ok(HttpResponse::InternalServerError()
|
|
.json(serde_json::json!({"error": e.to_string()})))
|
|
}
|
|
}
|
|
}
|
|
Err(_) => {
|
|
warn!("Invalid session ID format: {}", session_id);
|
|
Ok(HttpResponse::BadRequest().json(serde_json::json!({"error": "Invalid session ID"})))
|
|
}
|
|
}
|
|
}
|
|
|
|
#[actix_web::get("/api/sessions/{session_id}")]
|
|
async fn get_session_history(
|
|
data: web::Data<AppState>,
|
|
path: web::Path<String>,
|
|
) -> Result<HttpResponse> {
|
|
let session_id = path.into_inner();
|
|
let user_id = Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
|
|
|
|
match Uuid::parse_str(&session_id) {
|
|
Ok(session_uuid) => {
|
|
let orchestrator = BotOrchestrator::new(Arc::new(data.get_ref().clone()));
|
|
match orchestrator
|
|
.get_conversation_history(session_uuid, user_id)
|
|
.await
|
|
{
|
|
Ok(history) => {
|
|
info!(
|
|
"Retrieved {} history entries for session {}",
|
|
history.len(),
|
|
session_id
|
|
);
|
|
Ok(HttpResponse::Ok().json(history))
|
|
}
|
|
Err(e) => {
|
|
error!("Failed to get session history for {}: {}", session_id, e);
|
|
Ok(HttpResponse::InternalServerError()
|
|
.json(serde_json::json!({"error": e.to_string()})))
|
|
}
|
|
}
|
|
}
|
|
Err(_) => {
|
|
warn!("Invalid session ID format: {}", session_id);
|
|
Ok(HttpResponse::BadRequest().json(serde_json::json!({"error": "Invalid session ID"})))
|
|
}
|
|
}
|
|
}
|