zai_rs/batches/
create.rs

1use std::sync::Arc;
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use validator::Validate;
6
7use crate::{
8    ZaiResult,
9    client::{
10        endpoints::{ApiBase, EndpointConfig, paths},
11        http::{HttpClient, HttpClientConfig, parse_typed_response},
12    },
13};
14
15/// Endpoint for batch requests
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub enum BatchEndpoint {
18    /// Chat completions endpoint
19    #[serde(rename = "/v4/chat/completions")]
20    ChatCompletions,
21}
22
23/// Request body for creating a batch task
24#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
25pub struct CreateBatchBody {
26    /// ID of the uploaded .jsonl file (purpose must be "batch")
27    #[validate(length(min = 1))]
28    pub input_file_id: String,
29
30    /// Endpoint to be used for all requests in the batch
31    pub endpoint: BatchEndpoint,
32
33    /// Whether to auto delete input file after processing (default: true)
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub auto_delete_input_file: Option<bool>,
36
37    /// Arbitrary metadata for task management and tracking (up to 16 kv pairs)
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub metadata: Option<Value>,
40}
41
42impl CreateBatchBody {
43    pub fn new(input_file_id: impl Into<String>, endpoint: BatchEndpoint) -> Self {
44        Self {
45            input_file_id: input_file_id.into(),
46            endpoint,
47            auto_delete_input_file: Some(true),
48            metadata: None,
49        }
50    }
51
52    /// Set auto delete flag
53    pub fn with_auto_delete_input_file(mut self, v: bool) -> Self {
54        self.auto_delete_input_file = Some(v);
55        self
56    }
57
58    /// Set metadata object
59    pub fn with_metadata(mut self, v: Value) -> Self {
60        self.metadata = Some(v);
61        self
62    }
63}
64
65/// Create batch request (POST /paas/v4/batches)
66pub struct CreateBatchRequest {
67    pub key: String,
68    url: String,
69    endpoint_config: EndpointConfig,
70    api_base: ApiBase,
71    http_config: Arc<HttpClientConfig>,
72    pub body: CreateBatchBody,
73}
74
75impl CreateBatchRequest {
76    /// Build a new create-batch request with required fields
77    pub fn new(key: String, input_file_id: impl Into<String>, endpoint: BatchEndpoint) -> Self {
78        let body = CreateBatchBody::new(input_file_id, endpoint);
79        let endpoint_config = EndpointConfig::default();
80        let api_base = ApiBase::PaasV4;
81        let url = endpoint_config.url(&api_base, paths::BATCHES);
82        Self {
83            key,
84            url,
85            endpoint_config,
86            api_base,
87            http_config: Arc::new(HttpClientConfig::default()),
88            body,
89        }
90    }
91
92    fn rebuild_url(&mut self) {
93        self.url = self.endpoint_config.url(&self.api_base, paths::BATCHES);
94    }
95
96    pub fn with_base_url(mut self, base: impl Into<String>) -> Self {
97        self.api_base = ApiBase::Custom(base.into());
98        self.rebuild_url();
99        self
100    }
101
102    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
103        self.endpoint_config = endpoint_config;
104        self.rebuild_url();
105        self
106    }
107
108    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
109        self.http_config = Arc::new(config);
110        self
111    }
112
113    /// Set auto-delete flag (default true)
114    pub fn with_auto_delete_input_file(mut self, v: bool) -> Self {
115        self.body = self.body.with_auto_delete_input_file(v);
116        self
117    }
118
119    /// Set metadata object
120    pub fn with_metadata(mut self, v: serde_json::Value) -> Self {
121        self.body = self.body.with_metadata(v);
122        self
123    }
124
125    /// Validate body using `validator`
126    pub fn validate(&self) -> ZaiResult<()> {
127        self.body.validate().map_err(|e| e.into())
128    }
129
130    /// Send request and parse typed response
131    pub async fn send(&self) -> ZaiResult<CreateBatchResponse> {
132        self.validate()?;
133
134        let resp: reqwest::Response = self.post().await?;
135
136        let parsed = parse_typed_response::<CreateBatchResponse>(resp).await?;
137        Ok(parsed)
138    }
139}
140
141impl HttpClient for CreateBatchRequest {
142    type Body = CreateBatchBody;
143    type ApiUrl = String;
144    type ApiKey = String;
145
146    fn api_url(&self) -> &Self::ApiUrl {
147        &self.url
148    }
149    fn api_key(&self) -> &Self::ApiKey {
150        &self.key
151    }
152    fn body(&self) -> &Self::Body {
153        &self.body
154    }
155
156    fn http_config(&self) -> Arc<HttpClientConfig> {
157        Arc::clone(&self.http_config)
158    }
159}
160
161/// Response type for creating a batch task (same as a single item)
162pub type CreateBatchResponse = super::types::BatchItem;
zai_rs/batches/create.rs

zai_rs/batches/
create.rs
