Skip to main content

sim_lib_openai_server/storage/
table.rs

1use std::sync::{Arc, Mutex};
2
3use sim_kernel::{
4    ContentId, Cx, DefaultFactory, Error, Expr, Factory, NoopEvalPolicy, Object, ObjectCompat,
5    Result, Symbol, Table, Value,
6};
7
8use crate::{
9    objects::{
10        GatewayEvent, GatewayRequest, GatewayResponse, GatewayRun, content_id_expr, content_id_hex,
11    },
12    storage::{
13        GatewayBatch, GatewayFile, GatewayResponseObjectStore, GatewayStateStore, GatewayStore,
14        GatewayThread, GatewayThreadMessage, GatewayVectorStore, StoredGatewayResponse,
15    },
16};
17
18/// Gateway store that keeps every record kind in its own SIM [`Table`].
19///
20/// Each slot is a SIM table value guarded by a shared [`Cx`]; records are
21/// wrapped as opaque table values keyed by content id or string id. Implements
22/// [`GatewayStore`], [`GatewayResponseObjectStore`], and [`GatewayStateStore`].
23pub struct TableGatewayStore {
24    cx: Mutex<Cx>,
25    requests: Value,
26    runs: Value,
27    events: Value,
28    responses: Value,
29    response_objects: Value,
30    files: Value,
31    file_bytes: Value,
32    batches: Value,
33    threads: Value,
34    thread_messages: Value,
35    vector_stores: Value,
36}
37
38impl TableGatewayStore {
39    /// Creates a store with a fresh table for every record kind.
40    ///
41    /// Returns an error only if SIM table creation fails, which does not happen
42    /// for the in-memory tables used here.
43    pub fn new() -> Result<Self> {
44        let cx = Mutex::new(Cx::new(Arc::new(NoopEvalPolicy), Arc::new(DefaultFactory)));
45        Ok(Self {
46            requests: new_table()?,
47            runs: new_table()?,
48            events: new_table()?,
49            responses: new_table()?,
50            response_objects: new_table()?,
51            files: new_table()?,
52            file_bytes: new_table()?,
53            batches: new_table()?,
54            threads: new_table()?,
55            thread_messages: new_table()?,
56            vector_stores: new_table()?,
57            cx,
58        })
59    }
60
61    fn set<T>(&self, table: &Value, key: Symbol, value: T) -> Result<()>
62    where
63        T: TableStored,
64    {
65        let mut cx = self.cx()?;
66        let stored = cx.factory().opaque(Arc::new(StoredTableValue(value)))?;
67        table_impl(table)?.set(&mut cx, key, stored)
68    }
69
70    fn get<T>(&self, table: &Value, key: Symbol) -> Option<T>
71    where
72        T: TableStored,
73    {
74        let mut cx = self.cx().ok()?;
75        let value = table_impl(table).ok()?.get(&mut cx, key).ok()?;
76        value
77            .object()
78            .downcast_ref::<StoredTableValue<T>>()
79            .map(|stored| stored.0.clone())
80    }
81
82    fn cx(&self) -> Result<std::sync::MutexGuard<'_, Cx>> {
83        self.cx
84            .lock()
85            .map_err(|_| Error::PoisonedLock("openai gateway table store"))
86    }
87}
88
89impl Default for TableGatewayStore {
90    fn default() -> Self {
91        Self::new().expect("in-memory SIM table creation is infallible")
92    }
93}
94
95impl GatewayStore for TableGatewayStore {
96    fn put_request(&mut self, id: ContentId, request: GatewayRequest) -> Result<()> {
97        self.set(&self.requests, content_key(&id), request)
98    }
99
100    fn request(&self, id: &ContentId) -> Option<GatewayRequest> {
101        self.get(&self.requests, content_key(id))
102    }
103
104    fn put_run(&mut self, id: ContentId, run: GatewayRun) -> Result<()> {
105        self.set(&self.runs, content_key(&id), run)
106    }
107
108    fn run(&self, id: &ContentId) -> Option<GatewayRun> {
109        self.get(&self.runs, content_key(id))
110    }
111
112    fn put_event(&mut self, id: ContentId, event: GatewayEvent) -> Result<()> {
113        self.set(&self.events, content_key(&id), event)
114    }
115
116    fn event(&self, id: &ContentId) -> Option<GatewayEvent> {
117        self.get(&self.events, content_key(id))
118    }
119
120    fn put_response(&mut self, id: ContentId, response: GatewayResponse) -> Result<()> {
121        self.set(&self.responses, content_key(&id), response)
122    }
123
124    fn response(&self, id: &ContentId) -> Option<GatewayResponse> {
125        self.get(&self.responses, content_key(id))
126    }
127}
128
129impl GatewayResponseObjectStore for TableGatewayStore {
130    fn put_response_object(&mut self, response: StoredGatewayResponse) -> Result<()> {
131        self.set(
132            &self.responses,
133            content_key(response.content_id()),
134            response.response().clone(),
135        )?;
136        self.set(
137            &self.response_objects,
138            Symbol::new(response.response_id()),
139            response,
140        )
141    }
142
143    fn response_object(&self, response_id: &str) -> Option<StoredGatewayResponse> {
144        self.get(&self.response_objects, Symbol::new(response_id))
145    }
146}
147
148impl GatewayStateStore for TableGatewayStore {
149    fn put_file(&mut self, file: GatewayFile, bytes: Vec<u8>) -> Result<()> {
150        self.set(&self.file_bytes, Symbol::new(file.id()), bytes)?;
151        self.set(&self.files, Symbol::new(file.id()), file)
152    }
153
154    fn file(&self, file_id: &str) -> Option<GatewayFile> {
155        self.get(&self.files, Symbol::new(file_id))
156    }
157
158    fn file_bytes(&self, file_id: &str) -> Option<Vec<u8>> {
159        self.get(&self.file_bytes, Symbol::new(file_id))
160    }
161
162    fn put_batch(&mut self, batch: GatewayBatch) -> Result<()> {
163        self.set(&self.batches, Symbol::new(batch.id()), batch)
164    }
165
166    fn batch(&self, batch_id: &str) -> Option<GatewayBatch> {
167        self.get(&self.batches, Symbol::new(batch_id))
168    }
169
170    fn put_thread(&mut self, thread: GatewayThread) -> Result<()> {
171        self.set(&self.threads, Symbol::new(thread.id()), thread)
172    }
173
174    fn thread(&self, thread_id: &str) -> Option<GatewayThread> {
175        self.get(&self.threads, Symbol::new(thread_id))
176    }
177
178    fn put_thread_message(&mut self, message: GatewayThreadMessage) -> Result<()> {
179        let key = Symbol::new(message.thread_id());
180        let mut messages = self.thread_messages(message.thread_id());
181        messages.push(message);
182        self.set(&self.thread_messages, key, ThreadMessages(messages))
183    }
184
185    fn thread_messages(&self, thread_id: &str) -> Vec<GatewayThreadMessage> {
186        self.get::<ThreadMessages>(&self.thread_messages, Symbol::new(thread_id))
187            .map(|messages| messages.0)
188            .unwrap_or_default()
189    }
190
191    fn put_vector_store(&mut self, vector_store: GatewayVectorStore) -> Result<()> {
192        self.set(
193            &self.vector_stores,
194            Symbol::new(vector_store.id()),
195            vector_store,
196        )
197    }
198
199    fn vector_store(&self, vector_store_id: &str) -> Option<GatewayVectorStore> {
200        self.get(&self.vector_stores, Symbol::new(vector_store_id))
201    }
202}
203
204trait TableStored: Clone + Send + Sync + 'static {
205    fn label() -> &'static str;
206    fn to_expr(&self) -> Expr;
207}
208
209#[sim_citizen_derive::non_citizen(
210    reason = "OpenAI table storage wrapper; class-backed descriptor is the wrapped openai/* value",
211    kind = "marker",
212    descriptor = "openai/TableStored"
213)]
214#[derive(Clone)]
215struct StoredTableValue<T: TableStored>(T);
216
217impl<T: TableStored> Object for StoredTableValue<T> {
218    fn display(&self, _cx: &mut Cx) -> Result<String> {
219        Ok(format!("#<{}>", T::label()))
220    }
221
222    fn as_any(&self) -> &dyn std::any::Any {
223        self
224    }
225}
226
227impl<T: TableStored> ObjectCompat for StoredTableValue<T> {
228    fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
229        Ok(self.0.to_expr())
230    }
231}
232
233#[derive(Clone)]
234struct ThreadMessages(Vec<GatewayThreadMessage>);
235
236impl TableStored for GatewayRequest {
237    fn label() -> &'static str {
238        "openai-gateway-table-request"
239    }
240
241    fn to_expr(&self) -> Expr {
242        self.to_expr()
243    }
244}
245
246impl TableStored for GatewayRun {
247    fn label() -> &'static str {
248        "openai-gateway-table-run"
249    }
250
251    fn to_expr(&self) -> Expr {
252        self.to_expr()
253    }
254}
255
256impl TableStored for GatewayEvent {
257    fn label() -> &'static str {
258        "openai-gateway-table-event"
259    }
260
261    fn to_expr(&self) -> Expr {
262        self.to_expr()
263    }
264}
265
266impl TableStored for GatewayResponse {
267    fn label() -> &'static str {
268        "openai-gateway-table-response"
269    }
270
271    fn to_expr(&self) -> Expr {
272        self.to_expr()
273    }
274}
275
276impl TableStored for StoredGatewayResponse {
277    fn label() -> &'static str {
278        "openai-gateway-table-stored-response"
279    }
280
281    fn to_expr(&self) -> Expr {
282        Expr::Map(vec![
283            field("response-id", Expr::String(self.response_id().to_owned())),
284            field("content-id", content_id_expr(self.content_id())),
285            field("response", self.response().to_expr()),
286            optional_content_id_field("request-content-id", self.request_content_id.as_ref()),
287            optional_content_id_field("run-content-id", self.run_content_id.as_ref()),
288            field(
289                "event-content-ids",
290                Expr::Vector(
291                    self.event_content_ids
292                        .iter()
293                        .map(content_id_expr)
294                        .collect::<Vec<_>>(),
295                ),
296            ),
297            field(
298                "parent-response-id",
299                self.parent_response_id
300                    .as_ref()
301                    .map(|id| Expr::String(id.clone()))
302                    .unwrap_or(Expr::Nil),
303            ),
304            field(
305                "owner-key-id",
306                self.owner_key_id
307                    .as_ref()
308                    .map(|id| Expr::String(id.clone()))
309                    .unwrap_or(Expr::Nil),
310            ),
311        ])
312    }
313}
314
315impl TableStored for GatewayFile {
316    fn label() -> &'static str {
317        "openai-gateway-table-file"
318    }
319
320    fn to_expr(&self) -> Expr {
321        self.to_expr()
322    }
323}
324
325impl TableStored for Vec<u8> {
326    fn label() -> &'static str {
327        "openai-gateway-table-bytes"
328    }
329
330    fn to_expr(&self) -> Expr {
331        Expr::Bytes(self.clone())
332    }
333}
334
335impl TableStored for GatewayBatch {
336    fn label() -> &'static str {
337        "openai-gateway-table-batch"
338    }
339
340    fn to_expr(&self) -> Expr {
341        self.to_expr()
342    }
343}
344
345impl TableStored for GatewayThread {
346    fn label() -> &'static str {
347        "openai-gateway-table-thread"
348    }
349
350    fn to_expr(&self) -> Expr {
351        self.to_expr()
352    }
353}
354
355impl TableStored for GatewayThreadMessage {
356    fn label() -> &'static str {
357        "openai-gateway-table-thread-message"
358    }
359
360    fn to_expr(&self) -> Expr {
361        self.to_expr()
362    }
363}
364
365impl TableStored for GatewayVectorStore {
366    fn label() -> &'static str {
367        "openai-gateway-table-vector-store"
368    }
369
370    fn to_expr(&self) -> Expr {
371        self.to_expr()
372    }
373}
374
375impl TableStored for ThreadMessages {
376    fn label() -> &'static str {
377        "openai-gateway-table-thread-messages"
378    }
379
380    fn to_expr(&self) -> Expr {
381        Expr::Vector(
382            self.0
383                .iter()
384                .map(GatewayThreadMessage::to_expr)
385                .collect::<Vec<_>>(),
386        )
387    }
388}
389
390fn new_table() -> Result<Value> {
391    DefaultFactory.table(Vec::new())
392}
393
394fn table_impl(value: &Value) -> Result<&dyn Table> {
395    value
396        .object()
397        .as_table_impl()
398        .ok_or_else(|| Error::Eval("gateway table store slot is not a table".to_owned()))
399}
400
401fn content_key(id: &ContentId) -> Symbol {
402    Symbol::new(content_id_hex(id))
403}
404
405use sim_value::build::entry as field;
406
407fn optional_content_id_field(name: &str, value: Option<&ContentId>) -> (Expr, Expr) {
408    field(name, value.map(content_id_expr).unwrap_or(Expr::Nil))
409}