diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index 62d2e9e6..bfccf182 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py index ee600874..24dfc1a2 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py +++ b/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": diff --git a/haiku_rag_slim/haiku/rag/graph/research/state.py b/haiku_rag_slim/haiku/rag/graph/research/state.py index 8ac49c95..17e876ee 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/state.py +++ b/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( diff --git a/tests/graph/test_search_filter.py b/tests/graph/test_search_filter.py new file mode 100644 index 00000000..0c9e99bd --- /dev/null +++ b/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']}'"