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)]
213#[derive(Clone)]
214struct StoredTableValue<T: TableStored>(T);
215
216impl<T: TableStored> Object for StoredTableValue<T> {
217    fn display(&self, _cx: &mut Cx) -> Result<String> {
218        Ok(format!("#<{}>", T::label()))
219    }
220
221    fn as_any(&self) -> &dyn std::any::Any {
222        self
223    }
224}
225
226impl<T: TableStored> ObjectCompat for StoredTableValue<T> {
227    fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
228        Ok(self.0.to_expr())
229    }
230}
231
232#[derive(Clone)]
233struct ThreadMessages(Vec<GatewayThreadMessage>);
234
235impl TableStored for GatewayRequest {
236    fn label() -> &'static str {
237        "openai-gateway-table-request"
238    }
239
240    fn to_expr(&self) -> Expr {
241        self.to_expr()
242    }
243}
244
245impl TableStored for GatewayRun {
246    fn label() -> &'static str {
247        "openai-gateway-table-run"
248    }
249
250    fn to_expr(&self) -> Expr {
251        self.to_expr()
252    }
253}
254
255impl TableStored for GatewayEvent {
256    fn label() -> &'static str {
257        "openai-gateway-table-event"
258    }
259
260    fn to_expr(&self) -> Expr {
261        self.to_expr()
262    }
263}
264
265impl TableStored for GatewayResponse {
266    fn label() -> &'static str {
267        "openai-gateway-table-response"
268    }
269
270    fn to_expr(&self) -> Expr {
271        self.to_expr()
272    }
273}
274
275impl TableStored for StoredGatewayResponse {
276    fn label() -> &'static str {
277        "openai-gateway-table-stored-response"
278    }
279
280    fn to_expr(&self) -> Expr {
281        Expr::Map(vec![
282            field("response-id", Expr::String(self.response_id().to_owned())),
283            field("content-id", content_id_expr(self.content_id())),
284            field("response", self.response().to_expr()),
285            optional_content_id_field("request-content-id", self.request_content_id.as_ref()),
286            optional_content_id_field("run-content-id", self.run_content_id.as_ref()),
287            field(
288                "event-content-ids",
289                Expr::Vector(
290                    self.event_content_ids
291                        .iter()
292                        .map(content_id_expr)
293                        .collect::<Vec<_>>(),
294                ),
295            ),
296            field(
297                "parent-response-id",
298                self.parent_response_id
299                    .as_ref()
300                    .map(|id| Expr::String(id.clone()))
301                    .unwrap_or(Expr::Nil),
302            ),
303        ])
304    }
305}
306
307impl TableStored for GatewayFile {
308    fn label() -> &'static str {
309        "openai-gateway-table-file"
310    }
311
312    fn to_expr(&self) -> Expr {
313        self.to_expr()
314    }
315}
316
317impl TableStored for Vec<u8> {
318    fn label() -> &'static str {
319        "openai-gateway-table-bytes"
320    }
321
322    fn to_expr(&self) -> Expr {
323        Expr::Bytes(self.clone())
324    }
325}
326
327impl TableStored for GatewayBatch {
328    fn label() -> &'static str {
329        "openai-gateway-table-batch"
330    }
331
332    fn to_expr(&self) -> Expr {
333        self.to_expr()
334    }
335}
336
337impl TableStored for GatewayThread {
338    fn label() -> &'static str {
339        "openai-gateway-table-thread"
340    }
341
342    fn to_expr(&self) -> Expr {
343        self.to_expr()
344    }
345}
346
347impl TableStored for GatewayThreadMessage {
348    fn label() -> &'static str {
349        "openai-gateway-table-thread-message"
350    }
351
352    fn to_expr(&self) -> Expr {
353        self.to_expr()
354    }
355}
356
357impl TableStored for GatewayVectorStore {
358    fn label() -> &'static str {
359        "openai-gateway-table-vector-store"
360    }
361
362    fn to_expr(&self) -> Expr {
363        self.to_expr()
364    }
365}
366
367impl TableStored for ThreadMessages {
368    fn label() -> &'static str {
369        "openai-gateway-table-thread-messages"
370    }
371
372    fn to_expr(&self) -> Expr {
373        Expr::Vector(
374            self.0
375                .iter()
376                .map(GatewayThreadMessage::to_expr)
377                .collect::<Vec<_>>(),
378        )
379    }
380}
381
382fn new_table() -> Result<Value> {
383    DefaultFactory.table(Vec::new())
384}
385
386fn table_impl(value: &Value) -> Result<&dyn Table> {
387    value
388        .object()
389        .as_table_impl()
390        .ok_or_else(|| Error::Eval("gateway table store slot is not a table".to_owned()))
391}
392
393fn content_key(id: &ContentId) -> Symbol {
394    Symbol::new(content_id_hex(id))
395}
396
397use sim_value::build::entry as field;
398
399fn optional_content_id_field(name: &str, value: Option<&ContentId>) -> (Expr, Expr) {
400    field(name, value.map(content_id_expr).unwrap_or(Expr::Nil))
401}