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::{EntityTrait, QueryFilter, QueryOrder, QuerySelect};

use crate::compact::types::CompactLevel;
use crate::tokent::resolve_usage;
use crate::{AgentError, CompactSummary, MessageSummary};

impl super::CompactService {
    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 =
            (context_window_tokens as f32 * compaction_max_summary_ratio) as usize;

        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 =
            (context_window_tokens as f32 * compaction_max_summary_ratio) as usize;

        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),
        })
    }
}