ferrin_core/batch/
operations.rs

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

ferrin_core/batch/
operations.rs
