From b70fd8b7b46f141c34b57ffe3805f08c9d1d29e3 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 11 Dec 2025 12:44:14 +0200
Subject: [PATCH 1/6] Api for listing documents

---
 examples/ag-ui-research/backend/main.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py
index 32a2fec0..48593415 100644
--- a/examples/ag-ui-research/backend/main.py
+++ b/examples/ag-ui-research/backend/main.py
@@ -189,6 +189,19 @@ async def health_check(_: Request) -> JSONResponse:
     )
 
 
+async def list_documents(_: Request) -> JSONResponse:
+    """List all documents in the database."""
+    client = get_client(db_path)
+    docs = await client.document_repository.list_all()
+    return JSONResponse(
+        {
+            "documents": [
+                {"id": doc.id, "title": doc.title, "uri": doc.uri} for doc in docs
+            ]
+        }
+    )
+
+
 async def visualize_chunk(request: Request) -> JSONResponse:
     """Return visual grounding images for a chunk as base64."""
     import base64
@@ -228,6 +241,7 @@ async def visualize_chunk(request: Request) -> JSONResponse:
 app = Starlette(
     routes=[
         Route("/v1/research/stream", stream_research_agent, methods=["POST"]),
+        Route("/api/documents", list_documents, methods=["GET"]),
         Route("/api/visualize/{chunk_id}", visualize_chunk, methods=["GET"]),
         Route("/health", health_check, methods=["GET"]),
     ],

From 5f4f5f70f0d3d769c78db29699344a868e8b1cf7 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 11 Dec 2025 14:37:43 +0200
Subject: [PATCH 2/6] Search filter in qa & research graphs

---
 .../haiku/rag/graph/common/nodes.py           |  15 +-
 .../haiku/rag/graph/deep_qa/state.py          |   3 +
 .../haiku/rag/graph/research/state.py         |   3 +
 tests/graph/test_search_filter.py             | 187 ++++++++++++++++++
 4 files changed, 206 insertions(+), 2 deletions(-)
 create mode 100644 tests/graph/test_search_filter.py

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']}'"

From 39cbec1b88353036e293b201db36bcf99393464d Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 11 Dec 2025 14:38:23 +0200
Subject: [PATCH 3/6] Set search filters from ducment filter in example

---
 examples/ag-ui-research/backend/agent.py |  2 ++
 examples/ag-ui-research/backend/main.py  | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py
index 0182f1e6..3dce81ac 100644
--- a/examples/ag-ui-research/backend/agent.py
+++ b/examples/ag-ui-research/backend/agent.py
@@ -33,6 +33,7 @@ class AgentDeps:
 
     client: HaikuRAG
     agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
+    search_filter: str | None = None
 
 
 model = get_model(Config.research.model, Config)
@@ -74,6 +75,7 @@ async def run_research(ctx: RunContext[AgentDeps], question: str) -> str:
     graph = build_research_graph(Config)
     context = ResearchContext(original_question=question)
     state = ResearchState.from_config(context=context, config=Config)
+    state.search_filter = ctx.deps.search_filter
 
     graph_deps = ResearchDeps(
         client=ctx.deps.client,
diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py
index 48593415..3909c496 100644
--- a/examples/ag-ui-research/backend/main.py
+++ b/examples/ag-ui-research/backend/main.py
@@ -92,8 +92,19 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
                     effective_db_path = Path(effective_db_path)
                 client = get_client(effective_db_path)
 
+                # Build search filter from document IDs if provided
+                document_ids = input_data.state.get("documentFilter")
+                search_filter = None
+                if document_ids:
+                    ids_str = ", ".join(f"'{id}'" for id in document_ids)
+                    search_filter = f"id IN ({ids_str})"
+
                 # Create agent dependencies with shared emitter
-                agent_deps = AgentDeps(client=client, agui_emitter=emitter)
+                agent_deps = AgentDeps(
+                    client=client,
+                    agui_emitter=emitter,
+                    search_filter=search_filter,
+                )
 
                 # Start run with empty initial state
                 emitter.start_run(

From 9fd33a2d8337a00ffbe44c700202c6c4c73cfd74 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 11 Dec 2025 14:56:18 +0200
Subject: [PATCH 4/6] DocumentSelector in frontend to select documents to do
 research on.

---
 examples/ag-ui-research/backend/main.py       |   4 +-
 .../frontend/components/Agent.tsx             |  17 +-
 .../frontend/components/DocumentSelector.tsx  | 352 ++++++++++++++++++
 3 files changed, 370 insertions(+), 3 deletions(-)
 create mode 100644 examples/ag-ui-research/frontend/components/DocumentSelector.tsx

diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py
index 3909c496..d432a8ba 100644
--- a/examples/ag-ui-research/backend/main.py
+++ b/examples/ag-ui-research/backend/main.py
@@ -92,8 +92,8 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
                     effective_db_path = Path(effective_db_path)
                 client = get_client(effective_db_path)
 
-                # Build search filter from document IDs if provided
-                document_ids = input_data.state.get("documentFilter")
+                # Build search filter from document IDs (empty list = search all)
+                document_ids = input_data.state.get("documentFilter") or []
                 search_filter = None
                 if document_ids:
                     ids_str = ", ".join(f"'{id}'" for id in document_ids)
diff --git a/examples/ag-ui-research/frontend/components/Agent.tsx b/examples/ag-ui-research/frontend/components/Agent.tsx
index ac297cfc..a8dc5f59 100644
--- a/examples/ag-ui-research/frontend/components/Agent.tsx
+++ b/examples/ag-ui-research/frontend/components/Agent.tsx
@@ -3,6 +3,7 @@
 import { CopilotKit, useCoAgent } from "@copilotkit/react-core";
 import { CopilotChat } from "@copilotkit/react-ui";
 import "@copilotkit/react-ui/styles.css";
+import DocumentSelector from "./DocumentSelector";
 import StateDisplay from "./StateDisplay";
 
 interface InsightRecord {
@@ -82,10 +83,11 @@ interface ResearchState {
 	result?: ResearchReport;
 	current_activity?: string;
 	current_activity_message?: string;
+	documentFilter?: string[];
 }
 
 function AgentContent() {
-	const { state } = useCoAgent<ResearchState>({
+	const { state, setState, running } = useCoAgent<ResearchState>({
 		name: "research_agent",
 		initialState: {
 			context: {
@@ -101,9 +103,14 @@ function AgentContent() {
 			max_concurrency: 1,
 			last_eval: null,
 			last_analysis: null,
+			documentFilter: [],
 		},
 	});
 
+	const handleDocumentFilterChange = (ids: string[]) => {
+		setState({ ...state, documentFilter: ids });
+	};
+
 	return (
 		<>
 			<style>{`
@@ -163,6 +170,14 @@ function AgentContent() {
 							</p>
 						</header>
 
+						<div style={{ marginBottom: "1rem" }}>
+							<DocumentSelector
+								selected={state.documentFilter || []}
+								onChange={handleDocumentFilterChange}
+								disabled={running}
+							/>
+						</div>
+
 						<StateDisplay state={state} />
 					</div>
 				</div>
diff --git a/examples/ag-ui-research/frontend/components/DocumentSelector.tsx b/examples/ag-ui-research/frontend/components/DocumentSelector.tsx
new file mode 100644
index 00000000..33060026
--- /dev/null
+++ b/examples/ag-ui-research/frontend/components/DocumentSelector.tsx
@@ -0,0 +1,352 @@
+"use client";
+
+import { useCallback, useEffect, useMemo, useState } from "react";
+
+interface Document {
+	id: string;
+	title: string;
+	uri: string;
+}
+
+interface DocumentSelectorProps {
+	selected: string[];
+	onChange: (ids: string[]) => void;
+	disabled: boolean;
+}
+
+export default function DocumentSelector({
+	selected,
+	onChange,
+	disabled,
+}: DocumentSelectorProps) {
+	const [documents, setDocuments] = useState<Document[]>([]);
+	const [loading, setLoading] = useState(true);
+	const [error, setError] = useState<string | null>(null);
+	const [expanded, setExpanded] = useState(true);
+	const [searchQuery, setSearchQuery] = useState("");
+
+	useEffect(() => {
+		const fetchDocuments = async () => {
+			try {
+				const response = await fetch(
+					`${process.env.NEXT_PUBLIC_API_URL || "http://localhost:8000"}/api/documents`,
+				);
+				if (!response.ok) {
+					throw new Error("Failed to fetch documents");
+				}
+				const data = await response.json();
+				const docs = data.documents || [];
+				setDocuments(docs);
+				// Select all documents by default if none are selected
+				if (selected.length === 0 && docs.length > 0) {
+					onChange(docs.map((d: Document) => d.id));
+				}
+				setError(null);
+			} catch (err) {
+				setError(err instanceof Error ? err.message : "Unknown error");
+			} finally {
+				setLoading(false);
+			}
+		};
+
+		fetchDocuments();
+	}, []);
+
+	const filteredDocuments = useMemo(() => {
+		if (!searchQuery.trim()) return documents;
+		const query = searchQuery.toLowerCase();
+		return documents.filter(
+			(doc) =>
+				(doc.title || "").toLowerCase().includes(query) ||
+				(doc.uri || "").toLowerCase().includes(query),
+		);
+	}, [documents, searchQuery]);
+
+	const handleToggle = useCallback(
+		(id: string) => {
+			if (disabled) return;
+			if (selected.includes(id)) {
+				onChange(selected.filter((s) => s !== id));
+			} else {
+				onChange([...selected, id]);
+			}
+		},
+		[selected, onChange, disabled],
+	);
+
+	const handleSelectAll = useCallback(() => {
+		if (disabled) return;
+		const filteredIds = filteredDocuments.map((d) => d.id);
+		const allFilteredSelected = filteredIds.every((id) =>
+			selected.includes(id),
+		);
+		if (allFilteredSelected) {
+			// Deselect all filtered documents
+			onChange(selected.filter((id) => !filteredIds.includes(id)));
+		} else {
+			// Select all filtered documents (add to existing selection)
+			const newSelection = [...new Set([...selected, ...filteredIds])];
+			onChange(newSelection);
+		}
+	}, [selected, filteredDocuments, onChange, disabled]);
+
+	const selectedCount = selected.length;
+	const totalCount = documents.length;
+	const filterActive = selectedCount > 0 && selectedCount < totalCount;
+
+	return (
+		<div
+			style={{
+				background: "white",
+				borderRadius: "8px",
+				boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
+				overflow: "hidden",
+				opacity: disabled ? 0.6 : 1,
+				transition: "opacity 0.2s",
+			}}
+		>
+			<button
+				type="button"
+				onClick={() => setExpanded(!expanded)}
+				style={{
+					width: "100%",
+					display: "flex",
+					justifyContent: "space-between",
+					alignItems: "center",
+					padding: "0.75rem",
+					background: filterActive ? "#ebf8ff" : "#edf2f7",
+					border: filterActive ? "1px solid #90cdf4" : "1px solid #e2e8f0",
+					borderRadius: expanded ? "8px 8px 0 0" : "8px",
+					cursor: "pointer",
+					fontSize: "0.875rem",
+					fontWeight: "600",
+					color: filterActive ? "#2b6cb0" : "#2d3748",
+				}}
+			>
+				<span>
+					Document Filter
+					{filterActive && ` (${selectedCount}/${totalCount})`}
+					{!filterActive && selectedCount === 0 && " (All)"}
+				</span>
+				<span>{expanded ? "▼" : "▶"}</span>
+			</button>
+
+			{expanded && (
+				<div
+					style={{
+						padding: "0.75rem",
+						background: "#f7fafc",
+						borderLeft: "1px solid #e2e8f0",
+						borderRight: "1px solid #e2e8f0",
+						borderBottom: "1px solid #e2e8f0",
+						borderRadius: "0 0 8px 8px",
+					}}
+				>
+					{loading && (
+						<div
+							style={{
+								padding: "1rem",
+								textAlign: "center",
+								color: "#718096",
+								fontSize: "0.875rem",
+							}}
+						>
+							Loading documents...
+						</div>
+					)}
+
+					{error && (
+						<div
+							style={{
+								padding: "0.75rem",
+								background: "#fed7d7",
+								color: "#c53030",
+								borderRadius: "4px",
+								fontSize: "0.875rem",
+							}}
+						>
+							{error}
+						</div>
+					)}
+
+					{!loading && !error && documents.length === 0 && (
+						<div
+							style={{
+								padding: "1rem",
+								textAlign: "center",
+								color: "#718096",
+								fontSize: "0.875rem",
+							}}
+						>
+							No documents in database
+						</div>
+					)}
+
+					{!loading && !error && documents.length > 0 && (
+						<>
+							{/* Search Input */}
+							<div style={{ marginBottom: "0.5rem" }}>
+								<input
+									type="text"
+									placeholder="Search by title or URI..."
+									value={searchQuery}
+									onChange={(e) => setSearchQuery(e.target.value)}
+									disabled={disabled}
+									style={{
+										width: "100%",
+										padding: "0.5rem 0.75rem",
+										fontSize: "0.875rem",
+										border: "1px solid #e2e8f0",
+										borderRadius: "4px",
+										background: disabled ? "#f7fafc" : "white",
+										color: disabled ? "#a0aec0" : "#2d3748",
+										outline: "none",
+									}}
+								/>
+							</div>
+
+							{/* Select All / Clear */}
+							<div
+								style={{
+									marginBottom: "0.5rem",
+									paddingBottom: "0.5rem",
+									borderBottom: "1px solid #e2e8f0",
+									display: "flex",
+									alignItems: "center",
+									justifyContent: "space-between",
+								}}
+							>
+								<div>
+									<button
+										type="button"
+										onClick={handleSelectAll}
+										disabled={disabled}
+										style={{
+											padding: "0.375rem 0.75rem",
+											fontSize: "0.75rem",
+											background: disabled ? "#e2e8f0" : "#4299e1",
+											color: disabled ? "#a0aec0" : "white",
+											border: "none",
+											borderRadius: "4px",
+											cursor: disabled ? "not-allowed" : "pointer",
+										}}
+									>
+										{filteredDocuments.every((d) => selected.includes(d.id))
+											? "Clear Visible"
+											: "Select Visible"}
+									</button>
+									<span
+										style={{
+											marginLeft: "0.75rem",
+											fontSize: "0.75rem",
+											color: "#718096",
+										}}
+									>
+										{selectedCount} of {totalCount} selected
+									</span>
+								</div>
+								{searchQuery && (
+									<span
+										style={{
+											fontSize: "0.75rem",
+											color: "#718096",
+										}}
+									>
+										Showing {filteredDocuments.length} of {totalCount}
+									</span>
+								)}
+							</div>
+
+							{/* Document List */}
+							<div
+								style={{
+									maxHeight: "200px",
+									overflowY: "auto",
+									display: "flex",
+									flexDirection: "column",
+									gap: "0.25rem",
+								}}
+							>
+								{filteredDocuments.map((doc) => {
+									const isSelected = selected.includes(doc.id);
+									return (
+										<label
+											key={doc.id}
+											style={{
+												display: "flex",
+												alignItems: "center",
+												gap: "0.5rem",
+												padding: "0.5rem",
+												background: isSelected ? "#ebf8ff" : "white",
+												border: isSelected
+													? "1px solid #90cdf4"
+													: "1px solid #e2e8f0",
+												borderRadius: "4px",
+												cursor: disabled ? "not-allowed" : "pointer",
+												transition: "all 0.15s",
+											}}
+										>
+											<input
+												type="checkbox"
+												checked={isSelected}
+												onChange={() => handleToggle(doc.id)}
+												disabled={disabled}
+												style={{
+													width: "1rem",
+													height: "1rem",
+													cursor: disabled ? "not-allowed" : "pointer",
+												}}
+											/>
+											<div style={{ flex: 1, minWidth: 0 }}>
+												<div
+													style={{
+														fontSize: "0.875rem",
+														fontWeight: isSelected ? "600" : "400",
+														color: "#2d3748",
+														whiteSpace: "nowrap",
+														overflow: "hidden",
+														textOverflow: "ellipsis",
+													}}
+												>
+													{doc.title || "Untitled"}
+												</div>
+												<div
+													style={{
+														fontSize: "0.7rem",
+														color: "#718096",
+														whiteSpace: "nowrap",
+														overflow: "hidden",
+														textOverflow: "ellipsis",
+													}}
+												>
+													{doc.uri}
+												</div>
+											</div>
+										</label>
+									);
+								})}
+							</div>
+						</>
+					)}
+
+					{disabled && (
+						<div
+							style={{
+								marginTop: "0.5rem",
+								padding: "0.5rem",
+								background: "#fef5e7",
+								border: "1px solid #f6ad55",
+								borderRadius: "4px",
+								fontSize: "0.75rem",
+								color: "#744210",
+								textAlign: "center",
+							}}
+						>
+							Filter locked during research
+						</div>
+					)}
+				</div>
+			)}
+		</div>
+	);
+}

From ecba8cfd382b9946aa491fbc997de16dbf8ed208 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 11 Dec 2025 15:23:50 +0200
Subject: [PATCH 5/6] Docs and cli support for filter in deep QA, Research

---
 docs/agents.md                  | 14 ++++++++++++++
 docs/cli.md                     | 23 +++++++++++++++++++++--
 haiku_rag_slim/haiku/rag/app.py |  9 ++++++++-
 haiku_rag_slim/haiku/rag/cli.py | 18 ++++++++++++++++--
 4 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/docs/agents.md b/docs/agents.md
index 376ac14d..4d218383 100644
--- a/docs/agents.md
+++ b/docs/agents.md
@@ -229,6 +229,20 @@ async with HaikuRAG(path_to_db) as client:
     print(report.executive_summary)
 ```
 
+### Filtering Documents
+
+Both Research and Deep QA graphs support restricting searches to specific documents via the `search_filter` parameter. Set it to a SQL WHERE clause before running:
+
+```python
+state = ResearchState.from_config(context=context, config=Config)
+# Only search documents with these IDs
+state.search_filter = "id IN ('doc-123', 'doc-456')"
+
+result = await graph.run(state=state, deps=deps)
+```
+
+The filter applies to all search operations in the graph (context gathering and sub-question searches). See [Filtering Search Results](python.md#filtering-search-results) for available filter columns and syntax.
+
 Alternative usage with custom config:
 
 ```python
diff --git a/docs/cli.md b/docs/cli.md
index 1d432f7d..a1207ef4 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -149,8 +149,19 @@ Show verbose output with deep QA:
 haiku-rag ask "What are the main features and architecture of haiku.rag?" --deep --verbose
 ```
 
-The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used. With `--deep`, the question is decomposed into sub-questions that are answered in parallel before synthesizing a final answer. With `--verbose` (only with `--deep`), you'll see the planning, searching, evaluation, and synthesis steps as they happen.
-When available, citations use the document title; otherwise they fall back to the URI.
+Filter to specific documents with deep QA:
+```bash
+haiku-rag ask "What are the main findings?" --deep --filter "uri LIKE '%paper%'"
+```
+
+The QA agent searches your documents for relevant information and provides a comprehensive answer. When available, citations use the document title; otherwise they fall back to the URI.
+
+Flags:
+
+- `--cite`: Include citations showing which documents were used
+- `--deep`: Decompose the question into sub-questions answered in parallel before synthesizing a final answer
+- `--verbose`: Show planning, searching, evaluation, and synthesis steps (only with `--deep`)
+- `--filter`: Restrict searches to documents matching the filter (only with `--deep`)
 
 ## Research
 
@@ -166,8 +177,16 @@ With verbose output to see progress:
 haiku-rag research "How does haiku.rag organize and query documents?" --verbose
 ```
 
+Filter to specific documents:
+
+```bash
+haiku-rag research "What are the key findings?" --filter "uri LIKE '%paper%'"
+```
+
 Flags:
+
 - `--verbose`: Show planning, searching previews, evaluation summary, and stop reason
+- `--filter`: SQL WHERE clause to filter documents (see [Filtering Search Results](python.md#filtering-search-results))
 
 Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
 
diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py
index 37d87eac..c92206f4 100644
--- a/haiku_rag_slim/haiku/rag/app.py
+++ b/haiku_rag_slim/haiku/rag/app.py
@@ -314,6 +314,7 @@ class HaikuRAGApp:
         cite: bool = False,
         deep: bool = False,
         verbose: bool = False,
+        filter: str | None = None,
     ):
         """Ask a question using the RAG system.
 
@@ -322,6 +323,7 @@ class HaikuRAGApp:
             cite: Include citations in the answer
             deep: Use deep QA mode (multi-step reasoning)
             verbose: Show verbose output
+            filter: SQL WHERE clause to filter documents
         """
         async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client:
             try:
@@ -334,6 +336,7 @@ class HaikuRAGApp:
                     graph = build_deep_qa_graph(config=self.config)
                     context = DeepQAContext(original_question=question)
                     state = DeepQAState.from_config(context=context, config=self.config)
+                    state.search_filter = filter
                     deps = DeepQADeps(client=self.client)
 
                     if verbose:
@@ -371,12 +374,15 @@ class HaikuRAGApp:
             except Exception as e:
                 self.console.print(f"[red]Error: {e}[/red]")
 
-    async def research(self, question: str, verbose: bool = False):
+    async def research(
+        self, question: str, verbose: bool = False, filter: str | None = None
+    ):
         """Run research via the pydantic-graph pipeline.
 
         Args:
             question: The research question
             verbose: Show AG-UI event stream during execution
+            filter: SQL WHERE clause to filter documents
         """
         async with HaikuRAG(db_path=self.db_path, config=self.config) as client:
             try:
@@ -387,6 +393,7 @@ class HaikuRAGApp:
                 graph = build_research_graph(config=self.config)
                 context = ResearchContext(original_question=question)
                 state = ResearchState.from_config(context=context, config=self.config)
+                state.search_filter = filter
                 deps = ResearchDeps(client=client)
 
                 if verbose:
diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py
index b22fa1f8..84216409 100644
--- a/haiku_rag_slim/haiku/rag/cli.py
+++ b/haiku_rag_slim/haiku/rag/cli.py
@@ -305,9 +305,17 @@ def ask(
         "--verbose",
         help="Show verbose progress output (only with --deep)",
     ),
+    filter: str | None = typer.Option(
+        None,
+        "--filter",
+        "-f",
+        help="SQL WHERE clause to filter documents (only with --deep)",
+    ),
 ):
     app = create_app(db)
-    asyncio.run(app.ask(question=question, cite=cite, deep=deep, verbose=verbose))
+    asyncio.run(
+        app.ask(question=question, cite=cite, deep=deep, verbose=verbose, filter=filter)
+    )
 
 
 @cli.command("research", help="Run multi-agent research and output a concise report")
@@ -325,9 +333,15 @@ def research(
         "--verbose",
         help="Show planning, searching previews, evaluation summary, and stop reason",
     ),
+    filter: str | None = typer.Option(
+        None,
+        "--filter",
+        "-f",
+        help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
+    ),
 ):
     app = create_app(db)
-    asyncio.run(app.research(question=question, verbose=verbose))
+    asyncio.run(app.research(question=question, verbose=verbose, filter=filter))
 
 
 @cli.command("settings", help="Display current configuration settings")

From d7c755f4543c94d26062086e168e04c728a52a9f Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 11 Dec 2025 15:28:41 +0200
Subject: [PATCH 6/6] Add filter to simple qa agent

---
 CHANGELOG.md                         | 13 +++++++++++++
 docs/cli.md                          |  6 +++---
 haiku_rag_slim/haiku/rag/app.py      |  2 +-
 haiku_rag_slim/haiku/rag/cli.py      |  2 +-
 haiku_rag_slim/haiku/rag/client.py   |  8 ++++++--
 haiku_rag_slim/haiku/rag/qa/agent.py | 15 ++++++++++++---
 6 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 04335c3b..e595d314 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,19 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- **Search Filter for Graphs**: Research and Deep QA graphs now support `search_filter` parameter to restrict searches to specific documents
+  - Set `state.search_filter` to a SQL WHERE clause (e.g., `"id IN ('doc1', 'doc2')"`) before running the graph
+  - Enables document-scoped research workflows
+  - CLI: `haiku-rag research "question" --filter "uri LIKE '%paper%'"`
+  - CLI: `haiku-rag ask "question" --filter "title = 'My Doc'"`
+  - Python: `client.ask(question, filter="...")` and `agent.answer(question, filter="...")`
+- **AG-UI Research Example**: Added bidirectional state demonstration with document filter
+  - New `/api/documents` endpoint to list available documents
+  - Frontend document selector component with search and multi-select
+  - Demonstrates client-to-server state flow via AG-UI protocol
+
 ## [0.20.0] - 2025-12-10
 
 ### Added
diff --git a/docs/cli.md b/docs/cli.md
index a1207ef4..3e0edf0f 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -149,9 +149,9 @@ Show verbose output with deep QA:
 haiku-rag ask "What are the main features and architecture of haiku.rag?" --deep --verbose
 ```
 
-Filter to specific documents with deep QA:
+Filter to specific documents:
 ```bash
-haiku-rag ask "What are the main findings?" --deep --filter "uri LIKE '%paper%'"
+haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'"
 ```
 
 The QA agent searches your documents for relevant information and provides a comprehensive answer. When available, citations use the document title; otherwise they fall back to the URI.
@@ -161,7 +161,7 @@ Flags:
 - `--cite`: Include citations showing which documents were used
 - `--deep`: Decompose the question into sub-questions answered in parallel before synthesizing a final answer
 - `--verbose`: Show planning, searching, evaluation, and synthesis steps (only with `--deep`)
-- `--filter`: Restrict searches to documents matching the filter (only with `--deep`)
+- `--filter`: Restrict searches to documents matching the filter (see [Filtering Search Results](python.md#filtering-search-results))
 
 ## Research
 
diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py
index c92206f4..e6c95755 100644
--- a/haiku_rag_slim/haiku/rag/app.py
+++ b/haiku_rag_slim/haiku/rag/app.py
@@ -363,7 +363,7 @@ class HaikuRAGApp:
                         if cite:
                             citations = result.citations
                 else:
-                    answer, citations = await self.client.ask(question)
+                    answer, citations = await self.client.ask(question, filter=filter)
 
                 self.console.print(f"[bold blue]Question:[/bold blue] {question}")
                 self.console.print()
diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py
index 84216409..42f8fbd9 100644
--- a/haiku_rag_slim/haiku/rag/cli.py
+++ b/haiku_rag_slim/haiku/rag/cli.py
@@ -309,7 +309,7 @@ def ask(
         None,
         "--filter",
         "-f",
-        help="SQL WHERE clause to filter documents (only with --deep)",
+        help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
     ),
 ):
     app = create_app(db)
diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py
index afdc87af..53ceff30 100644
--- a/haiku_rag_slim/haiku/rag/client.py
+++ b/haiku_rag_slim/haiku/rag/client.py
@@ -1229,13 +1229,17 @@ class HaikuRAG:
         return final_results + passthrough
 
     async def ask(
-        self, question: str, system_prompt: str | None = None
+        self,
+        question: str,
+        system_prompt: str | None = None,
+        filter: str | None = None,
     ) -> "tuple[str, list[Citation]]":
         """Ask a question using the configured QA agent.
 
         Args:
             question: The question to ask.
             system_prompt: Optional custom system prompt for the QA agent.
+            filter: SQL WHERE clause to filter documents.
 
         Returns:
             Tuple of (answer text, list of resolved citations).
@@ -1243,7 +1247,7 @@ class HaikuRAG:
         from haiku.rag.qa import get_qa_agent
 
         qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt)
-        return await qa_agent.answer(question)
+        return await qa_agent.answer(question, filter=filter)
 
     async def visualize_chunk(self, chunk: Chunk) -> list:
         """Render page images with bounding box highlights for a chunk.
diff --git a/haiku_rag_slim/haiku/rag/qa/agent.py b/haiku_rag_slim/haiku/rag/qa/agent.py
index 5c37ffeb..03b6c0d5 100644
--- a/haiku_rag_slim/haiku/rag/qa/agent.py
+++ b/haiku_rag_slim/haiku/rag/qa/agent.py
@@ -15,6 +15,7 @@ class Dependencies(BaseModel):
     model_config = {"arbitrary_types_allowed": True}
     client: HaikuRAG
     search_results: list[SearchResult] = []
+    search_filter: str | None = None
 
 
 class QuestionAnswerAgent:
@@ -46,7 +47,9 @@ class QuestionAnswerAgent:
             Returns results with chunk IDs and relevance scores.
             Reference results by their chunk_id in cited_chunks.
             """
-            results = await ctx.deps.client.search(query, limit=limit)
+            results = await ctx.deps.client.search(
+                query, limit=limit, filter=ctx.deps.search_filter
+            )
             results = await ctx.deps.client.expand_context(results)
             # Store results for citation resolution
             ctx.deps.search_results = results
@@ -54,13 +57,19 @@ class QuestionAnswerAgent:
             parts = [r.format_for_agent() for r in results]
             return "\n\n".join(parts) if parts else "No results found."
 
-    async def answer(self, question: str) -> tuple[str, list[Citation]]:
+    async def answer(
+        self, question: str, filter: str | None = None
+    ) -> tuple[str, list[Citation]]:
         """Answer a question using the RAG system.
 
+        Args:
+            question: The question to answer
+            filter: SQL WHERE clause to filter documents
+
         Returns:
             Tuple of (answer text, list of resolved citations)
         """
-        deps = Dependencies(client=self._client)
+        deps = Dependencies(client=self._client, search_filter=filter)
         result = await self._agent.run(question, deps=deps)
         citations = resolve_citations(result.output.cited_chunks, deps.search_results)
         return result.output.answer, citations