Enhance OmniSearch for LLMs with limits and include_body flags

server/src/handlers/meta.rs

@@ -339,6 +339,8 @@ impl McpTool for OmniSearchHandler {
 
     async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
         let req: OmniSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let limit = req.limit.unwrap_or(5);
+        let include_body = req.include_body.unwrap_or(false);
         let matches = {
             let idx = state.get_search_index();
             idx.search(&req.query, req.namespace.as_deref())
@@ -348,11 +350,20 @@ impl McpTool for OmniSearchHandler {
 
         let kg_json = state.read_graph(|full| {
             let mut kg_entities = std::collections::HashMap::new();
+            let mut count = 0;
             for (id, doc_type, _, _, _) in &matches {
                 if doc_type == "entity"
                     && let Some(e) = full.entities.get(id)
                 {
-                    kg_entities.insert(id, e);
+                    if count >= limit { continue; }
+                    count += 1;
+                    if !include_body {
+                        let mut summary = e.clone();
+                        summary.observations = vec![];
+                        kg_entities.insert(id.clone(), summary);
+                    } else {
+                        kg_entities.insert(id.clone(), e.clone());
+                    }
                 }
             }
             serde_json::to_value(&kg_entities).unwrap_or_default()
@@ -380,6 +391,17 @@ impl McpTool for OmniSearchHandler {
             let filtered: Vec<_> = all_tasks
                 .iter()
                 .filter(|t| matched_tasks.contains(t.id.as_str()))
+                .take(limit)
+                .map(|t| {
+                    if !include_body {
+                        let mut summary = t.clone();
+                        summary.description = "".to_string();
+                        summary.acceptance_criteria = vec![];
+                        summary
+                    } else {
+                        t.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });
@@ -388,6 +410,16 @@ impl McpTool for OmniSearchHandler {
             let filtered: Vec<_> = all_snippets
                 .iter()
                 .filter(|s| matched_snippets.contains(s.name.as_str()))
+                .take(limit)
+                .map(|s| {
+                    if !include_body {
+                        let mut summary = s.clone();
+                        summary.code = "".to_string();
+                        summary
+                    } else {
+                        s.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });
@@ -396,6 +428,18 @@ impl McpTool for OmniSearchHandler {
             let filtered: Vec<_> = all_adrs
                 .iter()
                 .filter(|a| matched_adrs.contains(a.id.as_str()))
+                .take(limit)
+                .map(|a| {
+                    if !include_body {
+                        let mut summary = a.clone();
+                        summary.context = "".to_string();
+                        summary.decision = "".to_string();
+                        summary.consequence = "".to_string();
+                        summary
+                    } else {
+                        a.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });
@@ -409,6 +453,17 @@ impl McpTool for OmniSearchHandler {
                         && (contains_ignore_ascii_case(&d.description, &q)
                             || contains_ignore_ascii_case(&d.ideal_solution, &q))
                 })
+                .take(limit)
+                .map(|d| {
+                    if !include_body {
+                        let mut summary = d.clone();
+                        summary.description = "".to_string();
+                        summary.ideal_solution = "".to_string();
+                        summary
+                    } else {
+                        d.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });
@@ -420,6 +475,16 @@ impl McpTool for OmniSearchHandler {
                     req.namespace.as_ref().is_none_or(|ns| m.namespace == *ns)
                         && contains_ignore_ascii_case(&m.content, &q)
                 })
+                .take(limit)
+                .map(|m| {
+                    if !include_body {
+                        let mut summary = m.clone();
+                        summary.content = "".to_string();
+                        summary
+                    } else {
+                        m.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });