Skip to main content

zai_rs/tool/web_search/
data.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4
5use crate::{
6    ZaiResult,
7    client::{
8        endpoints::{ApiBase, EndpointConfig, paths},
9        http::{HttpClient, HttpClientConfig, parse_typed_response},
10    },
11    tool::web_search::{request::*, response::*},
12};
13
14/// Web search API client
15pub struct WebSearchRequest {
16    /// API key for authentication
17    pub key: String,
18    url: String,
19    endpoint_config: EndpointConfig,
20    api_base: ApiBase,
21    http_config: Arc<HttpClientConfig>,
22    /// Request body
23    body: WebSearchBody,
24}
25
26impl WebSearchRequest {
27    /// Create a new web search request
28    ///
29    /// # Arguments
30    /// * `key` - API key for authentication
31    /// * `search_query` - Search query content (max 70 characters)
32    /// * `search_engine` - Search engine to use
33    pub fn new(key: String, search_query: String, search_engine: SearchEngine) -> Self {
34        let body = WebSearchBody::new(search_query, search_engine);
35        Self::with_body(key, body)
36    }
37
38    /// Create a web search request with a pre-configured body
39    pub fn with_body(key: String, body: WebSearchBody) -> Self {
40        let endpoint_config = EndpointConfig::default();
41        let api_base = ApiBase::PaasV4;
42        let url = endpoint_config.url(&api_base, paths::WEB_SEARCH);
43        Self {
44            key,
45            url,
46            endpoint_config,
47            api_base,
48            http_config: Arc::new(HttpClientConfig::default()),
49            body,
50        }
51    }
52
53    fn rebuild_url(&mut self) {
54        self.url = self.endpoint_config.url(&self.api_base, paths::WEB_SEARCH);
55    }
56
57    pub fn with_base_url(mut self, base: impl Into<String>) -> Self {
58        self.api_base = ApiBase::Custom(base.into());
59        self.rebuild_url();
60        self
61    }
62
63    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
64        self.endpoint_config = endpoint_config;
65        self.rebuild_url();
66        self
67    }
68
69    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
70        self.http_config = Arc::new(config);
71        self
72    }
73
74    /// Enable search intent recognition
75    pub fn with_search_intent(mut self, enabled: bool) -> Self {
76        self.body = self.body.with_search_intent(enabled);
77        self
78    }
79
80    /// Set the number of results to return
81    pub fn with_count(mut self, count: i32) -> Self {
82        self.body = self.body.with_count(count);
83        self
84    }
85
86    /// Set domain filter for search results
87    pub fn with_domain_filter(mut self, domain: String) -> Self {
88        self.body = self.body.with_domain_filter(domain);
89        self
90    }
91
92    /// Set time range filter for search results
93    pub fn with_recency_filter(mut self, filter: SearchRecencyFilter) -> Self {
94        self.body = self.body.with_recency_filter(filter);
95        self
96    }
97
98    /// Set content size preference
99    pub fn with_content_size(mut self, size: ContentSize) -> Self {
100        self.body = self.body.with_content_size(size);
101        self
102    }
103
104    /// Set custom request ID
105    pub fn with_request_id(mut self, request_id: String) -> Self {
106        self.body = self.body.with_request_id(request_id);
107        self
108    }
109
110    /// Set user ID
111    pub fn with_user_id(mut self, user_id: String) -> Self {
112        self.body = self.body.with_user_id(user_id);
113        self
114    }
115
116    /// Validate the request
117    pub fn validate(&self) -> ZaiResult<()> {
118        self.body.validate_constraints()
119    }
120
121    /// Send the web search request and return the response
122    pub async fn send(&self) -> ZaiResult<WebSearchResponse> {
123        self.validate()?;
124        let resp: reqwest::Response = self.post().await?;
125        let parsed = parse_typed_response::<WebSearchResponse>(resp).await?;
126        Ok(parsed)
127    }
128}
129
130#[async_trait]
131impl HttpClient for WebSearchRequest {
132    type Body = WebSearchBody;
133    type ApiUrl = String;
134    type ApiKey = String;
135
136    fn api_url(&self) -> &Self::ApiUrl {
137        &self.url
138    }
139
140    fn api_key(&self) -> &Self::ApiKey {
141        &self.key
142    }
143
144    fn body(&self) -> &Self::Body {
145        &self.body
146    }
147
148    fn http_config(&self) -> Arc<HttpClientConfig> {
149        Arc::clone(&self.http_config)
150    }
151}