1use std::pin::Pin;
2
3use bytes::Bytes;
4use futures::{stream::StreamExt, Stream};
5use reqwest::multipart::Form;
6use reqwest_eventsource::{Event, EventSource, RequestBuilderExt};
7use serde::{de::DeserializeOwned, Serialize};
8
9use crate::{
10 config::{Config, OpenAIConfig},
11 error::{map_deserialization_error, ApiError, OpenAIError, WrappedError},
12 file::Files,
13 image::Images,
14 moderation::Moderations,
15 traits::AsyncTryFrom,
16 Assistants, Audio, AuditLogs, Batches, Chat, Completions, Embeddings, FineTuning, Invites,
17 Models, Projects, Responses, Threads, Uploads, Users, VectorStores,
18};
19
20#[derive(Debug, Clone, Default)]
21pub struct Client<C: Config> {
24 http_client: reqwest::Client,
25 config: C,
26 backoff: backoff::ExponentialBackoff,
27}
28
29impl Client<OpenAIConfig> {
30 pub fn new() -> Self {
32 Self::default()
33 }
34}
35
36impl<C: Config> Client<C> {
37 pub fn build(
39 http_client: reqwest::Client,
40 config: C,
41 backoff: backoff::ExponentialBackoff,
42 ) -> Self {
43 Self {
44 http_client,
45 config,
46 backoff,
47 }
48 }
49
50 pub fn with_config(config: C) -> Self {
52 Self {
53 http_client: reqwest::Client::new(),
54 config,
55 backoff: Default::default(),
56 }
57 }
58
59 pub fn with_http_client(mut self, http_client: reqwest::Client) -> Self {
63 self.http_client = http_client;
64 self
65 }
66
67 pub fn with_backoff(mut self, backoff: backoff::ExponentialBackoff) -> Self {
69 self.backoff = backoff;
70 self
71 }
72
73 pub fn models(&self) -> Models<C> {
77 Models::new(self)
78 }
79
80 pub fn completions(&self) -> Completions<C> {
82 Completions::new(self)
83 }
84
85 pub fn chat(&self) -> Chat<C> {
87 Chat::new(self)
88 }
89
90 pub fn images(&self) -> Images<C> {
92 Images::new(self)
93 }
94
95 pub fn moderations(&self) -> Moderations<C> {
97 Moderations::new(self)
98 }
99
100 pub fn files(&self) -> Files<C> {
102 Files::new(self)
103 }
104
105 pub fn uploads(&self) -> Uploads<C> {
107 Uploads::new(self)
108 }
109
110 pub fn fine_tuning(&self) -> FineTuning<C> {
112 FineTuning::new(self)
113 }
114
115 pub fn embeddings(&self) -> Embeddings<C> {
117 Embeddings::new(self)
118 }
119
120 pub fn audio(&self) -> Audio<C> {
122 Audio::new(self)
123 }
124
125 pub fn assistants(&self) -> Assistants<C> {
127 Assistants::new(self)
128 }
129
130 pub fn threads(&self) -> Threads<C> {
132 Threads::new(self)
133 }
134
135 pub fn vector_stores(&self) -> VectorStores<C> {
137 VectorStores::new(self)
138 }
139
140 pub fn batches(&self) -> Batches<C> {
142 Batches::new(self)
143 }
144
145 pub fn audit_logs(&self) -> AuditLogs<C> {
147 AuditLogs::new(self)
148 }
149
150 pub fn invites(&self) -> Invites<C> {
152 Invites::new(self)
153 }
154
155 pub fn users(&self) -> Users<C> {
157 Users::new(self)
158 }
159
160 pub fn projects(&self) -> Projects<C> {
162 Projects::new(self)
163 }
164
165 pub fn responses(&self) -> Responses<C> {
167 Responses::new(self)
168 }
169
170 pub fn config(&self) -> &C {
171 &self.config
172 }
173
174 pub(crate) async fn get<O>(&self, path: &str) -> Result<O, OpenAIError>
176 where
177 O: DeserializeOwned,
178 {
179 let request_maker = || async {
180 Ok(self
181 .http_client
182 .get(self.config.url(path))
183 .query(&self.config.query())
184 .headers(self.config.headers())
185 .build()?)
186 };
187
188 self.execute(request_maker).await
189 }
190
191 pub(crate) async fn get_with_query<Q, O>(&self, path: &str, query: &Q) -> Result<O, OpenAIError>
193 where
194 O: DeserializeOwned,
195 Q: Serialize + ?Sized,
196 {
197 let request_maker = || async {
198 Ok(self
199 .http_client
200 .get(self.config.url(path))
201 .query(&self.config.query())
202 .query(query)
203 .headers(self.config.headers())
204 .build()?)
205 };
206
207 self.execute(request_maker).await
208 }
209
210 pub(crate) async fn delete<O>(&self, path: &str) -> Result<O, OpenAIError>
212 where
213 O: DeserializeOwned,
214 {
215 let request_maker = || async {
216 Ok(self
217 .http_client
218 .delete(self.config.url(path))
219 .query(&self.config.query())
220 .headers(self.config.headers())
221 .build()?)
222 };
223
224 self.execute(request_maker).await
225 }
226
227 pub(crate) async fn get_raw(&self, path: &str) -> Result<Bytes, OpenAIError> {
229 let request_maker = || async {
230 Ok(self
231 .http_client
232 .get(self.config.url(path))
233 .query(&self.config.query())
234 .headers(self.config.headers())
235 .build()?)
236 };
237
238 self.execute_raw(request_maker).await
239 }
240
241 pub(crate) async fn post_raw<I>(&self, path: &str, request: I) -> Result<Bytes, OpenAIError>
243 where
244 I: Serialize,
245 {
246 let request_maker = || async {
247 Ok(self
248 .http_client
249 .post(self.config.url(path))
250 .query(&self.config.query())
251 .headers(self.config.headers())
252 .json(&request)
253 .build()?)
254 };
255
256 self.execute_raw(request_maker).await
257 }
258
259 pub(crate) async fn post<I, O>(&self, path: &str, request: I) -> Result<O, OpenAIError>
261 where
262 I: Serialize,
263 O: DeserializeOwned,
264 {
265 let request_maker = || async {
266 Ok(self
267 .http_client
268 .post(self.config.url(path))
269 .query(&self.config.query())
270 .headers(self.config.headers())
271 .json(&request)
272 .build()?)
273 };
274
275 self.execute(request_maker).await
276 }
277
278 pub(crate) async fn post_form_raw<F>(&self, path: &str, form: F) -> Result<Bytes, OpenAIError>
280 where
281 Form: AsyncTryFrom<F, Error = OpenAIError>,
282 F: Clone,
283 {
284 let request_maker = || async {
285 Ok(self
286 .http_client
287 .post(self.config.url(path))
288 .query(&self.config.query())
289 .headers(self.config.headers())
290 .multipart(<Form as AsyncTryFrom<F>>::try_from(form.clone()).await?)
291 .build()?)
292 };
293
294 self.execute_raw(request_maker).await
295 }
296
297 pub(crate) async fn post_form<O, F>(&self, path: &str, form: F) -> Result<O, OpenAIError>
299 where
300 O: DeserializeOwned,
301 Form: AsyncTryFrom<F, Error = OpenAIError>,
302 F: Clone,
303 {
304 let request_maker = || async {
305 Ok(self
306 .http_client
307 .post(self.config.url(path))
308 .query(&self.config.query())
309 .headers(self.config.headers())
310 .multipart(<Form as AsyncTryFrom<F>>::try_from(form.clone()).await?)
311 .build()?)
312 };
313
314 self.execute(request_maker).await
315 }
316
317 async fn execute_raw<M, Fut>(&self, request_maker: M) -> Result<Bytes, OpenAIError>
323 where
324 M: Fn() -> Fut,
325 Fut: core::future::Future<Output = Result<reqwest::Request, OpenAIError>>,
326 {
327 let client = self.http_client.clone();
328
329 backoff::future::retry(self.backoff.clone(), || async {
330 let request = request_maker().await.map_err(backoff::Error::Permanent)?;
331 let response = client
332 .execute(request)
333 .await
334 .map_err(OpenAIError::Reqwest)
335 .map_err(backoff::Error::Permanent)?;
336
337 let status = response.status();
338 let bytes = response
339 .bytes()
340 .await
341 .map_err(OpenAIError::Reqwest)
342 .map_err(backoff::Error::Permanent)?;
343
344 if status.is_server_error() {
345 let message: String = String::from_utf8_lossy(&bytes).into_owned();
347 tracing::warn!("Server error: {status} - {message}");
348 return Err(backoff::Error::Transient {
349 err: OpenAIError::ApiError(ApiError {
350 message,
351 r#type: None,
352 param: None,
353 code: None,
354 }),
355 retry_after: None,
356 });
357 }
358
359 if !status.is_success() {
361 let wrapped_error: WrappedError = serde_json::from_slice(bytes.as_ref())
362 .map_err(|e| map_deserialization_error(e, bytes.as_ref()))
363 .map_err(backoff::Error::Permanent)?;
364
365 if status.as_u16() == 429
366 && wrapped_error.error.r#type != Some("insufficient_quota".to_string())
369 {
370 tracing::warn!("Rate limited: {}", wrapped_error.error.message);
372 return Err(backoff::Error::Transient {
373 err: OpenAIError::ApiError(wrapped_error.error),
374 retry_after: None,
375 });
376 } else {
377 return Err(backoff::Error::Permanent(OpenAIError::ApiError(
378 wrapped_error.error,
379 )));
380 }
381 }
382
383 Ok(bytes)
384 })
385 .await
386 }
387
388 async fn execute<O, M, Fut>(&self, request_maker: M) -> Result<O, OpenAIError>
394 where
395 O: DeserializeOwned,
396 M: Fn() -> Fut,
397 Fut: core::future::Future<Output = Result<reqwest::Request, OpenAIError>>,
398 {
399 let bytes = self.execute_raw(request_maker).await?;
400
401 let response: O = serde_json::from_slice(bytes.as_ref())
402 .map_err(|e| map_deserialization_error(e, bytes.as_ref()))?;
403
404 Ok(response)
405 }
406
407 pub(crate) async fn post_stream<I, O>(
409 &self,
410 path: &str,
411 request: I,
412 ) -> Pin<Box<dyn Stream<Item = Result<O, OpenAIError>> + Send>>
413 where
414 I: Serialize,
415 O: DeserializeOwned + std::marker::Send + 'static,
416 {
417 let event_source = self
418 .http_client
419 .post(self.config.url(path))
420 .query(&self.config.query())
421 .headers(self.config.headers())
422 .json(&request)
423 .eventsource()
424 .unwrap();
425
426 stream(event_source).await
427 }
428
429 pub(crate) async fn post_stream_mapped_raw_events<I, O>(
430 &self,
431 path: &str,
432 request: I,
433 event_mapper: impl Fn(eventsource_stream::Event) -> Result<O, OpenAIError> + Send + 'static,
434 ) -> Pin<Box<dyn Stream<Item = Result<O, OpenAIError>> + Send>>
435 where
436 I: Serialize,
437 O: DeserializeOwned + std::marker::Send + 'static,
438 {
439 let event_source = self
440 .http_client
441 .post(self.config.url(path))
442 .query(&self.config.query())
443 .headers(self.config.headers())
444 .json(&request)
445 .eventsource()
446 .unwrap();
447
448 stream_mapped_raw_events(event_source, event_mapper).await
449 }
450
451 pub(crate) async fn _get_stream<Q, O>(
453 &self,
454 path: &str,
455 query: &Q,
456 ) -> Pin<Box<dyn Stream<Item = Result<O, OpenAIError>> + Send>>
457 where
458 Q: Serialize + ?Sized,
459 O: DeserializeOwned + std::marker::Send + 'static,
460 {
461 let event_source = self
462 .http_client
463 .get(self.config.url(path))
464 .query(query)
465 .query(&self.config.query())
466 .headers(self.config.headers())
467 .eventsource()
468 .unwrap();
469
470 stream(event_source).await
471 }
472}
473
474pub(crate) async fn stream<O>(
477 mut event_source: EventSource,
478) -> Pin<Box<dyn Stream<Item = Result<O, OpenAIError>> + Send>>
479where
480 O: DeserializeOwned + std::marker::Send + 'static,
481{
482 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
483
484 tokio::spawn(async move {
485 while let Some(ev) = event_source.next().await {
486 match ev {
487 Err(e) => {
488 if let Err(_e) = tx.send(Err(OpenAIError::StreamError(e.to_string()))) {
489 break;
491 }
492 }
493 Ok(event) => match event {
494 Event::Message(message) => {
495 if message.data == "[DONE]" {
496 break;
497 }
498
499 let response = match serde_json::from_str::<O>(&message.data) {
500 Err(e) => Err(map_deserialization_error(e, message.data.as_bytes())),
501 Ok(output) => Ok(output),
502 };
503
504 if let Err(_e) = tx.send(response) {
505 break;
507 }
508 }
509 Event::Open => continue,
510 },
511 }
512 }
513
514 event_source.close();
515 });
516
517 Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(rx))
518}
519
520pub(crate) async fn stream_mapped_raw_events<O>(
521 mut event_source: EventSource,
522 event_mapper: impl Fn(eventsource_stream::Event) -> Result<O, OpenAIError> + Send + 'static,
523) -> Pin<Box<dyn Stream<Item = Result<O, OpenAIError>> + Send>>
524where
525 O: DeserializeOwned + std::marker::Send + 'static,
526{
527 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
528
529 tokio::spawn(async move {
530 while let Some(ev) = event_source.next().await {
531 match ev {
532 Err(e) => {
533 if let Err(_e) = tx.send(Err(OpenAIError::StreamError(e.to_string()))) {
534 break;
536 }
537 }
538 Ok(event) => match event {
539 Event::Message(message) => {
540 let mut done = false;
541
542 if message.data == "[DONE]" {
543 done = true;
544 }
545
546 let response = event_mapper(message);
547
548 if let Err(_e) = tx.send(response) {
549 break;
551 }
552
553 if done {
554 break;
555 }
556 }
557 Event::Open => continue,
558 },
559 }
560 }
561
562 event_source.close();
563 });
564
565 Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(rx))
566}