omni_llm_kit/http_client/
http_client.rs

1
2pub use anyhow::{Result, anyhow};
3use derive_more::Deref;
4pub use http::{self, Method, Request, Response, StatusCode, Uri};
5
6use futures::future::BoxFuture;
7use http::request::Builder;
8#[cfg(feature = "test-support")]
9use std::fmt;
10use std::{
11    any::type_name,
12    sync::{Arc, Mutex},
13};
14pub use url::Url;
15use crate::http_client::AsyncBody;
16
17#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
18pub enum RedirectPolicy {
19    #[default]
20    NoFollow,
21    FollowLimit(u32),
22    FollowAll,
23}
24pub struct FollowRedirects(pub bool);
25
26pub trait HttpRequestExt {
27    /// Whether or not to follow redirects
28    fn follow_redirects(self, follow: RedirectPolicy) -> Self;
29}
30
31impl HttpRequestExt for http::request::Builder {
32    fn follow_redirects(self, follow: RedirectPolicy) -> Self {
33        self.extension(follow)
34    }
35}
36
37pub trait HttpClient: 'static + Send + Sync {
38    fn type_name(&self) -> &'static str;
39
40    fn send(
41        &self,
42        req: http::Request<AsyncBody>,
43    ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>>;
44
45    fn get<'a>(
46        &'a self,
47        uri: &str,
48        body: AsyncBody,
49        follow_redirects: bool,
50    ) -> BoxFuture<'a, anyhow::Result<Response<AsyncBody>>> {
51        let request = Builder::new()
52            .uri(uri)
53            .follow_redirects(if follow_redirects {
54                RedirectPolicy::FollowAll
55            } else {
56                RedirectPolicy::NoFollow
57            })
58            .body(body);
59
60        match request {
61            Ok(request) => Box::pin(async move { self.send(request).await }),
62            Err(e) => Box::pin(async move { Err(e.into()) }),
63        }
64    }
65
66    fn post_json<'a>(
67        &'a self,
68        uri: &str,
69        body: AsyncBody,
70    ) -> BoxFuture<'a, anyhow::Result<Response<AsyncBody>>> {
71        let request = Builder::new()
72            .uri(uri)
73            .method(Method::POST)
74            .header("Content-Type", "application/json")
75            .body(body);
76
77        match request {
78            Ok(request) => Box::pin(async move { self.send(request).await }),
79            Err(e) => Box::pin(async move { Err(e.into()) }),
80        }
81    }
82
83    fn proxy(&self) -> Option<&Url>;
84}
85
86/// An [`HttpClient`] that may have a proxy.
87#[derive(Deref)]
88pub struct HttpClientWithProxy {
89    #[deref]
90    client: Arc<dyn HttpClient>,
91    proxy: Option<Url>,
92}
93
94impl HttpClientWithProxy {
95    /// Returns a new [`HttpClientWithProxy`] with the given proxy URL.
96    pub fn new(client: Arc<dyn HttpClient>, proxy_url: Option<String>) -> Self {
97        let proxy_url = proxy_url
98            .and_then(|proxy| proxy.parse().ok())
99            .or_else(read_proxy_from_env);
100
101        Self::new_url(client, proxy_url)
102    }
103    pub fn new_url(client: Arc<dyn HttpClient>, proxy_url: Option<Url>) -> Self {
104        Self {
105            client,
106            proxy: proxy_url,
107        }
108    }
109}
110
111impl HttpClient for HttpClientWithProxy {
112    fn send(
113        &self,
114        req: Request<AsyncBody>,
115    ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {
116        self.client.send(req)
117    }
118
119    fn proxy(&self) -> Option<&Url> {
120        self.proxy.as_ref()
121    }
122
123    fn type_name(&self) -> &'static str {
124        self.client.type_name()
125    }
126}
127
128impl HttpClient for Arc<HttpClientWithProxy> {
129    fn send(
130        &self,
131        req: Request<AsyncBody>,
132    ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {
133        self.client.send(req)
134    }
135
136    fn proxy(&self) -> Option<&Url> {
137        self.proxy.as_ref()
138    }
139
140    fn type_name(&self) -> &'static str {
141        self.client.type_name()
142    }
143}
144
145/// An [`HttpClient`] that has a base URL.
146pub struct HttpClientWithUrl {
147    base_url: Mutex<String>,
148    client: HttpClientWithProxy,
149}
150
151impl std::ops::Deref for HttpClientWithUrl {
152    type Target = HttpClientWithProxy;
153
154    fn deref(&self) -> &Self::Target {
155        &self.client
156    }
157}
158
159impl HttpClientWithUrl {
160    /// Returns a new [`HttpClientWithUrl`] with the given base URL.
161    pub fn new(
162        client: Arc<dyn HttpClient>,
163        base_url: impl Into<String>,
164        proxy_url: Option<String>,
165    ) -> Self {
166        let client = HttpClientWithProxy::new(client, proxy_url);
167
168        Self {
169            base_url: Mutex::new(base_url.into()),
170            client,
171        }
172    }
173
174    pub fn new_url(
175        client: Arc<dyn HttpClient>,
176        base_url: impl Into<String>,
177        proxy_url: Option<Url>,
178    ) -> Self {
179        let client = HttpClientWithProxy::new_url(client, proxy_url);
180
181        Self {
182            base_url: Mutex::new(base_url.into()),
183            client,
184        }
185    }
186
187    /// Returns the base URL.
188    pub fn base_url(&self) -> String {
189        self.base_url
190            .lock()
191            .map_or_else(|_| Default::default(), |url| url.clone())
192    }
193
194    /// Sets the base URL.
195    pub fn set_base_url(&self, base_url: impl Into<String>) {
196        let base_url = base_url.into();
197        self.base_url
198            .lock()
199            .map(|mut url| {
200                *url = base_url;
201            })
202            .ok();
203    }
204
205    /// Builds a URL using the given path.
206    pub fn build_url(&self, path: &str) -> String {
207        format!("{}{}", self.base_url(), path)
208    }
209
210}
211
212impl HttpClient for Arc<HttpClientWithUrl> {
213    fn send(
214        &self,
215        req: Request<AsyncBody>,
216    ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {
217        self.client.send(req)
218    }
219
220    fn proxy(&self) -> Option<&Url> {
221        self.client.proxy.as_ref()
222    }
223
224    fn type_name(&self) -> &'static str {
225        self.client.type_name()
226    }
227}
228
229impl HttpClient for HttpClientWithUrl {
230    fn send(
231        &self,
232        req: Request<AsyncBody>,
233    ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {
234        self.client.send(req)
235    }
236
237    fn proxy(&self) -> Option<&Url> {
238        self.client.proxy.as_ref()
239    }
240
241    fn type_name(&self) -> &'static str {
242        self.client.type_name()
243    }
244}
245
246pub fn read_proxy_from_env() -> Option<Url> {
247    const ENV_VARS: &[&str] = &[
248        "ALL_PROXY",
249        "all_proxy",
250        "HTTPS_PROXY",
251        "https_proxy",
252        "HTTP_PROXY",
253        "http_proxy",
254    ];
255
256    ENV_VARS
257        .iter()
258        .find_map(|var| std::env::var(var).ok())
259        .and_then(|env| env.parse().ok())
260}
261
262pub struct BlockedHttpClient;
263
264impl BlockedHttpClient {
265    pub fn new() -> Self {
266        BlockedHttpClient
267    }
268}
269
270impl HttpClient for BlockedHttpClient {
271    fn send(
272        &self,
273        _req: Request<AsyncBody>,
274    ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {
275        Box::pin(async {
276            Err(std::io::Error::new(
277                std::io::ErrorKind::PermissionDenied,
278                "BlockedHttpClient disallowed request",
279            )
280            .into())
281        })
282    }
283
284    fn proxy(&self) -> Option<&Url> {
285        None
286    }
287
288    fn type_name(&self) -> &'static str {
289        type_name::<Self>()
290    }
291}
292
293#[cfg(feature = "test-support")]
294type FakeHttpHandler = Box<
295    dyn Fn(Request<AsyncBody>) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>>
296        + Send
297        + Sync
298        + 'static,
299>;
300
301#[cfg(feature = "test-support")]
302pub struct FakeHttpClient {
303    handler: FakeHttpHandler,
304}
305
306#[cfg(feature = "test-support")]
307impl FakeHttpClient {
308    pub fn create<Fut, F>(handler: F) -> Arc<HttpClientWithUrl>
309    where
310        Fut: futures::Future<Output = anyhow::Result<Response<AsyncBody>>> + Send + 'static,
311        F: Fn(Request<AsyncBody>) -> Fut + Send + Sync + 'static,
312    {
313        Arc::new(HttpClientWithUrl {
314            base_url: Mutex::new("http://test.example".into()),
315            client: HttpClientWithProxy {
316                client: Arc::new(Self {
317                    handler: Box::new(move |req| Box::pin(handler(req))),
318                }),
319                proxy: None,
320            },
321        })
322    }
323
324    pub fn with_404_response() -> Arc<HttpClientWithUrl> {
325        Self::create(|_| async move {
326            Ok(Response::builder()
327                .status(404)
328                .body(Default::default())
329                .unwrap())
330        })
331    }
332
333    pub fn with_200_response() -> Arc<HttpClientWithUrl> {
334        Self::create(|_| async move {
335            Ok(Response::builder()
336                .status(200)
337                .body(Default::default())
338                .unwrap())
339        })
340    }
341}
342
343#[cfg(feature = "test-support")]
344impl fmt::Debug for FakeHttpClient {
345    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
346        f.debug_struct("FakeHttpClient").finish()
347    }
348}
349
350#[cfg(feature = "test-support")]
351impl HttpClient for FakeHttpClient {
352    fn send(
353        &self,
354        req: Request<AsyncBody>,
355    ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {
356        let future = (self.handler)(req);
357        future
358    }
359
360    fn proxy(&self) -> Option<&Url> {
361        None
362    }
363
364    fn type_name(&self) -> &'static str {
365        type_name::<Self>()
366    }
367}
omni_llm_kit/http_client/http_client.rs

omni_llm_kit/http_client/
http_client.rs
