sim_lib_openai_server/storage/
objects.rs

1use sim_kernel::{ContentId, Expr, Result, Symbol};
2
3use crate::objects::{GatewayResponse, content_id_expr};
4
5use super::vector::GatewayVectorStore;
6
7/// Record kind tag designating a gateway file.
8pub const GATEWAY_FILE_KIND: &str = "openai-gateway/file";
9/// Record kind tag designating a gateway batch.
10pub const GATEWAY_BATCH_KIND: &str = "openai-gateway/batch";
11/// Record kind tag designating a gateway thread.
12pub const GATEWAY_THREAD_KIND: &str = "openai-gateway/thread";
13/// Record kind tag designating a gateway thread message.
14pub const GATEWAY_THREAD_MESSAGE_KIND: &str = "openai-gateway/thread-message";
15
16/// A stored response object linking a response id to its content-addressed
17/// [`GatewayResponse`] and the records that produced it.
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct StoredGatewayResponse {
20    response_id: String,
21    content_id: ContentId,
22    response: GatewayResponse,
23    pub(crate) request_content_id: Option<ContentId>,
24    pub(crate) run_content_id: Option<ContentId>,
25    pub(crate) event_content_ids: Vec<ContentId>,
26    pub(crate) parent_response_id: Option<String>,
27    pub(crate) owner_key_id: Option<String>,
28}
29
30impl StoredGatewayResponse {
31    /// Creates a stored response with no linked request, run, events, or parent.
32    pub fn new(
33        response_id: impl Into<String>,
34        content_id: ContentId,
35        response: GatewayResponse,
36    ) -> Self {
37        Self {
38            response_id: response_id.into(),
39            content_id,
40            response,
41            request_content_id: None,
42            run_content_id: None,
43            event_content_ids: Vec::new(),
44            parent_response_id: None,
45            owner_key_id: None,
46        }
47    }
48
49    /// Returns the public response identifier.
50    pub fn response_id(&self) -> &str {
51        &self.response_id
52    }
53
54    /// Returns the content id of the stored [`GatewayResponse`].
55    pub fn content_id(&self) -> &ContentId {
56        &self.content_id
57    }
58
59    /// Returns the stored response value.
60    pub fn response(&self) -> &GatewayResponse {
61        &self.response
62    }
63
64    /// Returns the gateway key id that created this response, if any.
65    pub fn owner_key_id(&self) -> Option<&str> {
66        self.owner_key_id.as_deref()
67    }
68}
69
70/// Store for response objects keyed by their public response id.
71pub trait GatewayResponseObjectStore {
72    /// Stores a response object, replacing any existing entry with the same id.
73    fn put_response_object(&mut self, response: StoredGatewayResponse) -> Result<()>;
74    /// Returns the response object with the given id, if present.
75    fn response_object(&self, response_id: &str) -> Option<StoredGatewayResponse>;
76}
77
78/// Lifecycle status of a [`GatewayBatch`].
79#[derive(Clone, Debug, PartialEq, Eq)]
80pub enum GatewayBatchStatus {
81    /// The batch has been accepted but not yet started.
82    Queued,
83    /// The batch is currently being processed.
84    InProgress,
85    /// The batch finished processing.
86    Completed,
87    /// The batch was cancelled before completion.
88    Cancelled,
89}
90
91impl GatewayBatchStatus {
92    /// Returns the OpenAI wire string for this status (e.g. `"in_progress"`).
93    pub fn as_str(&self) -> &'static str {
94        match self {
95            Self::Queued => "queued",
96            Self::InProgress => "in_progress",
97            Self::Completed => "completed",
98            Self::Cancelled => "cancelled",
99        }
100    }
101}
102
103/// Per-request tallies for a [`GatewayBatch`].
104#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
105pub struct GatewayBatchCounts {
106    total: u64,
107    completed: u64,
108    failed: u64,
109    cancelled: u64,
110}
111
112impl GatewayBatchCounts {
113    /// Creates a counts record from the total, completed, failed, and cancelled tallies.
114    pub fn new(total: u64, completed: u64, failed: u64, cancelled: u64) -> Self {
115        Self {
116            total,
117            completed,
118            failed,
119            cancelled,
120        }
121    }
122
123    /// Returns the total number of requests in the batch.
124    pub fn total(&self) -> u64 {
125        self.total
126    }
127
128    /// Returns the number of completed requests.
129    pub fn completed(&self) -> u64 {
130        self.completed
131    }
132
133    /// Returns the number of failed requests.
134    pub fn failed(&self) -> u64 {
135        self.failed
136    }
137
138    /// Returns the number of cancelled requests.
139    pub fn cancelled(&self) -> u64 {
140        self.cancelled
141    }
142
143    /// Encodes the counts as a SIM [`Expr`] map record.
144    pub fn to_expr(self) -> Expr {
145        Expr::Map(vec![
146            field("total", Expr::String(self.total.to_string())),
147            field("completed", Expr::String(self.completed.to_string())),
148            field("failed", Expr::String(self.failed.to_string())),
149            field("cancelled", Expr::String(self.cancelled.to_string())),
150        ])
151    }
152}
153
154/// A gateway batch job over an input file, with status and per-request counts.
155#[derive(Clone, Debug, PartialEq, Eq)]
156pub struct GatewayBatch {
157    id: String,
158    input_file_id: String,
159    endpoint: String,
160    status: GatewayBatchStatus,
161    output_file_id: Option<String>,
162    error_file_id: Option<String>,
163    created_at_ms: u64,
164    completed_at_ms: Option<u64>,
165    cancelled_at_ms: Option<u64>,
166    request_counts: GatewayBatchCounts,
167}
168
169impl GatewayBatch {
170    /// Creates a new batch in the [`GatewayBatchStatus::Queued`] state.
171    pub fn new(
172        id: impl Into<String>,
173        input_file_id: impl Into<String>,
174        endpoint: impl Into<String>,
175        created_at_ms: u64,
176        request_counts: GatewayBatchCounts,
177    ) -> Self {
178        Self {
179            id: id.into(),
180            input_file_id: input_file_id.into(),
181            endpoint: endpoint.into(),
182            status: GatewayBatchStatus::Queued,
183            output_file_id: None,
184            error_file_id: None,
185            created_at_ms,
186            completed_at_ms: None,
187            cancelled_at_ms: None,
188            request_counts,
189        }
190    }
191
192    /// Returns the batch transitioned to [`GatewayBatchStatus::Completed`] with
193    /// the given output/error files, completion time, and final counts.
194    pub fn complete(
195        mut self,
196        output_file_id: Option<String>,
197        error_file_id: Option<String>,
198        completed_at_ms: u64,
199        request_counts: GatewayBatchCounts,
200    ) -> Self {
201        self.status = GatewayBatchStatus::Completed;
202        self.output_file_id = output_file_id;
203        self.error_file_id = error_file_id;
204        self.completed_at_ms = Some(completed_at_ms);
205        self.request_counts = request_counts;
206        self
207    }
208
209    /// Returns the batch transitioned to [`GatewayBatchStatus::Cancelled`],
210    /// counting any still-queued requests as cancelled.
211    pub fn cancel(mut self, cancelled_at_ms: u64) -> Self {
212        let queued = self.request_counts.total.saturating_sub(
213            self.request_counts.completed
214                + self.request_counts.failed
215                + self.request_counts.cancelled,
216        );
217        self.status = GatewayBatchStatus::Cancelled;
218        self.cancelled_at_ms = Some(cancelled_at_ms);
219        self.request_counts.cancelled += queued;
220        self
221    }
222
223    /// Returns the batch identifier.
224    pub fn id(&self) -> &str {
225        &self.id
226    }
227
228    /// Returns the id of the input file backing this batch.
229    pub fn input_file_id(&self) -> &str {
230        &self.input_file_id
231    }
232
233    /// Returns the target API endpoint for the batched requests.
234    pub fn endpoint(&self) -> &str {
235        &self.endpoint
236    }
237
238    /// Returns the current lifecycle status.
239    pub fn status(&self) -> &GatewayBatchStatus {
240        &self.status
241    }
242
243    /// Returns the output file id, once the batch has produced one.
244    pub fn output_file_id(&self) -> Option<&str> {
245        self.output_file_id.as_deref()
246    }
247
248    /// Returns the error file id, if the batch produced one.
249    pub fn error_file_id(&self) -> Option<&str> {
250        self.error_file_id.as_deref()
251    }
252
253    /// Returns the creation timestamp in milliseconds since the Unix epoch.
254    pub fn created_at_ms(&self) -> u64 {
255        self.created_at_ms
256    }
257
258    /// Returns the completion timestamp in milliseconds, if completed.
259    pub fn completed_at_ms(&self) -> Option<u64> {
260        self.completed_at_ms
261    }
262
263    /// Returns the cancellation timestamp in milliseconds, if cancelled.
264    pub fn cancelled_at_ms(&self) -> Option<u64> {
265        self.cancelled_at_ms
266    }
267
268    /// Returns the current per-request counts.
269    pub fn request_counts(&self) -> GatewayBatchCounts {
270        self.request_counts
271    }
272
273    /// Encodes the batch as a SIM [`Expr`] map record.
274    pub fn to_expr(&self) -> Expr {
275        Expr::Map(vec![
276            field("kind", Expr::String(GATEWAY_BATCH_KIND.to_owned())),
277            field("id", Expr::String(self.id.clone())),
278            field("input-file-id", Expr::String(self.input_file_id.clone())),
279            field("endpoint", Expr::String(self.endpoint.clone())),
280            field("status", Expr::Symbol(Symbol::new(self.status.as_str()))),
281            optional_string_field("output-file-id", self.output_file_id.as_deref()),
282            optional_string_field("error-file-id", self.error_file_id.as_deref()),
283            field(
284                "created-at-ms",
285                Expr::String(self.created_at_ms.to_string()),
286            ),
287            optional_u64_field("completed-at-ms", self.completed_at_ms),
288            optional_u64_field("cancelled-at-ms", self.cancelled_at_ms),
289            field("request-counts", self.request_counts.to_expr()),
290        ])
291    }
292}
293
294/// Where a [`GatewayFile`]'s bytes live.
295#[derive(Clone, Debug, PartialEq, Eq)]
296pub enum GatewayFileStorageRef {
297    /// Bytes held in an in-memory content-addressed blob.
298    Memory {
299        /// Content id of the stored bytes.
300        content_id: ContentId,
301    },
302    /// Bytes held on the table-backed filesystem at the given path.
303    TableFs {
304        /// Filesystem path of the stored bytes.
305        path: String,
306    },
307}
308
309impl GatewayFileStorageRef {
310    /// Creates a [`GatewayFileStorageRef::Memory`] reference for `content_id`.
311    pub fn memory(content_id: ContentId) -> Self {
312        Self::Memory { content_id }
313    }
314
315    /// Creates a [`GatewayFileStorageRef::TableFs`] reference for `path`.
316    pub fn table_fs(path: impl Into<String>) -> Self {
317        Self::TableFs { path: path.into() }
318    }
319
320    /// Encodes the storage reference as a SIM [`Expr`] map record.
321    pub fn to_expr(&self) -> Expr {
322        match self {
323            Self::Memory { content_id } => Expr::Map(vec![
324                field("kind", Expr::Symbol(Symbol::new("memory"))),
325                field("content-id", content_id_expr(content_id)),
326            ]),
327            Self::TableFs { path } => Expr::Map(vec![
328                field("kind", Expr::Symbol(Symbol::new("table-fs"))),
329                field("path", Expr::String(path.clone())),
330            ]),
331        }
332    }
333}
334
335/// A gateway file record: its metadata plus a reference to where its bytes live.
336#[derive(Clone, Debug, PartialEq, Eq)]
337pub struct GatewayFile {
338    id: String,
339    filename: String,
340    purpose: String,
341    bytes: u64,
342    created_at_ms: u64,
343    storage_ref: GatewayFileStorageRef,
344}
345
346impl GatewayFile {
347    /// Creates a file record with the given metadata and storage reference.
348    pub fn new(
349        id: impl Into<String>,
350        filename: impl Into<String>,
351        purpose: impl Into<String>,
352        bytes: u64,
353        created_at_ms: u64,
354        storage_ref: GatewayFileStorageRef,
355    ) -> Self {
356        Self {
357            id: id.into(),
358            filename: filename.into(),
359            purpose: purpose.into(),
360            bytes,
361            created_at_ms,
362            storage_ref,
363        }
364    }
365
366    /// Returns the file identifier.
367    pub fn id(&self) -> &str {
368        &self.id
369    }
370
371    /// Returns the original filename.
372    pub fn filename(&self) -> &str {
373        &self.filename
374    }
375
376    /// Returns the declared purpose of the file (e.g. `"batch"`, `"assistants"`).
377    pub fn purpose(&self) -> &str {
378        &self.purpose
379    }
380
381    /// Returns the file size in bytes.
382    pub fn bytes(&self) -> u64 {
383        self.bytes
384    }
385
386    /// Returns the creation timestamp in milliseconds since the Unix epoch.
387    pub fn created_at_ms(&self) -> u64 {
388        self.created_at_ms
389    }
390
391    /// Returns a reference to where the file's bytes are stored.
392    pub fn storage_ref(&self) -> &GatewayFileStorageRef {
393        &self.storage_ref
394    }
395
396    /// Encodes the file record as a SIM [`Expr`] map record.
397    pub fn to_expr(&self) -> Expr {
398        Expr::Map(vec![
399            field("kind", Expr::String(GATEWAY_FILE_KIND.to_owned())),
400            field("id", Expr::String(self.id.clone())),
401            field("filename", Expr::String(self.filename.clone())),
402            field("purpose", Expr::String(self.purpose.clone())),
403            field("bytes", Expr::String(self.bytes.to_string())),
404            field(
405                "created-at-ms",
406                Expr::String(self.created_at_ms.to_string()),
407            ),
408            field("storage-ref", self.storage_ref.to_expr()),
409        ])
410    }
411}
412
413/// A gateway thread record with creation time and key/value metadata.
414#[derive(Clone, Debug, PartialEq, Eq)]
415pub struct GatewayThread {
416    id: String,
417    created_at_ms: u64,
418    metadata: Vec<(String, String)>,
419}
420
421impl GatewayThread {
422    /// Creates a thread record with the given id, creation time, and metadata.
423    pub fn new(id: impl Into<String>, created_at_ms: u64, metadata: Vec<(String, String)>) -> Self {
424        Self {
425            id: id.into(),
426            created_at_ms,
427            metadata,
428        }
429    }
430
431    /// Returns the thread identifier.
432    pub fn id(&self) -> &str {
433        &self.id
434    }
435
436    /// Returns the creation timestamp in milliseconds since the Unix epoch.
437    pub fn created_at_ms(&self) -> u64 {
438        self.created_at_ms
439    }
440
441    /// Returns the thread's key/value metadata pairs.
442    pub fn metadata(&self) -> &[(String, String)] {
443        &self.metadata
444    }
445
446    /// Encodes the thread record as a SIM [`Expr`] map record.
447    pub fn to_expr(&self) -> Expr {
448        Expr::Map(vec![
449            field("kind", Expr::String(GATEWAY_THREAD_KIND.to_owned())),
450            field("id", Expr::String(self.id.clone())),
451            field(
452                "created-at-ms",
453                Expr::String(self.created_at_ms.to_string()),
454            ),
455            field("metadata", metadata_expr(&self.metadata)),
456        ])
457    }
458}
459
460/// A single message belonging to a [`GatewayThread`].
461#[derive(Clone, Debug, PartialEq, Eq)]
462pub struct GatewayThreadMessage {
463    id: String,
464    thread_id: String,
465    role: String,
466    content: String,
467    created_at_ms: u64,
468}
469
470impl GatewayThreadMessage {
471    /// Creates a thread message with the given id, thread, role, content, and time.
472    pub fn new(
473        id: impl Into<String>,
474        thread_id: impl Into<String>,
475        role: impl Into<String>,
476        content: impl Into<String>,
477        created_at_ms: u64,
478    ) -> Self {
479        Self {
480            id: id.into(),
481            thread_id: thread_id.into(),
482            role: role.into(),
483            content: content.into(),
484            created_at_ms,
485        }
486    }
487
488    /// Returns the message identifier.
489    pub fn id(&self) -> &str {
490        &self.id
491    }
492
493    /// Returns the id of the thread this message belongs to.
494    pub fn thread_id(&self) -> &str {
495        &self.thread_id
496    }
497
498    /// Returns the message role (e.g. `"user"`, `"assistant"`).
499    pub fn role(&self) -> &str {
500        &self.role
501    }
502
503    /// Returns the message text content.
504    pub fn content(&self) -> &str {
505        &self.content
506    }
507
508    /// Returns the creation timestamp in milliseconds since the Unix epoch.
509    pub fn created_at_ms(&self) -> u64 {
510        self.created_at_ms
511    }
512
513    /// Encodes the thread message as a SIM [`Expr`] map record.
514    pub fn to_expr(&self) -> Expr {
515        Expr::Map(vec![
516            field("kind", Expr::String(GATEWAY_THREAD_MESSAGE_KIND.to_owned())),
517            field("id", Expr::String(self.id.clone())),
518            field("thread-id", Expr::String(self.thread_id.clone())),
519            field("role", Expr::String(self.role.clone())),
520            field("content", Expr::String(self.content.clone())),
521            field(
522                "created-at-ms",
523                Expr::String(self.created_at_ms.to_string()),
524            ),
525        ])
526    }
527}
528
529/// Store for the gateway's durable account state: files, batches, threads,
530/// thread messages, and vector stores, all keyed by their string ids.
531pub trait GatewayStateStore {
532    /// Stores a file record together with its raw bytes.
533    fn put_file(&mut self, file: GatewayFile, bytes: Vec<u8>) -> Result<()>;
534    /// Returns the file record with the given id, if present.
535    fn file(&self, file_id: &str) -> Option<GatewayFile>;
536    /// Returns the raw bytes for the given file id, if present.
537    fn file_bytes(&self, file_id: &str) -> Option<Vec<u8>>;
538
539    /// Stores a batch record, replacing any existing entry with the same id.
540    fn put_batch(&mut self, batch: GatewayBatch) -> Result<()>;
541    /// Returns the batch with the given id, if present.
542    fn batch(&self, batch_id: &str) -> Option<GatewayBatch>;
543
544    /// Stores a thread record, replacing any existing entry with the same id.
545    fn put_thread(&mut self, thread: GatewayThread) -> Result<()>;
546    /// Returns the thread with the given id, if present.
547    fn thread(&self, thread_id: &str) -> Option<GatewayThread>;
548
549    /// Appends a message to its thread.
550    fn put_thread_message(&mut self, message: GatewayThreadMessage) -> Result<()>;
551    /// Returns the messages for the given thread id, in insertion order.
552    fn thread_messages(&self, thread_id: &str) -> Vec<GatewayThreadMessage>;
553
554    /// Stores a vector store, replacing any existing entry with the same id.
555    fn put_vector_store(&mut self, vector_store: GatewayVectorStore) -> Result<()>;
556    /// Returns the vector store with the given id, if present.
557    fn vector_store(&self, vector_store_id: &str) -> Option<GatewayVectorStore>;
558}
559
560fn metadata_expr(metadata: &[(String, String)]) -> Expr {
561    Expr::Map(
562        metadata
563            .iter()
564            .map(|(key, value)| (Expr::String(key.clone()), Expr::String(value.clone())))
565            .collect(),
566    )
567}
568
569use sim_value::build::entry as field;
570
571fn optional_string_field(name: &str, value: Option<&str>) -> (Expr, Expr) {
572    field(
573        name,
574        value
575            .map(|value| Expr::String(value.to_owned()))
576            .unwrap_or(Expr::Nil),
577    )
578}
579
580fn optional_u64_field(name: &str, value: Option<u64>) -> (Expr, Expr) {
581    field(
582        name,
583        value
584            .map(|value| Expr::String(value.to_string()))
585            .unwrap_or(Expr::Nil),
586    )
587}
sim_lib_openai_server/storage/objects.rs

sim_lib_openai_server/storage/
objects.rs
