sim_lib_openai_server/routes/
embeddings.rs

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

sim_lib_openai_server/routes/
embeddings.rs
