Add low-latency memory search path

Author: strongkeep-debug

src/api/app.py

@@ -33,6 +33,7 @@
 from src.api.routes.enterprise import router as enterprise_router
 from src.api.routes.health import router as health_router
 from src.api.routes.memory import router as memory_router
+from src.api.routes.memory import search_router as memory_search_router
 from src.api.routes.memory import scrape_router as memory_scrape_router
 from src.api.routes.memory_graph import router as memory_graph_router
 from src.api.routes.scanner import router as scanner_router
@@ -155,6 +156,7 @@ async def lifespan(app: FastAPI):
     # ── Routes ────────────────────────────────────────────────────────
     app.include_router(health_router)
     app.include_router(memory_scrape_router)
+    app.include_router(memory_search_router)
     app.include_router(memory_router)
     app.include_router(memory_graph_router)
     app.include_router(code_router)

src/api/routes/__init__.py

@@ -1,4 +1,5 @@
 from .health import router as health_router
 from .memory import router as memory_router
+from .memory import search_router as memory_search_router
 
-__all__ = ["health_router", "memory_router"]
+__all__ = ["health_router", "memory_router", "memory_search_router"]

src/api/routes/memory.py

@@ -63,6 +63,11 @@
     dependencies=[Depends(enforce_rate_limit)],
 )
 
+search_router = APIRouter(
+    tags=["memory"],
+    dependencies=[Depends(require_ready), Depends(enforce_rate_limit)],
+)
+
 
 # Helpers
 def _model_name(model: Any) -> str:
@@ -667,6 +672,7 @@ async def retrieve_memory(req: RetrieveRequest, request: Request, user: dict = D
             confidence=result.confidence,
         )
         elapsed = round((time.perf_counter() - start) * 1000, 2)
+        pipeline.record_latency("agentic", elapsed)
         return _wrap(request, data, elapsed)
 
     except Exception as exc:
@@ -676,10 +682,15 @@ async def retrieve_memory(req: RetrieveRequest, request: Request, user: dict = D
 
 
 # POST /v1/memory/search
+@search_router.post(
+    "/search",
+    response_model=APIResponse,
+    summary="Raw semantic search across memory domains with optional answer synthesis",
+)
 @router.post(
     "/search",
     response_model=APIResponse,
-    summary="Raw semantic search across memory domains (no LLM answer)",
+    summary="Raw semantic search across memory domains with optional answer synthesis",
 )
 async def search_memory(req: SearchRequest, request: Request, user: dict = Depends(require_api_key)):
     start = time.perf_counter()
@@ -689,17 +700,34 @@ async def search_memory(req: SearchRequest, request: Request, user: dict = Depen
     user_id = user.get("username") or user.get("name") or user["id"]
 
     try:
-        all_results: List[SourceRecord] = []
-
-        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))
+        all_results = await pipeline.search_raw(
+            query=req.query,
+            user_id=user_id,
+            domains=req.domains,
+            top_k=req.top_k,
+        )
+        answer = ""
+        if req.answer:
+            answer = await pipeline.answer_from_sources(req.query, all_results)
 
-        data = SearchResponse(results=all_results, total=len(all_results))
         elapsed = round((time.perf_counter() - start) * 1000, 2)
+        pipeline.record_latency("answer" if req.answer else "raw", elapsed)
+        data = SearchResponse(
+            results=[
+                SourceRecord(
+                    domain=s.domain,
+                    content=s.content,
+                    score=round(s.score, 3),
+                    metadata=s.metadata,
+                )
+                for s in all_results
+            ],
+            total=len(all_results),
+            answer=answer,
+            model=_model_name(pipeline.model) if req.answer else "",
+            confidence=min(1.0, len(all_results) * 0.2) if answer else 0.0,
+            latency=pipeline.get_latency_snapshot(),
+        )
         return _wrap(request, data, elapsed)
 
     except Exception as exc:

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,15 +158,19 @@ 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"],
         description="Which memory domains to search",
     )
     top_k: int = Field(default=10, ge=1, le=100)
+    answer: bool = Field(
+        default=False,
+        description="When true, synthesize an answer from the raw hits without agentic tool selection.",
+    )
 
     @field_validator("domains")
     @classmethod
     def validate_domains(cls, v: List[str]) -> List[str]:
-        allowed = {"profile", "temporal", "summary"}
+        allowed = {"profile", "temporal", "summary", "snippet"}
         for d in v:
             if d not in allowed:
                 raise ValueError(f"Invalid domain '{d}'. Allowed: {allowed}")
@@ -177,6 +180,10 @@ def validate_domains(cls, v: List[str]) -> List[str]:
 class SearchResponse(BaseModel):
     results: List[SourceRecord] = Field(default_factory=list)
     total: int = 0
+    answer: str = ""
+    model: str = ""
+    confidence: float = 0.0
+    latency: Dict[str, Dict[str, float | int]] = Field(default_factory=dict)
 
 
 # ── Scrape (extract from shared chat links) ────────────────────────────────

src/pipelines/retrieval.py

@@ -22,6 +22,7 @@
 
 import asyncio
 import logging
+import time
 from typing import Any, Callable, Dict, List, Optional
 
 from dotenv import load_dotenv
@@ -42,6 +43,10 @@
 logger = logging.getLogger("xmem.pipelines.retrieval")
 
 
+_CACHE_TTL_SECONDS = 60.0
+_LATENCY_SAMPLE_LIMIT = 200
+
+
 # ═══════════════════════════════════════════════════════════════════════════
 # Tool schemas — These are the "function signatures" exposed to the LLM
 # ═══════════════════════════════════════════════════════════════════════════
@@ -133,6 +138,9 @@ 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[Any]]] = {}
+        self._retrieval_plan_cache: Dict[tuple[str, str, int, str], tuple[float, AIMessage]] = {}
+        self._latency_samples: Dict[str, List[float]] = {}
 
         logger.info("RetrievalPipeline initialized")
 
@@ -155,7 +163,7 @@ async def run(
         logger.info("=" * 60)
 
         # ── Step 0: Fetch available profile catalog for this user ─────
-        profile_catalog, profile_records = self._fetch_profile_catalog(user_id)
+        profile_catalog, profile_records = self._get_profile_catalog(user_id)
         catalog_text = self._format_catalog(profile_catalog)
         logger.info("Available profiles: %s", catalog_text)
 
@@ -169,7 +177,11 @@ async def run(
             HumanMessage(content=query),
         ]
 
-        ai_response: AIMessage = await self.model_with_tools.ainvoke(messages)
+        plan_key = (user_id, query.strip(), top_k, catalog_text)
+        ai_response = self._get_cached_retrieval_plan(plan_key)
+        if ai_response is None:
+            ai_response = await self.model_with_tools.ainvoke(messages)
+            self._cache_retrieval_plan(plan_key, ai_response)
         logger.info("LLM response received (tool_calls=%d)", len(ai_response.tool_calls or []))
 
         # ── Step 2: Execute tool calls ────────────────────────────────
@@ -237,16 +249,7 @@ async def _process_tool_call(tc):
             answer = ai_response.content
             logger.info("LLM answered without tool calls")
 
-        if isinstance(answer, list):
-            parts = []
-            for c in answer:
-                if isinstance(c, dict) and "text" in c:
-                    parts.append(c["text"])
-                elif isinstance(c, str):
-                    parts.append(c)
-                else:
-                    parts.append(str(c))
-            answer = "\n".join(parts)
+        answer = self._coerce_answer(answer)
 
         confidence = min(1.0, len(sources) * 0.2) if sources else 0.1
 
@@ -263,6 +266,52 @@ async def _process_tool_call(tc):
             confidence=confidence,
         )
 
+    async def search_raw(
+        self,
+        query: str,
+        user_id: str,
+        domains: List[str],
+        top_k: int = 10,
+    ) -> List[SourceRecord]:
+        """Return ranked memory hits without asking the LLM for a retrieval plan."""
+
+        domain_set = set(domains)
+        results: List[SourceRecord] = []
+
+        if "profile" in domain_set:
+            results.extend(await self._search_profile_raw(query, user_id, top_k))
+        if "temporal" in domain_set:
+            results.extend(await self._search_temporal(query, user_id, top_k))
+        if "summary" in domain_set:
+            results.extend(await self._search_summary(query, user_id, top_k))
+        if "snippet" in domain_set:
+            results.extend(await self._search_snippet(query, user_id, top_k))
+
+        return sorted(results, key=lambda record: record.score, reverse=True)
+
+    async def answer_from_sources(self, query: str, sources: List[SourceRecord]) -> str:
+        """Generate an answer from already-fetched sources without tool selection."""
+
+        context_text = self._format_tool_results(sources)
+        answer_prompt = ANSWER_PROMPT.format(context=context_text, query=query)
+        final_response = await self.model.ainvoke([HumanMessage(content=answer_prompt)])
+        return self._coerce_answer(final_response.content)
+
+    def record_latency(self, mode: str, elapsed_ms: float) -> None:
+        """Track bounded latency samples for raw, answer, and agentic modes."""
+
+        samples = self._latency_samples.setdefault(mode, [])
+        samples.append(float(elapsed_ms))
+        if len(samples) > _LATENCY_SAMPLE_LIMIT:
+            del samples[0 : len(samples) - _LATENCY_SAMPLE_LIMIT]
+
+    def get_latency_snapshot(self) -> Dict[str, Dict[str, float | int]]:
+        return {
+            mode: self._percentiles(samples)
+            for mode, samples in self._latency_samples.items()
+            if samples
+        }
+
     # ------------------------------------------------------------------
     # Tool execution
     # ------------------------------------------------------------------
@@ -349,6 +398,35 @@ def _search_profile(
 
     # -- Temporal: Neo4j semantic search ───────────────────────────────
 
+    async def _search_profile_raw(
+        self,
+        query: str,
+        user_id: str,
+        top_k: int = 10,
+    ) -> List[SourceRecord]:
+        """Semantic profile search for the low-latency raw search endpoint."""
+
+        try:
+            results = await self.vector_store.search_by_text(
+                query_text=query,
+                top_k=top_k,
+                filters={"user_id": user_id, "domain": "profile"},
+            )
+        except Exception as exc:
+            logger.warning("Profile raw search failed, using cached catalog: %s", exc)
+            _, results = self._get_profile_catalog(user_id)
+
+        records = []
+        for r in results[:top_k]:
+            records.append(SourceRecord(
+                domain="profile",
+                content=r.content,
+                score=r.score,
+                metadata={"id": r.id, **r.metadata},
+            ))
+        logger.info("Profile raw [%s]: %d results", query, len(records))
+        return records
+
     async def _search_temporal(
         self,
         query: str,
@@ -486,6 +564,20 @@ async def _search_snippet(
     # Profile catalog (tells the LLM what profile keys exist)
     # ------------------------------------------------------------------
 
+    def _get_profile_catalog(self, user_id: str):
+        cached = self._profile_catalog_cache.get(user_id)
+        now = time.monotonic()
+        if cached and cached[0] > now:
+            return cached[1], cached[2]
+
+        catalog, results = self._fetch_profile_catalog(user_id)
+        self._profile_catalog_cache[user_id] = (
+            now + _CACHE_TTL_SECONDS,
+            catalog,
+            results,
+        )
+        return catalog, results
+
     def _fetch_profile_catalog(self, user_id: str):
         """Fetch all profile entries for a user.
 
@@ -538,6 +630,58 @@ def _format_catalog(self, catalog: List[Dict[str, str]]) -> str:
             lines.append(f"  - {t} / {st}")
         return "\n".join(lines)
 
+    def _get_cached_retrieval_plan(
+        self,
+        key: tuple[str, str, int, str],
+    ) -> AIMessage | None:
+        cached = self._retrieval_plan_cache.get(key)
+        if not cached:
+            return None
+        expires_at, response = cached
+        if expires_at <= time.monotonic():
+            self._retrieval_plan_cache.pop(key, None)
+            return None
+        return response
+
+    def _cache_retrieval_plan(
+        self,
+        key: tuple[str, str, int, str],
+        response: AIMessage,
+    ) -> None:
+        self._retrieval_plan_cache[key] = (
+            time.monotonic() + _CACHE_TTL_SECONDS,
+            response,
+        )
+
+    def _coerce_answer(self, answer: Any) -> str:
+        if isinstance(answer, list):
+            parts = []
+            for c in answer:
+                if isinstance(c, dict) and "text" in c:
+                    parts.append(c["text"])
+                elif isinstance(c, str):
+                    parts.append(c)
+                else:
+                    parts.append(str(c))
+            return "\n".join(parts)
+        return str(answer)
+
+    def _percentiles(self, samples: List[float]) -> Dict[str, float | int]:
+        ordered = sorted(samples)
+
+        def pick(percentile: float) -> float:
+            if not ordered:
+                return 0.0
+            index = min(len(ordered) - 1, round((len(ordered) - 1) * percentile))
+            return round(ordered[index], 2)
+
+        return {
+            "count": len(ordered),
+            "p50": pick(0.50),
+            "p95": pick(0.95),
+            "p99": pick(0.99),
+        }
+
     # ------------------------------------------------------------------
     # Formatting helpers
     # ------------------------------------------------------------------