1use std::sync::{Arc, Mutex};
2
3use sim_citizen_derive::non_citizen;
4use sim_kernel::{Cx, Expr, Object, ObjectCompat, Ref, Result, Symbol};
5
6pub use crate::objects::{GatewayRequest, GatewayResponse, GatewayResponseValue};
7use crate::routes::{
8 admin::{
9 ADMIN_CACHE_STATS_PATH, ADMIN_CAPABILITY_REPORT_PATH, ADMIN_EVENTS_PATH,
10 ADMIN_MODEL_HEALTH_PATH, ADMIN_RUN_RETRIEVAL_PREFIX, ADMIN_RUN_RETRIEVAL_ROUTE,
11 ADMIN_RUNS_PATH, ADMIN_STORAGE_STATS_PATH, handle_admin_cache_stats,
12 handle_admin_capability_report, handle_admin_events, handle_admin_model_health,
13 handle_admin_run_get, handle_admin_runs, handle_admin_storage_stats,
14 },
15 audio::{
16 AUDIO_SPEECH_PATH, AUDIO_TRANSCRIPTIONS_PATH, handle_audio_speech,
17 handle_audio_transcriptions,
18 },
19 batches::{
20 BATCH_CANCEL_ROUTE, BATCH_RETRIEVAL_PREFIX, BATCH_RETRIEVAL_ROUTE, BATCHES_PATH,
21 handle_batch_cancel, handle_batch_retrieval, handle_batches,
22 },
23 chat_completions::{CHAT_COMPLETIONS_PATH, handle_chat_completions},
24 embeddings::{EMBEDDINGS_PATH, handle_embeddings},
25 files::{
26 FILE_RETRIEVAL_PREFIX, FILE_RETRIEVAL_ROUTE, FILES_PATH, handle_file_retrieval,
27 handle_files,
28 },
29 health::{HEALTH_PATH, handle_health},
30 images::{IMAGES_GENERATIONS_PATH, handle_image_generations},
31 models::{MODELS_PATH, handle_models},
32 replay::{
33 RESPONSE_EVENTS_ROUTE, RESPONSE_EVENTS_SUFFIX, RESPONSE_SIM_ROUTE, RESPONSE_SIM_SUFFIX,
34 SIM_FORK_PATH, SIM_REPLAY_PATH, handle_response_events, handle_response_sim,
35 handle_sim_fork, handle_sim_replay,
36 },
37 responses::{
38 RESPONSE_RETRIEVAL_PREFIX, RESPONSE_RETRIEVAL_ROUTE, RESPONSES_PATH,
39 handle_response_retrieval, handle_responses,
40 },
41 threads::{
42 THREAD_MESSAGES_ROUTE, THREAD_RETRIEVAL_PREFIX, THREAD_RETRIEVAL_ROUTE, THREADS_PATH,
43 handle_thread_get, handle_thread_post, handle_threads,
44 },
45 vector_stores::{
46 VECTOR_STORE_SEARCH_PREFIX, VECTOR_STORE_SEARCH_ROUTE, VECTOR_STORE_SEARCH_SUFFIX,
47 VECTOR_STORES_PATH, handle_vector_store_search, handle_vector_stores,
48 },
49};
50use crate::{
51 runtime::{
52 OpenAiFederation, OpenAiKeyTable, OpenAiPlanCache, OpenAiRunnerRegistry,
53 global_openai_key_table,
54 },
55 storage::MemoryGatewayStore,
56};
57
58type RouteHandler = fn(&GatewayRequest, &GatewayRouteState) -> GatewayResponse;
59
60#[derive(Clone, Copy)]
62pub struct GatewayRoute {
63 method: &'static str,
64 path: RoutePath,
65 handler: RouteHandler,
66}
67
68impl GatewayRoute {
69 pub fn new(method: &'static str, path: &'static str, handler: RouteHandler) -> Self {
71 Self {
72 method,
73 path: RoutePath::Exact(path),
74 handler,
75 }
76 }
77
78 pub fn prefix(
82 method: &'static str,
83 display_path: &'static str,
84 prefix: &'static str,
85 handler: RouteHandler,
86 ) -> Self {
87 Self {
88 method,
89 path: RoutePath::Prefix {
90 display_path,
91 prefix,
92 },
93 handler,
94 }
95 }
96
97 pub fn prefix_suffix(
101 method: &'static str,
102 display_path: &'static str,
103 prefix: &'static str,
104 suffix: &'static str,
105 handler: RouteHandler,
106 ) -> Self {
107 Self {
108 method,
109 path: RoutePath::PrefixSuffix {
110 display_path,
111 prefix,
112 suffix,
113 },
114 handler,
115 }
116 }
117
118 pub fn path(&self) -> &'static str {
120 self.path.display()
121 }
122
123 fn matches_path(&self, path: &str) -> bool {
124 self.path.matches(path)
125 }
126}
127
128#[derive(Clone, Copy)]
129enum RoutePath {
130 Exact(&'static str),
131 Prefix {
132 display_path: &'static str,
133 prefix: &'static str,
134 },
135 PrefixSuffix {
136 display_path: &'static str,
137 prefix: &'static str,
138 suffix: &'static str,
139 },
140}
141
142impl RoutePath {
143 fn display(self) -> &'static str {
144 match self {
145 Self::Exact(path) => path,
146 Self::Prefix { display_path, .. } => display_path,
147 Self::PrefixSuffix { display_path, .. } => display_path,
148 }
149 }
150
151 fn matches(self, path: &str) -> bool {
152 match self {
153 Self::Exact(route_path) => route_path == path,
154 Self::Prefix { prefix, .. } => path.starts_with(prefix),
155 Self::PrefixSuffix { prefix, suffix, .. } => {
156 path.starts_with(prefix) && path.ends_with(suffix)
157 }
158 }
159 }
160}
161
162#[derive(Clone)]
168pub struct GatewayRouteState {
169 store: Arc<Mutex<MemoryGatewayStore>>,
170 cache: Arc<Mutex<OpenAiPlanCache>>,
171 runners: OpenAiRunnerRegistry,
172 federation: OpenAiFederation,
173 keys: OpenAiKeyTable,
174}
175
176impl GatewayRouteState {
177 pub fn new(store: MemoryGatewayStore) -> Self {
179 Self {
180 store: Arc::new(Mutex::new(store)),
181 cache: Arc::new(Mutex::new(OpenAiPlanCache::new())),
182 runners: OpenAiRunnerRegistry::new(),
183 federation: OpenAiFederation::new(),
184 keys: global_openai_key_table().clone(),
185 }
186 }
187
188 pub fn with_cache(store: MemoryGatewayStore, cache: OpenAiPlanCache) -> Self {
190 Self {
191 store: Arc::new(Mutex::new(store)),
192 cache: Arc::new(Mutex::new(cache)),
193 runners: OpenAiRunnerRegistry::new(),
194 federation: OpenAiFederation::new(),
195 keys: global_openai_key_table().clone(),
196 }
197 }
198
199 pub fn with_runners(mut self, runners: OpenAiRunnerRegistry) -> Self {
201 self.runners = runners;
202 self
203 }
204
205 pub fn with_federation(mut self, federation: OpenAiFederation) -> Self {
207 self.federation = federation;
208 self
209 }
210
211 pub fn with_keys(mut self, keys: OpenAiKeyTable) -> Self {
213 self.keys = keys;
214 self
215 }
216
217 pub fn memory() -> Self {
219 Self::new(MemoryGatewayStore::new())
220 }
221
222 pub fn store(&self) -> &Arc<Mutex<MemoryGatewayStore>> {
224 &self.store
225 }
226
227 pub fn cache(&self) -> &Arc<Mutex<OpenAiPlanCache>> {
229 &self.cache
230 }
231
232 pub fn runners(&self) -> &OpenAiRunnerRegistry {
234 &self.runners
235 }
236
237 pub fn federation(&self) -> &OpenAiFederation {
239 &self.federation
240 }
241
242 pub fn keys(&self) -> &OpenAiKeyTable {
244 &self.keys
245 }
246}
247
248#[derive(Clone)]
250pub struct GatewayRoutes {
251 routes: Vec<GatewayRoute>,
252 state: GatewayRouteState,
253}
254
255impl GatewayRoutes {
256 pub fn new(routes: Vec<GatewayRoute>) -> Self {
258 Self::with_state(routes, GatewayRouteState::memory())
259 }
260
261 pub fn with_state(routes: Vec<GatewayRoute>, state: GatewayRouteState) -> Self {
263 Self { routes, state }
264 }
265
266 pub fn handle(&self, request: &GatewayRequest) -> GatewayResponse {
271 let mut path_exists = false;
272 for route in &self.routes {
273 if !route.matches_path(request.path()) {
274 continue;
275 }
276 path_exists = true;
277 if route.method == request.method() {
278 return (route.handler)(request, &self.state);
279 }
280 }
281 if path_exists {
282 GatewayResponse::text(405, "method not allowed")
283 } else {
284 GatewayResponse::text(404, "not found")
285 }
286 }
287
288 pub fn paths(&self) -> Vec<&'static str> {
290 self.routes.iter().map(GatewayRoute::path).collect()
291 }
292
293 pub fn state(&self) -> &GatewayRouteState {
295 &self.state
296 }
297}
298
299pub fn configure_routes() -> GatewayRoutes {
301 configure_routes_with_state(GatewayRouteState::memory())
302}
303
304pub fn configure_routes_with_state(state: GatewayRouteState) -> GatewayRoutes {
310 GatewayRoutes::with_state(
311 vec![
312 GatewayRoute::new("GET", HEALTH_PATH, handle_health),
313 GatewayRoute::new("GET", MODELS_PATH, handle_models),
314 GatewayRoute::new("POST", RESPONSES_PATH, handle_responses),
315 GatewayRoute::new("POST", SIM_REPLAY_PATH, handle_sim_replay),
316 GatewayRoute::new("POST", SIM_FORK_PATH, handle_sim_fork),
317 GatewayRoute::new("POST", CHAT_COMPLETIONS_PATH, handle_chat_completions),
318 GatewayRoute::new("POST", EMBEDDINGS_PATH, handle_embeddings),
319 GatewayRoute::new("POST", FILES_PATH, handle_files),
320 GatewayRoute::new("POST", BATCHES_PATH, handle_batches),
321 GatewayRoute::new(
322 "POST",
323 AUDIO_TRANSCRIPTIONS_PATH,
324 handle_audio_transcriptions,
325 ),
326 GatewayRoute::new("POST", AUDIO_SPEECH_PATH, handle_audio_speech),
327 GatewayRoute::new("POST", IMAGES_GENERATIONS_PATH, handle_image_generations),
328 GatewayRoute::new("POST", VECTOR_STORES_PATH, handle_vector_stores),
329 GatewayRoute::prefix_suffix(
330 "POST",
331 VECTOR_STORE_SEARCH_ROUTE,
332 VECTOR_STORE_SEARCH_PREFIX,
333 VECTOR_STORE_SEARCH_SUFFIX,
334 handle_vector_store_search,
335 ),
336 GatewayRoute::prefix(
337 "GET",
338 BATCH_RETRIEVAL_ROUTE,
339 BATCH_RETRIEVAL_PREFIX,
340 handle_batch_retrieval,
341 ),
342 GatewayRoute::prefix(
343 "POST",
344 BATCH_CANCEL_ROUTE,
345 BATCH_RETRIEVAL_PREFIX,
346 handle_batch_cancel,
347 ),
348 GatewayRoute::prefix(
349 "GET",
350 FILE_RETRIEVAL_ROUTE,
351 FILE_RETRIEVAL_PREFIX,
352 handle_file_retrieval,
353 ),
354 GatewayRoute::new("POST", THREADS_PATH, handle_threads),
355 GatewayRoute::prefix(
356 "POST",
357 THREAD_MESSAGES_ROUTE,
358 THREAD_RETRIEVAL_PREFIX,
359 handle_thread_post,
360 ),
361 GatewayRoute::prefix(
362 "GET",
363 THREAD_RETRIEVAL_ROUTE,
364 THREAD_RETRIEVAL_PREFIX,
365 handle_thread_get,
366 ),
367 GatewayRoute::prefix_suffix(
368 "GET",
369 RESPONSE_EVENTS_ROUTE,
370 RESPONSE_RETRIEVAL_PREFIX,
371 RESPONSE_EVENTS_SUFFIX,
372 handle_response_events,
373 ),
374 GatewayRoute::prefix_suffix(
375 "GET",
376 RESPONSE_SIM_ROUTE,
377 RESPONSE_RETRIEVAL_PREFIX,
378 RESPONSE_SIM_SUFFIX,
379 handle_response_sim,
380 ),
381 GatewayRoute::prefix(
382 "GET",
383 RESPONSE_RETRIEVAL_ROUTE,
384 RESPONSE_RETRIEVAL_PREFIX,
385 handle_response_retrieval,
386 ),
387 GatewayRoute::new("GET", ADMIN_RUNS_PATH, handle_admin_runs),
388 GatewayRoute::prefix(
389 "GET",
390 ADMIN_RUN_RETRIEVAL_ROUTE,
391 ADMIN_RUN_RETRIEVAL_PREFIX,
392 handle_admin_run_get,
393 ),
394 GatewayRoute::new("GET", ADMIN_EVENTS_PATH, handle_admin_events),
395 GatewayRoute::new("GET", ADMIN_STORAGE_STATS_PATH, handle_admin_storage_stats),
396 GatewayRoute::new("GET", ADMIN_MODEL_HEALTH_PATH, handle_admin_model_health),
397 GatewayRoute::new("GET", ADMIN_CACHE_STATS_PATH, handle_admin_cache_stats),
398 GatewayRoute::new(
399 "GET",
400 ADMIN_CAPABILITY_REPORT_PATH,
401 handle_admin_capability_report,
402 ),
403 ],
404 state,
405 )
406}
407
408#[derive(Clone)]
410#[non_citizen(
411 reason = "live OpenAI gateway route table handle; route requests use openai/GatewayRequest descriptor",
412 kind = "handle"
413)]
414pub struct GatewayRoutesValue {
415 routes: GatewayRoutes,
416}
417
418impl GatewayRoutesValue {
419 pub fn new(routes: GatewayRoutes) -> Self {
421 Self { routes }
422 }
423
424 pub fn routes(&self) -> &GatewayRoutes {
426 &self.routes
427 }
428}
429
430impl Object for GatewayRoutesValue {
431 fn display(&self, _cx: &mut Cx) -> Result<String> {
432 Ok("#<openai-gateway-routes>".to_owned())
433 }
434
435 fn as_any(&self) -> &dyn std::any::Any {
436 self
437 }
438}
439
440impl ObjectCompat for GatewayRoutesValue {
441 fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
442 Ok(Expr::Map(vec![(
443 Expr::Symbol(Symbol::new("routes")),
444 Expr::List(
445 self.routes
446 .paths()
447 .into_iter()
448 .map(|path| Expr::String(path.to_owned()))
449 .collect(),
450 ),
451 )]))
452 }
453}
454
455pub fn health_ref() -> Ref {
457 Ref::Symbol(Symbol::qualified("openai-gateway", "health"))
458}