gitdataai/libs/agent/compact/room_compactor.rs

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<Option<RoomCompactRecord>, 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::<Vec<_>>()
            })
            .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<RoomCompactRecord, AgentError> {
        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<RoomMessageModel>) -> Vec<RoomMessageModel> {
        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::<Vec<_>>()
                    .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<std::collections::HashMap<uuid::Uuid, String>>,
        config: CompactConfig,
    ) -> Result<RoomCompactContext, AgentError> {
        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::<Vec<_>>();
        let messages = Self::clean_dedupe_sort_messages(messages);

        let user_ids: Vec<uuid::Uuid> = messages
            .iter()
            .filter_map(|m| m.sender_id)
            .collect::<std::collections::HashSet<_>>()
            .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<uuid::Uuid> = 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::<Vec<_>>()
            .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<std::collections::HashMap<uuid::Uuid, String>>,
        requester_id: uuid::Uuid,
        context_window_tokens: i32,
        compaction_max_summary_ratio: f32,
    ) -> Result<CompactSummary, AgentError> {
        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<uuid::Uuid> = messages
            .iter()
            .filter_map(|m| m.sender_id)
            .collect::<std::collections::HashSet<_>>()
            .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<MessageSummary> = 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<MessageSummary> = 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::<Vec<_>>()
            .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<std::collections::HashMap<uuid::Uuid, String>>,
        context_window_tokens: i32,
        compaction_max_summary_ratio: f32,
    ) -> Result<CompactSummary, AgentError> {
        let messages: Vec<RoomMessageModel> = 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<uuid::Uuid> = messages
            .iter()
            .filter_map(|m| m.sender_id)
            .collect::<std::collections::HashSet<_>>()
            .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<MessageSummary> = 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<MessageSummary> = 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::<Vec<_>>()
            .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),
        })
    }
}