Skip to main content

zai_rs/tool/web_search/
data.rs

1use async_trait::async_trait;
2
3use crate::{
4    ZaiResult,
5    client::http::HttpClient,
6    tool::web_search::{request::*, response::*},
7};
8
9/// Web search API client
10pub struct WebSearchRequest {
11    /// API key for authentication
12    pub key: String,
13    /// Request body
14    body: WebSearchBody,
15}
16
17impl WebSearchRequest {
18    /// Create a new web search request
19    ///
20    /// # Arguments
21    /// * `key` - API key for authentication
22    /// * `search_query` - Search query content (max 70 characters)
23    /// * `search_engine` - Search engine to use
24    pub fn new(key: String, search_query: String, search_engine: SearchEngine) -> Self {
25        Self {
26            key,
27            body: WebSearchBody::new(search_query, search_engine),
28        }
29    }
30
31    /// Create a web search request with a pre-configured body
32    pub fn with_body(key: String, body: WebSearchBody) -> Self {
33        Self { key, body }
34    }
35
36    /// Enable search intent recognition
37    pub fn with_search_intent(mut self, enabled: bool) -> Self {
38        self.body = self.body.with_search_intent(enabled);
39        self
40    }
41
42    /// Set the number of results to return
43    pub fn with_count(mut self, count: i32) -> Self {
44        self.body = self.body.with_count(count);
45        self
46    }
47
48    /// Set domain filter for search results
49    pub fn with_domain_filter(mut self, domain: String) -> Self {
50        self.body = self.body.with_domain_filter(domain);
51        self
52    }
53
54    /// Set time range filter for search results
55    pub fn with_recency_filter(mut self, filter: SearchRecencyFilter) -> Self {
56        self.body = self.body.with_recency_filter(filter);
57        self
58    }
59
60    /// Set content size preference
61    pub fn with_content_size(mut self, size: ContentSize) -> Self {
62        self.body = self.body.with_content_size(size);
63        self
64    }
65
66    /// Set custom request ID
67    pub fn with_request_id(mut self, request_id: String) -> Self {
68        self.body = self.body.with_request_id(request_id);
69        self
70    }
71
72    /// Set user ID
73    pub fn with_user_id(mut self, user_id: String) -> Self {
74        self.body = self.body.with_user_id(user_id);
75        self
76    }
77
78    /// Validate the request
79    pub fn validate(&self) -> ZaiResult<()> {
80        self.body.validate_constraints()
81    }
82
83    /// Send the web search request and return the response
84    pub async fn send(&self) -> ZaiResult<WebSearchResponse> {
85        self.validate()?;
86        let resp: reqwest::Response = self.post().await?;
87        let parsed = resp.json::<WebSearchResponse>().await?;
88        Ok(parsed)
89    }
90}
91
92#[async_trait]
93impl HttpClient for WebSearchRequest {
94    type Body = WebSearchBody;
95    type ApiUrl = &'static str;
96    type ApiKey = String;
97
98    fn api_url(&self) -> &Self::ApiUrl {
99        &"https://open.bigmodel.cn/api/paas/v4/web_search"
100    }
101
102    fn api_key(&self) -> &Self::ApiKey {
103        &self.key
104    }
105
106    fn body(&self) -> &Self::Body {
107        &self.body
108    }
109}