zai_rs/model/
tools.rs

1//! # Tool Definitions & Configurations
2//!
3//! Defines the tool types that can be attached to chat requests, including
4//! function calling, web search integration, retrieval tools, and the
5//! [`ThinkingType`] configuration.
6//!
7//! # Key Types
8//!
9//! - [`ThinkingType`] — Controls reasoning mode for thinking-capable models
10//! - [`ReasoningEffort`] — Controls reasoning depth for GLM-5.2+ models
11//! - `Function` — Defines a callable function with JSON-schema parameters
12//! - `WebSearch` — Enables live web search within chat
13//! - [`Retrieval`] — Enables knowledge-base retrieval
14//! - [`ToolChoice`] — Enables the frozen automatic tool-selection policy
15
16use std::collections::BTreeMap;
17
18use serde::Serialize;
19use validator::Validate;
20
21use super::model_validate::validate_json_schema_value;
22use crate::tool::web_search::{ContentSize, SearchEngine, SearchRecencyFilter};
23
24/// Controls extended reasoning and whether reasoning context is cleared between
25/// turns.
26///
27/// # Examples
28///
29/// ```
30/// use zai_rs::model::{
31///     chat::ChatCompletion,
32///     chat_message_types::TextMessage,
33///     chat_models::GLM5_2,
34///     tools::ThinkingType,
35/// };
36///
37/// let request = ChatCompletion::new(GLM5_2 {}, TextMessage::user("Solve this"))
38///     .with_thinking(ThinkingType::enabled());
39///
40/// let request = ChatCompletion::new(GLM5_2 {}, TextMessage::user("Continue"))
41///     .with_thinking(ThinkingType::enabled().with_clear_thinking(false));
42/// ```
43///
44/// # Model compatibility
45///
46/// Thinking capabilities are available only on models that implement the
47/// `ThinkEnable` trait, such as GLM-5.2, GLM-5.1, GLM-5, GLM-4.7, and GLM-4.5
48/// series models.
49#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
50pub struct ThinkingType {
51    /// Whether thinking is enabled or disabled.
52    #[serde(rename = "type")]
53    pub mode: ThinkingMode,
54
55    /// Whether to clear historical `reasoning_content`.
56    ///
57    /// - `true` (default for standard API): Clears reasoning content each turn.
58    /// - `false` (recommended for Coding / Agent): Preserves reasoning content
59    ///   across turns, enabling better context for multi-step tool calls.
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub clear_thinking: Option<bool>,
62}
63
64/// Thinking mode variants.
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
66#[serde(rename_all = "lowercase")]
67pub enum ThinkingMode {
68    /// Extended reasoning is enabled; exposure of reasoning text depends on the
69    /// selected model and endpoint.
70    Enabled,
71    /// Extended reasoning is disabled.
72    Disabled,
73}
74
75impl ThinkingType {
76    /// Create a new thinking configuration with enabled mode.
77    pub fn enabled() -> Self {
78        Self {
79            mode: ThinkingMode::Enabled,
80            clear_thinking: None,
81        }
82    }
83
84    /// Create a new thinking configuration with disabled mode.
85    pub fn disabled() -> Self {
86        Self {
87            mode: ThinkingMode::Disabled,
88            clear_thinking: None,
89        }
90    }
91
92    /// Set whether to clear historical reasoning content.
93    ///
94    /// Use `false` for Coding / Agent scenarios where reasoning content should
95    /// be preserved across turns.
96    pub fn with_clear_thinking(mut self, clear: bool) -> Self {
97        self.clear_thinking = Some(clear);
98        self
99    }
100}
101
102/// Reasoning depth level for the `reasoning_effort` parameter (GLM-5.2+).
103///
104/// Controls how much reasoning the model invests when thinking mode is
105/// enabled. Higher levels yield deeper reasoning at the cost of latency and
106/// token usage; lower levels are faster and cheaper. Available only on
107/// GLM-5.2 and above (models implementing
108/// [`ReasoningEffortEnable`](super::traits::ReasoningEffortEnable)).
109///
110/// Levels, from highest to lowest reasoning depth:
111///
112/// | Variant | Description |
113/// |---------|-------------|
114/// | [`Max`](Self::Max) | Maximum reasoning depth; recommended for coding / architecture-level tasks |
115/// | [`Xhigh`](Self::Xhigh) | Extra-high reasoning |
116/// | [`High`](Self::High) | High reasoning (default mapping in many clients) |
117/// | [`Medium`](Self::Medium) | Balanced reasoning |
118/// | [`Low`](Self::Low) | Light reasoning |
119/// | [`Minimal`](Self::Minimal) | Minimal reasoning |
120/// | [`None`](Self::None) | No extra reasoning beyond base behaviour |
121///
122/// ## Usage
123///
124/// ```
125/// use zai_rs::model::{
126///     chat::ChatCompletion,
127///     chat_message_types::TextMessage,
128///     chat_models::GLM5_2,
129///     tools::{ReasoningEffort, ThinkingType},
130/// };
131///
132/// let request = ChatCompletion::new(GLM5_2 {}, TextMessage::user("Design an API"))
133///     .with_thinking(ThinkingType::enabled())
134///     .with_reasoning_effort(ReasoningEffort::Max);
135/// ```
136#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
137pub enum ReasoningEffort {
138    /// Maximum reasoning depth. Recommended for coding and architecture-level
139    /// tasks where correctness matters most.
140    #[serde(rename = "max")]
141    Max,
142    /// Extra-high reasoning depth.
143    #[serde(rename = "xhigh")]
144    Xhigh,
145    /// High reasoning depth.
146    #[serde(rename = "high")]
147    High,
148    /// Balanced reasoning depth.
149    #[serde(rename = "medium")]
150    Medium,
151    /// Light reasoning depth.
152    #[serde(rename = "low")]
153    Low,
154    /// Minimal reasoning depth.
155    #[serde(rename = "minimal")]
156    Minimal,
157    /// No extra reasoning beyond base behaviour.
158    #[serde(rename = "none")]
159    None,
160}
161
162/// External capability attached to a chat request.
163///
164/// Variants cover caller-defined functions, knowledge retrieval, web search,
165/// and Model Context Protocol servers.
166///
167/// # Examples
168///
169/// ```
170/// use zai_rs::model::tools::{Function, Tools, WebSearch};
171/// use zai_rs::tool::web_search::SearchEngine;
172///
173/// let function_tool = Tools::Function {
174///     function: Function::new(
175///         "get_weather",
176///         "Get weather data",
177///         serde_json::json!({"type": "object"}),
178///     ),
179/// };
180///
181/// let search_tool = Tools::WebSearch {
182///     web_search: WebSearch::new(SearchEngine::SearchPro)
183///         .with_enable(true)
184///         .with_count(10)
185/// };
186/// ```
187#[derive(Debug, Clone, Serialize)]
188#[serde(tag = "type")]
189#[serde(rename_all = "snake_case")]
190pub enum Tools {
191    /// Custom function calling tool with parameters.
192    ///
193    /// Allows the AI to invoke user-defined functions with structured
194    /// arguments. Functions must be pre-defined with JSON schemas for
195    /// parameter validation.
196    Function {
197        /// The function definition (name, description, parameter schema).
198        function: Function,
199    },
200
201    /// Knowledge retrieval system access tools.
202    ///
203    /// Provides access to knowledge bases, document collections, or other
204    /// structured information sources that the AI can query.
205    Retrieval {
206        /// The retrieval-tool descriptor.
207        retrieval: Retrieval,
208    },
209
210    /// Web search capabilities for internet access.
211    ///
212    /// Enables the AI to perform web searches and access current information
213    /// from the internet. Supports various search engines and configurations.
214    WebSearch {
215        /// The web-search-tool descriptor.
216        web_search: WebSearch,
217    },
218
219    /// A provider-hosted MCP tool attached to a chat request.
220    ///
221    /// This config asks the chat service to connect to an MCP server. It is
222    /// distinct from [`crate::mcp::McpClient`], which connects from this SDK.
223    #[serde(rename = "mcp")]
224    MCP {
225        /// The MCP-tool descriptor.
226        mcp: MCP,
227    },
228}
229
230impl Validate for Tools {
231    fn validate(&self) -> Result<(), validator::ValidationErrors> {
232        match self {
233            Self::Function { function } => function.validate(),
234            Self::Retrieval { retrieval } => retrieval.validate(),
235            Self::WebSearch { web_search } => web_search.validate(),
236            Self::MCP { mcp } => mcp.validate(),
237        }
238    }
239}
240
241impl From<Function> for Tools {
242    fn from(function: Function) -> Self {
243        Self::Function { function }
244    }
245}
246
247/// Definition of a caller-provided function that the model may invoke.
248///
249/// # Validation
250///
251/// * `name` - Must be 1 to 64 ASCII letters, digits, underscores, or hyphens
252/// * `parameters` - Must be a valid JSON schema
253#[derive(Clone, Serialize, Validate)]
254#[validate(schema(function = "validate_function"))]
255pub struct Function {
256    /// The name of the function. Must match `[A-Za-z0-9_-]{1,64}`.
257    #[validate(length(min = 1, max = 64))]
258    pub name: String,
259
260    /// A description of what the function does.
261    pub description: String,
262
263    /// JSON Schema object describing the function parameters.
264    #[serde(skip_serializing_if = "Option::is_none")]
265    #[validate(custom(function = "validate_json_schema_value"))]
266    pub parameters: Option<serde_json::Value>,
267}
268
269impl std::fmt::Debug for Function {
270    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
271        formatter
272            .debug_struct("Function")
273            .field("name", &self.name)
274            .field("description", &"[REDACTED]")
275            .field("parameters_configured", &self.parameters.is_some())
276            .finish()
277    }
278}
279
280fn validate_function(function: &Function) -> Result<(), validator::ValidationError> {
281    if function.name.trim().is_empty() {
282        return Err(validator::ValidationError::new("name_must_not_be_blank"));
283    }
284    if !function
285        .name
286        .bytes()
287        .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
288    {
289        return Err(validator::ValidationError::new("invalid_function_name"));
290    }
291    if function
292        .parameters
293        .as_ref()
294        .is_some_and(|parameters| !parameters.is_object())
295    {
296        return Err(validator::ValidationError::new(
297            "function_parameters_must_be_an_object_schema",
298        ));
299    }
300    Ok(())
301}
302
303impl Function {
304    /// Create a function definition with a JSON Schema parameter object.
305    ///
306    /// # Examples
307    ///
308    /// ```
309    /// use zai_rs::model::tools::Function;
310    ///
311    /// let func = Function::new(
312    ///     "get_weather",
313    ///     "Get current weather for a location",
314    ///     serde_json::json!({
315    ///         "type": "object",
316    ///         "properties": { "location": { "type": "string" } }
317    ///     })
318    /// );
319    /// ```
320    pub fn new(
321        name: impl Into<String>,
322        description: impl Into<String>,
323        parameters: serde_json::Value,
324    ) -> Self {
325        Self {
326            name: name.into(),
327            description: description.into(),
328            parameters: Some(parameters),
329        }
330    }
331}
332
333/// Knowledge-base retrieval tool configuration.
334///
335/// Attaches a Zhipu AI knowledge base to a chat completion so the model can
336/// retrieve relevant passages from it before answering. Create the knowledge
337/// base (and obtain its `knowledge_id`) in the BigModel console, then pass the
338/// id here via the [`Retrieval`] tool.
339///
340/// See the official guide:
341/// <https://docs.bigmodel.cn/cn/guide/tools/knowledge/retrieval>.
342///
343/// # Wire form
344///
345/// Serializes as `{"knowledge_id":"…","prompt_template":"…"}` (the
346/// `prompt_template` field is omitted when not set).
347///
348/// # Usage
349///
350/// ```rust
351/// use zai_rs::model::tools::{Retrieval, Tools};
352///
353/// let tool = Tools::Retrieval {
354///     retrieval: Retrieval::new("kb_1234567890"),
355/// };
356/// // Or attach a custom prompt template:
357/// let tool = Tools::Retrieval {
358///     retrieval: Retrieval::new("kb_1234567890")
359///         .with_prompt_template("仅依据知识库回答：{knowledge}"),
360/// };
361/// ```
362#[derive(Clone, Serialize, Validate)]
363#[validate(schema(function = "validate_retrieval"))]
364pub struct Retrieval {
365    /// Knowledge-base id (required). Obtain it from the BigModel console after
366    /// creating and populating a knowledge base.
367    #[validate(length(min = 1))]
368    pub knowledge_id: String,
369    /// Optional prompt template applied when the model consumes retrieved
370    /// knowledge. Serialized as `prompt_template`; omitted when `None`.
371    #[serde(skip_serializing_if = "Option::is_none")]
372    #[validate(length(min = 1))]
373    pub prompt_template: Option<String>,
374}
375
376impl std::fmt::Debug for Retrieval {
377    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
378        formatter
379            .debug_struct("Retrieval")
380            .field("knowledge_id", &"[REDACTED]")
381            .field(
382                "prompt_template_configured",
383                &self.prompt_template.is_some(),
384            )
385            .finish()
386    }
387}
388
389fn validate_retrieval(retrieval: &Retrieval) -> Result<(), validator::ValidationError> {
390    if retrieval.knowledge_id.trim().is_empty() {
391        return Err(validator::ValidationError::new(
392            "knowledge_id_must_not_be_blank",
393        ));
394    }
395    if retrieval
396        .prompt_template
397        .as_deref()
398        .is_some_and(|template| template.trim().is_empty())
399    {
400        return Err(validator::ValidationError::new(
401            "prompt_template_must_not_be_blank",
402        ));
403    }
404    Ok(())
405}
406
407impl Retrieval {
408    /// Create a retrieval tool bound to `knowledge_id`.
409    ///
410    pub fn new(knowledge_id: impl Into<String>) -> Self {
411        Self {
412            knowledge_id: knowledge_id.into(),
413            prompt_template: None,
414        }
415    }
416
417    /// Attach a custom prompt template to the retrieval tool.
418    pub fn with_prompt_template(mut self, prompt_template: impl Into<String>) -> Self {
419        self.prompt_template = Some(prompt_template.into());
420        self
421    }
422}
423
424/// The order in which search results are returned.
425#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
426#[serde(rename_all = "snake_case")]
427pub enum ResultSequence {
428    /// Search results appear before the model's answer.
429    Before,
430    /// Search results appear after the model's answer.
431    After,
432}
433
434/// Web-search configuration attached to a chat request.
435#[derive(Clone, Serialize, Validate)]
436#[validate(schema(function = "validate_web_search"))]
437pub struct WebSearch {
438    /// Search engine type (required). Supported: search_std, search_pro,
439    /// search_pro_sogou, search_pro_quark.
440    pub search_engine: SearchEngine,
441
442    /// Whether to enable web search. Default is false.
443    #[serde(skip_serializing_if = "Option::is_none")]
444    pub enable: Option<bool>,
445
446    /// Force-triggered search query string.
447    #[serde(skip_serializing_if = "Option::is_none")]
448    pub search_query: Option<String>,
449
450    /// Whether to perform search intent detection. true: execute only when
451    /// intent is detected; false: skip detection and search directly.
452    #[serde(skip_serializing_if = "Option::is_none")]
453    pub search_intent: Option<bool>,
454
455    /// Number of results to return (1-50).
456    #[serde(skip_serializing_if = "Option::is_none")]
457    #[validate(range(min = 1, max = 50))]
458    pub count: Option<u32>,
459
460    /// Whitelist domain filter, e.g., "www.example.com".
461    #[serde(skip_serializing_if = "Option::is_none")]
462    pub search_domain_filter: Option<String>,
463
464    /// Time range filter.
465    #[serde(skip_serializing_if = "Option::is_none")]
466    pub search_recency_filter: Option<SearchRecencyFilter>,
467
468    /// Snippet summary size: medium or high.
469    #[serde(skip_serializing_if = "Option::is_none")]
470    pub content_size: Option<ContentSize>,
471
472    /// Return sequence for search results: before or after.
473    #[serde(skip_serializing_if = "Option::is_none")]
474    pub result_sequence: Option<ResultSequence>,
475
476    /// Whether to include detailed search source information.
477    #[serde(skip_serializing_if = "Option::is_none")]
478    pub search_result: Option<bool>,
479
480    /// Whether an answer requires search results to be returned.
481    #[serde(skip_serializing_if = "Option::is_none")]
482    pub require_search: Option<bool>,
483
484    /// Custom prompt to post-process search results.
485    #[serde(skip_serializing_if = "Option::is_none")]
486    pub search_prompt: Option<String>,
487}
488
489impl std::fmt::Debug for WebSearch {
490    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
491        formatter
492            .debug_struct("WebSearch")
493            .field("search_engine", &self.search_engine)
494            .field("enable", &self.enable)
495            .field("search_query_configured", &self.search_query.is_some())
496            .field("search_intent", &self.search_intent)
497            .field("count", &self.count)
498            .field(
499                "search_domain_filter_configured",
500                &self.search_domain_filter.is_some(),
501            )
502            .field("search_recency_filter", &self.search_recency_filter)
503            .field("content_size", &self.content_size)
504            .field("result_sequence", &self.result_sequence)
505            .field("search_result", &self.search_result)
506            .field("require_search", &self.require_search)
507            .field("search_prompt_configured", &self.search_prompt.is_some())
508            .finish()
509    }
510}
511
512fn validate_web_search(web_search: &WebSearch) -> Result<(), validator::ValidationError> {
513    if [
514        web_search.search_query.as_deref(),
515        web_search.search_domain_filter.as_deref(),
516        web_search.search_prompt.as_deref(),
517    ]
518    .into_iter()
519    .flatten()
520    .any(|value| value.trim().is_empty())
521    {
522        Err(validator::ValidationError::new(
523            "search_text_must_not_be_blank",
524        ))
525    } else {
526        Ok(())
527    }
528}
529
530impl WebSearch {
531    /// Create a WebSearch config with the required search engine; other fields
532    /// are optional.
533    pub fn new(search_engine: SearchEngine) -> Self {
534        Self {
535            search_engine,
536            enable: None,
537            search_query: None,
538            search_intent: None,
539            count: None,
540            search_domain_filter: None,
541            search_recency_filter: None,
542            content_size: None,
543            result_sequence: None,
544            search_result: None,
545            require_search: None,
546            search_prompt: None,
547        }
548    }
549
550    /// Enable or disable web search.
551    pub fn with_enable(mut self, enable: bool) -> Self {
552        self.enable = Some(enable);
553        self
554    }
555    /// Set a forced search query.
556    pub fn with_search_query(mut self, query: impl Into<String>) -> Self {
557        self.search_query = Some(query.into());
558        self
559    }
560    /// Set search intent detection behavior.
561    pub fn with_search_intent(mut self, search_intent: bool) -> Self {
562        self.search_intent = Some(search_intent);
563        self
564    }
565    /// Set results count (1-50).
566    pub fn with_count(mut self, count: u32) -> Self {
567        self.count = Some(count);
568        self
569    }
570    /// Restrict to a whitelist domain.
571    pub fn with_search_domain_filter(mut self, domain: impl Into<String>) -> Self {
572        self.search_domain_filter = Some(domain.into());
573        self
574    }
575    /// Set time range filter.
576    pub fn with_search_recency_filter(mut self, filter: SearchRecencyFilter) -> Self {
577        self.search_recency_filter = Some(filter);
578        self
579    }
580    /// Set content size.
581    pub fn with_content_size(mut self, size: ContentSize) -> Self {
582        self.content_size = Some(size);
583        self
584    }
585    /// Set result sequence.
586    pub fn with_result_sequence(mut self, seq: ResultSequence) -> Self {
587        self.result_sequence = Some(seq);
588        self
589    }
590    /// Toggle returning detailed search source info.
591    pub fn with_search_result(mut self, enable: bool) -> Self {
592        self.search_result = Some(enable);
593        self
594    }
595    /// Require search results for answering.
596    pub fn with_require_search(mut self, require: bool) -> Self {
597        self.require_search = Some(require);
598        self
599    }
600    /// Set a custom prompt to post-process search results.
601    pub fn with_search_prompt(mut self, prompt: impl Into<String>) -> Self {
602        self.search_prompt = Some(prompt.into());
603        self
604    }
605}
606/// Provider-side MCP connection configuration embedded in a chat request.
607///
608/// This does not create a local MCP connection. For Zhipu-hosted MCP servers,
609/// put the MCP code in `server_label` and leave `server_url` unset. To call MCP
610/// tools directly from this SDK, use [`crate::mcp::McpClient`].
611#[derive(Clone, Serialize, Validate)]
612#[validate(schema(function = "validate_mcp"))]
613pub struct MCP {
614    /// MCP server identifier (required). If connecting to Zhipu MCP via code,
615    /// put the code here.
616    #[validate(length(min = 1))]
617    pub server_label: String,
618
619    /// MCP server URL.
620    #[serde(skip_serializing_if = "Option::is_none")]
621    #[validate(url)]
622    pub server_url: Option<String>,
623
624    /// Transport type. Default: streamable-http.
625    #[serde(skip_serializing_if = "Option::is_none")]
626    pub transport_type: Option<MCPTransportType>,
627
628    /// Allowed tool names.
629    #[serde(skip_serializing_if = "Vec::is_empty")]
630    pub allowed_tools: Vec<String>,
631
632    /// Authentication headers required by the MCP server.
633    #[serde(skip_serializing_if = "Option::is_none")]
634    pub headers: Option<BTreeMap<String, String>>,
635}
636
637fn validate_mcp(mcp: &MCP) -> Result<(), validator::ValidationError> {
638    if mcp.server_label.trim().is_empty() {
639        return Err(validator::ValidationError::new(
640            "server_label_must_not_be_blank",
641        ));
642    }
643    if let Some(server_url) = mcp.server_url.as_deref() {
644        let Ok(url) = server_url.parse::<url::Url>() else {
645            return Err(validator::ValidationError::new("invalid_server_url"));
646        };
647        if !matches!(url.scheme(), "http" | "https")
648            || !url.username().is_empty()
649            || url.password().is_some()
650            || url.fragment().is_some()
651        {
652            return Err(validator::ValidationError::new("invalid_server_url"));
653        }
654    }
655    if mcp.allowed_tools.iter().any(|tool| tool.trim().is_empty()) {
656        return Err(validator::ValidationError::new(
657            "allowed_tool_must_not_be_blank",
658        ));
659    }
660    if mcp.headers.as_ref().is_some_and(|headers| {
661        headers.iter().any(|(name, value)| {
662            !valid_forwarded_header_name(name) || !valid_forwarded_header_value(value)
663        })
664    }) {
665        return Err(validator::ValidationError::new("invalid_forwarded_header"));
666    }
667    Ok(())
668}
669
670fn valid_forwarded_header_name(name: &str) -> bool {
671    !name.is_empty()
672        && name.bytes().all(|byte| {
673            byte.is_ascii_alphanumeric()
674                || matches!(
675                    byte,
676                    b'!' | b'#'
677                        | b'$'
678                        | b'%'
679                        | b'&'
680                        | b'\''
681                        | b'*'
682                        | b'+'
683                        | b'-'
684                        | b'.'
685                        | b'^'
686                        | b'_'
687                        | b'`'
688                        | b'|'
689                        | b'~'
690                )
691        })
692}
693
694fn valid_forwarded_header_value(value: &str) -> bool {
695    !value.trim().is_empty()
696        && value
697            .chars()
698            .all(|character| character == '\t' || !character.is_control())
699}
700
701impl std::fmt::Debug for MCP {
702    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
703        formatter
704            .debug_struct("MCP")
705            .field("server_label", &"[REDACTED]")
706            .field("server_url_configured", &self.server_url.is_some())
707            .field("transport_type", &self.transport_type)
708            .field("allowed_tool_count", &self.allowed_tools.len())
709            .field("header_count", &self.headers.as_ref().map(BTreeMap::len))
710            .finish()
711    }
712}
713
714impl MCP {
715    /// Create a new MCP config with required server_label and default transport
716    /// type.
717    pub fn new(server_label: impl Into<String>) -> Self {
718        Self {
719            server_label: server_label.into(),
720            server_url: None,
721            transport_type: Some(MCPTransportType::StreamableHttp),
722            allowed_tools: Vec::new(),
723            headers: None,
724        }
725    }
726
727    /// Set the MCP server URL.
728    pub fn with_server_url(mut self, url: impl Into<String>) -> Self {
729        self.server_url = Some(url.into());
730        self
731    }
732    /// Set the MCP transport type.
733    pub fn with_transport_type(mut self, transport: MCPTransportType) -> Self {
734        self.transport_type = Some(transport);
735        self
736    }
737    /// Replace the allowed tool list.
738    pub fn with_allowed_tools(mut self, tools: impl Into<Vec<String>>) -> Self {
739        self.allowed_tools = tools.into();
740        self
741    }
742    /// Add a single allowed tool.
743    pub fn add_allowed_tool(mut self, tool: impl Into<String>) -> Self {
744        self.allowed_tools.push(tool.into());
745        self
746    }
747    /// Set authentication headers map.
748    pub fn with_headers(mut self, headers: BTreeMap<String, String>) -> Self {
749        self.headers = Some(headers);
750        self
751    }
752    /// Add or update a single header entry.
753    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
754        self.headers
755            .get_or_insert_with(BTreeMap::new)
756            .insert(key.into(), value.into());
757        self
758    }
759}
760
761/// Allowed MCP transport types.
762#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
763#[serde(rename_all = "kebab-case")]
764pub enum MCPTransportType {
765    /// Server-Sent Events transport.
766    Sse,
767    /// Streamable HTTP transport.
768    StreamableHttp,
769}
770
771/// Requested response representation.
772#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
773#[serde(rename_all = "snake_case")]
774#[serde(tag = "type")]
775pub enum ResponseFormat {
776    /// Plain text response format.
777    Text,
778    /// Structured JSON object response format.
779    JsonObject,
780}
781
782/// Controls how the model selects tools during a chat completion.
783///
784/// This is the value carried by the `tool_choice` request parameter. It is only
785/// meaningful when [`Tools`] are also attached to the request. The frozen API
786/// schema accepts only the bare string `"auto"`; the older `"none"` and forced
787/// function object forms were never part of this operation's contract.
788///
789/// # Usage
790///
791/// ```rust
792/// use zai_rs::model::tools::ToolChoice;
793///
794/// // Let the model decide (default behaviour):
795/// let choice = ToolChoice::auto();
796/// ```
797#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
798#[serde(rename_all = "lowercase")]
799pub enum ToolChoice {
800    /// The model decides whether to call a tool (`"auto"` on the wire).
801    Auto,
802}
803
804impl ToolChoice {
805    /// Let the model decide whether to call a tool (`"auto"`).
806    pub const fn auto() -> Self {
807        Self::Auto
808    }
809}
810
811#[cfg(test)]
812mod tests {
813    use super::*;
814
815    // ThinkingType tests
816    #[test]
817    fn test_thinking_type_enabled_serialization() {
818        let thinking = ThinkingType::enabled();
819        let json = serde_json::to_string(&thinking).unwrap();
820        assert!(json.contains("\"type\":\"enabled\""));
821        assert!(!json.contains("clear_thinking"));
822    }
823
824    #[test]
825    fn test_thinking_type_disabled_serialization() {
826        let thinking = ThinkingType::disabled();
827        let json = serde_json::to_string(&thinking).unwrap();
828        assert!(json.contains("\"type\":\"disabled\""));
829        assert!(!json.contains("clear_thinking"));
830    }
831
832    #[test]
833    fn test_thinking_type_with_clear_thinking_serialization() {
834        let thinking = ThinkingType::enabled().with_clear_thinking(false);
835        let json = serde_json::to_string(&thinking).unwrap();
836        assert!(json.contains("\"type\":\"enabled\""));
837        assert!(json.contains("\"clear_thinking\":false"));
838    }
839
840    #[test]
841    fn test_thinking_type_disabled_with_clear_thinking() {
842        let thinking = ThinkingType::disabled().with_clear_thinking(true);
843        let json = serde_json::to_string(&thinking).unwrap();
844        assert!(json.contains("\"type\":\"disabled\""));
845        assert!(json.contains("\"clear_thinking\":true"));
846    }
847
848    // Function tests
849    #[test]
850    fn test_function_new() {
851        let params = serde_json::json!({
852            "type": "object",
853            "properties": {
854                "name": {"type": "string"}
855            }
856        });
857        let func = Function::new("test_func", "A test function", params);
858
859        assert_eq!(func.name, "test_func");
860        assert_eq!(func.description, "A test function");
861        assert!(func.parameters.is_some());
862    }
863
864    #[test]
865    fn test_function_serialization() {
866        let params = serde_json::json!({
867            "type": "object",
868            "properties": {
869                "value": {"type": "number"}
870            }
871        });
872        let func = Function::new("test_func", "A test function", params);
873        let json = serde_json::to_string(&func).unwrap();
874
875        assert!(json.contains("\"name\":\"test_func\""));
876        assert!(json.contains("\"description\":\"A test function\""));
877        assert!(json.contains("\"properties\""));
878    }
879
880    #[test]
881    fn test_function_validation() {
882        let params = serde_json::json!({
883            "type": "object",
884            "properties": {}
885        });
886        let func = Function::new("valid_name", "Description", params.clone());
887
888        // Name length validation: 1-64 characters
889        assert!(func.validate().is_ok());
890        assert!(
891            Function::new("valid-name", "Description", params.clone())
892                .validate()
893                .is_ok()
894        );
895
896        let invalid_name = Function::new("", "Description", params.clone());
897        assert!(invalid_name.validate().is_err());
898        let blank_name = Function::new("   ", "Description", params.clone());
899        assert!(blank_name.validate().is_err());
900        let punctuation = Function::new("invalid!", "Description", params.clone());
901        assert!(punctuation.validate().is_err());
902
903        let long_name = Function::new("a".repeat(65), "Description", params);
904        assert!(long_name.validate().is_err());
905    }
906
907    // Retrieval tests
908    #[test]
909    fn test_retrieval_new() {
910        let retrieval = Retrieval::new("kb_123").with_prompt_template("template");
911        assert_eq!(retrieval.knowledge_id, "kb_123");
912        assert_eq!(retrieval.prompt_template, Some("template".to_string()));
913    }
914
915    #[test]
916    fn test_retrieval_new_without_template() {
917        let retrieval = Retrieval::new("kb_456");
918        assert_eq!(retrieval.knowledge_id, "kb_456");
919        assert!(retrieval.prompt_template.is_none());
920    }
921
922    #[test]
923    fn test_retrieval_serialization() {
924        let retrieval = Retrieval::new("kb_789");
925        let json = serde_json::to_string(&retrieval).unwrap();
926        assert!(json.contains("\"knowledge_id\":\"kb_789\""));
927        // prompt_template should be omitted when None
928        assert!(!json.contains("prompt_template"));
929    }
930
931    #[test]
932    fn retrieval_prompt_template_serializes() {
933        let retrieval = Retrieval::new("kb_builder").with_prompt_template("ctx: {knowledge}");
934        let json = serde_json::to_value(&retrieval).unwrap();
935        assert_eq!(json["knowledge_id"], "kb_builder");
936        assert_eq!(json["prompt_template"], "ctx: {knowledge}");
937    }
938
939    #[test]
940    fn tool_validation_rejects_blank_optional_values() {
941        assert!(Retrieval::new(" ").validate().is_err());
942        assert!(
943            WebSearch::new(SearchEngine::SearchPro)
944                .with_search_query(" ")
945                .validate()
946                .is_err()
947        );
948        assert!(MCP::new(" ").validate().is_err());
949        assert!(
950            MCP::new("server")
951                .with_server_url("ftp://example.com")
952                .validate()
953                .is_err()
954        );
955        assert!(
956            MCP::new("server")
957                .with_header("Authorization\r\nInjected", "secret")
958                .validate()
959                .is_err()
960        );
961    }
962
963    // WebSearch tests
964    #[test]
965    fn test_web_search_new() {
966        let web_search = WebSearch::new(SearchEngine::SearchPro);
967        assert_eq!(web_search.search_engine, SearchEngine::SearchPro);
968        assert!(web_search.enable.is_none());
969    }
970
971    #[test]
972    fn test_web_search_with_enable() {
973        let web_search = WebSearch::new(SearchEngine::SearchPro).with_enable(true);
974        assert_eq!(web_search.enable, Some(true));
975    }
976
977    #[test]
978    fn test_web_search_with_search_query() {
979        let web_search = WebSearch::new(SearchEngine::SearchPro).with_search_query("test query");
980        assert_eq!(web_search.search_query, Some("test query".to_string()));
981    }
982
983    #[test]
984    fn test_web_search_with_search_intent() {
985        let web_search = WebSearch::new(SearchEngine::SearchPro).with_search_intent(true);
986        assert_eq!(web_search.search_intent, Some(true));
987    }
988
989    #[test]
990    fn test_web_search_with_count() {
991        let web_search = WebSearch::new(SearchEngine::SearchPro).with_count(10);
992        assert_eq!(web_search.count, Some(10));
993    }
994
995    #[test]
996    fn test_web_search_with_search_domain_filter() {
997        let web_search =
998            WebSearch::new(SearchEngine::SearchPro).with_search_domain_filter("example.com");
999        assert_eq!(
1000            web_search.search_domain_filter,
1001            Some("example.com".to_string())
1002        );
1003    }
1004
1005    #[test]
1006    fn test_web_search_with_search_recency_filter() {
1007        let filter = SearchRecencyFilter::OneDay;
1008        let web_search = WebSearch::new(SearchEngine::SearchPro).with_search_recency_filter(filter);
1009        assert_eq!(web_search.search_recency_filter, Some(filter));
1010    }
1011
1012    #[test]
1013    fn test_web_search_with_content_size() {
1014        let size = ContentSize::Medium;
1015        let web_search = WebSearch::new(SearchEngine::SearchPro).with_content_size(size);
1016        assert_eq!(web_search.content_size, Some(size));
1017    }
1018
1019    #[test]
1020    fn test_web_search_with_result_sequence() {
1021        let seq = ResultSequence::After;
1022        let web_search = WebSearch::new(SearchEngine::SearchPro).with_result_sequence(seq);
1023        assert_eq!(web_search.result_sequence, Some(seq));
1024    }
1025
1026    #[test]
1027    fn test_web_search_with_search_result() {
1028        let web_search = WebSearch::new(SearchEngine::SearchPro).with_search_result(true);
1029        assert_eq!(web_search.search_result, Some(true));
1030    }
1031
1032    #[test]
1033    fn test_web_search_with_require_search() {
1034        let web_search = WebSearch::new(SearchEngine::SearchPro).with_require_search(true);
1035        assert_eq!(web_search.require_search, Some(true));
1036    }
1037
1038    #[test]
1039    fn test_web_search_with_search_prompt() {
1040        let web_search =
1041            WebSearch::new(SearchEngine::SearchPro).with_search_prompt("custom prompt");
1042        assert_eq!(web_search.search_prompt, Some("custom prompt".to_string()));
1043    }
1044
1045    #[test]
1046    fn test_web_search_serialization() {
1047        let web_search = WebSearch::new(SearchEngine::SearchPro)
1048            .with_enable(true)
1049            .with_count(5);
1050        let json = serde_json::to_string(&web_search).unwrap();
1051        assert!(json.contains("\"search_engine\""));
1052        assert!(json.contains("\"enable\":true"));
1053        assert!(json.contains("\"count\":5"));
1054    }
1055
1056    // MCP tests
1057    #[test]
1058    fn test_mcp_new() {
1059        let mcp = MCP::new("server_label");
1060        assert_eq!(mcp.server_label, "server_label");
1061        assert_eq!(mcp.transport_type, Some(MCPTransportType::StreamableHttp));
1062        assert!(mcp.allowed_tools.is_empty());
1063    }
1064
1065    #[test]
1066    fn test_mcp_with_server_url() {
1067        let mcp = MCP::new("server_label").with_server_url("https://example.com");
1068        assert_eq!(mcp.server_url, Some("https://example.com".to_string()));
1069    }
1070
1071    #[test]
1072    fn test_mcp_with_transport_type() {
1073        let mcp = MCP::new("server_label").with_transport_type(MCPTransportType::Sse);
1074        assert_eq!(mcp.transport_type, Some(MCPTransportType::Sse));
1075    }
1076
1077    #[test]
1078    fn test_mcp_with_allowed_tools() {
1079        let mcp = MCP::new("server_label")
1080            .with_allowed_tools(vec!["tool1".to_string(), "tool2".to_string()]);
1081        assert_eq!(mcp.allowed_tools.len(), 2);
1082        assert!(mcp.allowed_tools.contains(&"tool1".to_string()));
1083    }
1084
1085    #[test]
1086    fn test_mcp_add_allowed_tool() {
1087        let mcp = MCP::new("server_label")
1088            .add_allowed_tool("tool1")
1089            .add_allowed_tool("tool2");
1090        assert_eq!(mcp.allowed_tools.len(), 2);
1091    }
1092
1093    #[test]
1094    fn test_mcp_with_headers() {
1095        let mut headers = BTreeMap::new();
1096        headers.insert("Authorization".to_string(), "Bearer token".to_string());
1097        let mcp = MCP::new("server_label").with_headers(headers.clone());
1098        assert_eq!(mcp.headers, Some(headers));
1099    }
1100
1101    #[test]
1102    fn test_mcp_with_header() {
1103        let mcp = MCP::new("server_label").with_header("Authorization", "Bearer token");
1104        let debug = format!("{mcp:?}");
1105        assert!(!debug.contains("Bearer token"));
1106        assert!(!debug.contains("Authorization"));
1107        assert!(debug.contains("header_count"));
1108        let headers = mcp.headers.unwrap();
1109        assert_eq!(
1110            headers.get("Authorization"),
1111            Some(&"Bearer token".to_string())
1112        );
1113    }
1114
1115    #[test]
1116    fn test_mcp_serialization() {
1117        let mcp = MCP::new("server_label")
1118            .with_server_url("https://example.com")
1119            .with_transport_type(MCPTransportType::Sse);
1120        let json = serde_json::to_string(&mcp).unwrap();
1121        assert!(json.contains("\"server_label\":\"server_label\""));
1122        assert!(json.contains("\"server_url\":\"https://example.com\""));
1123        assert!(json.contains("\"transport_type\":\"sse\""));
1124        // allowed_tools should be omitted when empty
1125        assert!(!json.contains("allowed_tools"));
1126    }
1127
1128    #[test]
1129    fn mcp_validation_rejects_credentialed_or_fragmented_urls() {
1130        assert!(
1131            MCP::new("server")
1132                .with_server_url("https://user:secret@example.com/mcp")
1133                .validate()
1134                .is_err()
1135        );
1136        assert!(
1137            MCP::new("server")
1138                .with_server_url("https://example.com/mcp#secret")
1139                .validate()
1140                .is_err()
1141        );
1142    }
1143
1144    #[test]
1145    fn tool_debug_output_redacts_caller_content_and_remote_configuration() {
1146        let function = Function::new(
1147            "lookup",
1148            "private-description",
1149            serde_json::json!({"private-schema-key": {"type": "string"}}),
1150        );
1151        let retrieval =
1152            Retrieval::new("private-knowledge").with_prompt_template("private-template");
1153        let search = WebSearch::new(SearchEngine::SearchPro)
1154            .with_search_query("private-query")
1155            .with_search_domain_filter("private.example")
1156            .with_search_prompt("private-search-prompt");
1157        let mcp = MCP::new("private-server")
1158            .with_server_url("https://private.example/mcp?token=private-token")
1159            .add_allowed_tool("private-tool")
1160            .with_header("Authorization", "private-header");
1161        let debug = format!("{function:?} {retrieval:?} {search:?} {mcp:?}");
1162        for secret in [
1163            "private-description",
1164            "private-schema-key",
1165            "private-knowledge",
1166            "private-template",
1167            "private-query",
1168            "private.example",
1169            "private-search-prompt",
1170            "private-server",
1171            "private-token",
1172            "private-tool",
1173            "Authorization",
1174            "private-header",
1175        ] {
1176            assert!(!debug.contains(secret), "Debug leaked {secret}");
1177        }
1178    }
1179
1180    // MCPTransportType tests
1181    #[test]
1182    fn test_mcp_transport_type_sse_serialization() {
1183        let transport = MCPTransportType::Sse;
1184        let json = serde_json::to_string(&transport).unwrap();
1185        assert!(json.contains("\"sse\""));
1186    }
1187
1188    #[test]
1189    fn test_mcp_transport_type_streamable_http_serialization() {
1190        let transport = MCPTransportType::StreamableHttp;
1191        let json = serde_json::to_string(&transport).unwrap();
1192        assert!(json.contains("\"streamable-http\""));
1193    }
1194
1195    // ResponseFormat tests
1196    #[test]
1197    fn test_response_format_text_serialization() {
1198        let format = ResponseFormat::Text;
1199        let json = serde_json::to_string(&format).unwrap();
1200        assert!(json.contains("\"type\":\"text\""));
1201    }
1202
1203    #[test]
1204    fn test_response_format_json_object_serialization() {
1205        let format = ResponseFormat::JsonObject;
1206        let json = serde_json::to_string(&format).unwrap();
1207        assert!(json.contains("\"type\":\"json_object\""));
1208    }
1209
1210    // ToolChoice tests
1211    #[test]
1212    fn test_tool_choice_auto_serializes_as_bare_string() {
1213        let json = serde_json::to_value(ToolChoice::auto()).unwrap();
1214        assert_eq!(json, serde_json::json!("auto"));
1215    }
1216
1217    // Tools enum tests
1218    #[test]
1219    fn test_tools_function_serialization() {
1220        let func = Function::new("test_func", "test", serde_json::json!({}));
1221        let tools = Tools::Function { function: func };
1222        let json = serde_json::to_string(&tools).unwrap();
1223        assert!(json.contains("\"type\":\"function\""));
1224        assert!(json.contains("\"name\":\"test_func\""));
1225    }
1226
1227    #[test]
1228    fn test_tools_retrieval_serialization() {
1229        let retrieval = Retrieval::new("kb_123");
1230        let tools = Tools::Retrieval { retrieval };
1231        let json = serde_json::to_string(&tools).unwrap();
1232        assert!(json.contains("\"type\":\"retrieval\""));
1233        assert!(json.contains("\"knowledge_id\":\"kb_123\""));
1234    }
1235
1236    #[test]
1237    fn test_tools_web_search_serialization() {
1238        let web_search = WebSearch::new(SearchEngine::SearchPro);
1239        let tools = Tools::WebSearch { web_search };
1240        let json = serde_json::to_string(&tools).unwrap();
1241        assert!(json.contains("\"type\":\"web_search\""));
1242        assert!(json.contains("\"search_engine\""));
1243    }
1244
1245    #[test]
1246    fn test_tools_mcp_serialization() {
1247        let mcp = MCP::new("server_label");
1248        let tools = Tools::MCP { mcp };
1249        let json = serde_json::to_string(&tools).unwrap();
1250        assert!(json.contains("\"type\":\"mcp\""));
1251        assert!(json.contains("\"server_label\":\"server_label\""));
1252    }
1253
1254    // ResultSequence tests
1255    #[test]
1256    fn test_result_sequence_before_serialization() {
1257        let seq = ResultSequence::Before;
1258        let json = serde_json::to_string(&seq).unwrap();
1259        assert!(json.contains("\"before\""));
1260    }
1261
1262    #[test]
1263    fn test_result_sequence_after_serialization() {
1264        let seq = ResultSequence::After;
1265        let json = serde_json::to_string(&seq).unwrap();
1266        assert!(json.contains("\"after\""));
1267    }
1268}
zai_rs/model/tools.rs

zai_rs/model/
tools.rs
