Skip to main content

sie_sdk/types/
resources.rs

1//! Files, batches, connections and pools.
2
3// Wire-mirror types: field names are the API contract itself, and the ones whose
4// meaning is not obvious carry their own doc comment.
5#![allow(missing_docs)]
6
7use std::collections::HashMap;
8
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11
12/// An uploaded file.
13#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
14pub struct File {
15    #[serde(default, deserialize_with = "crate::types::null_as_default")]
16    pub id: String,
17    #[serde(
18        default,
19        deserialize_with = "crate::types::null_as_default",
20        rename = "object"
21    )]
22    pub object_kind: String,
23    #[serde(default, deserialize_with = "crate::types::null_as_default")]
24    pub bytes: u64,
25    #[serde(default, deserialize_with = "crate::types::null_as_default")]
26    pub created_at: i64,
27    #[serde(default, deserialize_with = "crate::types::null_as_default")]
28    pub filename: String,
29    #[serde(default, deserialize_with = "crate::types::null_as_default")]
30    pub purpose: String,
31    #[serde(default, deserialize_with = "crate::types::null_as_default")]
32    pub status: String,
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub expires_at: Option<i64>,
35}
36
37/// One page of the file list.
38#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
39pub struct FileList {
40    #[serde(
41        default,
42        deserialize_with = "crate::types::null_as_default",
43        rename = "object"
44    )]
45    pub object_kind: String,
46    #[serde(default, deserialize_with = "crate::types::null_as_default")]
47    pub data: Vec<File>,
48    #[serde(default)]
49    pub first_id: Option<String>,
50    #[serde(default)]
51    pub last_id: Option<String>,
52    #[serde(default, deserialize_with = "crate::types::null_as_default")]
53    pub has_more: bool,
54}
55
56/// Confirmation that a file was deleted.
57#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
58pub struct FileDeleted {
59    #[serde(default, deserialize_with = "crate::types::null_as_default")]
60    pub id: String,
61    #[serde(
62        default,
63        deserialize_with = "crate::types::null_as_default",
64        rename = "object"
65    )]
66    pub object_kind: String,
67    #[serde(default, deserialize_with = "crate::types::null_as_default")]
68    pub deleted: bool,
69}
70
71/// How many requests in a batch are in each state.
72#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
73pub struct BatchRequestCounts {
74    #[serde(default, deserialize_with = "crate::types::null_as_default")]
75    pub total: u64,
76    #[serde(default, deserialize_with = "crate::types::null_as_default")]
77    pub completed: u64,
78    #[serde(default, deserialize_with = "crate::types::null_as_default")]
79    pub failed: u64,
80}
81
82/// An offline batch of requests.
83#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
84pub struct Batch {
85    #[serde(default, deserialize_with = "crate::types::null_as_default")]
86    pub id: String,
87    #[serde(
88        default,
89        deserialize_with = "crate::types::null_as_default",
90        rename = "object"
91    )]
92    pub object_kind: String,
93    #[serde(default, deserialize_with = "crate::types::null_as_default")]
94    pub endpoint: String,
95    #[serde(default, deserialize_with = "crate::types::null_as_default")]
96    pub input_file_id: String,
97    #[serde(default, deserialize_with = "crate::types::null_as_default")]
98    pub completion_window: String,
99    #[serde(default, deserialize_with = "crate::types::null_as_default")]
100    pub status: String,
101    #[serde(default)]
102    pub output_file_id: Option<String>,
103    #[serde(default)]
104    pub error_file_id: Option<String>,
105    #[serde(default, deserialize_with = "crate::types::null_as_default")]
106    pub created_at: i64,
107    #[serde(default)]
108    pub finished_at: Option<i64>,
109    #[serde(default)]
110    pub in_progress_at: Option<i64>,
111    #[serde(default)]
112    pub completed_at: Option<i64>,
113    #[serde(default)]
114    pub failed_at: Option<i64>,
115    #[serde(default)]
116    pub expired_at: Option<i64>,
117    #[serde(default)]
118    pub cancelled_at: Option<i64>,
119    #[serde(default, deserialize_with = "crate::types::null_as_default")]
120    pub request_counts: BatchRequestCounts,
121    #[serde(default, skip_serializing_if = "Option::is_none")]
122    pub errors: Option<Value>,
123    #[serde(default)]
124    pub metadata: Option<Value>,
125}
126
127/// One page of the batch list.
128#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
129pub struct BatchList {
130    #[serde(
131        default,
132        deserialize_with = "crate::types::null_as_default",
133        rename = "object"
134    )]
135    pub object_kind: String,
136    #[serde(default, deserialize_with = "crate::types::null_as_default")]
137    pub data: Vec<Batch>,
138    #[serde(default)]
139    pub first_id: Option<String>,
140    #[serde(default)]
141    pub last_id: Option<String>,
142    #[serde(default, deserialize_with = "crate::types::null_as_default")]
143    pub has_more: bool,
144}
145
146/// A stored credential for a data store a connector job can read or write.
147///
148/// The secret itself is never returned.
149#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
150pub struct Connection {
151    #[serde(default, deserialize_with = "crate::types::null_as_default")]
152    pub id: i64,
153    #[serde(
154        default,
155        deserialize_with = "crate::types::null_as_default",
156        rename = "type"
157    )]
158    pub kind: String,
159    #[serde(default, deserialize_with = "crate::types::null_as_default")]
160    pub name: String,
161    /// Bumped whenever the secret is rotated.
162    #[serde(default, deserialize_with = "crate::types::null_as_default")]
163    pub authorization_generation: i64,
164    #[serde(default)]
165    pub source_schema: Option<String>,
166    #[serde(default)]
167    pub sink_schema: Option<String>,
168    #[serde(default, deserialize_with = "crate::types::null_as_default")]
169    pub created_at: f64,
170}
171
172/// A newly created connection, with the org it belongs to.
173#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
174pub struct ConnectionCreated {
175    #[serde(default, deserialize_with = "crate::types::null_as_default")]
176    pub org: String,
177    #[serde(default, deserialize_with = "crate::types::null_as_default")]
178    pub account_id: i64,
179    #[serde(flatten)]
180    pub connection: Connection,
181}
182
183/// Confirmation that a connection was revoked.
184#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
185pub struct ConnectionRevoked {
186    #[serde(default, deserialize_with = "crate::types::null_as_default")]
187    pub org: String,
188    #[serde(default, deserialize_with = "crate::types::null_as_default")]
189    pub account_id: i64,
190    #[serde(default, deserialize_with = "crate::types::null_as_default")]
191    pub name: String,
192    #[serde(default, deserialize_with = "crate::types::null_as_default")]
193    pub state: String,
194}
195
196/// A worker assigned to a pool.
197#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
198pub struct AssignedWorkerInfo {
199    #[serde(default, deserialize_with = "crate::types::null_as_default")]
200    pub name: String,
201    #[serde(default, deserialize_with = "crate::types::null_as_default")]
202    pub url: String,
203    #[serde(default, deserialize_with = "crate::types::null_as_default")]
204    pub gpu: String,
205    #[serde(default, deserialize_with = "crate::types::null_as_default")]
206    pub bundle: String,
207}
208
209/// A pool's live state.
210#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
211pub struct PoolStatusInfo {
212    /// `pending`, `active` or `expired`.
213    #[serde(default, deserialize_with = "crate::types::null_as_default")]
214    pub state: String,
215    #[serde(default, deserialize_with = "crate::types::null_as_default")]
216    pub assigned_workers: Vec<AssignedWorkerInfo>,
217    #[serde(default, deserialize_with = "crate::types::null_as_default")]
218    pub created_at: f64,
219    #[serde(default, deserialize_with = "crate::types::null_as_default")]
220    pub last_renewed: f64,
221}
222
223/// A pool's requested shape, as the server echoes it back.
224#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
225pub struct PoolSpecResponse {
226    #[serde(default, deserialize_with = "crate::types::null_as_default")]
227    pub name: String,
228    #[serde(default, deserialize_with = "crate::types::null_as_default")]
229    pub queue_pool: String,
230    #[serde(default, deserialize_with = "crate::types::null_as_default")]
231    pub gpus: HashMap<String, u32>,
232    /// Ceiling on assigned workers per GPU type.
233    #[serde(default, deserialize_with = "crate::types::null_as_default")]
234    pub gpu_caps: HashMap<String, u32>,
235    #[serde(default)]
236    pub bundle: Option<String>,
237    #[serde(default, deserialize_with = "crate::types::null_as_default")]
238    pub minimum_worker_count: u32,
239    #[serde(default, deserialize_with = "crate::types::null_as_default")]
240    pub pinned_models: Vec<String>,
241}
242
243/// A pool, with its spec and its status.
244#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
245pub struct PoolInfo {
246    #[serde(default, deserialize_with = "crate::types::null_as_default")]
247    pub name: String,
248    #[serde(default, deserialize_with = "crate::types::null_as_default")]
249    pub spec: PoolSpecResponse,
250    #[serde(default, deserialize_with = "crate::types::null_as_default")]
251    pub status: PoolStatusInfo,
252}
253
254/// A pool as it appears in a listing.
255#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
256pub struct PoolListItem {
257    #[serde(default, deserialize_with = "crate::types::null_as_default")]
258    pub name: String,
259    #[serde(default, deserialize_with = "crate::types::null_as_default")]
260    pub state: String,
261    #[serde(default, deserialize_with = "crate::types::null_as_default")]
262    pub gpus: HashMap<String, u32>,
263    #[serde(default, deserialize_with = "crate::types::null_as_default")]
264    pub worker_count: u32,
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270    use serde_json::json;
271
272    #[test]
273    fn connection_created_flattens_the_connection_it_wraps() {
274        let created: ConnectionCreated = serde_json::from_value(json!({
275            "org": "acme", "account_id": 7, "id": 1, "type": "postgres", "name": "warehouse",
276            "authorization_generation": 1, "created_at": 1.0
277        }))
278        .unwrap();
279        assert_eq!(created.org, "acme");
280        assert_eq!(created.connection.name, "warehouse");
281        assert_eq!(created.connection.kind, "postgres");
282        assert!(created.connection.source_schema.is_none());
283    }
284
285    #[test]
286    fn batches_tolerate_the_null_timestamp_fields() {
287        let batch: Batch = serde_json::from_value(json!({
288            "id": "batch_1", "object": "batch", "endpoint": "/v1/embeddings",
289            "input_file_id": "file_1", "completion_window": "24h", "status": "in_progress",
290            "output_file_id": null, "failed_at": null, "created_at": 100,
291            "request_counts": {"total": 10, "completed": 3, "failed": 0}
292        }))
293        .unwrap();
294        assert_eq!(batch.status, "in_progress");
295        assert_eq!(batch.output_file_id, None);
296        assert_eq!(batch.request_counts.completed, 3);
297    }
298
299    #[test]
300    fn a_file_without_the_optional_ttl_fields_decodes() {
301        let file: File = serde_json::from_value(json!({
302            "id": "file_1", "object": "file", "bytes": 42, "created_at": 1,
303            "filename": "in.jsonl", "purpose": "batch", "status": "processed"
304        }))
305        .unwrap();
306        assert_eq!(file.bytes, 42);
307        assert!(file.expires_at.is_none());
308    }
309}