rig-core 0.35.0

use std::{convert::Infallible, str::FromStr};

use crate::OneOrMany;
use serde::{Deserialize, Serialize};
use thiserror::Error;

use super::CompletionError;

// ================================================================
// Message models
// ================================================================

/// A useful trait to help convert `rig::completion::Message` to your own message type.
///
/// Particularly useful if you don't want to create a free-standing function as
/// when trying to use `TryFrom<T>`, you would normally run into the orphan rule as Vec is
/// technically considered a foreign type (it's owned by stdlib).
pub trait ConvertMessage: Sized + Send + Sync {
    type Error: std::error::Error + Send;

    fn convert_from_message(message: Message) -> Result<Vec<Self>, Self::Error>;
}

/// A message represents a run of input (user) and output (assistant).
/// Each message type (based on it's `role`) can contain a atleast one bit of content such as text,
///  images, audio, documents, or tool related information. While each message type can contain
///  multiple content, most often, you'll only see one content type per message
///  (an image w/ a description, etc).
///
/// Each provider is responsible with converting the generic message into it's provider specific
///  type using `From` or `TryFrom` traits. Since not every provider supports every feature, the
///  conversion can be lossy (providing an image might be discarded for a non-image supporting
///  provider) though the message being converted back and forth should always be the same.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(tag = "role", rename_all = "lowercase")]
pub enum Message {
    /// System message containing instruction text.
    System { content: String },

    /// User message containing one or more content types defined by `UserContent`.
    User { content: OneOrMany<UserContent> },

    /// Assistant message containing one or more content types defined by `AssistantContent`.
    Assistant {
        id: Option<String>,
        content: OneOrMany<AssistantContent>,
    },
}

/// Describes the content of a message, which can be text, a tool result, an image, audio, or
///  a document. Dependent on provider supporting the content type. Multimedia content is generally
///  base64 (defined by it's format) encoded but additionally supports urls (for some providers).
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum UserContent {
    Text(Text),
    ToolResult(ToolResult),
    Image(Image),
    Audio(Audio),
    Video(Video),
    Document(Document),
}

/// Describes responses from a provider which is either text or a tool call.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(untagged)]
pub enum AssistantContent {
    Text(Text),
    ToolCall(ToolCall),
    Reasoning(Reasoning),
    Image(Image),
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(tag = "type", content = "content", rename_all = "snake_case")]
#[non_exhaustive]
/// A typed reasoning block used by providers that emit structured thinking data.
pub enum ReasoningContent {
    /// Plain reasoning text with an optional provider signature.
    Text {
        text: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        signature: Option<String>,
    },
    /// Provider-encrypted reasoning payload.
    Encrypted(String),
    /// Redacted reasoning payload preserved as opaque data.
    Redacted { data: String },
    /// Provider-generated reasoning summary text.
    Summary(String),
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[non_exhaustive]
/// Assistant reasoning payload with an optional provider-supplied identifier.
pub struct Reasoning {
    /// Provider reasoning identifier, when supplied by the upstream API.
    pub id: Option<String>,
    /// Ordered reasoning content blocks.
    pub content: Vec<ReasoningContent>,
}

impl Reasoning {
    /// Create a new reasoning item from a single item
    pub fn new(input: &str) -> Self {
        Self::new_with_signature(input, None)
    }

    /// Create a new reasoning item from a single text item and optional signature.
    pub fn new_with_signature(input: &str, signature: Option<String>) -> Self {
        Self {
            id: None,
            content: vec![ReasoningContent::Text {
                text: input.to_string(),
                signature,
            }],
        }
    }

    /// Set or clear the provider reasoning ID.
    pub fn optional_id(mut self, id: Option<String>) -> Self {
        self.id = id;
        self
    }

    /// Set a provider reasoning ID.
    pub fn with_id(mut self, id: String) -> Self {
        self.id = Some(id);
        self
    }

    #[deprecated(note = "Use `new_with_signature` or content-level signature handling instead")]
    pub fn with_signature(mut self, signature: Option<String>) -> Self {
        if let Some(existing_signature) =
            self.content.iter_mut().find_map(|content| match content {
                ReasoningContent::Text { signature, .. } => Some(signature),
                _ => None,
            })
        {
            *existing_signature = signature;
        }
        self
    }

    /// Create reasoning content from multiple text blocks.
    pub fn multi(input: Vec<String>) -> Self {
        Self {
            id: None,
            content: input
                .into_iter()
                .map(|text| ReasoningContent::Text {
                    text,
                    signature: None,
                })
                .collect(),
        }
    }

    /// Create a redacted reasoning block.
    pub fn redacted(data: impl Into<String>) -> Self {
        Self {
            id: None,
            content: vec![ReasoningContent::Redacted { data: data.into() }],
        }
    }

    /// Create an encrypted reasoning block.
    pub fn encrypted(data: impl Into<String>) -> Self {
        Self {
            id: None,
            content: vec![ReasoningContent::Encrypted(data.into())],
        }
    }

    /// Create one reasoning block containing summary items.
    pub fn summaries(input: Vec<String>) -> Self {
        Self {
            id: None,
            content: input.into_iter().map(ReasoningContent::Summary).collect(),
        }
    }

    /// Render reasoning as displayable text by joining text-like blocks with newlines.
    pub fn display_text(&self) -> String {
        self.content
            .iter()
            .filter_map(|content| match content {
                ReasoningContent::Text { text, .. } => Some(text.as_str()),
                ReasoningContent::Summary(summary) => Some(summary.as_str()),
                ReasoningContent::Redacted { data } => Some(data.as_str()),
                ReasoningContent::Encrypted(_) => None,
            })
            .collect::<Vec<_>>()
            .join("\n")
    }

    /// Return the first text reasoning block, if present.
    pub fn first_text(&self) -> Option<&str> {
        self.content.iter().find_map(|content| match content {
            ReasoningContent::Text { text, .. } => Some(text.as_str()),
            _ => None,
        })
    }

    /// Return the first signature from text reasoning, if present.
    pub fn first_signature(&self) -> Option<&str> {
        self.content.iter().find_map(|content| match content {
            ReasoningContent::Text {
                signature: Some(signature),
                ..
            } => Some(signature.as_str()),
            _ => None,
        })
    }

    /// Return the first encrypted reasoning payload, if present.
    pub fn encrypted_content(&self) -> Option<&str> {
        self.content.iter().find_map(|content| match content {
            ReasoningContent::Encrypted(data) => Some(data.as_str()),
            _ => None,
        })
    }
}

/// Tool result content containing information about a tool call and it's resulting content.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct ToolResult {
    pub id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub call_id: Option<String>,
    pub content: OneOrMany<ToolResultContent>,
}

/// Describes the content of a tool result, which can be text or an image.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum ToolResultContent {
    Text(Text),
    Image(Image),
}

/// Describes a tool call with an id and function to call, generally produced by a provider.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct ToolCall {
    pub id: String,
    pub call_id: Option<String>,
    pub function: ToolFunction,
    /// Optional cryptographic signature for the tool call.
    ///
    /// This field is used by some providers (e.g., Google) to provide a signature
    /// that can verify the authenticity and integrity of the tool call. When present,
    /// it allows verification that the tool call was actually generated by the model
    /// and has not been tampered with.
    ///
    /// This is an optional, provider-specific feature and will be `None` for providers
    /// that don't support tool call signatures.
    pub signature: Option<String>,
    /// Additional provider-specific parameters to be sent to the completion model provider
    pub additional_params: Option<serde_json::Value>,
}

impl ToolCall {
    pub fn new(id: String, function: ToolFunction) -> Self {
        Self {
            id,
            call_id: None,
            function,
            signature: None,
            additional_params: None,
        }
    }

    pub fn with_call_id(mut self, call_id: String) -> Self {
        self.call_id = Some(call_id);
        self
    }

    pub fn with_signature(mut self, signature: Option<String>) -> Self {
        self.signature = signature;
        self
    }

    pub fn with_additional_params(mut self, additional_params: Option<serde_json::Value>) -> Self {
        self.additional_params = additional_params;
        self
    }
}

/// Describes a tool function to call with a name and arguments, generally produced by a provider.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct ToolFunction {
    pub name: String,
    pub arguments: serde_json::Value,
}

impl ToolFunction {
    pub fn new(name: String, arguments: serde_json::Value) -> Self {
        Self { name, arguments }
    }
}

// ================================================================
// Base content models
// ================================================================

/// Basic text content.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Text {
    pub text: String,
}

impl Text {
    pub fn text(&self) -> &str {
        &self.text
    }
}

impl std::fmt::Display for Text {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let Self { text } = self;
        write!(f, "{text}")
    }
}

/// Image content containing image data and metadata about it.
#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Image {
    pub data: DocumentSourceKind,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub media_type: Option<ImageMediaType>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub detail: Option<ImageDetail>,
    #[serde(flatten, skip_serializing_if = "Option::is_none")]
    pub additional_params: Option<serde_json::Value>,
}

impl Image {
    pub fn try_into_url(self) -> Result<String, MessageError> {
        match self.data {
            DocumentSourceKind::Url(url) => Ok(url),
            DocumentSourceKind::Base64(data) => {
                let Some(media_type) = self.media_type else {
                    return Err(MessageError::ConversionError(
                        "A media type is required to create a valid base64-encoded image URL"
                            .to_string(),
                    ));
                };

                Ok(format!(
                    "data:image/{ty};base64,{data}",
                    ty = media_type.to_mime_type()
                ))
            }
            unknown => Err(MessageError::ConversionError(format!(
                "Tried to convert unknown type to a URL: {unknown:?}"
            ))),
        }
    }
}

/// The kind of image source (to be used).
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Default)]
#[serde(tag = "type", content = "value", rename_all = "camelCase")]
#[non_exhaustive]
pub enum DocumentSourceKind {
    /// A file URL/URI.
    Url(String),
    /// A base-64 encoded string.
    Base64(String),
    /// Raw bytes
    Raw(Vec<u8>),
    /// A string (or a string literal).
    String(String),
    #[default]
    /// An unknown file source (there's nothing there).
    Unknown,
}

impl DocumentSourceKind {
    pub fn url(url: &str) -> Self {
        Self::Url(url.to_string())
    }

    pub fn base64(base64_string: &str) -> Self {
        Self::Base64(base64_string.to_string())
    }

    pub fn raw(bytes: impl Into<Vec<u8>>) -> Self {
        Self::Raw(bytes.into())
    }

    pub fn string(input: &str) -> Self {
        Self::String(input.into())
    }

    pub fn unknown() -> Self {
        Self::Unknown
    }

    pub fn try_into_inner(self) -> Option<String> {
        match self {
            Self::Url(s) | Self::Base64(s) => Some(s),
            _ => None,
        }
    }
}

impl std::fmt::Display for DocumentSourceKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Url(string) => write!(f, "{string}"),
            Self::Base64(string) => write!(f, "{string}"),
            Self::String(string) => write!(f, "{string}"),
            Self::Raw(_) => write!(f, "<binary data>"),
            Self::Unknown => write!(f, "<unknown>"),
        }
    }
}

/// Audio content containing audio data and metadata about it.
#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Audio {
    pub data: DocumentSourceKind,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub media_type: Option<AudioMediaType>,
    #[serde(flatten, skip_serializing_if = "Option::is_none")]
    pub additional_params: Option<serde_json::Value>,
}

/// Video content containing video data and metadata about it.
#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Video {
    pub data: DocumentSourceKind,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub media_type: Option<VideoMediaType>,
    #[serde(flatten, skip_serializing_if = "Option::is_none")]
    pub additional_params: Option<serde_json::Value>,
}

/// Document content containing document data and metadata about it.
#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Document {
    pub data: DocumentSourceKind,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub media_type: Option<DocumentMediaType>,
    #[serde(flatten, skip_serializing_if = "Option::is_none")]
    pub additional_params: Option<serde_json::Value>,
}

/// Describes the format of the content, which can be base64 or string.
#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ContentFormat {
    #[default]
    Base64,
    String,
    Url,
}

/// Helper enum that tracks the media type of the content.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub enum MediaType {
    Image(ImageMediaType),
    Audio(AudioMediaType),
    Document(DocumentMediaType),
    Video(VideoMediaType),
}

/// Describes the image media type of the content. Not every provider supports every media type.
/// Convertible to and from MIME type strings.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ImageMediaType {
    JPEG,
    PNG,
    GIF,
    WEBP,
    HEIC,
    HEIF,
    SVG,
}

/// Describes the document media type of the content. Not every provider supports every media type.
/// Includes also programming languages as document types for providers who support code running.
/// Convertible to and from MIME type strings.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum DocumentMediaType {
    PDF,
    TXT,
    RTF,
    HTML,
    CSS,
    MARKDOWN,
    CSV,
    XML,
    Javascript,
    Python,
}

impl DocumentMediaType {
    pub fn is_code(&self) -> bool {
        matches!(self, Self::Javascript | Self::Python)
    }
}

/// Describes the audio media type of the content. Not every provider supports every media type.
/// Convertible to and from MIME type strings.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum AudioMediaType {
    WAV,
    MP3,
    AIFF,
    AAC,
    OGG,
    FLAC,
    M4A,
    PCM16,
    PCM24,
}

/// Describes the video media type of the content. Not every provider supports every media type.
/// Convertible to and from MIME type strings.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum VideoMediaType {
    AVI,
    MP4,
    MPEG,
    MOV,
    WEBM,
}

/// Describes the detail of the image content, which can be low, high, or auto (open-ai specific).
#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ImageDetail {
    Low,
    High,
    #[default]
    Auto,
}

// ================================================================
// Impl. for message models
// ================================================================

impl Message {
    /// This helper method is primarily used to extract the first string prompt from a `Message`.
    /// Since `Message` might have more than just text content, we need to find the first text.
    pub(crate) fn rag_text(&self) -> Option<String> {
        match self {
            Message::User { content } => {
                for item in content.iter() {
                    if let UserContent::Text(Text { text }) = item {
                        return Some(text.clone());
                    }
                }
                None
            }
            Message::System { .. } => None,
            _ => None,
        }
    }

    /// Helper constructor to make creating system messages easier.
    pub fn system(text: impl Into<String>) -> Self {
        Message::System {
            content: text.into(),
        }
    }

    /// Helper constructor to make creating user messages easier.
    pub fn user(text: impl Into<String>) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::text(text)),
        }
    }

    /// Helper constructor to make creating assistant messages easier.
    pub fn assistant(text: impl Into<String>) -> Self {
        Message::Assistant {
            id: None,
            content: OneOrMany::one(AssistantContent::text(text)),
        }
    }

    /// Helper constructor to make creating assistant messages easier.
    pub fn assistant_with_id(id: String, text: impl Into<String>) -> Self {
        Message::Assistant {
            id: Some(id),
            content: OneOrMany::one(AssistantContent::text(text)),
        }
    }

    /// Helper constructor to make creating tool result messages easier.
    pub fn tool_result(id: impl Into<String>, content: impl Into<String>) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::ToolResult(ToolResult {
                id: id.into(),
                call_id: None,
                content: OneOrMany::one(ToolResultContent::text(content)),
            })),
        }
    }

    pub fn tool_result_with_call_id(
        id: impl Into<String>,
        call_id: Option<String>,
        content: impl Into<String>,
    ) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::ToolResult(ToolResult {
                id: id.into(),
                call_id,
                content: OneOrMany::one(ToolResultContent::text(content)),
            })),
        }
    }
}

impl UserContent {
    /// Helper constructor to make creating user text content easier.
    pub fn text(text: impl Into<String>) -> Self {
        UserContent::Text(text.into().into())
    }

    /// Helper constructor to make creating user image content easier.
    pub fn image_base64(
        data: impl Into<String>,
        media_type: Option<ImageMediaType>,
        detail: Option<ImageDetail>,
    ) -> Self {
        UserContent::Image(Image {
            data: DocumentSourceKind::Base64(data.into()),
            media_type,
            detail,
            additional_params: None,
        })
    }

    /// Helper constructor to make creating user image content from raw unencoded bytes easier.
    pub fn image_raw(
        data: impl Into<Vec<u8>>,
        media_type: Option<ImageMediaType>,
        detail: Option<ImageDetail>,
    ) -> Self {
        UserContent::Image(Image {
            data: DocumentSourceKind::Raw(data.into()),
            media_type,
            detail,
            ..Default::default()
        })
    }

    /// Helper constructor to make creating user image content easier.
    pub fn image_url(
        url: impl Into<String>,
        media_type: Option<ImageMediaType>,
        detail: Option<ImageDetail>,
    ) -> Self {
        UserContent::Image(Image {
            data: DocumentSourceKind::Url(url.into()),
            media_type,
            detail,
            additional_params: None,
        })
    }

    /// Helper constructor to make creating user audio content easier.
    pub fn audio(data: impl Into<String>, media_type: Option<AudioMediaType>) -> Self {
        UserContent::Audio(Audio {
            data: DocumentSourceKind::Base64(data.into()),
            media_type,
            additional_params: None,
        })
    }

    /// Helper constructor to make creating user audio content from raw unencoded bytes easier.
    pub fn audio_raw(data: impl Into<Vec<u8>>, media_type: Option<AudioMediaType>) -> Self {
        UserContent::Audio(Audio {
            data: DocumentSourceKind::Raw(data.into()),
            media_type,
            ..Default::default()
        })
    }

    /// Helper to create an audio resource from a URL
    pub fn audio_url(url: impl Into<String>, media_type: Option<AudioMediaType>) -> Self {
        UserContent::Audio(Audio {
            data: DocumentSourceKind::Url(url.into()),
            media_type,
            ..Default::default()
        })
    }

    /// Helper constructor to make creating user document content easier.
    /// This creates a document that assumes the data being passed in is a raw string.
    pub fn document(data: impl Into<String>, media_type: Option<DocumentMediaType>) -> Self {
        let data: String = data.into();
        UserContent::Document(Document {
            data: DocumentSourceKind::string(&data),
            media_type,
            additional_params: None,
        })
    }

    /// Helper to create a document from raw unencoded bytes
    pub fn document_raw(data: impl Into<Vec<u8>>, media_type: Option<DocumentMediaType>) -> Self {
        UserContent::Document(Document {
            data: DocumentSourceKind::Raw(data.into()),
            media_type,
            ..Default::default()
        })
    }

    /// Helper to create a document from a URL
    pub fn document_url(url: impl Into<String>, media_type: Option<DocumentMediaType>) -> Self {
        UserContent::Document(Document {
            data: DocumentSourceKind::Url(url.into()),
            media_type,
            ..Default::default()
        })
    }

    /// Helper constructor to make creating user tool result content easier.
    pub fn tool_result(id: impl Into<String>, content: OneOrMany<ToolResultContent>) -> Self {
        UserContent::ToolResult(ToolResult {
            id: id.into(),
            call_id: None,
            content,
        })
    }

    /// Helper constructor to make creating user tool result content easier.
    pub fn tool_result_with_call_id(
        id: impl Into<String>,
        call_id: String,
        content: OneOrMany<ToolResultContent>,
    ) -> Self {
        UserContent::ToolResult(ToolResult {
            id: id.into(),
            call_id: Some(call_id),
            content,
        })
    }
}

impl AssistantContent {
    /// Helper constructor to make creating assistant text content easier.
    pub fn text(text: impl Into<String>) -> Self {
        AssistantContent::Text(text.into().into())
    }

    /// Helper constructor to make creating assistant image content easier.
    pub fn image_base64(
        data: impl Into<String>,
        media_type: Option<ImageMediaType>,
        detail: Option<ImageDetail>,
    ) -> Self {
        AssistantContent::Image(Image {
            data: DocumentSourceKind::Base64(data.into()),
            media_type,
            detail,
            additional_params: None,
        })
    }

    /// Helper constructor to make creating assistant tool call content easier.
    pub fn tool_call(
        id: impl Into<String>,
        name: impl Into<String>,
        arguments: serde_json::Value,
    ) -> Self {
        AssistantContent::ToolCall(ToolCall::new(
            id.into(),
            ToolFunction {
                name: name.into(),
                arguments,
            },
        ))
    }

    pub fn tool_call_with_call_id(
        id: impl Into<String>,
        call_id: String,
        name: impl Into<String>,
        arguments: serde_json::Value,
    ) -> Self {
        AssistantContent::ToolCall(
            ToolCall::new(
                id.into(),
                ToolFunction {
                    name: name.into(),
                    arguments,
                },
            )
            .with_call_id(call_id),
        )
    }

    pub fn reasoning(reasoning: impl AsRef<str>) -> Self {
        AssistantContent::Reasoning(Reasoning::new(reasoning.as_ref()))
    }
}

impl ToolResultContent {
    /// Helper constructor to make creating tool result text content easier.
    pub fn text(text: impl Into<String>) -> Self {
        ToolResultContent::Text(text.into().into())
    }

    /// Helper constructor to make tool result images from a base64-encoded string.
    pub fn image_base64(
        data: impl Into<String>,
        media_type: Option<ImageMediaType>,
        detail: Option<ImageDetail>,
    ) -> Self {
        ToolResultContent::Image(Image {
            data: DocumentSourceKind::Base64(data.into()),
            media_type,
            detail,
            additional_params: None,
        })
    }

    /// Helper constructor to make tool result images from a base64-encoded string.
    pub fn image_raw(
        data: impl Into<Vec<u8>>,
        media_type: Option<ImageMediaType>,
        detail: Option<ImageDetail>,
    ) -> Self {
        ToolResultContent::Image(Image {
            data: DocumentSourceKind::Raw(data.into()),
            media_type,
            detail,
            ..Default::default()
        })
    }

    /// Helper constructor to make tool result images from a URL.
    pub fn image_url(
        url: impl Into<String>,
        media_type: Option<ImageMediaType>,
        detail: Option<ImageDetail>,
    ) -> Self {
        ToolResultContent::Image(Image {
            data: DocumentSourceKind::Url(url.into()),
            media_type,
            detail,
            additional_params: None,
        })
    }

    /// Parse a tool output string into appropriate ToolResultContent(s).
    ///
    /// Supports three formats:
    /// 1. Simple text: Any string → `OneOrMany::one(Text)`
    /// 2. Image JSON: `{"type": "image", "data": "...", "mimeType": "..."}` → `OneOrMany::one(Image)`
    /// 3. Hybrid JSON: `{"response": {...}, "parts": [...]}` → `OneOrMany::many([Text, Image, ...])`
    ///
    /// If JSON parsing fails, treats the entire string as text.
    pub fn from_tool_output(output: impl Into<String>) -> OneOrMany<ToolResultContent> {
        let output_str = output.into();

        if let Ok(json) = serde_json::from_str::<serde_json::Value>(&output_str) {
            if json.get("response").is_some() || json.get("parts").is_some() {
                let mut results: Vec<ToolResultContent> = Vec::new();

                if let Some(response) = json.get("response") {
                    results.push(ToolResultContent::Text(Text {
                        text: response.to_string(),
                    }));
                }

                if let Some(parts) = json.get("parts").and_then(|p| p.as_array()) {
                    for part in parts {
                        let is_image = part
                            .get("type")
                            .and_then(|t| t.as_str())
                            .is_some_and(|t| t == "image");

                        if !is_image {
                            continue;
                        }

                        if let (Some(data), Some(mime_type)) = (
                            part.get("data").and_then(|v| v.as_str()),
                            part.get("mimeType").and_then(|v| v.as_str()),
                        ) {
                            let data_kind =
                                if data.starts_with("http://") || data.starts_with("https://") {
                                    DocumentSourceKind::Url(data.to_string())
                                } else {
                                    DocumentSourceKind::Base64(data.to_string())
                                };

                            results.push(ToolResultContent::Image(Image {
                                data: data_kind,
                                media_type: ImageMediaType::from_mime_type(mime_type),
                                detail: None,
                                additional_params: None,
                            }));
                        }
                    }
                }

                if !results.is_empty() {
                    return OneOrMany::many(results).unwrap_or_else(|_| {
                        OneOrMany::one(ToolResultContent::Text(output_str.into()))
                    });
                }
            }

            let is_image = json
                .get("type")
                .and_then(|v| v.as_str())
                .is_some_and(|t| t == "image");

            if is_image
                && let (Some(data), Some(mime_type)) = (
                    json.get("data").and_then(|v| v.as_str()),
                    json.get("mimeType").and_then(|v| v.as_str()),
                )
            {
                let data_kind = if data.starts_with("http://") || data.starts_with("https://") {
                    DocumentSourceKind::Url(data.to_string())
                } else {
                    DocumentSourceKind::Base64(data.to_string())
                };

                return OneOrMany::one(ToolResultContent::Image(Image {
                    data: data_kind,
                    media_type: ImageMediaType::from_mime_type(mime_type),
                    detail: None,
                    additional_params: None,
                }));
            }
        }

        OneOrMany::one(ToolResultContent::Text(output_str.into()))
    }
}

/// Trait for converting between MIME types and media types.
pub trait MimeType {
    fn from_mime_type(mime_type: &str) -> Option<Self>
    where
        Self: Sized;
    fn to_mime_type(&self) -> &'static str;
}

impl MimeType for MediaType {
    fn from_mime_type(mime_type: &str) -> Option<Self> {
        ImageMediaType::from_mime_type(mime_type)
            .map(MediaType::Image)
            .or_else(|| {
                DocumentMediaType::from_mime_type(mime_type)
                    .map(MediaType::Document)
                    .or_else(|| {
                        AudioMediaType::from_mime_type(mime_type)
                            .map(MediaType::Audio)
                            .or_else(|| {
                                VideoMediaType::from_mime_type(mime_type).map(MediaType::Video)
                            })
                    })
            })
    }

    fn to_mime_type(&self) -> &'static str {
        match self {
            MediaType::Image(media_type) => media_type.to_mime_type(),
            MediaType::Audio(media_type) => media_type.to_mime_type(),
            MediaType::Document(media_type) => media_type.to_mime_type(),
            MediaType::Video(media_type) => media_type.to_mime_type(),
        }
    }
}

impl MimeType for ImageMediaType {
    fn from_mime_type(mime_type: &str) -> Option<Self> {
        match mime_type {
            "image/jpeg" => Some(ImageMediaType::JPEG),
            "image/png" => Some(ImageMediaType::PNG),
            "image/gif" => Some(ImageMediaType::GIF),
            "image/webp" => Some(ImageMediaType::WEBP),
            "image/heic" => Some(ImageMediaType::HEIC),
            "image/heif" => Some(ImageMediaType::HEIF),
            "image/svg+xml" => Some(ImageMediaType::SVG),
            _ => None,
        }
    }

    fn to_mime_type(&self) -> &'static str {
        match self {
            ImageMediaType::JPEG => "image/jpeg",
            ImageMediaType::PNG => "image/png",
            ImageMediaType::GIF => "image/gif",
            ImageMediaType::WEBP => "image/webp",
            ImageMediaType::HEIC => "image/heic",
            ImageMediaType::HEIF => "image/heif",
            ImageMediaType::SVG => "image/svg+xml",
        }
    }
}

impl MimeType for DocumentMediaType {
    fn from_mime_type(mime_type: &str) -> Option<Self> {
        match mime_type {
            "application/pdf" => Some(DocumentMediaType::PDF),
            "text/plain" => Some(DocumentMediaType::TXT),
            "text/rtf" => Some(DocumentMediaType::RTF),
            "text/html" => Some(DocumentMediaType::HTML),
            "text/css" => Some(DocumentMediaType::CSS),
            "text/md" | "text/markdown" => Some(DocumentMediaType::MARKDOWN),
            "text/csv" => Some(DocumentMediaType::CSV),
            "text/xml" => Some(DocumentMediaType::XML),
            "application/x-javascript" | "text/x-javascript" => Some(DocumentMediaType::Javascript),
            "application/x-python" | "text/x-python" => Some(DocumentMediaType::Python),
            _ => None,
        }
    }

    fn to_mime_type(&self) -> &'static str {
        match self {
            DocumentMediaType::PDF => "application/pdf",
            DocumentMediaType::TXT => "text/plain",
            DocumentMediaType::RTF => "text/rtf",
            DocumentMediaType::HTML => "text/html",
            DocumentMediaType::CSS => "text/css",
            DocumentMediaType::MARKDOWN => "text/markdown",
            DocumentMediaType::CSV => "text/csv",
            DocumentMediaType::XML => "text/xml",
            DocumentMediaType::Javascript => "application/x-javascript",
            DocumentMediaType::Python => "application/x-python",
        }
    }
}

impl MimeType for AudioMediaType {
    fn from_mime_type(mime_type: &str) -> Option<Self> {
        match mime_type {
            "audio/wav" => Some(AudioMediaType::WAV),
            "audio/mp3" => Some(AudioMediaType::MP3),
            "audio/aiff" => Some(AudioMediaType::AIFF),
            "audio/aac" => Some(AudioMediaType::AAC),
            "audio/ogg" => Some(AudioMediaType::OGG),
            "audio/flac" => Some(AudioMediaType::FLAC),
            "audio/m4a" => Some(AudioMediaType::M4A),
            "audio/pcm16" => Some(AudioMediaType::PCM16),
            "audio/pcm24" => Some(AudioMediaType::PCM24),
            _ => None,
        }
    }

    fn to_mime_type(&self) -> &'static str {
        match self {
            AudioMediaType::WAV => "audio/wav",
            AudioMediaType::MP3 => "audio/mp3",
            AudioMediaType::AIFF => "audio/aiff",
            AudioMediaType::AAC => "audio/aac",
            AudioMediaType::OGG => "audio/ogg",
            AudioMediaType::FLAC => "audio/flac",
            AudioMediaType::M4A => "audio/m4a",
            AudioMediaType::PCM16 => "audio/pcm16",
            AudioMediaType::PCM24 => "audio/pcm24",
        }
    }
}

impl MimeType for VideoMediaType {
    fn from_mime_type(mime_type: &str) -> Option<Self>
    where
        Self: Sized,
    {
        match mime_type {
            "video/avi" => Some(VideoMediaType::AVI),
            "video/mp4" => Some(VideoMediaType::MP4),
            "video/mpeg" => Some(VideoMediaType::MPEG),
            "video/mov" => Some(VideoMediaType::MOV),
            "video/webm" => Some(VideoMediaType::WEBM),
            &_ => None,
        }
    }

    fn to_mime_type(&self) -> &'static str {
        match self {
            VideoMediaType::AVI => "video/avi",
            VideoMediaType::MP4 => "video/mp4",
            VideoMediaType::MPEG => "video/mpeg",
            VideoMediaType::MOV => "video/mov",
            VideoMediaType::WEBM => "video/webm",
        }
    }
}

impl std::str::FromStr for ImageDetail {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "low" => Ok(ImageDetail::Low),
            "high" => Ok(ImageDetail::High),
            "auto" => Ok(ImageDetail::Auto),
            _ => Err(()),
        }
    }
}

// ================================================================
// FromStr, From<String>, and From<&str> impls
// ================================================================

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

impl From<&String> for Text {
    fn from(text: &String) -> Self {
        text.to_owned().into()
    }
}

impl From<&str> for Text {
    fn from(text: &str) -> Self {
        text.to_owned().into()
    }
}

impl FromStr for Text {
    type Err = Infallible;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(s.into())
    }
}

impl From<&Message> for Message {
    fn from(msg: &Message) -> Self {
        msg.clone()
    }
}

impl From<String> for Message {
    fn from(text: String) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::Text(text.into())),
        }
    }
}

impl From<&str> for Message {
    fn from(text: &str) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::Text(text.into())),
        }
    }
}

impl From<&String> for Message {
    fn from(text: &String) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::Text(text.into())),
        }
    }
}

impl From<Text> for Message {
    fn from(text: Text) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::Text(text)),
        }
    }
}

impl From<Image> for Message {
    fn from(image: Image) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::Image(image)),
        }
    }
}

impl From<Audio> for Message {
    fn from(audio: Audio) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::Audio(audio)),
        }
    }
}

impl From<Document> for Message {
    fn from(document: Document) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::Document(document)),
        }
    }
}

impl From<String> for ToolResultContent {
    fn from(text: String) -> Self {
        ToolResultContent::text(text)
    }
}

impl From<String> for AssistantContent {
    fn from(text: String) -> Self {
        AssistantContent::text(text)
    }
}

impl From<String> for UserContent {
    fn from(text: String) -> Self {
        UserContent::text(text)
    }
}

impl From<AssistantContent> for Message {
    fn from(content: AssistantContent) -> Self {
        Message::Assistant {
            id: None,
            content: OneOrMany::one(content),
        }
    }
}

impl From<UserContent> for Message {
    fn from(content: UserContent) -> Self {
        Message::User {
            content: OneOrMany::one(content),
        }
    }
}

impl From<OneOrMany<AssistantContent>> for Message {
    fn from(content: OneOrMany<AssistantContent>) -> Self {
        Message::Assistant { id: None, content }
    }
}

impl From<OneOrMany<UserContent>> for Message {
    fn from(content: OneOrMany<UserContent>) -> Self {
        Message::User { content }
    }
}

impl From<ToolCall> for Message {
    fn from(tool_call: ToolCall) -> Self {
        Message::Assistant {
            id: None,
            content: OneOrMany::one(AssistantContent::ToolCall(tool_call)),
        }
    }
}

impl From<ToolResult> for Message {
    fn from(tool_result: ToolResult) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::ToolResult(tool_result)),
        }
    }
}

impl From<ToolResultContent> for Message {
    fn from(tool_result_content: ToolResultContent) -> Self {
        Message::User {
            content: OneOrMany::one(UserContent::ToolResult(ToolResult {
                id: String::new(),
                call_id: None,
                content: OneOrMany::one(tool_result_content),
            })),
        }
    }
}

#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum ToolChoice {
    #[default]
    Auto,
    None,
    Required,
    Specific {
        function_names: Vec<String>,
    },
}

// ================================================================
// Error types
// ================================================================

/// Error type to represent issues with converting messages to and from specific provider messages.
#[derive(Debug, Error)]
pub enum MessageError {
    #[error("Message conversion error: {0}")]
    ConversionError(String),
}

impl From<MessageError> for CompletionError {
    fn from(error: MessageError) -> Self {
        CompletionError::RequestError(error.into())
    }
}

#[cfg(test)]
mod tests {
    use super::{Message, Reasoning, ReasoningContent};

    #[test]
    fn reasoning_constructors_and_accessors_work() {
        let single = Reasoning::new("think");
        assert_eq!(single.first_text(), Some("think"));
        assert_eq!(single.first_signature(), None);

        let signed = Reasoning::new_with_signature("signed", Some("sig-1".to_string()));
        assert_eq!(signed.first_text(), Some("signed"));
        assert_eq!(signed.first_signature(), Some("sig-1"));

        let multi = Reasoning::multi(vec!["a".to_string(), "b".to_string()]);
        assert_eq!(multi.display_text(), "a\nb");
        assert_eq!(multi.first_text(), Some("a"));

        let redacted = Reasoning::redacted("redacted-value");
        assert_eq!(redacted.display_text(), "redacted-value");
        assert_eq!(redacted.first_text(), None);

        let encrypted = Reasoning::encrypted("enc");
        assert_eq!(encrypted.encrypted_content(), Some("enc"));
        assert_eq!(encrypted.display_text(), "");

        let summaries = Reasoning::summaries(vec!["s1".to_string(), "s2".to_string()]);
        assert_eq!(summaries.display_text(), "s1\ns2");
        assert_eq!(summaries.encrypted_content(), None);
    }

    #[test]
    fn reasoning_content_serde_roundtrip() {
        let variants = vec![
            ReasoningContent::Text {
                text: "plain".to_string(),
                signature: Some("sig".to_string()),
            },
            ReasoningContent::Encrypted("opaque".to_string()),
            ReasoningContent::Redacted {
                data: "redacted".to_string(),
            },
            ReasoningContent::Summary("summary".to_string()),
        ];

        for variant in variants {
            let json = serde_json::to_string(&variant).expect("serialize");
            let roundtrip: ReasoningContent = serde_json::from_str(&json).expect("deserialize");
            assert_eq!(roundtrip, variant);
        }
    }

    #[test]
    fn system_message_constructor_and_serde_roundtrip() {
        let message = Message::system("You are concise.");

        match &message {
            Message::System { content } => assert_eq!(content, "You are concise."),
            _ => panic!("Expected system message"),
        }

        let json = serde_json::to_string(&message).expect("serialize");
        let roundtrip: Message = serde_json::from_str(&json).expect("deserialize");
        assert_eq!(roundtrip, message);
    }
}