jutella 0.8.1

// Copyright (c) 2024 Dmitry Markin
//
// SPDX-License-Identifier: MIT
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

//! OpenAI API Message types.

use serde::{
    de::Error as _,
    ser::{SerializeMap, SerializeSeq},
    Deserialize, Deserializer, Serialize, Serializer,
};
use serde_json::value::Value;

/// Conversation message.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Message {
    /// System message.
    System(SystemMessage),
    /// User message.
    User(UserMessage),
    /// Assistant message.
    Assistant(AssistantMessage),
    /// Tool message.
    Tool(ToolMessage),
}

/// System message.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SystemMessage {
    /// The contents of the message.
    pub content: String,
    /// An optional name for the participant. Provides the model information
    /// to differentiate between participants of the same role.
    pub name: Option<String>,
}

impl SystemMessage {
    pub fn new(content: String) -> Self {
        Self {
            content,
            name: None,
        }
    }
}

/// User message.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UserMessage {
    /// The contents of the message.
    pub content: Content,
    /// An optional name for the participant. Provides the model information
    /// to differentiate between participants of the same role.
    pub name: Option<String>,
}

impl UserMessage {
    pub fn new(content: Content) -> Self {
        Self {
            content,
            name: None,
        }
    }

    #[cfg(test)]
    pub fn new_from_str(content: &str) -> Self {
        Self {
            content: Content::Text(content.to_string()),
            name: None,
        }
    }
}

/// Assistant message.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AssistantMessage {
    /// The contents of the message.
    pub content: Content,
    /// An optional name for the participant. Provides the model information
    /// to differentiate between participants of the same role.
    pub name: Option<String>,
    /// The refusal message by the assistant.
    pub refusal: Option<String>,
    /// The tool calls generated by the model, such as function calls.
    pub tool_calls: Option<Value>,
}

impl AssistantMessage {
    pub fn new(content: Content) -> Self {
        Self {
            content,
            name: None,
            refusal: None,
            tool_calls: None,
        }
    }

    #[cfg(test)]
    pub fn new_from_str(content: &str) -> Self {
        Self {
            content: Content::Text(content.to_string()),
            name: None,
            refusal: None,
            tool_calls: None,
        }
    }
}

/// Tool message.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolMessage {
    /// The contents of the message.
    pub content: String,
    /// Tool call that this message is responding to.
    pub tool_call_id: String,
}

impl From<SystemMessage> for Message {
    fn from(message: SystemMessage) -> Self {
        Self::System(message)
    }
}

impl From<UserMessage> for Message {
    fn from(message: UserMessage) -> Self {
        Self::User(message)
    }
}

impl From<AssistantMessage> for Message {
    fn from(message: AssistantMessage) -> Self {
        Self::Assistant(message)
    }
}

impl From<ToolMessage> for Message {
    fn from(message: ToolMessage) -> Self {
        Self::Tool(message)
    }
}

/// The role of the message author.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    /// System message.
    System,
    /// User message.
    User,
    /// Assistant message.
    Assistant,
    /// Tool message.
    Tool,
}

/// Single-part generic message. Used to parse model responses.
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
pub struct ResponseGenericMessage {
    /// The role of the message author.
    role: Role,
    /// The contents of the message.
    content: Option<String>,
    /// An optional name for the participant. Provides the model information
    /// to differentiate between participants of the same role.
    name: Option<String>,
    /// The refusal message by the assistant.
    refusal: Option<String>,
    /// The tool calls generated by the model, such as function calls.
    tool_calls: Option<Value>,
    /// Tool call that this message is responding to.
    tool_call_id: Option<String>,

    // OpenRouter specific fields.
    /// Reasoning performaed by model.
    reasoning: Option<String>,
    /// Generated images.
    images: Option<Vec<ImagePart>>,
}

/// Multi-part generic message. Used to pass user messages.
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct RequestGenericMessage {
    /// The role of the message author.
    role: Role,
    /// The contents of the message.
    content: Content,
    /// An optional name for the participant. Provides the model information
    /// to differentiate between participants of the same role.
    #[serde(skip_serializing_if = "Option::is_none")]
    name: Option<String>,
    /// Tool call that this message is responding to.
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_call_id: Option<String>,
}

impl From<Message> for RequestGenericMessage {
    fn from(message: Message) -> Self {
        match message {
            Message::System(m) => m.into(),
            Message::User(m) => m.into(),
            Message::Assistant(m) => m.into(),
            Message::Tool(m) => m.into(),
        }
    }
}

impl From<SystemMessage> for RequestGenericMessage {
    fn from(SystemMessage { content, name }: SystemMessage) -> Self {
        Self {
            role: Role::System,
            content: Content::Text(content),
            name,
            tool_call_id: None,
        }
    }
}

impl From<UserMessage> for RequestGenericMessage {
    fn from(UserMessage { content, name }: UserMessage) -> Self {
        Self {
            role: Role::User,
            content,
            name,
            tool_call_id: None,
        }
    }
}

impl From<AssistantMessage> for RequestGenericMessage {
    fn from(
        AssistantMessage {
            content,
            name,
            refusal: _,
            tool_calls: _,
        }: AssistantMessage,
    ) -> Self {
        Self {
            role: Role::Assistant,
            content,
            name,
            tool_call_id: None,
        }
    }
}

impl From<ToolMessage> for RequestGenericMessage {
    fn from(
        ToolMessage {
            content,
            tool_call_id,
        }: ToolMessage,
    ) -> Self {
        Self {
            role: Role::Tool,
            content: Content::Text(content),
            name: None,
            tool_call_id: Some(tool_call_id),
        }
    }
}

/// Image part in a multi-part user message.
#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
pub struct ImagePart {
    /// Image URL or base64-encoded image data in the following format:
    /// `data:{mime_type};base64,{base64_image}`.
    pub url: String,
    /// Image detail level. Typically one of `auto`, `low`, `high`.
    pub detail: Option<String>,
}

/// File part in a multi-part user message.
#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
pub struct FilePart {
    /// File URL (only supported by OpenRouter) or base64-encoded file data in the following
    /// format: `data:application/pdf;base64,{base64_pdf}`.
    pub file_data: String,
    /// File name to pass to the model.
    pub filename: Option<String>,
}

/// Content part of a multi-part user message.
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum ContentPart {
    /// Text content part.
    Text(String),
    /// Image content part.
    Image(ImagePart),
    /// File content part.
    File(FilePart),
}

impl Serialize for ContentPart {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match self {
            ContentPart::Text(text) => {
                let mut map = serializer.serialize_map(Some(2))?;
                map.serialize_entry("type", "text")?;
                map.serialize_entry("text", text)?;
                map.end()
            }
            ContentPart::Image(image) => {
                let mut map = serializer.serialize_map(Some(2))?;
                map.serialize_entry("type", "image_url")?;
                map.serialize_entry("image_url", image)?;
                map.end()
            }
            ContentPart::File(file) => {
                let mut map = serializer.serialize_map(Some(2))?;
                map.serialize_entry("type", "file")?;
                map.serialize_entry("file", file)?;
                map.end()
            }
        }
    }
}

/// Content of the message. Either a string or array of content parts.
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum Content {
    /// Text message content.
    Text(String),
    /// Multipart message content.
    ContentParts(Vec<ContentPart>),
}

impl Serialize for Content {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match self {
            Content::Text(s) => serializer.serialize_str(s),
            Content::ContentParts(parts) => {
                let mut seq = serializer.serialize_seq(Some(parts.len()))?;
                for part in parts {
                    seq.serialize_element(part)?;
                }
                seq.end()
            }
        }
    }
}

/// Intermediate image content part. Used for `Deserialize` implementation only.
#[derive(Debug, Clone, Eq, PartialEq, serde_query::Deserialize)]
pub struct IntermediateImagePart {
    /// Type. Either `image_url` or `file`.
    #[query(".type")]
    ty: String,
    /// Image URL or base64 encoded data.
    #[query(".image_url.url")]
    url: String,
}

impl<'de> Deserialize<'de> for ImagePart {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let part = IntermediateImagePart::deserialize(deserializer)?;

        match part.ty.as_ref() {
            "image_url" => Ok(ImagePart {
                url: part.url,
                detail: None,
            }),
            ty => Err(D::Error::custom(format!("unsupported type `{}`", ty))),
        }
    }
}

/// Error when converting messages
#[derive(Debug, thiserror::Error)]
pub enum Error {
    /// Missing mandatory field
    #[error("Missing mandatory field `{0}`")]
    MissingField(&'static str),
    /// Invalid role
    #[error("Expected role {0:?}, got {1:?}")]
    RoleMismatch(Role, Role),
}

/// Assistant message that can hold text content only.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResponseAssistantMessage {
    /// The contents of the message.
    pub content: Option<String>,
    /// An optional name for the participant. Provides the model information
    /// to differentiate between participants of the same role.
    pub name: Option<String>,
    /// The refusal message by the assistant.
    pub refusal: Option<String>,
    /// The tool calls generated by the model, such as function calls.
    pub tool_calls: Option<Value>,
    /// Reasoning performaed by model.
    pub reasoning: Option<String>,
    /// Generated images.
    pub images: Option<Vec<ImagePart>>,
}

impl TryFrom<ResponseGenericMessage> for ResponseAssistantMessage {
    type Error = Error;

    fn try_from(
        ResponseGenericMessage {
            role,
            content,
            name,
            refusal,
            tool_calls,
            tool_call_id: _,
            reasoning,
            images,
        }: ResponseGenericMessage,
    ) -> Result<Self, Error> {
        if role == Role::Assistant {
            Ok(Self {
                content,
                name,
                refusal,
                tool_calls,
                reasoning,
                images,
            })
        } else {
            Err(Error::RoleMismatch(Role::Assistant, role))
        }
    }
}