grok_api 0.1.6

//! Data models for Grok API requests and responses

use serde::{Deserialize, Serialize};
use serde_json::Value;

/// Represents an image URL
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ImageUrl {
    /// The URL of the image
    pub url: String,
    /// The detail level of the image (low, high, auto)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub detail: Option<String>,
}

/// Represents a video URL
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct VideoUrl {
    /// The URL of the video
    pub url: String,
    /// The detail level (optional)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub detail: Option<String>,
}

/// A part of a message content (text, image, or video)
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentPart {
    /// Text content
    Text { text: String },
    /// Image content
    ImageUrl { image_url: ImageUrl },
    /// Video content
    VideoUrl { video_url: VideoUrl },
}

/// The content of a message, which can be a simple string or a list of parts
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum MessageContent {
    /// Simple text content
    Text(String),
    /// Complex content with multiple parts
    Parts(Vec<ContentPart>),
}

impl MessageContent {
    /// Get the content as a text string if possible
    pub fn as_text(&self) -> Option<&str> {
        match self {
            MessageContent::Text(s) => Some(s),
            MessageContent::Parts(parts) => {
                // Return the first text part, or None
                for part in parts {
                    if let ContentPart::Text { text } = part {
                        return Some(text);
                    }
                }
                None
            }
        }
    }

    /// Get the text content, returning empty string if none
    pub fn text(&self) -> &str {
        match self {
            MessageContent::Text(s) => s,
            MessageContent::Parts(parts) => {
                // Return the first text part, or empty string
                for part in parts {
                    if let ContentPart::Text { text } = part {
                        return text;
                    }
                }
                ""
            }
        }
    }
}

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

impl From<&str> for MessageContent {
    fn from(s: &str) -> Self {
        MessageContent::Text(s.to_string())
    }
}

impl std::fmt::Display for MessageContent {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MessageContent::Text(s) => write!(f, "{}", s),
            MessageContent::Parts(parts) => {
                for part in parts {
                    match part {
                        ContentPart::Text { text } => write!(f, "{}", text)?,
                        ContentPart::ImageUrl { image_url } => {
                            write!(f, "[Image: {}]", image_url.url)?
                        }
                        ContentPart::VideoUrl { video_url } => {
                            write!(f, "[Video: {}]", video_url.url)?
                        }
                    }
                }
                Ok(())
            }
        }
    }
}

/// Represents a chat message in a conversation
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ChatMessage {
    /// The role of the message sender
    pub role: String,

    /// The content of the message
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<MessageContent>,

    /// Tool calls made by the assistant (if any)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,

    /// Tool call ID (required for role: "tool")
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

impl ChatMessage {
    /// Create a system message
    pub fn system<S: Into<MessageContent>>(content: S) -> Self {
        Self {
            role: "system".to_string(),
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: None,
        }
    }

    /// Create a user message
    pub fn user<S: Into<MessageContent>>(content: S) -> Self {
        Self {
            role: "user".to_string(),
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: None,
        }
    }

    /// Create a user message with multiple parts
    pub fn user_parts(parts: Vec<ContentPart>) -> Self {
        Self {
            role: "user".to_string(),
            content: Some(MessageContent::Parts(parts)),
            tool_calls: None,
            tool_call_id: None,
        }
    }

    /// Create an assistant message
    pub fn assistant<S: Into<MessageContent>>(content: S) -> Self {
        Self {
            role: "assistant".to_string(),
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: None,
        }
    }

    /// Create an assistant message with tool calls
    pub fn assistant_with_tools<S: Into<MessageContent>>(
        content: Option<S>,
        tool_calls: Vec<ToolCall>,
    ) -> Self {
        Self {
            role: "assistant".to_string(),
            content: content.map(|c| c.into()),
            tool_calls: Some(tool_calls),
            tool_call_id: None,
        }
    }

    /// Create a tool result message
    pub fn tool(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
        Self {
            role: "tool".to_string(),
            content: Some(MessageContent::Text(content.into())),
            tool_calls: None,
            tool_call_id: Some(tool_call_id.into()),
        }
    }
}

/// A chat completion request
#[derive(Debug, Clone, Serialize)]
pub struct ChatRequest {
    /// The model to use for completion
    pub model: String,

    /// The conversation messages
    pub messages: Vec<ChatMessage>,

    /// Sampling temperature (0.0 to 2.0)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,

    /// Maximum number of tokens to generate
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,

    /// Whether to stream the response
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,

    /// Tools available for the model to call
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<Value>>,

    /// Top-p sampling parameter
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f32>,

    /// Frequency penalty (-2.0 to 2.0).
    ///
    /// **⚠ Not supported by reasoning models** (Grok 4, Grok 4.20 variants).
    /// Sending this field to a reasoning model will result in an API error.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f32>,

    /// Presence penalty (-2.0 to 2.0).
    ///
    /// **⚠ Not supported by reasoning models** (Grok 4, Grok 4.20 variants).
    /// Sending this field to a reasoning model will result in an API error.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f32>,
}

/// Response from a chat completion request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatResponse {
    /// Unique identifier for the response
    pub id: String,

    /// Object type (always "chat.completion")
    pub object: String,

    /// Unix timestamp of when the response was created
    pub created: u64,

    /// The model used for completion
    pub model: String,

    /// The completion choices
    pub choices: Vec<Choice>,

    /// Token usage information
    pub usage: Usage,
}

impl ChatResponse {
    /// Get the primary response content (from first choice)
    pub fn content(&self) -> Option<&str> {
        self.choices
            .first()
            .and_then(|choice| choice.message.content.as_ref())
            .and_then(|content| content.as_text())
    }

    /// Get the primary response message
    pub fn message(&self) -> Option<&Message> {
        self.choices.first().map(|choice| &choice.message)
    }

    /// Get tool calls from the primary response
    pub fn tool_calls(&self) -> Option<&[ToolCall]> {
        self.choices
            .first()
            .and_then(|choice| choice.message.tool_calls.as_deref())
    }

    /// Check if the response includes tool calls
    pub fn has_tool_calls(&self) -> bool {
        self.tool_calls().is_some_and(|calls| !calls.is_empty())
    }
}

/// A single completion choice
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Choice {
    /// Index of this choice
    pub index: u32,

    /// The message content
    pub message: Message,

    /// Reason why the completion finished
    pub finish_reason: Option<String>,
}

/// A message in a chat completion
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// Role of the message sender
    pub role: String,

    /// Content of the message
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<MessageContent>,

    /// Tool/function calls made by the assistant
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
}

/// A tool/function call made by the assistant
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolCall {
    /// Unique identifier for the tool call
    pub id: String,

    /// Type of tool call (usually "function")
    #[serde(rename = "type")]
    pub call_type: String,

    /// The function being called
    pub function: FunctionCall,
}

/// Details of a function call
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct FunctionCall {
    /// Name of the function
    pub name: String,

    /// JSON string of function arguments
    pub arguments: String,
}

impl FunctionCall {
    /// Parse the arguments as JSON
    pub fn parse_arguments(&self) -> Result<Value, serde_json::Error> {
        serde_json::from_str(&self.arguments)
    }
}

/// Token usage information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Usage {
    /// Number of tokens in the prompt
    pub prompt_tokens: u32,

    /// Number of tokens in the completion
    pub completion_tokens: u32,

    /// Total number of tokens used
    pub total_tokens: u32,

    /// Number of prompt tokens served from cache (reduces cost on repeated prompts).
    /// Only present when prompt caching was used.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cached_prompt_tokens: Option<u32>,

    /// Tokens used internally by the model for reasoning (not in the final response).
    /// Only present for reasoning models.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_tokens: Option<u32>,
}

/// Available Grok models
///
/// ## Grok 4.20 Notes
/// - **Flagship model** with industry-leading speed and agentic tool calling.
/// - Context window: **2,000,000 tokens**.
/// - Does **not** support the `logprobs` field (ignored if sent).
/// - Reasoning models do **not** support `presence_penalty`, `frequency_penalty`, or `stop`.
/// - No `reasoning_effort` parameter (returns error if sent).
///
/// ## Migrating from Grok 3 → Grok 4 / 4.20
/// Grok 4 is a **reasoning-only** model; there is no non-reasoning mode.
/// Remove any `presence_penalty`, `frequency_penalty`, `stop`, or `reasoning_effort`
/// fields from your requests before switching.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Model {
    // ── Grok 4.20 (flagship, 2M context) ──────────────────────────────────────
    /// Grok 4.20 Reasoning (0309) — flagship reasoning model, 2M context window
    Grok4_20_0309Reasoning,

    /// Grok 4.20 Non-Reasoning (0309) — flagship standard model, 2M context window
    Grok4_20_0309NonReasoning,

    /// Grok 4.20 Multi-Agent (0309) — optimised for agentic tool-calling pipelines
    Grok4_20MultiAgent0309,

    // ── Grok 4.1 Fast ─────────────────────────────────────────────────────────
    /// Grok 4.1 Fast Reasoning — fast reasoning model, 2M context window
    Grok4_1FastReasoning,

    /// Grok 4.1 Fast Non-Reasoning — fast standard model, 2M context window
    Grok4_1FastNonReasoning,

    // ── Grok 4 ────────────────────────────────────────────────────────────────
    /// Grok 4 (0709) — reasoning-only; no `presence_penalty`/`frequency_penalty`
    Grok4_0709,

    // ── Grok 3 ────────────────────────────────────────────────────────────────
    /// Grok 3 — previous flagship model (knowledge cut-off: November 2024)
    Grok3,

    /// Grok 3 Mini — efficient smaller model
    Grok3Mini,

    // ── Code ──────────────────────────────────────────────────────────────────
    /// Grok Code Fast 1 — optimised for code generation
    GrokCodeFast1,

    // ── Image generation ──────────────────────────────────────────────────────
    /// Grok Imagine Image Pro — high-quality image generation model
    GrokImagineImagePro,

    /// Grok Imagine Image — standard image generation model
    GrokImagineImage,

    // ── Video generation ──────────────────────────────────────────────────────
    /// Grok Imagine Video — video generation model
    GrokImagineVideo,
}

impl Model {
    /// Get the model identifier string used in API requests.
    pub fn as_str(&self) -> &'static str {
        match self {
            // Grok 4.20
            Model::Grok4_20_0309Reasoning => "grok-4.20-0309-reasoning",
            Model::Grok4_20_0309NonReasoning => "grok-4.20-0309-non-reasoning",
            Model::Grok4_20MultiAgent0309 => "grok-4.20-multi-agent-0309",
            // Grok 4.1 Fast
            Model::Grok4_1FastReasoning => "grok-4-1-fast-reasoning",
            Model::Grok4_1FastNonReasoning => "grok-4-1-fast-non-reasoning",
            // Grok 4
            Model::Grok4_0709 => "grok-4-0709",
            // Grok 3
            Model::Grok3 => "grok-3",
            Model::Grok3Mini => "grok-3-mini",
            // Code
            Model::GrokCodeFast1 => "grok-code-fast-1",
            // Image generation
            Model::GrokImagineImagePro => "grok-imagine-image-pro",
            Model::GrokImagineImage => "grok-imagine-image",
            // Video generation
            Model::GrokImagineVideo => "grok-imagine-video",
        }
    }

    /// Parse a model from its API identifier string.
    pub fn parse(s: &str) -> Option<Self> {
        match s {
            // Grok 4.20
            "grok-4.20-0309-reasoning" => Some(Model::Grok4_20_0309Reasoning),
            "grok-4.20-0309-non-reasoning" => Some(Model::Grok4_20_0309NonReasoning),
            "grok-4.20-multi-agent-0309" => Some(Model::Grok4_20MultiAgent0309),
            // Grok 4.1 Fast
            "grok-4-1-fast-reasoning" => Some(Model::Grok4_1FastReasoning),
            "grok-4-1-fast-non-reasoning" => Some(Model::Grok4_1FastNonReasoning),
            // Grok 4
            "grok-4-0709" => Some(Model::Grok4_0709),
            // Grok 3
            "grok-3" => Some(Model::Grok3),
            "grok-3-mini" => Some(Model::Grok3Mini),
            // Code
            "grok-code-fast-1" => Some(Model::GrokCodeFast1),
            // Image generation
            "grok-imagine-image-pro" => Some(Model::GrokImagineImagePro),
            "grok-imagine-image" => Some(Model::GrokImagineImage),
            // Video generation
            "grok-imagine-video" => Some(Model::GrokImagineVideo),
            _ => None,
        }
    }

    /// Get all available models, ordered from newest to oldest.
    pub fn all() -> Vec<Self> {
        vec![
            // Grok 4.20 — flagship
            Model::Grok4_20_0309Reasoning,
            Model::Grok4_20_0309NonReasoning,
            Model::Grok4_20MultiAgent0309,
            // Grok 4.1 Fast
            Model::Grok4_1FastReasoning,
            Model::Grok4_1FastNonReasoning,
            // Grok 4
            Model::Grok4_0709,
            // Grok 3
            Model::Grok3,
            Model::Grok3Mini,
            // Code
            Model::GrokCodeFast1,
            // Image generation
            Model::GrokImagineImagePro,
            Model::GrokImagineImage,
            // Video generation
            Model::GrokImagineVideo,
        ]
    }

    /// Returns `true` if this model is a reasoning model.
    ///
    /// Reasoning models do **not** support `presence_penalty`, `frequency_penalty`,
    /// `stop`, or `reasoning_effort` parameters. Sending those fields will cause an
    /// API error.
    pub fn is_reasoning_model(&self) -> bool {
        matches!(
            self,
            Model::Grok4_20_0309Reasoning
                | Model::Grok4_20MultiAgent0309
                | Model::Grok4_1FastReasoning
                | Model::Grok4_0709
        )
    }

    /// Returns `true` if this model supports `presence_penalty` and
    /// `frequency_penalty` request fields.
    ///
    /// Reasoning models (Grok 4 and above reasoning variants) do not support
    /// these parameters. Sending them will return an API error.
    pub fn supports_frequency_presence_penalty(&self) -> bool {
        !self.is_reasoning_model()
    }

    /// Returns `true` if this model supports the `logprobs` response field.
    ///
    /// Grok 4.20 models ignore `logprobs` if specified. All other text models
    /// support it.
    pub fn supports_logprobs(&self) -> bool {
        !matches!(
            self,
            Model::Grok4_20_0309Reasoning
                | Model::Grok4_20_0309NonReasoning
                | Model::Grok4_20MultiAgent0309
        )
    }

    /// Returns the context window size in tokens, if known.
    ///
    /// Returns `None` for image/video generation models which do not use a
    /// token-based context window.
    pub fn context_window(&self) -> Option<u32> {
        match self {
            // 2 million token context
            Model::Grok4_20_0309Reasoning
            | Model::Grok4_20_0309NonReasoning
            | Model::Grok4_20MultiAgent0309
            | Model::Grok4_1FastReasoning
            | Model::Grok4_1FastNonReasoning
            | Model::Grok4_0709 => Some(2_000_000),
            // Grok 3 — 131K context
            Model::Grok3 | Model::Grok3Mini => Some(131_072),
            // Code model
            Model::GrokCodeFast1 => Some(131_072),
            // Generation models — no text context window
            Model::GrokImagineImagePro | Model::GrokImagineImage | Model::GrokImagineVideo => None,
        }
    }

    /// Returns `true` if this is a text / language model (not image or video generation).
    pub fn is_language_model(&self) -> bool {
        !matches!(
            self,
            Model::GrokImagineImagePro | Model::GrokImagineImage | Model::GrokImagineVideo
        )
    }

    /// Returns `true` if this model can generate images.
    pub fn is_image_model(&self) -> bool {
        matches!(self, Model::GrokImagineImagePro | Model::GrokImagineImage)
    }

    /// Returns `true` if this model can generate video.
    pub fn is_video_model(&self) -> bool {
        matches!(self, Model::GrokImagineVideo)
    }
}

impl std::fmt::Display for Model {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

impl From<Model> for String {
    fn from(model: Model) -> Self {
        model.as_str().to_string()
    }
}

#[cfg(test)]
mod tests {
    // ... (rest of the tests remain unchanged)
}