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