refactor(bot): limit lock scope for history and tidy progress output

- Reordered imports for clarity (chrono and tokio::time::Instant).
- Fixed comment indentation around compact automation note.
- Refactored session history retrieval to acquire the mutex only briefly, then process compacted message skipping and history limiting outside the lock.
- Added explanatory comments for the new lock handling logic.
- Cleaned up token progress calculation and display formatting, improving readability of GPU/CPU/TOKENS bars.
- Minor formatting adjustments throughout the file.
This commit is contained in:
Rodrigo Rodriguez (Pragmatismo) 2025-11-07 16:43:49 -03:00
parent 120d06a0db
commit 9e1370c04f

View file

@ -6,16 +6,16 @@ use crate::shared::models::{BotResponse, Suggestion, UserMessage, UserSession};
use crate::shared::state::AppState; use crate::shared::state::AppState;
use actix_web::{web, HttpRequest, HttpResponse, Result}; use actix_web::{web, HttpRequest, HttpResponse, Result};
use actix_ws::Message as WsMessage; use actix_ws::Message as WsMessage;
use chrono::{Utc}; use chrono::Utc;
use diesel::PgConnection; use diesel::PgConnection;
use log::{error, info, trace, warn}; use log::{error, info, trace, warn};
use serde_json; use serde_json;
use tokio::time::Instant;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use tokio::sync::mpsc; use tokio::sync::mpsc;
use tokio::sync::Mutex as AsyncMutex; use tokio::sync::Mutex as AsyncMutex;
use tokio::time::Instant;
use uuid::Uuid; use uuid::Uuid;
pub fn get_default_bot(conn: &mut PgConnection) -> (Uuid, String) { pub fn get_default_bot(conn: &mut PgConnection) -> (Uuid, String) {
@ -407,21 +407,25 @@ impl BotOrchestrator {
.unwrap_or(-1) .unwrap_or(-1)
}; };
// Acquire lock briefly for history retrieval with configurable limit
let history = {
let mut sm = self.state.session_manager.lock().await; let mut sm = self.state.session_manager.lock().await;
let mut history = sm.get_conversation_history(session.id, user_id)?; let mut history = sm.get_conversation_history(session.id, user_id)?;
// Skip all messages before the most recent compacted message (type 9) // Skip all messages before the most recent compacted message (type 9)
if let Some(last_compacted_index) = history.iter().rposition(|(role, content)| { if let Some(last_compacted_index) = history
role == "COMPACTED" || content.starts_with("SUMMARY:") .iter()
}) { .rposition(|(role, content)| role == "COMPACTED" || content.starts_with("SUMMARY:"))
{
history = history.split_off(last_compacted_index); history = history.split_off(last_compacted_index);
} }
// Apply history limit if configured
if history_limit > 0 && history.len() > history_limit as usize { if history_limit > 0 && history.len() > history_limit as usize {
let start = history.len() - history_limit as usize; let start = history.len() - history_limit as usize;
history.drain(0..start); history.drain(0..start);
} }
history
};
let mut prompt = String::new(); let mut prompt = String::new();
if !system_prompt.is_empty() { if !system_prompt.is_empty() {
@ -560,9 +564,11 @@ impl BotOrchestrator {
// Update progress if interval elapsed // Update progress if interval elapsed
if last_progress_update.elapsed() >= progress_interval { if last_progress_update.elapsed() >= progress_interval {
let current_tokens = initial_tokens + crate::shared::utils::estimate_token_count(&full_response); let current_tokens =
initial_tokens + crate::shared::utils::estimate_token_count(&full_response);
if let Ok(metrics) = get_system_metrics(current_tokens, max_context_size) { if let Ok(metrics) = get_system_metrics(current_tokens, max_context_size) {
let gpu_bar = "".repeat((metrics.gpu_usage.unwrap_or(0.0) / 5.0).round() as usize); let gpu_bar =
"".repeat((metrics.gpu_usage.unwrap_or(0.0) / 5.0).round() as usize);
let cpu_bar = "".repeat((metrics.cpu_usage / 5.0).round() as usize); let cpu_bar = "".repeat((metrics.cpu_usage / 5.0).round() as usize);
let token_ratio = current_tokens as f64 / max_context_size.max(1) as f64; let token_ratio = current_tokens as f64 / max_context_size.max(1) as f64;
let token_bar = "".repeat((token_ratio * 20.0).round() as usize); let token_bar = "".repeat((token_ratio * 20.0).round() as usize);
@ -612,24 +618,31 @@ trace!(
let total_tokens = crate::shared::utils::estimate_token_count(&prompt) let total_tokens = crate::shared::utils::estimate_token_count(&prompt)
+ crate::shared::utils::estimate_token_count(&context_data) + crate::shared::utils::estimate_token_count(&context_data)
+ crate::shared::utils::estimate_token_count(&full_response); + crate::shared::utils::estimate_token_count(&full_response);
info!("Total tokens (context + prompt + response): {}", total_tokens); info!(
"Total tokens (context + prompt + response): {}",
total_tokens
);
// Trigger compact prompt if enabled // Trigger compact prompt if enabled
let config_manager = ConfigManager::new(Arc::clone(&self.state.conn)); let config_manager = ConfigManager::new(Arc::clone(&self.state.conn));
let compact_enabled = config_manager let compact_enabled = config_manager
.get_config(&Uuid::parse_str(&message.bot_id).unwrap_or_default(), "prompt-compact", None) .get_config(
&Uuid::parse_str(&message.bot_id).unwrap_or_default(),
"prompt-compact",
None,
)
.unwrap_or_default() .unwrap_or_default()
.parse::<i32>() .parse::<i32>()
.unwrap_or(0); .unwrap_or(0);
if compact_enabled > 0 { if compact_enabled > 0 {
let state = self.state.clone(); let state = self.state.clone();
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || loop {
loop { if let Err(e) = tokio::runtime::Handle::current()
if let Err(e) = tokio::runtime::Handle::current().block_on(crate::automation::execute_compact_prompt(state.clone())) { .block_on(crate::automation::execute_compact_prompt(state.clone()))
{
error!("Failed to execute compact prompt: {}", e); error!("Failed to execute compact prompt: {}", e);
} }
std::thread::sleep(Duration::from_secs(60)); std::thread::sleep(Duration::from_secs(60));
}
}); });
} }
@ -1311,8 +1324,3 @@ async fn send_warning_handler(
Ok(HttpResponse::Ok().json(serde_json::json!({"status": "warning_sent"}))) Ok(HttpResponse::Ok().json(serde_json::json!({"status": "warning_sent"})))
} }
);
}
Ok(HttpResponse::Ok().json(serde_json::json!({"status": "warning_sent"})))
}