1use std::sync::Arc;
2#[cfg(not(target_family = "wasm"))]
3use std::sync::Mutex;
4
5use bytes::Bytes;
6use futures::stream::StreamExt;
7use reqwest::{header::HeaderMap, multipart::Form, Response};
8use serde::{de::DeserializeOwned, Serialize};
9
10use crate::error::StreamError;
11#[cfg(feature = "middleware")]
12use crate::executor::TowerExecutor;
13use crate::{
14 config::{Config, OpenAIConfig},
15 error::{map_deserialization_error, ApiError, ApiErrorResponse, OpenAIError, WrappedError},
16 executor::{HttpRequestFactory, ReqwestExecutor, SharedExecutor},
17 traits::AsyncTryFrom,
18 RequestOptions,
19};
20
21struct RequestParts {
22 request_client: reqwest::Client,
23 method: reqwest::Method,
24 url: String,
25 headers: HeaderMap,
26 query: Vec<(String, String)>,
27}
28
29impl RequestParts {
30 fn build_request_builder(&self) -> reqwest::RequestBuilder {
31 self.request_client
32 .request(self.method.clone(), self.url.clone())
33 .query(&self.query)
34 .headers(self.headers.clone())
35 }
36}
37
38#[cfg(feature = "administration")]
39use crate::admin::Admin;
40#[cfg(feature = "chatkit")]
41use crate::chatkit::Chatkit;
42#[cfg(feature = "file")]
43use crate::file::Files;
44#[cfg(feature = "image")]
45use crate::image::Images;
46#[cfg(feature = "moderation")]
47use crate::moderation::Moderations;
48#[cfg(feature = "audio")]
49use crate::Audio;
50#[cfg(feature = "batch")]
51use crate::Batches;
52#[cfg(feature = "chat-completion")]
53use crate::Chat;
54#[cfg(feature = "completions")]
55use crate::Completions;
56#[cfg(feature = "container")]
57use crate::Containers;
58#[cfg(feature = "content-provenance-checks")]
59use crate::ContentProvenanceChecks;
60#[cfg(feature = "responses")]
61use crate::Conversations;
62#[cfg(feature = "embedding")]
63use crate::Embeddings;
64#[cfg(feature = "evals")]
65use crate::Evals;
66#[cfg(feature = "finetuning")]
67use crate::FineTuning;
68#[cfg(feature = "model")]
69use crate::Models;
70#[cfg(feature = "realtime")]
71use crate::Realtime;
72#[cfg(feature = "responses")]
73use crate::Responses;
74#[cfg(feature = "safety")]
75use crate::Safety;
76#[cfg(feature = "skill")]
77use crate::Skills;
78#[cfg(feature = "upload")]
79use crate::Uploads;
80#[cfg(feature = "vectorstore")]
81use crate::VectorStores;
82#[cfg(feature = "video")]
83#[allow(deprecated)]
84use crate::Videos;
85
86#[derive(Clone)]
87pub struct Client<C: Config> {
90 request_client: reqwest::Client,
91 executor: SharedExecutor,
92 config: C,
93}
94
95impl<C> std::fmt::Debug for Client<C>
96where
97 C: Config + std::fmt::Debug,
98{
99 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100 f.debug_struct("Client")
101 .field("request_client", &self.request_client)
102 .field("config", &self.config)
103 .finish()
104 }
105}
106
107impl<C: Config> Default for Client<C>
108where
109 C: Default,
110{
111 fn default() -> Self {
112 let request_client = reqwest::Client::new();
113 Self {
114 executor: Arc::new(ReqwestExecutor::new(request_client.clone())),
115 request_client,
116 config: C::default(),
117 }
118 }
119}
120
121impl Client<OpenAIConfig> {
122 pub fn new() -> Self {
124 Self::default()
125 }
126}
127
128impl<C: Config> Client<C> {
129 pub fn build(http_client: reqwest::Client, config: C) -> Self {
131 Self {
132 executor: Arc::new(ReqwestExecutor::new(http_client.clone())),
133 request_client: http_client,
134 config,
135 }
136 }
137
138 pub fn with_config(config: C) -> Self {
140 let request_client = reqwest::Client::new();
141 Self {
142 executor: Arc::new(ReqwestExecutor::new(request_client.clone())),
143 request_client,
144 config,
145 }
146 }
147
148 pub fn with_http_client(mut self, http_client: reqwest::Client) -> Self {
152 self.executor = Arc::new(ReqwestExecutor::new(http_client.clone()));
153 self.request_client = http_client;
154 self
155 }
156
157 #[cfg(all(feature = "middleware", not(target_family = "wasm")))]
159 pub fn with_http_service<S>(mut self, service: S) -> Self
160 where
161 S: tower::Service<HttpRequestFactory, Response = Response> + Clone + Send + Sync + 'static,
162 S::Future: Send + 'static,
163 S::Error: Into<OpenAIError> + Send + Sync + 'static,
164 {
165 self.executor = Arc::new(TowerExecutor::new(service));
170 self
171 }
172
173 #[cfg(all(feature = "middleware", target_family = "wasm"))]
175 pub fn with_http_service<S>(mut self, service: S) -> Self
176 where
177 S: tower::Service<HttpRequestFactory, Response = Response> + Clone + 'static,
178 S::Future: 'static,
179 S::Error: Into<OpenAIError> + 'static,
180 {
181 self.executor = Arc::new(TowerExecutor::new(service));
186 self
187 }
188
189 #[cfg(feature = "model")]
193 pub fn models(&self) -> Models<'_, C> {
194 Models::new(self)
195 }
196
197 #[cfg(feature = "completions")]
199 pub fn completions(&self) -> Completions<'_, C> {
200 Completions::new(self)
201 }
202
203 #[cfg(feature = "chat-completion")]
205 pub fn chat(&self) -> Chat<'_, C> {
206 Chat::new(self)
207 }
208
209 #[cfg(feature = "image")]
211 pub fn images(&self) -> Images<'_, C> {
212 Images::new(self)
213 }
214
215 #[cfg(feature = "moderation")]
217 pub fn moderations(&self) -> Moderations<'_, C> {
218 Moderations::new(self)
219 }
220
221 #[cfg(feature = "file")]
223 pub fn files(&self) -> Files<'_, C> {
224 Files::new(self)
225 }
226
227 #[cfg(feature = "upload")]
229 pub fn uploads(&self) -> Uploads<'_, C> {
230 Uploads::new(self)
231 }
232
233 #[cfg(feature = "finetuning")]
235 pub fn fine_tuning(&self) -> FineTuning<'_, C> {
236 FineTuning::new(self)
237 }
238
239 #[cfg(feature = "embedding")]
241 pub fn embeddings(&self) -> Embeddings<'_, C> {
242 Embeddings::new(self)
243 }
244
245 #[cfg(feature = "audio")]
247 pub fn audio(&self) -> Audio<'_, C> {
248 Audio::new(self)
249 }
250
251 #[cfg(feature = "video")]
253 #[deprecated(note = "The Videos API is deprecated.")]
254 #[allow(deprecated)]
255 pub fn videos(&self) -> Videos<'_, C> {
256 Videos::new(self)
257 }
258
259 #[cfg(feature = "vectorstore")]
261 pub fn vector_stores(&self) -> VectorStores<'_, C> {
262 VectorStores::new(self)
263 }
264
265 #[cfg(feature = "batch")]
267 pub fn batches(&self) -> Batches<'_, C> {
268 Batches::new(self)
269 }
270
271 #[cfg(feature = "administration")]
274 pub fn admin(&self) -> Admin<'_, C> {
275 Admin::new(self)
276 }
277
278 #[cfg(feature = "responses")]
280 pub fn responses(&self) -> Responses<'_, C> {
281 Responses::new(self)
282 }
283
284 #[cfg(feature = "responses")]
286 pub fn conversations(&self) -> Conversations<'_, C> {
287 Conversations::new(self)
288 }
289
290 #[cfg(feature = "container")]
292 pub fn containers(&self) -> Containers<'_, C> {
293 Containers::new(self)
294 }
295
296 #[cfg(feature = "skill")]
298 pub fn skills(&self) -> Skills<'_, C> {
299 Skills::new(self)
300 }
301
302 #[cfg(feature = "evals")]
304 pub fn evals(&self) -> Evals<'_, C> {
305 Evals::new(self)
306 }
307
308 #[cfg(feature = "chatkit")]
309 pub fn chatkit(&self) -> Chatkit<'_, C> {
310 Chatkit::new(self)
311 }
312
313 #[cfg(feature = "realtime")]
315 pub fn realtime(&self) -> Realtime<'_, C> {
316 Realtime::new(self)
317 }
318
319 #[cfg(feature = "decisions")]
321 pub fn decisions(&self) -> crate::Decisions<'_, C> {
322 crate::Decisions::new(self)
323 }
324
325 #[cfg(feature = "safety")]
327 pub fn safety(&self) -> crate::Safety<'_, C> {
328 Safety::new(self)
329 }
330
331 #[cfg(feature = "content-provenance-checks")]
333 pub fn content_provenance_checks(&self) -> crate::ContentProvenanceChecks<'_, C> {
334 ContentProvenanceChecks::new(self)
335 }
336
337 pub fn config(&self) -> &C {
338 &self.config
339 }
340
341 fn build_request_parts(
342 &self,
343 method: reqwest::Method,
344 path: &str,
345 request_options: &RequestOptions,
346 ) -> Arc<RequestParts> {
347 let url = if let Some(path) = request_options.path() {
348 self.config.url(path.as_str())
349 } else {
350 self.config.url(path)
351 };
352 let mut headers = self.config.headers();
353 if let Some(request_headers) = request_options.headers() {
354 headers.extend(request_headers.clone());
355 }
356
357 let mut query = self
358 .config
359 .query()
360 .into_iter()
361 .map(|(key, value)| (key.to_string(), value.to_string()))
362 .collect::<Vec<_>>();
363 query.extend_from_slice(request_options.query());
364
365 Arc::new(RequestParts {
366 request_client: self.request_client.clone(),
367 method,
368 url,
369 headers,
370 query,
371 })
372 }
373
374 fn build_request_factory(
375 &self,
376 method: reqwest::Method,
377 path: &str,
378 request_options: &RequestOptions,
379 ) -> HttpRequestFactory {
380 let request_parts = self.build_request_parts(method, path, request_options);
381
382 HttpRequestFactory::new(move || {
383 let request_parts = request_parts.clone();
384
385 async move {
386 let request = request_parts.build_request_builder().build()?;
387 Ok(request)
388 }
389 })
390 }
391
392 fn build_request_factory_with_json<I>(
393 &self,
394 method: reqwest::Method,
395 path: &str,
396 request: I,
397 request_options: &RequestOptions,
398 ) -> Result<HttpRequestFactory, OpenAIError>
399 where
400 I: Serialize,
401 {
402 let request = Bytes::from(serde_json::to_vec(&request).map_err(|error| {
405 OpenAIError::InvalidArgument(format!("failed to serialize request: {error}"))
406 })?);
407 let request_parts = self.build_request_parts(method, path, request_options);
408
409 Ok(HttpRequestFactory::new(move || {
410 let request_parts = request_parts.clone();
411 let request = request.clone();
412
413 async move {
414 let request_builder = request_parts
415 .build_request_builder()
416 .header(reqwest::header::CONTENT_TYPE, "application/json")
417 .body(request.clone());
418
419 Ok(request_builder.build()?)
420 }
421 }))
422 }
423
424 fn build_request_factory_with_form<F>(
425 &self,
426 method: reqwest::Method,
427 path: &str,
428 form: F,
429 request_options: &RequestOptions,
430 ) -> Result<HttpRequestFactory, OpenAIError>
431 where
432 F: Clone + crate::traits::MaybeSend + 'static,
433 Form: AsyncTryFrom<F, Error = OpenAIError>,
434 {
435 #[cfg(not(target_family = "wasm"))]
439 let form = Arc::new(Mutex::new(form));
440 let request_parts = self.build_request_parts(method, path, request_options);
441
442 Ok(HttpRequestFactory::new(move || {
443 let request_parts = request_parts.clone();
444 let form = form.clone();
445
446 async move {
447 #[cfg(not(target_family = "wasm"))]
448 let form = form
449 .lock()
450 .expect("multipart request factory mutex poisoned")
451 .clone();
452 #[cfg(target_family = "wasm")]
453 let form = form.clone();
454 let form = <Form as AsyncTryFrom<F>>::try_from(form).await?;
455 let request_builder = request_parts.build_request_builder().multipart(form);
456
457 Ok(request_builder.build()?)
458 }
459 }))
460 }
461
462 #[allow(unused)]
464 pub(crate) async fn get<O>(
465 &self,
466 path: &str,
467 request_options: &RequestOptions,
468 ) -> Result<O, OpenAIError>
469 where
470 O: DeserializeOwned,
471 {
472 let request_factory =
473 self.build_request_factory(reqwest::Method::GET, path, request_options);
474 self.execute(request_factory).await
475 }
476
477 #[allow(unused)]
479 pub(crate) async fn delete<O>(
480 &self,
481 path: &str,
482 request_options: &RequestOptions,
483 ) -> Result<O, OpenAIError>
484 where
485 O: DeserializeOwned,
486 {
487 let request_factory =
488 self.build_request_factory(reqwest::Method::DELETE, path, request_options);
489 self.execute(request_factory).await
490 }
491
492 #[allow(unused)]
494 pub(crate) async fn get_raw(
495 &self,
496 path: &str,
497 request_options: &RequestOptions,
498 ) -> Result<(Bytes, HeaderMap), OpenAIError> {
499 let request_factory =
500 self.build_request_factory(reqwest::Method::GET, path, request_options);
501 self.execute_raw(request_factory).await
502 }
503
504 #[allow(unused)]
506 pub(crate) async fn post_raw<I>(
507 &self,
508 path: &str,
509 request: I,
510 request_options: &RequestOptions,
511 ) -> Result<(Bytes, HeaderMap), OpenAIError>
512 where
513 I: Serialize,
514 {
515 let request_factory = self.build_request_factory_with_json(
516 reqwest::Method::POST,
517 path,
518 request,
519 request_options,
520 )?;
521 self.execute_raw(request_factory).await
522 }
523
524 #[allow(unused)]
526 pub(crate) async fn post<I, O>(
527 &self,
528 path: &str,
529 request: I,
530 request_options: &RequestOptions,
531 ) -> Result<O, OpenAIError>
532 where
533 I: Serialize,
534 O: DeserializeOwned,
535 {
536 let request_factory = self.build_request_factory_with_json(
537 reqwest::Method::POST,
538 path,
539 request,
540 request_options,
541 )?;
542 self.execute(request_factory).await
543 }
544
545 #[allow(unused)]
547 pub(crate) async fn post_form_raw<F>(
548 &self,
549 path: &str,
550 form: F,
551 request_options: &RequestOptions,
552 ) -> Result<(Bytes, HeaderMap), OpenAIError>
553 where
554 F: Clone + crate::traits::MaybeSend + 'static,
555 Form: AsyncTryFrom<F, Error = OpenAIError>,
556 {
557 let request_factory = self.build_request_factory_with_form(
558 reqwest::Method::POST,
559 path,
560 form,
561 request_options,
562 )?;
563 self.execute_raw(request_factory).await
564 }
565
566 #[allow(unused)]
568 pub(crate) async fn post_form<O, F>(
569 &self,
570 path: &str,
571 form: F,
572 request_options: &RequestOptions,
573 ) -> Result<O, OpenAIError>
574 where
575 O: DeserializeOwned,
576 F: Clone + crate::traits::MaybeSend + 'static,
577 Form: AsyncTryFrom<F, Error = OpenAIError>,
578 {
579 let request_factory = self.build_request_factory_with_form(
580 reqwest::Method::POST,
581 path,
582 form,
583 request_options,
584 )?;
585 self.execute(request_factory).await
586 }
587
588 #[allow(unused)]
589 pub(crate) async fn post_form_stream<O, F>(
590 &self,
591 path: &str,
592 form: F,
593 request_options: &RequestOptions,
594 ) -> Result<crate::types::stream::StreamResponse<O>, OpenAIError>
595 where
596 F: Clone + crate::traits::MaybeSend + 'static,
597 Form: AsyncTryFrom<F, Error = OpenAIError>,
598 O: DeserializeOwned + crate::traits::MaybeSend + 'static,
599 {
600 let request_factory = self.build_request_factory_with_form(
601 reqwest::Method::POST,
602 path,
603 form,
604 request_options,
605 )?;
606
607 self.execute_stream(request_factory).await
608 }
609
610 async fn execute_raw(
611 &self,
612 request_factory: HttpRequestFactory,
613 ) -> Result<(Bytes, HeaderMap), OpenAIError> {
614 let response = self.execute_response(request_factory).await?;
615 read_response(response).await
616 }
617
618 async fn execute<O>(&self, request_factory: HttpRequestFactory) -> Result<O, OpenAIError>
619 where
620 O: DeserializeOwned,
621 {
622 let (bytes, _headers) = self.execute_raw(request_factory).await?;
623
624 let response: O = serde_json::from_slice(bytes.as_ref())
625 .map_err(|e| map_deserialization_error(e, bytes.as_ref()))?;
626
627 Ok(response)
628 }
629
630 async fn execute_response(
631 &self,
632 request_factory: HttpRequestFactory,
633 ) -> Result<Response, OpenAIError> {
634 let response = self.executor.execute(request_factory).await?;
635 if !response.status().is_success() {
636 return Err(read_error_response(response).await);
637 }
638 Ok(response)
639 }
640
641 async fn execute_stream<O>(
642 &self,
643 request_factory: HttpRequestFactory,
644 ) -> Result<crate::types::stream::StreamResponse<O>, OpenAIError>
645 where
646 O: DeserializeOwned + crate::traits::MaybeSend + 'static,
647 {
648 let response = self.execute_response(request_factory).await?;
649 Ok(stream(response).await)
650 }
651
652 async fn execute_stream_mapped_raw_events<O>(
653 &self,
654 request_factory: HttpRequestFactory,
655 event_mapper: impl Fn(eventsource_stream::Event) -> Result<O, OpenAIError>
656 + crate::traits::MaybeSend
657 + 'static,
658 ) -> Result<crate::types::stream::StreamResponse<O>, OpenAIError>
659 where
660 O: DeserializeOwned + crate::traits::MaybeSend + 'static,
661 {
662 let response = self.execute_response(request_factory).await?;
663 Ok(stream_mapped_raw_events(response, event_mapper).await)
664 }
665
666 #[allow(unused)]
668 pub(crate) async fn post_stream<I, O>(
669 &self,
670 path: &str,
671 request: I,
672 request_options: &RequestOptions,
673 ) -> Result<crate::types::stream::StreamResponse<O>, OpenAIError>
674 where
675 I: Serialize,
676 O: DeserializeOwned + crate::traits::MaybeSend + 'static,
677 {
678 let request_factory = self.build_request_factory_with_json(
679 reqwest::Method::POST,
680 path,
681 request,
682 request_options,
683 )?;
684 self.execute_stream(request_factory).await
687 }
688
689 #[allow(unused)]
690 pub(crate) async fn post_stream_mapped_raw_events<I, O>(
691 &self,
692 path: &str,
693 request: I,
694 request_options: &RequestOptions,
695 event_mapper: impl Fn(eventsource_stream::Event) -> Result<O, OpenAIError>
696 + crate::traits::MaybeSend
697 + 'static,
698 ) -> Result<crate::types::stream::StreamResponse<O>, OpenAIError>
699 where
700 I: Serialize,
701 O: DeserializeOwned + crate::traits::MaybeSend + 'static,
702 {
703 let request_factory = self.build_request_factory_with_json(
704 reqwest::Method::POST,
705 path,
706 request,
707 request_options,
708 )?;
709 self.execute_stream_mapped_raw_events(request_factory, event_mapper)
710 .await
711 }
712
713 #[allow(unused)]
715 pub(crate) async fn get_stream<O>(
716 &self,
717 path: &str,
718 request_options: &RequestOptions,
719 ) -> Result<crate::types::stream::StreamResponse<O>, OpenAIError>
720 where
721 O: DeserializeOwned + crate::traits::MaybeSend + 'static,
722 {
723 let request_factory =
724 self.build_request_factory(reqwest::Method::GET, path, request_options);
725 self.execute_stream(request_factory).await
726 }
727}
728
729async fn read_response(response: Response) -> Result<(Bytes, HeaderMap), OpenAIError> {
730 let headers = response.headers().clone();
731 let bytes = response.bytes().await.map_err(OpenAIError::Reqwest)?;
732 Ok((bytes, headers))
733}
734
735async fn read_error_response(response: Response) -> OpenAIError {
736 let status = response.status();
737 let bytes = match response.bytes().await {
738 Ok(b) => b,
739 Err(e) => return OpenAIError::Reqwest(e),
740 };
741
742 if status.is_server_error() {
743 let message: String = String::from_utf8_lossy(&bytes).into_owned();
745 tracing::warn!("Server error: {status} - {message}");
746 return OpenAIError::ApiError(ApiErrorResponse {
747 status_code: status,
748 api_error: ApiError {
749 misalignment: None,
750 message,
751 r#type: None,
752 param: None,
753 code: None,
754 },
755 });
756 }
757
758 match serde_json::from_slice::<WrappedError>(bytes.as_ref()) {
760 Ok(wrapped) => OpenAIError::ApiError(ApiErrorResponse {
761 status_code: status,
762 api_error: wrapped.error,
763 }),
764 Err(e) => map_deserialization_error(e, bytes.as_ref()),
765 }
766}
767
768pub(crate) async fn stream<O>(response: Response) -> crate::types::stream::StreamResponse<O>
771where
772 O: DeserializeOwned + crate::traits::MaybeSend + 'static,
773{
774 stream_mapped_raw_events(response, |event| {
775 serde_json::from_str::<O>(&event.data)
776 .map_err(|error| map_deserialization_error(error, event.data.as_bytes()))
777 })
778 .await
779}
780
781#[cfg(target_family = "wasm")]
782pub(crate) async fn stream_mapped_raw_events<O>(
783 response: Response,
784 event_mapper: impl Fn(eventsource_stream::Event) -> Result<O, OpenAIError> + 'static,
785) -> crate::types::stream::StreamResponse<O>
786where
787 O: DeserializeOwned + 'static,
788{
789 let byte_stream = response
790 .bytes_stream()
791 .map(|result| result.map_err(std::io::Error::other));
792 let event_stream = Box::pin(eventsource_stream::EventStream::new(byte_stream));
793
794 Box::pin(futures::stream::unfold(
795 (event_stream, event_mapper),
796 |(mut event_stream, event_mapper)| async move {
797 loop {
798 let event = match event_stream.next().await {
799 Some(Ok(event)) => event,
800 Some(Err(error)) => {
801 return Some((
802 Err(OpenAIError::StreamError(Box::new(
803 StreamError::EventStream(error.to_string()),
804 ))),
805 (event_stream, event_mapper),
806 ));
807 }
808 None => return None,
809 };
810
811 if event.data == "[DONE]" {
812 return None;
813 }
814
815 if event.event == "keepalive" {
816 continue;
817 }
818
819 let response = event_mapper(event);
820 return Some((response, (event_stream, event_mapper)));
821 }
822 },
823 ))
824}
825
826#[cfg(not(target_family = "wasm"))]
827pub(crate) async fn stream_mapped_raw_events<O>(
828 response: Response,
829 event_mapper: impl Fn(eventsource_stream::Event) -> Result<O, OpenAIError> + Send + 'static,
830) -> crate::types::stream::StreamResponse<O>
831where
832 O: DeserializeOwned + std::marker::Send + 'static,
833{
834 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
835
836 tokio::spawn(async move {
837 let byte_stream = response
838 .bytes_stream()
839 .map(|r| r.map_err(std::io::Error::other));
840 let mut event_stream = std::pin::pin!(eventsource_stream::EventStream::new(byte_stream));
841
842 while let Some(ev) = tokio::select! {
847 biased;
848 _ = tx.closed() => None,
849 ev = event_stream.next() => ev,
850 } {
851 let event = match ev {
852 Ok(e) => e,
853 Err(e) => {
854 let _ = tx.send(Err(OpenAIError::StreamError(Box::new(
855 StreamError::EventStream(e.to_string()),
856 ))));
857 break;
858 }
859 };
860 if event.data == "[DONE]" {
861 break;
862 }
863
864 if event.event == "keepalive" {
865 continue;
866 }
867
868 let response = event_mapper(event);
869
870 if tx.send(response).is_err() {
871 break;
872 }
873 }
874 });
875
876 Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(rx))
877}
878
879#[cfg(all(test, feature = "middleware", not(target_family = "wasm")))]
880mod tests {
881 use std::sync::{
882 atomic::{AtomicBool, AtomicUsize, Ordering},
883 Arc,
884 };
885
886 use futures::StreamExt;
887 use http::Response as HttpResponse;
888 use serde_json::json;
889 use tower::{service_fn, ServiceBuilder};
890
891 use super::Client;
892 use crate::{
893 config::OpenAIConfig, error::OpenAIError, executor::HttpRequestFactory,
894 retry::SimpleRetryPolicy, traits::AsyncTryFrom, RequestOptions,
895 };
896
897 #[tokio::test]
898 async fn unary_requests_dispatch_through_middleware_service() {
899 let request_count = Arc::new(AtomicUsize::new(0));
900 let service = {
901 let request_count = request_count.clone();
902 ServiceBuilder::new()
903 .concurrency_limit(1)
904 .service(service_fn(move |factory: HttpRequestFactory| {
905 let request_count = request_count.clone();
906 async move {
907 let request = factory.build().await?;
908 assert_eq!(request.url().path(), "/models");
909 request_count.fetch_add(1, Ordering::SeqCst);
910 Ok::<reqwest::Response, OpenAIError>(
911 HttpResponse::builder()
912 .status(200)
913 .header("content-type", "application/json")
914 .body(reqwest::Body::from(
915 "{\"object\":\"list\",\"data\":[{\"id\":\"model\"}]}",
916 ))
917 .unwrap()
918 .into(),
919 )
920 }
921 }))
922 };
923
924 let client = Client::with_config(
925 OpenAIConfig::new()
926 .with_api_base("http://example.test")
927 .with_api_key("test-key"),
928 )
929 .with_http_service(service);
930
931 let value: serde_json::Value = client.get("/models", &RequestOptions::new()).await.unwrap();
932
933 assert_eq!(value["object"], "list");
934 assert_eq!(request_count.load(Ordering::SeqCst), 1);
935 }
936
937 #[tokio::test]
938 async fn stream_requests_open_through_middleware_service() {
939 let request_count = Arc::new(AtomicUsize::new(0));
940 let service = {
941 let request_count = request_count.clone();
942 ServiceBuilder::new()
943 .concurrency_limit(1)
944 .service(service_fn(move |factory: HttpRequestFactory| {
945 let request_count = request_count.clone();
946 async move {
947 let request = factory.build().await?;
948 assert_eq!(request.url().path(), "/responses");
949 request_count.fetch_add(1, Ordering::SeqCst);
950 Ok::<reqwest::Response, OpenAIError>(
951 HttpResponse::builder()
952 .status(200)
953 .header("content-type", "text/event-stream")
954 .body(reqwest::Body::from(
955 "data: {\"ok\":true}\n\ndata: [DONE]\n\n",
956 ))
957 .unwrap()
958 .into(),
959 )
960 }
961 }))
962 };
963
964 let client = Client::with_config(
965 OpenAIConfig::new()
966 .with_api_base("http://example.test")
967 .with_api_key("test-key"),
968 )
969 .with_http_service(service);
970
971 let mut stream = client
972 .post_stream::<_, serde_json::Value>(
973 "/responses",
974 json!({ "stream": true }),
975 &RequestOptions::new(),
976 )
977 .await
978 .unwrap();
979
980 let first = stream.next().await.unwrap().unwrap();
981
982 assert_eq!(first, json!({ "ok": true }));
983 assert_eq!(request_count.load(Ordering::SeqCst), 1);
984 }
985
986 #[tokio::test]
987 async fn middleware_retry_policy_retries_429_responses() {
988 let request_count = Arc::new(AtomicUsize::new(0));
989 let service = {
990 let request_count = request_count.clone();
991 ServiceBuilder::new()
992 .retry(SimpleRetryPolicy::default())
993 .service(service_fn(move |factory: HttpRequestFactory| {
994 let request_count = request_count.clone();
995 async move {
996 let request = factory.build().await?;
997 assert_eq!(request.url().path(), "/models");
998 let attempt = request_count.fetch_add(1, Ordering::SeqCst);
999
1000 let response = if attempt == 0 {
1001 HttpResponse::builder()
1002 .status(429)
1003 .header("content-type", "application/json")
1004 .body(reqwest::Body::from(
1005 r#"{"error":{"message":"retry me","type":"rate_limit_error","param":null,"code":null}}"#,
1006 ))
1007 .unwrap()
1008 } else {
1009 HttpResponse::builder()
1010 .status(200)
1011 .header("content-type", "application/json")
1012 .body(reqwest::Body::from(
1013 r#"{"object":"list","data":[{"id":"retry-model"}]}"#,
1014 ))
1015 .unwrap()
1016 };
1017
1018 Ok::<reqwest::Response, OpenAIError>(response.into())
1019 }
1020 }))
1021 };
1022
1023 let client = Client::with_config(
1024 OpenAIConfig::new()
1025 .with_api_base("http://example.test")
1026 .with_api_key("test-key"),
1027 )
1028 .with_http_service(service);
1029
1030 let value: serde_json::Value = client.get("/models", &RequestOptions::new()).await.unwrap();
1031
1032 assert_eq!(value["data"][0]["id"], "retry-model");
1033 assert_eq!(request_count.load(Ordering::SeqCst), 2);
1034 }
1035
1036 #[derive(Clone)]
1037 struct RetryableMultipartInput {
1038 conversions: Arc<AtomicUsize>,
1039 }
1040
1041 impl AsyncTryFrom<RetryableMultipartInput> for reqwest::multipart::Form {
1042 type Error = OpenAIError;
1043
1044 async fn try_from(value: RetryableMultipartInput) -> Result<Self, Self::Error> {
1045 value.conversions.fetch_add(1, Ordering::SeqCst);
1046 Ok(reqwest::multipart::Form::new().text("field", "value"))
1047 }
1048 }
1049
1050 #[tokio::test]
1051 async fn middleware_retry_policy_rebuilds_multipart_form_per_attempt() {
1052 let request_count = Arc::new(AtomicUsize::new(0));
1053 let conversion_count = Arc::new(AtomicUsize::new(0));
1054
1055 let service = {
1056 let request_count = request_count.clone();
1057 ServiceBuilder::new()
1058 .retry(SimpleRetryPolicy::default())
1059 .service(service_fn(move |factory: HttpRequestFactory| {
1060 let request_count = request_count.clone();
1061 async move {
1062 let request = factory.build().await?;
1063 assert_eq!(request.method(), reqwest::Method::POST);
1064 assert_eq!(request.url().path(), "/files");
1065 let attempt = request_count.fetch_add(1, Ordering::SeqCst);
1066
1067 let response = if attempt == 0 {
1068 HttpResponse::builder()
1069 .status(429)
1070 .header("content-type", "application/json")
1071 .body(reqwest::Body::from(
1072 r#"{"error":{"message":"retry me","type":"rate_limit_error","param":null,"code":null}}"#,
1073 ))
1074 .unwrap()
1075 } else {
1076 HttpResponse::builder()
1077 .status(200)
1078 .header("content-type", "application/json")
1079 .body(reqwest::Body::from(r#"{"ok":true}"#))
1080 .unwrap()
1081 };
1082
1083 Ok::<reqwest::Response, OpenAIError>(response.into())
1084 }
1085 }))
1086 };
1087
1088 let client = Client::with_config(
1089 OpenAIConfig::new()
1090 .with_api_base("http://example.test")
1091 .with_api_key("test-key"),
1092 )
1093 .with_http_service(service);
1094
1095 let value: serde_json::Value = client
1096 .post_form(
1097 "/files",
1098 RetryableMultipartInput {
1099 conversions: conversion_count.clone(),
1100 },
1101 &RequestOptions::new(),
1102 )
1103 .await
1104 .unwrap();
1105
1106 assert_eq!(value, json!({ "ok": true }));
1107 assert_eq!(request_count.load(Ordering::SeqCst), 2);
1108 assert_eq!(conversion_count.load(Ordering::SeqCst), 2);
1109 }
1110
1111 #[tokio::test]
1112 async fn dropping_stream_releases_idle_upstream_response() {
1113 struct DropGuard(Arc<AtomicBool>);
1114
1115 impl Drop for DropGuard {
1116 fn drop(&mut self) {
1117 self.0.store(true, Ordering::SeqCst);
1118 }
1119 }
1120
1121 let upstream_dropped = Arc::new(AtomicBool::new(false));
1122
1123 let dropped = upstream_dropped.clone();
1124 let service = ServiceBuilder::new()
1125 .concurrency_limit(1)
1126 .service(service_fn(move |factory: HttpRequestFactory| {
1127 let guard = DropGuard(dropped.clone());
1128 async move {
1129 factory.build().await?;
1130
1131 let body = futures::stream::once(async {
1134 Ok::<_, std::io::Error>(bytes::Bytes::from_static(
1135 b"data: {\"ok\":true}\n\n",
1136 ))
1137 })
1138 .chain(futures::stream::unfold(
1139 guard,
1140 |guard| async move {
1141 futures::future::pending::<()>().await;
1142 Some((Ok(bytes::Bytes::new()), guard))
1143 },
1144 ));
1145
1146 Ok::<reqwest::Response, OpenAIError>(
1147 HttpResponse::builder()
1148 .status(200)
1149 .header("content-type", "text/event-stream")
1150 .body(reqwest::Body::wrap_stream(body))
1151 .unwrap()
1152 .into(),
1153 )
1154 }
1155 }));
1156
1157 let client = Client::with_config(
1158 OpenAIConfig::new()
1159 .with_api_base("http://example.test")
1160 .with_api_key("test-key"),
1161 )
1162 .with_http_service(service);
1163
1164 let mut stream = client
1165 .post_stream::<_, serde_json::Value>(
1166 "/responses",
1167 json!({ "stream": true }),
1168 &RequestOptions::new(),
1169 )
1170 .await
1171 .unwrap();
1172
1173 assert_eq!(stream.next().await.unwrap().unwrap(), json!({ "ok": true }));
1174 assert!(!upstream_dropped.load(Ordering::SeqCst));
1175
1176 drop(stream);
1177
1178 for _ in 0..100 {
1181 if upstream_dropped.load(Ordering::SeqCst) {
1182 break;
1183 }
1184 tokio::task::yield_now().await;
1185 }
1186
1187 assert!(
1188 upstream_dropped.load(Ordering::SeqCst),
1189 "reader task leaked the upstream response after the stream was dropped"
1190 );
1191 }
1192}