Skip to main content

engram_core/
store_handler.rs

1use std::sync::Arc;
2
3use serde::Deserialize;
4use serde_json::{Value, json};
5
6use engram_storage::Memory;
7
8use crate::error::CoreError;
9use crate::persistence::f32_vec_to_bytes;
10use crate::server::ServerState;
11use crate::tags_normalize::{TagsInput, normalize_tags};
12use crate::timestamp::current_utc_timestamp;
13
14const MAX_FIELD_LENGTH: usize = 10_000;
15
16#[derive(Deserialize)]
17struct StoreParams {
18    memory_type: String,
19    context: String,
20    action: String,
21    result: String,
22    tags: Option<TagsInput>,
23    project: Option<String>,
24}
25
26pub async fn handle(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
27    let mut parsed: StoreParams = serde_json::from_value(params)
28        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
29    validate_field_length("context", &parsed.context)?;
30    validate_field_length("action", &parsed.action)?;
31    validate_field_length("result", &parsed.result)?;
32    let normalized_tags = normalize_tags(parsed.tags.take());
33    let memory_id = uuid::Uuid::new_v4().to_string();
34    let timestamp = current_utc_timestamp();
35    let embedding = compute_embedding(state, &parsed).await?;
36    let memory = build_memory(&memory_id, &timestamp, &parsed, normalized_tags, &embedding);
37    persist_memory(state, &memory, &memory_id, &embedding).await?;
38    Ok(json!({ "id": memory_id, "indexed": true }))
39}
40
41async fn compute_embedding(
42    state: &Arc<ServerState>,
43    params: &StoreParams,
44) -> Result<engram_embeddings::ThreeFieldEmbedding, CoreError> {
45    let context = params.context.clone();
46    let action = params.action.clone();
47    let result = params.result.clone();
48    let state_clone = Arc::clone(state);
49    tokio::task::spawn_blocking(move || {
50        let provider = state_clone.embedding_provider.as_ref();
51        let text_gen_ref = state_clone
52            .text_generator
53            .as_deref()
54            .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
55        state_clone
56            .embedder
57            .embed_fields(&context, &action, &result, provider, text_gen_ref)
58            .map_err(|error| {
59                CoreError::Api(engram_llm_client::ApiError::EmbeddingApiUnavailable(
60                    error.to_string(),
61                ))
62            })
63    })
64    .await
65    .map_err(|error| CoreError::SocketError(error.to_string()))?
66}
67
68fn build_memory(
69    memory_id: &str,
70    timestamp: &str,
71    params: &StoreParams,
72    normalized_tags: Option<String>,
73    embedding: &engram_embeddings::ThreeFieldEmbedding,
74) -> Memory {
75    Memory {
76        id: memory_id.to_string(),
77        memory_type: params.memory_type.clone(),
78        context: params.context.clone(),
79        action: params.action.clone(),
80        result: params.result.clone(),
81        score: 0.0,
82        embedding_context: Some(f32_vec_to_bytes(&embedding.context)),
83        embedding_action: Some(f32_vec_to_bytes(&embedding.action)),
84        embedding_result: Some(f32_vec_to_bytes(&embedding.result)),
85        indexed: true,
86        tags: normalized_tags,
87        project: params.project.clone(),
88        parent_id: None,
89        source_ids: None,
90        insight_type: None,
91        created_at: timestamp.to_string(),
92        updated_at: timestamp.to_string(),
93        used_count: 0,
94        last_used_at: None,
95        superseded_by: None,
96    }
97}
98
99async fn persist_memory(
100    state: &Arc<ServerState>,
101    memory: &Memory,
102    memory_id: &str,
103    embedding: &engram_embeddings::ThreeFieldEmbedding,
104) -> Result<(), CoreError> {
105    let memory_owned = memory.clone();
106    let state_db = Arc::clone(state);
107    tokio::task::spawn_blocking(move || {
108        let database = state_db.database.lock().unwrap();
109        database.insert_memory(&memory_owned)?;
110        Ok::<(), CoreError>(())
111    })
112    .await
113    .map_err(|error| CoreError::SocketError(error.to_string()))??;
114
115    let hashed_id = hash_string_to_u64(memory_id);
116    let rng_value = deterministic_rng(hashed_id);
117    let memory_id_owned = memory_id.to_string();
118    let embedding_owned = engram_embeddings::ThreeFieldEmbedding {
119        context: embedding.context.clone(),
120        action: embedding.action.clone(),
121        result: embedding.result.clone(),
122    };
123    let state_idx = Arc::clone(state);
124    tokio::task::spawn_blocking(move || {
125        let mut indexes = state_idx.indexes.write().unwrap();
126        indexes.insert(hashed_id, &memory_id_owned, &embedding_owned, rng_value)?;
127        Ok::<(), CoreError>(())
128    })
129    .await
130    .map_err(|error| CoreError::SocketError(error.to_string()))??;
131
132    Ok(())
133}
134
135const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
136const FNV_PRIME: u64 = 0x0100_0000_01b3;
137const DETERMINISTIC_RNG_MULTIPLIER: u64 = 0x9e37_79b9_7f4a_7c15;
138
139fn hash_string_to_u64(value: &str) -> u64 {
140    let mut hash: u64 = FNV_OFFSET_BASIS;
141    for byte in value.as_bytes() {
142        hash ^= *byte as u64;
143        hash = hash.wrapping_mul(FNV_PRIME);
144    }
145    hash
146}
147
148fn deterministic_rng(id: u64) -> f64 {
149    let mixed = id.wrapping_mul(DETERMINISTIC_RNG_MULTIPLIER);
150    (mixed >> 11) as f64 / (1u64 << 53) as f64
151}
152
153fn validate_field_length(field_name: &str, value: &str) -> Result<(), CoreError> {
154    if value.len() > MAX_FIELD_LENGTH {
155        return Err(CoreError::DispatchError(format!(
156            "{field_name} exceeds maximum length of {MAX_FIELD_LENGTH} bytes"
157        )));
158    }
159    Ok(())
160}