pe_core/

message.rs

//! Message type hierarchy — flows through every agent and LLM call.
//!
//! Based on Group 16 of the pre-plan. Every message has content, optional id,
//! optional name, and provider-specific metadata passthrough.

use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;

/// Top-level message enum — all message types the system handles.
///
/// # REVIEW(002): `#[non_exhaustive]` per library-design.md
/// This enum will grow (e.g. FunctionMessage, ToolResultMessage, RemoveMessage).
/// Without non_exhaustive, adding a variant is a breaking change for downstream
/// match arms. With it, users must have a `_ =>` arm, enabling safe evolution.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "role")]
#[non_exhaustive]
pub enum Message {
    Human(HumanMessage),
    Ai(AiMessage),
    System(SystemMessage),
    Tool(ToolMessage),
}

/// User input — can contain text, images, audio, files (multimodal)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HumanMessage {
    pub content: MessageContent,
    #[serde(default)]
    pub id: Option<String>,
    #[serde(default)]
    pub name: Option<String>,
}

/// LLM output — contains tool calls, usage metadata, provider metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AiMessage {
    pub content: MessageContent,
    #[serde(default)]
    pub tool_calls: Vec<ToolCall>,
    #[serde(default)]
    pub invalid_tool_calls: Vec<InvalidToolCall>,
    #[serde(default)]
    pub usage_metadata: Option<UsageMetadata>,
    #[serde(default)]
    pub response_metadata: HashMap<String, serde_json::Value>,
    #[serde(default)]
    pub id: Option<String>,
}

/// Instructions injected before conversation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemMessage {
    pub content: String,
    #[serde(default)]
    pub id: Option<String>,
}

/// Result of a tool execution — must match tool_call_id from AiMessage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolMessage {
    pub content: String,
    pub tool_call_id: String,
    #[serde(default)]
    pub name: Option<String>,
    #[serde(default)]
    pub id: Option<String>,
    /// Tool-specific metadata from execution (result count, source, etc.).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, Value>>,
    /// Execution duration in milliseconds.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub duration_ms: Option<u64>,
}

/// Message content — either plain text or multimodal content blocks
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
#[non_exhaustive]
pub enum MessageContent {
    Text(String),
    Blocks(Vec<ContentBlock>),
}

/// Content block types for multimodal messages
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
#[non_exhaustive]
pub enum ContentBlock {
    Text { text: String },
    Image { url: String },
    Audio { data: String },
    File { path: String, mime_type: String },
    Reasoning { content: String },
    Citation { source: String, quote: String },
}

/// Structured tool invocation request from LLM
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub id: String,
    pub name: String,
    pub args: serde_json::Value,
}

/// Malformed tool call — LLM tried but failed schema validation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InvalidToolCall {
    pub id: String,
    pub name: String,
    pub args: String,
    pub error: String,
}

/// Token usage from a single LLM call
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UsageMetadata {
    pub input_tokens: u32,
    pub output_tokens: u32,
    pub total_tokens: u32,
    #[serde(default)]
    pub input_token_details: Option<InputTokenDetails>,
    #[serde(default)]
    pub output_token_details: Option<OutputTokenDetails>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct InputTokenDetails {
    #[serde(default)]
    pub cache_read: u32,
    #[serde(default)]
    pub cache_write: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct OutputTokenDetails {
    #[serde(default)]
    pub reasoning_tokens: u32,
}

/// Sentinel value to clear a message list entirely
pub const REMOVE_ALL_MESSAGES: &str = "__REMOVE_ALL_MESSAGES__";

impl Message {
    pub fn human(content: impl Into<String>) -> Self {
        Message::Human(HumanMessage {
            content: MessageContent::Text(content.into()),
            id: None,
            name: None,
        })
    }

    pub fn system(content: impl Into<String>) -> Self {
        Message::System(SystemMessage {
            content: content.into(),
            id: None,
        })
    }

    pub fn ai(content: impl Into<String>) -> Self {
        Message::Ai(AiMessage {
            content: MessageContent::Text(content.into()),
            tool_calls: vec![],
            invalid_tool_calls: vec![],
            usage_metadata: None,
            response_metadata: HashMap::new(),
            id: None,
        })
    }

    pub fn tool(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
        Message::Tool(ToolMessage {
            content: content.into(),
            tool_call_id: tool_call_id.into(),
            name: None,
            id: None,
            metadata: None,
            duration_ms: None,
        })
    }

    /// Create a tool message with metadata and duration.
    pub fn tool_with_metadata(
        content: impl Into<String>,
        tool_call_id: impl Into<String>,
        metadata: Option<HashMap<String, Value>>,
        duration_ms: Option<u64>,
    ) -> Self {
        Message::Tool(ToolMessage {
            content: content.into(),
            tool_call_id: tool_call_id.into(),
            name: None,
            id: None,
            metadata,
            duration_ms,
        })
    }

    pub fn id(&self) -> Option<&str> {
        match self {
            Message::Human(m) => m.id.as_deref(),
            Message::Ai(m) => m.id.as_deref(),
            Message::System(m) => m.id.as_deref(),
            Message::Tool(m) => m.id.as_deref(),
        }
    }
}

impl MessageContent {
    pub fn as_text(&self) -> Option<&str> {
        match self {
            MessageContent::Text(t) => Some(t),
            MessageContent::Blocks(_) => None,
        }
    }
}

// NOTE: add_messages reducer lives in reducers.rs (single source of truth).
// The canonical version is pe_core::reducers::add_messages, re-exported at crate root.