xai-grok 0.1.0

//! Wire types for the xAI Grok chat-completions API.
//!
//! xAI's `/v1/chat/completions` endpoint is `OpenAI`-compatible, so these types
//! mirror the `OpenAI` chat schema: `messages` (text or multimodal content
//! parts), `tools` / `tool_calls` for function calling, `response_format` for
//! structured outputs, and the streaming `chat.completion.chunk` delta shape.
//!
//! They are intentionally `serde`-only data holders — all transport and
//! streaming logic lives in [`crate::client`].

use serde::{Deserialize, Serialize};

/// The role of a chat message.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    /// System / developer instruction.
    System,
    /// End-user input.
    User,
    /// Model output.
    Assistant,
    /// The result of a tool/function call, fed back to the model.
    Tool,
}

/// A single piece of message content.
///
/// A message's content is either a plain string or an array of typed parts
/// (text + images) for vision requests. Both serialize to the shapes xAI
/// accepts.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Content {
    /// Plain text content (the common case).
    Text(String),
    /// Multimodal content: a mix of text and image parts.
    Parts(Vec<ContentPart>),
}

impl Content {
    /// Convenience constructor for plain text.
    pub fn text(s: impl Into<String>) -> Self {
        Self::Text(s.into())
    }
}

impl<T: Into<String>> From<T> for Content {
    fn from(s: T) -> Self {
        Self::Text(s.into())
    }
}

/// One part of a multimodal message.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentPart {
    /// A text fragment.
    Text {
        /// The text.
        text: String,
    },
    /// An image, referenced by URL or `data:` URI.
    ImageUrl {
        /// The image reference.
        image_url: ImageUrl,
    },
}

impl ContentPart {
    /// A text part.
    pub fn text(s: impl Into<String>) -> Self {
        Self::Text { text: s.into() }
    }

    /// An image part from a URL or `data:` URI.
    pub fn image(url: impl Into<String>) -> Self {
        Self::ImageUrl {
            image_url: ImageUrl {
                url: url.into(),
                detail: None,
            },
        }
    }
}

/// An image reference for a vision request.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ImageUrl {
    /// A public URL or a base64 `data:image/...;base64,...` URI.
    pub url: String,
    /// Optional detail hint (`"low"`, `"high"`, `"auto"`).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub detail: Option<String>,
}

/// A chat message.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Message {
    /// Who authored the message.
    pub role: Role,
    /// The message content. Optional because an assistant message that only
    /// makes tool calls carries `null` content.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<Content>,
    /// Tool calls requested by the assistant (assistant messages only).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    /// For `role = "tool"`: which tool call this message answers.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

impl Message {
    /// A system message.
    pub fn system(content: impl Into<Content>) -> Self {
        Self::simple(Role::System, content)
    }

    /// A user message.
    pub fn user(content: impl Into<Content>) -> Self {
        Self::simple(Role::User, content)
    }

    /// An assistant message.
    pub fn assistant(content: impl Into<Content>) -> Self {
        Self::simple(Role::Assistant, content)
    }

    /// A tool-result message answering `tool_call_id`.
    pub fn tool(tool_call_id: impl Into<String>, content: impl Into<Content>) -> Self {
        Self {
            role: Role::Tool,
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: Some(tool_call_id.into()),
        }
    }

    fn simple(role: Role, content: impl Into<Content>) -> Self {
        Self {
            role,
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: None,
        }
    }
}

/// A tool the model may call. Currently only function tools exist.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Tool {
    /// A function tool, described by a JSON Schema for its parameters.
    Function {
        /// The function definition.
        function: FunctionDef,
    },
}

impl Tool {
    /// Build a function tool from a name, description, and JSON-Schema
    /// parameters object.
    pub fn function(
        name: impl Into<String>,
        description: impl Into<String>,
        parameters: serde_json::Value,
    ) -> Self {
        Self::Function {
            function: FunctionDef {
                name: name.into(),
                description: Some(description.into()),
                parameters: Some(parameters),
            },
        }
    }
}

/// A function tool definition.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct FunctionDef {
    /// The function name the model uses when calling it.
    pub name: String,
    /// A natural-language description of what the function does.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    /// A JSON Schema object describing the parameters.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parameters: Option<serde_json::Value>,
}

/// A tool call emitted by the model.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolCall {
    /// The unique id of this call (echo it back on the tool result message).
    pub id: String,
    /// The tool type (always `"function"` today).
    #[serde(rename = "type", default = "default_tool_type")]
    pub kind: String,
    /// The called function and its arguments.
    pub function: FunctionCall,
}

fn default_tool_type() -> String {
    "function".to_string()
}

/// The function part of a [`ToolCall`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct FunctionCall {
    /// The function name.
    pub name: String,
    /// The arguments as a JSON string (parse with `serde_json`).
    pub arguments: String,
}

/// Controls how the model selects tools.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ToolChoice {
    /// `"auto"`, `"none"`, or `"required"`.
    Mode(String),
    /// Force a specific named function.
    Named {
        /// Always `"function"`.
        #[serde(rename = "type")]
        kind: String,
        /// The forced function.
        function: NamedFunction,
    },
}

/// The function targeted by a forced [`ToolChoice`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct NamedFunction {
    /// The function name to force.
    pub name: String,
}

/// Structured-output specification (`response_format`).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ResponseFormat {
    /// Free-form text (the default).
    Text,
    /// Any valid JSON object.
    JsonObject,
    /// JSON constrained to a schema.
    JsonSchema {
        /// The schema wrapper.
        json_schema: JsonSchema,
    },
}

impl ResponseFormat {
    /// Build a strict JSON-schema response format.
    pub fn json_schema(name: impl Into<String>, schema: serde_json::Value) -> Self {
        Self::JsonSchema {
            json_schema: JsonSchema {
                name: name.into(),
                schema,
                strict: Some(true),
            },
        }
    }
}

/// A named JSON Schema for structured outputs.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct JsonSchema {
    /// A name for the schema.
    pub name: String,
    /// The JSON Schema itself.
    pub schema: serde_json::Value,
    /// Whether the model must adhere strictly to the schema.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub strict: Option<bool>,
}

/// A chat-completions request body.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatRequest {
    /// The model id, e.g. `"grok-4.3"`.
    pub model: String,
    /// The conversation so far.
    pub messages: Vec<Message>,
    /// Whether to stream the response as SSE.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,
    /// Sampling temperature.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    /// Maximum tokens to generate.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,
    /// Available tools.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<Tool>>,
    /// Tool-selection policy.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ToolChoice>,
    /// Structured-output specification.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,
}

impl ChatRequest {
    /// Start a request for `model` with the given messages.
    pub fn new(model: impl Into<String>, messages: Vec<Message>) -> Self {
        Self {
            model: model.into(),
            messages,
            stream: None,
            temperature: None,
            max_tokens: None,
            tools: None,
            tool_choice: None,
            response_format: None,
        }
    }

    /// Set the sampling temperature (builder style).
    #[must_use]
    pub fn temperature(mut self, t: f32) -> Self {
        self.temperature = Some(t);
        self
    }

    /// Set the max output tokens (builder style).
    #[must_use]
    pub fn max_tokens(mut self, n: u32) -> Self {
        self.max_tokens = Some(n);
        self
    }

    /// Attach available tools (builder style).
    #[must_use]
    pub fn tools(mut self, tools: Vec<Tool>) -> Self {
        self.tools = Some(tools);
        self
    }

    /// Set the tool-selection policy (builder style).
    #[must_use]
    pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
        self.tool_choice = Some(choice);
        self
    }

    /// Set the structured-output format (builder style).
    #[must_use]
    pub fn response_format(mut self, format: ResponseFormat) -> Self {
        self.response_format = Some(format);
        self
    }
}

/// Token-usage accounting returned with a completion.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
    /// Tokens in the prompt.
    pub prompt_tokens: u32,
    /// Tokens in the completion.
    pub completion_tokens: u32,
    /// Total tokens billed.
    pub total_tokens: u32,
}

/// A non-streamed chat completion response.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatResponse {
    /// The completion id.
    pub id: String,
    /// The model that produced it.
    pub model: String,
    /// The returned choices.
    pub choices: Vec<ResponseChoice>,
    /// Usage accounting, if reported.
    #[serde(default)]
    pub usage: Option<Usage>,
}

impl ChatResponse {
    /// The first choice's message content as text, if any.
    pub fn content(&self) -> Option<&str> {
        self.choices
            .first()
            .and_then(|c| c.message.content.as_ref())
            .and_then(|c| match c {
                Content::Text(s) => Some(s.as_str()),
                Content::Parts(_) => None,
            })
    }

    /// The tool calls from the first choice, if any.
    pub fn tool_calls(&self) -> &[ToolCall] {
        self.choices
            .first()
            .and_then(|c| c.message.tool_calls.as_deref())
            .unwrap_or(&[])
    }
}

/// One choice in a non-streamed response.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ResponseChoice {
    /// The choice index.
    pub index: u32,
    /// The assistant message.
    pub message: Message,
    /// Why generation stopped (`"stop"`, `"tool_calls"`, `"length"`, ...).
    #[serde(default)]
    pub finish_reason: Option<String>,
}

// ---------------------------------------------------------------------------
// Streaming chunk types (`chat.completion.chunk`)
// ---------------------------------------------------------------------------

/// One streamed `chat.completion.chunk` object (parsed from an SSE `data:`
/// payload). The `[DONE]` sentinel is handled by the client, not represented
/// here.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatChunk {
    /// The completion id (stable across chunks).
    #[serde(default)]
    pub id: String,
    /// The per-choice deltas in this chunk.
    pub choices: Vec<ChunkChoice>,
    /// Usage, present only on the final chunk when requested.
    #[serde(default)]
    pub usage: Option<Usage>,
}

/// A per-choice delta within a [`ChatChunk`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChunkChoice {
    /// The choice index this delta applies to.
    pub index: u32,
    /// The incremental delta.
    pub delta: Delta,
    /// The finish reason, set on the choice's final chunk.
    #[serde(default)]
    pub finish_reason: Option<String>,
}

/// The incremental delta carried by a streaming chunk.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct Delta {
    /// The role, sent on the first delta for a choice.
    #[serde(default)]
    pub role: Option<String>,
    /// A content fragment to append.
    #[serde(default)]
    pub content: Option<String>,
    /// Tool-call fragments to merge by `index`.
    #[serde(default)]
    pub tool_calls: Option<Vec<ToolCallDelta>>,
}

/// A streamed tool-call fragment.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolCallDelta {
    /// Which tool call (by position) this fragment belongs to.
    pub index: u32,
    /// The tool-call id, sent on the first fragment.
    #[serde(default)]
    pub id: Option<String>,
    /// The function name/arguments fragment.
    #[serde(default)]
    pub function: Option<FunctionDelta>,
}

/// The function part of a [`ToolCallDelta`].
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct FunctionDelta {
    /// The function name, sent on the first fragment.
    #[serde(default)]
    pub name: Option<String>,
    /// An arguments-JSON fragment to concatenate.
    #[serde(default)]
    pub arguments: Option<String>,
}