aisdk 0.5.2

//! Message types for the `aisdk` library.

use crate::core::{
    language_model::{LanguageModelResponseContentType, Usage},
    tools::{ToolCallInfo, ToolResultInfo},
};

/// The role of a participant in a conversation.
#[derive(Debug, Clone)]
pub enum Role {
    /// System-level instructions or context.
    System,
    /// Human user input.
    User,
    /// AI assistant response.
    Assistant,
}

/// A message in a conversation with a language model.
#[derive(Debug, Clone)]
pub enum Message {
    /// A system message providing context or instructions.
    System(SystemMessage),
    /// A user message containing input from the human.
    User(UserMessage),
    /// An assistant message containing the model's response.
    Assistant(AssistantMessage),
    /// A tool result message from executing a tool call.
    Tool(ToolResultInfo),
    /// A developer-specific message for advanced use cases.
    Developer(String),
}

/// A List of `Message`s
pub type Messages = Vec<Message>;

impl Message {
    /// Start a new conversation with an empty message list.
    ///
    /// Returns a `MessageBuilder<Conversation>`, allowing any number of
    /// `user`/`assistant` calls without forcing an initial system or user message.
    ///
    /// # Example
    /// ```
    /// use aisdk::core::Message;
    ///
    /// let mut msg = Message::conversation_builder();
    /// for _ in 0..10 {
    ///     msg = msg.user("hello");
    /// }
    /// let messages = msg.build();
    /// ```
    pub fn conversation_builder() -> MessageBuilder<Conversation> {
        MessageBuilder::conversation_builder()
    }

    /// Create a new message builder in the initial state.
    ///
    /// Returns a `MessageBuilder<Initial>` that enforces type-safe order:
    /// the first message **must** be either a system prompt or a user message.
    /// After that, the builder transitions to `Conversation` state and allows
    /// free mixing of user/assistant messages but not system prompts.
    ///
    /// # Example
    /// ```
    /// use aisdk::core::Message;
    ///
    /// let msgs = Message::builder()
    ///     .system("You are helpful.")
    ///     .user("Hello!")
    ///     .assistant("Hi there.")
    ///     .build();
    /// ```
    pub fn builder() -> MessageBuilder<Initial> {
        MessageBuilder::default()
    }
}

/// A system message that provides context or instructions to the model.
#[derive(Debug, Clone)]
pub struct SystemMessage {
    /// The text content of the system message.
    pub content: String,
}

impl SystemMessage {
    /// Creates a new system message with the given content.
    pub fn new(content: impl Into<String>) -> Self {
        Self {
            content: content.into(),
        }
    }
}

impl From<String> for SystemMessage {
    fn from(value: String) -> Self {
        Self::new(value)
    }
}

impl From<&str> for SystemMessage {
    fn from(value: &str) -> Self {
        Self::new(value)
    }
}

/// A user message containing input from the human participant.
#[derive(Debug, Clone)]
pub struct UserMessage {
    /// The text content of the user message.
    pub content: String,
}

impl UserMessage {
    /// Creates a new user message with the given content.
    pub fn new(content: impl Into<String>) -> Self {
        Self {
            content: content.into(),
        }
    }
}

impl From<String> for UserMessage {
    fn from(value: String) -> Self {
        Self::new(value)
    }
}

impl From<&str> for UserMessage {
    fn from(value: &str) -> Self {
        Self::new(value)
    }
}

/// A message generated by the language model assistant.
#[derive(Default, Debug, Clone)]
pub struct AssistantMessage {
    /// The content of the assistant's response.
    pub content: LanguageModelResponseContentType,
    /// Optional usage statistics for the response.
    pub usage: Option<Usage>,
}

impl From<String> for AssistantMessage {
    fn from(value: String) -> Self {
        Self {
            content: value.into(),
            usage: None,
        }
    }
}

impl AssistantMessage {
    /// Creates a new assistant message with the given content and usage.
    pub fn new(content: LanguageModelResponseContentType, usage: Option<Usage>) -> Self {
        Self { content, usage }
    }
}

/// Type-state marker for the initial message builder state.
///
/// In this state, the first message must be either a system prompt or a user message.
#[derive(Debug, Clone)]
pub struct Initial;

/// Type-state marker for the conversation message builder state.
///
/// In this state, only user and assistant messages can be added.
#[derive(Debug, Clone)]
pub struct Conversation;

/// A type-state builder for constructing message lists safely.
///
/// This builder ensures that messages are added in a valid order,
/// preventing invalid conversation structures.
#[derive(Debug, Clone)]
pub struct MessageBuilder<State = Initial> {
    messages: Messages,
    state: std::marker::PhantomData<State>,
}

impl MessageBuilder {
    /// Creates a new message builder starting in the conversation state.
    ///
    /// This allows building conversations without requiring an initial system or user message.
    pub fn conversation_builder() -> MessageBuilder<Conversation> {
        MessageBuilder {
            messages: Vec::new(),
            state: std::marker::PhantomData,
        }
    }
}

impl Default for MessageBuilder {
    fn default() -> Self {
        MessageBuilder {
            messages: Vec::new(),
            state: std::marker::PhantomData,
        }
    }
}

impl<State> MessageBuilder<State> {
    /// Builds the message list.
    pub fn build(self) -> Messages {
        self.messages
    }
}

impl MessageBuilder<Initial> {
    /// Adds a system message and transitions to the conversation state.
    ///
    /// # Parameters
    ///
    /// * `content` - The system message content.
    ///
    /// # Returns
    ///
    /// The builder in the conversation state.
    pub fn system(mut self, content: impl Into<String>) -> MessageBuilder<Conversation> {
        self.messages.push(Message::System(content.into().into()));
        MessageBuilder {
            messages: self.messages,
            state: std::marker::PhantomData,
        }
    }

    /// Adds a user message and transitions to the conversation state.
    ///
    /// # Parameters
    ///
    /// * `content` - The user message content.
    ///
    /// # Returns
    ///
    /// The builder in the conversation state.
    pub fn user(mut self, content: impl Into<String>) -> MessageBuilder<Conversation> {
        self.messages.push(Message::User(content.into().into()));
        MessageBuilder {
            messages: self.messages,
            state: std::marker::PhantomData,
        }
    }
}

impl MessageBuilder<Conversation> {
    /// Adds a user message to the conversation.
    ///
    /// # Parameters
    ///
    /// * `content` - The user message content.
    ///
    /// # Returns
    ///
    /// The builder with the message added.
    pub fn user(mut self, content: impl Into<String>) -> MessageBuilder<Conversation> {
        self.messages.push(Message::User(content.into().into()));
        MessageBuilder {
            messages: self.messages,
            state: std::marker::PhantomData,
        }
    }

    /// Adds an assistant message to the conversation.
    ///
    /// # Parameters
    ///
    /// * `content` - The assistant message content.
    ///
    /// # Returns
    ///
    /// The builder with the message added.
    pub fn assistant(mut self, content: impl Into<String>) -> MessageBuilder<Conversation> {
        self.messages
            .push(Message::Assistant(content.into().into()));
        MessageBuilder {
            messages: self.messages,
            state: std::marker::PhantomData,
        }
    }
}

/// A message tagged with its step id in a list of messages
/// used for tracking steps in a conversation
#[derive(Debug, Clone)]
pub(crate) struct TaggedMessage {
    pub step_id: usize,
    pub message: Message,
}

impl TaggedMessage {
    pub fn new(step_id: usize, message: Message) -> Self {
        Self { step_id, message }
    }

    pub fn initial_step_msg(message: Message) -> Self {
        Self {
            step_id: 0,
            message,
        }
    }
}

// conversions assume message is in initial step
impl From<Message> for TaggedMessage {
    fn from(value: Message) -> Self {
        Self::initial_step_msg(value)
    }
}

// conversions disregard tagging information
impl From<TaggedMessage> for Message {
    fn from(value: TaggedMessage) -> Self {
        value.message
    }
}

/// Helper trait for extracting messages from TaggedMessage collections
pub(crate) trait TaggedMessageHelpers {
    fn extract_tool_calls(&self) -> Option<Vec<ToolCallInfo>>;
    fn extract_tool_results(&self) -> Option<Vec<ToolResultInfo>>;
}

impl TaggedMessageHelpers for [TaggedMessage] {
    fn extract_tool_calls(&self) -> Option<Vec<ToolCallInfo>> {
        let calls: Vec<ToolCallInfo> = self
            .iter()
            .filter_map(|msg| match msg.message {
                Message::Assistant(AssistantMessage {
                    content: LanguageModelResponseContentType::ToolCall(ref tool_info),
                    ..
                }) => Some(tool_info.clone()),
                _ => None,
            })
            .collect();
        if calls.is_empty() { None } else { Some(calls) }
    }

    fn extract_tool_results(&self) -> Option<Vec<ToolResultInfo>> {
        let results: Vec<ToolResultInfo> = self
            .iter()
            .filter_map(|msg| match msg.message {
                Message::Tool(ref info) => Some(info.clone()),
                _ => None,
            })
            .collect();
        if results.is_empty() {
            None
        } else {
            Some(results)
        }
    }
}