use async_trait::async_trait;
use crate::{
ZaiResult,
client::http::HttpClient,
tool::web_search::{request::*, response::*},
};
pub struct WebSearchRequest {
pub key: String,
body: WebSearchBody,
}
impl WebSearchRequest {
pub fn new(key: String, search_query: String, search_engine: SearchEngine) -> Self {
Self {
key,
body: WebSearchBody::new(search_query, search_engine),
}
}
pub fn with_body(key: String, body: WebSearchBody) -> Self {
Self { key, body }
}
pub fn with_search_intent(mut self, enabled: bool) -> Self {
self.body = self.body.with_search_intent(enabled);
self
}
pub fn with_count(mut self, count: i32) -> Self {
self.body = self.body.with_count(count);
self
}
pub fn with_domain_filter(mut self, domain: String) -> Self {
self.body = self.body.with_domain_filter(domain);
self
}
pub fn with_recency_filter(mut self, filter: SearchRecencyFilter) -> Self {
self.body = self.body.with_recency_filter(filter);
self
}
pub fn with_content_size(mut self, size: ContentSize) -> Self {
self.body = self.body.with_content_size(size);
self
}
pub fn with_request_id(mut self, request_id: String) -> Self {
self.body = self.body.with_request_id(request_id);
self
}
pub fn with_user_id(mut self, user_id: String) -> Self {
self.body = self.body.with_user_id(user_id);
self
}
pub fn validate(&self) -> ZaiResult<()> {
self.body.validate_constraints()
}
pub async fn send(&self) -> ZaiResult<WebSearchResponse> {
self.validate()?;
let resp: reqwest::Response = self.post().await?;
let parsed = resp.json::<WebSearchResponse>().await?;
Ok(parsed)
}
}
#[async_trait]
impl HttpClient for WebSearchRequest {
type Body = WebSearchBody;
type ApiUrl = &'static str;
type ApiKey = String;
fn api_url(&self) -> &Self::ApiUrl {
&"https://open.bigmodel.cn/api/paas/v4/web_search"
}
fn api_key(&self) -> &Self::ApiKey {
&self.key
}
fn body(&self) -> &Self::Body {
&self.body
}
}