zai-rs 0.5.1

//! # Chat Request Body
//!
//! Provides the core `ChatBody` structure shared by all chat-completion
//! endpoints. The generic type parameters enforce compile-time compatibility
//! between model and message types.
//!
//! # Type Parameters
//!
//! - `N` — Model identifier type (must implement `ModelName`)
//! - `M` — Message type (must form a `Bounded` pair with `N`)

use serde::Serialize;
use validator::*;

use super::{tools::*, traits::*};

/// Main request body structure for chat API calls.
///
/// This structure represents a complete chat request with all possible
/// configuration options. It uses generic types to support different model
/// names and message types while maintaining type safety through trait bounds.
///
/// # Type Parameters
///
/// * `N` - The model name type, must implement `ModelName`
/// * `M` - The message type, must form a `Bounded` pair with `N`
///
/// # Examples
///
/// ```
/// use zai_rs::model::{
///     chat_base_request::ChatBody,
///     chat_message_types::TextMessage,
///     chat_models::GLM5_2,
/// };
///
/// // Create a basic chat request
/// let chat_body = ChatBody::new(GLM5_2 {}, TextMessage::user("Hello"))
///     .add_message(TextMessage::assistant("How can I help?"))
///     .with_temperature(0.7)
///     .with_max_tokens(1_000);
/// ```
#[derive(Debug, Clone, Validate, Serialize)]
pub struct ChatBody<N, M>
where
    N: ModelName,
    (N, M): Bounded,
{
    /// The model to use for the chat completion.
    pub model: N,

    /// A list of messages comprising the conversation so far.
    pub messages: Vec<M>,

    /// A unique identifier for the request. Optional field that will be omitted
    /// from serialization if not provided.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub request_id: Option<String>,

    /// Thinking-mode configuration for models that support extended reasoning.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thinking: Option<ThinkingType>,

    /// Controls the depth of reasoning when thinking mode is enabled. Only
    /// available for GLM-5.2 and above (models that implement
    /// `ReasoningEffortEnable`). See [`ReasoningEffort`] for the available
    /// levels.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_effort: Option<ReasoningEffort>,

    /// Whether to use sampling during generation. When `true`, the model will
    /// use probabilistic sampling; when `false`, it will use deterministic
    /// generation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub do_sample: Option<bool>,

    /// Whether to stream back partial message deltas as they are generated.
    /// When `true`, responses will be sent as server-sent events.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,

    /// Whether to enable streaming of tool calls (streaming function call
    /// parameters). Supported by GLM-5.2, GLM-5.1, GLM-5, GLM-5-Turbo, GLM-4.7,
    /// and GLM-4.6 models. Defaults to false when omitted.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_stream: Option<bool>,

    /// Controls randomness in the output. Higher values (closer to 1.0) make
    /// the output more random, while lower values (closer to 0.0) make it
    /// more deterministic. Must be between 0.0 and 1.0.
    #[serde(skip_serializing_if = "Option::is_none")]
    #[validate(range(min = 0.0, max = 1.0))]
    pub temperature: Option<f64>,

    /// Controls diversity via nucleus sampling. Only tokens with cumulative
    /// probability up to `top_p` are considered. Must be between 0.0 and
    /// 1.0.
    #[serde(skip_serializing_if = "Option::is_none")]
    #[validate(range(min = 0.0, max = 1.0))]
    pub top_p: Option<f64>,

    /// The maximum number of tokens to generate in the completion.
    /// Must be between 1 and 98304.
    #[serde(skip_serializing_if = "Option::is_none")]
    #[validate(range(min = 1, max = 98304))]
    pub max_tokens: Option<u32>,

    /// A list of tools the model may call. Currently supports function calling,
    /// web search, and retrieval tools.
    /// Note: server expects an array; we model this as a vector of tool items.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<Tools>>,

    /// Controls whether the model calls a tool, and which one. Only meaningful
    /// when [`tools`](Self::tools) is set. See [`ToolChoice`] for the wire
    /// forms (`"auto"`, `"none"`, or a specific function).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ToolChoice>,

    /// A unique identifier representing your end-user, which can help monitor
    /// and detect abuse. Must be between 6 and 128 characters long.
    #[serde(skip_serializing_if = "Option::is_none")]
    #[validate(length(min = 6, max = 128))]
    pub user_id: Option<String>,

    /// Up to 1 sequence where the API will stop generating further tokens.
    #[serde(skip_serializing_if = "Option::is_none")]
    #[validate(length(min = 1, max = 1))]
    pub stop: Option<Vec<String>>,

    /// An object specifying the format that the model must output.
    /// Can be either text or JSON object format.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,
}

impl<N, M> ChatBody<N, M>
where
    N: ModelName,
    (N, M): Bounded,
{
    /// Create a new chat request body from a model and the first message batch.
    pub fn new(model: N, messages: M) -> Self {
        Self {
            model,
            messages: vec![messages],
            request_id: None,
            thinking: None,
            reasoning_effort: None,
            do_sample: None,
            stream: None,
            tool_stream: None,
            temperature: None,
            top_p: None,
            max_tokens: None,
            tools: None,
            tool_choice: None,
            user_id: None,
            stop: None,
            response_format: None,
        }
    }

    /// Append another message batch to the conversation.
    pub fn add_messages(mut self, messages: M) -> Self {
        self.messages.push(messages);
        self
    }
    /// Add a single message to the conversation (preferred over add_messages
    /// for clarity when adding one message).
    pub fn add_message(mut self, message: M) -> Self {
        self.messages.push(message);
        self
    }
    /// Add multiple messages to the conversation at once.
    pub fn extend_messages(mut self, messages: impl IntoIterator<Item = M>) -> Self {
        self.messages.extend(messages);
        self
    }
    /// Set the client-side request id (used for tracing/dedup).
    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
        self.request_id = Some(request_id.into());
        self
    }
    /// Enable/disable sampling (`do_sample`).
    pub fn with_do_sample(mut self, do_sample: bool) -> Self {
        self.do_sample = Some(do_sample);
        self
    }
    /// Enable/disable streaming responses.
    pub fn with_stream(mut self, stream: bool) -> Self {
        self.stream = Some(stream);
        self
    }
    /// Set the sampling temperature.
    pub fn with_temperature(mut self, temperature: f64) -> Self {
        self.temperature = Some(temperature);
        self
    }
    /// Set the nucleus-sampling probability (`top_p`).
    pub fn with_top_p(mut self, top_p: f64) -> Self {
        self.top_p = Some(top_p);
        self
    }
    /// Set the maximum number of tokens to generate.
    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }
    /// Deprecated: use `add_tool` (single) or `extend_tools` (Vec) on
    /// ChatBody, or prefer ChatCompletion::add_tool / add_tools at the
    /// client layer.
    #[deprecated(note = "with_tools is deprecated; use add_tool/add_tools instead")]
    pub fn with_tools(mut self, tools: impl Into<Vec<Tools>>) -> Self {
        self.tools = Some(tools.into());
        self
    }
    /// Add a single tool to the request.
    pub fn add_tool(mut self, tool: Tools) -> Self {
        self.tools.get_or_insert(Vec::new()).push(tool);
        self
    }
    /// Add multiple tools to the request at once.
    pub fn extend_tools(mut self, tools: Vec<Tools>) -> Self {
        self.tools.get_or_insert(Vec::new()).extend(tools);
        self
    }
    /// Set the `tool_choice` policy (`auto` / `none` / a specific function).
    /// Only meaningful when [`tools`](Self::tools) is attached.
    pub fn with_tool_choice(mut self, tool_choice: ToolChoice) -> Self {
        self.tool_choice = Some(tool_choice);
        self
    }
    /// Set the response format (plain text or JSON object).
    pub fn with_response_format(mut self, format: ResponseFormat) -> Self {
        self.response_format = Some(format);
        self
    }
    /// Set the end-user id (used for abuse monitoring).
    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
        self.user_id = Some(user_id.into());
        self
    }
    /// Add a stop sequence that halts generation when encountered.
    pub fn with_stop(mut self, stop: String) -> Self {
        self.stop.get_or_insert_with(Vec::new).push(stop);
        self
    }
}

impl<N, M> ChatBody<N, M>
where
    N: ModelName + ThinkEnable,
    (N, M): Bounded,
{
    /// Configure thinking mode for a model that supports extended reasoning.
    ///
    /// This method is only available for models that implement the
    /// `ThinkEnable` trait, ensuring type safety for thinking-enabled
    /// models.
    ///
    /// # Arguments
    ///
    /// * `thinking` - Whether thinking is enabled and whether prior reasoning is
    ///   cleared
    ///
    /// # Returns
    ///
    /// Returns `self` with the thinking field set, allowing for method
    /// chaining.
    ///
    /// # Examples
    ///
    /// ```
    /// use zai_rs::model::{
    ///     chat_base_request::ChatBody,
    ///     chat_message_types::TextMessage,
    ///     chat_models::GLM5_2,
    ///     tools::ThinkingType,
    /// };
    ///
    /// let chat_body = ChatBody::new(GLM5_2 {}, TextMessage::user("Solve this"))
    ///     .with_thinking(ThinkingType::enabled());
    /// ```
    pub fn with_thinking(mut self, thinking: ThinkingType) -> Self {
        self.thinking = Some(thinking);
        self
    }
}

// Only available for models that support the reasoning_effort parameter
// (GLM-5.2 and above).
impl<N, M> ChatBody<N, M>
where
    N: ModelName + ReasoningEffortEnable,
    (N, M): Bounded,
{
    /// Sets the `reasoning_effort` level that controls how much reasoning the
    /// model performs when thinking mode is enabled.
    ///
    /// Only available on GLM-5.2 and above (models implementing
    /// `ReasoningEffortEnable`). Typically combined with
    /// [`with_thinking`](Self::with_thinking) to enable thinking first.
    ///
    /// # Examples
    ///
    /// ```text
    /// use zai_rs::model::tools::ReasoningEffort;
    ///
    /// let chat_body = ChatBody::new(GLM5_2 {}, messages)
    ///     .with_thinking(ThinkingType::enabled())
    ///     .with_reasoning_effort(ReasoningEffort::Max);
    /// ```
    pub fn with_reasoning_effort(mut self, effort: ReasoningEffort) -> Self {
        self.reasoning_effort = Some(effort);
        self
    }
}

// Only available when the model supports streaming tool calls (GLM-4.6)
impl<N, M> ChatBody<N, M>
where
    N: ModelName + ToolStreamEnable,
    (N, M): Bounded,
{
    /// Enables streaming tool calls. Supported by GLM-5.2, GLM-5.1, GLM-5,
    /// GLM-5-Turbo, GLM-4.7, and GLM-4.6 models. Default is false when omitted.
    pub fn with_tool_stream(mut self, tool_stream: bool) -> Self {
        self.tool_stream = Some(tool_stream);
        if tool_stream {
            // Enabling tool_stream implies stream=true
            self.stream = Some(true);
        }
        self
    }
}

// Preserve the historical `Into<Vec<Tools>>` convenience used by `with_tools`.
impl From<Tools> for Vec<Tools> {
    fn from(tool: Tools) -> Self {
        vec![tool]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::{
        chat_message_types::TextMessage,
        chat_models::{GLM4_5, GLM4_6, GLM5_2},
    };

    #[test]
    fn test_with_tool_stream_sets_both_fields() {
        let body: ChatBody<GLM4_6, TextMessage> =
            ChatBody::new(GLM4_6 {}, TextMessage::user("test"));
        let body = body.with_tool_stream(true);
        assert_eq!(body.tool_stream, Some(true));
        assert_eq!(body.stream, Some(true));
    }

    #[test]
    fn test_with_tool_stream_false_does_not_force_stream() {
        let body: ChatBody<GLM4_6, TextMessage> =
            ChatBody::new(GLM4_6 {}, TextMessage::user("test"));
        let body = body.with_tool_stream(false);
        assert_eq!(body.tool_stream, Some(false));
        // stream should not be forced to true when tool_stream is false
        assert_ne!(body.stream, Some(true));
    }

    #[test]
    fn test_add_tool_accumulates() {
        let body: ChatBody<GLM4_6, TextMessage> =
            ChatBody::new(GLM4_6 {}, TextMessage::user("test"));
        let tool = crate::model::tools::Function::new(
            "test_fn",
            "A test function",
            serde_json::json!({"type": "object"}),
        );
        let body = body.add_tool(crate::model::tools::Tools::Function { function: tool });
        assert!(body.tools.is_some());
        assert_eq!(body.tools.unwrap().len(), 1);
    }

    #[test]
    fn test_extend_messages() {
        let body: ChatBody<GLM4_6, TextMessage> =
            ChatBody::new(GLM4_6 {}, TextMessage::user("first"));
        let body = body.extend_messages(vec![
            TextMessage::assistant("second"),
            TextMessage::user("third"),
        ]);
        assert_eq!(body.messages.len(), 3);
        match &body.messages[0] {
            TextMessage::User { content } => assert_eq!(content, "first"),
            _ => panic!("Expected User message"),
        }
    }

    #[test]
    fn test_add_message() {
        let body: ChatBody<GLM4_6, TextMessage> =
            ChatBody::new(GLM4_6 {}, TextMessage::user("first"));
        let body = body.add_message(TextMessage::assistant("second"));
        assert_eq!(body.messages.len(), 2);
    }

    #[test]
    fn test_glm52_reasoning_effort_builder() {
        // reasoning_effort is only available on GLM-5.2+ (ReasoningEffortEnable)
        let body: ChatBody<GLM5_2, TextMessage> = ChatBody::new(GLM5_2 {}, TextMessage::user("hi"))
            .with_thinking(ThinkingType::enabled())
            .with_reasoning_effort(ReasoningEffort::Max);
        assert_eq!(body.reasoning_effort, Some(ReasoningEffort::Max));
        assert!(body.thinking.is_some());
    }

    #[test]
    fn test_glm52_serializes_model_name() {
        let body: ChatBody<GLM5_2, TextMessage> = ChatBody::new(GLM5_2 {}, TextMessage::user("hi"));
        let json = serde_json::to_value(&body).unwrap();
        assert_eq!(json["model"], "glm-5.2");
        // reasoning_effort must be omitted when not set
        assert!(json.get("reasoning_effort").is_none());
    }

    #[test]
    fn test_sampling_parameters_serialize_without_f32_expansion() {
        let body: ChatBody<GLM4_5, TextMessage> = ChatBody::new(GLM4_5 {}, TextMessage::user("hi"))
            .with_temperature(0.7)
            .with_top_p(0.9);

        let json = serde_json::to_value(&body).unwrap();
        assert_eq!(json["temperature"], serde_json::json!(0.7));
        assert_eq!(json["top_p"], serde_json::json!(0.9));

        let serialized = serde_json::to_string(&body).unwrap();
        assert!(serialized.contains("\"temperature\":0.7"));
        assert!(serialized.contains("\"top_p\":0.9"));
        assert!(!serialized.contains("0.699999"));
        assert!(!serialized.contains("0.899999"));
    }

    #[test]
    fn test_reasoning_effort_serializes_level() {
        let body: ChatBody<GLM5_2, TextMessage> = ChatBody::new(GLM5_2 {}, TextMessage::user("hi"))
            .with_reasoning_effort(ReasoningEffort::Xhigh);
        let json = serde_json::to_value(&body).unwrap();
        assert_eq!(json["reasoning_effort"], "xhigh");
    }

    #[test]
    fn test_tool_choice_is_omitted_by_default() {
        let body: ChatBody<GLM4_6, TextMessage> = ChatBody::new(GLM4_6 {}, TextMessage::user("hi"));
        assert!(body.tool_choice.is_none());
        let json = serde_json::to_value(&body).unwrap();
        assert!(json.get("tool_choice").is_none());
    }

    #[test]
    fn test_with_tool_choice_serializes_each_form() {
        // auto -> bare "auto"
        let body: ChatBody<GLM4_6, TextMessage> = ChatBody::new(GLM4_6 {}, TextMessage::user("hi"))
            .with_tool_choice(crate::model::tools::ToolChoice::auto());
        let json = serde_json::to_value(&body).unwrap();
        assert_eq!(json["tool_choice"], serde_json::json!("auto"));

        // none -> bare "none"
        let body: ChatBody<GLM4_6, TextMessage> = ChatBody::new(GLM4_6 {}, TextMessage::user("hi"))
            .with_tool_choice(crate::model::tools::ToolChoice::none());
        let json = serde_json::to_value(&body).unwrap();
        assert_eq!(json["tool_choice"], serde_json::json!("none"));

        // function -> object form
        let body: ChatBody<GLM4_6, TextMessage> = ChatBody::new(GLM4_6 {}, TextMessage::user("hi"))
            .with_tool_choice(crate::model::tools::ToolChoice::function("get_weather"));
        let json = serde_json::to_value(&body).unwrap();
        assert_eq!(
            json["tool_choice"],
            serde_json::json!({
                "type": "function",
                "function": { "name": "get_weather" }
            })
        );
    }

    #[test]
    fn test_with_response_format_serializes() {
        use crate::model::tools::ResponseFormat;
        let body: ChatBody<GLM4_6, TextMessage> = ChatBody::new(GLM4_6 {}, TextMessage::user("hi"))
            .with_response_format(ResponseFormat::JsonObject);
        assert_eq!(body.response_format, Some(ResponseFormat::JsonObject));
        let json = serde_json::to_value(&body).unwrap();
        assert_eq!(json["response_format"]["type"], "json_object");
    }
}