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

ferrin_google/batch/
mod.rs
