sim_lib_openai_server/routes/
embeddings.rs

1use serde_json::{Map, Value, json};
2use sim_kernel::{Error, Expr, Symbol};
3
4use crate::{
5    clock::{GatewayClock, SystemGatewayClock},
6    content_id::content_id_for_expr,
7    objects::{GatewayRequest, GatewayResponse},
8    plan::{check_plan, parse_plan, resolve_atom_address, shape::plan_parts},
9    server::GatewayRouteState,
10    storage::GatewayStore,
11};
12
13use super::{
14    errors::OpenAiRouteError,
15    execution_record::{EventInput, EventLog, RunPrologue, append_event, begin_run},
16};
17
18/// The embeddings engine shares the gateway execution-record substrate; its id
19/// generators and execution outcome are the shared types under route-local
20/// names (OVERLAP9.04).
21pub use super::execution_record::{
22    GatewayRunExecution as EmbeddingExecution, GatewayRunIdGenerators as EmbeddingIdGenerators,
23};
24
25/// Route path for the OpenAI-compatible `POST /v1/embeddings` endpoint.
26pub const EMBEDDINGS_PATH: &str = "/v1/embeddings";
27/// Model id of the built-in small fixed-dimension f64 embedding backend.
28pub const TENSOR_F64_SMALL_EMBEDDING_MODEL: &str = "sim/embed/tensor-f64-small";
29
30const TENSOR_F64_SMALL_DIMENSION: usize = 8;
31const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
32const FNV_PRIME: u64 = 0x100000001b3;
33const EMBEDDING_SCALE: u64 = 1_000_000;
34
35type RouteResult<T> = std::result::Result<T, OpenAiRouteError>;
36
37#[derive(Clone, Debug)]
38struct EmbeddingModel {
39    id: String,
40    dimension: usize,
41    runner: Symbol,
42}
43
44#[derive(Clone, Debug)]
45struct EmbeddingUsage {
46    prompt_tokens: u64,
47    total_tokens: u64,
48}
49
50/// Handles `POST /v1/embeddings`, executing the request against the gateway
51/// store and returning the OpenAI-shaped embeddings response.
52pub fn handle_embeddings(request: &GatewayRequest, state: &GatewayRouteState) -> GatewayResponse {
53    let mut clock = SystemGatewayClock;
54    let seed = clock.now_ms().unwrap_or(1);
55    let mut ids = EmbeddingIdGenerators::deterministic(seed);
56    match state.store().lock() {
57        Ok(mut store) => execute_embedding_request(&mut *store, &mut ids, &mut clock, request)
58            .response()
59            .clone(),
60        Err(err) => OpenAiRouteError::internal_message(format!("gateway store lock failed: {err}"))
61            .into_response(),
62    }
63}
64
65/// Executes an embedding request end to end, recording request, run, and
66/// events in the store and returning the [`EmbeddingExecution`] outcome.
67///
68/// Any failure is captured as an error response inside the returned execution
69/// rather than propagated.
70pub fn execute_embedding_request<S, C>(
71    store: &mut S,
72    ids: &mut EmbeddingIdGenerators,
73    clock: &mut C,
74    request: &GatewayRequest,
75) -> EmbeddingExecution
76where
77    S: GatewayStore,
78    C: GatewayClock,
79{
80    match try_execute_embedding_request(store, ids, clock, request) {
81        Ok(execution) => execution,
82        Err(error) => EmbeddingExecution::error(error),
83    }
84}
85
86fn try_execute_embedding_request<S, C>(
87    store: &mut S,
88    ids: &mut EmbeddingIdGenerators,
89    clock: &mut C,
90    request: &GatewayRequest,
91) -> RouteResult<EmbeddingExecution>
92where
93    S: GatewayStore,
94    C: GatewayClock,
95{
96    let object = request_object(request.body())?;
97    let model = required_string(&object, "model")?.to_owned();
98    let inputs = embedding_inputs(&object)?;
99    let record_execution = object.get("store").and_then(Value::as_bool).unwrap_or(true);
100
101    let plan = parse_plan(&model).map_err(OpenAiRouteError::bad_model_from_error)?;
102    check_plan(&plan).map_err(OpenAiRouteError::bad_model_from_error)?;
103    let embedding_model = embedding_model_from_plan(&plan, &model)?;
104
105    let RunPrologue {
106        recorded_request,
107        request_content_id,
108        run_id,
109        run_content_id,
110    } = begin_run(store, ids, clock, request, None, record_execution)?;
111
112    let embeddings = inputs
113        .iter()
114        .map(|input| embedding_for_text(&embedding_model.id, input, embedding_model.dimension))
115        .collect::<Vec<_>>();
116    let usage = embedding_usage(&inputs);
117
118    let mut event_log = EventLog::default();
119    append_event(
120        store,
121        ids,
122        clock,
123        &run_id,
124        EventInput::new(0, "request-start", recorded_request.to_expr()),
125        record_execution,
126        &mut event_log,
127    )?;
128    append_event(
129        store,
130        ids,
131        clock,
132        &run_id,
133        EventInput::new(1, "plan-start", plan.clone()),
134        record_execution,
135        &mut event_log,
136    )?;
137    append_event(
138        store,
139        ids,
140        clock,
141        &run_id,
142        EventInput::new(2, "model-start", Expr::String(embedding_model.id.clone())),
143        record_execution,
144        &mut event_log,
145    )?;
146    append_event(
147        store,
148        ids,
149        clock,
150        &run_id,
151        EventInput::new(
152            3,
153            "embedding",
154            embedding_event_expr(&embedding_model, inputs.len()),
155        ),
156        record_execution,
157        &mut event_log,
158    )?;
159    append_event(
160        store,
161        ids,
162        clock,
163        &run_id,
164        EventInput::new(4, "usage", usage_expr(&usage)),
165        record_execution,
166        &mut event_log,
167    )?;
168
169    let response_body = embedding_response_body(&embedding_model.id, &embeddings, &usage)?;
170    let response = GatewayResponse::json(200, response_body);
171    append_event(
172        store,
173        ids,
174        clock,
175        &run_id,
176        EventInput::new(
177            5,
178            "final",
179            final_event_expr(&embedding_model, inputs.len(), &usage),
180        ),
181        record_execution,
182        &mut event_log,
183    )?;
184    let response_content_id = if record_execution {
185        let id = content_id_for_expr(&response.to_expr()).map_err(OpenAiRouteError::internal)?;
186        store
187            .put_response(id.clone(), response.clone())
188            .map_err(OpenAiRouteError::internal)?;
189        Some(id)
190    } else {
191        None
192    };
193
194    Ok(EmbeddingExecution {
195        response,
196        request_content_id: Some(request_content_id),
197        run_content_id: Some(run_content_id),
198        event_content_ids: event_log.content_ids,
199        events: event_log.events,
200        response_id: None,
201        response_created_at_ms: None,
202        response_content_id,
203    })
204}
205
206use crate::routes::request_json::{request_object, required_string};
207
208fn embedding_inputs(object: &Map<String, Value>) -> RouteResult<Vec<String>> {
209    match object.get("input") {
210        Some(Value::String(input)) => Ok(vec![input.clone()]),
211        Some(Value::Array(inputs)) => inputs
212            .iter()
213            .map(|input| match input {
214                Value::String(text) => Ok(text.clone()),
215                _ => Err(OpenAiRouteError::bad_request(
216                    "embeddings input list must contain only strings",
217                    Some("input"),
218                    "invalid_input",
219                )),
220            })
221            .collect(),
222        Some(_) => Err(OpenAiRouteError::bad_request(
223            "embeddings input must be a string or list of strings",
224            Some("input"),
225            "invalid_input",
226        )),
227        None => Err(OpenAiRouteError::missing_required("input")),
228    }
229}
230
231fn embedding_model_from_plan(plan: &Expr, model: &str) -> RouteResult<EmbeddingModel> {
232    let (name, args) = plan_parts(plan).map_err(OpenAiRouteError::bad_model_from_error)?;
233    if name != "atom" {
234        return Err(OpenAiRouteError::bad_request(
235            "embeddings model must be a plan atom",
236            Some("model"),
237            "invalid_model",
238        ));
239    }
240    let [Expr::String(address)] = args else {
241        return Err(OpenAiRouteError::bad_model_from_error(Error::Eval(
242            "plan/atom expects one address".to_owned(),
243        )));
244    };
245    let descriptor =
246        resolve_atom_address(address).map_err(|err| OpenAiRouteError::model(err, model))?;
247    if !descriptor.address.starts_with("sim/embed/") {
248        return Err(model_not_found(model));
249    }
250    let dimension = match descriptor.address.as_str() {
251        TENSOR_F64_SMALL_EMBEDDING_MODEL => TENSOR_F64_SMALL_DIMENSION,
252        _ => return Err(model_not_found(model)),
253    };
254    Ok(EmbeddingModel {
255        id: descriptor.address,
256        dimension,
257        runner: descriptor.runner,
258    })
259}
260
261fn model_not_found(model: &str) -> OpenAiRouteError {
262    OpenAiRouteError::model(Error::Eval(format!("model_not_found: {model}")), model)
263}
264
265fn embedding_for_text(model: &str, input: &str, dimension: usize) -> Vec<f64> {
266    (0..dimension)
267        .map(|index| hash_to_unit(stable_embedding_hash(model, input, index)))
268        .collect()
269}
270
271fn stable_embedding_hash(model: &str, input: &str, dimension_index: usize) -> u64 {
272    let mut hash = FNV_OFFSET_BASIS;
273    mix_bytes(&mut hash, model.as_bytes());
274    mix_byte(&mut hash, 0xff);
275    mix_bytes(&mut hash, input.as_bytes());
276    mix_byte(&mut hash, 0xfe);
277    mix_bytes(&mut hash, &dimension_index.to_le_bytes());
278    hash
279}
280
281fn mix_bytes(hash: &mut u64, bytes: &[u8]) {
282    for byte in bytes {
283        mix_byte(hash, *byte);
284    }
285}
286fn mix_byte(hash: &mut u64, byte: u8) {
287    *hash ^= u64::from(byte);
288    *hash = hash.wrapping_mul(FNV_PRIME);
289}
290fn hash_to_unit(hash: u64) -> f64 {
291    let bucket = hash % (EMBEDDING_SCALE * 2 + 1);
292    (bucket as f64 / EMBEDDING_SCALE as f64) - 1.0
293}
294fn embedding_usage(inputs: &[String]) -> EmbeddingUsage {
295    let prompt_tokens = inputs
296        .iter()
297        .map(|input| input.split_whitespace().count() as u64)
298        .sum();
299    EmbeddingUsage {
300        prompt_tokens,
301        total_tokens: prompt_tokens,
302    }
303}
304
305fn embedding_response_body(
306    model: &str,
307    embeddings: &[Vec<f64>],
308    usage: &EmbeddingUsage,
309) -> RouteResult<Vec<u8>> {
310    serde_json::to_vec(&json!({
311        "object": "list",
312        "data": embeddings
313            .iter()
314            .enumerate()
315            .map(|(index, embedding)| {
316                json!({
317                    "object": "embedding",
318                    "embedding": embedding,
319                    "index": index,
320                })
321            })
322            .collect::<Vec<_>>(),
323        "model": model,
324        "usage": {
325            "prompt_tokens": usage.prompt_tokens,
326            "total_tokens": usage.total_tokens,
327        },
328    }))
329    .map_err(|err| {
330        OpenAiRouteError::internal_message(format!("failed to encode embeddings response: {err}"))
331    })
332}
333
334fn embedding_event_expr(model: &EmbeddingModel, input_count: usize) -> Expr {
335    Expr::Map(vec![
336        field("model", Expr::String(model.id.clone())),
337        field("runner", Expr::Symbol(model.runner.clone())),
338        field("input-count", Expr::String(input_count.to_string())),
339        field("dimension", Expr::String(model.dimension.to_string())),
340    ])
341}
342
343// Token counts are Expr::String by design here, CONSISTENT with the sibling
344// embeddings event fields (input-count, dimension). Do not 'fix' to numbers --
345// that would make this record internally inconsistent. See OVERLAP6.03e.
346fn usage_expr(usage: &EmbeddingUsage) -> Expr {
347    Expr::Map(vec![
348        field(
349            "prompt-tokens",
350            Expr::String(usage.prompt_tokens.to_string()),
351        ),
352        field("total-tokens", Expr::String(usage.total_tokens.to_string())),
353    ])
354}
355
356fn final_event_expr(model: &EmbeddingModel, input_count: usize, usage: &EmbeddingUsage) -> Expr {
357    Expr::Map(vec![
358        field("model", Expr::String(model.id.clone())),
359        field("object", Expr::String("list".to_owned())),
360        field("input-count", Expr::String(input_count.to_string())),
361        field("dimension", Expr::String(model.dimension.to_string())),
362        field("usage", usage_expr(usage)),
363    ])
364}
365
366use sim_value::build::entry as field;
sim_lib_openai_server/routes/embeddings.rs

sim_lib_openai_server/routes/
embeddings.rs
