Skip to main content

zai_rs/batches/
create.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4use validator::Validate;
5
6use crate::{ZaiResult, client::ZaiClient};
7
8/// Endpoint for batch requests
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10pub enum BatchEndpoint {
11    /// Chat completions endpoint
12    #[serde(rename = "/v4/chat/completions")]
13    ChatCompletions,
14}
15
16/// Request body for creating a batch task
17#[derive(Clone, Serialize, Deserialize, Validate)]
18pub struct BatchCreateBody {
19    /// ID of the uploaded .jsonl file (purpose must be "batch")
20    #[validate(length(min = 1))]
21    pub input_file_id: String,
22
23    /// Endpoint to be used for all requests in the batch
24    pub endpoint: BatchEndpoint,
25
26    /// Whether to auto delete input file after processing (default: true)
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub auto_delete_input_file: Option<bool>,
29
30    /// String metadata for task management and tracking (up to 16 key-value pairs).
31    #[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    /// Create a new batch body from an input file id and target endpoint
52    /// (auto-delete defaults to `true`).
53    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    /// Set auto delete flag
63    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    /// Set string metadata.
69    pub fn with_metadata(mut self, v: BTreeMap<String, String>) -> Self {
70        self.metadata = Some(v);
71        self
72    }
73}
74
75/// Create batch request (POST /paas/v4/batches)
76pub struct BatchCreateRequest {
77    /// Request body.
78    body: BatchCreateBody,
79}
80
81impl BatchCreateRequest {
82    /// Build a new create-batch request with required fields
83    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    /// Set auto-delete flag (default true)
89    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    /// Set string metadata.
95    pub fn with_metadata(mut self, v: BTreeMap<String, String>) -> Self {
96        self.body = self.body.with_metadata(v);
97        self
98    }
99
100    /// Validate body using `validator`
101    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    /// Send request via a [`ZaiClient`] and parse typed response
131    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
142/// Response type for creating a batch task (same as a single item)
143pub 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}