kagi_api/
client.rs

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    /// Create a [KagiClient] for a given [KagiConfig] and [reqwest::Client].
16    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    /// Create a [KagiClient] for a given [KagiConfig] and [reqwest::Client].
28    ///
29    /// Use [rvcr::VCRMiddleware] to log request/response to disk.
30    ///
31    /// FIXME: Mention `KAGI_API_LOGDIR`.
32    ///
33    /// FIXME: Avoid panics from `rvcr` by fixing them upstream.
34    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        // FIXME: Use `url::Url` combinators instead of `format!("{}{}")`.
66        // https://docs.rs/url/latest/url/struct.Url.html
67        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        // FIXME: Don't `.unwrap()`, use `KagiError::HeaderError(String)`.
73        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}