use serde::{Deserialize, Serialize};
use serde_json::Value;
use validator::Validate;
use crate::{ZaiResult, client::ZaiClient};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BatchEndpoint {
#[serde(rename = "/v4/chat/completions")]
ChatCompletions,
}
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
pub struct CreateBatchBody {
#[validate(length(min = 1))]
pub input_file_id: String,
pub endpoint: BatchEndpoint,
#[serde(skip_serializing_if = "Option::is_none")]
pub auto_delete_input_file: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
}
impl CreateBatchBody {
pub fn new(input_file_id: impl Into<String>, endpoint: BatchEndpoint) -> Self {
Self {
input_file_id: input_file_id.into(),
endpoint,
auto_delete_input_file: Some(true),
metadata: None,
}
}
pub fn with_auto_delete_input_file(mut self, v: bool) -> Self {
self.auto_delete_input_file = Some(v);
self
}
pub fn with_metadata(mut self, v: Value) -> Self {
self.metadata = Some(v);
self
}
}
pub struct CreateBatchRequest {
pub body: CreateBatchBody,
}
impl CreateBatchRequest {
pub fn new(input_file_id: impl Into<String>, endpoint: BatchEndpoint) -> Self {
let body = CreateBatchBody::new(input_file_id, endpoint);
Self { body }
}
pub fn with_auto_delete_input_file(mut self, v: bool) -> Self {
self.body = self.body.with_auto_delete_input_file(v);
self
}
pub fn with_metadata(mut self, v: serde_json::Value) -> Self {
self.body = self.body.with_metadata(v);
self
}
pub fn validate(&self) -> ZaiResult<()> {
self.body.validate().map_err(std::convert::Into::into)
}
pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<CreateBatchResponse> {
self.validate()?;
let route = crate::client::routes::BATCHES_CREATE;
let url = client.endpoints().resolve_route(route, &[])?;
client
.send_json::<_, CreateBatchResponse>(route.method(), url, &self.body)
.await
}
}
pub type CreateBatchResponse = super::types::BatchItem;