danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #189 from ggozad/feat/bidirectional-agui

Browse source 
Filtering in graphs & cli
This commit is contained in:
Yiorgis Gozadinos 2025-12-11 15:33:59 +02:00 committed by GitHub
commit 1f7f27e68f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 693 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
CHANGELOG.md
View file

@ -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

14
docs/agents.md
View file

@ -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

23
docs/cli.md
View file

@ -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:
```bash
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.
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 (see [Filtering Search Results](python.md#filtering-search-results))
## 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.

2
examples/ag-ui-research/backend/agent.py
View file

@ -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,

27
examples/ag-ui-research/backend/main.py
View file

@ -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 (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)
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(
@ -189,6 +200,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 +252,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"]),
],

17
examples/ag-ui-research/frontend/components/Agent.tsx
View file

@ -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>

352
examples/ag-ui-research/frontend/components/DocumentSelector.tsx Normal file
View file

@ -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>
);
}

11
haiku_rag_slim/haiku/rag/app.py
View file

@ -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:
@ -360,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()
@ -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:

18
haiku_rag_slim/haiku/rag/cli.py
View file

@ -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 (e.g., \"uri LIKE '%arxiv%'\")",
),
):
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")

8
haiku_rag_slim/haiku/rag/client.py
View file

@ -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.

15
haiku_rag_slim/haiku/rag/graph/common/nodes.py
View file

@ -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

3
haiku_rag_slim/haiku/rag/graph/deep_qa/state.py
View file

@ -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":

3
haiku_rag_slim/haiku/rag/graph/research/state.py
View file

@ -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(

15
haiku_rag_slim/haiku/rag/qa/agent.py
View file

@ -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

187
tests/graph/test_search_filter.py Normal file
View file

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