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//! - [`FunctionTool`] — Defines a callable function with JSON-schema parameters
12//! - [`WebSearchTool`] — Enables live web search within chat
13//! - [`Retrieval`] — Enables knowledge-base retrieval
14//! - [`ToolChoice`] — Controls tool-selection behaviour (`auto`, `none`, or
15//!   specific function)
16
17use std::collections::HashMap;
18
19use serde::Serialize;
20use validator::*;
21
22use super::model_validate::validate_json_schema_value;
23use crate::tool::web_search::request::{ContentSize, SearchEngine, SearchRecencyFilter};
24
25/// Controls thinking/reasoning capabilities in AI models.
26///
27/// This structure determines whether a model should engage in step-by-step
28/// reasoning when processing requests, and whether to preserve reasoning
29/// content across turns via `clear_thinking`. Thinking mode can improve
30/// accuracy for complex tasks but may increase response time and token usage.
31///
32/// ## Fields
33///
34/// - `mode` - Whether thinking is enabled or disabled
35/// - `clear_thinking` - When `false`, preserves `reasoning_content` across
36///   turns (recommended for Coding / Agent scenarios)
37///
38/// ## Usage
39///
40/// ```rust,ignore
41/// let client = ChatCompletion::new(model, messages, api_key)
42///     .with_thinking(ThinkingType::enabled());
43///
44/// // Preserve reasoning content across turns (Coding / Agent)
45/// let client = ChatCompletion::new(model, messages, api_key)
46///     .with_thinking(ThinkingType::enabled().with_clear_thinking(false));
47/// ```
48///
49/// ## Model Compatibility
50///
51/// Thinking capabilities are available only on models that implement the
52/// `ThinkEnable` trait, such as GLM-5.2, GLM-5.1, GLM-5, GLM-4.7, and GLM-4.5
53/// series models.
54#[derive(Debug, Clone, Serialize)]
55pub struct ThinkingType {
56    /// Whether thinking is enabled or disabled.
57    #[serde(rename = "type")]
58    pub mode: ThinkingMode,
59
60    /// Whether to clear historical `reasoning_content`.
61    ///
62    /// - `true` (default for standard API): Clears reasoning content each turn.
63    /// - `false` (recommended for Coding / Agent): Preserves reasoning content
64    ///   across turns, enabling better context for multi-step tool calls.
65    #[serde(skip_serializing_if = "Option::is_none")]
66    pub clear_thinking: Option<bool>,
67}
68
69/// Thinking mode variants.
70#[derive(Debug, Clone, Serialize)]
71#[serde(rename_all = "lowercase")]
72pub enum ThinkingMode {
73    Enabled,
74    Disabled,
75}
76
77impl ThinkingType {
78    /// Create a new thinking configuration with enabled mode.
79    pub fn enabled() -> Self {
80        Self {
81            mode: ThinkingMode::Enabled,
82            clear_thinking: None,
83        }
84    }
85
86    /// Create a new thinking configuration with disabled mode.
87    pub fn disabled() -> Self {
88        Self {
89            mode: ThinkingMode::Disabled,
90            clear_thinking: None,
91        }
92    }
93
94    /// Set whether to clear historical reasoning content.
95    ///
96    /// Use `false` for Coding / Agent scenarios where reasoning content should
97    /// be preserved across turns.
98    pub fn with_clear_thinking(mut self, clear: bool) -> Self {
99        self.clear_thinking = Some(clear);
100        self
101    }
102}
103
104/// Reasoning depth level for the `reasoning_effort` parameter (GLM-5.2+).
105///
106/// Controls how much reasoning the model invests when thinking mode is
107/// enabled. Higher levels yield deeper reasoning at the cost of latency and
108/// token usage; lower levels are faster and cheaper. Available only on
109/// GLM-5.2 and above (models implementing
110/// [`ReasoningEffortEnable`](super::traits::ReasoningEffortEnable)).
111///
112/// Levels, from highest to lowest reasoning depth:
113///
114/// | Variant | Description |
115/// |---------|-------------|
116/// | [`Max`](Self::Max) | Maximum reasoning depth; recommended for coding / architecture-level tasks |
117/// | [`Xhigh`](Self::Xhigh) | Extra-high reasoning |
118/// | [`High`](Self::High) | High reasoning (default mapping in many clients) |
119/// | [`Medium`](Self::Medium) | Balanced reasoning |
120/// | [`Low`](Self::Low) | Light reasoning |
121/// | [`Minimal`](Self::Minimal) | Minimal reasoning |
122/// | [`None`](Self::None) | No extra reasoning beyond base behaviour |
123///
124/// ## Usage
125///
126/// ```rust,ignore
127/// use zai_rs::model::tools::ReasoningEffort;
128///
129/// let client = ChatCompletion::new(GLM5_2 {}, messages, api_key)
130///     .with_thinking(ThinkingType::enabled())
131///     .with_reasoning_effort(ReasoningEffort::Max);
132/// ```
133#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
134pub enum ReasoningEffort {
135    /// Maximum reasoning depth. Recommended for coding and architecture-level
136    /// tasks where correctness matters most.
137    #[serde(rename = "max")]
138    Max,
139    /// Extra-high reasoning depth.
140    #[serde(rename = "xhigh")]
141    Xhigh,
142    /// High reasoning depth.
143    #[serde(rename = "high")]
144    High,
145    /// Balanced reasoning depth.
146    #[serde(rename = "medium")]
147    Medium,
148    /// Light reasoning depth.
149    #[serde(rename = "low")]
150    Low,
151    /// Minimal reasoning depth.
152    #[serde(rename = "minimal")]
153    Minimal,
154    /// No extra reasoning beyond base behaviour.
155    #[serde(rename = "none")]
156    None,
157}
158
159/// Available tools that AI assistants can invoke during conversations.
160///
161/// This enum defines the different categories of external tools and
162/// capabilities that can be made available to AI models. Each tool type serves
163/// specific purposes and has its own configuration requirements.
164///
165/// ## Tool Categories
166///
167/// ### Function Tools
168/// Custom user-defined functions that the AI can call with structured
169/// parameters. Useful for integrating external APIs, databases, or business
170/// logic.
171///
172/// ### Retrieval Tools
173/// Access to knowledge bases, document collections, or information retrieval
174/// systems. Enables the AI to query structured knowledge sources.
175///
176/// ### Web Search Tools
177/// Internet search capabilities for accessing current information.
178/// Allows the AI to perform web searches and retrieve up-to-date information.
179///
180/// ### MCP Tools
181/// Model Context Protocol tools for standardized tool integration.
182/// Provides a standardized interface for tool communication.
183///
184/// ## Usage
185///
186/// ```rust,ignore
187/// // Function tool
188/// let function_tool = Tools::Function {
189///     function: Function::new("get_weather", "Get weather data", parameters)
190/// };
191///
192/// // Web search tool
193/// let search_tool = Tools::WebSearch {
194///     web_search: WebSearch::new(SearchEngine::SearchPro)
195///         .with_enable(true)
196///         .with_count(10)
197/// };
198/// ```
199#[derive(Debug, Clone, Serialize)]
200#[serde(tag = "type")]
201#[serde(rename_all = "snake_case")]
202pub enum Tools {
203    /// Custom function calling tool with parameters.
204    ///
205    /// Allows the AI to invoke user-defined functions with structured
206    /// arguments. Functions must be pre-defined with JSON schemas for
207    /// parameter validation.
208    Function { function: Function },
209
210    /// Knowledge retrieval system access tools.
211    ///
212    /// Provides access to knowledge bases, document collections, or other
213    /// structured information sources that the AI can query.
214    Retrieval { retrieval: Retrieval },
215
216    /// Web search capabilities for internet access.
217    ///
218    /// Enables the AI to perform web searches and access current information
219    /// from the internet. Supports various search engines and configurations.
220    WebSearch { web_search: WebSearch },
221
222    /// Model Context Protocol (MCP) tools.
223    ///
224    /// Standardized tools that follow the Model Context Protocol specification,
225    /// providing a consistent interface for tool integration and communication.
226    #[serde(rename = "mcp")]
227    MCP { mcp: MCP },
228}
229
230/// Definition of a callable function tool.
231///
232/// This structure defines a function that can be called by the assistant,
233/// including its name, description, and parameter schema.
234///
235/// # Validation
236///
237/// * `name` - Must be between 1 and 64 characters
238/// * `parameters` - Must be a valid JSON schema
239#[derive(Debug, Clone, Serialize, Validate)]
240pub struct Function {
241    /// The name of the function. Must be between 1 and 64 characters.
242    #[validate(length(min = 1, max = 64))]
243    pub name: String,
244
245    /// A description of what the function does.
246    pub description: String,
247
248    /// JSON schema describing the function's parameters.
249    /// Server expects an object; keep as Value to avoid double-encoding
250    /// strings.
251    #[serde(skip_serializing_if = "Option::is_none")]
252    #[validate(custom(function = "validate_json_schema_value"))]
253    pub parameters: Option<serde_json::Value>,
254}
255
256impl Function {
257    /// Creates a new function call definition.
258    ///
259    /// # Arguments
260    ///
261    /// * `name` - The name of the function
262    /// * `description` - A description of what the function does
263    /// * `parameters` - JSON schema string describing the function parameters
264    ///
265    /// # Returns
266    ///
267    /// A new `Function` instance.
268    ///
269    /// # Examples
270    ///
271    /// ```rust,ignore
272    /// let func = Function::new(
273    ///     "get_weather",
274    ///     "Get current weather for a location",
275    ///     r#"{"type": "object", "properties": {"location": {"type": "string"}}}"#
276    /// );
277    /// ```
278    pub fn new(
279        name: impl Into<String>,
280        description: impl Into<String>,
281        parameters: serde_json::Value,
282    ) -> Self {
283        Self {
284            name: name.into(),
285            description: description.into(),
286            parameters: Some(parameters),
287        }
288    }
289}
290
291/// Configuration for retrieval tool capabilities.
292///
293/// This structure represents a retrieval tool that can access knowledge bases
294/// or document collections. Currently a placeholder for future expansion.
295#[derive(Debug, Clone, Serialize)]
296pub struct Retrieval {
297    knowledge_id: String,
298    #[serde(skip_serializing_if = "Option::is_none")]
299    prompt_template: Option<String>,
300}
301
302impl Retrieval {
303    /// Creates a new `Retrieval` instance.
304    pub fn new(knowledge_id: impl Into<String>, prompt_template: Option<String>) -> Self {
305        Self {
306            knowledge_id: knowledge_id.into(),
307            prompt_template,
308        }
309    }
310}
311
312/// Configuration for web search tool capabilities.
313///
314/// The order in which search results are returned.
315#[derive(Debug, Clone, Serialize, PartialEq)]
316#[serde(rename_all = "snake_case")]
317pub enum ResultSequence {
318    Before,
319    After,
320}
321
322/// This structure represents a web search tool that can perform internet
323/// searches. Fields mirror the external web_search schema.
324#[derive(Debug, Clone, Serialize, Validate)]
325pub struct WebSearch {
326    /// Search engine type (required). Supported: search_std, search_pro,
327    /// search_pro_sogou, search_pro_quark.
328    pub search_engine: SearchEngine,
329
330    /// Whether to enable web search. Default is false.
331    #[serde(skip_serializing_if = "Option::is_none")]
332    pub enable: Option<bool>,
333
334    /// Force-triggered search query string.
335    #[serde(skip_serializing_if = "Option::is_none")]
336    pub search_query: Option<String>,
337
338    /// Whether to perform search intent detection. true: execute only when
339    /// intent is detected; false: skip detection and search directly.
340    #[serde(skip_serializing_if = "Option::is_none")]
341    pub search_intent: Option<bool>,
342
343    /// Number of results to return (1-50).
344    #[serde(skip_serializing_if = "Option::is_none")]
345    #[validate(range(min = 1, max = 50))]
346    pub count: Option<u32>,
347
348    /// Whitelist domain filter, e.g., "www.example.com".
349    #[serde(skip_serializing_if = "Option::is_none")]
350    pub search_domain_filter: Option<String>,
351
352    /// Time range filter.
353    #[serde(skip_serializing_if = "Option::is_none")]
354    pub search_recency_filter: Option<SearchRecencyFilter>,
355
356    /// Snippet summary size: medium or high.
357    #[serde(skip_serializing_if = "Option::is_none")]
358    pub content_size: Option<ContentSize>,
359
360    /// Return sequence for search results: before or after.
361    #[serde(skip_serializing_if = "Option::is_none")]
362    pub result_sequence: Option<ResultSequence>,
363
364    /// Whether to include detailed search source information.
365    #[serde(skip_serializing_if = "Option::is_none")]
366    pub search_result: Option<bool>,
367
368    /// Whether an answer requires search results to be returned.
369    #[serde(skip_serializing_if = "Option::is_none")]
370    pub require_search: Option<bool>,
371
372    /// Custom prompt to post-process search results.
373    #[serde(skip_serializing_if = "Option::is_none")]
374    pub search_prompt: Option<String>,
375}
376
377impl WebSearch {
378    /// Create a WebSearch config with the required search engine; other fields
379    /// are optional.
380    pub fn new(search_engine: SearchEngine) -> Self {
381        Self {
382            search_engine,
383            enable: None,
384            search_query: None,
385            search_intent: None,
386            count: None,
387            search_domain_filter: None,
388            search_recency_filter: None,
389            content_size: None,
390            result_sequence: None,
391            search_result: None,
392            require_search: None,
393            search_prompt: None,
394        }
395    }
396
397    /// Enable or disable web search.
398    pub fn with_enable(mut self, enable: bool) -> Self {
399        self.enable = Some(enable);
400        self
401    }
402    /// Set a forced search query.
403    pub fn with_search_query(mut self, query: impl Into<String>) -> Self {
404        self.search_query = Some(query.into());
405        self
406    }
407    /// Set search intent detection behavior.
408    pub fn with_search_intent(mut self, search_intent: bool) -> Self {
409        self.search_intent = Some(search_intent);
410        self
411    }
412    /// Set results count (1-50).
413    pub fn with_count(mut self, count: u32) -> Self {
414        self.count = Some(count);
415        self
416    }
417    /// Restrict to a whitelist domain.
418    pub fn with_search_domain_filter(mut self, domain: impl Into<String>) -> Self {
419        self.search_domain_filter = Some(domain.into());
420        self
421    }
422    /// Set time range filter.
423    pub fn with_search_recency_filter(mut self, filter: SearchRecencyFilter) -> Self {
424        self.search_recency_filter = Some(filter);
425        self
426    }
427    /// Set content size.
428    pub fn with_content_size(mut self, size: ContentSize) -> Self {
429        self.content_size = Some(size);
430        self
431    }
432    /// Set result sequence.
433    pub fn with_result_sequence(mut self, seq: ResultSequence) -> Self {
434        self.result_sequence = Some(seq);
435        self
436    }
437    /// Toggle returning detailed search source info.
438    pub fn with_search_result(mut self, enable: bool) -> Self {
439        self.search_result = Some(enable);
440        self
441    }
442    /// Require search results for answering.
443    pub fn with_require_search(mut self, require: bool) -> Self {
444        self.require_search = Some(require);
445        self
446    }
447    /// Set a custom prompt to post-process search results.
448    pub fn with_search_prompt(mut self, prompt: impl Into<String>) -> Self {
449        self.search_prompt = Some(prompt.into());
450        self
451    }
452}
453/// Represents the MCP connection configuration. When connecting to Zhipu's MCP
454/// server using an MCP code, fill `server_label` with that code and leave
455/// `server_url` empty.
456#[derive(Debug, Clone, Serialize, Validate)]
457pub struct MCP {
458    /// MCP server identifier (required). If connecting to Zhipu MCP via code,
459    /// put the code here.
460    #[validate(length(min = 1))]
461    pub server_label: String,
462
463    /// MCP server URL.
464    #[serde(skip_serializing_if = "Option::is_none")]
465    #[validate(url)]
466    pub server_url: Option<String>,
467
468    /// Transport type. Default: streamable-http.
469    #[serde(skip_serializing_if = "Option::is_none")]
470    pub transport_type: Option<MCPTransportType>,
471
472    /// Allowed tool names.
473    #[serde(skip_serializing_if = "Vec::is_empty")]
474    pub allowed_tools: Vec<String>,
475
476    /// Authentication headers required by the MCP server.
477    #[serde(skip_serializing_if = "Option::is_none")]
478    pub headers: Option<HashMap<String, String>>,
479}
480
481impl MCP {
482    /// Create a new MCP config with required server_label and default transport
483    /// type.
484    pub fn new(server_label: impl Into<String>) -> Self {
485        Self {
486            server_label: server_label.into(),
487            server_url: None,
488            transport_type: Some(MCPTransportType::StreamableHttp),
489            allowed_tools: Vec::new(),
490            headers: None,
491        }
492    }
493
494    /// Set the MCP server URL.
495    pub fn with_server_url(mut self, url: impl Into<String>) -> Self {
496        self.server_url = Some(url.into());
497        self
498    }
499    /// Set the MCP transport type.
500    pub fn with_transport_type(mut self, transport: MCPTransportType) -> Self {
501        self.transport_type = Some(transport);
502        self
503    }
504    /// Replace the allowed tool list.
505    pub fn with_allowed_tools(mut self, tools: impl Into<Vec<String>>) -> Self {
506        self.allowed_tools = tools.into();
507        self
508    }
509    /// Add a single allowed tool.
510    pub fn add_allowed_tool(mut self, tool: impl Into<String>) -> Self {
511        self.allowed_tools.push(tool.into());
512        self
513    }
514    /// Set authentication headers map.
515    pub fn with_headers(mut self, headers: HashMap<String, String>) -> Self {
516        self.headers = Some(headers);
517        self
518    }
519    /// Add or update a single header entry.
520    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
521        let mut map = self.headers.unwrap_or_default();
522        map.insert(key.into(), value.into());
523        self.headers = Some(map);
524        self
525    }
526}
527
528/// Allowed MCP transport types.
529#[derive(Debug, Clone, Serialize, PartialEq)]
530#[serde(rename_all = "kebab-case")]
531pub enum MCPTransportType {
532    Sse,
533    StreamableHttp,
534}
535
536/// Specifies the format for the model's response.
537///
538/// This enum controls how the model should structure its output, either as
539/// plain text or as a structured JSON object.
540///
541/// # Variants
542///
543/// * `Text` - Plain text response format
544/// * `JsonObject` - Structured JSON object response format
545#[derive(Debug, Clone, Copy, Serialize)]
546#[serde(rename_all = "snake_case")]
547#[serde(tag = "type")]
548pub enum ResponseFormat {
549    /// Plain text response format.
550    Text,
551    /// Structured JSON object response format.
552    JsonObject,
553}
554
555#[cfg(test)]
556mod tests {
557    use super::*;
558
559    // ThinkingType tests
560    #[test]
561    fn test_thinking_type_enabled_serialization() {
562        let thinking = ThinkingType::enabled();
563        let json = serde_json::to_string(&thinking).unwrap();
564        assert!(json.contains("\"type\":\"enabled\""));
565        assert!(!json.contains("clear_thinking"));
566    }
567
568    #[test]
569    fn test_thinking_type_disabled_serialization() {
570        let thinking = ThinkingType::disabled();
571        let json = serde_json::to_string(&thinking).unwrap();
572        assert!(json.contains("\"type\":\"disabled\""));
573        assert!(!json.contains("clear_thinking"));
574    }
575
576    #[test]
577    fn test_thinking_type_with_clear_thinking_serialization() {
578        let thinking = ThinkingType::enabled().with_clear_thinking(false);
579        let json = serde_json::to_string(&thinking).unwrap();
580        assert!(json.contains("\"type\":\"enabled\""));
581        assert!(json.contains("\"clear_thinking\":false"));
582    }
583
584    #[test]
585    fn test_thinking_type_disabled_with_clear_thinking() {
586        let thinking = ThinkingType::disabled().with_clear_thinking(true);
587        let json = serde_json::to_string(&thinking).unwrap();
588        assert!(json.contains("\"type\":\"disabled\""));
589        assert!(json.contains("\"clear_thinking\":true"));
590    }
591
592    // Function tests
593    #[test]
594    fn test_function_new() {
595        let params = serde_json::json!({
596            "type": "object",
597            "properties": {
598                "name": {"type": "string"}
599            }
600        });
601        let func = Function::new("test_func", "A test function", params);
602
603        assert_eq!(func.name, "test_func");
604        assert_eq!(func.description, "A test function");
605        assert!(func.parameters.is_some());
606    }
607
608    #[test]
609    fn test_function_serialization() {
610        let params = serde_json::json!({
611            "type": "object",
612            "properties": {
613                "value": {"type": "number"}
614            }
615        });
616        let func = Function::new("test_func", "A test function", params);
617        let json = serde_json::to_string(&func).unwrap();
618
619        assert!(json.contains("\"name\":\"test_func\""));
620        assert!(json.contains("\"description\":\"A test function\""));
621        assert!(json.contains("\"properties\""));
622    }
623
624    #[test]
625    fn test_function_validation() {
626        let params = serde_json::json!({
627            "type": "object",
628            "properties": {}
629        });
630        let func = Function::new("valid_name", "Description", params.clone());
631
632        // Name length validation: 1-64 characters
633        assert!(func.validate().is_ok());
634
635        let invalid_name = Function::new("", "Description", params.clone());
636        assert!(invalid_name.validate().is_err());
637
638        let long_name = Function::new("a".repeat(65), "Description", params);
639        assert!(long_name.validate().is_err());
640    }
641
642    // Retrieval tests
643    #[test]
644    fn test_retrieval_new() {
645        let retrieval = Retrieval::new("kb_123", Some("template".to_string()));
646        assert_eq!(retrieval.knowledge_id, "kb_123");
647        assert_eq!(retrieval.prompt_template, Some("template".to_string()));
648    }
649
650    #[test]
651    fn test_retrieval_new_without_template() {
652        let retrieval = Retrieval::new("kb_456", None);
653        assert_eq!(retrieval.knowledge_id, "kb_456");
654        assert!(retrieval.prompt_template.is_none());
655    }
656
657    #[test]
658    fn test_retrieval_serialization() {
659        let retrieval = Retrieval::new("kb_789", None);
660        let json = serde_json::to_string(&retrieval).unwrap();
661        assert!(json.contains("\"knowledge_id\":\"kb_789\""));
662        // prompt_template should be omitted when None
663        assert!(!json.contains("prompt_template"));
664    }
665
666    // WebSearch tests
667    #[test]
668    fn test_web_search_new() {
669        let web_search = WebSearch::new(SearchEngine::SearchPro);
670        assert_eq!(web_search.search_engine, SearchEngine::SearchPro);
671        assert!(web_search.enable.is_none());
672    }
673
674    #[test]
675    fn test_web_search_with_enable() {
676        let web_search = WebSearch::new(SearchEngine::SearchPro).with_enable(true);
677        assert_eq!(web_search.enable, Some(true));
678    }
679
680    #[test]
681    fn test_web_search_with_search_query() {
682        let web_search = WebSearch::new(SearchEngine::SearchPro).with_search_query("test query");
683        assert_eq!(web_search.search_query, Some("test query".to_string()));
684    }
685
686    #[test]
687    fn test_web_search_with_search_intent() {
688        let web_search = WebSearch::new(SearchEngine::SearchPro).with_search_intent(true);
689        assert_eq!(web_search.search_intent, Some(true));
690    }
691
692    #[test]
693    fn test_web_search_with_count() {
694        let web_search = WebSearch::new(SearchEngine::SearchPro).with_count(10);
695        assert_eq!(web_search.count, Some(10));
696    }
697
698    #[test]
699    fn test_web_search_with_search_domain_filter() {
700        let web_search =
701            WebSearch::new(SearchEngine::SearchPro).with_search_domain_filter("example.com");
702        assert_eq!(
703            web_search.search_domain_filter,
704            Some("example.com".to_string())
705        );
706    }
707
708    #[test]
709    fn test_web_search_with_search_recency_filter() {
710        let filter = SearchRecencyFilter::OneDay;
711        let web_search =
712            WebSearch::new(SearchEngine::SearchPro).with_search_recency_filter(filter.clone());
713        assert_eq!(web_search.search_recency_filter, Some(filter));
714    }
715
716    #[test]
717    fn test_web_search_with_content_size() {
718        let size = ContentSize::Medium;
719        let web_search = WebSearch::new(SearchEngine::SearchPro).with_content_size(size.clone());
720        assert_eq!(web_search.content_size, Some(size));
721    }
722
723    #[test]
724    fn test_web_search_with_result_sequence() {
725        let seq = ResultSequence::After;
726        let web_search = WebSearch::new(SearchEngine::SearchPro).with_result_sequence(seq.clone());
727        assert_eq!(web_search.result_sequence, Some(seq));
728    }
729
730    #[test]
731    fn test_web_search_with_search_result() {
732        let web_search = WebSearch::new(SearchEngine::SearchPro).with_search_result(true);
733        assert_eq!(web_search.search_result, Some(true));
734    }
735
736    #[test]
737    fn test_web_search_with_require_search() {
738        let web_search = WebSearch::new(SearchEngine::SearchPro).with_require_search(true);
739        assert_eq!(web_search.require_search, Some(true));
740    }
741
742    #[test]
743    fn test_web_search_with_search_prompt() {
744        let web_search =
745            WebSearch::new(SearchEngine::SearchPro).with_search_prompt("custom prompt");
746        assert_eq!(web_search.search_prompt, Some("custom prompt".to_string()));
747    }
748
749    #[test]
750    fn test_web_search_serialization() {
751        let web_search = WebSearch::new(SearchEngine::SearchPro)
752            .with_enable(true)
753            .with_count(5);
754        let json = serde_json::to_string(&web_search).unwrap();
755        assert!(json.contains("\"search_engine\""));
756        assert!(json.contains("\"enable\":true"));
757        assert!(json.contains("\"count\":5"));
758    }
759
760    // MCP tests
761    #[test]
762    fn test_mcp_new() {
763        let mcp = MCP::new("server_label");
764        assert_eq!(mcp.server_label, "server_label");
765        assert_eq!(mcp.transport_type, Some(MCPTransportType::StreamableHttp));
766        assert!(mcp.allowed_tools.is_empty());
767    }
768
769    #[test]
770    fn test_mcp_with_server_url() {
771        let mcp = MCP::new("server_label").with_server_url("https://example.com");
772        assert_eq!(mcp.server_url, Some("https://example.com".to_string()));
773    }
774
775    #[test]
776    fn test_mcp_with_transport_type() {
777        let mcp = MCP::new("server_label").with_transport_type(MCPTransportType::Sse);
778        assert_eq!(mcp.transport_type, Some(MCPTransportType::Sse));
779    }
780
781    #[test]
782    fn test_mcp_with_allowed_tools() {
783        let mcp = MCP::new("server_label")
784            .with_allowed_tools(vec!["tool1".to_string(), "tool2".to_string()]);
785        assert_eq!(mcp.allowed_tools.len(), 2);
786        assert!(mcp.allowed_tools.contains(&"tool1".to_string()));
787    }
788
789    #[test]
790    fn test_mcp_add_allowed_tool() {
791        let mcp = MCP::new("server_label")
792            .add_allowed_tool("tool1")
793            .add_allowed_tool("tool2");
794        assert_eq!(mcp.allowed_tools.len(), 2);
795    }
796
797    #[test]
798    fn test_mcp_with_headers() {
799        let mut headers = HashMap::new();
800        headers.insert("Authorization".to_string(), "Bearer token".to_string());
801        let mcp = MCP::new("server_label").with_headers(headers.clone());
802        assert_eq!(mcp.headers, Some(headers));
803    }
804
805    #[test]
806    fn test_mcp_with_header() {
807        let mcp = MCP::new("server_label").with_header("Authorization", "Bearer token");
808        let headers = mcp.headers.unwrap();
809        assert_eq!(
810            headers.get("Authorization"),
811            Some(&"Bearer token".to_string())
812        );
813    }
814
815    #[test]
816    fn test_mcp_serialization() {
817        let mcp = MCP::new("server_label")
818            .with_server_url("https://example.com")
819            .with_transport_type(MCPTransportType::Sse);
820        let json = serde_json::to_string(&mcp).unwrap();
821        assert!(json.contains("\"server_label\":\"server_label\""));
822        assert!(json.contains("\"server_url\":\"https://example.com\""));
823        assert!(json.contains("\"transport_type\":\"sse\""));
824        // allowed_tools should be omitted when empty
825        assert!(!json.contains("allowed_tools"));
826    }
827
828    // MCPTransportType tests
829    #[test]
830    fn test_mcp_transport_type_sse_serialization() {
831        let transport = MCPTransportType::Sse;
832        let json = serde_json::to_string(&transport).unwrap();
833        assert!(json.contains("\"sse\""));
834    }
835
836    #[test]
837    fn test_mcp_transport_type_streamable_http_serialization() {
838        let transport = MCPTransportType::StreamableHttp;
839        let json = serde_json::to_string(&transport).unwrap();
840        assert!(json.contains("\"streamable-http\""));
841    }
842
843    // ResponseFormat tests
844    #[test]
845    fn test_response_format_text_serialization() {
846        let format = ResponseFormat::Text;
847        let json = serde_json::to_string(&format).unwrap();
848        assert!(json.contains("\"type\":\"text\""));
849    }
850
851    #[test]
852    fn test_response_format_json_object_serialization() {
853        let format = ResponseFormat::JsonObject;
854        let json = serde_json::to_string(&format).unwrap();
855        assert!(json.contains("\"type\":\"json_object\""));
856    }
857
858    // Tools enum tests
859    #[test]
860    fn test_tools_function_serialization() {
861        let func = Function::new("test_func", "test", serde_json::json!({}));
862        let tools = Tools::Function { function: func };
863        let json = serde_json::to_string(&tools).unwrap();
864        assert!(json.contains("\"type\":\"function\""));
865        assert!(json.contains("\"name\":\"test_func\""));
866    }
867
868    #[test]
869    fn test_tools_retrieval_serialization() {
870        let retrieval = Retrieval::new("kb_123", None);
871        let tools = Tools::Retrieval { retrieval };
872        let json = serde_json::to_string(&tools).unwrap();
873        assert!(json.contains("\"type\":\"retrieval\""));
874        assert!(json.contains("\"knowledge_id\":\"kb_123\""));
875    }
876
877    #[test]
878    fn test_tools_web_search_serialization() {
879        let web_search = WebSearch::new(SearchEngine::SearchPro);
880        let tools = Tools::WebSearch { web_search };
881        let json = serde_json::to_string(&tools).unwrap();
882        assert!(json.contains("\"type\":\"web_search\""));
883        assert!(json.contains("\"search_engine\""));
884    }
885
886    #[test]
887    fn test_tools_mcp_serialization() {
888        let mcp = MCP::new("server_label");
889        let tools = Tools::MCP { mcp };
890        let json = serde_json::to_string(&tools).unwrap();
891        assert!(json.contains("\"type\":\"mcp\""));
892        assert!(json.contains("\"server_label\":\"server_label\""));
893    }
894
895    // ResultSequence tests
896    #[test]
897    fn test_result_sequence_before_serialization() {
898        let seq = ResultSequence::Before;
899        let json = serde_json::to_string(&seq).unwrap();
900        assert!(json.contains("\"before\""));
901    }
902
903    #[test]
904    fn test_result_sequence_after_serialization() {
905        let seq = ResultSequence::After;
906        let json = serde_json::to_string(&seq).unwrap();
907        assert!(json.contains("\"after\""));
908    }
909}
zai_rs/model/tools.rs

zai_rs/model/
tools.rs
