flashmind-types 0.1.0

//! LLM wire-format types — [`Role`], [`Message`], [`ToolCall`], and [`ContentPart`].
//!
//! These are the types serialized into JSON for API calls and received in responses.
//! They are distinct from the richer intermediate representation in `flashmind-core`
//! (`Conversation` / `ConversationEntry`).
//!
//! # Key types
//!
//! | Type | Role |
//! |------|------|
//! | [`Message`] | Single message in API request/response (role, content, tool calls, parts) |
//! | [`Role`] | Message role: System, User, Assistant, Developer, Tool |
//! | [`ToolCall`] | Tool invocation requested by the assistant |
//! | [`ToolResult`] | Result of a tool call (returned as `role: "tool"` message) |
//! | [`ContentPart`] | Multimodal content block: text, image, document, video, audio |

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

/// OpenAI-compatible message role.
///
/// Serialises to lowercase strings (`"system"`, `"user"`, `"assistant"`, `"developer"`, `"tool"`).
#[derive(
    Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, strum::Display, strum::EnumString,
)]
#[strum(serialize_all = "lowercase")]
#[serde(rename_all = "lowercase")]
pub enum Role {
    /// System-level instructions that set the model's behavior and constraints.
    System,
    /// Messages from the human user (questions, commands, feedback).
    User,
    /// Responses generated by the assistant/model.
    Assistant,
    /// Developer-level system messages (OpenAI-specific; hidden from the model's view in some providers).
    Developer,
    /// Results returned by tool execution, linked back to an assistant's [`ToolCall`] via `tool_call_id`.
    Tool,
}

/// A tool call requested by the assistant in its response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    /// Server-assigned ID that must be echoed back in the corresponding tool result.
    pub id: String,
    /// Name of the tool to invoke.
    pub name: String,
    /// JSON object of argument values.
    pub arguments: serde_json::Value,
}

/// Result returned by a tool, attached as a `role: "tool"` message in the
/// conversation history.
///
/// This is the wire-format representation of a tool result — the version that
/// gets serialized into the LLM API request. It is distinct from
/// [`crate::tool::ToolResult`], which is the richer, internal representation
/// used during tool execution (carrying sources, diffs, and interrupt state).
///
/// The conversion from [`crate::tool::ToolResult`] to this type happens in the
/// agent loop when building messages for the LLM provider.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
    /// ID of the tool call this result corresponds to. Must match the
    /// `id` field of the original [`ToolCall`] so the LLM can pair them.
    pub tool_call_id: String,
    /// Text output of the tool execution. For failures this contains the
    /// error message; for successes it contains the tool's return value.
    pub output: String,
    /// Whether the tool executed successfully (`true`) or failed (`false`).
    pub success: bool,
}

/// Multimodal content block within a user message.
///
/// Serialised with a `type` discriminant matching the OpenAI `ContentPart` schema:
/// `"text"`, `"image"`, `"document"`, `"video"`, `"audio"`.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum ContentPart {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "image")]
    Image {
        media_type: String,
        /// Base64-encoded image data.
        data: String,
    },
    #[serde(rename = "image_url")]
    ImageUrl { url: String },
    #[serde(rename = "document")]
    Document {
        media_type: String,
        filename: String,
        data: String,
    },
    #[serde(rename = "video")]
    Video {
        media_type: String,
        filename: String,
        data: String,
    },
    #[serde(rename = "video_url")]
    VideoUrl { url: String },
    #[serde(rename = "audio")]
    Audio {
        media_type: String,
        filename: String,
        data: String,
    },
}

/// A single message in the LLM API request / response payload.
///
/// For tool results set `tool_call_id` **in addition to** `content` (which holds the tool's output).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// Message role (`system`, `user`, `assistant`, `developer`, `tool`).
    pub role: Role,
    /// Text content of the message.
    pub content: String,
    /// Tool calls requested by the assistant (only on `assistant` messages).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    /// ID of the tool call this is a result for (only on `tool` messages).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
    /// Multimodal content parts (images, documents, audio). When present, `content` is also included as a text part.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parts: Option<Vec<ContentPart>>,
    /// Timestamp when the message was created.
    #[serde(default = "chrono::Utc::now")]
    pub timestamp: DateTime<Utc>,
}

impl Message {
    fn new(role: Role, content: String) -> Self {
        Self {
            role,
            content,
            tool_calls: None,
            tool_call_id: None,
            parts: None,
            timestamp: Utc::now(),
        }
    }

    /// Create a plain-text user message.
    pub fn user(content: impl Into<String>) -> Self {
        Self::new(Role::User, content.into())
    }

    /// Create an assistant text message (no tool calls).
    pub fn assistant(content: impl Into<String>) -> Self {
        Self::new(Role::Assistant, content.into())
    }

    /// Create an assistant message that includes tool call requests.
    pub fn assistant_with_tool_calls(
        content: impl Into<String>,
        tool_calls: Vec<ToolCall>,
    ) -> Self {
        Self {
            tool_calls: Some(tool_calls),
            ..Self::new(Role::Assistant, content.into())
        }
    }

    /// Create a system prompt message.
    pub fn system(content: impl Into<String>) -> Self {
        Self::new(Role::System, content.into())
    }

    /// Create a developer message (system-level instructions using the `developer` role).
    pub fn developer(content: impl Into<String>) -> Self {
        Self::new(Role::Developer, content.into())
    }

    /// Create a tool-result message linked to a prior `ToolCall.id`.
    pub fn tool_result(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            tool_call_id: Some(tool_call_id.into()),
            ..Self::new(Role::Tool, content.into())
        }
    }

    /// Create a multimodal user message with inline attachments.
    pub fn user_with_parts(content: impl Into<String>, parts: Vec<ContentPart>) -> Self {
        Self {
            parts: if parts.is_empty() { None } else { Some(parts) },
            ..Self::new(Role::User, content.into())
        }
    }

    /// Replace binary content parts (image, document, video, audio) with
    /// text placeholders so the message can be sent to models that don't
    /// support those modalities.
    ///
    /// Example:
    /// - `Image { media_type: "image/jpeg", .. }` → `Text { text: "[Image: image/jpeg]" }`
    pub fn strip_binary_data(&self) -> Self {
        let parts = self.parts.as_ref().map(|parts| {
            parts
                .iter()
                .map(|p| match p {
                    ContentPart::Image { media_type, .. } => ContentPart::Text {
                        text: format!("[Image: {}]", media_type),
                    },
                    ContentPart::ImageUrl { url } => ContentPart::Text {
                        text: format!("[Image: {}]", url),
                    },
                    ContentPart::Document {
                        media_type,
                        filename,
                        ..
                    } => ContentPart::Text {
                        text: format!("[Document: {} ({})]", filename, media_type),
                    },
                    ContentPart::Video {
                        media_type,
                        filename,
                        ..
                    } => ContentPart::Text {
                        text: format!("[Video: {} ({})]", filename, media_type),
                    },
                    ContentPart::VideoUrl { url } => ContentPart::Text {
                        text: format!("[Video: {}]", url),
                    },
                    ContentPart::Audio {
                        media_type,
                        filename,
                        ..
                    } => ContentPart::Text {
                        text: format!("[Audio: {} ({})]", filename, media_type),
                    },
                    other => other.clone(),
                })
                .collect()
        });

        Self {
            parts,
            ..self.clone()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_message_serialization() {
        let msg = Message::user("Hello");
        let json = serde_json::to_string(&msg).unwrap();
        assert!(json.contains("\"role\":\"user\""));
        assert!(json.contains("Hello"));
    }

    #[test]
    fn test_message_constructors() {
        let user = Message::user("hi");
        assert_eq!(user.role, Role::User);
        assert_eq!(user.content, "hi");

        let assistant = Message::assistant("hello");
        assert_eq!(assistant.role, Role::Assistant);

        let system = Message::system("you are helpful");
        assert_eq!(system.role, Role::System);

        let developer = Message::developer("injected context");
        assert_eq!(developer.role, Role::Developer);
        assert_eq!(developer.content, "injected context");

        let tool = Message::tool_result("call-1", "result data");
        assert_eq!(tool.role, Role::Tool);
        assert_eq!(tool.tool_call_id.as_deref(), Some("call-1"));
    }

    #[test]
    fn test_content_part_serde() {
        let text = ContentPart::Text {
            text: "hello".into(),
        };
        let json = serde_json::to_string(&text).unwrap();
        assert!(json.contains("\"type\":\"text\""));
        let roundtrip: ContentPart = serde_json::from_str(&json).unwrap();
        match roundtrip {
            ContentPart::Text { text } => assert_eq!(text, "hello"),
            _ => panic!("expected text"),
        }

        let img = ContentPart::Image {
            media_type: "image/jpeg".into(),
            data: "abc123".into(),
        };
        let json = serde_json::to_string(&img).unwrap();
        assert!(json.contains("\"type\":\"image\""));
    }

    #[test]
    fn test_strip_binary_data() {
        let msg = Message::user_with_parts(
            "photo",
            vec![
                ContentPart::Text {
                    text: "a cat".into(),
                },
                ContentPart::Image {
                    media_type: "image/jpeg".into(),
                    data: "bigdata".into(),
                },
            ],
        );
        let stripped = msg.strip_binary_data();
        let parts = stripped.parts.unwrap();
        assert_eq!(parts.len(), 2);
        match &parts[1] {
            ContentPart::Text { text } => assert!(text.contains("image/jpeg")),
            _ => panic!("expected text placeholder for image"),
        }
    }

    #[test]
    fn test_role_serde() {
        let role = Role::Assistant;
        let json = serde_json::to_string(&role).unwrap();
        assert_eq!(json, "\"assistant\"");

        let parsed: Role = serde_json::from_str("\"tool\"").unwrap();
        assert_eq!(parsed, Role::Tool);
    }
}