sim_lib_openai_server/routes/
embeddings.rs

1use serde_json::{Map, Value, json};
2use sim_kernel::{ContentId, Error, Expr, Symbol};
3
4use crate::{
5    clock::{GatewayClock, SystemGatewayClock},
6    content_id::{content_id_for_expr, request_content_id},
7    ids::GatewayIdGenerator,
8    objects::{GatewayEvent, GatewayRequest, GatewayResponse, GatewayRun},
9    plan::{check_plan, parse_plan, resolve_atom_address, shape::plan_parts},
10    runtime::redacted_gateway_request,
11    server::GatewayRouteState,
12    storage::GatewayStore,
13};
14
15use super::errors::OpenAiRouteError;
16
17/// Route path for the OpenAI-compatible `POST /v1/embeddings` endpoint.
18pub const EMBEDDINGS_PATH: &str = "/v1/embeddings";
19/// Model id of the built-in small fixed-dimension f64 embedding backend.
20pub const TENSOR_F64_SMALL_EMBEDDING_MODEL: &str = "sim/embed/tensor-f64-small";
21
22const TENSOR_F64_SMALL_DIMENSION: usize = 8;
23const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
24const FNV_PRIME: u64 = 0x100000001b3;
25const EMBEDDING_SCALE: u64 = 1_000_000;
26
27type RouteResult<T> = std::result::Result<T, OpenAiRouteError>;
28
29/// Holds the per-kind id generators used to mint request, run, and event ids
30/// for a single embedding execution.
31#[derive(Clone, Debug)]
32pub struct EmbeddingIdGenerators {
33    request: GatewayIdGenerator,
34    run: GatewayIdGenerator,
35    event: GatewayIdGenerator,
36}
37
38impl EmbeddingIdGenerators {
39    /// Builds id generators seeded deterministically from `start`.
40    pub fn deterministic(start: u64) -> Self {
41        Self {
42            request: GatewayIdGenerator::deterministic("gwreq", start),
43            run: GatewayIdGenerator::deterministic("gwrun", start),
44            event: GatewayIdGenerator::deterministic("gwevt", start),
45        }
46    }
47}
48
49/// Captures the outcome of an embedding request: the wire response plus the
50/// content-addressed ledger ids and events it produced.
51#[derive(Clone, Debug)]
52pub struct EmbeddingExecution {
53    response: GatewayResponse,
54    request_content_id: Option<ContentId>,
55    run_content_id: Option<ContentId>,
56    event_content_ids: Vec<ContentId>,
57    events: Vec<GatewayEvent>,
58    response_content_id: Option<ContentId>,
59}
60
61impl EmbeddingExecution {
62    /// Returns the wire response produced by the embedding execution.
63    pub fn response(&self) -> &GatewayResponse {
64        &self.response
65    }
66
67    /// Returns the content id of the stored request, if it was recorded.
68    pub fn request_content_id(&self) -> Option<&ContentId> {
69        self.request_content_id.as_ref()
70    }
71
72    /// Returns the content id of the stored run, if it was recorded.
73    pub fn run_content_id(&self) -> Option<&ContentId> {
74        self.run_content_id.as_ref()
75    }
76
77    /// Returns the content ids of the stored events, in sequence order.
78    pub fn event_content_ids(&self) -> &[ContentId] {
79        &self.event_content_ids
80    }
81
82    /// Returns the events emitted during the embedding execution.
83    pub fn events(&self) -> &[GatewayEvent] {
84        &self.events
85    }
86
87    /// Returns the content id of the stored response, if it was recorded.
88    pub fn response_content_id(&self) -> Option<&ContentId> {
89        self.response_content_id.as_ref()
90    }
91
92    fn error(error: OpenAiRouteError) -> Self {
93        Self {
94            response: error.into_response(),
95            request_content_id: None,
96            run_content_id: None,
97            event_content_ids: Vec::new(),
98            events: Vec::new(),
99            response_content_id: None,
100        }
101    }
102}
103
104#[derive(Clone, Debug)]
105struct EmbeddingModel {
106    id: String,
107    dimension: usize,
108    runner: Symbol,
109}
110
111#[derive(Clone, Debug)]
112struct EmbeddingUsage {
113    prompt_tokens: u64,
114    total_tokens: u64,
115}
116
117/// Handles `POST /v1/embeddings`, executing the request against the gateway
118/// store and returning the OpenAI-shaped embeddings response.
119pub fn handle_embeddings(request: &GatewayRequest, state: &GatewayRouteState) -> GatewayResponse {
120    let mut clock = SystemGatewayClock;
121    let seed = clock.now_ms().unwrap_or(1);
122    let mut ids = EmbeddingIdGenerators::deterministic(seed);
123    match state.store().lock() {
124        Ok(mut store) => execute_embedding_request(&mut *store, &mut ids, &mut clock, request)
125            .response()
126            .clone(),
127        Err(err) => OpenAiRouteError::internal_message(format!("gateway store lock failed: {err}"))
128            .into_response(),
129    }
130}
131
132/// Executes an embedding request end to end, recording request, run, and
133/// events in the store and returning the [`EmbeddingExecution`] outcome.
134///
135/// Any failure is captured as an error response inside the returned execution
136/// rather than propagated.
137pub fn execute_embedding_request<S, C>(
138    store: &mut S,
139    ids: &mut EmbeddingIdGenerators,
140    clock: &mut C,
141    request: &GatewayRequest,
142) -> EmbeddingExecution
143where
144    S: GatewayStore,
145    C: GatewayClock,
146{
147    match try_execute_embedding_request(store, ids, clock, request) {
148        Ok(execution) => execution,
149        Err(error) => EmbeddingExecution::error(error),
150    }
151}
152
153fn try_execute_embedding_request<S, C>(
154    store: &mut S,
155    ids: &mut EmbeddingIdGenerators,
156    clock: &mut C,
157    request: &GatewayRequest,
158) -> RouteResult<EmbeddingExecution>
159where
160    S: GatewayStore,
161    C: GatewayClock,
162{
163    let object = request_object(request.body())?;
164    let model = required_string(&object, "model")?.to_owned();
165    let inputs = embedding_inputs(&object)?;
166    let record_execution = object.get("store").and_then(Value::as_bool).unwrap_or(true);
167
168    let plan = parse_plan(&model).map_err(OpenAiRouteError::bad_model_from_error)?;
169    check_plan(&plan).map_err(OpenAiRouteError::bad_model_from_error)?;
170    let embedding_model = embedding_model_from_plan(&plan, &model)?;
171
172    let recorded_request = redacted_gateway_request(request).with_metadata(
173        ids.request.next_id().map_err(OpenAiRouteError::internal)?,
174        clock.now_ms().map_err(OpenAiRouteError::internal)?,
175    );
176    let request_content_id =
177        request_content_id(&recorded_request).map_err(OpenAiRouteError::internal)?;
178    if record_execution {
179        store
180            .put_request(request_content_id.clone(), recorded_request.clone())
181            .map_err(OpenAiRouteError::internal)?;
182    }
183
184    let run_id = ids.run.next_id().map_err(OpenAiRouteError::internal)?;
185    let run = GatewayRun::new(
186        run_id.clone(),
187        request_content_id.clone(),
188        clock.now_ms().map_err(OpenAiRouteError::internal)?,
189    );
190    let run_content_id = content_id_for_expr(&run.to_expr()).map_err(OpenAiRouteError::internal)?;
191    if record_execution {
192        store
193            .put_run(run_content_id.clone(), run)
194            .map_err(OpenAiRouteError::internal)?;
195    }
196
197    let embeddings = inputs
198        .iter()
199        .map(|input| embedding_for_text(&embedding_model.id, input, embedding_model.dimension))
200        .collect::<Vec<_>>();
201    let usage = embedding_usage(&inputs);
202
203    let mut event_log = EventLog::default();
204    append_event(
205        store,
206        ids,
207        clock,
208        &run_id,
209        EventInput(0, "request-start", recorded_request.to_expr()),
210        record_execution,
211        &mut event_log,
212    )?;
213    append_event(
214        store,
215        ids,
216        clock,
217        &run_id,
218        EventInput(1, "plan-start", plan.clone()),
219        record_execution,
220        &mut event_log,
221    )?;
222    append_event(
223        store,
224        ids,
225        clock,
226        &run_id,
227        EventInput(2, "model-start", Expr::String(embedding_model.id.clone())),
228        record_execution,
229        &mut event_log,
230    )?;
231    append_event(
232        store,
233        ids,
234        clock,
235        &run_id,
236        EventInput(
237            3,
238            "embedding",
239            embedding_event_expr(&embedding_model, inputs.len()),
240        ),
241        record_execution,
242        &mut event_log,
243    )?;
244    append_event(
245        store,
246        ids,
247        clock,
248        &run_id,
249        EventInput(4, "usage", usage_expr(&usage)),
250        record_execution,
251        &mut event_log,
252    )?;
253
254    let response_body = embedding_response_body(&embedding_model.id, &embeddings, &usage)?;
255    let response = GatewayResponse::json(200, response_body);
256    append_event(
257        store,
258        ids,
259        clock,
260        &run_id,
261        EventInput(
262            5,
263            "final",
264            final_event_expr(&embedding_model, inputs.len(), &usage),
265        ),
266        record_execution,
267        &mut event_log,
268    )?;
269    let response_content_id = if record_execution {
270        let id = content_id_for_expr(&response.to_expr()).map_err(OpenAiRouteError::internal)?;
271        store
272            .put_response(id.clone(), response.clone())
273            .map_err(OpenAiRouteError::internal)?;
274        Some(id)
275    } else {
276        None
277    };
278
279    Ok(EmbeddingExecution {
280        response,
281        request_content_id: Some(request_content_id),
282        run_content_id: Some(run_content_id),
283        event_content_ids: event_log.content_ids,
284        events: event_log.events,
285        response_content_id,
286    })
287}
288
289struct EventInput(u64, &'static str, Expr);
290
291#[derive(Default)]
292struct EventLog {
293    content_ids: Vec<ContentId>,
294    events: Vec<GatewayEvent>,
295}
296
297fn append_event<S, C>(
298    store: &mut S,
299    ids: &mut EmbeddingIdGenerators,
300    clock: &mut C,
301    run_id: &str,
302    input: EventInput,
303    store_event: bool,
304    event_log: &mut EventLog,
305) -> RouteResult<()>
306where
307    S: GatewayStore,
308    C: GatewayClock,
309{
310    let event = GatewayEvent::new(
311        ids.event.next_id().map_err(OpenAiRouteError::internal)?,
312        run_id,
313        input.0,
314        Symbol::new(input.1),
315        input.2,
316        clock.now_ms().map_err(OpenAiRouteError::internal)?,
317    );
318    let id = content_id_for_expr(&event.to_expr()).map_err(OpenAiRouteError::internal)?;
319    if store_event {
320        store
321            .put_event(id.clone(), event.clone())
322            .map_err(OpenAiRouteError::internal)?;
323    }
324    event_log.content_ids.push(id);
325    event_log.events.push(event);
326    Ok(())
327}
328
329use crate::routes::request_json::{request_object, required_string};
330
331fn embedding_inputs(object: &Map<String, Value>) -> RouteResult<Vec<String>> {
332    match object.get("input") {
333        Some(Value::String(input)) => Ok(vec![input.clone()]),
334        Some(Value::Array(inputs)) => inputs
335            .iter()
336            .map(|input| match input {
337                Value::String(text) => Ok(text.clone()),
338                _ => Err(OpenAiRouteError::bad_request(
339                    "embeddings input list must contain only strings",
340                    Some("input"),
341                    "invalid_input",
342                )),
343            })
344            .collect(),
345        Some(_) => Err(OpenAiRouteError::bad_request(
346            "embeddings input must be a string or list of strings",
347            Some("input"),
348            "invalid_input",
349        )),
350        None => Err(OpenAiRouteError::missing_required("input")),
351    }
352}
353
354fn embedding_model_from_plan(plan: &Expr, model: &str) -> RouteResult<EmbeddingModel> {
355    let (name, args) = plan_parts(plan).map_err(OpenAiRouteError::bad_model_from_error)?;
356    if name != "atom" {
357        return Err(OpenAiRouteError::bad_request(
358            "embeddings model must be a plan atom",
359            Some("model"),
360            "invalid_model",
361        ));
362    }
363    let [Expr::String(address)] = args else {
364        return Err(OpenAiRouteError::bad_model_from_error(Error::Eval(
365            "plan/atom expects one address".to_owned(),
366        )));
367    };
368    let descriptor =
369        resolve_atom_address(address).map_err(|err| OpenAiRouteError::model(err, model))?;
370    if !descriptor.address.starts_with("sim/embed/") {
371        return Err(model_not_found(model));
372    }
373    let dimension = match descriptor.address.as_str() {
374        TENSOR_F64_SMALL_EMBEDDING_MODEL => TENSOR_F64_SMALL_DIMENSION,
375        _ => return Err(model_not_found(model)),
376    };
377    Ok(EmbeddingModel {
378        id: descriptor.address,
379        dimension,
380        runner: descriptor.runner,
381    })
382}
383
384fn model_not_found(model: &str) -> OpenAiRouteError {
385    OpenAiRouteError::model(Error::Eval(format!("model_not_found: {model}")), model)
386}
387
388fn embedding_for_text(model: &str, input: &str, dimension: usize) -> Vec<f64> {
389    (0..dimension)
390        .map(|index| hash_to_unit(stable_embedding_hash(model, input, index)))
391        .collect()
392}
393
394fn stable_embedding_hash(model: &str, input: &str, dimension_index: usize) -> u64 {
395    let mut hash = FNV_OFFSET_BASIS;
396    mix_bytes(&mut hash, model.as_bytes());
397    mix_byte(&mut hash, 0xff);
398    mix_bytes(&mut hash, input.as_bytes());
399    mix_byte(&mut hash, 0xfe);
400    mix_bytes(&mut hash, &dimension_index.to_le_bytes());
401    hash
402}
403
404fn mix_bytes(hash: &mut u64, bytes: &[u8]) {
405    for byte in bytes {
406        mix_byte(hash, *byte);
407    }
408}
409fn mix_byte(hash: &mut u64, byte: u8) {
410    *hash ^= u64::from(byte);
411    *hash = hash.wrapping_mul(FNV_PRIME);
412}
413fn hash_to_unit(hash: u64) -> f64 {
414    let bucket = hash % (EMBEDDING_SCALE * 2 + 1);
415    (bucket as f64 / EMBEDDING_SCALE as f64) - 1.0
416}
417fn embedding_usage(inputs: &[String]) -> EmbeddingUsage {
418    let prompt_tokens = inputs
419        .iter()
420        .map(|input| input.split_whitespace().count() as u64)
421        .sum();
422    EmbeddingUsage {
423        prompt_tokens,
424        total_tokens: prompt_tokens,
425    }
426}
427
428fn embedding_response_body(
429    model: &str,
430    embeddings: &[Vec<f64>],
431    usage: &EmbeddingUsage,
432) -> RouteResult<Vec<u8>> {
433    serde_json::to_vec(&json!({
434        "object": "list",
435        "data": embeddings
436            .iter()
437            .enumerate()
438            .map(|(index, embedding)| {
439                json!({
440                    "object": "embedding",
441                    "embedding": embedding,
442                    "index": index,
443                })
444            })
445            .collect::<Vec<_>>(),
446        "model": model,
447        "usage": {
448            "prompt_tokens": usage.prompt_tokens,
449            "total_tokens": usage.total_tokens,
450        },
451    }))
452    .map_err(|err| {
453        OpenAiRouteError::internal_message(format!("failed to encode embeddings response: {err}"))
454    })
455}
456
457fn embedding_event_expr(model: &EmbeddingModel, input_count: usize) -> Expr {
458    Expr::Map(vec![
459        field("model", Expr::String(model.id.clone())),
460        field("runner", Expr::Symbol(model.runner.clone())),
461        field("input-count", Expr::String(input_count.to_string())),
462        field("dimension", Expr::String(model.dimension.to_string())),
463    ])
464}
465
466// Token counts are Expr::String by design here, CONSISTENT with the sibling
467// embeddings event fields (input-count, dimension). Do not 'fix' to numbers --
468// that would make this record internally inconsistent. See OVERLAP6.03e.
469fn usage_expr(usage: &EmbeddingUsage) -> Expr {
470    Expr::Map(vec![
471        field(
472            "prompt-tokens",
473            Expr::String(usage.prompt_tokens.to_string()),
474        ),
475        field("total-tokens", Expr::String(usage.total_tokens.to_string())),
476    ])
477}
478
479fn final_event_expr(model: &EmbeddingModel, input_count: usize, usage: &EmbeddingUsage) -> Expr {
480    Expr::Map(vec![
481        field("model", Expr::String(model.id.clone())),
482        field("object", Expr::String("list".to_owned())),
483        field("input-count", Expr::String(input_count.to_string())),
484        field("dimension", Expr::String(model.dimension.to_string())),
485        field("usage", usage_expr(usage)),
486    ])
487}
488
489use sim_value::build::entry as field;
sim_lib_openai_server/routes/embeddings.rs

sim_lib_openai_server/routes/
embeddings.rs
