Include code annotations in raw search

Author: strongkeep-debug

src/api/schemas.py

@@ -158,7 +158,7 @@ class SearchRequest(BaseModel):
         ..., min_length=1, max_length=256, pattern=r"^[\w.\-@]+$",
     )
     domains: List[str] = Field(
-        default=["profile", "temporal", "summary", "snippet"],
+        default=["profile", "temporal", "summary", "snippet", "code"],
         description="Which memory domains to search",
     )
     top_k: int = Field(default=10, ge=1, le=100)
@@ -170,7 +170,7 @@ class SearchRequest(BaseModel):
     @field_validator("domains")
     @classmethod
     def validate_domains(cls, v: List[str]) -> List[str]:
-        allowed = {"profile", "temporal", "summary", "snippet"}
+        allowed = {"profile", "temporal", "summary", "snippet", "code"}
         for d in v:
             if d not in allowed:
                 raise ValueError(f"Invalid domain '{d}'. Allowed: {allowed}")

src/pipelines/retrieval.py

@@ -317,6 +317,8 @@ async def search_raw(
             tasks.append(self._search_summary(query, user_id, top_k))
         if "snippet" in domain_set:
             tasks.append(self._search_snippet(query, user_id, top_k))
+        if "code" in domain_set:
+            tasks.append(self._search_code(query, user_id, top_k))
 
         if not tasks:
             return []
@@ -561,6 +563,53 @@ async def _search_summary(
         logger.info("  → Summary [%s]: %d results", query, len(records))
         return records
 
+    # -- Code: Pinecone semantic search --------------------------------
+
+    async def _search_code(
+        self,
+        query: str,
+        user_id: str,
+        top_k: int = 10,
+    ) -> List[SourceRecord]:
+        """Semantic search over stored code annotations."""
+
+        results = await self.vector_store.search_by_text(
+            query_text=query,
+            top_k=top_k,
+            filters={
+                "user_id": user_id,
+                "domain": "code",
+            },
+        )
+
+        records = []
+        for r in results:
+            metadata = dict(r.metadata)
+            detail_parts = []
+            for label, key in (
+                ("repo", "repo"),
+                ("file", "target_file"),
+                ("symbol", "target_symbol"),
+                ("type", "annotation_type"),
+                ("severity", "severity"),
+            ):
+                value = metadata.get(key)
+                if value:
+                    detail_parts.append(f"{label}={value}")
+
+            prefix = f"[{'; '.join(detail_parts)}] " if detail_parts else ""
+            records.append(
+                SourceRecord(
+                    domain="code",
+                    content=f"{prefix}{r.content}",
+                    score=r.score,
+                    metadata={"id": r.id, **metadata},
+                )
+            )
+
+        logger.info("  -> Code [%s]: %d results", query, len(records))
+        return records
+
     # -- Snippet: Pinecone semantic search ─────────────────────────────
 
     def _get_snippet_store(self, user_id: str) -> BaseVectorStore:

src/schemas/retrieval.py

@@ -26,7 +26,7 @@ def source_count(self) -> int:
 class SourceRecord:
     """A single piece of evidence fetched from a data store."""
 
-    domain: str              # "profile", "temporal", "summary", "snippet"
+    domain: str              # "profile", "temporal", "summary", "snippet", "code"
     content: str             # the actual text
     score: float = 0.0       # similarity score (if applicable)
     metadata: Dict[str, Any] = field(default_factory=dict)

tests/api/test_memory_search_routes.py

@@ -18,17 +18,38 @@ class FakeSearchPipeline:
 
     def __init__(self) -> None:
         self.answer_calls = 0
+        self.search_calls: list[dict[str, object]] = []
         self.latencies: dict[str, list[float]] = {}
 
-    async def search_raw(self, query: str, user_id: str, domains: list[str], top_k: int):
+    async def search_raw(
+        self, query: str, user_id: str, domains: list[str], top_k: int
+    ):
         assert query == "latency"
         assert user_id == "Static Key User"
-        assert domains == ["profile", "summary"]
         assert top_k == 3
-        return [
-            SourceRecord(domain="summary", content="Low-latency summary", score=0.9),
-            SourceRecord(domain="profile", content="work / company = XMem", score=0.7),
-        ]
+        self.search_calls.append(
+            {"query": query, "user_id": user_id, "domains": domains, "top_k": top_k}
+        )
+
+        fixtures = {
+            "summary": SourceRecord(
+                domain="summary",
+                content="Low-latency summary",
+                score=0.9,
+            ),
+            "profile": SourceRecord(
+                domain="profile",
+                content="work / company = XMem",
+                score=0.7,
+            ),
+            "code": SourceRecord(
+                domain="code",
+                content="[file=src/retry.py; symbol=RetryLoop] Timeout retry note",
+                score=0.8,
+                metadata={"target_file": "src/retry.py"},
+            ),
+        }
+        return [fixtures[domain] for domain in domains if domain in fixtures]
 
     async def answer_from_sources(self, query: str, sources: list[SourceRecord]) -> str:
         self.answer_calls += 1
@@ -39,7 +60,12 @@ def record_latency(self, mode: str, elapsed_ms: float) -> None:
 
     def get_latency_snapshot(self):
         return {
-            mode: {"count": len(samples), "p50": samples[-1], "p95": samples[-1], "p99": samples[-1]}
+            mode: {
+                "count": len(samples),
+                "p50": samples[-1],
+                "p95": samples[-1],
+                "p99": samples[-1],
+            }
             for mode, samples in self.latencies.items()
         }
 
@@ -78,6 +104,7 @@ def test_memory_search_route_returns_raw_hits_without_answer(memory_search_app):
     assert payload["data"]["total"] == 2
     assert payload["data"]["answer"] == ""
     assert payload["data"]["latency"]["raw"]["count"] == 1
+    assert pipeline.search_calls[0]["domains"] == ["profile", "summary"]
     assert pipeline.answer_calls == 0
 
 
@@ -102,3 +129,25 @@ def test_root_search_alias_can_synthesize_answer(memory_search_app):
     assert payload["data"]["model"] == "fake-retrieval"
     assert payload["data"]["latency"]["answer"]["count"] == 1
     assert pipeline.answer_calls == 1
+
+
+def test_memory_search_route_accepts_code_domain(memory_search_app):
+    app, pipeline = memory_search_app
+    response = TestClient(app).post(
+        "/v1/memory/search",
+        headers={"Authorization": "Bearer test-static-key"},
+        json={
+            "query": "latency",
+            "user_id": "ignored-by-auth",
+            "domains": ["code"],
+            "top_k": 3,
+        },
+    )
+
+    payload = response.json()
+
+    assert response.status_code == 200
+    assert payload["data"]["total"] == 1
+    assert payload["data"]["results"][0]["domain"] == "code"
+    assert payload["data"]["results"][0]["metadata"]["target_file"] == "src/retry.py"
+    assert pipeline.search_calls[0]["domains"] == ["code"]

tests/integration/test_retrieval_pipeline.py

@@ -121,6 +121,20 @@ async def test_raw_search_returns_ranked_hits_without_tool_selection(
         {"user_id": "alice", "domain": "summary"},
         score=0.9,
     )
+    vector_store.seed(
+        "code-1",
+        "RetryLoop can spin when the first retrieval attempt times out.",
+        {
+            "user_id": "alice",
+            "domain": "code",
+            "annotation_type": "bug_report",
+            "target_symbol": "RetryLoop",
+            "target_file": "src/retry.py",
+            "repo": "xmem",
+            "severity": "high",
+        },
+        score=0.95,
+    )
     neo4j_client.seed_event(
         user_id="alice",
         date="05-11",
@@ -137,15 +151,23 @@ async def test_raw_search_returns_ranked_hits_without_tool_selection(
     results = await pipeline.search_raw(
         "latency",
         "alice",
-        ["profile", "temporal", "summary"],
+        ["profile", "temporal", "summary", "code"],
         top_k=5,
     )
 
     assert [record.score for record in results] == sorted(
         [record.score for record in results],
         reverse=True,
     )
-    assert {record.domain for record in results} == {"profile", "temporal", "summary"}
+    assert {record.domain for record in results} == {
+        "profile",
+        "temporal",
+        "summary",
+        "code",
+    }
+    code_hit = next(record for record in results if record.domain == "code")
+    assert "file=src/retry.py" in code_hit.content
+    assert code_hit.metadata["target_symbol"] == "RetryLoop"
     assert not pipeline.model_with_tools.calls