Address raw search review feedback

Author: cnguyen14

src/api/routes/memory.py

@@ -53,6 +53,7 @@
 
 _ingest_semaphore = asyncio.Semaphore(5)
 _latency_samples: dict[str, deque[float]] = defaultdict(lambda: deque(maxlen=200))
+_latency_lock = threading.Lock()
 
 router = APIRouter(
     prefix="/v1/memory",
@@ -112,7 +113,8 @@ def _error(request: Request, detail: str, code: int, elapsed_ms: float = 0) -> J
 
 
 def _record_latency(mode: str, elapsed_ms: float) -> None:
-    _latency_samples[mode].append(elapsed_ms)
+    with _latency_lock:
+        _latency_samples[mode].append(elapsed_ms)
 
 
 def _percentile(sorted_values: List[float], percentile: float) -> float:
@@ -123,8 +125,11 @@ def _percentile(sorted_values: List[float], percentile: float) -> float:
 
 
 def _latency_stats() -> Dict[str, Dict[str, float]]:
+    with _latency_lock:
+        snapshot = {mode: list(samples) for mode, samples in _latency_samples.items()}
+
     stats: Dict[str, Dict[str, float]] = {}
-    for mode, samples in _latency_samples.items():
+    for mode, samples in snapshot.items():
         values = sorted(samples)
         stats[mode] = {
             "count": len(values),
@@ -135,11 +140,14 @@ def _latency_stats() -> Dict[str, Dict[str, float]]:
     return stats
 
 
-async def _timed(mode: str, func, *args, **kwargs):
+async def _timed(mode: str, func, *args, threaded: bool = False, **kwargs):
     start = time.perf_counter()
-    result = func(*args, **kwargs)
-    if hasattr(result, "__await__"):
-        result = await result
+    if threaded:
+        result = await asyncio.to_thread(func, *args, **kwargs)
+    else:
+        result = func(*args, **kwargs)
+        if hasattr(result, "__await__"):
+            result = await result
     elapsed_ms = round((time.perf_counter() - start) * 1000, 2)
     _record_latency(mode, elapsed_ms)
     return result, elapsed_ms
@@ -727,27 +735,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")