1use serde::{Deserialize, Serialize};
2use snafu::Snafu;
3use time::OffsetDateTime;
4
5use crate::common::serde::*;
6use crate::generation::{GenerateContentRequest, GenerationResponse};
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct BatchRequestFileItem {
11 pub request: GenerateContentRequest,
13 #[serde(with = "key_as_string")]
15 pub key: usize,
16}
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct BatchResponseFileItem {
21 #[serde(flatten)]
23 pub response: BatchGenerateContentResponseItem,
24 #[serde(with = "key_as_string")]
26 pub key: usize,
27}
28
29impl From<BatchGenerateContentResponseItem> for Result<GenerationResponse, IndividualRequestError> {
30 fn from(response: BatchGenerateContentResponseItem) -> Self {
31 match response {
32 BatchGenerateContentResponseItem::Response(r) => Ok(r),
33 BatchGenerateContentResponseItem::Error(err) => Err(err),
34 }
35 }
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize)]
40#[serde(untagged)]
41pub enum BatchOperationResponse {
42 #[serde(rename_all = "camelCase")]
44 InlinedResponses { inlined_responses: InlinedResponses },
45 #[serde(rename_all = "camelCase")]
47 ResponsesFile { responses_file: String },
48}
49
50#[derive(Debug, Clone, Serialize, Deserialize)]
52#[serde(rename_all = "camelCase")]
53pub struct InlinedResponses {
54 pub inlined_responses: Vec<InlinedBatchGenerationResponseItem>,
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize)]
63#[serde(rename_all = "camelCase")]
64pub struct InlinedBatchGenerationResponseItem {
65 pub metadata: RequestMetadata,
67 #[serde(flatten)]
69 pub result: BatchGenerateContentResponseItem,
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize)]
74#[serde(rename_all = "camelCase")]
75pub enum BatchGenerateContentResponseItem {
76 Response(GenerationResponse),
78 Error(IndividualRequestError),
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
84pub struct IndividualRequestError {
85 pub code: i32,
87 pub message: String,
89 #[serde(skip_serializing_if = "Option::is_none")]
91 pub details: Option<serde_json::Value>,
92}
93
94#[derive(Debug, Clone, Serialize, Deserialize)]
96#[serde(rename_all = "camelCase")]
97pub struct BatchGenerateContentResponse {
98 pub name: String,
100 pub metadata: BatchMetadata,
102}
103
104#[derive(Debug, Clone, Serialize, Deserialize)]
106#[serde(rename_all = "camelCase")]
107pub struct BatchMetadata {
108 #[serde(rename = "@type")]
110 pub type_annotation: String,
111 pub model: String,
113 pub display_name: String,
115 #[serde(with = "time::serde::rfc3339")]
117 pub create_time: OffsetDateTime,
118 #[serde(with = "time::serde::rfc3339")]
120 pub update_time: OffsetDateTime,
121 pub batch_stats: BatchStats,
123 pub state: BatchState,
125 pub name: String,
127}
128
129#[derive(Debug, Clone, Serialize, Deserialize)]
131#[serde(rename_all = "camelCase")]
132pub struct BatchRequestItem {
133 pub request: GenerateContentRequest,
135 pub metadata: RequestMetadata,
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize)]
141#[serde(rename_all = "camelCase")]
142pub struct BatchGenerateContentRequest {
143 pub batch: BatchConfig,
145}
146
147#[derive(Debug, Clone, Serialize, Deserialize)]
149#[serde(rename_all = "camelCase")]
150pub struct BatchConfig {
151 pub display_name: String,
153 pub input_config: InputConfig,
155}
156
157#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
159#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
160#[allow(clippy::enum_variant_names)]
161pub enum BatchState {
162 BatchStateUnspecified,
164 BatchStatePending,
166 BatchStateRunning,
168 BatchStateSucceeded,
170 BatchStateFailed,
172 BatchStateCancelled,
174 BatchStateExpired,
176}
177
178#[derive(Debug, Clone, Serialize, Deserialize)]
180#[serde(rename_all = "camelCase")]
181pub struct BatchStats {
182 #[serde(with = "i64_as_string")]
184 pub request_count: i64,
185 #[serde(default, with = "i64_as_string::optional")]
187 pub pending_request_count: Option<i64>,
188 #[serde(default, with = "i64_as_string::optional")]
190 pub completed_request_count: Option<i64>,
191 #[serde(default, with = "i64_as_string::optional")]
193 pub failed_request_count: Option<i64>,
194 #[serde(default, with = "i64_as_string::optional")]
196 pub successful_request_count: Option<i64>,
197}
198
199#[derive(Debug, Serialize, Deserialize)]
201pub struct BatchOperation {
202 pub name: String,
204 pub metadata: BatchMetadata,
206 #[serde(default)]
208 pub done: bool,
209 #[serde(flatten)]
211 pub result: Option<OperationResult>,
212}
213
214#[derive(Debug, Snafu, serde::Deserialize, serde::Serialize)]
216pub struct OperationError {
217 pub code: i32,
219 pub message: String,
221 }
223
224#[derive(Debug, Serialize, Deserialize)]
226#[serde(rename_all = "camelCase")]
227pub enum OperationResult {
228 Response(BatchOperationResponse),
230 Error(OperationError),
232}
233
234impl From<OperationResult> for Result<BatchOperationResponse, OperationError> {
235 fn from(operation: OperationResult) -> Self {
236 match operation {
237 OperationResult::Response(response) => Ok(response),
238 OperationResult::Error(error) => Err(error),
239 }
240 }
241}
242
243#[derive(Debug, serde::Deserialize)]
246#[serde(rename_all = "camelCase")]
247pub struct ListBatchesResponse {
248 pub operations: Vec<BatchOperation>,
250 pub next_page_token: Option<String>,
252}
253
254#[derive(Debug, Clone, Serialize, Deserialize)]
256#[serde(rename_all = "camelCase")]
257pub enum InputConfig {
258 Requests(RequestsContainer),
260 FileName(String),
262}
263
264#[derive(Debug, Clone, Serialize, Deserialize)]
266#[serde(rename_all = "camelCase")]
267pub struct RequestsContainer {
268 pub requests: Vec<BatchRequestItem>,
270}
271
272#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
274#[serde(rename_all = "camelCase")]
275pub struct RequestMetadata {
276 #[serde(with = "key_as_string")]
278 pub key: usize,
279}