ferrin_spec/
batch.rs

1//! Provider batch processing interface.
2
3use std::future::Future;
4
5use chrono::DateTime;
6use chrono::Utc;
7use serde::Deserialize;
8use serde::Serialize;
9use tokio_util::sync::CancellationToken;
10use url::Url;
11
12use crate::dynamic::BoxStream;
13use crate::error::ProviderError;
14use crate::image_model::AspectRatio;
15use crate::image_model::ImageFile;
16use crate::image_model::ImageResult;
17use crate::image_model::ImageSize;
18use crate::language_model::GenerateResult;
19use crate::language_model::Prompt;
20use crate::language_model::ReasoningEffort;
21use crate::language_model::ResponseFormat;
22use crate::language_model::SupportedUrls;
23use crate::language_model::ToolChoice;
24use crate::language_model::ToolDefinition;
25use crate::shared::BatchId;
26use crate::shared::Headers;
27use crate::shared::ModelId;
28use crate::shared::ProviderId;
29use crate::shared::ProviderMetadata;
30use crate::shared::ProviderOptions;
31use crate::shared::Warning;
32
33/// Stream of per-request results of a finished batch.
34pub type BatchResultStream = BoxStream<'static, Result<BatchItemResult, ProviderError>>;
35
36/// Submit many text or image requests as one provider batch job.
37pub trait Batch: Send + Sync + 'static {
38    /// Provider identifier.
39    fn provider(&self) -> &ProviderId;
40
41    /// URL patterns the provider can fetch itself (see `LanguageModel`).
42    fn supported_urls(&self) -> impl Future<Output = SupportedUrls> + Send;
43
44    /// Starts a batch job.
45    fn do_start_batch(
46        &self,
47        options: BatchStartOptions,
48    ) -> impl Future<Output = Result<BatchStartResult, ProviderError>> + Send;
49
50    /// Returns the normalized status of a batch job.
51    fn do_get_batch_status(
52        &self,
53        options: BatchOperationOptions,
54    ) -> impl Future<Output = Result<BatchStatus, ProviderError>> + Send;
55
56    /// Streams the results of a finished batch job.
57    fn do_get_batch_results(
58        &self,
59        options: BatchOperationOptions,
60    ) -> impl Future<Output = Result<BatchResultStream, ProviderError>> + Send;
61
62    /// Whether [`do_cancel_batch`](Self::do_cancel_batch) is implemented.
63    fn supports_cancel_batch(&self) -> bool {
64        false
65    }
66
67    /// Cancels a batch job.
68    fn do_cancel_batch(
69        &self,
70        options: BatchOperationOptions,
71    ) -> impl Future<Output = Result<BatchCancelResult, ProviderError>> + Send {
72        let _ = options;
73        std::future::ready(Err(ProviderError::unsupported("cancel_batch")))
74    }
75
76    /// Whether [`do_list_batches`](Self::do_list_batches) is implemented.
77    fn supports_list_batches(&self) -> bool {
78        false
79    }
80
81    /// Lists batch jobs.
82    fn do_list_batches(
83        &self,
84        options: BatchListOptions,
85    ) -> impl Future<Output = Result<BatchListResult, ProviderError>> + Send {
86        let _ = options;
87        std::future::ready(Err(ProviderError::unsupported("list_batches")))
88    }
89}
90
91/// Settings of a text request inside a batch (no headers, no cancellation).
92#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
93pub struct TextBatchRequestOptions {
94    /// The prompt.
95    pub prompt: Prompt,
96    /// Maximum output tokens.
97    #[serde(default, skip_serializing_if = "Option::is_none")]
98    pub max_output_tokens: Option<u32>,
99    /// Temperature.
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub temperature: Option<f64>,
102    /// Stop sequences.
103    #[serde(default, skip_serializing_if = "Option::is_none")]
104    pub stop_sequences: Option<Vec<String>>,
105    /// Top-p.
106    #[serde(default, skip_serializing_if = "Option::is_none")]
107    pub top_p: Option<f64>,
108    /// Top-k.
109    #[serde(default, skip_serializing_if = "Option::is_none")]
110    pub top_k: Option<u32>,
111    /// Presence penalty.
112    #[serde(default, skip_serializing_if = "Option::is_none")]
113    pub presence_penalty: Option<f64>,
114    /// Frequency penalty.
115    #[serde(default, skip_serializing_if = "Option::is_none")]
116    pub frequency_penalty: Option<f64>,
117    /// Seed.
118    #[serde(default, skip_serializing_if = "Option::is_none")]
119    pub seed: Option<u64>,
120    /// Reasoning effort.
121    #[serde(default)]
122    pub reasoning: ReasoningEffort,
123    /// Response format.
124    #[serde(default, skip_serializing_if = "Option::is_none")]
125    pub response_format: Option<ResponseFormat>,
126    /// Tool choice.
127    #[serde(default, skip_serializing_if = "Option::is_none")]
128    pub tool_choice: Option<ToolChoice>,
129    /// Tools.
130    #[serde(default, skip_serializing_if = "Vec::is_empty")]
131    pub tools: Vec<ToolDefinition>,
132    /// Provider-specific options.
133    #[serde(default, skip_serializing_if = "ProviderOptions::is_empty")]
134    pub provider_options: ProviderOptions,
135}
136
137/// Settings of an image request inside a batch.
138#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
139pub struct ImageBatchRequestOptions {
140    /// Prompt.
141    #[serde(default, skip_serializing_if = "Option::is_none")]
142    pub prompt: Option<String>,
143    /// Number of images.
144    pub n: u32,
145    /// Size.
146    #[serde(default, skip_serializing_if = "Option::is_none")]
147    pub size: Option<ImageSize>,
148    /// Aspect ratio.
149    #[serde(default, skip_serializing_if = "Option::is_none")]
150    pub aspect_ratio: Option<AspectRatio>,
151    /// Seed.
152    #[serde(default, skip_serializing_if = "Option::is_none")]
153    pub seed: Option<u64>,
154    /// Input files.
155    #[serde(default, skip_serializing_if = "Vec::is_empty")]
156    pub files: Vec<ImageFile>,
157    /// Mask.
158    #[serde(default, skip_serializing_if = "Option::is_none")]
159    pub mask: Option<ImageFile>,
160    /// Provider-specific options.
161    #[serde(default, skip_serializing_if = "ProviderOptions::is_empty")]
162    pub provider_options: ProviderOptions,
163}
164
165/// One request of a batch, tagged by `type`.
166#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
167#[serde(tag = "type", rename_all = "lowercase")]
168#[non_exhaustive]
169pub enum BatchRequest {
170    /// A text generation request.
171    Text {
172        /// Caller-assigned request id, echoed in results.
173        id: String,
174        /// Model to use.
175        model_id: ModelId,
176        /// Settings.
177        options: TextBatchRequestOptions,
178    },
179    /// An image generation request.
180    Image {
181        /// Caller-assigned request id, echoed in results.
182        id: String,
183        /// Model to use.
184        model_id: ModelId,
185        /// Settings.
186        options: ImageBatchRequestOptions,
187    },
188}
189
190impl BatchRequest {
191    /// Returns the caller-assigned request id.
192    #[must_use]
193    pub fn id(&self) -> &str {
194        match self {
195            Self::Text { id, .. } | Self::Image { id, .. } => id,
196        }
197    }
198}
199
200/// Options for starting a batch.
201#[derive(Debug, Clone)]
202pub struct BatchStartOptions {
203    /// Requests to submit.
204    pub requests: Vec<BatchRequest>,
205    /// Webhook URL to notify on completion.
206    pub webhook_url: Option<Url>,
207    /// Provider-specific options keyed by provider name.
208    pub provider_options: ProviderOptions,
209    /// Additional request headers.
210    pub headers: Headers,
211    /// Cancellation token.
212    pub cancellation: CancellationToken,
213}
214
215/// Options for status, results and cancel calls.
216#[derive(Debug, Clone)]
217pub struct BatchOperationOptions {
218    /// The batch job.
219    pub batch_id: BatchId,
220    /// Provider-specific options keyed by provider name.
221    pub provider_options: ProviderOptions,
222    /// Additional request headers.
223    pub headers: Headers,
224    /// Cancellation token.
225    pub cancellation: CancellationToken,
226}
227
228impl BatchOperationOptions {
229    /// Creates options for `batch_id`.
230    #[must_use]
231    pub fn new(batch_id: impl Into<BatchId>) -> Self {
232        Self {
233            batch_id: batch_id.into(),
234            provider_options: ProviderOptions::new(),
235            headers: Headers::new(),
236            cancellation: CancellationToken::new(),
237        }
238    }
239}
240
241/// Options for listing batches.
242#[derive(Debug, Clone, Default)]
243pub struct BatchListOptions {
244    /// Page size.
245    pub limit: Option<usize>,
246    /// Cursor from a previous page.
247    pub cursor: Option<String>,
248    /// Provider-specific options keyed by provider name.
249    pub provider_options: ProviderOptions,
250    /// Additional request headers.
251    pub headers: Headers,
252    /// Cancellation token.
253    pub cancellation: CancellationToken,
254}
255
256/// Normalized batch state.
257#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
258#[serde(rename_all = "lowercase")]
259#[non_exhaustive]
260pub enum BatchState {
261    /// Queued or running.
262    Pending,
263    /// Finished; results are available.
264    Completed,
265    /// Failed as a whole.
266    Failed,
267}
268
269/// Error reported for a batch or a batch item.
270#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
271pub struct BatchError {
272    /// Message.
273    pub message: String,
274    /// Provider error type.
275    #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
276    pub error_type: Option<String>,
277    /// Provider error code.
278    #[serde(default, skip_serializing_if = "Option::is_none")]
279    pub code: Option<String>,
280    /// HTTP status code.
281    #[serde(default, skip_serializing_if = "Option::is_none")]
282    pub status_code: Option<u16>,
283}
284
285/// Request counts of a batch.
286#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
287pub struct BatchRequestCounts {
288    /// Total requests.
289    pub total: u64,
290    /// Requests not finished yet.
291    pub pending: u64,
292    /// Requests finished successfully.
293    pub completed: u64,
294    /// Requests that failed.
295    pub failed: u64,
296}
297
298/// Normalized status of a batch.
299#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
300pub struct BatchStatus {
301    /// Normalized state.
302    pub status: BatchState,
303    /// Provider's raw status string.
304    #[serde(default, skip_serializing_if = "Option::is_none")]
305    pub raw_status: Option<String>,
306    /// Request counts, if reported.
307    #[serde(default, skip_serializing_if = "Option::is_none")]
308    pub request_counts: Option<BatchRequestCounts>,
309    /// Batch-level error, if failed.
310    #[serde(default, skip_serializing_if = "Option::is_none")]
311    pub error: Option<BatchError>,
312    /// Creation time.
313    #[serde(default, skip_serializing_if = "Option::is_none")]
314    pub created_at: Option<DateTime<Utc>>,
315    /// Expiry time.
316    #[serde(default, skip_serializing_if = "Option::is_none")]
317    pub expires_at: Option<DateTime<Utc>>,
318    /// Provider-specific metadata.
319    #[serde(default, skip_serializing_if = "Option::is_none")]
320    pub provider_metadata: Option<ProviderMetadata>,
321}
322
323impl BatchStatus {
324    /// Creates a status with only the normalized state set.
325    #[must_use]
326    pub fn new(status: BatchState) -> Self {
327        Self {
328            status,
329            raw_status: None,
330            request_counts: None,
331            error: None,
332            created_at: None,
333            expires_at: None,
334            provider_metadata: None,
335        }
336    }
337}
338
339/// A warning attached to a batch start, optionally scoped to one request.
340#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
341pub struct BatchWarning {
342    /// Request the warning applies to, or `None` for the whole batch.
343    #[serde(default, skip_serializing_if = "Option::is_none")]
344    pub request_id: Option<String>,
345    /// The warning.
346    pub warning: Warning,
347}
348
349/// Result of starting a batch.
350#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
351pub struct BatchStartResult {
352    /// Provider batch id.
353    pub batch_id: BatchId,
354    /// Initial status.
355    #[serde(flatten)]
356    pub status: BatchStatus,
357    /// Warnings.
358    #[serde(default)]
359    pub warnings: Vec<BatchWarning>,
360}
361
362/// Result of cancelling a batch.
363#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
364pub struct BatchCancelResult {
365    /// Provider-specific metadata.
366    #[serde(default, skip_serializing_if = "Option::is_none")]
367    pub provider_metadata: Option<ProviderMetadata>,
368}
369
370/// One entry of a batch listing.
371#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
372pub struct BatchListItem {
373    /// Provider batch id.
374    pub batch_id: BatchId,
375    /// Status.
376    #[serde(flatten)]
377    pub status: BatchStatus,
378}
379
380/// Result of listing batches.
381#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
382pub struct BatchListResult {
383    /// Batches on this page.
384    pub batches: Vec<BatchListItem>,
385    /// Cursor for the next page.
386    #[serde(default, skip_serializing_if = "Option::is_none")]
387    pub next_cursor: Option<String>,
388    /// Provider-specific metadata.
389    #[serde(default, skip_serializing_if = "Option::is_none")]
390    pub provider_metadata: Option<ProviderMetadata>,
391}
392
393/// Outcome of one request, tagged by `status`.
394#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
395#[serde(tag = "status", rename_all = "lowercase")]
396#[non_exhaustive]
397pub enum BatchItem<R> {
398    /// The request succeeded.
399    Succeeded {
400        /// Caller-assigned request id.
401        id: String,
402        /// The result.
403        result: R,
404    },
405    /// The request failed.
406    Failed {
407        /// Caller-assigned request id.
408        id: String,
409        /// The error.
410        error: BatchError,
411        /// Provider-specific metadata.
412        #[serde(default, skip_serializing_if = "Option::is_none")]
413        provider_metadata: Option<ProviderMetadata>,
414    },
415    /// The request was cancelled.
416    Cancelled {
417        /// Caller-assigned request id.
418        id: String,
419        /// Error details, if any.
420        #[serde(default, skip_serializing_if = "Option::is_none")]
421        error: Option<BatchError>,
422        /// Provider-specific metadata.
423        #[serde(default, skip_serializing_if = "Option::is_none")]
424        provider_metadata: Option<ProviderMetadata>,
425    },
426    /// The request expired before completion.
427    Expired {
428        /// Caller-assigned request id.
429        id: String,
430        /// Error details, if any.
431        #[serde(default, skip_serializing_if = "Option::is_none")]
432        error: Option<BatchError>,
433        /// Provider-specific metadata.
434        #[serde(default, skip_serializing_if = "Option::is_none")]
435        provider_metadata: Option<ProviderMetadata>,
436    },
437}
438
439impl<R> BatchItem<R> {
440    /// Returns the caller-assigned request id.
441    #[must_use]
442    pub fn id(&self) -> &str {
443        match self {
444            Self::Succeeded { id, .. }
445            | Self::Failed { id, .. }
446            | Self::Cancelled { id, .. }
447            | Self::Expired { id, .. } => id,
448        }
449    }
450}
451
452/// Outcome of one request, tagged by request `type`.
453#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
454#[serde(tag = "type", rename_all = "lowercase")]
455#[non_exhaustive]
456pub enum BatchItemResult {
457    /// A text request.
458    Text(Box<BatchItem<GenerateResult>>),
459    /// An image request.
460    Image(Box<BatchItem<ImageResult>>),
461}
462
463impl BatchItemResult {
464    /// Returns the caller-assigned request id.
465    #[must_use]
466    pub fn id(&self) -> &str {
467        match self {
468            Self::Text(item) => item.id(),
469            Self::Image(item) => item.id(),
470        }
471    }
472}
ferrin_spec/batch.rs

ferrin_spec/
batch.rs
