ai_sdk_provider/language_model/

prompt.rs

use super::content::*;
use crate::Error;
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::Value;
use std::ops::{Deref, DerefMut};

/// A container for conversation messages to send to a language model.
///
/// The `Prompt` struct wraps a vector of `Message` objects and provides flexible
/// conversion traits to support multiple input formats. This enables you to construct
/// prompts from simple strings, structured message arrays, or individual messages.
///
/// # Construction Methods
///
/// `Prompt` can be created in several ways:
/// - From a simple string: `Prompt::from("Hello")`
/// - From a slice: `Prompt::from("Hello")`
/// - From a Vec of messages: `Prompt::from(vec![message1, message2])`
/// - From a single message: `Prompt::from(message)`
/// - From JSON via TryFrom: `json_value.try_into()`
///
/// # Flexibility
///
/// The struct implements `Deref` and `DerefMut` to provide Vec-like access to the
/// underlying messages. You can use standard Vec methods like `len()`, `push()`, and `iter()`.
///
/// # Usage
///
/// ```ignore
/// // Simple text prompt
/// let prompt: Prompt = "What is 2+2?".into();
///
/// // Structured messages
/// let messages = vec![
///     Message::System { content: "You are helpful".into() },
///     Message::User { content: vec![UserContentPart::Text { text: "Hello".into() }] },
/// ];
/// let prompt: Prompt = messages.into();
///
/// // JSON conversion
/// let json = json!({"role": "user", "content": "Hello"});
/// let prompt: Prompt = json.try_into()?;
/// ```
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct Prompt(Vec<Message>);

impl Prompt {
    /// Creates a new empty prompt with no messages.
    ///
    /// # Returns
    ///
    /// A prompt containing an empty message vector that can be populated with `push()`.
    pub fn new() -> Self {
        Self(Vec::new())
    }

    /// Creates a prompt from a single message.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to initialize the prompt with
    ///
    /// # Returns
    ///
    /// A prompt containing the provided message as its sole element.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let msg = Message::User { content: vec![UserContentPart::Text { text: "Hello".into() }] };
    /// let prompt = Prompt::from_message(msg);
    /// ```
    pub fn from_message(message: Message) -> Self {
        Self(vec![message])
    }
}

/// Enables Vec-like access to the underlying message vector.
///
/// Through this implementation, you can use standard Vec methods directly on Prompt,
/// such as `len()`, `push()`, `iter()`, and indexing operations.
impl Deref for Prompt {
    type Target = Vec<Message>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Enables mutable Vec-like access to the underlying message vector.
impl DerefMut for Prompt {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

/// Creates a prompt from a String, automatically treating it as a user message.
///
/// The string is converted into a User message with a single text content part.
/// This provides a convenient way to create simple text prompts.
impl From<String> for Prompt {
    fn from(s: String) -> Self {
        Prompt(vec![Message::User {
            content: vec![UserContentPart::Text { text: s }],
        }])
    }
}

/// Creates a prompt from a string slice, automatically treating it as a user message.
impl From<&str> for Prompt {
    fn from(s: &str) -> Self {
        Prompt::from(s.to_string())
    }
}

/// Creates a prompt from a vector of messages.
///
/// Useful when you have multiple messages with different roles that you want to
/// form a complete conversation history.
impl From<Vec<Message>> for Prompt {
    fn from(v: Vec<Message>) -> Self {
        Prompt(v)
    }
}

/// Creates a prompt from a single message.
///
/// Wraps the message in a vector to form a single-message prompt.
impl From<Message> for Prompt {
    fn from(m: Message) -> Self {
        Prompt(vec![m])
    }
}

/// Creates a prompt from a JSON value with flexible input formats.
///
/// This trait supports converting from various JSON structures:
/// - A JSON string: converted to a User message with text content
/// - A JSON object: deserialized as a single Message
/// - A JSON array: deserialized as an array of Message objects
///
/// # Errors
///
/// Returns an error if the JSON structure doesn't match one of the supported formats
/// or if deserialization fails.
impl TryFrom<Value> for Prompt {
    type Error = Error;

    fn try_from(value: Value) -> std::result::Result<Self, Self::Error> {
        match value {
            Value::String(s) => Ok(Prompt::from(s)),

            Value::Array(arr) => {
                let messages: Vec<Message> = serde_json::from_value(Value::Array(arr))
                    .map_err(|e| format!("Invalid prompt array: {}", e))?;
                Ok(Prompt(messages))
            }

            Value::Object(obj) => {
                let message: Message = serde_json::from_value(Value::Object(obj))
                    .map_err(|e| format!("Invalid prompt object: {}", e))?;
                Ok(Prompt(vec![message]))
            }

            _ => Err("JSON must be a string, object, or array of messages".into()),
        }
    }
}

/// Custom deserializer that handles both string and array representations of user content.
///
/// This deserializer provides flexibility when deserializing user message content.
/// It can handle either a simple string (converted to a text part) or an array of
/// structured content parts. This allows both `"Hello"` and `[{"type": "text", "text": "Hello"}]`
/// to deserialize successfully.
fn deserialize_user_content<'de, D>(
    deserializer: D,
) -> std::result::Result<Vec<UserContentPart>, D::Error>
where
    D: Deserializer<'de>,
{
    let value = Value::deserialize(deserializer)?;
    match value {
        Value::String(s) => Ok(vec![UserContentPart::Text { text: s }]),
        Value::Array(arr) => {
            serde_json::from_value(Value::Array(arr)).map_err(serde::de::Error::custom)
        }
        _ => Err(serde::de::Error::custom(
            "User content must be a string or an array of content parts",
        )),
    }
}

/// A message in a conversation with a language model.
///
/// Messages represent different roles in the conversation: system instructions, user inputs,
/// assistant responses, and tool results. Each message contains content appropriate to its role.
///
/// # Variants
///
/// * `System` - Provides instructions and context for how the model should behave
/// * `User` - Input from the user/human, typically containing text and/or files
/// * `Assistant` - Output from the language model, typically containing text and/or tool calls
/// * `Tool` - Results from executing tools that the model requested
///
/// # Serialization
///
/// Messages are serialized with a `role` field that identifies their type. User message
/// content is deserialized flexibly to accept either a string or an array of content parts.
///
/// # Usage in Prompts
///
/// Messages form the conversation history passed to the model. A typical conversation
/// might follow this pattern:
/// 1. System message (optional) - sets context
/// 2. User message - the human's question
/// 3. Assistant message - the model's response
/// 4. Repeat steps 2-3 as needed for multi-turn conversations
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "role", rename_all = "lowercase")]
pub enum Message {
    /// System message providing instructions and context for model behavior.
    ///
    /// System messages are typically placed at the beginning of a conversation
    /// to establish tone, guidelines, or special instructions for the model.
    System {
        /// Plain text instructions for the model
        content: String,
    },
    /// Message from the human user.
    ///
    /// User messages contain the actual input - questions, requests, or statements
    /// from the human. Content can be text, files, or a mix of both.
    User {
        /// Content parts, typically text and/or file references.
        /// Deserialized flexibly to accept string or array formats.
        #[serde(deserialize_with = "deserialize_user_content")]
        content: Vec<UserContentPart>,
    },
    /// Message from the assistant (language model).
    ///
    /// Assistant messages contain the model's responses, which can include text,
    /// reasoning, tool calls, or references to files.
    Assistant {
        /// Content parts generated by the model
        content: Vec<AssistantContentPart>,
    },
    /// Message containing tool execution results.
    ///
    /// After the model requests a tool call, tool results are provided in a Tool
    /// message so the model can see what happened and make follow-up decisions.
    Tool {
        /// Results from tool executions
        content: Vec<ToolResultPart>,
    },
}

/// Represents file content as either raw binary data or a URL reference.
///
/// This enum provides flexibility in how file content is transmitted. Binary data
/// is useful for small files or when you have the data in memory, while URLs are
/// more efficient for large files or remote resources.
///
/// # Variants
///
/// * `Binary` - Raw file bytes embedded in the message
/// * `Url` - URL pointing to the file location (HTTP/HTTPS or other schemes)
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum FileData {
    /// Raw file bytes embedded in the message
    Binary(Vec<u8>),
    /// URL pointing to the file location
    Url(String),
}

/// A content element within a user message.
///
/// User messages can contain multiple content parts, allowing for complex inputs
/// that mix text with various types of media like images or audio files.
///
/// # Variants
///
/// * `Text` - Plain text content
/// * `File` - Binary or URL-based file content with MIME type
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum UserContentPart {
    /// Plain text content from the user
    Text {
        /// The actual text content
        text: String,
    },
    /// File or media content from the user (images, audio, video, documents, etc.)
    File {
        /// The file data, either as binary bytes or as a URL
        data: FileData,
        /// MIME type identifying the file format (e.g., "image/jpeg", "audio/mp3")
        media_type: String,
    },
}

/// A content element within an assistant (model) message.
///
/// Assistant messages contain the model's response, which can include various types
/// of content: text answers, reasoning processes, generated files, and tool invocations.
///
/// # Variants
///
/// * `Text` - Plain text response from the model
/// * `Reasoning` - Internal reasoning from specialized models (e.g., o1)
/// * `File` - Generated files or media
/// * `ToolCall` - A request to invoke a tool or function
/// * `ToolResult` - Results from a tool execution
///
/// # Processing Assistant Messages
///
/// When receiving an assistant message, you should examine each content part to determine
/// what action to take. For example, if you encounter a ToolCall, you should execute the
/// tool and include the result in a subsequent message.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "kebab-case")]
pub enum AssistantContentPart {
    /// Text response from the model
    Text(TextPart),
    /// Reasoning process from the model (for reasoning models like o1)
    Reasoning(ReasoningPart),
    /// File or media generated by the model
    File(FilePart),
    /// Tool invocation request from the model
    ToolCall(ToolCallPart),
    /// Tool execution result provided by the model
    ToolResult(ToolResultPart),
}

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

    #[test]
    fn test_message_system() {
        let msg = Message::System {
            content: "You are helpful".into(),
        };
        let json = serde_json::to_value(&msg).unwrap();
        assert_eq!(json["role"], "system");
        assert_eq!(json["content"], "You are helpful");
    }

    #[test]
    fn test_message_user() {
        let msg = Message::User {
            content: vec![UserContentPart::Text {
                text: "Hello".into(),
            }],
        };
        let json = serde_json::to_value(&msg).unwrap();
        assert_eq!(json["role"], "user");
        assert_eq!(json["content"][0]["type"], "text");
    }

    #[test]
    fn test_user_content_file_binary() {
        let part = UserContentPart::File {
            data: FileData::Binary(vec![1, 2, 3]),
            media_type: "image/png".into(),
        };
        let json = serde_json::to_value(&part).unwrap();
        assert_eq!(json["type"], "file");
        assert_eq!(json["media_type"], "image/png");
    }

    #[test]
    fn test_user_content_file_url() {
        let part = UserContentPart::File {
            data: FileData::Url("https://example.com/image.jpg".into()),
            media_type: "image/jpeg".into(),
        };
        let json = serde_json::to_value(&part).unwrap();
        assert_eq!(json["type"], "file");
        assert_eq!(json["media_type"], "image/jpeg");
    }

    #[test]
    fn test_prompt_struct() {
        let mut prompt = Prompt::new();
        assert!(prompt.is_empty());

        prompt.push(Message::System {
            content: "test".into(),
        });
        assert_eq!(prompt.len(), 1);
    }

    #[test]
    fn test_prompt_from_string() {
        let prompt: Prompt = "Hello".into();
        assert_eq!(prompt.len(), 1);
        match &prompt[0] {
            Message::User { content } => match &content[0] {
                UserContentPart::Text { text } => assert_eq!(text, "Hello"),
                _ => panic!("Expected text content"),
            },
            _ => panic!("Expected user message"),
        }
    }

    #[test]
    fn test_prompt_from_vec() {
        let v = vec![Message::System {
            content: "s".into(),
        }];
        let prompt: Prompt = v.into();
        assert_eq!(prompt.len(), 1);
    }

    #[test]
    fn test_try_from_json_string() {
        let json = json!("Hello world");
        let prompt: Prompt = json.try_into().unwrap();
        assert_eq!(prompt.len(), 1);
        match &prompt[0] {
            Message::User { content } => match &content[0] {
                UserContentPart::Text { text } => assert_eq!(text, "Hello world"),
                _ => panic!("Expected text content"),
            },
            _ => panic!("Expected user message"),
        }
    }

    #[test]
    fn test_try_from_json_object() {
        let json = json!({
            "role": "system",
            "content": "You are helpful"
        });
        let prompt: Prompt = json.try_into().unwrap();
        assert_eq!(prompt.len(), 1);
        match &prompt[0] {
            Message::System { content } => assert_eq!(content, "You are helpful"),
            _ => panic!("Expected system message"),
        }
    }

    #[test]
    fn test_try_from_json_array() {
        let json = json!([
            {
                "role": "system",
                "content": "System"
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "User"
                    }
                ]
            }
        ]);
        let prompt: Prompt = json.try_into().unwrap();
        assert_eq!(prompt.len(), 2);
    }

    #[test]
    fn test_try_from_invalid_json() {
        let json = json!(123); // Invalid type
        let result: Result<Prompt, _> = json.try_into();
        assert!(result.is_err());
    }

    #[test]
    fn test_message_user_string_content() {
        let json = json!({
            "role": "user",
            "content": "Hello world"
        });
        let msg: Message = serde_json::from_value(json).unwrap();
        match msg {
            Message::User { content } => {
                assert_eq!(content.len(), 1);
                match &content[0] {
                    UserContentPart::Text { text } => assert_eq!(text, "Hello world"),
                    _ => panic!("Expected text content"),
                }
            }
            _ => panic!("Expected user message"),
        }
    }

    #[test]
    fn test_message_user_array_content() {
        let json = json!({
            "role": "user",
            "content": [
                { "type": "text", "text": "Hello" },
                { "type": "text", "text": "World" }
            ]
        });
        let msg: Message = serde_json::from_value(json).unwrap();
        match msg {
            Message::User { content } => {
                assert_eq!(content.len(), 2);
            }
            _ => panic!("Expected user message"),
        }
    }
}