1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4use validator::Validate;
5
6use crate::{ZaiResult, client::ZaiClient};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10pub enum BatchEndpoint {
11 #[serde(rename = "/v4/chat/completions")]
13 ChatCompletions,
14}
15
16#[derive(Clone, Serialize, Deserialize, Validate)]
18pub struct BatchCreateBody {
19 #[validate(length(min = 1))]
21 pub input_file_id: String,
22
23 pub endpoint: BatchEndpoint,
25
26 #[serde(skip_serializing_if = "Option::is_none")]
28 pub auto_delete_input_file: Option<bool>,
29
30 #[serde(skip_serializing_if = "Option::is_none")]
32 pub metadata: Option<BTreeMap<String, String>>,
33}
34
35impl std::fmt::Debug for BatchCreateBody {
36 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37 formatter
38 .debug_struct("BatchCreateBody")
39 .field("input_file_id", &"[REDACTED]")
40 .field("endpoint", &self.endpoint)
41 .field("auto_delete_input_file", &self.auto_delete_input_file)
42 .field(
43 "metadata_entries",
44 &self.metadata.as_ref().map(BTreeMap::len),
45 )
46 .finish()
47 }
48}
49
50impl BatchCreateBody {
51 pub fn new(input_file_id: impl Into<String>, endpoint: BatchEndpoint) -> Self {
54 Self {
55 input_file_id: input_file_id.into(),
56 endpoint,
57 auto_delete_input_file: Some(true),
58 metadata: None,
59 }
60 }
61
62 pub fn with_auto_delete_input_file(mut self, v: bool) -> Self {
64 self.auto_delete_input_file = Some(v);
65 self
66 }
67
68 pub fn with_metadata(mut self, v: BTreeMap<String, String>) -> Self {
70 self.metadata = Some(v);
71 self
72 }
73}
74
75pub struct BatchCreateRequest {
77 body: BatchCreateBody,
79}
80
81impl BatchCreateRequest {
82 pub fn new(input_file_id: impl Into<String>, endpoint: BatchEndpoint) -> Self {
84 let body = BatchCreateBody::new(input_file_id, endpoint);
85 Self { body }
86 }
87
88 pub fn with_auto_delete_input_file(mut self, v: bool) -> Self {
90 self.body = self.body.with_auto_delete_input_file(v);
91 self
92 }
93
94 pub fn with_metadata(mut self, v: BTreeMap<String, String>) -> Self {
96 self.body = self.body.with_metadata(v);
97 self
98 }
99
100 pub fn validate(&self) -> ZaiResult<()> {
102 self.body.validate()?;
103 if self.body.input_file_id.trim().is_empty() {
104 return Err(crate::client::validation::invalid(
105 "input_file_id cannot be blank",
106 ));
107 }
108 if let Some(metadata) = self.body.metadata.as_ref() {
109 if metadata.len() > 16 {
110 return Err(crate::client::validation::invalid(
111 "metadata cannot contain more than 16 entries",
112 ));
113 }
114 for (key, value) in metadata {
115 if key.trim().is_empty() || key.chars().count() > 64 {
116 return Err(crate::client::validation::invalid(
117 "metadata keys must contain between 1 and 64 non-blank characters",
118 ));
119 }
120 if value.chars().count() > 512 {
121 return Err(crate::client::validation::invalid(
122 "metadata values cannot exceed 512 characters",
123 ));
124 }
125 }
126 }
127 Ok(())
128 }
129
130 pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<BatchCreateResponse> {
132 self.validate()?;
133
134 let route = crate::client::routes::BATCHES_CREATE;
135 let url = client.endpoints().resolve_route(route, &[])?;
136 client
137 .send_json::<_, BatchCreateResponse>(route.method(), url, &self.body)
138 .await
139 }
140}
141
142pub type BatchCreateResponse = super::types::BatchItem;
144
145#[cfg(test)]
146mod tests {
147 use super::*;
148
149 #[test]
150 fn metadata_is_a_bounded_string_map() {
151 let metadata = BTreeMap::from([("customer".to_owned(), "acme".to_owned())]);
152 let request = BatchCreateRequest::new("file-1", BatchEndpoint::ChatCompletions)
153 .with_metadata(metadata);
154 assert!(request.validate().is_ok());
155 assert_eq!(
156 serde_json::to_value(&request.body).unwrap()["metadata"]["customer"],
157 "acme"
158 );
159
160 let too_many = (0..17)
161 .map(|index| (format!("key-{index}"), "value".to_owned()))
162 .collect();
163 assert!(
164 BatchCreateRequest::new("file-1", BatchEndpoint::ChatCompletions)
165 .with_metadata(too_many)
166 .validate()
167 .is_err()
168 );
169 assert!(
170 BatchCreateRequest::new("file-1", BatchEndpoint::ChatCompletions)
171 .with_metadata(BTreeMap::from([(" ".to_owned(), "value".to_owned())]))
172 .validate()
173 .is_err()
174 );
175 }
176
177 #[test]
178 fn debug_redacts_file_id_and_metadata() {
179 let body =
180 BatchCreateBody::new("private-file-id", BatchEndpoint::ChatCompletions).with_metadata(
181 BTreeMap::from([("private-key".to_owned(), "private-value".to_owned())]),
182 );
183 let debug = format!("{body:?}");
184 for secret in ["private-file-id", "private-key", "private-value"] {
185 assert!(!debug.contains(secret));
186 }
187 assert!(debug.contains("metadata_entries: Some(1)"));
188 }
189}