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

sim_lib_openai_server/storage/
objects.rs
