Skip to main content

zai_rs/tool/web_search/
request.rs

1use serde::{Deserialize, Serialize};
2use validator::Validate;
3
4/// Web search engine options supported by the API
5#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
6#[serde(rename_all = "snake_case")]
7pub enum SearchEngine {
8    /// Zhipu basic search engine
9    SearchStd,
10    /// Zhipu advanced search engine
11    SearchPro,
12    /// Sogou search engine
13    SearchProSogou,
14    /// Quark search engine
15    SearchProQuark,
16}
17
18/// Search result recency filter options
19#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
20#[serde(rename_all = "camelCase")]
21pub enum SearchRecencyFilter {
22    /// Search within one day
23    OneDay,
24    /// Search within one week
25    OneWeek,
26    /// Search within one month
27    OneMonth,
28    /// Search within one year
29    OneYear,
30    /// No time limit (default)
31    NoLimit,
32}
33
34/// Content size options for search results
35#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
36#[serde(rename_all = "snake_case")]
37pub enum ContentSize {
38    /// Medium content size with summary information for basic reasoning needs
39    Medium,
40    /// High content size with maximized context and detailed information
41    High,
42}
43
44/// Web search request body
45#[derive(Clone, Serialize, Validate)]
46pub struct WebSearchBody {
47    /// Search query content (max 70 characters)
48    #[validate(length(
49        min = 1,
50        max = 70,
51        message = "search_query must be between 1 and 70 characters"
52    ))]
53    pub search_query: String,
54
55    /// Search engine to use
56    pub search_engine: SearchEngine,
57
58    /// Whether to perform search intent recognition
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub search_intent: Option<bool>,
61
62    /// Number of results to return (1-50)
63    #[validate(range(min = 1, max = 50, message = "count must be between 1 and 50"))]
64    #[serde(skip_serializing_if = "Option::is_none")]
65    pub count: Option<i32>,
66
67    /// Domain filter for search results (whitelist)
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub search_domain_filter: Option<String>,
70
71    /// Time range filter for search results
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub search_recency_filter: Option<SearchRecencyFilter>,
74
75    /// Content size control
76    #[serde(skip_serializing_if = "Option::is_none")]
77    pub content_size: Option<ContentSize>,
78
79    /// Unique request identifier
80    #[validate(length(
81        min = 6,
82        max = 64,
83        message = "request_id must be between 6 and 64 characters"
84    ))]
85    #[serde(skip_serializing_if = "Option::is_none")]
86    pub request_id: Option<String>,
87
88    /// End user unique ID (6-128 characters)
89    #[validate(length(
90        min = 6,
91        max = 128,
92        message = "user_id must be between 6 and 128 characters"
93    ))]
94    #[serde(skip_serializing_if = "Option::is_none")]
95    pub user_id: Option<String>,
96}
97
98impl std::fmt::Debug for WebSearchBody {
99    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        formatter
101            .debug_struct("WebSearchBody")
102            .field("search_query", &"[REDACTED]")
103            .field("search_engine", &self.search_engine)
104            .field("search_intent", &self.search_intent)
105            .field("count", &self.count)
106            .field(
107                "search_domain_filter",
108                &self.search_domain_filter.as_ref().map(|_| "[REDACTED]"),
109            )
110            .field("search_recency_filter", &self.search_recency_filter)
111            .field("content_size", &self.content_size)
112            .field(
113                "request_id",
114                &self.request_id.as_ref().map(|_| "[REDACTED]"),
115            )
116            .field("user_id", &self.user_id.as_ref().map(|_| "[REDACTED]"))
117            .finish()
118    }
119}
120
121impl WebSearchBody {
122    /// Create a new web search request body with required parameters
123    pub fn new(search_query: impl Into<String>, search_engine: SearchEngine) -> Self {
124        Self {
125            search_query: search_query.into(),
126            search_engine,
127            search_intent: Some(false),
128            count: None,
129            search_domain_filter: None,
130            search_recency_filter: None,
131            content_size: None,
132            request_id: None,
133            user_id: None,
134        }
135    }
136
137    /// Enable search intent recognition
138    pub fn with_search_intent(mut self, enabled: bool) -> Self {
139        self.search_intent = Some(enabled);
140        self
141    }
142
143    /// Set the number of results to return
144    pub fn with_count(mut self, count: i32) -> Self {
145        self.count = Some(count);
146        self
147    }
148
149    /// Set domain filter for search results
150    pub fn with_domain_filter(mut self, domain: impl Into<String>) -> Self {
151        self.search_domain_filter = Some(domain.into());
152        self
153    }
154
155    /// Set time range filter for search results
156    pub fn with_recency_filter(mut self, filter: SearchRecencyFilter) -> Self {
157        self.search_recency_filter = Some(filter);
158        self
159    }
160
161    /// Set content size preference
162    pub fn with_content_size(mut self, size: ContentSize) -> Self {
163        self.content_size = Some(size);
164        self
165    }
166
167    /// Set custom request ID
168    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
169        self.request_id = Some(request_id.into());
170        self
171    }
172
173    /// Set user ID
174    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
175        self.user_id = Some(user_id.into());
176        self
177    }
178
179    /// Validate the request body constraints
180    pub fn validate_constraints(&self) -> crate::ZaiResult<()> {
181        self.validate().map_err(crate::ZaiError::from)?;
182
183        if self.search_query.trim().is_empty() {
184            return Err(crate::client::validation::invalid(
185                "search_query cannot be blank",
186            ));
187        }
188        if self.search_intent.is_none() {
189            return Err(crate::client::validation::invalid(
190                "search_intent is required",
191            ));
192        }
193        if self
194            .search_domain_filter
195            .as_deref()
196            .is_some_and(|domain| domain.trim().is_empty())
197        {
198            return Err(crate::client::validation::invalid(
199                "search_domain_filter cannot be blank",
200            ));
201        }
202
203        // Additional validation for count based on search engine
204        if let Some(count) = self.count
205            && matches!(self.search_engine, SearchEngine::SearchProSogou)
206        {
207            match count {
208                10 | 20 | 30 | 40 | 50 => {},
209                _ => {
210                    return Err(crate::client::validation::invalid(
211                        "search_pro_sogou only supports count values: 10, 20, 30, 40, 50",
212                    ));
213                },
214            }
215        }
216
217        if matches!(self.search_engine, SearchEngine::SearchProQuark)
218            && (self.count.is_some() || self.search_domain_filter.is_some())
219        {
220            return Err(crate::client::validation::invalid(
221                "search_pro_quark does not support count or search_domain_filter",
222            ));
223        }
224
225        Ok(())
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232
233    #[test]
234    fn defaults_and_recency_use_official_wire_values() {
235        let body = WebSearchBody::new("query".to_string(), SearchEngine::SearchStd)
236            .with_recency_filter(SearchRecencyFilter::OneWeek);
237        let value = serde_json::to_value(body).unwrap();
238
239        assert_eq!(value["search_intent"], false);
240        assert_eq!(value["search_recency_filter"], "oneWeek");
241    }
242
243    #[test]
244    fn debug_redacts_query_domain_and_identifiers() {
245        let body = WebSearchBody::new("private query".to_owned(), SearchEngine::SearchStd)
246            .with_domain_filter("private.example")
247            .with_request_id("private-request")
248            .with_user_id("private-user");
249        let debug = format!("{body:?}");
250        for secret in [
251            "private query",
252            "private.example",
253            "private-request",
254            "private-user",
255        ] {
256            assert!(!debug.contains(secret));
257        }
258    }
259}