ferrin_google/batch/
mod.rs

1//! Batch generation (`batchGenerateContent`): inline or file-based input,
2//! status, results, cancel and list.
3//!
4//! Derived from the Vercel AI SDK (Apache-2.0, Copyright 2023 Vercel, Inc.),
5//! translated from TypeScript to Rust and modified; see `NOTICE`.
6
7mod results;
8
9use std::time::Duration;
10
11use ferrin_provider_util::batch::normalize_batch_request_counts;
12use ferrin_provider_util::http::ResponseHandlers;
13use ferrin_provider_util::http::get;
14use ferrin_provider_util::http::json_lines_response_handler;
15use ferrin_provider_util::http::json_response_handler;
16use ferrin_provider_util::http::post_json;
17use ferrin_provider_util::http::text_response_handler;
18use ferrin_spec::BatchId;
19use ferrin_spec::JsonObject;
20use ferrin_spec::JsonValue;
21use ferrin_spec::ModelId;
22use ferrin_spec::ProviderId;
23use ferrin_spec::batch::Batch;
24use ferrin_spec::batch::BatchCancelResult;
25use ferrin_spec::batch::BatchError;
26use ferrin_spec::batch::BatchListItem;
27use ferrin_spec::batch::BatchListOptions;
28use ferrin_spec::batch::BatchListResult;
29use ferrin_spec::batch::BatchOperationOptions;
30use ferrin_spec::batch::BatchRequest;
31use ferrin_spec::batch::BatchResultStream;
32use ferrin_spec::batch::BatchStartOptions;
33use ferrin_spec::batch::BatchStartResult;
34use ferrin_spec::batch::BatchState;
35use ferrin_spec::batch::BatchStatus;
36use ferrin_spec::batch::BatchWarning;
37use ferrin_spec::error::InvalidArgumentError;
38use ferrin_spec::error::InvalidResponseDataError;
39use ferrin_spec::error::ProviderError;
40use ferrin_spec::image_model::ImageOptions;
41use ferrin_spec::language_model::CallOptions;
42use ferrin_spec::language_model::SupportedUrls;
43use ferrin_spec::shared::Warning;
44use futures_util::StreamExt;
45use serde::Deserialize;
46use serde_json::json;
47
48use crate::api_types::RpcStatus;
49use crate::api_types::deserialize_count;
50use crate::config::DOWNLOAD_PATH_PREFIX;
51use crate::config::SharedConfig;
52use crate::error::failed_response_handler;
53use crate::files::GoogleFiles;
54use crate::files::UploadRequest;
55use crate::image::GoogleImageModel;
56use crate::language_model::base_supported_urls;
57use crate::output::OutputMapper;
58use crate::request::prepare_request;
59
60pub use self::results::BatchResultLine;
61pub use self::results::convert_line;
62
63/// Provider id family.
64pub const FAMILY: &str = "batch";
65
66/// Inline request bodies stay below this many bytes; larger batches are
67/// uploaded as a JSONL file.
68pub const INLINE_MAX_BYTES: usize = 20_000_000;
69
70/// Maximum size of an uploaded JSONL input file.
71pub const INPUT_FILE_MAX_BYTES: u64 = 2_000_000_000;
72
73/// Prefix of the generated batch display names.
74pub const DISPLAY_NAME_PREFIX: &str = "ferrin-batch-";
75
76/// A long-running batch operation.
77#[derive(Debug, Clone, Default, PartialEq, Deserialize)]
78#[serde(rename_all = "camelCase")]
79pub struct BatchOperation {
80    /// Operation name (`batches/...`), used as the batch id.
81    pub name: String,
82    /// Metadata (state, counts, output).
83    #[serde(default)]
84    pub metadata: Option<BatchOperationMetadata>,
85    /// Whether the operation finished.
86    #[serde(default)]
87    pub done: Option<bool>,
88    /// Error of a failed operation.
89    #[serde(default)]
90    pub error: Option<RpcStatus>,
91    /// Response of a finished operation (same shape as `metadata.output`).
92    #[serde(default)]
93    pub response: Option<JsonValue>,
94}
95
96/// Metadata of a batch operation.
97#[derive(Debug, Clone, Default, PartialEq, Deserialize)]
98#[serde(rename_all = "camelCase")]
99pub struct BatchOperationMetadata {
100    /// Raw state (`BATCH_STATE_PENDING`, `JOB_STATE_SUCCEEDED`, ...).
101    #[serde(default)]
102    pub state: Option<String>,
103    /// Creation time (RFC 3339).
104    #[serde(default)]
105    pub create_time: Option<String>,
106    /// Request counts.
107    #[serde(default)]
108    pub batch_stats: Option<BatchStats>,
109    /// Output (`{responsesFile}` or `{inlinedResponses: {inlinedResponses}}`).
110    #[serde(default)]
111    pub output: Option<JsonValue>,
112}
113
114/// Request counts of a batch.
115#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
116#[serde(rename_all = "camelCase")]
117pub struct BatchStats {
118    /// Total requests.
119    #[serde(default, deserialize_with = "deserialize_count")]
120    pub request_count: Option<u64>,
121    /// Finished successfully.
122    #[serde(default = "zero_count", deserialize_with = "deserialize_zero_count")]
123    pub successful_request_count: Option<u64>,
124    /// Failed.
125    #[serde(default = "zero_count", deserialize_with = "deserialize_zero_count")]
126    pub failed_request_count: Option<u64>,
127    /// Not finished yet.
128    #[serde(default = "zero_count", deserialize_with = "deserialize_zero_count")]
129    pub pending_request_count: Option<u64>,
130}
131
132fn zero_count() -> Option<u64> {
133    Some(0)
134}
135
136fn deserialize_zero_count<'de, D: serde::Deserializer<'de>>(
137    deserializer: D,
138) -> Result<Option<u64>, D::Error> {
139    let value = Option::<JsonValue>::deserialize(deserializer)?;
140    Ok(match value {
141        None => Some(0),
142        Some(JsonValue::Number(value)) => value.as_u64(),
143        Some(JsonValue::String(value))
144            if !value.is_empty() && value.bytes().all(|byte| byte.is_ascii_digit()) =>
145        {
146            value.parse().ok()
147        }
148        _ => None,
149    })
150}
151
152#[derive(Debug, Deserialize)]
153#[serde(rename_all = "camelCase")]
154struct ListResponse {
155    #[serde(default)]
156    operations: Vec<BatchOperation>,
157    #[serde(default)]
158    next_page_token: Option<String>,
159}
160
161/// Maps an RPC status to a batch error.
162#[must_use]
163pub fn rpc_error(status: &RpcStatus, fallback: &str) -> BatchError {
164    BatchError {
165        message: status
166            .message
167            .clone()
168            .unwrap_or_else(|| fallback.to_owned()),
169        error_type: status.status.clone(),
170        code: status.code.map(|code| code.to_string()),
171        status_code: None,
172    }
173}
174
175/// Maps an operation to the normalized status.
176#[must_use]
177pub fn map_status(operation: &BatchOperation) -> BatchStatus {
178    let metadata = operation.metadata.as_ref();
179    let raw_state = metadata.and_then(|metadata| metadata.state.clone());
180    let normalized = raw_state.as_deref().map(|state| {
181        state
182            .trim_start_matches("BATCH_STATE_")
183            .trim_start_matches("JOB_STATE_")
184    });
185    let state = if operation.error.is_some() {
186        BatchState::Failed
187    } else {
188        match normalized {
189            Some("SUCCEEDED") => BatchState::Completed,
190            Some("FAILED" | "CANCELLED" | "EXPIRED") => BatchState::Failed,
191            Some(_) => BatchState::Pending,
192            None if operation.done == Some(true) => BatchState::Completed,
193            None => BatchState::Pending,
194        }
195    };
196    let mut status = BatchStatus::new(state);
197    status.raw_status = raw_state;
198    status.error = operation
199        .error
200        .as_ref()
201        .map(|error| rpc_error(error, "Google batch failed"));
202    status.request_counts = metadata
203        .and_then(|metadata| metadata.batch_stats.as_ref())
204        .and_then(|stats| {
205            normalize_batch_request_counts(
206                stats.request_count,
207                stats.pending_request_count,
208                stats.successful_request_count,
209                stats.failed_request_count,
210            )
211        });
212    status.created_at = metadata
213        .and_then(|metadata| metadata.create_time.as_deref())
214        .and_then(|time| chrono::DateTime::parse_from_rfc3339(time).ok())
215        .map(|time| time.with_timezone(&chrono::Utc));
216    status
217}
218
219fn batch_model_id(requests: &[BatchRequest]) -> Result<ModelId, ProviderError> {
220    let mut model_id: Option<&ModelId> = None;
221    for request in requests {
222        let id = match request {
223            BatchRequest::Text { model_id, .. } | BatchRequest::Image { model_id, .. } => model_id,
224            #[allow(unreachable_patterns, reason = "BatchRequest is non-exhaustive")]
225            _ => {
226                return Err(ProviderError::unsupported(
227                    "batch requests other than text and image",
228                ));
229            }
230        };
231        match model_id {
232            None => model_id = Some(id),
233            Some(first) if first != id => {
234                return Err(InvalidArgumentError::new(
235                    "requests",
236                    "google batches require every request to use the same model because the model is part of the batch endpoint",
237                )
238                .into());
239            }
240            Some(_) => {}
241        }
242    }
243    model_id.cloned().ok_or_else(|| {
244        InvalidArgumentError::new("requests", "google batches require at least one request").into()
245    })
246}
247
248/// A prepared batch request: the `generateContent` body and its warnings.
249#[derive(Debug, Clone, PartialEq)]
250pub struct PreparedBatchRequest {
251    /// Caller-assigned request id.
252    pub id: String,
253    /// Request body.
254    pub body: JsonObject,
255    /// Warnings.
256    pub warnings: Vec<Warning>,
257}
258
259/// Batch service backed by `batchGenerateContent`.
260#[derive(Debug, Clone)]
261pub struct GoogleBatch {
262    config: SharedConfig,
263    provider: ProviderId,
264}
265
266impl GoogleBatch {
267    /// Creates the service.
268    #[must_use]
269    pub fn new(config: SharedConfig) -> Self {
270        Self {
271            provider: config.provider_id(FAMILY),
272            config,
273        }
274    }
275
276    /// Builds the `generateContent` body of one batch request.
277    ///
278    /// # Errors
279    ///
280    /// Returns the language model's preparation errors; image requests with a
281    /// mask or `n > 1` fail with [`ProviderError::InvalidArgument`].
282    pub fn prepare_request(
283        &self,
284        request: &BatchRequest,
285    ) -> Result<PreparedBatchRequest, ProviderError> {
286        match request {
287            BatchRequest::Text {
288                id,
289                model_id,
290                options,
291            } => {
292                let mut call = CallOptions::new(options.prompt.clone());
293                call.max_output_tokens = options.max_output_tokens;
294                call.temperature = options.temperature;
295                call.stop_sequences = options.stop_sequences.clone();
296                call.top_p = options.top_p;
297                call.top_k = options.top_k;
298                call.presence_penalty = options.presence_penalty;
299                call.frequency_penalty = options.frequency_penalty;
300                call.seed = options.seed;
301                call.reasoning = options.reasoning;
302                call.response_format = options.response_format.clone();
303                call.tool_choice = options.tool_choice.clone();
304                call.tools = options.tools.clone();
305                call.provider_options = options.provider_options.clone();
306                let prepared = prepare_request(&self.config, model_id.as_str(), &call)?;
307                Ok(PreparedBatchRequest {
308                    id: id.clone(),
309                    body: prepared.body,
310                    warnings: prepared.warnings,
311                })
312            }
313            BatchRequest::Image {
314                id,
315                model_id,
316                options,
317            } => {
318                let image_options = ImageOptions {
319                    prompt: options.prompt.clone(),
320                    n: options.n,
321                    size: options.size,
322                    aspect_ratio: options.aspect_ratio,
323                    seed: options.seed,
324                    files: options.files.clone(),
325                    mask: options.mask.clone(),
326                    provider_options: options.provider_options.clone(),
327                    ..ImageOptions::default()
328                };
329                let (call, mut warnings) =
330                    GoogleImageModel::new(self.config.clone(), model_id.clone())
331                        .prepare_call(&image_options)?;
332                let prepared = prepare_request(&self.config, model_id.as_str(), &call)?;
333                warnings.extend(prepared.warnings);
334                Ok(PreparedBatchRequest {
335                    id: id.clone(),
336                    body: prepared.body,
337                    warnings,
338                })
339            }
340            #[allow(unreachable_patterns, reason = "BatchRequest is non-exhaustive")]
341            _ => Err(ProviderError::unsupported(
342                "batch requests other than text and image",
343            )),
344        }
345    }
346
347    async fn retrieve(
348        &self,
349        options: &BatchOperationOptions,
350    ) -> Result<(BatchOperation, Option<JsonValue>), ProviderError> {
351        let handlers = ResponseHandlers::new(
352            json_response_handler::<BatchOperation>(),
353            failed_response_handler(),
354        );
355        let response = get(
356            self.config.transport.as_ref(),
357            self.config.url(options.batch_id.as_str()),
358            self.config.headers(&options.headers)?,
359            &handlers,
360            options.cancellation.clone(),
361        )
362        .await?;
363        Ok((response.value, response.raw))
364    }
365
366    fn line(prepared: &PreparedBatchRequest) -> String {
367        json!({"key": prepared.id, "request": prepared.body}).to_string()
368    }
369
370    fn inlined(prepared: &PreparedBatchRequest) -> JsonValue {
371        json!({"request": prepared.body, "metadata": {"key": prepared.id}})
372    }
373
374    #[allow(
375        clippy::too_many_lines,
376        reason = "inline/file input selection is one sequential flow"
377    )]
378    async fn build_start_body(
379        &self,
380        options: &BatchStartOptions,
381        model_id: &ModelId,
382        display_name: &str,
383    ) -> Result<(JsonObject, Vec<BatchWarning>, Option<JsonObject>), ProviderError> {
384        let mut warnings = Vec::new();
385        let mut batch = JsonObject::new();
386        batch.insert("displayName".to_owned(), JsonValue::from(display_name));
387        if let Some(webhook) = &options.webhook_url {
388            batch.insert(
389                "webhookConfig".to_owned(),
390                json!({"uris": [webhook.as_str()]}),
391            );
392        }
393        let mut probe = batch.clone();
394        probe.insert(
395            "inputConfig".to_owned(),
396            json!({"requests": {"requests": []}}),
397        );
398        let mut inline_bytes = json!({"batch": probe}).to_string().len();
399        let mut inlined: Vec<PreparedBatchRequest> = Vec::new();
400        let mut file_lines: Option<Vec<String>> = None;
401        for request in &options.requests {
402            let mut prepared = self.prepare_request(request)?;
403            warnings.extend(prepared.warnings.drain(..).map(|warning| BatchWarning {
404                request_id: Some(prepared.id.clone()),
405                warning,
406            }));
407            if let Some(lines) = &mut file_lines {
408                lines.push(Self::line(&prepared));
409                continue;
410            }
411            let request_bytes = Self::inlined(&prepared).to_string().len();
412            let next = inline_bytes + request_bytes + usize::from(!inlined.is_empty());
413            if next < INLINE_MAX_BYTES {
414                inlined.push(prepared);
415                inline_bytes = next;
416            } else {
417                let mut lines: Vec<String> = inlined.iter().map(Self::line).collect();
418                lines.push(Self::line(&prepared));
419                inlined.clear();
420                file_lines = Some(lines);
421            }
422        }
423        let mut metadata = None;
424        match file_lines {
425            None => {
426                let requests: Vec<JsonValue> = inlined.iter().map(Self::inlined).collect();
427                batch.insert(
428                    "inputConfig".to_owned(),
429                    json!({"requests": {"requests": requests}}),
430                );
431            }
432            Some(lines) => {
433                let data = lines.join("\n");
434                if data.len() as u64 > INPUT_FILE_MAX_BYTES {
435                    return Err(InvalidArgumentError::new(
436                        "requests",
437                        "google batch input files must not exceed 2 GB",
438                    )
439                    .into());
440                }
441                let mut upload = UploadRequest::new(data.into(), "application/jsonl");
442                upload.display_name = Some(format!("{display_name}-input"));
443                upload.headers = options.headers.clone();
444                upload.cancellation = options.cancellation.clone();
445                upload.poll_interval =
446                    Duration::from_millis(crate::files::DEFAULT_POLL_INTERVAL_MS);
447                let file = GoogleFiles::new(self.config.clone())
448                    .upload_bytes(upload)
449                    .await?;
450                batch.insert("inputConfig".to_owned(), json!({"fileName": file.name}));
451                let mut object = JsonObject::new();
452                object.insert("inputFileId".to_owned(), JsonValue::from(file.name));
453                if let Some(expires) = file.expiration_time {
454                    object.insert("inputFileExpiresAt".to_owned(), JsonValue::from(expires));
455                }
456                metadata = Some(object);
457            }
458        }
459        let _ = model_id;
460        Ok((batch, warnings, metadata))
461    }
462}
463
464impl Batch for GoogleBatch {
465    fn provider(&self) -> &ProviderId {
466        &self.provider
467    }
468
469    async fn supported_urls(&self) -> SupportedUrls {
470        base_supported_urls(&self.config.base_url)
471    }
472
473    #[tracing::instrument(skip_all, fields(requests = options.requests.len()))]
474    async fn do_start_batch(
475        &self,
476        options: BatchStartOptions,
477    ) -> Result<BatchStartResult, ProviderError> {
478        let model_id = batch_model_id(&options.requests)?;
479        let display_name = format!("{DISPLAY_NAME_PREFIX}{}", self.config.generate_id());
480        let (batch, warnings, metadata) = self
481            .build_start_body(&options, &model_id, &display_name)
482            .await?;
483        let handlers = ResponseHandlers::new(
484            json_response_handler::<BatchOperation>(),
485            failed_response_handler(),
486        );
487        let response = post_json(
488            self.config.transport.as_ref(),
489            self.config
490                .model_url(model_id.as_str(), "batchGenerateContent"),
491            self.config.headers(&options.headers)?,
492            &json!({"batch": batch}),
493            &handlers,
494            options.cancellation.clone(),
495        )
496        .await?;
497        let operation = response.value;
498        let mut status = map_status(&operation);
499        if let Some(metadata) = metadata {
500            let mapper = OutputMapper::new(self.config.clone(), Default::default());
501            status.provider_metadata = Some(mapper.metadata(metadata));
502        }
503        Ok(BatchStartResult {
504            batch_id: BatchId::new(operation.name),
505            status,
506            warnings,
507        })
508    }
509
510    #[tracing::instrument(skip_all, fields(batch = %options.batch_id))]
511    async fn do_get_batch_status(
512        &self,
513        options: BatchOperationOptions,
514    ) -> Result<BatchStatus, ProviderError> {
515        let (operation, _) = self.retrieve(&options).await?;
516        Ok(map_status(&operation))
517    }
518
519    #[tracing::instrument(skip_all, fields(batch = %options.batch_id))]
520    async fn do_get_batch_results(
521        &self,
522        options: BatchOperationOptions,
523    ) -> Result<BatchResultStream, ProviderError> {
524        let (operation, raw) = self.retrieve(&options).await?;
525        let status = map_status(&operation);
526        if status.status == BatchState::Pending {
527            return Err(InvalidArgumentError::new(
528                "batch_id",
529                format!("google batch \"{}\" is not complete", options.batch_id),
530            )
531            .into());
532        }
533        let output = operation
534            .metadata
535            .and_then(|metadata| metadata.output)
536            .or(operation.response)
537            .unwrap_or(JsonValue::Null);
538        if let Some(inlined) = output
539            .get("inlinedResponses")
540            .and_then(|value| value.get("inlinedResponses"))
541            .and_then(JsonValue::as_array)
542        {
543            let config = self.config.clone();
544            let items: Vec<Result<_, ProviderError>> = inlined
545                .iter()
546                .map(|item| {
547                    let line = serde_json::from_value::<results::InlinedResponse>(item.clone())
548                        .map_err(|error| {
549                            ProviderError::InvalidResponseData(Box::new(
550                                InvalidResponseDataError::new(
551                                    format!(
552                                        "google returned an invalid inlined batch response: {error}"
553                                    ),
554                                    item.clone(),
555                                ),
556                            ))
557                        })?;
558                    Ok(convert_line(&config, line.into_line()))
559                })
560                .collect();
561            return Ok(Box::pin(futures_util::stream::iter(items)));
562        }
563        let Some(file) = output.get("responsesFile").and_then(JsonValue::as_str) else {
564            if status.status == BatchState::Completed {
565                return Err(ProviderError::InvalidResponseData(Box::new(
566                    InvalidResponseDataError::new(
567                        format!(
568                            "google batch \"{}\" completed without batch output",
569                            options.batch_id
570                        ),
571                        raw.unwrap_or(JsonValue::Null),
572                    ),
573                )));
574            }
575            return Ok(Box::pin(futures_util::stream::empty()));
576        };
577        let mut url = self
578            .config
579            .origin_url(&format!("{DOWNLOAD_PATH_PREFIX}{file}:download"));
580        url.set_query(Some("alt=media"));
581        let handlers = ResponseHandlers::new(
582            json_lines_response_handler::<BatchResultLine>(),
583            failed_response_handler(),
584        );
585        let response = get(
586            self.config.transport.as_ref(),
587            url,
588            self.config.headers(&options.headers)?,
589            &handlers,
590            options.cancellation.clone(),
591        )
592        .await?;
593        let config = self.config.clone();
594        Ok(Box::pin(response.value.map(move |parsed| {
595            parsed.into_result().map(|line| convert_line(&config, line))
596        })))
597    }
598
599    fn supports_cancel_batch(&self) -> bool {
600        true
601    }
602
603    #[tracing::instrument(skip_all, fields(batch = %options.batch_id))]
604    async fn do_cancel_batch(
605        &self,
606        options: BatchOperationOptions,
607    ) -> Result<BatchCancelResult, ProviderError> {
608        let handlers = ResponseHandlers::new(text_response_handler(), failed_response_handler());
609        post_json(
610            self.config.transport.as_ref(),
611            self.config
612                .url(&format!("{}:cancel", options.batch_id.as_str())),
613            self.config.headers(&options.headers)?,
614            &json!({}),
615            &handlers,
616            options.cancellation.clone(),
617        )
618        .await?;
619        Ok(BatchCancelResult::default())
620    }
621
622    fn supports_list_batches(&self) -> bool {
623        true
624    }
625
626    #[tracing::instrument(skip_all)]
627    async fn do_list_batches(
628        &self,
629        options: BatchListOptions,
630    ) -> Result<BatchListResult, ProviderError> {
631        let mut url = self.config.url("batches");
632        {
633            let mut query = url.query_pairs_mut();
634            if let Some(limit) = options.limit {
635                query.append_pair("pageSize", &limit.to_string());
636            }
637            if let Some(cursor) = &options.cursor {
638                query.append_pair("pageToken", cursor);
639            }
640        }
641        let handlers = ResponseHandlers::new(
642            json_response_handler::<ListResponse>(),
643            failed_response_handler(),
644        );
645        let response = get(
646            self.config.transport.as_ref(),
647            url,
648            self.config.headers(&options.headers)?,
649            &handlers,
650            options.cancellation.clone(),
651        )
652        .await?;
653        let batches = response
654            .value
655            .operations
656            .iter()
657            .map(|operation| BatchListItem {
658                batch_id: BatchId::new(operation.name.clone()),
659                status: map_status(operation),
660            })
661            .collect();
662        Ok(BatchListResult {
663            batches,
664            next_cursor: response.value.next_page_token,
665            provider_metadata: None,
666        })
667    }
668}
ferrin_google/batch/mod.rs

ferrin_google/batch/
mod.rs
