1use 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#[derive(Debug, Clone)]
36pub struct RecallResult {
37 pub items: Vec<RecallItem>,
40 pub search_type_used: Option<SearchType>,
42 pub auto_routed: bool,
44 pub search_response: Option<SearchResponse>,
46}
47
48#[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 let normalized: Vec<RecallScope> = match scope {
81 None => vec![RecallScope::Auto],
83 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 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 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 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 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 for src in &sources {
153 if auto_fallthrough && *src == RecallScope::Graph && !merged.is_empty() {
156 break;
157 }
158
159 let part: Vec<RecallItem> = match src {
160 RecallScope::Auto => continue, 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 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 #[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 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}