gitdataai/libs/agent/chat/context.rs

use crate::client::ChatRequestMessage;
use crate::compact::MessageSummary;
use chrono::{DateTime, Utc};
use models::rooms::MessageSenderType;
use models::rooms::room_message::Model as RoomMessageModel;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use uuid::Uuid;

/// Sender type for AI context, supporting all roles in the chat.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum AiContextSenderType {
    /// Regular user message
    User,
    /// AI assistant message
    Ai,
    /// System message (e.g., summary, notification)
    System,
    /// Developer message (for system-level instructions)
    Developer,
    /// Tool call message
    Function,
    /// Tool result message
    FunctionResult,
    Webhook,
}

impl AiContextSenderType {
    pub fn from_sender_type(sender_type: &models::rooms::MessageSenderType) -> Self {
        match sender_type {
            models::rooms::MessageSenderType::Ai => Self::Ai,
            models::rooms::MessageSenderType::System => Self::System,
            models::rooms::MessageSenderType::Tool => Self::FunctionResult,
            MessageSenderType::User => Self::User,
            MessageSenderType::Webhook => Self::Webhook,
        }
    }
}

/// Room message context for AI processing.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RoomMessageContext {
    pub uid: Uuid,
    pub sender_type: AiContextSenderType,
    pub sender_uid: Option<Uuid>,
    pub sender_name: Option<String>,
    pub content: String,
    pub content_type: models::rooms::MessageContentType,
    pub send_at: DateTime<Utc>,
    /// Tool call ID for FunctionResult messages, used to associate tool results with their calls.
    pub tool_call_id: Option<String>,
}

impl RoomMessageContext {
    pub fn from_model(model: &RoomMessageModel, sender_name: Option<String>) -> Self {
        Self {
            uid: model.id,
            sender_type: AiContextSenderType::from_sender_type(&model.sender_type),
            sender_uid: model.sender_id,
            sender_name,
            content: model.content.clone(),
            content_type: model.content_type.clone(),
            send_at: model.send_at,
            tool_call_id: Self::extract_tool_call_id(&model.content),
        }
    }

    fn extract_tool_call_id(content: &str) -> Option<String> {
        let content = content.trim();
        if let Ok(v) = serde_json::from_str::<Value>(content) {
            v.get("tool_call_id")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string())
        } else {
            None
        }
    }

    pub fn from_model_with_names(
        model: &RoomMessageModel,
        user_names: &HashMap<Uuid, String>,
    ) -> Self {
        let sender_name = model
            .sender_id
            .and_then(|uid| user_names.get(&uid).cloned());
        Self::from_model(model, sender_name)
    }

    pub fn to_message(&self) -> ChatRequestMessage {
        let name_str = self.sender_name.clone();
        match self.sender_type {
            AiContextSenderType::User => {
                let mut msg = ChatRequestMessage::user(self.display_content());
                if let Some(n) = name_str {
                    msg = msg.with_name(n);
                }
                msg
            }
            AiContextSenderType::Ai => {
                let mut msg = ChatRequestMessage::assistant(Some(self.display_content()), None);
                if let Some(n) = name_str {
                    msg = msg.with_name(n);
                }
                msg
            }
            AiContextSenderType::System => {
                let mut msg = ChatRequestMessage::system(&self.display_content());
                if let Some(n) = name_str {
                    msg = msg.with_name(n);
                }
                msg
            }
            AiContextSenderType::Developer => {
                let mut msg = ChatRequestMessage::developer(&self.display_content());
                if let Some(n) = name_str {
                    msg = msg.with_name(n);
                }
                msg
            }
            AiContextSenderType::Function => ChatRequestMessage::user(&self.content)
                .with_name(self.sender_name.as_deref().unwrap_or("unknown")),
            AiContextSenderType::FunctionResult => {
                let id = self
                    .tool_call_id
                    .clone()
                    .unwrap_or_else(|| "unknown".to_string());
                ChatRequestMessage::tool(id, self.display_content())
            }
            AiContextSenderType::Webhook => {
                let mut msg = ChatRequestMessage::user(&self.content);
                msg = msg.with_name(name_str.as_deref().unwrap_or("webhook"));
                msg
            }
        }
    }

    fn display_content(&self) -> String {
        let mut content = self.content.trim().to_string();
        if content.is_empty() {
            content = match self.content_type {
                models::rooms::MessageContentType::Text => "[empty]".to_string(),
                models::rooms::MessageContentType::Image => "[image]".to_string(),
                models::rooms::MessageContentType::Audio => "[audio]".to_string(),
                models::rooms::MessageContentType::Video => "[video]".to_string(),
                models::rooms::MessageContentType::File => "[file]".to_string(),
            };
        }

        if let Some(sender_name) = &self.sender_name {
            content = format!("[{}] {}", sender_name, content);
        }

        content
    }
}

impl From<&RoomMessageModel> for RoomMessageContext {
    fn from(model: &RoomMessageModel) -> Self {
        RoomMessageContext::from_model(model, None)
    }
}

impl From<MessageSummary> for RoomMessageContext {
    fn from(summary: MessageSummary) -> Self {
        // Map MessageSenderType to AiContextSenderType
        let sender_type = AiContextSenderType::from_sender_type(&summary.sender_type);
        // For FunctionResult (tool results), ensure tool_call_id is set
        let tool_call_id = if sender_type == AiContextSenderType::FunctionResult {
            summary.tool_call_id
        } else {
            None
        };
        Self {
            uid: summary.id,
            sender_type,
            sender_uid: summary.sender_id,
            sender_name: Some(summary.sender_name),
            content: summary.content,
            content_type: summary.content_type,
            send_at: summary.send_at,
            tool_call_id,
        }
    }
}