muna 0.0.17

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

use serde::{Deserialize, Serialize};

use crate::types::Acceleration;

/// Reason that the model stopped generating.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StopReason {
    EndTurn,
    MaxTokens,
    StopSequence,
    ToolUse,
    PauseTurn,
    Refusal,
}

/// Billing usage.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Usage {
    /// Number of input tokens which were used, excluding cached tokens.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub input_tokens: Option<u64>,
    /// Number of output tokens which were used.
    pub output_tokens: u64,
    /// Number of input tokens used to create the cache entry.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cache_creation_input_tokens: Option<u64>,
    /// Number of input tokens read from the cache.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cache_read_input_tokens: Option<u64>,
}

/// Content block generated by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
    /// Text content block.
    Text {
        /// Response text.
        text: String,
    },
    /// Thinking content block.
    Thinking {
        /// Reasoning contents preceding the final answer.
        thinking: String,
        /// Signature verifying the thinking block. Always empty for Muna predictions.
        #[serde(default)]
        signature: String,
    },
}

/// Message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// Unique object identifier.
    pub id: String,
    /// Object type, always `message`.
    #[serde(rename = "type")]
    pub r#type: String,
    /// Conversational role of the generated message, always `assistant`.
    pub role: String,
    /// Content generated by the model.
    pub content: Vec<ContentBlock>,
    /// The model that completed the prompt.
    pub model: String,
    /// Reason that the model stopped generating.
    #[serde(default)]
    pub stop_reason: Option<StopReason>,
    /// Which custom stop sequence was generated, if any.
    #[serde(default)]
    pub stop_sequence: Option<String>,
    /// Billing and rate-limit usage.
    pub usage: Usage,
}

/// Content block delta.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlockDelta {
    /// Text content delta.
    TextDelta {
        /// Text fragment.
        text: String,
    },
    /// Thinking content delta.
    ThinkingDelta {
        /// Reasoning text fragment.
        thinking: String,
    },
}

/// Changes to the top-level message in a `message_delta` event.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MessageDelta {
    /// Reason that the model stopped generating.
    #[serde(default)]
    pub stop_reason: Option<StopReason>,
    /// Which custom stop sequence was generated, if any.
    #[serde(default)]
    pub stop_sequence: Option<String>,
}

/// Raw message stream event.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum RawMessageStreamEvent {
    /// Message start event.
    MessageStart {
        /// Message with empty content.
        message: Message,
    },
    /// Content block start event.
    ContentBlockStart {
        /// Content block index in the message.
        index: usize,
        /// Content block with empty contents.
        content_block: ContentBlock,
    },
    /// Content block delta event.
    ContentBlockDelta {
        /// Content block index in the message.
        index: usize,
        /// Content block delta.
        delta: ContentBlockDelta,
    },
    /// Content block stop event.
    ContentBlockStop {
        /// Content block index in the message.
        index: usize,
    },
    /// Message delta event.
    MessageDelta {
        /// Changes to the top-level message.
        delta: MessageDelta,
        /// Cumulative billing and rate-limit usage.
        usage: Usage,
    },
    /// Message stop event.
    MessageStop,
}

impl RawMessageStreamEvent {

    /// Event type name, as used for named server-sent events.
    pub fn event_type(&self) -> &'static str {
        match self {
            Self::MessageStart { .. } => "message_start",
            Self::ContentBlockStart { .. } => "content_block_start",
            Self::ContentBlockDelta { .. } => "content_block_delta",
            Self::ContentBlockStop { .. } => "content_block_stop",
            Self::MessageDelta { .. } => "message_delta",
            Self::MessageStop => "message_stop",
        }
    }
}

/// Input content block.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlockParam {
    /// Text content block.
    Text {
        /// Block text.
        text: String,
    },
}

/// Input message content, either a string or a list of content blocks.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessageContent {
    /// Plain text content.
    Text(String),
    /// Content block list.
    Blocks(Vec<ContentBlockParam>),
}

impl MessageContent {

    /// Flatten the content into a plain string.
    pub fn flatten(&self) -> String {
        match self {
            Self::Text(text) => text.clone(),
            Self::Blocks(blocks) => blocks
                .iter()
                .map(|block| match block {
                    ContentBlockParam::Text { text } => text.as_str(),
                })
                .collect(),
        }
    }
}

/// Input message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageParam {
    /// Message role.
    pub role: String,
    /// Message content.
    pub content: MessageContent,
}

/// Parameters for creating a message.
#[derive(Debug, Clone, Default)]
pub struct MessageCreateParams {
    /// Model predictor tag.
    pub model: String,
    /// The maximum number of tokens to generate before stopping.
    pub max_tokens: i32,
    /// Input messages.
    pub messages: Vec<MessageParam>,
    /// Custom text sequences that will cause the model to stop generating.
    /// Ignored unless the predictor natively supports it.
    pub stop_sequences: Option<Vec<String>>,
    /// System prompt.
    pub system: Option<MessageContent>,
    /// Amount of randomness injected into the response.
    pub temperature: Option<f32>,
    /// Only sample from the top K options for each subsequent token.
    /// Ignored unless the predictor natively supports it.
    pub top_k: Option<i32>,
    /// Nucleus sampling coefficient.
    pub top_p: Option<f32>,
    /// Prediction acceleration.
    pub acceleration: Option<Acceleration>,
}