Address raw search review feedback

Author: cnguyen14

src/api/routes/memory.py

@@ -961,27 +961,39 @@ async def search_memory(req: SearchRequest, request: Request, user: dict = Depen
         all_results: List[SourceRecord] = []
         latency_ms: Dict[str, float] = {}
         plan = pipeline.raw_retrieval_plan(req.domains, answer=req.answer)
+        raw_tasks = []
 
         if "profile" in plan:
-            results, elapsed = await _timed("profile", _search_profile, pipeline, user_id)
-            latency_ms["profile"] = elapsed
-            all_results.extend(results)
+            raw_tasks.append((
+                "profile",
+                _timed("profile", _search_profile, pipeline, user_id, threaded=True),
+            ))
         if "temporal" in plan:
-            results, elapsed = await _timed("temporal", _search_temporal, pipeline, req.query, user_id, req.top_k)
-            latency_ms["temporal"] = elapsed
-            all_results.extend(results)
+            raw_tasks.append((
+                "temporal",
+                _timed("temporal", _search_temporal, pipeline, req.query, user_id, req.top_k, threaded=True),
+            ))
         if "summary" in plan:
-            results, elapsed = await _timed("summary", _search_summary, pipeline, req.query, user_id, req.top_k)
-            latency_ms["summary"] = elapsed
-            all_results.extend(results)
+            raw_tasks.append((
+                "summary",
+                _timed("summary", _search_summary, pipeline, req.query, user_id, req.top_k),
+            ))
         if "snippet" in plan:
-            results, elapsed = await _timed("snippet", _search_snippet, pipeline, req.query, user_id, req.top_k)
-            latency_ms["snippet"] = elapsed
-            all_results.extend(results)
+            raw_tasks.append((
+                "snippet",
+                _timed("snippet", _search_snippet, pipeline, req.query, user_id, req.top_k),
+            ))
         if "code" in plan:
-            results, elapsed = await _timed("code", _search_code, pipeline, req.query, user_id, req.top_k)
-            latency_ms["code"] = elapsed
-            all_results.extend(results)
+            raw_tasks.append((
+                "code",
+                _timed("code", _search_code, pipeline, req.query, user_id, req.top_k),
+            ))
+
+        if raw_tasks:
+            raw_results = await asyncio.gather(*(task for _, task in raw_tasks))
+            for (domain, _), (results, elapsed) in zip(raw_tasks, raw_results):
+                latency_ms[domain] = elapsed
+                all_results.extend(results)
 
         all_results.sort(key=lambda record: record.score, reverse=True)
 

src/pipelines/retrieval.py

@@ -137,6 +137,7 @@ def __init__(
         self._profile_catalog_cache: Dict[str, tuple[float, List[Dict[str, str]], list]] = {}
         self._raw_retrieval_plan_cache: Dict[tuple[tuple[str, ...], bool], tuple[str, ...]] = {}
         self._cache_ttl_seconds = 60.0
+        self._profile_catalog_cache_max_users = 256
 
         logger.info("RetrievalPipeline initialized")
 
@@ -499,8 +500,11 @@ def _fetch_profile_catalog(self, user_id: str):
             raw_results — the full SearchResult list, cached for _search_profile
         """
         now = time.monotonic()
+        self._prune_profile_catalog_cache(now)
+
         cached = self._profile_catalog_cache.get(user_id)
         if cached and now - cached[0] < self._cache_ttl_seconds:
+            self._profile_catalog_cache[user_id] = (now, cached[1], cached[2])
             return cached[1], cached[2]
 
         try:
@@ -536,6 +540,20 @@ def _fetch_profile_catalog(self, user_id: str):
         self._profile_catalog_cache[user_id] = (now, catalog, results)
         return catalog, results
 
+    def _prune_profile_catalog_cache(self, now: float) -> None:
+        """Bound profile catalog cache by TTL and number of cached users."""
+        expired_user_ids = [
+            cached_user_id
+            for cached_user_id, (cached_at, _, _) in self._profile_catalog_cache.items()
+            if now - cached_at >= self._cache_ttl_seconds
+        ]
+        for cached_user_id in expired_user_ids:
+            self._profile_catalog_cache.pop(cached_user_id, None)
+
+        while len(self._profile_catalog_cache) >= self._profile_catalog_cache_max_users:
+            oldest_user_id = next(iter(self._profile_catalog_cache))
+            self._profile_catalog_cache.pop(oldest_user_id, None)
+
     def raw_retrieval_plan(self, domains: List[str], answer: bool = False) -> tuple[str, ...]:
         """Return a cached deterministic raw-search plan for the requested domains."""
         ordered_allowed = ("profile", "temporal", "summary", "snippet", "code")