grok_api 0.1.71

//! 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>,

    /// Reasoning effort level for models that support extended thinking.
    /// Valid values: `"low"`, `"medium"`, `"high"`.
    /// Only send this for models that support it (e.g. `grok-4.3`, `grok-3-mini`).
    /// Sending it to unsupported models returns an API error.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_effort: Option<String>,
}

/// 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>>,

    /// Reasoning / chain-of-thought content produced by the model when
    /// `reasoning_effort` was set on the request.  Only present for
    /// reasoning-capable models (e.g. grok-4.3, grok-3-mini).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_content: Option<String>,
}

/// 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.3 — New Flagship (May 2026)
/// - **Fastest, most intelligent** model from xAI.
/// - Context window: **1,000,000 tokens**.
/// - Supports `reasoning_effort` with three levels: `low`, `medium`, `high`.
/// - Priced at $1.25 / 1 M input tokens, $2.50 / 1 M output tokens.
/// - **Recommended replacement** for all retiring Grok 4.1 Fast, Grok 4-0709,
///   Grok Code Fast 1, and Grok 3 workloads.
///
/// ## Grok 4.20 Notes
/// - **Flagship** with industry-leading speed and agentic tool calling.
/// - Context window: **2,000,000 tokens**.
/// - Does **not** support the `logprobs` field (ignored if sent).
/// - `Grok4_20_0309Reasoning` / `Grok4_20MultiAgent0309` do **not** support
///   `presence_penalty`, `frequency_penalty`, or `stop`.
/// - No `reasoning_effort` parameter on these variants (returns error if sent).
///
/// ## Retirements — effective May 15, 2026
/// The following variants are **deprecated** and will stop working on
/// 2026-05-15 12:00 PT. Compiler deprecation warnings guide you to the
/// recommended replacement.
///
/// | Deprecated variant | Retires | Use instead |
/// |---|---|---|
/// | `Grok4_1FastReasoning` | 2026-05-15 | [`Model::Grok4_3`] |
/// | `Grok4_1FastNonReasoning` | 2026-05-15 | [`Model::Grok4_20NonReasoning`] |
/// | `Grok4_0709` | 2026-05-15 | [`Model::Grok4_3`] |
/// | `Grok3` | 2026-05-15 | [`Model::Grok4_3`] |
/// | `GrokCodeFast1` | 2026-05-15 | [`Model::Grok4_3`] |
/// | `GrokImagineImagePro` | 2026-05-15 | [`Model::GrokImagineImage`] |
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Model {
    // ── Grok 4.3 (new flagship, 1M context, May 2026) ─────────────────────────
    /// Grok 4.3 — fastest and most intelligent xAI model; 1 M token context
    /// window; configurable reasoning effort (`low` / `medium` / `high`).
    ///
    /// **Recommended default for all new and migrated workloads.**
    Grok4_3,

    // ── 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 — stable alias for the 4.20 fast standard model,
    /// 2M context window. Recommended replacement for `grok-4-1-fast-non-reasoning`.
    Grok4_20NonReasoning,

    /// Grok 4.20 Non-Reasoning (0309) — dated variant of the 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 (DEPRECATED — retires 2026-05-15) ───────────────────────
    /// ⚠️ **Deprecated** — retires 2026-05-15. Use [`Model::Grok4_3`] instead.
    #[deprecated(
        since = "0.1.7",
        note = "Retires 2026-05-15. Migrate to Model::Grok4_3 (grok-4.3)."
    )]
    Grok4_1FastReasoning,

    /// ⚠️ **Deprecated** — retires 2026-05-15. Use [`Model::Grok4_20NonReasoning`] instead.
    #[deprecated(
        since = "0.1.7",
        note = "Retires 2026-05-15. Migrate to Model::Grok4_20NonReasoning (grok-4.20-non-reasoning)."
    )]
    Grok4_1FastNonReasoning,

    // ── Grok 4 (DEPRECATED — retires 2026-05-15) ──────────────────────────────
    /// ⚠️ **Deprecated** — retires 2026-05-15. Use [`Model::Grok4_3`] instead.
    #[deprecated(
        since = "0.1.7",
        note = "Retires 2026-05-15. Migrate to Model::Grok4_3 (grok-4.3)."
    )]
    Grok4_0709,

    // ── Grok 3 ────────────────────────────────────────────────────────────────
    /// ⚠️ **Deprecated** — retires 2026-05-15. Use [`Model::Grok4_3`] instead.
    #[deprecated(
        since = "0.1.7",
        note = "Retires 2026-05-15. Migrate to Model::Grok4_3 (grok-4.3)."
    )]
    Grok3,

    /// Grok 3 Mini — efficient smaller model (not retiring May 2026)
    Grok3Mini,

    // ── Code (DEPRECATED — retires 2026-05-15) ────────────────────────────────
    /// ⚠️ **Deprecated** — retires 2026-05-15. Use [`Model::Grok4_3`] instead.
    #[deprecated(
        since = "0.1.7",
        note = "Retires 2026-05-15. Migrate to Model::Grok4_3 (grok-4.3)."
    )]
    GrokCodeFast1,

    // ── Image generation ──────────────────────────────────────────────────────
    /// ⚠️ **Deprecated** — retires 2026-05-15. Use [`Model::GrokImagineImage`] instead.
    #[deprecated(
        since = "0.1.7",
        note = "Retires 2026-05-15. Migrate to Model::GrokImagineImage (grok-imagine-image)."
    )]
    GrokImagineImagePro,

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

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

#[allow(deprecated)]
impl Model {
    /// Get the model identifier string used in API requests.
    pub fn as_str(&self) -> &'static str {
        match self {
            // Grok 4.3
            Model::Grok4_3 => "grok-4.3",
            // Grok 4.20
            Model::Grok4_20_0309Reasoning => "grok-4.20-0309-reasoning",
            Model::Grok4_20NonReasoning => "grok-4.20-non-reasoning",
            Model::Grok4_20_0309NonReasoning => "grok-4.20-0309-non-reasoning",
            Model::Grok4_20MultiAgent0309 => "grok-4.20-multi-agent-0309",
            // Grok 4.1 Fast (deprecated)
            Model::Grok4_1FastReasoning => "grok-4-1-fast-reasoning",
            Model::Grok4_1FastNonReasoning => "grok-4-1-fast-non-reasoning",
            // Grok 4 (deprecated)
            Model::Grok4_0709 => "grok-4-0709",
            // Grok 3
            Model::Grok3 => "grok-3",
            Model::Grok3Mini => "grok-3-mini",
            // Code (deprecated)
            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.
    ///
    /// Returns `None` for unknown strings. Deprecated model strings are still
    /// recognised so that existing config files and env vars continue to work
    /// until the xAI API retires them on 2026-05-15.
    pub fn parse(s: &str) -> Option<Self> {
        match s {
            // Grok 4.3
            "grok-4.3" => Some(Model::Grok4_3),
            // Grok 4.20
            "grok-4.20-0309-reasoning" => Some(Model::Grok4_20_0309Reasoning),
            "grok-4.20-non-reasoning" => Some(Model::Grok4_20NonReasoning),
            "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 (deprecated — recognised until 2026-05-15)
            #[allow(deprecated)]
            "grok-4-1-fast-reasoning" => Some(Model::Grok4_1FastReasoning),
            #[allow(deprecated)]
            "grok-4-1-fast-non-reasoning" => Some(Model::Grok4_1FastNonReasoning),
            // Grok 4 (deprecated — recognised until 2026-05-15)
            #[allow(deprecated)]
            "grok-4-0709" => Some(Model::Grok4_0709),
            // Grok 3
            #[allow(deprecated)]
            "grok-3" => Some(Model::Grok3),
            "grok-3-mini" => Some(Model::Grok3Mini),
            // Code (deprecated — recognised until 2026-05-15)
            #[allow(deprecated)]
            "grok-code-fast-1" => Some(Model::GrokCodeFast1),
            // Image generation
            #[allow(deprecated)]
            "grok-imagine-image-pro" => Some(Model::GrokImagineImagePro),
            "grok-imagine-image" => Some(Model::GrokImagineImage),
            // Video generation
            "grok-imagine-video" => Some(Model::GrokImagineVideo),
            _ => None,
        }
    }

    /// All **active** (non-deprecated) models, ordered newest → oldest.
    ///
    /// Deprecated variants are intentionally excluded so that UI lists,
    /// help text, and auto-complete only surface models that still work.
    pub fn all() -> Vec<Self> {
        vec![
            // Grok 4.3 — new flagship
            Model::Grok4_3,
            // Grok 4.20 — active variants
            Model::Grok4_20_0309Reasoning,
            Model::Grok4_20NonReasoning,
            Model::Grok4_20_0309NonReasoning,
            Model::Grok4_20MultiAgent0309,
            // Grok 3 Mini (Grok3 itself is deprecated)
            Model::Grok3Mini,
            // Image generation (GrokImagineImagePro is deprecated)
            Model::GrokImagineImage,
            // Video generation
            Model::GrokImagineVideo,
        ]
    }

    /// All models including deprecated ones — useful for migration tooling
    /// or exhaustive API-string lookup tables.
    pub fn all_including_deprecated() -> Vec<Self> {
        vec![
            // Grok 4.3
            Model::Grok4_3,
            // Grok 4.20
            Model::Grok4_20_0309Reasoning,
            Model::Grok4_20NonReasoning,
            Model::Grok4_20_0309NonReasoning,
            Model::Grok4_20MultiAgent0309,
            // Grok 4.1 Fast (deprecated)
            #[allow(deprecated)]
            Model::Grok4_1FastReasoning,
            #[allow(deprecated)]
            Model::Grok4_1FastNonReasoning,
            // Grok 4 (deprecated)
            #[allow(deprecated)]
            Model::Grok4_0709,
            // Grok 3
            #[allow(deprecated)]
            Model::Grok3,
            Model::Grok3Mini,
            // Code (deprecated)
            #[allow(deprecated)]
            Model::GrokCodeFast1,
            // Image generation
            #[allow(deprecated)]
            Model::GrokImagineImagePro,
            Model::GrokImagineImage,
            // Video generation
            Model::GrokImagineVideo,
        ]
    }

    /// Returns `true` if this model is a **pure reasoning model** that does
    /// **not** support `presence_penalty`, `frequency_penalty`, `stop`, or
    /// `reasoning_effort` parameters. Sending those fields will cause an API error.
    ///
    /// Note: [`Model::Grok4_3`] is **not** a pure reasoning model — it
    /// supports the `reasoning_effort` parameter (low / medium / high) but
    /// also works without it.
    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.
    ///
    /// Pure reasoning models (Grok 4.20 reasoning variants, Grok 4-0709)
    /// do not support these parameters.
    pub fn supports_frequency_presence_penalty(&self) -> bool {
        !self.is_reasoning_model()
    }

    /// Returns `true` if this model supports the `reasoning_effort` request
    /// field (`"low"` / `"medium"` / `"high"`).
    ///
    /// Currently only [`Model::Grok4_3`] supports configurable reasoning effort.
    pub fn supports_reasoning_effort(&self) -> bool {
        matches!(self, Model::Grok4_3)
    }

    /// Returns `true` if this model supports the `logprobs` response field.
    ///
    /// Grok 4.20 family and Grok 4.3 do not support `logprobs`.
    pub fn supports_logprobs(&self) -> bool {
        !matches!(
            self,
            Model::Grok4_3
                | Model::Grok4_20_0309Reasoning
                | Model::Grok4_20NonReasoning
                | 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 {
            // 1 million token context (Grok 4.3)
            Model::Grok4_3 => Some(1_000_000),
            // 2 million token context (Grok 4.20 family)
            Model::Grok4_20_0309Reasoning
            | Model::Grok4_20NonReasoning
            | Model::Grok4_20_0309NonReasoning
            | Model::Grok4_20MultiAgent0309
            | Model::Grok4_1FastReasoning
            | Model::Grok4_1FastNonReasoning
            | Model::Grok4_0709 => Some(2_000_000),
            // Grok 3 — 131 K 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)
}