zai_rs/tool/web_search/
data.rs1use async_trait::async_trait;
2
3use crate::{
4 ZaiResult,
5 client::http::HttpClient,
6 tool::web_search::{request::*, response::*},
7};
8
9pub struct WebSearchRequest {
11 pub key: String,
13 body: WebSearchBody,
15}
16
17impl WebSearchRequest {
18 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 pub fn with_body(key: String, body: WebSearchBody) -> Self {
33 Self { key, body }
34 }
35
36 pub fn with_search_intent(mut self, enabled: bool) -> Self {
38 self.body = self.body.with_search_intent(enabled);
39 self
40 }
41
42 pub fn with_count(mut self, count: i32) -> Self {
44 self.body = self.body.with_count(count);
45 self
46 }
47
48 pub fn with_domain_filter(mut self, domain: String) -> Self {
50 self.body = self.body.with_domain_filter(domain);
51 self
52 }
53
54 pub fn with_recency_filter(mut self, filter: SearchRecencyFilter) -> Self {
56 self.body = self.body.with_recency_filter(filter);
57 self
58 }
59
60 pub fn with_content_size(mut self, size: ContentSize) -> Self {
62 self.body = self.body.with_content_size(size);
63 self
64 }
65
66 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 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 pub fn validate(&self) -> ZaiResult<()> {
80 self.body.validate_constraints()
81 }
82
83 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}