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)]
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}