rune-chain-core 0.1.2

use serde::{Deserialize, Serialize};

use crate::ToolCall;

/// The originator of a [`Message`] in a conversation turn.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    /// Instructions that frame the model's behaviour for the whole conversation.
    System,
    /// A turn sent by the end user or caller.
    Human,
    /// A turn generated by the AI model.
    Ai,
    /// A synthetic turn injected by a tool or function-call result.
    Tool,
}

impl std::fmt::Display for Role {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Role::System => write!(f, "system"),
            Role::Human => write!(f, "human"),
            Role::Ai => write!(f, "ai"),
            Role::Tool => write!(f, "tool"),
        }
    }
}

/// Raw image payload that can be embedded in a message.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum ImageData {
    /// A publicly accessible image URL.
    Url { url: String },
    /// A base64-encoded image with an explicit MIME type.
    Base64 { mime_type: String, data: String },
}

/// A single part within a multi-part message content.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum ContentPart {
    /// A plain-text fragment.
    Text { text: String },
    /// An image fragment (URL or base64).
    Image { image: ImageData },
}

impl ContentPart {
    /// Convenience constructor for a text part.
    pub fn text(text: impl Into<String>) -> Self {
        ContentPart::Text { text: text.into() }
    }

    /// Convenience constructor for an image-URL part.
    pub fn image_url(url: impl Into<String>) -> Self {
        ContentPart::Image {
            image: ImageData::Url { url: url.into() },
        }
    }

    /// Convenience constructor for a base64-encoded image part.
    pub fn image_base64(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
        ContentPart::Image {
            image: ImageData::Base64 {
                mime_type: mime_type.into(),
                data: data.into(),
            },
        }
    }
}

/// The body of a [`Message`]: either plain text or a list of rich content parts.
///
/// `MessageContent` serialises as a JSON string for `Text` and as a JSON array
/// for `Parts`, matching the format expected by OpenAI, Anthropic, and Ollama.
///
/// # Constructing
///
/// Plain-text messages work through `From<String>` / `From<&str>`:
/// ```rust
/// use rune_chain_core::MessageContent;
///
/// let c: MessageContent = "hello".into();
/// assert_eq!(c.as_text(), "hello");
/// ```
///
/// Multi-modal messages use [`MessageContent::Parts`]:
/// ```rust
/// use rune_chain_core::{MessageContent, ContentPart};
///
/// let c = MessageContent::Parts(vec![
///     ContentPart::text("What is in this image?"),
///     ContentPart::image_url("https://example.com/cat.jpg"),
/// ]);
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessageContent {
    /// Plain UTF-8 text.
    Text(String),
    /// An ordered sequence of rich content parts (text + images).
    Parts(Vec<ContentPart>),
}

impl MessageContent {
    /// Return the text content.
    ///
    /// - For `Text(s)`: borrows `s` directly.
    /// - For `Parts(...)`: returns `""` — use [`text_content`](Self::text_content) to get all text.
    pub fn as_text(&self) -> &str {
        match self {
            MessageContent::Text(s) => s.as_str(),
            MessageContent::Parts(_) => "",
        }
    }

    /// Collect all text fragments into an owned `String`.
    ///
    /// For `Text(s)`: clones `s`.
    /// For `Parts(...)`: joins all `Text` parts with a space.
    pub fn text_content(&self) -> String {
        match self {
            MessageContent::Text(s) => s.clone(),
            MessageContent::Parts(parts) => parts
                .iter()
                .filter_map(|p| {
                    if let ContentPart::Text { text } = p {
                        Some(text.as_str())
                    } else {
                        None
                    }
                })
                .collect::<Vec<_>>()
                .join(" "),
        }
    }

    /// Return `true` if this is a plain-text message with no image parts.
    pub fn is_text(&self) -> bool {
        matches!(self, MessageContent::Text(_))
    }
}

impl From<String> for MessageContent {
    fn from(s: String) -> Self {
        MessageContent::Text(s)
    }
}

impl From<&str> for MessageContent {
    fn from(s: &str) -> Self {
        MessageContent::Text(s.to_string())
    }
}

impl From<&String> for MessageContent {
    fn from(s: &String) -> Self {
        MessageContent::Text(s.clone())
    }
}

impl std::fmt::Display for MessageContent {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.text_content())
    }
}

impl PartialEq<str> for MessageContent {
    fn eq(&self, other: &str) -> bool {
        self.as_text() == other
    }
}

impl PartialEq<&str> for MessageContent {
    fn eq(&self, other: &&str) -> bool {
        self.as_text() == *other
    }
}

impl PartialEq<String> for MessageContent {
    fn eq(&self, other: &String) -> bool {
        self.as_text() == other.as_str()
    }
}

/// A single turn in a conversation, carrying a [`Role`], rich [`MessageContent`],
/// and optional tool-call metadata for function-calling workflows.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Message {
    /// Who produced this message.
    pub role: Role,
    /// The message body — plain text or multi-modal parts.
    pub content: MessageContent,
    /// Tool calls the model wants to make (populated on [`Role::Ai`] turns when
    /// the model requests native function calling).
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_calls: Vec<ToolCall>,
    /// The call ID this result belongs to (populated on [`Role::Tool`] turns).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

impl Message {
    /// Create a new [`Message`] with an explicit role and content.
    ///
    /// # Example
    ///
    /// ```rust
    /// use rune_chain_core::{Message, Role};
    ///
    /// let msg = Message::new(Role::Human, "Hello!");
    /// assert_eq!(msg.role, Role::Human);
    /// assert_eq!(msg.content.as_text(), "Hello!");
    /// ```
    pub fn new(role: Role, content: impl Into<MessageContent>) -> Self {
        Self {
            role,
            content: content.into(),
            tool_calls: Vec::new(),
            tool_call_id: None,
        }
    }

    /// Shorthand for a [`Role::System`] message.
    pub fn system(content: impl Into<MessageContent>) -> Self {
        Self::new(Role::System, content)
    }

    /// Shorthand for a [`Role::Human`] message.
    pub fn human(content: impl Into<MessageContent>) -> Self {
        Self::new(Role::Human, content)
    }

    /// Shorthand for a [`Role::Ai`] message.
    pub fn ai(content: impl Into<MessageContent>) -> Self {
        Self::new(Role::Ai, content)
    }

    /// Create an AI message that carries pending tool calls.
    ///
    /// Used internally by function-calling agents to represent the model's
    /// "I want to call these tools" turn before observations are appended.
    ///
    /// # Example
    ///
    /// ```rust
    /// use rune_chain_core::{Message, ToolCall};
    ///
    /// let msg = Message::ai_with_tool_calls(
    ///     "",
    ///     vec![ToolCall::new("call_1", "upper_case", r#"{"input":"hello"}"#)],
    /// );
    /// assert_eq!(msg.tool_calls.len(), 1);
    /// ```
    pub fn ai_with_tool_calls(
        content: impl Into<MessageContent>,
        tool_calls: Vec<ToolCall>,
    ) -> Self {
        Self {
            role: Role::Ai,
            content: content.into(),
            tool_calls,
            tool_call_id: None,
        }
    }

    /// Create a tool-result message to feed back an observation from a tool call.
    ///
    /// # Example
    ///
    /// ```rust
    /// use rune_chain_core::Message;
    ///
    /// let msg = Message::tool_result("call_1", "HELLO");
    /// ```
    pub fn tool_result(call_id: impl Into<String>, content: impl Into<MessageContent>) -> Self {
        Self {
            role: Role::Tool,
            content: content.into(),
            tool_calls: Vec::new(),
            tool_call_id: Some(call_id.into()),
        }
    }
}