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, RouterDecision};
9use engram_storage::routing_log::RoutingLogEntry;
10
11use crate::error::CoreError;
12use crate::indexes::instrumentation::ReaderTracker;
13use crate::lock_helpers;
14use crate::rank_fusion::{SHADOW_K_SET, limit_results, merge_results, shadow_rewards_for_k_set};
15use crate::server::ServerState;
16use crate::timestamp::current_utc_timestamp;
17
18const MAX_QUERY_LENGTH: usize = 5_000;
19
20#[derive(Deserialize)]
21#[allow(dead_code)]
22struct SearchParams {
23    query: String,
24    limit: Option<usize>,
25    mode: Option<String>,
26    project: Option<String>,
27    tags: Option<Vec<String>>,
28}
29
30pub async fn handle(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
31    let parsed: SearchParams = serde_json::from_value(params)
32        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
33    if parsed.query.len() > MAX_QUERY_LENGTH {
34        return Err(CoreError::DispatchError(format!(
35            "query exceeds maximum length of {MAX_QUERY_LENGTH} bytes"
36        )));
37    }
38    let query_id = uuid::Uuid::new_v4().to_string();
39    let detected_mode = resolve_mode(&parsed);
40    let decision = resolve_decision(&parsed, state, detected_mode);
41    let top_k = decision.top_k;
42    match embed_query(state, &parsed.query).await {
43        Ok(query_embedding) => {
44            let vector_results = search_vector_index(state, &query_embedding, top_k).await?;
45            let filtered =
46                rank_and_load(state, &vector_results, &parsed, &decision, &query_id).await?;
47            Ok(search_response(filtered, false))
48        }
49        Err(CoreError::Api(ApiError::EmbeddingApiUnavailable(_))) => {
50            let filtered = rank_and_load(state, &[], &parsed, &decision, &query_id).await?;
51            Ok(search_response(filtered, true))
52        }
53        Err(other) => Err(other),
54    }
55}
56
57async fn rank_and_load(
58    state: &Arc<ServerState>,
59    vector_results: &[(String, f32)],
60    params: &SearchParams,
61    decision: &RouterDecision,
62    query_id: &str,
63) -> Result<Vec<Value>, CoreError> {
64    let top_k = decision.top_k;
65    let sparse_results = search_fts(state, &params.query, top_k).await?;
66    let merged = merge_results(vector_results, &sparse_results, &state.config.search);
67    log_routing(state, decision, query_id, &merged);
68    let limited = limit_results(merged, top_k);
69    let memories = load_memories(state, &limited, query_id).await?;
70    Ok(filter_by_tags(memories, &params.tags))
71}
72
73/// Side-effect-only instrumentation: records the served router decision plus the
74/// counterfactual shadow rewards for offline analysis. Best-effort — a logging
75/// failure never affects the served `{results, degraded}` response.
76fn log_routing(
77    state: &Arc<ServerState>,
78    decision: &RouterDecision,
79    query_id: &str,
80    merged: &[(String, f64)],
81) {
82    let shadow_rewards = shadow_rewards_for_k_set(merged, &SHADOW_K_SET);
83    let shadow_rewards_json = serde_json::to_string(
84        &shadow_rewards
85            .iter()
86            .map(|(k, reward)| json!({ "k": k, "reward": reward }))
87            .collect::<Vec<Value>>(),
88    )
89    .ok();
90    let created_at = current_utc_timestamp();
91    let entry = RoutingLogEntry {
92        query_id,
93        mode: decision.mode.as_str(),
94        search_strategy: decision.search_strategy.as_str(),
95        llm_selection: decision.llm_selection.as_str(),
96        contextualization: decision.contextualization.as_str(),
97        proactivity: decision.proactivity.as_str(),
98        top_k: decision.top_k,
99        shadow_rewards_json: shadow_rewards_json.as_deref(),
100        created_at: &created_at,
101    };
102    let _ = {
103        let database = lock_helpers::lock_db(state);
104        database.log_routing_decision(&entry)
105    };
106}
107
108fn search_response(results: Vec<Value>, degraded: bool) -> Value {
109    json!({ "results": results, "degraded": degraded })
110}
111
112fn resolve_mode(params: &SearchParams) -> Mode {
113    match &params.mode {
114        Some(mode_string) => {
115            Mode::parse(mode_string).unwrap_or_else(|_| Mode::detect(&params.query))
116        }
117        None => Mode::detect(&params.query),
118    }
119}
120
121/// Resolves the full router decision for this search. The router lock is dropped
122/// before returning (the owned `RouterDecision` is moved out). An explicit
123/// `params.limit` overrides only the served `top_k`; every other level is taken
124/// straight from `router.decide` so serving stays byte-identical to the prior
125/// `resolve_top_k` behaviour.
126fn resolve_decision(params: &SearchParams, state: &Arc<ServerState>, mode: Mode) -> RouterDecision {
127    let mut decision = {
128        let router = lock_helpers::lock_router(state);
129        router.decide(mode, 0.5)
130    };
131    if let Some(limit) = params.limit {
132        decision.top_k = limit;
133    }
134    decision
135}
136
137async fn embed_query(state: &Arc<ServerState>, query: &str) -> Result<Vec<f32>, CoreError> {
138    let query_owned = query.to_string();
139    let state_clone = Arc::clone(state);
140    tokio::task::spawn_blocking(move || {
141        let provider = state_clone.embedding_provider.as_ref();
142        let text_gen_ref = state_clone
143            .text_generator
144            .as_deref()
145            .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
146        state_clone
147            .embedder
148            .embed_query(&query_owned, provider, text_gen_ref)
149            .map_err(|error| {
150                CoreError::Api(engram_llm_client::ApiError::EmbeddingApiUnavailable(
151                    error.to_string(),
152                ))
153            })
154    })
155    .await
156    .map_err(|error| CoreError::SocketError(error.to_string()))?
157}
158
159async fn search_vector_index(
160    state: &Arc<ServerState>,
161    query_embedding: &[f32],
162    top_k: usize,
163) -> Result<Vec<(String, f32)>, CoreError> {
164    let embedding_owned = query_embedding.to_vec();
165    let state_clone = Arc::clone(state);
166    tokio::task::spawn_blocking(move || {
167        let _tracker = ReaderTracker::new();
168        let indexes = lock_helpers::read_indexes(&state_clone);
169        let raw_results = indexes
170            .search(&embedding_owned, top_k)
171            .map_err(CoreError::Hnsw)?;
172        let resolved: Vec<(String, f32)> = raw_results
173            .into_iter()
174            .filter_map(|(node_id, score)| {
175                indexes
176                    .resolve_node_id(node_id)
177                    .map(|memory_id| (memory_id.to_string(), score))
178            })
179            .collect();
180        Ok(resolved)
181    })
182    .await
183    .map_err(|error| CoreError::SocketError(error.to_string()))?
184}
185
186async fn search_fts(
187    state: &Arc<ServerState>,
188    query: &str,
189    top_k: usize,
190) -> Result<Vec<(String, f64)>, CoreError> {
191    let query_owned = query.to_string();
192    let state_clone = Arc::clone(state);
193    tokio::task::spawn_blocking(move || {
194        let database = lock_helpers::lock_db(&state_clone);
195        let fts_results = database.search_fts(&query_owned, top_k)?;
196        let scored: Vec<(String, f64)> = fts_results
197            .into_iter()
198            .map(|fts| (fts.memory.id, fts.rank.abs()))
199            .collect();
200        Ok::<_, CoreError>(scored)
201    })
202    .await
203    .map_err(|error| CoreError::SocketError(error.to_string()))?
204}
205
206fn filter_by_tags(memories: Vec<Value>, required_tags: &Option<Vec<String>>) -> Vec<Value> {
207    let required = match required_tags {
208        Some(tags) if !tags.is_empty() => tags,
209        _ => return memories,
210    };
211    let required_lower: Vec<String> = required.iter().map(|tag| tag.to_lowercase()).collect();
212    memories
213        .into_iter()
214        .filter(|memory| memory_has_all_tags(memory, &required_lower))
215        .collect()
216}
217
218fn memory_has_all_tags(memory: &Value, required_lower: &[String]) -> bool {
219    let raw = match memory.get("tags").and_then(Value::as_str) {
220        Some(value) if !value.is_empty() => value,
221        _ => return false,
222    };
223    let stored_tags: Vec<String> = match serde_json::from_str(raw) {
224        Ok(parsed) => parsed,
225        Err(_) => return false,
226    };
227    let stored_lower: Vec<String> = stored_tags.iter().map(|tag| tag.to_lowercase()).collect();
228    required_lower
229        .iter()
230        .all(|required_tag| stored_lower.contains(required_tag))
231}
232
233async fn load_memories(
234    state: &Arc<ServerState>,
235    scored_results: &[(String, f64)],
236    query_id: &str,
237) -> Result<Vec<Value>, CoreError> {
238    let memory_ids: Vec<String> = scored_results.iter().map(|(id, _)| id.clone()).collect();
239    let scores: HashMap<String, f64> = scored_results.iter().cloned().collect();
240    let query_id = query_id.to_string();
241    let state_clone = Arc::clone(state);
242    tokio::task::spawn_blocking(move || {
243        let database = lock_helpers::lock_db(&state_clone);
244        let timestamp = current_utc_timestamp();
245        let mut results = Vec::new();
246        for memory_id in &memory_ids {
247            let memory = match database.get_memory(memory_id) {
248                Ok(m) => m,
249                Err(_) => continue,
250            };
251            if memory.superseded_by.is_some() {
252                continue;
253            }
254            let _ = database.track_search(memory_id, &query_id, &timestamp);
255            let score = scores.get(memory_id.as_str()).copied().unwrap_or(0.0);
256            results.push(json!({
257                "id": memory.id,
258                "score": score,
259                "context": memory.context,
260                "action": memory.action,
261                "result": memory.result,
262                "memory_type": memory.memory_type,
263                "tags": memory.tags,
264            }));
265        }
266        Ok::<Vec<Value>, CoreError>(results)
267    })
268    .await
269    .map_err(|error| CoreError::SocketError(error.to_string()))?
270}
engram_core/search_handler.rs

engram_core/
search_handler.rs
