gitdataai/libs/agent/compact/summarizer.rs

use models::rooms::room_message::Model as RoomMessageModel;
use models::users::user::{Column as UserCol, Entity as User};
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};

use crate::AgentError;
use crate::client::call_with_params;
use crate::client::types::ChatRequestMessage;
use crate::compact::types::MessageSummary;
use crate::tokent::TokenUsage;

impl super::CompactService {
    pub async fn summarize_room_increment(
        &self,
        previous_summary: Option<&str>,
        messages: &[RoomMessageModel],
        max_summary_tokens: usize,
    ) -> Result<(String, Option<TokenUsage>), AgentError> {
        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 = 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 body = crate::compact::helpers::messages_to_text(messages, sender_mapper);
        let previous = previous_summary
            .filter(|s| !s.trim().is_empty())
            .map(|s| format!("Previous compressed room summary:\n{}\n\n", s.trim()))
            .unwrap_or_default();

        let user_msg = ChatRequestMessage::user(format!(
            "Create an incremental room summary. Start from the previous summary if present, \
             then merge the new messages below. Deduplicate repeated messages, clean noise, \
             keep chronological order, and preserve decisions, facts, assignments/owners, \
             unresolved questions, and concrete next steps. The result MUST NOT exceed {} tokens.\n\n\
             Format:\n\
             **Summary:** <compact overview>\n\
             **Decisions:** <bullets or 'none'>\n\
             **Owners:** <bullets with owner -> task or 'none'>\n\
             **Open items:** <bullets or 'none'>\n\n\
             {}New messages:\n\n{}",
            max_summary_tokens, previous, body
        ));

        let response = call_with_params(
            &[user_msg],
            &self.model,
            &self.ai_client_config,
            0.2,
            max_summary_tokens.min(4096) as u32,
            None,
            None,
            None,
        )
        .await
        .map_err(|e| AgentError::OpenAi(e.to_string()))?;

        let remote_usage =
            TokenUsage::from_remote(response.input_tokens as u32, response.output_tokens as u32);

        Ok((response.content, remote_usage))
    }

    pub async fn summarize_messages(
        &self,
        messages: &[RoomMessageModel],
        max_summary_tokens: usize,
    ) -> Result<(String, Option<TokenUsage>), AgentError> {
        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 = 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 body = crate::compact::helpers::messages_to_text(messages, sender_mapper);

        let user_msg = ChatRequestMessage::user(format!(
            "Summarise the following conversation concisely, preserving all key facts, \
             decisions, and any pending or in-progress work. \
             The summary MUST NOT exceed {} tokens. \
             Use this format:\n\n\
             **Summary:** <one-paragraph overview>\n\
             **Key decisions:** <bullet list or 'none'>\n\
             **Open items:** <bullet list or 'none'>\n\n\
             Conversation:\n\n{}",
            max_summary_tokens, body
        ));

        let response = call_with_params(
            &[user_msg],
            &self.model,
            &self.ai_client_config,
            0.3,
            2048,
            None,
            None,
            None,
        )
        .await
        .map_err(|e| AgentError::OpenAi(e.to_string()))?;

        let remote_usage =
            TokenUsage::from_remote(response.input_tokens as u32, response.output_tokens as u32);

        Ok((response.content, remote_usage))
    }

    pub fn message_to_summary(
        m: &RoomMessageModel,
        user_name_map: &std::collections::HashMap<uuid::Uuid, String>,
    ) -> MessageSummary {
        let sender_name = if let Some(user_id) = m.sender_id {
            user_name_map
                .get(&user_id)
                .cloned()
                .unwrap_or_else(|| m.sender_type.to_string())
        } else {
            m.sender_type.to_string()
        };
        MessageSummary {
            id: m.id,
            sender_type: m.sender_type.clone(),
            sender_id: m.sender_id,
            sender_name,
            content: m.content.clone(),
            content_type: m.content_type.clone(),
            tool_call_id: None,
            send_at: m.send_at,
        }
    }

    pub async fn get_user_name_map(
        &self,
        user_ids: &[uuid::Uuid],
    ) -> Result<std::collections::HashMap<uuid::Uuid, String>, AgentError> {
        use std::collections::HashMap;
        let mut map = HashMap::new();
        if !user_ids.is_empty() {
            let users = User::find()
                .filter(UserCol::Uid.is_in(user_ids.to_vec()))
                .all(&self.db)
                .await
                .map_err(|e| AgentError::Internal(e.to_string()))?;
            for user in users {
                map.insert(user.uid, user.username);
            }
        }
        Ok(map)
    }
}