1use reqwest::header::HeaderValue;
2use serde::Serialize;
3
4use crate::{
5 config::{RecordReplayMode, REPLAY_LOG},
6 KagiConfig, KagiError, KagiResult,
7};
8
9pub struct KagiClient {
10 pub config: KagiConfig,
11 pub inner_client: reqwest_middleware::ClientWithMiddleware,
12}
13
14impl KagiClient {
15 pub fn new(config: KagiConfig) -> Self {
17 let reqwest_client = reqwest::Client::new();
18 let client_builder = reqwest_middleware::ClientBuilder::new(reqwest_client);
19 let inner_client = client_builder.build();
20
21 Self {
22 config,
23 inner_client,
24 }
25 }
26
27 pub fn with_replay(
35 config: KagiConfig,
36 reqwest_client: reqwest::Client,
37 record_replay_mode: RecordReplayMode,
38 ) -> KagiResult<Self> {
39 let client_builder = reqwest_middleware::ClientBuilder::new(reqwest_client);
40 let Some(vcr_logfile) = config.logdir.as_ref().map(|path| path.join(REPLAY_LOG)) else {
41 let inner_client = client_builder.build();
42 let client = KagiClient {
43 config,
44 inner_client,
45 };
46 return Ok(client);
47 };
48
49 let middleware = rvcr::VCRMiddleware::try_from(vcr_logfile.to_path_buf())
50 .map_err(KagiError::VCRMiddlewareError)?
51 .with_mode(rvcr::VCRMode::from(record_replay_mode));
52
53 let inner_client = client_builder.with(middleware).build();
54 let client = KagiClient {
55 config,
56 inner_client,
57 };
58 Ok(client)
59 }
60
61 pub async fn query<Q>(&self, endpoint: &str, request_body: Q) -> KagiResult<reqwest::Response>
62 where
63 Q: Serialize,
64 {
65 let request_url = reqwest::Url::parse(&format!("{}{}", self.config.api_base, endpoint))
68 .map_err(KagiError::UrlError)?;
69 let content_type: Result<HeaderValue, _> = "application/json".parse();
70 let authorization: Result<HeaderValue, _> = format!("Bot {}", self.config.api_key).parse();
71
72 let request = self
74 .inner_client
75 .post(request_url)
76 .header("Content-Type", content_type.unwrap())
77 .header("Authorization", authorization.unwrap())
78 .json(&request_body);
79
80 let response = request
81 .send()
82 .await
83 .map_err(KagiError::ReqwestMiddlewareError)?;
84
85 if response.status() != reqwest::StatusCode::OK {
86 return Err(KagiError::StatusCodeError(response.status()));
87 }
88
89 Ok(response)
90 }
91}