gitdataai/libs/agent/chat/context.rs

use async_openai::types::chat::{
    ChatCompletionRequestAssistantMessage, ChatCompletionRequestAssistantMessageContent,
    ChatCompletionRequestDeveloperMessage, ChatCompletionRequestDeveloperMessageContent,
    ChatCompletionRequestFunctionMessage, ChatCompletionRequestMessage,
    ChatCompletionRequestSystemMessage, ChatCompletionRequestSystemMessageContent,
    ChatCompletionRequestToolMessage, ChatCompletionRequestToolMessageContent,
    ChatCompletionRequestUserMessage, ChatCompletionRequestUserMessageContent,
};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use uuid::Uuid;

use crate::compact::MessageSummary;
use models::rooms::room_message::Model as RoomMessageModel;

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

impl AiContextSenderType {
    pub fn from_sender_type(sender_type: &models::rooms::MessageSenderType) -> Self {
        match sender_type {
            models::rooms::MessageSenderType::Member => Self::User,
            models::rooms::MessageSenderType::Admin => Self::User,
            models::rooms::MessageSenderType::Owner => Self::User,
            models::rooms::MessageSenderType::Ai => Self::Ai,
            models::rooms::MessageSenderType::System => Self::System,
            models::rooms::MessageSenderType::Tool => Self::Function,
            models::rooms::MessageSenderType::Guest => Self::User,
        }
    }
}

/// 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) -> ChatCompletionRequestMessage {
        match self.sender_type {
            AiContextSenderType::User => {
                ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage {
                    content: ChatCompletionRequestUserMessageContent::Text(self.display_content()),
                    name: self.sender_name.clone(),
                })
            }
            AiContextSenderType::Ai => {
                ChatCompletionRequestMessage::Assistant(ChatCompletionRequestAssistantMessage {
                    content: Some(ChatCompletionRequestAssistantMessageContent::Text(
                        self.display_content(),
                    )),
                    name: self.sender_name.clone(),
                    refusal: None,
                    audio: None,
                    tool_calls: None,
                    #[allow(deprecated)]
                    function_call: None,
                })
            }
            AiContextSenderType::System => {
                ChatCompletionRequestMessage::System(ChatCompletionRequestSystemMessage {
                    content: ChatCompletionRequestSystemMessageContent::Text(
                        self.display_content(),
                    ),
                    name: self.sender_name.clone(),
                })
            }
            AiContextSenderType::Developer => {
                ChatCompletionRequestMessage::Developer(ChatCompletionRequestDeveloperMessage {
                    content: ChatCompletionRequestDeveloperMessageContent::Text(
                        self.display_content(),
                    ),
                    name: self.sender_name.clone(),
                })
            }
            AiContextSenderType::Function => {
                ChatCompletionRequestMessage::Function(ChatCompletionRequestFunctionMessage {
                    content: Some(self.content.clone()),
                    name: self.display_content(), // Function name is stored in content
                })
            }
            AiContextSenderType::FunctionResult => {
                ChatCompletionRequestMessage::Tool(ChatCompletionRequestToolMessage {
                    content: ChatCompletionRequestToolMessageContent::Text(self.display_content()),
                    tool_call_id: self
                        .tool_call_id
                        .clone()
                        .unwrap_or_else(|| "unknown".to_string()),
                })
            }
        }
    }

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