use std::collections::BTreeSet;
use std::num::NonZeroUsize;
use serde::{Deserialize, Deserializer, Serialize};
use crate::KitError;
use crate::error::RemoteFailure;
use crate::error::ValidationViolation;
use crate::generation::{GenerateRequest, GenerateResponse, GenerationContinuation, Usage};
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(bound(serialize = ""), deny_unknown_fields)]
pub struct BatchInput<C = GenerationContinuation> {
custom_id: String,
request: GenerateRequest<C>,
}
impl<'de, C> Deserialize<'de> for BatchInput<C> {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(bound(deserialize = ""), deny_unknown_fields)]
struct Wire<C> {
custom_id: String,
request: GenerateRequest<C>,
}
let wire = Wire::deserialize(deserializer)?;
Self::new(wire.custom_id, wire.request).map_err(serde::de::Error::custom)
}
}
impl<C> BatchInput<C> {
pub fn new(
custom_id: impl Into<String>,
request: GenerateRequest<C>,
) -> Result<Self, KitError> {
let custom_id = custom_id.into();
if custom_id.trim().is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "batch.custom_id",
},
"batch custom id must not be empty",
));
}
Ok(Self { custom_id, request })
}
#[must_use]
pub fn custom_id(&self) -> &str {
&self.custom_id
}
#[must_use]
pub const fn request(&self) -> &GenerateRequest<C> {
&self.request
}
fn try_map_continuation<D>(
self,
map: &mut impl FnMut(C) -> Result<D, KitError>,
) -> Result<BatchInput<D>, KitError> {
Ok(BatchInput {
custom_id: self.custom_id,
request: self.request.try_map_continuation(map)?,
})
}
}
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(
bound(serialize = ""),
tag = "source",
rename_all = "snake_case",
deny_unknown_fields
)]
pub enum BatchSubmitRequest<C = GenerationContinuation> {
Inline {
rows: Vec<BatchInput<C>>,
},
ProviderFile {
file_id: String,
},
}
impl<'de, C> Deserialize<'de> for BatchSubmitRequest<C> {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(
bound(deserialize = ""),
tag = "source",
rename_all = "snake_case",
deny_unknown_fields
)]
enum Wire<C> {
Inline { rows: Vec<BatchInput<C>> },
ProviderFile { file_id: String },
}
let request = match Wire::deserialize(deserializer)? {
Wire::Inline { rows } => Self::Inline { rows },
Wire::ProviderFile { file_id } => Self::ProviderFile { file_id },
};
request.validate().map_err(serde::de::Error::custom)?;
Ok(request)
}
}
impl<C> BatchSubmitRequest<C> {
pub fn try_map_continuation<D>(
self,
mut map: impl FnMut(C) -> Result<D, KitError>,
) -> Result<BatchSubmitRequest<D>, KitError> {
Ok(match self {
Self::Inline { rows } => BatchSubmitRequest::Inline {
rows: rows
.into_iter()
.map(|row| row.try_map_continuation(&mut map))
.collect::<Result<_, _>>()?,
},
Self::ProviderFile { file_id } => BatchSubmitRequest::ProviderFile { file_id },
})
}
pub(crate) fn validate(&self) -> Result<(), KitError> {
match self {
Self::Inline { rows } => {
if rows.is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "batch.rows",
},
"inline batch rows must not be empty",
));
}
let mut ids = BTreeSet::new();
for row in rows {
if !ids.insert(row.custom_id()) {
return Err(KitError::invalid_input(
ValidationViolation::Duplicate {
field: "batch.custom_id",
},
format!("duplicate batch custom id `{}`", row.custom_id()),
));
}
}
Ok(())
}
Self::ProviderFile { file_id } if file_id.trim().is_empty() => {
Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "batch.file_id",
},
"batch file id must not be empty",
))
}
Self::ProviderFile { .. } => Ok(()),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "status", rename_all = "snake_case")]
pub enum BatchStatus {
Queued,
Running,
Completed,
Failed {
errors: Vec<RemoteFailure>,
},
Cancelled,
Expired,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct BatchJob {
pub id: String,
pub status: BatchStatus,
pub created_at: Option<String>,
pub completed_at: Option<String>,
pub usage: Option<Usage>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(
bound(serialize = "", deserialize = ""),
tag = "status",
rename_all = "snake_case"
)]
pub enum BatchResult<C = GenerationContinuation> {
Succeeded {
custom_id: String,
output: Box<GenerateResponse<C>>,
},
Failed {
custom_id: String,
error: RemoteFailure,
},
Cancelled {
custom_id: String,
},
Expired {
custom_id: String,
},
}
impl<C> BatchResult<C> {
pub fn try_map_continuation<D>(
self,
map: impl FnMut(C) -> Result<D, KitError>,
) -> Result<BatchResult<D>, KitError> {
Ok(match self {
Self::Succeeded { custom_id, output } => BatchResult::Succeeded {
custom_id,
output: Box::new((*output).try_map_continuation(map)?),
},
Self::Failed { custom_id, error } => BatchResult::Failed { custom_id, error },
Self::Cancelled { custom_id } => BatchResult::Cancelled { custom_id },
Self::Expired { custom_id } => BatchResult::Expired { custom_id },
})
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct BatchListRequest {
limit: Option<NonZeroUsize>,
cursor: Option<String>,
}
impl BatchListRequest {
#[must_use]
pub fn with_limit(mut self, limit: NonZeroUsize) -> Self {
self.limit = Some(limit);
self
}
pub fn with_cursor(mut self, cursor: impl Into<String>) -> Result<Self, KitError> {
let cursor = cursor.into();
if cursor.trim().is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "batch.cursor",
},
"batch cursor must not be empty",
));
}
self.cursor = Some(cursor);
Ok(self)
}
#[must_use]
pub const fn limit(&self) -> Option<NonZeroUsize> {
self.limit
}
#[must_use]
pub fn cursor(&self) -> Option<&str> {
self.cursor.as_deref()
}
#[must_use]
pub fn into_parts(self) -> (Option<NonZeroUsize>, Option<String>) {
(self.limit, self.cursor)
}
}
impl<'de> Deserialize<'de> for BatchListRequest {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Wire {
limit: Option<NonZeroUsize>,
cursor: Option<String>,
}
let wire = Wire::deserialize(deserializer)?;
let mut request = Self {
limit: wire.limit,
cursor: None,
};
if let Some(cursor) = wire.cursor {
request = request
.with_cursor(cursor)
.map_err(serde::de::Error::custom)?;
}
Ok(request)
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct BatchPage {
pub jobs: Vec<BatchJob>,
pub next_cursor: Option<String>,
}