sim_lib_openai_server/storage/
table.rs1use 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
18pub 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 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}