engram_core/
search_handler.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use serde::Deserialize;
5use serde_json::{Value, json};
6
7use engram_router::Mode;
8
9use crate::error::CoreError;
10use crate::indexes::instrumentation::ReaderTracker;
11use crate::server::ServerState;
12use crate::timestamp::current_utc_timestamp;
13
14const VECTOR_WEIGHT: f64 = 0.7;
15const SPARSE_WEIGHT: f64 = 0.3;
16const MAX_QUERY_LENGTH: usize = 5_000;
17
18#[derive(Deserialize)]
19#[allow(dead_code)]
20struct SearchParams {
21    query: String,
22    limit: Option<usize>,
23    mode: Option<String>,
24    project: Option<String>,
25    tags: Option<Vec<String>>,
26}
27
28pub async fn handle(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
29    let parsed: SearchParams = serde_json::from_value(params)
30        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
31    if parsed.query.len() > MAX_QUERY_LENGTH {
32        return Err(CoreError::DispatchError(format!(
33            "query exceeds maximum length of {MAX_QUERY_LENGTH} bytes"
34        )));
35    }
36    let detected_mode = resolve_mode(&parsed);
37    let top_k = resolve_top_k(&parsed, state, detected_mode);
38    let query_embedding = embed_query(state, &parsed.query).await?;
39    let vector_results = search_vector_index(state, &query_embedding, top_k).await?;
40    let sparse_results = search_fts(state, &parsed.query, top_k).await?;
41    let merged = merge_results(&vector_results, &sparse_results);
42    let limited = limit_results(merged, top_k);
43    let memories = load_memories(state, &limited).await?;
44    let filtered = filter_by_tags(memories, &parsed.tags);
45    Ok(json!(filtered))
46}
47
48fn resolve_mode(params: &SearchParams) -> Mode {
49    match &params.mode {
50        Some(mode_string) => {
51            Mode::parse(mode_string).unwrap_or_else(|_| Mode::detect(&params.query))
52        }
53        None => Mode::detect(&params.query),
54    }
55}
56
57fn resolve_top_k(params: &SearchParams, state: &Arc<ServerState>, mode: Mode) -> usize {
58    if let Some(limit) = params.limit {
59        return limit;
60    }
61    let router = state.router.lock().unwrap();
62    let decision = router.decide(mode, 0.5);
63    decision.top_k
64}
65
66async fn embed_query(state: &Arc<ServerState>, query: &str) -> Result<Vec<f32>, CoreError> {
67    let query_owned = query.to_string();
68    let state_clone = Arc::clone(state);
69    tokio::task::spawn_blocking(move || {
70        let provider = state_clone.embedding_provider.as_ref();
71        let text_gen_ref = state_clone
72            .text_generator
73            .as_deref()
74            .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
75        state_clone
76            .embedder
77            .embed_query(&query_owned, provider, text_gen_ref)
78            .map_err(|error| {
79                CoreError::Api(engram_llm_client::ApiError::EmbeddingApiUnavailable(
80                    error.to_string(),
81                ))
82            })
83    })
84    .await
85    .map_err(|error| CoreError::SocketError(error.to_string()))?
86}
87
88async fn search_vector_index(
89    state: &Arc<ServerState>,
90    query_embedding: &[f32],
91    top_k: usize,
92) -> Result<Vec<(String, f32)>, CoreError> {
93    let embedding_owned = query_embedding.to_vec();
94    let state_clone = Arc::clone(state);
95    tokio::task::spawn_blocking(move || {
96        let _tracker = ReaderTracker::new();
97        let indexes = state_clone.indexes.read().unwrap();
98        let raw_results = indexes
99            .search(&embedding_owned, top_k)
100            .map_err(CoreError::Hnsw)?;
101        let resolved: Vec<(String, f32)> = raw_results
102            .into_iter()
103            .filter_map(|(node_id, score)| {
104                indexes
105                    .resolve_node_id(node_id)
106                    .map(|memory_id| (memory_id.to_string(), score))
107            })
108            .collect();
109        Ok(resolved)
110    })
111    .await
112    .map_err(|error| CoreError::SocketError(error.to_string()))?
113}
114
115async fn search_fts(
116    state: &Arc<ServerState>,
117    query: &str,
118    top_k: usize,
119) -> Result<Vec<(String, f64)>, CoreError> {
120    let query_owned = query.to_string();
121    let state_clone = Arc::clone(state);
122    tokio::task::spawn_blocking(move || {
123        let database = state_clone.database.lock().unwrap();
124        let fts_results = database.search_fts(&query_owned, top_k)?;
125        let scored: Vec<(String, f64)> = fts_results
126            .into_iter()
127            .map(|fts| (fts.memory.id, fts.rank.abs()))
128            .collect();
129        Ok::<_, CoreError>(scored)
130    })
131    .await
132    .map_err(|error| CoreError::SocketError(error.to_string()))?
133}
134
135fn merge_results(
136    vector_results: &[(String, f32)],
137    sparse_results: &[(String, f64)],
138) -> Vec<(String, f64)> {
139    let mut combined: HashMap<String, f64> = HashMap::new();
140    let vector_max = vector_results
141        .iter()
142        .map(|(_, score)| *score as f64)
143        .fold(0.0f64, f64::max)
144        .max(1.0);
145    for (memory_id, score) in vector_results {
146        let normalized = *score as f64 / vector_max;
147        combined
148            .entry(memory_id.clone())
149            .and_modify(|existing| *existing = existing.max(normalized * VECTOR_WEIGHT))
150            .or_insert(normalized * VECTOR_WEIGHT);
151    }
152    let sparse_max = sparse_results
153        .iter()
154        .map(|(_, rank)| *rank)
155        .fold(0.0f64, f64::max)
156        .max(1.0);
157    for (memory_id, rank) in sparse_results {
158        let normalized = *rank / sparse_max;
159        combined
160            .entry(memory_id.clone())
161            .and_modify(|existing| *existing += normalized * SPARSE_WEIGHT)
162            .or_insert(normalized * SPARSE_WEIGHT);
163    }
164    let mut results: Vec<(String, f64)> = combined.into_iter().collect();
165    results.sort_by(|a, b| b.1.total_cmp(&a.1));
166    results
167}
168
169fn limit_results(results: Vec<(String, f64)>, top_k: usize) -> Vec<(String, f64)> {
170    results.into_iter().take(top_k).collect()
171}
172
173fn filter_by_tags(memories: Vec<Value>, required_tags: &Option<Vec<String>>) -> Vec<Value> {
174    let required = match required_tags {
175        Some(tags) if !tags.is_empty() => tags,
176        _ => return memories,
177    };
178    let required_lower: Vec<String> = required.iter().map(|tag| tag.to_lowercase()).collect();
179    memories
180        .into_iter()
181        .filter(|memory| memory_has_all_tags(memory, &required_lower))
182        .collect()
183}
184
185fn memory_has_all_tags(memory: &Value, required_lower: &[String]) -> bool {
186    let raw = match memory.get("tags").and_then(Value::as_str) {
187        Some(value) if !value.is_empty() => value,
188        _ => return false,
189    };
190    let stored_tags: Vec<String> = match serde_json::from_str(raw) {
191        Ok(parsed) => parsed,
192        Err(_) => return false,
193    };
194    let stored_lower: Vec<String> = stored_tags.iter().map(|tag| tag.to_lowercase()).collect();
195    required_lower
196        .iter()
197        .all(|required_tag| stored_lower.contains(required_tag))
198}
199
200async fn load_memories(
201    state: &Arc<ServerState>,
202    scored_results: &[(String, f64)],
203) -> Result<Vec<Value>, CoreError> {
204    let memory_ids: Vec<String> = scored_results.iter().map(|(id, _)| id.clone()).collect();
205    let scores: HashMap<String, f64> = scored_results.iter().cloned().collect();
206    let state_clone = Arc::clone(state);
207    tokio::task::spawn_blocking(move || {
208        let database = state_clone.database.lock().unwrap();
209        let timestamp = current_utc_timestamp();
210        let mut results = Vec::new();
211        for memory_id in &memory_ids {
212            let memory = match database.get_memory(memory_id) {
213                Ok(m) => m,
214                Err(_) => continue,
215            };
216            let _ = database.track_search(memory_id, &timestamp);
217            let score = scores.get(memory_id.as_str()).copied().unwrap_or(0.0);
218            results.push(json!({
219                "id": memory.id,
220                "score": score,
221                "context": memory.context,
222                "action": memory.action,
223                "result": memory.result,
224                "memory_type": memory.memory_type,
225                "tags": memory.tags,
226            }));
227        }
228        Ok::<Vec<Value>, CoreError>(results)
229    })
230    .await
231    .map_err(|error| CoreError::SocketError(error.to_string()))?
232}
engram_core/search_handler.rs

engram_core/
search_handler.rs
