Add fast raw memory search path

src/api/routes/memory.py

@@ -7,10 +7,12 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
 import logging
 import threading
 import time
-from typing import Any, Dict, List
+from collections import defaultdict, deque
+from typing import Any, Callable, Dict, List
 
 from fastapi import APIRouter, Depends, Request, UploadFile, File
 from fastapi.responses import JSONResponse
@@ -58,6 +60,8 @@
 logger = logging.getLogger("xmem.api.routes.memory")
 
 _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",
@@ -233,6 +237,53 @@ def _schedule_job(job: Dict[str, Any], handler) -> None:
         asyncio.create_task(run_job(get_default_job_store(), job["job_id"], handler))
 
 
+def _record_latency(mode: str, elapsed_ms: float) -> None:
+    with _latency_lock:
+        _latency_samples[mode].append(elapsed_ms)
+
+
+def _percentile(sorted_values: List[float], percentile: float) -> float:
+    if not sorted_values:
+        return 0.0
+    index = min(len(sorted_values) - 1, int(round((len(sorted_values) - 1) * percentile)))
+    return round(sorted_values[index], 2)
+
+
+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 snapshot.items():
+        values = sorted(samples)
+        stats[mode] = {
+            "count": len(values),
+            "p50_ms": _percentile(values, 0.50),
+            "p95_ms": _percentile(values, 0.95),
+            "p99_ms": _percentile(values, 0.99),
+        }
+    return stats
+
+
+async def _timed(
+    mode: str,
+    func: Callable[..., Any],
+    *args: Any,
+    threaded: bool = False,
+    **kwargs: Any,
+) -> tuple[Any, float]:
+    start = time.perf_counter()
+    if threaded:
+        result = await asyncio.to_thread(func, *args, **kwargs)
+    else:
+        result = func(*args, **kwargs)
+        if inspect.isawaitable(result):
+            result = await result
+    elapsed_ms = round((time.perf_counter() - start) * 1000, 2)
+    _record_latency(mode, elapsed_ms)
+    return result, elapsed_ms
+
+
 def _detect_chat_provider(*urls: str) -> str:
     for url in urls:
         lowered = (url or "").lower()
@@ -915,15 +966,69 @@ async def search_memory(req: SearchRequest, request: Request, user: dict = Depen
 
     try:
         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:
+            raw_tasks.append((
+                "profile",
+                _timed("profile", _search_profile, pipeline, user_id, threaded=True),
+            ))
+        if "temporal" in plan:
+            raw_tasks.append((
+                "temporal",
+                _timed("temporal", _search_temporal, pipeline, req.query, user_id, req.top_k, threaded=True),
+            ))
+        if "summary" in plan:
+            raw_tasks.append((
+                "summary",
+                _timed("summary", _search_summary, pipeline, req.query, user_id, req.top_k),
+            ))
+        if "snippet" in plan:
+            raw_tasks.append((
+                "snippet",
+                _timed("snippet", _search_snippet, pipeline, req.query, user_id, req.top_k),
+            ))
+        if "code" in plan:
+            raw_tasks.append((
+                "code",
+                _timed("code", _search_code, pipeline, req.query, user_id, req.top_k),
+            ))
 
-        if "profile" in req.domains:
-            all_results.extend(_search_profile(pipeline, user_id))
-        if "temporal" in req.domains:
-            all_results.extend(_search_temporal(pipeline, req.query, user_id, req.top_k))
-        if "summary" in req.domains:
-            all_results.extend(await _search_summary(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)
+
+        answer = None
+        answer_sources: List[SourceRecord] = []
+        confidence = 0.0
+        if req.answer:
+            answer_result, elapsed = await _timed("answer", pipeline.run, req.query, user_id, req.top_k)
+            latency_ms["answer"] = elapsed
+            answer = answer_result.answer
+            confidence = answer_result.confidence
+            answer_sources = [
+                SourceRecord(
+                    domain=s.domain, content=s.content,
+                    score=round(s.score, 3), metadata=s.metadata,
+                )
+                for s in answer_result.sources
+            ]
 
-        data = SearchResponse(results=all_results, total=len(all_results))
+        data = SearchResponse(
+            results=all_results,
+            total=len(all_results),
+            answer=answer,
+            answer_sources=answer_sources,
+            confidence=confidence,
+            latency_ms=latency_ms,
+            latency_stats=_latency_stats(),
+        )
         elapsed = round((time.perf_counter() - start) * 1000, 2)
         return _wrap(request, data, elapsed)
 
@@ -988,6 +1093,34 @@ async def _search_summary(pipeline: RetrievalPipeline, query: str, user_id: str,
         return []
 
 
+async def _search_snippet(pipeline: RetrievalPipeline, query: str, user_id: str, top_k: int) -> List[SourceRecord]:
+    try:
+        raw = await pipeline._search_snippet(query=query, user_id=user_id, top_k=top_k)
+        return [
+            SourceRecord(domain=r.domain, content=r.content, score=round(r.score, 3), metadata=r.metadata)
+            for r in raw
+        ]
+    except Exception as exc:
+        logger.warning("Snippet search error: %s", exc)
+        return []
+
+
+async def _search_code(pipeline: RetrievalPipeline, query: str, user_id: str, top_k: int) -> List[SourceRecord]:
+    try:
+        raw = await pipeline.vector_store.search_by_text(
+            query_text=query,
+            top_k=top_k,
+            filters={"user_id": user_id, "domain": "code"},
+        )
+        return [
+            SourceRecord(domain="code", content=r.content, score=round(r.score, 3), metadata={"id": r.id, **r.metadata})
+            for r in raw
+        ]
+    except Exception as exc:
+        logger.warning("Code search error: %s", exc)
+        return []
+
+
 # POST /v1/memory/scrape
 @scrape_router.post(
     "/scrape",

src/api/schemas.py

@@ -7,7 +7,6 @@
 
 from __future__ import annotations
 
-from datetime import datetime
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
@@ -159,24 +158,40 @@ class SearchRequest(BaseModel):
         ..., min_length=1, max_length=256, pattern=r"^[\w.\-@]+$",
     )
     domains: List[str] = Field(
-        default=["profile", "temporal", "summary"],
+        default=["profile", "temporal", "summary", "snippet", "code"],
         description="Which memory domains to search",
     )
     top_k: int = Field(default=10, ge=1, le=100)
+    answer: bool = Field(
+        default=False,
+        description="When true, also generate a synthesized answer after returning raw ranked hits.",
+    )
 
     @field_validator("domains")
     @classmethod
     def validate_domains(cls, v: List[str]) -> List[str]:
-        allowed = {"profile", "temporal", "summary"}
+        allowed = {"profile", "temporal", "summary", "snippet", "code"}
         for d in v:
             if d not in allowed:
                 raise ValueError(f"Invalid domain '{d}'. Allowed: {allowed}")
-        return v
+        return list(dict.fromkeys(v))
+
+
+class SearchLatencySummary(BaseModel):
+    count: int = 0
+    p50_ms: float = 0.0
+    p95_ms: float = 0.0
+    p99_ms: float = 0.0
 
 
 class SearchResponse(BaseModel):
     results: List[SourceRecord] = Field(default_factory=list)
     total: int = 0
+    answer: Optional[str] = None
+    answer_sources: List[SourceRecord] = Field(default_factory=list)
+    confidence: float = 0.0
+    latency_ms: Dict[str, float] = Field(default_factory=dict)
+    latency_stats: Dict[str, SearchLatencySummary] = Field(default_factory=dict)
 
 
 # ── Scrape (extract from shared chat links) ────────────────────────────────

src/pipelines/retrieval.py

@@ -22,6 +22,8 @@
 
 import asyncio
 import logging
+import threading
+import time
 from typing import Any, Callable, Dict, List, Optional
 
 from dotenv import load_dotenv
@@ -133,6 +135,12 @@ def __init__(
 
         self.embed_fn = embed_fn
         self._snippet_stores: Dict[str, BaseVectorStore] = {}
+        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
+        self._profile_catalog_cache_lock = threading.Lock()
+        self._raw_retrieval_plan_cache_lock = threading.Lock()
 
         logger.info("RetrievalPipeline initialized")
 
@@ -494,6 +502,17 @@ def _fetch_profile_catalog(self, user_id: str):
             catalog  — list of {topic, sub_topic} for the prompt
             raw_results — the full SearchResult list, cached for _search_profile
         """
+        now = time.monotonic()
+        with self._profile_catalog_cache_lock:
+            self._prune_profile_catalog_cache(now)
+
+            cached = self._profile_catalog_cache.get(user_id)
+            if cached and now - cached[0] < self._cache_ttl_seconds:
+                catalog, results = cached[1], cached[2]
+                self._profile_catalog_cache.pop(user_id)
+                self._profile_catalog_cache[user_id] = (now, catalog, results)
+                return catalog, results
+
         try:
             results = self.vector_store.search_by_metadata(
                 filters={"user_id": user_id, "domain": "profile"},
@@ -524,8 +543,35 @@ def _fetch_profile_catalog(self, user_id: str):
                     "sub_topic": "",
                 })
 
+        with self._profile_catalog_cache_lock:
+            self._prune_profile_catalog_cache(now)
+            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")
+        normalized = tuple(d for d in ordered_allowed if d in set(domains))
+        key = (normalized, answer)
+        with self._raw_retrieval_plan_cache_lock:
+            if key not in self._raw_retrieval_plan_cache:
+                self._raw_retrieval_plan_cache[key] = normalized
+            return self._raw_retrieval_plan_cache[key]
+
     def _format_catalog(self, catalog: List[Dict[str, str]]) -> str:
         """Format profile catalog for the system prompt."""
         if not catalog: