1use async_trait::async_trait;
5use serde::{Deserialize, Serialize};
6use std::fmt;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum BatchEndpoint {
11 ChatCompletions,
12}
13
14impl BatchEndpoint {
15 pub fn as_str(self) -> &'static str {
16 match self {
17 Self::ChatCompletions => "/v1/chat/completions",
18 }
19 }
20}
21
22impl fmt::Display for BatchEndpoint {
23 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24 f.write_str(self.as_str())
25 }
26}
27
28#[derive(Debug, Clone, Serialize)]
31pub struct BatchItem {
32 pub custom_id: String,
33 pub body: serde_json::Value,
34}
35
36#[derive(Debug, Clone, Serialize)]
42pub struct SubmitBatch {
43 endpoint: &'static str,
44 model: String,
45 requests: Vec<BatchItem>,
46}
47
48#[derive(Debug, Clone, Serialize)]
50pub struct EmbeddingRequest {
51 model: String,
52 input: Vec<String>,
53 encoding_format: &'static str,
54}
55
56impl EmbeddingRequest {
57 pub fn new(model: impl Into<String>, input: Vec<String>) -> Self {
58 Self {
59 model: model.into(),
60 input,
61 encoding_format: "float",
62 }
63 }
64}
65
66#[derive(Debug, Clone, Deserialize)]
68pub struct EmbeddingData {
69 pub index: i64,
70 pub embedding: Vec<f32>,
71}
72
73#[derive(Debug, Clone, Deserialize)]
77pub struct EmbeddingResponseBody {
78 pub data: Vec<EmbeddingData>,
79}
80
81#[derive(Debug, Clone)]
83pub struct RawEmbeddingResponse {
84 pub raw: serde_json::Value,
85 pub provider_request_id: Option<String>,
86}
87
88impl SubmitBatch {
89 pub fn new(
90 endpoint: BatchEndpoint,
91 model: impl Into<String>,
92 requests: Vec<BatchItem>,
93 ) -> Self {
94 Self {
95 endpoint: endpoint.as_str(),
96 model: model.into(),
97 requests,
98 }
99 }
100
101 pub fn requests(&self) -> &[BatchItem] {
104 &self.requests
105 }
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
112#[serde(rename_all = "snake_case")]
113pub enum BatchStatus {
114 Validating,
115 InProgress,
116 Finalizing,
117 Completed,
118 Failed,
119 Expired,
120 Cancelling,
121 Cancelled,
122 #[serde(other)]
125 Unknown,
126}
127
128impl BatchStatus {
129 pub fn is_terminal(self) -> bool {
130 matches!(
131 self,
132 Self::Completed | Self::Failed | Self::Expired | Self::Cancelled
133 )
134 }
135}
136
137#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)]
140pub struct RequestCounts {
141 #[serde(default)]
142 pub total: i64,
143 #[serde(default)]
144 pub completed: i64,
145 #[serde(default)]
146 pub failed: i64,
147}
148
149#[derive(Debug, Clone, Copy, Default, Deserialize)]
152pub struct BatchUsage {
153 #[serde(default)]
154 pub prompt_tokens: Option<i64>,
155 #[serde(default)]
156 pub completion_tokens: Option<i64>,
157 #[serde(default)]
158 pub total_tokens: Option<i64>,
159 #[serde(default)]
160 pub cost: Option<f64>,
161 #[serde(default)]
162 pub is_byok: Option<bool>,
163}
164
165#[derive(Debug, Clone, Serialize, Deserialize)]
172#[serde(untagged)]
173pub enum BatchResultLine {
174 Parsed(ResultLine),
175 Undecodable(serde_json::Value),
176}
177
178#[derive(Debug, Clone, Serialize, Deserialize)]
181pub struct ResultLine {
182 #[serde(default)]
183 pub id: Option<String>,
184 pub custom_id: String,
185 #[serde(default)]
186 pub response: Option<ResultResponse>,
187 #[serde(default)]
188 pub error: Option<serde_json::Value>,
189 #[serde(flatten)]
190 pub extra: serde_json::Map<String, serde_json::Value>,
191}
192
193#[derive(Debug, thiserror::Error)]
196pub enum BatchLineError {
197 #[error("result envelope decode: line does not match the documented result envelope")]
200 Undecodable,
201 #[error("line carries an error instead of a response")]
202 NoResponse,
203 #[error("provider returned status {0}")]
204 ProviderStatus(u16),
205 #[error("response body has no assistant message content")]
206 NoContent,
207}
208
209impl BatchResultLine {
210 pub fn custom_id(&self) -> Option<&str> {
213 match self {
214 Self::Parsed(line) => Some(&line.custom_id),
215 Self::Undecodable(raw) => raw.get("custom_id").and_then(serde_json::Value::as_str),
216 }
217 }
218
219 pub fn raw(&self) -> serde_json::Value {
222 match self {
223 Self::Parsed(line) => {
224 serde_json::to_value(line).expect("result line re-serialization cannot fail")
225 }
226 Self::Undecodable(raw) => raw.clone(),
227 }
228 }
229
230 pub fn chat_content(&self) -> Result<&str, BatchLineError> {
233 let line = match self {
234 Self::Parsed(line) => line,
235 Self::Undecodable(_) => return Err(BatchLineError::Undecodable),
236 };
237 let response = line.response.as_ref().ok_or(BatchLineError::NoResponse)?;
238 if response.status_code != 200 {
239 return Err(BatchLineError::ProviderStatus(response.status_code));
240 }
241 chat_completion_content(&response.body).ok_or(BatchLineError::NoContent)
242 }
243}
244
245pub fn chat_completion_content(body: &serde_json::Value) -> Option<&str> {
249 body.get("choices")
250 .and_then(|choices| choices.get(0))
251 .and_then(|choice| choice.get("message"))
252 .and_then(|message| message.get("content"))
253 .and_then(serde_json::Value::as_str)
254}
255
256#[derive(Debug, Clone, Serialize, Deserialize)]
259pub struct ResultResponse {
260 pub status_code: u16,
261 #[serde(default)]
262 pub request_id: Option<String>,
263 pub body: serde_json::Value,
264 #[serde(flatten)]
265 pub extra: serde_json::Map<String, serde_json::Value>,
266}
267
268#[derive(Debug, Clone)]
273pub struct RawChatResponse {
274 pub raw: serde_json::Value,
275 pub provider_request_id: Option<String>,
276}
277
278impl RawChatResponse {
279 pub fn content(&self) -> Result<&str, BatchLineError> {
281 chat_completion_content(&self.raw).ok_or(BatchLineError::NoContent)
282 }
283}
284
285#[derive(Debug, Clone, Deserialize)]
287pub struct BatchObject {
288 pub id: String,
289 #[serde(default)]
290 pub endpoint: Option<String>,
291 #[serde(default)]
292 pub model: Option<String>,
293 pub status: BatchStatus,
294 #[serde(default)]
295 pub request_counts: Option<RequestCounts>,
296 #[serde(default)]
297 pub usage: Option<BatchUsage>,
298 #[serde(default)]
302 pub results: Option<Vec<BatchResultLine>>,
303 #[serde(default)]
304 pub error: Option<serde_json::Value>,
305 #[serde(default)]
306 pub created_at: Option<i64>,
307 #[serde(default)]
308 pub finalized_at: Option<i64>,
309 #[serde(skip)]
311 pub provider_request_id: Option<String>,
312}
313
314impl BatchObject {
315 pub fn batch_age_seconds(&self) -> Option<i64> {
318 let created = self.created_at?;
319 let end = self
320 .finalized_at
321 .unwrap_or_else(|| finalized_or_now(None).unwrap_or(0));
322 Some(end - created)
323 }
324}
325
326fn finalized_or_now(finalized_at: Option<i64>) -> Option<i64> {
327 finalized_at.or_else(|| {
328 std::time::SystemTime::now()
329 .duration_since(std::time::UNIX_EPOCH)
330 .ok()
331 .map(|d| d.as_secs() as i64)
332 })
333}
334
335#[derive(Debug, thiserror::Error)]
339pub enum TransportError {
340 #[error("provider rejected the request ({status}): {message}")]
343 Api {
344 status: u16,
345 message: String,
346 provider_request_id: Option<String>,
347 },
348 #[error("custom_id exceeds the provider's 64-character cap: {custom_id}")]
351 CustomIdTooLong { custom_id: String },
352 #[error("http transport failure: {source}")]
353 Http {
354 #[source]
355 source: reqwest::Error,
356 provider_request_id: Option<String>,
357 },
358 #[error("response decode failure: {source}")]
359 Decode {
360 #[source]
361 source: serde_json::Error,
362 provider_request_id: Option<String>,
363 },
364}
365
366impl From<reqwest::Error> for TransportError {
367 fn from(source: reqwest::Error) -> Self {
368 Self::Http {
369 source,
370 provider_request_id: None,
371 }
372 }
373}
374
375impl From<serde_json::Error> for TransportError {
376 fn from(source: serde_json::Error) -> Self {
377 Self::Decode {
378 source,
379 provider_request_id: None,
380 }
381 }
382}
383
384impl TransportError {
385 pub fn retryable(&self) -> bool {
386 fn status_retryable(status: u16) -> bool {
387 status == 408 || status == 429 || status >= 500
388 }
389 match self {
390 Self::Api { status, .. } => status_retryable(*status),
391 Self::CustomIdTooLong { .. } => false,
392 Self::Http { source, .. } => {
393 source.is_timeout()
394 || source.is_connect()
395 || source.is_body()
396 || source.is_decode()
397 || source
398 .status()
399 .is_some_and(|status| status_retryable(status.as_u16()))
400 }
401 Self::Decode { .. } => false,
402 }
403 }
404
405 pub fn provider_request_id(&self) -> Option<&str> {
407 match self {
408 Self::CustomIdTooLong { .. } => None,
409 Self::Api {
410 provider_request_id,
411 ..
412 } => provider_request_id.as_deref(),
413 Self::Http {
414 provider_request_id,
415 ..
416 }
417 | Self::Decode {
418 provider_request_id,
419 ..
420 } => provider_request_id.as_deref(),
421 }
422 }
423
424 pub fn not_found(&self) -> bool {
428 matches!(self, Self::Api { status: 404, .. })
429 }
430}
431
432#[async_trait]
435pub trait BatchTransport: Send + Sync {
436 async fn submit(
437 &self,
438 batch: &SubmitBatch,
439 idempotency_key: &str,
440 ) -> Result<BatchObject, TransportError>;
441 async fn get(&self, batch_id: &str) -> Result<BatchObject, TransportError>;
442}
443
444#[async_trait]
446pub trait EmbeddingTransport: Send + Sync {
447 async fn embed(
448 &self,
449 model: &str,
450 inputs: &[String],
451 idempotency_key: &str,
452 ) -> Result<RawEmbeddingResponse, TransportError>;
453}