sie_sdk/client/
batches.rs

1//! `/v1/batches`: OpenAI-compatible offline batches.
2
3use reqwest::Method;
4use serde_json::{Value, json};
5
6use crate::client::files::TRANSFER_TIMEOUT_FLOOR;
7use crate::client::{Client, meta::parse_json};
8use crate::error::Result;
9use crate::http::{PreparedRequest, headers};
10use crate::retry::RetryPolicy;
11use crate::types::{Batch, BatchList};
12
13/// The only endpoint batches can target today.
14pub const DEFAULT_BATCH_ENDPOINT: &str = "/v1/embeddings";
15/// The only completion window the server accepts today.
16pub const DEFAULT_COMPLETION_WINDOW: &str = "24h";
17
18/// The batches namespace. Obtain one with [`Client::batches`].
19#[derive(Debug, Clone)]
20pub struct Batches {
21    client: Client,
22}
23
24impl Client {
25    /// Operations on offline batches.
26    pub fn batches(&self) -> Batches {
27        Batches {
28            client: self.clone(),
29        }
30    }
31}
32
33impl Batches {
34    /// Queue a batch over a previously uploaded JSONL file.
35    pub fn create(&self, input_file_id: impl Into<String>) -> BatchCreate {
36        BatchCreate {
37            client: self.client.clone(),
38            input_file_id: input_file_id.into(),
39            endpoint: DEFAULT_BATCH_ENDPOINT.to_string(),
40            completion_window: DEFAULT_COMPLETION_WINDOW.to_string(),
41            metadata: None,
42        }
43    }
44
45    /// Fetch one batch.
46    pub async fn retrieve(&self, batch_id: &str) -> Result<Batch> {
47        let request = self
48            .client
49            .request(Method::GET, &format!("/v1/batches/{batch_id}"))?
50            .header("accept", headers::JSON_CONTENT_TYPE);
51        let response = self.client.send_once(request, RetryPolicy::NONE).await?;
52        parse_json(&response, "batch")
53    }
54
55    /// Ask the server to stop a batch.
56    pub async fn cancel(&self, batch_id: &str) -> Result<Batch> {
57        let request = self
58            .client
59            .request(Method::POST, &format!("/v1/batches/{batch_id}/cancel"))?
60            .json_headers();
61        let response = self.client.send_once(request, RetryPolicy::NONE).await?;
62        parse_json(&response, "batch")
63    }
64
65    /// List batches.
66    pub fn list(&self) -> BatchListRequest {
67        BatchListRequest {
68            client: self.client.clone(),
69            after: None,
70            limit: None,
71        }
72    }
73}
74
75/// Creates a batch. Build with [`Batches::create`].
76#[derive(Debug, Clone)]
77pub struct BatchCreate {
78    client: Client,
79    input_file_id: String,
80    endpoint: String,
81    completion_window: String,
82    metadata: Option<Value>,
83}
84
85impl BatchCreate {
86    /// Which endpoint the batch's lines target. Defaults to `/v1/embeddings`.
87    pub fn endpoint(mut self, endpoint: impl Into<String>) -> Self {
88        self.endpoint = endpoint.into();
89        self
90    }
91
92    /// How long the server may take. Defaults to `24h`.
93    pub fn completion_window(mut self, window: impl Into<String>) -> Self {
94        self.completion_window = window.into();
95        self
96    }
97
98    /// Caller-defined metadata stored with the batch.
99    pub fn metadata(mut self, metadata: Value) -> Self {
100        self.metadata = Some(metadata);
101        self
102    }
103
104    /// Send the request.
105    pub async fn send(self) -> Result<Batch> {
106        let mut body = json!({
107            "input_file_id": self.input_file_id,
108            "endpoint": self.endpoint,
109            "completion_window": self.completion_window,
110        });
111        if let Some(metadata) = self.metadata
112            && let Some(object) = body.as_object_mut()
113        {
114            object.insert("metadata".to_string(), metadata);
115        }
116
117        let request = self
118            .client
119            .request(Method::POST, "/v1/batches")?
120            .json_headers()
121            .body(serde_json::to_vec(&body).unwrap_or_default());
122        let response = self
123            .client
124            .send_with_timeout(request, RetryPolicy::NONE, TRANSFER_TIMEOUT_FLOOR)
125            .await?;
126        parse_json(&response, "batch")
127    }
128}
129
130/// Lists batches. Build with [`Batches::list`].
131#[derive(Debug, Clone)]
132pub struct BatchListRequest {
133    client: Client,
134    after: Option<String>,
135    limit: Option<u32>,
136}
137
138impl BatchListRequest {
139    /// Start after this batch id.
140    pub fn after(mut self, after: impl Into<String>) -> Self {
141        self.after = Some(after.into());
142        self
143    }
144
145    /// Cap on the page size.
146    pub fn limit(mut self, limit: u32) -> Self {
147        self.limit = Some(limit);
148        self
149    }
150
151    /// Fetch one page, with its pagination cursors.
152    pub async fn page(self) -> Result<BatchList> {
153        let mut url = self.client.url("/v1/batches")?;
154        {
155            let mut query = url.query_pairs_mut();
156            if let Some(after) = &self.after {
157                query.append_pair("after", after);
158            }
159            if let Some(limit) = self.limit {
160                query.append_pair("limit", &limit.to_string());
161            }
162        }
163        let request =
164            PreparedRequest::new(Method::GET, url).header("accept", headers::JSON_CONTENT_TYPE);
165        let response = self.client.send_once(request, RetryPolicy::NONE).await?;
166        parse_json(&response, "batch list")
167    }
168
169    /// Fetch one page and return only its batches.
170    pub async fn send(self) -> Result<Vec<Batch>> {
171        Ok(self.page().await?.data)
172    }
173}
sie_sdk/client/batches.rs

sie_sdk/client/
batches.rs
