botserver/src/automation/compact_prompt.rs

use crate::config::ConfigManager;
use crate::llm_models;
use crate::shared::models::Automation;
use crate::shared::state::AppState;
use diesel::prelude::*;
use log::{error,  trace};
use std::collections::HashSet;
use std::sync::Arc;
use tokio::time::{interval, Duration};
use uuid::Uuid;
pub fn start_compact_prompt_scheduler(state: Arc<AppState>) {
 tokio::spawn(async move {
 tokio::time::sleep(Duration::from_secs(30)).await;
 let mut interval = interval(Duration::from_secs(60));
 loop {
 interval.tick().await;
 if let Err(e) = execute_compact_prompt(Arc::clone(&state)).await {
 error!("Prompt compaction failed: {}", e);
 }
 }
 });
}
async fn execute_compact_prompt(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 use crate::shared::models::system_automations::dsl::{is_active, system_automations};
 let automations: Vec<Automation> = {
 let mut conn = state.conn.get().map_err(|e| format!("Failed to acquire lock: {}", e))?;
 system_automations.filter(is_active.eq(true)).load::<Automation>(&mut *conn)?
 };
 for automation in automations {
 if let Err(e) = compact_prompt_for_bot(&state, &automation).await {
 error!("Failed to compact prompt for bot {}: {}", automation.bot_id, e);
 }
 }
 Ok(())
}
async fn compact_prompt_for_bot(state: &Arc<AppState>, automation: &Automation) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 use once_cell::sync::Lazy;
 use scopeguard::guard;
 static IN_PROGRESS: Lazy<tokio::sync::Mutex<HashSet<Uuid>>> = Lazy::new(|| tokio::sync::Mutex::new(HashSet::new()));
 {
 let mut in_progress = IN_PROGRESS.lock().await;
 if in_progress.contains(&automation.bot_id) {
 return Ok(());
 }
 in_progress.insert(automation.bot_id);
 }
 let bot_id = automation.bot_id;
 let _cleanup = guard((), |_| {
 tokio::spawn(async move {
 let mut in_progress = IN_PROGRESS.lock().await;
 in_progress.remove(&bot_id);
 });
 });
 let config_manager = ConfigManager::new(state.conn.clone());
 let compact_threshold = config_manager
 .get_config(&automation.bot_id, "prompt-compact", None)?
 .parse::<i32>()
 .unwrap_or(0);
 if compact_threshold == 0 {
 return Ok(());
 } else if compact_threshold < 0 {
 trace!("Negative compact threshold detected for bot {}, skipping", automation.bot_id);
 }
 let sessions = {
 let mut session_manager = state.session_manager.lock().await;
 session_manager.get_user_sessions(Uuid::nil())?
 };
 for session in sessions {
 if session.bot_id != automation.bot_id {
 continue;
 }
 let history = {
 let mut session_manager = state.session_manager.lock().await;
 session_manager.get_conversation_history(session.id, session.user_id)?
 };
 trace!("Compacting prompt for session {}: {} messages", session.id, history.len());
 let mut compacted = String::new();
 for (role, content) in &history {
 compacted.push_str(&format!("{}: {}\n", role, content));
 }
 let llm_provider = state.llm_provider.clone();
 let compacted_clone = compacted.clone();
 let summarized = match llm_provider.summarize(&compacted_clone).await {
 Ok(summary) => {
 trace!("Successfully summarized conversation for session {}, summary length: {}", session.id, summary.len());
 let handler = llm_models::get_handler(&config_manager.get_config(&automation.bot_id, "llm-model", None).unwrap_or_default());
 let filtered = handler.process_content(&summary);
 format!("SUMMARY: {}", filtered)
 },
 Err(e) => {
 error!("Failed to summarize conversation for session {}: {}", session.id, e);
 format!("SUMMARY: {}", compacted)
 }
 };
 trace!("Prompt compacted {}: {} messages", session.id, history.len());
 {
 let mut session_manager = state.session_manager.lock().await;
 session_manager.save_message(session.id, session.user_id, 9, &summarized, 1)?;
 }
 }
 Ok(())
}