zai_rs/tool/web_search/
data.rs

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