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;
330
331fn required_string<'a>(object: &'a Map<String, Value>, name: &'static str) -> RouteResult<&'a str> {
332    object
333        .get(name)
334        .and_then(Value::as_str)
335        .ok_or_else(|| OpenAiRouteError::missing_required(name))
336}
337
338fn embedding_inputs(object: &Map<String, Value>) -> RouteResult<Vec<String>> {
339    match object.get("input") {
340        Some(Value::String(input)) => Ok(vec![input.clone()]),
341        Some(Value::Array(inputs)) => inputs
342            .iter()
343            .map(|input| match input {
344                Value::String(text) => Ok(text.clone()),
345                _ => Err(OpenAiRouteError::bad_request(
346                    "embeddings input list must contain only strings",
347                    Some("input"),
348                    "invalid_input",
349                )),
350            })
351            .collect(),
352        Some(_) => Err(OpenAiRouteError::bad_request(
353            "embeddings input must be a string or list of strings",
354            Some("input"),
355            "invalid_input",
356        )),
357        None => Err(OpenAiRouteError::missing_required("input")),
358    }
359}
360
361fn embedding_model_from_plan(plan: &Expr, model: &str) -> RouteResult<EmbeddingModel> {
362    let (name, args) = plan_parts(plan).map_err(OpenAiRouteError::bad_model_from_error)?;
363    if name != "atom" {
364        return Err(OpenAiRouteError::bad_request(
365            "embeddings model must be a plan atom",
366            Some("model"),
367            "invalid_model",
368        ));
369    }
370    let [Expr::String(address)] = args else {
371        return Err(OpenAiRouteError::bad_model_from_error(Error::Eval(
372            "plan/atom expects one address".to_owned(),
373        )));
374    };
375    let descriptor =
376        resolve_atom_address(address).map_err(|err| OpenAiRouteError::model(err, model))?;
377    if !descriptor.address.starts_with("sim/embed/") {
378        return Err(model_not_found(model));
379    }
380    let dimension = match descriptor.address.as_str() {
381        TENSOR_F64_SMALL_EMBEDDING_MODEL => TENSOR_F64_SMALL_DIMENSION,
382        _ => return Err(model_not_found(model)),
383    };
384    Ok(EmbeddingModel {
385        id: descriptor.address,
386        dimension,
387        runner: descriptor.runner,
388    })
389}
390
391fn model_not_found(model: &str) -> OpenAiRouteError {
392    OpenAiRouteError::model(Error::Eval(format!("model_not_found: {model}")), model)
393}
394
395fn embedding_for_text(model: &str, input: &str, dimension: usize) -> Vec<f64> {
396    (0..dimension)
397        .map(|index| hash_to_unit(stable_embedding_hash(model, input, index)))
398        .collect()
399}
400
401fn stable_embedding_hash(model: &str, input: &str, dimension_index: usize) -> u64 {
402    let mut hash = FNV_OFFSET_BASIS;
403    mix_bytes(&mut hash, model.as_bytes());
404    mix_byte(&mut hash, 0xff);
405    mix_bytes(&mut hash, input.as_bytes());
406    mix_byte(&mut hash, 0xfe);
407    mix_bytes(&mut hash, &dimension_index.to_le_bytes());
408    hash
409}
410
411fn mix_bytes(hash: &mut u64, bytes: &[u8]) {
412    for byte in bytes {
413        mix_byte(hash, *byte);
414    }
415}
416fn mix_byte(hash: &mut u64, byte: u8) {
417    *hash ^= u64::from(byte);
418    *hash = hash.wrapping_mul(FNV_PRIME);
419}
420fn hash_to_unit(hash: u64) -> f64 {
421    let bucket = hash % (EMBEDDING_SCALE * 2 + 1);
422    (bucket as f64 / EMBEDDING_SCALE as f64) - 1.0
423}
424fn embedding_usage(inputs: &[String]) -> EmbeddingUsage {
425    let prompt_tokens = inputs
426        .iter()
427        .map(|input| input.split_whitespace().count() as u64)
428        .sum();
429    EmbeddingUsage {
430        prompt_tokens,
431        total_tokens: prompt_tokens,
432    }
433}
434
435fn embedding_response_body(
436    model: &str,
437    embeddings: &[Vec<f64>],
438    usage: &EmbeddingUsage,
439) -> RouteResult<Vec<u8>> {
440    serde_json::to_vec(&json!({
441        "object": "list",
442        "data": embeddings
443            .iter()
444            .enumerate()
445            .map(|(index, embedding)| {
446                json!({
447                    "object": "embedding",
448                    "embedding": embedding,
449                    "index": index,
450                })
451            })
452            .collect::<Vec<_>>(),
453        "model": model,
454        "usage": {
455            "prompt_tokens": usage.prompt_tokens,
456            "total_tokens": usage.total_tokens,
457        },
458    }))
459    .map_err(|err| {
460        OpenAiRouteError::internal_message(format!("failed to encode embeddings response: {err}"))
461    })
462}
463
464fn embedding_event_expr(model: &EmbeddingModel, input_count: usize) -> Expr {
465    Expr::Map(vec![
466        field("model", Expr::String(model.id.clone())),
467        field("runner", Expr::Symbol(model.runner.clone())),
468        field("input-count", Expr::String(input_count.to_string())),
469        field("dimension", Expr::String(model.dimension.to_string())),
470    ])
471}
472
473// Token counts are Expr::String by design here, CONSISTENT with the sibling
474// embeddings event fields (input-count, dimension). Do not 'fix' to numbers --
475// that would make this record internally inconsistent. See OVERLAP6.03e.
476fn usage_expr(usage: &EmbeddingUsage) -> Expr {
477    Expr::Map(vec![
478        field(
479            "prompt-tokens",
480            Expr::String(usage.prompt_tokens.to_string()),
481        ),
482        field("total-tokens", Expr::String(usage.total_tokens.to_string())),
483    ])
484}
485
486fn final_event_expr(model: &EmbeddingModel, input_count: usize, usage: &EmbeddingUsage) -> Expr {
487    Expr::Map(vec![
488        field("model", Expr::String(model.id.clone())),
489        field("object", Expr::String("list".to_owned())),
490        field("input-count", Expr::String(input_count.to_string())),
491        field("dimension", Expr::String(model.dimension.to_string())),
492        field("usage", usage_expr(usage)),
493    ])
494}
495
496use sim_value::build::entry as field;
sim_lib_openai_server/routes/embeddings.rs

sim_lib_openai_server/routes/
embeddings.rs
