Skip to main content

cognee_lib/api/
recall.rs

1//! Smart search with session routing -- `recall()`.
2//!
3//! Wraps the standard search pipeline with two additional capabilities:
4//! 1. Session-first routing: check session Q&A entries by keyword overlap
5//!    before falling through to graph search.
6//! 2. Auto query-type selection: use [`route_query()`] to pick the best
7//!    [`SearchType`] based on query text patterns.
8//!
9//! Equivalent to Python's `cognee.api.v1.recall.recall()` and
10//! `cognee.memory.entries.normalize_scope()`.
11//!
12//! As of LIB-08 the scope-routing primitives (`RecallScope`, `RecallSource`,
13//! `ScopeInput`, `RecallItem`, `normalize_scope`) and the four source helpers
14//! (`search_session`, `search_trace`, `fetch_graph_context`, `run_graph`) live
15//! in `cognee_search::recall_scope`. They are re-exported here so existing
16//! call sites (`api::mod.rs`, `lib.rs::prelude`, integration tests) compile
17//! unchanged.
18
19use cognee_search::observability::{
20    COGNEE_RECALL_SCOPE, COGNEE_RECALL_SOURCE, COGNEE_RESULT_COUNT, COGNEE_SEARCH_QUERY,
21    COGNEE_SESSION_ENTRY_COUNT, COGNEE_SESSION_ID,
22};
23use cognee_search::recall_scope::{fetch_graph_context, run_graph, search_session, search_trace};
24use cognee_search::{SearchOrchestrator, SearchResponse, SearchType};
25use cognee_session::{SessionManager, SessionStore};
26use tracing::{field, info};
27
28use super::error::ApiError;
29
30pub use cognee_search::recall_scope::{
31    RecallItem, RecallOptions, RecallScope, RecallSource, ScopeInput, normalize_scope,
32};
33
34/// Full recall result.
35#[derive(Debug, Clone)]
36pub struct RecallResult {
37    /// Source-tagged results -- order matches the iteration order of the
38    /// resolved `sources` list (Python `recall.py:503-513`).
39    pub items: Vec<RecallItem>,
40    /// The search type that was used (if graph search ran).
41    pub search_type_used: Option<SearchType>,
42    /// Whether auto-routing was applied (only relevant when graph ran).
43    pub auto_routed: bool,
44    /// The raw graph search response (if graph search ran).
45    pub search_response: Option<SearchResponse>,
46}
47
48/// Smart search with optional session routing, auto query-type selection, and
49/// configurable source fan-out. Byte-for-byte parity with Python
50/// `cognee.api.v1.recall.recall()` (`cognee/api/v1/recall/recall.py:317-531`).
51///
52/// # Behavior
53/// 1. Normalize `scope` via [`normalize_scope`] (`recall.py:373`).
54/// 2. Resolve `[Auto]` per `(session_id, datasets, query_type)` triple
55///    (`recall.py:374-386`):
56///    - `session_id` set + no datasets + no query_type -> `[Session, Graph]`
57///      with `auto_fallthrough=true` (graph short-circuited on session hit).
58///    - `session_id` set otherwise -> `[Session, Graph]`, both contribute.
59///    - No `session_id` -> `[Graph]`.
60/// 3. Iterate sources in caller-supplied order (`recall.py:503`); each source
61///    appends to a flat merged list of [`RecallItem`]s.
62/// 4. Telemetry attrs: `COGNEE_RECALL_SCOPE`, `COGNEE_RECALL_SOURCE`,
63///    `COGNEE_RESULT_COUNT`, `COGNEE_SESSION_ENTRY_COUNT` (`recall.py:515-522`).
64#[allow(clippy::too_many_arguments)]
65pub async fn recall(
66    query_text: &str,
67    query_type: Option<SearchType>,
68    datasets: Option<Vec<String>>,
69    top_k: usize,
70    auto_route: bool,
71    session_id: Option<&str>,
72    user_id: Option<&str>,
73    search_orchestrator: &SearchOrchestrator,
74    session_store: Option<&dyn SessionStore>,
75    session_manager: Option<&SessionManager>,
76    scope: Option<Vec<RecallScope>>,
77    options: Option<RecallOptions>,
78) -> Result<RecallResult, ApiError> {
79    // -- Resolve scope to a concrete source list (Python recall.py:373-386). --
80    let normalized: Vec<RecallScope> = match scope {
81        // None means caller did not supply a scope at all -> "auto".
82        None => vec![RecallScope::Auto],
83        // Caller-supplied. We do NOT re-validate via normalize_scope here
84        // because the type system already constrained values to RecallScope;
85        // unknown strings are rejected at the HTTP/CLI boundary before they
86        // reach this function.
87        Some(v) if v.is_empty() => vec![RecallScope::Auto],
88        Some(v) => v,
89    };
90
91    let auto_mode = normalized.as_slice() == [RecallScope::Auto];
92    let (sources, auto_fallthrough): (Vec<RecallScope>, bool) = if auto_mode {
93        // Python recall.py:374-386.
94        match (session_id, datasets.as_ref(), query_type) {
95            (Some(_), None, None) => (vec![RecallScope::Session, RecallScope::Graph], true),
96            (Some(_), _, _) => (vec![RecallScope::Session, RecallScope::Graph], false),
97            (None, _, _) => (vec![RecallScope::Graph], false),
98        }
99    } else {
100        (normalized, false)
101    };
102
103    // Comma-joined source names, mirroring Python's `span_scope` (`recall.py:388`).
104    let span_scope: String = sources
105        .iter()
106        .filter_map(|s| s.as_source().map(|src| src.as_str()))
107        .collect::<Vec<_>>()
108        .join(",");
109
110    // Truncate query preview to 500 chars for PII control, mirroring Python's
111    // `span.set_attribute(COGNEE_SEARCH_QUERY, query_text[:500])`.
112    let query_preview: &str = {
113        let mut end = query_text.len();
114        if query_text.chars().count() > 500 {
115            let mut idx = 0usize;
116            for (count, (byte_idx, _)) in query_text.char_indices().enumerate() {
117                if count == 500 {
118                    idx = byte_idx;
119                    break;
120                }
121            }
122            if idx > 0 {
123                end = idx;
124            }
125        }
126        &query_text[..end]
127    };
128
129    let span = tracing::info_span!(
130        "cognee.api.recall",
131        { COGNEE_SEARCH_QUERY } = query_preview,
132        { COGNEE_RECALL_SCOPE } = span_scope.as_str(),
133        { COGNEE_SESSION_ID } = session_id.unwrap_or(""),
134        "cognee.recall.top_k" = top_k,
135        { cognee_search::observability::COGNEE_SEARCH_TYPE } = field::Empty,
136        { COGNEE_RECALL_SOURCE } = field::Empty,
137        { COGNEE_RESULT_COUNT } = field::Empty,
138        { COGNEE_SESSION_ENTRY_COUNT } = field::Empty,
139    );
140    let _enter = span.enter();
141
142    // Track the captured graph response (if any) so the result struct can
143    // surface `search_type_used` / `auto_routed` / `search_response` to
144    // callers that still rely on those fields.
145    let mut merged: Vec<RecallItem> = Vec::new();
146    let mut graph_search_type: Option<SearchType> = None;
147    let mut graph_auto_routed = false;
148    let mut graph_response: Option<SearchResponse> = None;
149    let mut session_result_count: usize = 0;
150
151    // -- Iterate sources in caller-supplied order (Python recall.py:503-513). --
152    for src in &sources {
153        // Auto-mode short-circuit: a session hit skips the graph runner.
154        // (Python recall.py:508-509)
155        if auto_fallthrough && *src == RecallScope::Graph && !merged.is_empty() {
156            break;
157        }
158
159        let part: Vec<RecallItem> = match src {
160            RecallScope::Auto => continue, // sentinel -- already resolved above
161            RecallScope::Session => {
162                search_session(query_text, session_id, user_id, top_k, session_store)
163                    .await
164                    .map_err(|e| ApiError::Search(e.to_string()))?
165            }
166            RecallScope::Trace => {
167                search_trace(query_text, session_id, user_id, top_k, session_manager)
168                    .await
169                    .map_err(|e| ApiError::Search(e.to_string()))?
170            }
171            RecallScope::GraphContext => fetch_graph_context(session_id, user_id, session_manager)
172                .await
173                .map_err(|e| ApiError::Search(e.to_string()))?,
174            RecallScope::Graph => {
175                let (items, used_type, was_auto, response) = run_graph(
176                    query_text,
177                    query_type,
178                    datasets.clone(),
179                    top_k,
180                    auto_route,
181                    session_id,
182                    search_orchestrator,
183                    &span,
184                    options.as_ref(),
185                )
186                .await
187                .map_err(|e| ApiError::Search(e.to_string()))?;
188                graph_search_type = Some(used_type);
189                graph_auto_routed = was_auto;
190                graph_response = Some(response);
191                items
192            }
193        };
194
195        if *src == RecallScope::Session {
196            session_result_count = part.len();
197        }
198        merged.extend(part);
199    }
200
201    // -- Telemetry (Python recall.py:515-522). --
202    let source_label: &str = if sources.iter().filter(|s| s.as_source().is_some()).count() == 1 {
203        sources
204            .iter()
205            .find_map(|s| s.as_source())
206            .map(|s| s.as_str())
207            .unwrap_or("graph")
208    } else {
209        "multi"
210    };
211    span.record(COGNEE_RECALL_SOURCE, source_label);
212    span.record(COGNEE_RESULT_COUNT, merged.len());
213    if session_result_count > 0 {
214        span.record(COGNEE_SESSION_ENTRY_COUNT, session_result_count);
215    }
216
217    info!(
218        results = merged.len(),
219        sources = ?sources,
220        session_id = session_id.unwrap_or("-"),
221        "recall: completed"
222    );
223
224    // Mirrors Python `send_telemetry("cognee.recall", ...)` from
225    // cognee/api/v1/recall/recall.py:402.
226    #[cfg(feature = "telemetry")]
227    {
228        let search_type_label = graph_search_type
229            .or(query_type)
230            .map(|t| format!("{t:?}"))
231            .unwrap_or_default();
232        cognee_telemetry::send_telemetry(
233            "cognee.recall",
234            user_id.unwrap_or("sdk"),
235            Some(serde_json::json!({
236                "query_length": query_text.len(),
237                "scope": span_scope,
238                "auto_route": auto_route,
239                "top_k": top_k,
240                "search_type": search_type_label,
241                "session_id": session_id,
242                "datasets": datasets,
243                "dataset_ids": serde_json::Value::Null,
244            })),
245        );
246    }
247
248    Ok(RecallResult {
249        items: merged,
250        search_type_used: graph_search_type,
251        auto_routed: graph_auto_routed,
252        search_response: graph_response,
253    })
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn recall_options_fields_are_set_correctly() {
262        let opts = RecallOptions {
263            triplet_distance_penalty: Some(6.5),
264            node_name: Some(vec!["Alice".to_string()]),
265            ..Default::default()
266        };
267        assert_eq!(opts.triplet_distance_penalty, Some(6.5));
268        assert_eq!(
269            opts.node_name.as_deref(),
270            Some(["Alice".to_string()].as_slice())
271        );
272        // Un-set fields default to None.
273        assert!(opts.system_prompt.is_none());
274        assert!(opts.wide_search_top_k.is_none());
275        assert!(opts.feedback_influence.is_none());
276        assert!(opts.neighborhood_depth.is_none());
277        assert!(opts.neighborhood_seed_top_k.is_none());
278    }
279}