Search code annotations from code store

Author: strongkeep-debug

src/pipelines/retrieval.py

@@ -37,7 +37,7 @@
 from src.graph.neo4j_client import Neo4jClient
 from src.prompts.retrieval import ANSWER_PROMPT, build_system_prompt
 from src.schemas.retrieval import RetrievalResult, SourceRecord
-from src.schemas.code import snippets_namespace
+from src.schemas.code import annotations_namespace, snippets_namespace
 from src.storage.base import BaseVectorStore
 from src.storage.factory import get_vector_store
 
@@ -121,6 +121,8 @@ def __init__(
         model: Optional[BaseChatModel] = None,
         vector_store: Optional[BaseVectorStore] = None,
         neo4j_client: Optional[Neo4jClient] = None,
+        code_vector_store: Optional[BaseVectorStore] = None,
+        org_id: str = "default",
     ) -> None:
         # ── LLM ───────────────────────────────────────────────────────
         if model is None:
@@ -139,6 +141,15 @@ def __init__(
             self.vector_store = get_vector_store()
         else:
             self.vector_store = vector_store
+        if code_vector_store is not None:
+            self.code_vector_store = code_vector_store
+        elif vector_store is None:
+            self.code_vector_store = get_vector_store(
+                namespace=annotations_namespace(org_id),
+                create_if_not_exists=False,
+            )
+        else:
+            self.code_vector_store = self.vector_store
 
         # ── Graph store (Neo4j) ───────────────────────────────────────
         embed_fn = _get_embed_fn()
@@ -579,7 +590,7 @@ async def _search_code(
     ) -> List[SourceRecord]:
         """Semantic search over stored code annotations."""
 
-        results = await self.vector_store.search_by_text(
+        results = await self.code_vector_store.search_by_text(
             query_text=query,
             top_k=top_k,
             filters={

tests/integration/test_retrieval_pipeline.py

@@ -171,6 +171,50 @@ async def test_raw_search_returns_ranked_hits_without_tool_selection(
     assert not pipeline.model_with_tools.calls
 
 
+@pytest.mark.asyncio
+async def test_raw_code_search_uses_code_annotation_store(vector_store, neo4j_client):
+    code_vector_store = type(vector_store)()
+    vector_store.seed(
+        "default-code-1",
+        "This stale default-store code record must not be returned.",
+        {
+            "user_id": "alice",
+            "domain": "code",
+            "target_symbol": "WrongStore",
+            "target_file": "src/wrong.py",
+        },
+        score=0.99,
+    )
+    code_vector_store.seed(
+        "annotation-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.8,
+    )
+    model = FakeChatModel()
+    pipeline = RetrievalPipeline(
+        model=model,
+        vector_store=vector_store,
+        code_vector_store=code_vector_store,
+        neo4j_client=neo4j_client,
+    )
+
+    results = await pipeline.search_raw("retry", "alice", ["code"], top_k=5)
+
+    assert [record.metadata["id"] for record in results] == ["annotation-code-1"]
+    assert "symbol=RetryLoop" in results[0].content
+    assert "WrongStore" not in results[0].content
+    assert not pipeline.model_with_tools.calls
+
+
 @pytest.mark.asyncio
 async def test_raw_search_runs_requested_domains_concurrently(
     vector_store, neo4j_client