gemini_genai_rs/batches/
mod.rs

1//! Batches API — create, list, get, cancel, delete batch prediction jobs.
2//!
3//! Feature-gated behind `batches`.
4
5use serde::{Deserialize, Serialize};
6
7use crate::client::Client;
8use crate::client::http::HttpError;
9use crate::transport::auth::ServiceEndpoint;
10
11/// State of a batch job.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
14pub enum BatchJobState {
15    /// State not set by the server.
16    StateUnspecified,
17    /// Queued, not yet started.
18    Pending,
19    /// Currently executing.
20    Running,
21    /// Completed successfully.
22    Succeeded,
23    /// Terminated with an error.
24    Failed,
25    /// Cancellation requested, still winding down.
26    Cancelling,
27    /// Cancelled before completion.
28    Cancelled,
29}
30
31/// Source configuration for a batch job.
32#[derive(Debug, Clone, Serialize, Deserialize)]
33#[serde(rename_all = "camelCase")]
34pub struct BatchJobSource {
35    /// GCS URI of the input file (JSONL).
36    pub gcs_uri: Option<String>,
37    /// BigQuery input table.
38    pub bigquery_source: Option<String>,
39    /// Format of the input (e.g., "bigquery", "jsonl").
40    #[serde(default)]
41    pub format: Option<String>,
42}
43
44/// Destination configuration for a batch job.
45#[derive(Debug, Clone, Serialize, Deserialize)]
46#[serde(rename_all = "camelCase")]
47pub struct BatchJobDestination {
48    /// GCS URI prefix for output.
49    pub gcs_uri: Option<String>,
50    /// BigQuery output table.
51    pub bigquery_destination: Option<String>,
52}
53
54/// A batch prediction job resource.
55#[derive(Debug, Clone, Serialize, Deserialize)]
56#[serde(rename_all = "camelCase")]
57pub struct BatchJob {
58    /// Resource name.
59    #[serde(default)]
60    pub name: String,
61    /// Display name.
62    #[serde(default)]
63    pub display_name: Option<String>,
64    /// Model used for batch prediction.
65    #[serde(default)]
66    pub model: Option<String>,
67    /// State of the batch job.
68    #[serde(default)]
69    pub state: Option<BatchJobState>,
70    /// Input source.
71    #[serde(default)]
72    pub source: Option<BatchJobSource>,
73    /// Output destination.
74    #[serde(default)]
75    pub destination: Option<BatchJobDestination>,
76    /// Creation time (RFC3339).
77    #[serde(default)]
78    pub create_time: Option<String>,
79    /// Update time (RFC3339).
80    #[serde(default)]
81    pub update_time: Option<String>,
82    /// Completion time (RFC3339).
83    #[serde(default)]
84    pub completion_time: Option<String>,
85    /// Error details if state is Failed.
86    #[serde(default)]
87    pub error: Option<serde_json::Value>,
88}
89
90/// Configuration for creating a batch job.
91#[derive(Debug, Clone)]
92pub struct CreateBatchJobConfig {
93    /// Model for batch prediction.
94    pub model: String,
95    /// Display name.
96    pub display_name: Option<String>,
97    /// Input source configuration.
98    pub source: BatchJobSource,
99    /// Output destination configuration.
100    pub destination: BatchJobDestination,
101}
102
103/// Response from listBatchJobs.
104#[derive(Debug, Clone, Serialize, Deserialize)]
105#[serde(rename_all = "camelCase")]
106pub struct ListBatchJobsResponse {
107    /// List of batch jobs.
108    #[serde(default)]
109    pub batch_jobs: Vec<BatchJob>,
110    /// Pagination token for the next page.
111    #[serde(default)]
112    pub next_page_token: Option<String>,
113}
114
115/// Errors from the Batches API.
116#[derive(Debug, thiserror::Error)]
117pub enum BatchesError {
118    #[error(transparent)]
119    /// Transport-level HTTP failure.
120    Http(#[from] HttpError),
121    #[error("Failed to parse response: {0}")]
122    /// Response body failed to parse.
123    Parse(#[from] serde_json::Error),
124    #[error("Auth error: {0}")]
125    /// Authentication/authorization failure.
126    Auth(#[from] crate::session::AuthError),
127}
128
129impl Client {
130    /// List batch jobs.
131    pub async fn list_batch_jobs(&self) -> Result<ListBatchJobsResponse, BatchesError> {
132        let url = self.rest_url(ServiceEndpoint::BatchJobs);
133        let headers = self.auth_headers().await?;
134        let json = self.http_client().get_json(&url, headers).await?;
135        if json.is_null() {
136            return Ok(ListBatchJobsResponse {
137                batch_jobs: vec![],
138                next_page_token: None,
139            });
140        }
141        Ok(serde_json::from_value(json)?)
142    }
143
144    /// Create a batch prediction job.
145    pub async fn create_batch_job(
146        &self,
147        config: CreateBatchJobConfig,
148    ) -> Result<BatchJob, BatchesError> {
149        let url = self.rest_url(ServiceEndpoint::BatchJobs);
150        let headers = self.auth_headers().await?;
151
152        let mut body = serde_json::json!({
153            "model": config.model,
154            "source": config.source,
155            "destination": config.destination,
156        });
157
158        if let Some(name) = config.display_name {
159            body["displayName"] = serde_json::Value::String(name);
160        }
161
162        let json = self.http_client().post_json(&url, headers, &body).await?;
163        Ok(serde_json::from_value(json)?)
164    }
165
166    /// Create a batch embeddings job.
167    pub async fn create_batch_embeddings(
168        &self,
169        config: CreateBatchJobConfig,
170    ) -> Result<BatchJob, BatchesError> {
171        // Same endpoint — the model determines the operation type
172        self.create_batch_job(config).await
173    }
174
175    /// Get a batch job by name.
176    pub async fn get_batch_job(&self, name: &str) -> Result<BatchJob, BatchesError> {
177        let base_url = self.rest_url(ServiceEndpoint::BatchJobs);
178        let url = format!("{base_url}/{name}");
179        let headers = self.auth_headers().await?;
180        let json = self.http_client().get_json(&url, headers).await?;
181        Ok(serde_json::from_value(json)?)
182    }
183
184    /// Cancel a batch job by name.
185    pub async fn cancel_batch_job(&self, name: &str) -> Result<(), BatchesError> {
186        let base_url = self.rest_url(ServiceEndpoint::BatchJobs);
187        let url = format!("{base_url}/{name}:cancel");
188        let headers = self.auth_headers().await?;
189        self.http_client()
190            .post_json(&url, headers, &serde_json::json!({}))
191            .await?;
192        Ok(())
193    }
194
195    /// Delete a batch job by name.
196    pub async fn delete_batch_job(&self, name: &str) -> Result<(), BatchesError> {
197        let base_url = self.rest_url(ServiceEndpoint::BatchJobs);
198        let url = format!("{base_url}/{name}");
199        let headers = self.auth_headers().await?;
200        self.http_client().delete(&url, headers).await?;
201        Ok(())
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    #[test]
210    fn parse_batch_job() {
211        let json = serde_json::json!({
212            "name": "batchJobs/123",
213            "displayName": "My Batch",
214            "model": "models/gemini-1.5-flash",
215            "state": "RUNNING",
216            "createTime": "2026-03-01T00:00:00Z"
217        });
218        let job: BatchJob = serde_json::from_value(json).unwrap();
219        assert_eq!(job.name, "batchJobs/123");
220        assert_eq!(job.state, Some(BatchJobState::Running));
221    }
222
223    #[test]
224    fn parse_list_batch_jobs_response() {
225        let json = serde_json::json!({
226            "batchJobs": [
227                {"name": "batchJobs/1", "state": "PENDING"},
228                {"name": "batchJobs/2", "state": "SUCCEEDED"}
229            ],
230            "nextPageToken": "page2"
231        });
232        let resp: ListBatchJobsResponse = serde_json::from_value(json).unwrap();
233        assert_eq!(resp.batch_jobs.len(), 2);
234        assert_eq!(resp.next_page_token, Some("page2".to_string()));
235    }
236
237    #[test]
238    fn batch_job_state_serialization() {
239        assert_eq!(
240            serde_json::to_value(BatchJobState::Running).unwrap(),
241            "RUNNING"
242        );
243        assert_eq!(
244            serde_json::to_value(BatchJobState::Succeeded).unwrap(),
245            "SUCCEEDED"
246        );
247        assert_eq!(
248            serde_json::to_value(BatchJobState::Cancelled).unwrap(),
249            "CANCELLED"
250        );
251    }
252
253    #[test]
254    fn batch_source_serialization() {
255        let source = BatchJobSource {
256            gcs_uri: Some("gs://bucket/input.jsonl".to_string()),
257            bigquery_source: None,
258            format: Some("jsonl".to_string()),
259        };
260        let json = serde_json::to_value(&source).unwrap();
261        assert_eq!(json["gcsUri"], "gs://bucket/input.jsonl");
262    }
263
264    #[test]
265    fn batch_destination_serialization() {
266        let dest = BatchJobDestination {
267            gcs_uri: Some("gs://bucket/output/".to_string()),
268            bigquery_destination: None,
269        };
270        let json = serde_json::to_value(&dest).unwrap();
271        assert_eq!(json["gcsUri"], "gs://bucket/output/");
272    }
273
274    #[test]
275    fn empty_list_response() {
276        let json = serde_json::json!({"batchJobs": []});
277        let resp: ListBatchJobsResponse = serde_json::from_value(json).unwrap();
278        assert!(resp.batch_jobs.is_empty());
279    }
280
281    #[test]
282    fn batch_job_with_error() {
283        let json = serde_json::json!({
284            "name": "batchJobs/bad",
285            "state": "FAILED",
286            "error": {"code": 400, "message": "Invalid input"}
287        });
288        let job: BatchJob = serde_json::from_value(json).unwrap();
289        assert_eq!(job.state, Some(BatchJobState::Failed));
290        assert!(job.error.is_some());
291    }
292}
gemini_genai_rs/batches/mod.rs

gemini_genai_rs/batches/
mod.rs
