Search filter in qa & research graphs

haiku_rag_slim/haiku/rag/graph/common/nodes.py

@@ -34,6 +34,7 @@ class GraphState(Protocol):
 
     context: GraphContext
     max_concurrency: int
+    search_filter: str | None
 
 
 class GraphDeps(Protocol):
@@ -99,11 +100,16 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
 
             plan_agent = Agent(**agent_config)
 
+            # Capture search filter for use in tool
+            search_filter = state.search_filter
+
             @plan_agent.tool
             async def gather_context(
                 ctx2: RunContext[AgentDepsT], query: str, limit: int | None = None
             ) -> str:
-                results = await ctx2.deps.client.search(query, limit=limit)
+                results = await ctx2.deps.client.search(
+                    query, limit=limit, filter=search_filter
+                )
                 results = await ctx2.deps.client.expand_context(results)
                 return "\n\n".join(r.content for r in results)
 
@@ -225,6 +231,9 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
         deps_type=deps_type,
     )
 
+    # Capture search filter for use in tool
+    search_filter = state.search_filter
+
     @agent.tool
     async def search_and_answer(
         ctx2: RunContext[AgentDepsT], query: str, limit: int | None = None
@@ -234,7 +243,9 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
         Returns results with chunk IDs and relevance scores.
         Reference results by their chunk_id in cited_chunks.
         """
-        results = await ctx2.deps.client.search(query, limit=limit)
+        results = await ctx2.deps.client.search(
+            query, limit=limit, filter=search_filter
+        )
         results = await ctx2.deps.client.expand_context(results)
         # Store results for citation resolution
         ctx2.deps.search_results = results

haiku_rag_slim/haiku/rag/graph/deep_qa/state.py

@@ -36,6 +36,9 @@ class DeepQAState(BaseModel):
         default=1, description="Maximum parallel sub-question searches"
     )
     iterations: int = Field(default=0, description="Current iteration number")
+    search_filter: str | None = Field(
+        default=None, description="SQL WHERE clause to filter search results"
+    )
 
     @classmethod
     def from_config(cls, context: DeepQAContext, config: "AppConfig") -> "DeepQAState":

haiku_rag_slim/haiku/rag/graph/research/state.py

@@ -63,6 +63,9 @@ class ResearchState(BaseModel):
     last_analysis: InsightAnalysis | None = Field(
         default=None, description="Last insight analysis"
     )
+    search_filter: str | None = Field(
+        default=None, description="SQL WHERE clause to filter search results"
+    )
 
     @classmethod
     def from_config(

tests/graph/test_search_filter.py

@@ -0,0 +1,187 @@
+import pytest
+from pydantic_ai.models.test import TestModel
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
+from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
+from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
+from haiku.rag.graph.research.dependencies import ResearchContext
+from haiku.rag.graph.research.graph import build_research_graph
+from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+
+
+@pytest.fixture
+async def client_with_docs(temp_db_path):
+    """Create a client with two distinct documents."""
+    client = HaikuRAG(temp_db_path, create=True)
+
+    # Add two documents with distinct content
+    doc1 = await client.create_document(
+        "Document about cats: Cats are small furry mammals that purr.",
+        title="Cat Facts",
+    )
+    doc2 = await client.create_document(
+        "Document about dogs: Dogs are loyal companions that bark.",
+        title="Dog Facts",
+    )
+
+    yield client, doc1.id, doc2.id
+
+    client.close()
+
+
+@pytest.mark.asyncio
+async def test_search_filter_restricts_results(client_with_docs):
+    """Test that search_filter restricts search to specified documents."""
+    client, doc1_id, doc2_id = client_with_docs
+
+    # Search without filter - should find both
+    results_all = await client.search("animals mammals companions")
+    assert len(results_all) >= 1
+
+    # Search with filter for doc1 only
+    filter_doc1 = f"id = '{doc1_id}'"
+    results_filtered = await client.search(
+        "animals mammals companions", filter=filter_doc1
+    )
+
+    # All results should be from doc1
+    for result in results_filtered:
+        assert result.document_id == doc1_id
+
+
+@pytest.mark.asyncio
+async def test_research_graph_uses_search_filter(monkeypatch, client_with_docs):
+    """Test that research graph passes search_filter to search operations."""
+    client, doc1_id, doc2_id = client_with_docs
+
+    # Track search calls to verify filter is passed
+    search_calls = []
+    original_search = client.search
+
+    async def tracking_search(query, limit=None, search_type="hybrid", filter=None):
+        search_calls.append({"query": query, "filter": filter})
+        return await original_search(query, limit, search_type, filter)
+
+    client.search = tracking_search
+
+    # Mock get_model to return TestModel
+    def test_model_factory(_provider, _model, _config=None):
+        return TestModel()
+
+    monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
+
+    graph = build_research_graph()
+
+    # Create state with search_filter
+    filter_clause = f"id = '{doc1_id}'"
+    state = ResearchState(
+        context=ResearchContext(original_question="Tell me about animals"),
+        max_iterations=1,
+        confidence_threshold=0.5,
+        search_filter=filter_clause,
+    )
+
+    deps = ResearchDeps(client=client)
+
+    await graph.run(state=state, deps=deps)
+
+    # Verify search was called with the filter
+    assert len(search_calls) > 0, "Expected search to be called"
+    for call in search_calls:
+        assert call["filter"] == filter_clause, (
+            f"Expected filter '{filter_clause}', got '{call['filter']}'"
+        )
+
+
+@pytest.mark.asyncio
+async def test_deep_qa_graph_uses_search_filter(monkeypatch, client_with_docs):
+    """Test that deep QA graph passes search_filter to search operations."""
+    client, doc1_id, doc2_id = client_with_docs
+
+    # Track search calls to verify filter is passed
+    search_calls = []
+    original_search = client.search
+
+    async def tracking_search(query, limit=None, search_type="hybrid", filter=None):
+        search_calls.append({"query": query, "filter": filter})
+        return await original_search(query, limit, search_type, filter)
+
+    client.search = tracking_search
+
+    # Mock get_model to return TestModel
+    def test_model_factory(_provider, _model, _config=None):
+        return TestModel()
+
+    monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
+
+    graph = build_deep_qa_graph()
+
+    # Create state with search_filter
+    filter_clause = f"id = '{doc2_id}'"
+    state = DeepQAState(
+        context=DeepQAContext(original_question="Tell me about animals"),
+        max_sub_questions=2,
+        search_filter=filter_clause,
+    )
+
+    deps = DeepQADeps(client=client)
+
+    await graph.run(state=state, deps=deps)
+
+    # Verify search was called with the filter
+    assert len(search_calls) > 0, "Expected search to be called"
+    for call in search_calls:
+        assert call["filter"] == filter_clause, (
+            f"Expected filter '{filter_clause}', got '{call['filter']}'"
+        )
+
+
+@pytest.mark.asyncio
+async def test_search_filter_none_searches_all(monkeypatch, client_with_docs):
+    """Test that search_filter=None searches all documents."""
+    client, doc1_id, doc2_id = client_with_docs
+
+    # Track search calls
+    search_calls = []
+    original_search = client.search
+
+    async def tracking_search(query, limit=None, search_type="hybrid", filter=None):
+        search_calls.append({"query": query, "filter": filter})
+        return await original_search(query, limit, search_type, filter)
+
+    client.search = tracking_search
+
+    # Mock get_model
+    def test_model_factory(_provider, _model, _config=None):
+        return TestModel()
+
+    monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
+
+    graph = build_research_graph()
+
+    # Create state without search_filter (None)
+    state = ResearchState(
+        context=ResearchContext(original_question="Tell me about animals"),
+        max_iterations=1,
+        confidence_threshold=0.5,
+        search_filter=None,
+    )
+
+    deps = ResearchDeps(client=client)
+
+    await graph.run(state=state, deps=deps)
+
+    # Verify search was called with None filter
+    assert len(search_calls) > 0, "Expected search to be called"
+    for call in search_calls:
+        assert call["filter"] is None, f"Expected filter None, got '{call['filter']}'"