ferrin_core/batch/
operations.rs

1//! `get_batch_status`, `get_batch_results`, `cancel_batch` and `list_batches`.
2//!
3//! Derived from the Vercel AI SDK (Apache-2.0, Copyright 2023 Vercel, Inc.),
4//! translated from TypeScript to Rust and modified; see `NOTICE`.
5
6use std::future::IntoFuture;
7use std::sync::Arc;
8
9pub(super) use ferrin_spec::BatchId;
10use ferrin_spec::BatchRef;
11use ferrin_spec::BoxFuture;
12use ferrin_spec::batch::BatchCancelResult;
13use ferrin_spec::batch::BatchListOptions;
14use ferrin_spec::batch::BatchListResult;
15use ferrin_spec::batch::BatchOperationOptions;
16use ferrin_spec::batch::BatchStatus;
17use ferrin_spec::error::ProviderError;
18use ferrin_tool::ToolSet;
19use futures_util::StreamExt;
20use futures_util::stream;
21use tracing::Instrument;
22
23use crate::error::Error;
24use crate::modality::ModalityOptions;
25use crate::modality::impl_modality_builder;
26use crate::modality_stream::StreamDeadline;
27use crate::retry::retry;
28
29use super::result::BatchResults;
30use super::result::convert_item;
31use super::service_identity;
32use crate::telemetry::spans;
33
34macro_rules! batch_operation {
35    ($(#[$meta:meta])* $name:ident, $ty:ident) => {
36        $(#[$meta])*
37        #[must_use]
38        pub fn $name(batch: impl Into<BatchRef>, batch_id: impl Into<BatchId>) -> $ty {
39            $ty {
40                batch: batch.into(),
41                batch_id: batch_id.into(),
42                base: ModalityOptions::default(),
43            }
44        }
45    };
46}
47
48batch_operation!(
49    /// Fetches the normalized status of a batch.
50    get_batch_status,
51    GetBatchStatus
52);
53
54/// Builder returned by [`get_batch_status`]; `.await` runs the call.
55#[derive(Debug)]
56pub struct GetBatchStatus {
57    batch: BatchRef,
58    batch_id: BatchId,
59    base: ModalityOptions,
60}
61
62impl_modality_builder!(GetBatchStatus);
63
64impl IntoFuture for GetBatchStatus {
65    type Output = Result<BatchStatus, Error>;
66    type IntoFuture = BoxFuture<'static, Self::Output>;
67
68    fn into_future(self) -> Self::IntoFuture {
69        Box::pin(async move {
70            let identity = service_identity(&self.batch);
71            let span = spans::modality_span("get_batch_status", &identity);
72            let batch = self.batch.clone();
73            let base = self.base.clone();
74            let batch_id = self.batch_id;
75            base.run(|base, token| {
76                async move {
77                    let headers = base.request_headers();
78                    retry(&base.retry_policy, &token, |_| {
79                        let options = BatchOperationOptions {
80                            batch_id: batch_id.clone(),
81                            provider_options: base.provider_options.clone(),
82                            headers: headers.clone(),
83                            cancellation: token.child_token(),
84                        };
85                        let batch = &batch;
86                        async move {
87                            batch
88                                .do_get_batch_status(options)
89                                .await
90                                .map_err(Error::from)
91                        }
92                    })
93                    .await
94                }
95                .instrument(span)
96            })
97            .await
98        })
99    }
100}
101
102/// Streams the results of a finished batch.
103#[must_use]
104pub fn get_batch_results(
105    batch: impl Into<BatchRef>,
106    batch_id: impl Into<BatchId>,
107) -> GetBatchResults {
108    GetBatchResults {
109        batch: batch.into(),
110        batch_id: batch_id.into(),
111        tools: ToolSet::new(),
112        base: ModalityOptions::default(),
113    }
114}
115
116/// Builder returned by [`get_batch_results`]; `.await` opens the stream.
117#[derive(Debug)]
118pub struct GetBatchResults {
119    batch: BatchRef,
120    batch_id: BatchId,
121    tools: ToolSet,
122    base: ModalityOptions,
123}
124
125impl GetBatchResults {
126    /// Tools used to parse and validate tool calls in text results.
127    #[must_use]
128    pub fn tools(mut self, tools: ToolSet) -> Self {
129        self.tools = tools;
130        self
131    }
132}
133
134impl_modality_builder!(GetBatchResults);
135
136impl IntoFuture for GetBatchResults {
137    type Output = Result<BatchResults, Error>;
138    type IntoFuture = BoxFuture<'static, Self::Output>;
139
140    fn into_future(self) -> Self::IntoFuture {
141        Box::pin(async move {
142            let identity = service_identity(&self.batch);
143            let span = spans::modality_span("get_batch_results", &identity);
144            let batch = self.batch.clone();
145            let base = self.base.clone();
146            let batch_id = self.batch_id;
147            let tools = Arc::new(self.tools);
148            let deadline = StreamDeadline::new(&base.cancellation, base.timeout);
149            let token = deadline.cancellation.clone();
150            let stream = deadline
151                .run(
152                    async move {
153                        let headers = base.request_headers();
154                        retry(&base.retry_policy, &token, |_| {
155                            let options = BatchOperationOptions {
156                                batch_id: batch_id.clone(),
157                                provider_options: base.provider_options.clone(),
158                                headers: headers.clone(),
159                                cancellation: token.child_token(),
160                            };
161                            let batch = &batch;
162                            async move {
163                                batch
164                                    .do_get_batch_results(options)
165                                    .await
166                                    .map_err(Error::from)
167                            }
168                        })
169                        .await
170                    }
171                    .instrument(span),
172                )
173                .await?;
174            let converted = stream::unfold(Some((deadline, stream, tools)), |state| async move {
175                let (deadline, mut stream, tools) = state?;
176                let item = deadline
177                    .run(async {
178                        match stream.next().await {
179                            Some(Ok(item)) => Ok(Some(convert_item(item, &tools).await)),
180                            Some(Err(error)) => Err(Error::from(error)),
181                            None => Ok(None),
182                        }
183                    })
184                    .await;
185                match item {
186                    Ok(Some(item)) => Some((Ok(item), Some((deadline, stream, tools)))),
187                    Ok(None) => None,
188                    Err(error) => Some((Err(error), None)),
189                }
190            });
191            Ok(Box::pin(converted) as BatchResults)
192        })
193    }
194}
195
196batch_operation!(
197    /// Cancels a batch. Fails with an unsupported functionality error when
198    /// the provider does not implement cancellation.
199    cancel_batch,
200    CancelBatch
201);
202
203/// Builder returned by [`cancel_batch`]; `.await` runs the call.
204#[derive(Debug)]
205pub struct CancelBatch {
206    batch: BatchRef,
207    batch_id: BatchId,
208    base: ModalityOptions,
209}
210
211impl_modality_builder!(@no_retry CancelBatch);
212
213impl IntoFuture for CancelBatch {
214    type Output = Result<BatchCancelResult, Error>;
215    type IntoFuture = BoxFuture<'static, Self::Output>;
216
217    fn into_future(self) -> Self::IntoFuture {
218        Box::pin(async move {
219            let identity = service_identity(&self.batch);
220            let span = spans::modality_span("cancel_batch", &identity);
221            if !self.batch.supports_cancel_batch() {
222                return Err(Error::from(ProviderError::unsupported(
223                    "batch cancellation",
224                )));
225            }
226            let batch = self.batch.clone();
227            let base = self.base.clone();
228            let batch_id = self.batch_id;
229            base.run(|base, token| {
230                async move {
231                    batch
232                        .do_cancel_batch(BatchOperationOptions {
233                            batch_id,
234                            provider_options: base.provider_options.clone(),
235                            headers: base.request_headers(),
236                            cancellation: token,
237                        })
238                        .await
239                        .map_err(Error::from)
240                }
241                .instrument(span)
242            })
243            .await
244        })
245    }
246}
247
248/// Lists batches. Fails with an unsupported functionality error when the
249/// provider does not implement listing.
250#[must_use]
251pub fn list_batches(batch: impl Into<BatchRef>) -> ListBatches {
252    ListBatches {
253        batch: batch.into(),
254        limit: None,
255        cursor: None,
256        base: ModalityOptions::default(),
257    }
258}
259
260/// Builder returned by [`list_batches`]; `.await` runs the call.
261#[derive(Debug)]
262pub struct ListBatches {
263    batch: BatchRef,
264    limit: Option<usize>,
265    cursor: Option<String>,
266    base: ModalityOptions,
267}
268
269impl ListBatches {
270    /// Page size.
271    #[must_use]
272    pub fn limit(mut self, limit: usize) -> Self {
273        self.limit = Some(limit);
274        self
275    }
276
277    /// Cursor of the next page.
278    #[must_use]
279    pub fn cursor(mut self, cursor: impl Into<String>) -> Self {
280        self.cursor = Some(cursor.into());
281        self
282    }
283}
284
285impl_modality_builder!(ListBatches);
286
287impl IntoFuture for ListBatches {
288    type Output = Result<BatchListResult, Error>;
289    type IntoFuture = BoxFuture<'static, Self::Output>;
290
291    fn into_future(self) -> Self::IntoFuture {
292        Box::pin(async move {
293            let identity = service_identity(&self.batch);
294            let span = spans::modality_span("list_batches", &identity);
295            if !self.batch.supports_list_batches() {
296                return Err(Error::from(ProviderError::unsupported("batch listing")));
297            }
298            let batch = self.batch.clone();
299            let base = self.base.clone();
300            let limit = self.limit;
301            let cursor = self.cursor;
302            base.run(|base, token| {
303                async move {
304                    let headers = base.request_headers();
305                    retry(&base.retry_policy, &token, |_| {
306                        let options = BatchListOptions {
307                            limit,
308                            cursor: cursor.clone(),
309                            provider_options: base.provider_options.clone(),
310                            headers: headers.clone(),
311                            cancellation: token.child_token(),
312                        };
313                        let batch = &batch;
314                        async move { batch.do_list_batches(options).await.map_err(Error::from) }
315                    })
316                    .await
317                }
318                .instrument(span)
319            })
320            .await
321        })
322    }
323}
ferrin_core/batch/operations.rs

ferrin_core/batch/
operations.rs
