feat: add filter_docs to HybridRetriever and get_retriever_filtered helper

server/retriever.py

@@ -57,6 +57,7 @@ class HybridRetriever(BaseRetriever):
     workspace_id: str = "default"
     use_hyde: bool = False
     use_multi_query: bool = False
+    filter_docs: list[str] | None = None
 
     class Config:
         arbitrary_types_allowed = True
@@ -121,7 +122,10 @@ class HybridRetriever(BaseRetriever):
             return query
 
     def _dense_retrieve(self, query: str, k: int) -> list[dict]:
-        results = self.vectorstore.similarity_search_with_relevance_scores(query, k=k)
+        filter_arg = {"source": {"$in": self.filter_docs}} if self.filter_docs else None
+        results = self.vectorstore.similarity_search_with_relevance_scores(
+            query, k=k, filter=filter_arg
+        )
         output = []
         for doc, score in results:
             output.append({
@@ -173,7 +177,10 @@ class HybridRetriever(BaseRetriever):
         for q in queries:
             dense_query = self._hyde_expand(q) if self.use_hyde else q
             d_results = self._dense_retrieve(dense_query, k=self.retrieve_k)
-            s_results = get_index(self.workspace_id).search(q, k=self.retrieve_k)
+            filter_sources = set(self.filter_docs) if self.filter_docs else None
+            s_results = get_index(self.workspace_id).search(
+                q, k=self.retrieve_k, filter_sources=filter_sources
+            )
 
             for rank, doc in enumerate(d_results):
                 key = doc["content"][:120]
@@ -226,6 +233,23 @@ def get_retriever(workspace_id: str = DEFAULT_WORKSPACE) -> HybridRetriever:
     return _retriever_cache[workspace_id]
 
 
+def get_retriever_filtered(workspace_id: str, filter_docs: list[str]) -> HybridRetriever:
+    """One-off retriever with doc filter applied. Reuses cached vectorstore; does NOT cache itself."""
+    config = load_config()
+    retrieval_cfg = config.get("retrieval", {})
+    return HybridRetriever(
+        vectorstore=get_vectorstore(workspace_id),
+        dense_weight=retrieval_cfg.get("dense_weight", 0.7),
+        sparse_weight=retrieval_cfg.get("sparse_weight", 0.3),
+        retrieve_k=retrieval_cfg.get("retrieve_k", 10),
+        rerank_k=retrieval_cfg.get("rerank_k", 5),
+        workspace_id=workspace_id,
+        use_hyde=retrieval_cfg.get("hyde_enabled", False),
+        use_multi_query=retrieval_cfg.get("multi_query_enabled", False),
+        filter_docs=filter_docs,
+    )
+
+
 def retrieve_with_scores(query: str, k: int = 5) -> list[dict]:
     """Compatibility shim for precision eval — returns top-k chunks as dicts."""
     retriever = get_retriever()

tests/test_retriever_filter.py

@@ -0,0 +1,74 @@
+from unittest.mock import MagicMock, patch
+from langchain_core.documents import Document
+from server.retriever import HybridRetriever, get_retriever_filtered
+
+
+def _make_retriever(filter_docs=None):
+    mock_vs = MagicMock()
+    mock_vs.similarity_search_with_relevance_scores.return_value = [
+        (Document(page_content="UPI grew in 2024", metadata={"source": "rbi.pdf", "page": 1, "chunk_index": 0}), 0.85),
+        (Document(page_content="NPCI payments record", metadata={"source": "npci.pdf", "page": 1, "chunk_index": 0}), 0.72),
+    ]
+    return HybridRetriever(
+        vectorstore=mock_vs,
+        dense_weight=0.7,
+        sparse_weight=0.3,
+        retrieve_k=10,
+        rerank_k=5,
+        workspace_id="test",
+        use_hyde=False,
+        use_multi_query=False,
+        filter_docs=filter_docs,
+    )
+
+
+def test_filter_docs_field_defaults_to_none():
+    retriever = _make_retriever()
+    assert retriever.filter_docs is None
+
+
+def test_filter_docs_field_set():
+    retriever = _make_retriever(filter_docs=["rbi.pdf"])
+    assert retriever.filter_docs == ["rbi.pdf"]
+
+
+def test_dense_retrieve_passes_filter_to_chroma():
+    retriever = _make_retriever(filter_docs=["rbi.pdf"])
+    retriever._dense_retrieve("UPI", k=5)
+    call_kwargs = retriever.vectorstore.similarity_search_with_relevance_scores.call_args
+    assert call_kwargs.kwargs.get("filter") == {"source": {"$in": ["rbi.pdf"]}}
+
+
+def test_dense_retrieve_no_filter_when_none():
+    retriever = _make_retriever(filter_docs=None)
+    retriever._dense_retrieve("UPI", k=5)
+    call_kwargs = retriever.vectorstore.similarity_search_with_relevance_scores.call_args
+    assert call_kwargs.kwargs.get("filter") is None
+
+
+def test_get_retriever_filtered_returns_retriever_with_filter():
+    mock_vs = MagicMock()
+    with patch("server.retriever.get_vectorstore", return_value=mock_vs), \
+         patch("server.retriever.load_config", return_value={
+             "retrieval": {"dense_weight": 0.7, "sparse_weight": 0.3,
+                           "retrieve_k": 10, "rerank_k": 5,
+                           "hyde_enabled": False, "multi_query_enabled": False}
+         }):
+        r = get_retriever_filtered("default", ["rbi.pdf"])
+    assert r.filter_docs == ["rbi.pdf"]
+    assert r.vectorstore is mock_vs
+
+
+def test_get_retriever_filtered_does_not_modify_singleton_cache():
+    from server.retriever import _retriever_cache
+    initial_keys = set(_retriever_cache.keys())
+    mock_vs = MagicMock()
+    with patch("server.retriever.get_vectorstore", return_value=mock_vs), \
+         patch("server.retriever.load_config", return_value={
+             "retrieval": {"dense_weight": 0.7, "sparse_weight": 0.3,
+                           "retrieve_k": 10, "rerank_k": 5,
+                           "hyde_enabled": False, "multi_query_enabled": False}
+         }):
+        get_retriever_filtered("workspace-filtered", ["doc.pdf"])
+    assert "workspace-filtered" not in _retriever_cache
+    assert set(_retriever_cache.keys()) == initial_keys