Skip to main content

async_openai/
client.rs

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)]
87/// Client is a container for config and HTTP execution
88/// used to make API calls.
89pub 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    /// Client with default [OpenAIConfig]
123    pub fn new() -> Self {
124        Self::default()
125    }
126}
127
128impl<C: Config> Client<C> {
129    /// Create client with a custom HTTP client and config.
130    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    /// Create client with [OpenAIConfig] or [crate::config::AzureConfig]
139    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    /// Provide your own [client] to make HTTP requests with.
149    ///
150    /// [client]: reqwest::Client
151    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    /// Provide your own tower-compatible service to execute HTTP requests.
158    #[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        // This is the public middleware escape hatch. We erase the concrete
166        // tower stack here so the rest of the client does not become generic
167        // over the service type, which would otherwise leak through every API
168        // group and make the crate much harder to use.
169        self.executor = Arc::new(TowerExecutor::new(service));
170        self
171    }
172
173    /// Provide your own tower-compatible service to execute HTTP requests.
174    #[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        // wasm futures produced by reqwest are not `Send`, so the wasm version
182        // intentionally avoids native thread-safety bounds. Users are still
183        // responsible for choosing tower layers that work in their wasm
184        // runtime.
185        self.executor = Arc::new(TowerExecutor::new(service));
186        self
187    }
188
189    // API groups
190
191    /// To call [Models] group related APIs using this client.
192    #[cfg(feature = "model")]
193    pub fn models(&self) -> Models<'_, C> {
194        Models::new(self)
195    }
196
197    /// To call [Completions] group related APIs using this client.
198    #[cfg(feature = "completions")]
199    pub fn completions(&self) -> Completions<'_, C> {
200        Completions::new(self)
201    }
202
203    /// To call [Chat] group related APIs using this client.
204    #[cfg(feature = "chat-completion")]
205    pub fn chat(&self) -> Chat<'_, C> {
206        Chat::new(self)
207    }
208
209    /// To call [Images] group related APIs using this client.
210    #[cfg(feature = "image")]
211    pub fn images(&self) -> Images<'_, C> {
212        Images::new(self)
213    }
214
215    /// To call [Moderations] group related APIs using this client.
216    #[cfg(feature = "moderation")]
217    pub fn moderations(&self) -> Moderations<'_, C> {
218        Moderations::new(self)
219    }
220
221    /// To call [Files] group related APIs using this client.
222    #[cfg(feature = "file")]
223    pub fn files(&self) -> Files<'_, C> {
224        Files::new(self)
225    }
226
227    /// To call [Uploads] group related APIs using this client.
228    #[cfg(feature = "upload")]
229    pub fn uploads(&self) -> Uploads<'_, C> {
230        Uploads::new(self)
231    }
232
233    /// To call [FineTuning] group related APIs using this client.
234    #[cfg(feature = "finetuning")]
235    pub fn fine_tuning(&self) -> FineTuning<'_, C> {
236        FineTuning::new(self)
237    }
238
239    /// To call [Embeddings] group related APIs using this client.
240    #[cfg(feature = "embedding")]
241    pub fn embeddings(&self) -> Embeddings<'_, C> {
242        Embeddings::new(self)
243    }
244
245    /// To call [Audio] group related APIs using this client.
246    #[cfg(feature = "audio")]
247    pub fn audio(&self) -> Audio<'_, C> {
248        Audio::new(self)
249    }
250
251    /// To call [Videos] group related APIs using this client.
252    #[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    /// To call [VectorStores] group related APIs using this client.
260    #[cfg(feature = "vectorstore")]
261    pub fn vector_stores(&self) -> VectorStores<'_, C> {
262        VectorStores::new(self)
263    }
264
265    /// To call [Batches] group related APIs using this client.
266    #[cfg(feature = "batch")]
267    pub fn batches(&self) -> Batches<'_, C> {
268        Batches::new(self)
269    }
270
271    /// To call [Admin] group related APIs using this client.
272    /// This groups together admin API keys, invites, users, projects, audit logs, and certificates.
273    #[cfg(feature = "administration")]
274    pub fn admin(&self) -> Admin<'_, C> {
275        Admin::new(self)
276    }
277
278    /// To call [Responses] group related APIs using this client.
279    #[cfg(feature = "responses")]
280    pub fn responses(&self) -> Responses<'_, C> {
281        Responses::new(self)
282    }
283
284    /// To call [Conversations] group related APIs using this client.
285    #[cfg(feature = "responses")]
286    pub fn conversations(&self) -> Conversations<'_, C> {
287        Conversations::new(self)
288    }
289
290    /// To call [Containers] group related APIs using this client.
291    #[cfg(feature = "container")]
292    pub fn containers(&self) -> Containers<'_, C> {
293        Containers::new(self)
294    }
295
296    /// To call [Skills] group related APIs using this client.
297    #[cfg(feature = "skill")]
298    pub fn skills(&self) -> Skills<'_, C> {
299        Skills::new(self)
300    }
301
302    /// To call [Evals] group related APIs using this client.
303    #[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    /// To call [Realtime] group related APIs using this client.
314    #[cfg(feature = "realtime")]
315    pub fn realtime(&self) -> Realtime<'_, C> {
316        Realtime::new(self)
317    }
318
319    /// Classify or score inputs with the Decisions API.
320    #[cfg(feature = "decisions")]
321    pub fn decisions(&self) -> crate::Decisions<'_, C> {
322        crate::Decisions::new(self)
323    }
324
325    /// To call [Safety] group related APIs.
326    #[cfg(feature = "safety")]
327    pub fn safety(&self) -> crate::Safety<'_, C> {
328        Safety::new(self)
329    }
330
331    /// To call [ContentProvenanceChecks] group related APIs
332    #[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        // JSON bodies are materialized once so the base BYOT path can keep
403        // accepting borrowed inputs.
404        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        // Multipart is the reason the factory exists.
436        //
437        // `Mutex` is only here to make the captured state `Sync` on native targets.
438        #[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    /// Make a GET request to {path} and deserialize the response body
463    #[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    /// Make a DELETE request to {path} and deserialize the response body
478    #[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    /// Make a GET request to {path} and return the response body
493    #[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    /// Make a POST request to {path} and return the response body
505    #[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    /// Make a POST request to {path} and deserialize the response body
525    #[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    /// POST a form at {path} and return the response body
546    #[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    /// POST a form at {path} and deserialize the response body
567    #[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    /// Make HTTP POST request to receive SSE
667    #[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        // Stream setup is still request/response first. We only create the SSE
685        // stream after the HTTP layer has returned a response object.
686        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    /// Make HTTP GET request to receive SSE
714    #[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        // OpenAI does not guarantee server errors are returned as JSON so we cannot deserialize them.
744        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    // Deserialize response body from the error object
759    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
768/// Request which responds with SSE.
769/// [server-sent events](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#event_stream_format)
770pub(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        // Also observe the consumer dropping the stream: relying on
843        // `tx.send(..).is_err()` alone would keep this task - and the upstream
844        // response it holds - alive until the next event arrives, which never
845        // happens when upstream is open but idle.
846        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                    // One event, then an upstream which stays open but never sends
1132                    // again - no further events, no `[DONE]`, no error.
1133                    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        // The reader task should observe the dropped consumer and release the
1179        // response without waiting for another event from upstream.
1180        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}