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

engram_core/
search_handler.rs
