muna 0.0.18

/*
*   Muna
*   Copyright © 2026 NatML Inc. All Rights Reserved.
*/

use serde::{Deserialize, Serialize};

use crate::types::{Acceleration, Image};

/// Chat message content: a plain string or a list of content parts.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ChatCompletionContent {
    /// Plain text content.
    Text(String),
    /// Content part list.
    Parts(Vec<ChatCompletionContentPart>),
}

impl ChatCompletionContent {

    /// Flatten textual content into a plain string, joining parts with a
    /// newline. Media parts contribute nothing; callers that support media
    /// must translate parts before flattening.
    pub fn flatten(&self) -> String {
        match self {
            Self::Text(text) => text.clone(),
            Self::Parts(parts) => parts
                .iter()
                .filter_map(|part| match part {
                    ChatCompletionContentPart::Text { text } => Some(text.as_str()),
                    ChatCompletionContentPart::Refusal { refusal } => Some(refusal.as_str()),
                    _ => None,
                })
                .collect::<Vec<_>>()
                .join("\n"),
        }
    }

    /// Whether the content is textual only (plain string, or parts that
    /// all flatten to text).
    pub fn is_text(&self) -> bool {
        match self {
            Self::Text(_) => true,
            Self::Parts(parts) => parts.iter().all(|part| matches!(
                part,
                ChatCompletionContentPart::Text { .. } |
                ChatCompletionContentPart::Refusal { .. }
            )),
        }
    }
}

/// OpenAI content part. All five official types parse; support is
/// per-model, gated on the predictor signature.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ChatCompletionContentPart {
    /// Text content part.
    Text {
        /// The text content.
        text: String,
    },
    /// Image content part.
    ImageUrl {
        /// Image URL payload.
        image_url: ChatCompletionContentPartImageUrl,
    },
    /// Audio content part.
    InputAudio {
        /// Audio payload.
        input_audio: ChatCompletionContentPartInputAudio,
    },
    /// File content part.
    File {
        /// File payload.
        file: ChatCompletionContentPartFile,
    },
    /// Assistant refusal part; appears when clients replay history
    /// containing refusals. Flattens as text.
    Refusal {
        /// The refusal message generated by the model.
        refusal: String,
    },
}

/// Image URL payload for an image content part.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionContentPartImageUrl {
    /// Either a URL of the image or a base64 data URL.
    pub url: String,
    /// Detail level of the image.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub detail: Option<ChatCompletionImageDetail>,
}

/// Detail level of an image content part.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ChatCompletionImageDetail {
    Auto,
    Low,
    High,
}

/// Audio payload for an audio content part.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionContentPartInputAudio {
    /// Base64 encoded audio data.
    pub data: String,
    /// The format of the encoded audio data.
    pub format: ChatCompletionInputAudioFormat,
}

/// Encoded input audio format.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ChatCompletionInputAudioFormat {
    Wav,
    Mp3,
}

impl ChatCompletionInputAudioFormat {

    pub fn as_str(self) -> &'static str {
        match self {
            Self::Wav => "wav",
            Self::Mp3 => "mp3",
        }
    }
}

/// File payload for a file content part.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ChatCompletionContentPartFile {
    /// The base64 encoded file data.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub file_data: Option<String>,
    /// The ID of an uploaded file to use as input.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub file_id: Option<String>,
    /// The name of the file.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub filename: Option<String>,
}

/// Definition of a function the model may call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionDefinition {
    /// The name of the function to be called.
    pub name: String,
    /// A description of what the function does, used by the model to
    /// choose when and how to call the function.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    /// The parameters the function accepts, described as a JSON Schema object.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub parameters: Option<serde_json::Value>,
    /// Whether to enable strict schema adherence when generating the
    /// function call.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub strict: Option<bool>,
}

/// A function tool the model may call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionFunctionTool {
    /// Tool type, always `function`.
    pub r#type: String,
    /// The function definition.
    pub function: FunctionDefinition,
}

/// A call to a function tool created by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionMessageFunctionToolCall {
    /// The ID of the tool call.
    pub id: String,
    /// Tool type, always `function`.
    pub r#type: String,
    /// The function the model called.
    pub function: ChatCompletionToolCallFunction,
}

/// Function name and arguments on a completed tool call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionToolCallFunction {
    /// The name of the function to call.
    pub name: String,
    /// The arguments to call the function with, as a JSON-encoded string.
    pub arguments: String,
}

/// Streamed tool call fragment, accumulated by `index`: the first
/// fragment carries `id` and the function name; subsequent fragments
/// append to the JSON-encoded `arguments` string.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChoiceDeltaToolCall {
    /// Index of the tool call in the message's tool call list.
    pub index: usize,
    /// The ID of the tool call. Present on the first fragment only.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    /// Tool type, always `function`. Present on the first fragment only.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub r#type: Option<String>,
    /// Function name and argument fragments.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub function: Option<ChoiceDeltaToolCallFunction>,
}

/// Function fragments on a streamed tool call.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ChoiceDeltaToolCallFunction {
    /// The name of the function to call. Present on the first fragment only.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    /// Fragment of the JSON-encoded arguments string.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub arguments: Option<String>,
}

/// Tool choice mode. `required` and named-function forcing need
/// constrained decoding and are not yet supported.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ChatCompletionToolChoice {
    Auto,
    None,
}

/// Chat message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionMessage {
    /// Message role.
    pub role: String,
    /// Message content.
    #[serde(default)]
    pub content: Option<ChatCompletionContent>,
    /// Reasoning contents of the message, before the final answer.
    /// Follows the DeepSeek convention for reasoning models.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reasoning_content: Option<String>,
    /// Tool calls generated by the model.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ChatCompletionMessageFunctionToolCall>>,
    /// Tool call that this message is responding to.
    /// Only present on `tool` role messages.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

/// Breakdown of tokens used in the prompt.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct PromptTokensDetails {
    /// Audio input tokens present in the prompt.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub audio_tokens: Option<u64>,
    /// The unadjusted number of prompt tokens written to cache.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cache_write_tokens: Option<u64>,
    /// Cached tokens present in the prompt.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cached_tokens: Option<u64>,
}

/// Breakdown of tokens used in the completion.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CompletionTokensDetails {
    /// Tokens generated by the model for reasoning.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reasoning_tokens: Option<u64>,
}

/// Usage information for a chat completion request.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ChatCompletionUsage {
    /// Number of tokens in the prompt.
    pub prompt_tokens: u64,
    /// Number of tokens in the generated completion.
    pub completion_tokens: u64,
    /// Total number of tokens used in the request.
    pub total_tokens: u64,
    /// Breakdown of tokens used in the prompt.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub prompt_tokens_details: Option<PromptTokensDetails>,
    /// Breakdown of tokens used in the completion.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub completion_tokens_details: Option<CompletionTokensDetails>,
}

/// Chat completion choice.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionChoice {
    /// Index of the choice in the list of choices.
    pub index: usize,
    /// Chat completion message generated by the model.
    pub message: ChatCompletionMessage,
    /// Reason the model stopped generating tokens.
    #[serde(default)]
    pub finish_reason: Option<String>,
    /// Log probability information for the choice.
    #[serde(default)]
    pub logprobs: Option<serde_json::Value>,
}

/// Chat completion response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletion {
    /// Object type, always `chat.completion`.
    pub object: String,
    /// Unique identifier for the chat completion.
    pub id: String,
    /// Model used for the chat completion.
    pub model: String,
    /// Generated chat completion choices.
    pub choices: Vec<ChatCompletionChoice>,
    /// Unix timestamp, in seconds, when the completion was created.
    pub created: u64,
    /// Usage statistics for the completion request.
    #[serde(default)]
    pub usage: Option<ChatCompletionUsage>,
}

/// Chat completion chunk delta.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ChatCompletionDelta {
    /// Role of the author of this message delta.
    #[serde(default)]
    pub role: Option<String>,
    /// Content of the message delta.
    #[serde(default)]
    pub content: Option<String>,
    /// Reasoning contents of the message delta, before the final answer.
    /// Follows the DeepSeek convention for reasoning models.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reasoning_content: Option<String>,
    /// Streamed tool call fragments, accumulated by `index`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ChoiceDeltaToolCall>>,
}

/// Chat completion chunk choice.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionChunkChoice {
    /// Index of the choice in the list of choices.
    pub index: usize,
    /// Chat completion delta generated by the model.
    #[serde(default)]
    pub delta: Option<ChatCompletionDelta>,
    /// Reason the model stopped generating tokens.
    #[serde(default)]
    pub finish_reason: Option<String>,
    /// Log probability information for the choice.
    #[serde(default)]
    pub logprobs: Option<serde_json::Value>,
}

/// Chat completion chunk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionChunk {
    /// Object type, always `chat.completion.chunk`.
    pub object: String,
    /// Unique identifier for the chat completion. Each chunk has the same ID.
    pub id: String,
    /// Model used for the chat completion.
    pub model: String,
    /// Generated chat completion chunk choices.
    pub choices: Vec<ChatCompletionChunkChoice>,
    /// Unix timestamp, in seconds, when the chunk was created.
    pub created: u64,
    /// Usage statistics for the completion request.
    #[serde(default)]
    pub usage: Option<ChatCompletionUsage>,
}

/// Reasoning effort for reasoning models.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ChatCompletionReasoningEffort {
    #[serde(rename = "minimal")]
    Minimal,
    #[serde(rename = "low")]
    Low,
    #[serde(rename = "medium")]
    Medium,
    #[serde(rename = "high")]
    High,
    #[serde(rename = "xhigh")]
    XHigh,
}

impl ChatCompletionReasoningEffort {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Minimal => "minimal",
            Self::Low => "low",
            Self::Medium => "medium",
            Self::High => "high",
            Self::XHigh => "xhigh",
        }
    }
}

/// Parameters for creating a chat completion.
#[derive(Debug, Clone, Default)]
pub struct ChatCompletionCreateParams {
    /// Chat predictor tag.
    pub model: String,
    /// Messages comprising the conversation so far.
    pub messages: Vec<ChatCompletionMessage>,
    /// Tools the model may call.
    pub tools: Option<Vec<ChatCompletionFunctionTool>>,
    /// Tool choice mode. Defaults to `auto`.
    pub tool_choice: Option<ChatCompletionToolChoice>,
    /// Response format.
    pub response_format: Option<serde_json::Map<String, serde_json::Value>>,
    /// Reasoning effort for reasoning models.
    pub reasoning_effort: Option<ChatCompletionReasoningEffort>,
    /// Maximum completion tokens.
    pub max_completion_tokens: Option<i32>,
    /// Sampling temperature to use.
    pub temperature: Option<f32>,
    /// Nucleus sampling coefficient.
    pub top_p: Option<f32>,
    /// Token frequency penalty.
    pub frequency_penalty: Option<f32>,
    /// Token presence penalty.
    pub presence_penalty: Option<f32>,
    /// Prediction acceleration.
    pub acceleration: Option<Acceleration>,
}

/// Embedding data, either as float values or base64-encoded bytes.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum EmbeddingData {
    Float(Vec<f32>),
    Base64(String),
}

/// Embedding vector.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Embedding {
    /// Object type, always `embedding`.
    pub object: String,
    /// Embedding vector as float values or a base64-encoded string.
    pub embedding: EmbeddingData,
    /// Index of the embedding in the response list.
    pub index: usize,
}

/// Usage information for an embedding request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingUsage {
    /// Number of tokens in the input prompt.
    pub prompt_tokens: u64,
    /// Total number of tokens used in the request.
    pub total_tokens: u64,
}

/// Response from creating embeddings.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingCreateResponse {
    pub object: String,
    pub model: String,
    pub data: Vec<Embedding>,
    pub usage: EmbeddingUsage,
}

/// Generated image.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageData {
    /// Base64-encoded image data.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub b64_json: Option<String>,
    /// Raw image. Only populated for the `raw` output format.
    #[serde(skip)]
    pub image: Option<Image>,
}

/// Token usage for an image generation request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageUsage {
    /// Number of tokens (images and text) in the input prompt.
    pub input_tokens: u64,
    /// Number of output tokens generated by the model.
    pub output_tokens: u64,
    /// Total number of tokens (images and text) used for the image generation.
    pub total_tokens: u64,
}

/// Image generation response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageResponse {
    /// The list of generated images.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub data: Option<Vec<ImageData>>,
    /// The background parameter used for the image generation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub background: Option<String>,
    /// Unix timestamp, in seconds, when the image was created.
    pub created: i64,
    /// Token usage information for the image generation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub usage: Option<ImageUsage>,
}