use models::rooms::room_message::{ Column as RmCol, Entity as RoomMessage, Model as RoomMessageModel, }; use sea_orm::ColumnTrait; use sea_orm::{ConnectionTrait, EntityTrait, QueryFilter, QueryOrder, QuerySelect}; use crate::compact::types::{CompactConfig, CompactLevel, RoomCompactContext, RoomCompactRecord}; use crate::tokent::resolve_usage; use crate::{AgentError, CompactSummary, MessageSummary}; impl super::CompactService { pub async fn latest_room_compact_record( &self, room_id: uuid::Uuid, ) -> Result, AgentError> { let stmt = sea_orm::Statement::from_sql_and_values( sea_orm::DbBackend::Postgres, "SELECT id, room, from_seq, to_seq, summary, message_count, source_message_ids, created_at \ FROM room_compact_summary WHERE room = $1 ORDER BY to_seq DESC, created_at DESC LIMIT 1", vec![room_id.into()], ); let Some(row) = self .db .query_one_raw(stmt) .await .map_err(|e| AgentError::Internal(e.to_string()))? else { return Ok(None); }; let source_json: serde_json::Value = row .try_get("", "source_message_ids") .map_err(|e| AgentError::Internal(e.to_string()))?; let source_message_ids = source_json .as_array() .map(|ids| { ids.iter() .filter_map(|v| v.as_str()) .filter_map(|s| uuid::Uuid::parse_str(s).ok()) .collect::>() }) .unwrap_or_default(); Ok(Some(RoomCompactRecord { id: row .try_get("", "id") .map_err(|e| AgentError::Internal(e.to_string()))?, room_id: row .try_get("", "room") .map_err(|e| AgentError::Internal(e.to_string()))?, from_seq: row .try_get("", "from_seq") .map_err(|e| AgentError::Internal(e.to_string()))?, to_seq: row .try_get("", "to_seq") .map_err(|e| AgentError::Internal(e.to_string()))?, summary: row .try_get("", "summary") .map_err(|e| AgentError::Internal(e.to_string()))?, message_count: row .try_get("", "message_count") .map_err(|e| AgentError::Internal(e.to_string()))?, source_message_ids, created_at: row .try_get("", "created_at") .map_err(|e| AgentError::Internal(e.to_string()))?, })) } async fn insert_room_compact_record( &self, room_id: uuid::Uuid, from_seq: i64, to_seq: i64, summary: &str, source_message_ids: &[uuid::Uuid], ) -> Result { let id = uuid::Uuid::new_v4(); let now = chrono::Utc::now(); let source_json = serde_json::Value::Array( source_message_ids .iter() .map(|id| serde_json::Value::String(id.to_string())) .collect(), ); let stmt = sea_orm::Statement::from_sql_and_values( sea_orm::DbBackend::Postgres, "INSERT INTO room_compact_summary \ (id, room, from_seq, to_seq, summary, message_count, source_message_ids, created_at, updated_at) \ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)", vec![ id.into(), room_id.into(), from_seq.into(), to_seq.into(), summary.to_string().into(), (source_message_ids.len() as i32).into(), source_json.into(), now.into(), now.into(), ], ); self.db .execute_raw(stmt) .await .map_err(|e| AgentError::Internal(e.to_string()))?; Ok(RoomCompactRecord { id, room_id, from_seq, to_seq, summary: summary.to_string(), message_count: source_message_ids.len() as i32, source_message_ids: source_message_ids.to_vec(), created_at: now, }) } fn clean_dedupe_sort_messages(mut messages: Vec) -> Vec { messages.retain(|m| { m.revoked.is_none() && !m.content.trim().is_empty() && matches!(m.content_type, models::rooms::MessageContentType::Text) }); messages.sort_by_key(|m| (m.seq, m.send_at)); let mut seen = std::collections::HashSet::new(); messages .into_iter() .filter(|m| { let normalized = m .content .split_whitespace() .collect::>() .join(" ") .to_lowercase(); let key = format!("{}:{:?}:{}", m.sender_type, m.sender_id, normalized); seen.insert(key) }) .collect() } fn resolve_retain_count(config: CompactConfig, estimated_tokens: usize) -> usize { let level = if config.auto_level { CompactLevel::auto_select(estimated_tokens, config.token_threshold) } else { config.default_level }; level.retain_count() } pub async fn prepare_room_compact_context( &self, room_id: uuid::Uuid, requester_id: uuid::Uuid, user_names: Option>, config: CompactConfig, ) -> Result { let latest = self.latest_room_compact_record(room_id).await?; let cutoff_seq = latest.as_ref().map(|r| r.to_seq); let previous_summary = latest.as_ref().map(|r| r.summary.as_str()); let messages = self .fetch_room_messages_secure(room_id, requester_id) .await?; let messages = messages .into_iter() .filter(|m| cutoff_seq.map(|seq| m.seq > seq).unwrap_or(true)) .collect::>(); let messages = Self::clean_dedupe_sort_messages(messages); let user_ids: Vec = messages .iter() .filter_map(|m| m.sender_id) .collect::>() .into_iter() .collect(); let user_name_map = match user_names { Some(map) => map, None => self.get_user_name_map(&user_ids).await?, }; let sender_mapper = |m: &RoomMessageModel| { if let Some(user_id) = m.sender_id { if let Some(username) = user_name_map.get(&user_id) { return username.clone(); } } m.sender_type.to_string() }; let incremental_text = crate::compact::helpers::messages_to_text(&messages, sender_mapper); let estimate_input = match previous_summary { Some(summary) if !summary.is_empty() => format!("{}\n{}", summary, incremental_text), _ => incremental_text.clone(), }; let estimated_tokens = crate::tokent::count_message_text(&estimate_input, &self.model) .unwrap_or_else(|_| estimate_input.len() / 4); let retain_count = Self::resolve_retain_count(config, estimated_tokens); if estimated_tokens >= config.token_threshold && messages.len() > retain_count { let split_index = messages.len().saturating_sub(retain_count); let (to_summarize, retained_messages) = messages.split_at(split_index); let from_seq = to_summarize .first() .map(|m| m.seq) .unwrap_or(cutoff_seq.unwrap_or(0) + 1); let to_seq = to_summarize.last().map(|m| m.seq).unwrap_or(from_seq); let source_ids: Vec = to_summarize.iter().map(|m| m.id).collect(); let (summary, _usage) = self .summarize_room_increment(previous_summary, to_summarize, config.max_summary_tokens) .await?; let record = self .insert_room_compact_record(room_id, from_seq, to_seq, &summary, &source_ids) .await?; let retained = retained_messages .iter() .map(|m| Self::message_to_summary(m, &user_name_map)) .collect(); return Ok(RoomCompactContext { room_id, cutoff_seq: Some(record.to_seq), summary: Some(record.summary), retained, estimated_tokens, compacted: true, }); } let retained = messages .iter() .rev() .take(50) .collect::>() .into_iter() .rev() .map(|m| Self::message_to_summary(m, &user_name_map)) .collect(); Ok(RoomCompactContext { room_id, cutoff_seq, summary: latest.map(|r| r.summary), retained, estimated_tokens, compacted: false, }) } pub async fn compact_room( &self, room_id: uuid::Uuid, level: CompactLevel, user_names: Option>, requester_id: uuid::Uuid, context_window_tokens: i32, compaction_max_summary_ratio: f32, ) -> Result { let messages = self .fetch_room_messages_secure(room_id, requester_id) .await?; if messages.is_empty() { let room_exists = models::rooms::room::Entity::find_by_id(room_id) .one(&self.db) .await .map_err(|e| AgentError::Internal(e.to_string()))? .is_some(); if room_exists { return Err(AgentError::Internal("Access denied or room empty".into())); } else { return Err(AgentError::Internal("Room not found".into())); } } let user_ids: Vec = messages .iter() .filter_map(|m| m.sender_id) .collect::>() .into_iter() .collect(); let user_name_map = match user_names { Some(map) => map, None => self.get_user_name_map(&user_ids).await?, }; if messages.len() <= level.retain_count() { let retained: Vec = messages .iter() .map(|m| Self::message_to_summary(m, &user_name_map)) .collect(); return Ok(CompactSummary { session_id: uuid::Uuid::new_v4(), room_id, retained, summary: String::new(), compacted_at: chrono::Utc::now(), messages_compressed: 0, usage: None, }); } let retain_count = level.retain_count(); let split_index = messages.len().saturating_sub(retain_count); let (to_summarize, retained_messages) = messages.split_at(split_index); let retained: Vec = retained_messages .iter() .map(|m| Self::message_to_summary(m, &user_name_map)) .collect(); let max_summary_tokens = CompactConfig::summary_token_budget( context_window_tokens.max(0) as usize, compaction_max_summary_ratio, ); let (summary, remote_usage) = self .summarize_messages(to_summarize, max_summary_tokens) .await?; let summarized_text = to_summarize .iter() .map(|m| m.content.as_str()) .collect::>() .join("\n"); let usage = resolve_usage(remote_usage, &self.model, &summarized_text, &summary); Ok(CompactSummary { session_id: uuid::Uuid::new_v4(), room_id, retained, summary, compacted_at: chrono::Utc::now(), messages_compressed: to_summarize.len(), usage: Some(usage), }) } pub async fn compact_session( &self, session_id: uuid::Uuid, level: CompactLevel, user_names: Option>, context_window_tokens: i32, compaction_max_summary_ratio: f32, ) -> Result { let messages: Vec = RoomMessage::find() .filter(RmCol::Room.eq(session_id)) .order_by_asc(RmCol::Seq) .limit(10000) .all(&self.db) .await .map_err(|e| AgentError::Internal(e.to_string()))?; if messages.is_empty() { return Err(AgentError::Internal("session has no messages".into())); } let user_ids: Vec = messages .iter() .filter_map(|m| m.sender_id) .collect::>() .into_iter() .collect(); let user_name_map = match user_names { Some(map) => map, None => self.get_user_name_map(&user_ids).await?, }; if messages.len() <= level.retain_count() { let retained: Vec = messages .iter() .map(|m| Self::message_to_summary(m, &user_name_map)) .collect(); return Ok(CompactSummary { session_id, room_id: uuid::Uuid::nil(), retained, summary: String::new(), compacted_at: chrono::Utc::now(), messages_compressed: 0, usage: None, }); } let retain_count = level.retain_count(); let split_index = messages.len().saturating_sub(retain_count); let (to_summarize, retained_messages) = messages.split_at(split_index); let retained: Vec = retained_messages .iter() .map(|m| Self::message_to_summary(m, &user_name_map)) .collect(); let max_summary_tokens = CompactConfig::summary_token_budget( context_window_tokens.max(0) as usize, compaction_max_summary_ratio, ); let (summary, remote_usage) = self .summarize_messages(to_summarize, max_summary_tokens) .await?; let summarized_text = to_summarize .iter() .map(|m| m.content.as_str()) .collect::>() .join("\n"); let usage = resolve_usage(remote_usage, &self.model, &summarized_text, &summary); Ok(CompactSummary { session_id, room_id: uuid::Uuid::nil(), retained, summary, compacted_at: chrono::Utc::now(), messages_compressed: to_summarize.len(), usage: Some(usage), }) } }