1use serde::{Deserialize, Serialize};
2use snafu::Snafu;
3use time::OffsetDateTime;
4
5use crate::common::serde::*;
6use crate::generation::{GenerateContentRequest, GenerationResponse};
7use crate::Model;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct BatchRequestFileItem {
12 pub request: GenerateContentRequest,
14 #[serde(with = "key_as_string")]
16 pub key: usize,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct BatchResponseFileItem {
22 #[serde(flatten)]
24 pub response: BatchGenerateContentResponseItem,
25 #[serde(with = "key_as_string")]
27 pub key: usize,
28}
29
30impl From<BatchGenerateContentResponseItem> for Result<GenerationResponse, IndividualRequestError> {
31 fn from(response: BatchGenerateContentResponseItem) -> Self {
32 match response {
33 BatchGenerateContentResponseItem::Response(r) => Ok(r),
34 BatchGenerateContentResponseItem::Error(err) => Err(err),
35 }
36 }
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
41#[serde(untagged)]
42pub enum BatchOperationResponse {
43 #[serde(rename_all = "camelCase")]
45 InlinedResponses { inlined_responses: InlinedResponses },
46 #[serde(rename_all = "camelCase")]
48 ResponsesFile { responses_file: String },
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct InlinedResponses {
55 pub inlined_responses: Vec<InlinedBatchGenerationResponseItem>,
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(rename_all = "camelCase")]
65pub struct InlinedBatchGenerationResponseItem {
66 pub metadata: RequestMetadata,
68 #[serde(flatten)]
70 pub result: BatchGenerateContentResponseItem,
71}
72
73#[derive(Debug, Clone, Serialize, Deserialize)]
75#[serde(rename_all = "camelCase")]
76pub enum BatchGenerateContentResponseItem {
77 Response(GenerationResponse),
79 Error(IndividualRequestError),
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
85pub struct IndividualRequestError {
86 pub code: i32,
88 pub message: String,
90 #[serde(skip_serializing_if = "Option::is_none")]
92 pub details: Option<serde_json::Value>,
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
97#[serde(rename_all = "camelCase")]
98pub struct BatchGenerateContentResponse {
99 pub name: String,
101 pub metadata: BatchMetadata,
103}
104
105#[derive(Debug, Clone, Serialize, Deserialize)]
107#[serde(rename_all = "camelCase")]
108pub struct BatchMetadata {
109 #[serde(rename = "@type")]
111 pub type_annotation: String,
112 pub model: Model,
114 pub display_name: String,
116 #[serde(with = "time::serde::rfc3339")]
118 pub create_time: OffsetDateTime,
119 #[serde(with = "time::serde::rfc3339")]
121 pub update_time: OffsetDateTime,
122 pub batch_stats: BatchStats,
124 pub state: BatchState,
126 pub name: String,
128}
129
130#[derive(Debug, Clone, Serialize, Deserialize)]
132#[serde(rename_all = "camelCase")]
133pub struct BatchRequestItem {
134 pub request: GenerateContentRequest,
136 pub metadata: RequestMetadata,
138}
139
140#[derive(Debug, Clone, Serialize, Deserialize)]
142#[serde(rename_all = "camelCase")]
143pub struct BatchGenerateContentRequest {
144 pub batch: BatchConfig,
146}
147
148#[derive(Debug, Clone, Serialize, Deserialize)]
150#[serde(rename_all = "camelCase")]
151pub struct BatchConfig {
152 pub display_name: String,
154 pub input_config: InputConfig,
156}
157
158#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
160#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
161#[allow(clippy::enum_variant_names)]
162pub enum BatchState {
163 BatchStateUnspecified,
165 BatchStatePending,
167 BatchStateRunning,
169 BatchStateSucceeded,
171 BatchStateFailed,
173 BatchStateCancelled,
175 BatchStateExpired,
177}
178
179#[derive(Debug, Clone, Serialize, Deserialize)]
181#[serde(rename_all = "camelCase")]
182pub struct BatchStats {
183 #[serde(with = "i64_as_string")]
185 pub request_count: i64,
186 #[serde(default, with = "i64_as_string::optional")]
188 pub pending_request_count: Option<i64>,
189 #[serde(default, with = "i64_as_string::optional")]
191 pub completed_request_count: Option<i64>,
192 #[serde(default, with = "i64_as_string::optional")]
194 pub failed_request_count: Option<i64>,
195 #[serde(default, with = "i64_as_string::optional")]
197 pub successful_request_count: Option<i64>,
198}
199
200#[derive(Debug, Serialize, Deserialize)]
202#[serde(rename_all = "camelCase")]
203pub struct BatchOperation {
204 pub name: String,
206 pub metadata: BatchMetadata,
208 #[serde(default)]
210 pub done: bool,
211 #[serde(flatten)]
213 pub result: Option<OperationResult>,
214}
215
216#[derive(Debug, Snafu, serde::Deserialize, serde::Serialize)]
218pub struct OperationError {
219 pub code: i32,
221 pub message: String,
223 }
225
226#[derive(Debug, Serialize, Deserialize)]
228#[serde(rename_all = "camelCase")]
229pub enum OperationResult {
230 Response(BatchOperationResponse),
232 Error(OperationError),
234}
235
236impl From<OperationResult> for Result<BatchOperationResponse, OperationError> {
237 fn from(operation: OperationResult) -> Self {
238 match operation {
239 OperationResult::Response(response) => Ok(response),
240 OperationResult::Error(error) => Err(error),
241 }
242 }
243}
244
245#[derive(Debug, serde::Deserialize)]
248#[serde(rename_all = "camelCase")]
249pub struct ListBatchesResponse {
250 pub operations: Vec<BatchOperation>,
252 pub next_page_token: Option<String>,
254}
255
256#[derive(Debug, Clone, Serialize, Deserialize)]
258#[serde(rename_all = "camelCase")]
259pub enum InputConfig {
260 Requests(RequestsContainer),
262 FileName(String),
264}
265
266impl InputConfig {
267 pub fn batch_size(&self) -> Option<usize> {
271 match self {
272 InputConfig::Requests(container) => Some(container.requests.len()),
273 InputConfig::FileName(_) => None,
274 }
275 }
276}
277
278#[derive(Debug, Clone, Serialize, Deserialize)]
280#[serde(rename_all = "camelCase")]
281pub struct RequestsContainer {
282 pub requests: Vec<BatchRequestItem>,
284}
285
286#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
288#[serde(rename_all = "camelCase")]
289pub struct RequestMetadata {
290 #[serde(with = "key_as_string")]
292 pub key: usize,
293}