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

Add filter param to search(). When filtering, we search in two steps: First filter documents, then search on the chunks with document_id IN the filtered documents

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-30 13:03:23 +02:00
parent f9b6894974
commit 7d57fd6c97
No known key found for this signature in database
3 changed files with 237 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/haiku/rag/client.py
View file

@ -424,7 +424,11 @@ class HaikuRAG:
return await self.document_repository.list_all(limit=limit, offset=offset)
async def search(
self, query: str, limit: int = 5, search_type: str = "hybrid"
self,
query: str,
limit: int = 5,
search_type: str = "hybrid",
filter: str | None = None,
) -> list[tuple[Chunk, float]]:
"""Search for relevant chunks using the specified search method with optional reranking.
@ -432,6 +436,7 @@ class HaikuRAG:
query: The search query string.
limit: Maximum number of results to return.
search_type: Type of search - "vector", "fts", or "hybrid" (default).
filter: Optional SQL WHERE clause to filter documents before searching chunks.
Returns:
List of (chunk, score) tuples ordered by relevance.
@ -441,12 +446,12 @@ class HaikuRAG:
if reranker is None:
# No reranking - return direct search results
return await self.chunk_repository.search(query, limit, search_type)
return await self.chunk_repository.search(query, limit, search_type, filter)
# Get more initial results (3X) for reranking
search_limit = limit * 3
search_results = await self.chunk_repository.search(
query, search_limit, search_type
query, search_limit, search_type, filter
)
# Apply reranking

55
src/haiku/rag/store/repositories/chunk.py
View file

@ -230,7 +230,11 @@ class ChunkRepository:
return True
async def search(
self, query: str, limit: int = 5, search_type: str = "hybrid"
self,
query: str,
limit: int = 5,
search_type: str = "hybrid",
filter: str | None = None,
) -> list[tuple[Chunk, float]]:
"""Search for relevant chunks using the specified search method.
@ -238,6 +242,7 @@ class ChunkRepository:
query: The search query string.
limit: Maximum number of results to return.
search_type: Type of search - "vector", "fts", or "hybrid" (default).
filter: Optional SQL WHERE clause to filter documents before searching chunks.
Returns:
List of (chunk, score) tuples ordered by relevance.
@ -245,19 +250,42 @@ class ChunkRepository:
if not query.strip():
return []
chunk_where_clause = None
if filter:
# We perform filtering as a two-step process, first filtering documents, then
# filtering chunks based on those document IDs.
# This is because LanceDB does not support joins directly in search queries.
matching_doc_ids = self._get_filtered_document_ids(filter)
if not matching_doc_ids:
return []
# Build WHERE clause for chunks table
# Use IN clause with document IDs
id_list = "', '".join(matching_doc_ids)
chunk_where_clause = f"document_id IN ('{id_list}')"
if search_type == "vector":
query_embedding = await self.embedder.embed(query)
results = self.store.chunks_table.search(
query_embedding, query_type="vector", vector_column_name="vector"
).limit(limit)
)
if chunk_where_clause:
results = results.where(chunk_where_clause)
results = results.limit(limit)
return await self._process_search_results(results)
elif search_type == "fts":
results = self.store.chunks_table.search(query, query_type="fts").limit(
limit
)
results = self.store.chunks_table.search(query, query_type="fts")
if chunk_where_clause:
results = results.where(chunk_where_clause)
results = results.limit(limit)
return await self._process_search_results(results)
else: # hybrid (default)
@ -267,9 +295,13 @@ class ChunkRepository:
reranker = RRFReranker()
# Perform native hybrid search with RRF reranking
results = self.store.chunks_table.search(query_type="hybrid")
if chunk_where_clause:
results = results.where(chunk_where_clause)
results = (
self.store.chunks_table.search(query_type="hybrid")
.vector(query_embedding)
results.vector(query_embedding)
.text(query)
.rerank(reranker)
.limit(limit)
@ -332,6 +364,15 @@ class ChunkRepository:
return adjacent_chunks
def _get_filtered_document_ids(self, filter: str) -> list[str]:
"""Query documents table with filter and return matching document IDs."""
filtered_docs = (
self.store.documents_table.search()
.where(filter)
.to_pydantic(DocumentRecord)
)
return [doc.id for doc in filtered_docs]
async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]:
"""Process search results into chunks with document info and scores."""
chunks_with_scores = []

181
tests/test_filter.py Normal file
View file

@ -0,0 +1,181 @@
import pytest
from haiku.rag.client import HaikuRAG
@pytest.mark.asyncio
async def test_search_with_uri_filter(temp_db_path):
"""Test filtering by document URI."""
async with HaikuRAG(db_path=temp_db_path) as client:
# Add multiple test documents
await client.create_document(
content="Python tutorial content",
uri="https://example.com/python.html",
title="Python Guide",
)
await client.create_document(
content=" Java tutorial content",
uri="https://other.com/java.html",
title="Java Guide",
)
# Filter by URI pattern
results = await client.search(
"tutorial", limit=5, filter="uri LIKE '%example.com%'"
)
assert len(results) > 0
for chunk, _ in results:
assert chunk.document_uri is not None
assert "example.com" in chunk.document_uri
# Filter by exact URI
results = await client.search(
"tutorial", limit=5, filter="uri = 'https://other.com/java.html'"
)
assert len(results) > 0
for chunk, _ in results:
assert chunk.document_uri == "https://other.com/java.html"
@pytest.mark.asyncio
async def test_search_with_title_filter(temp_db_path):
"""Test filtering by document title."""
async with HaikuRAG(db_path=temp_db_path) as client:
# Add test documents
await client.create_document(
content="Programming content",
uri="https://example.com/doc1.html",
title="Python Programming",
)
await client.create_document(
content="Programming content",
uri="https://example.com/doc2.html",
title="Java Programming",
)
# Filter by title pattern
results = await client.search(
"programming", limit=5, filter="title LIKE '%Python%'"
)
assert len(results) > 0
for chunk, _ in results:
assert chunk.document_title is not None
assert "Python" in chunk.document_title
@pytest.mark.asyncio
async def test_search_with_combined_filters(temp_db_path):
"""Test filtering with AND/OR conditions."""
async with HaikuRAG(db_path=temp_db_path) as client:
# Add test documents
await client.create_document(
content="Content about AI",
uri="https://arxiv.org/paper1.pdf",
title="Machine Learning Paper",
)
await client.create_document(
content="Content about AI",
uri="https://example.com/tutorial.html",
title="AI Tutorial",
)
await client.create_document(
content="Content about AI",
uri="https://arxiv.org/paper2.pdf",
title="Deep Learning Paper",
)
# Filter with AND condition
results = await client.search(
"AI", limit=5, filter="uri LIKE '%arxiv%' AND title LIKE '%Machine%'"
)
assert len(results) > 0
for chunk, _ in results:
assert chunk.document_uri is not None
assert chunk.document_title is not None
assert "arxiv" in chunk.document_uri
assert "Machine" in chunk.document_title
# Filter with OR condition (if supported)
results = await client.search(
"AI", limit=5, filter="title LIKE '%Tutorial%' OR title LIKE '%Deep%'"
)
assert len(results) > 0
@pytest.mark.asyncio
async def test_search_with_no_matching_filter(temp_db_path):
"""Test that search returns empty results when filter matches no documents."""
async with HaikuRAG(db_path=temp_db_path) as client:
# Add a test document
await client.create_document(
content="Test content",
uri="https://example.com/test.html",
title="Test Document",
)
# Search with non-matching filter
results = await client.search(
"test", limit=5, filter="uri = 'https://nonexistent.com/doc.html'"
)
assert len(results) == 0
@pytest.mark.asyncio
async def test_search_with_invalid_filter(temp_db_path):
"""Test that invalid filter syntax raises an appropriate error."""
async with HaikuRAG(db_path=temp_db_path) as client:
# Add a test document
await client.create_document(
content="Test content",
uri="https://example.com/test.html",
title="Test Document",
)
# Invalid filter should raise RuntimeError
with pytest.raises(RuntimeError, match="No field named invalid"):
await client.search("test", limit=5, filter="invalid = 'value'")
@pytest.mark.asyncio
async def test_search_filter_with_all_search_types(temp_db_path):
"""Test that filtering works with all search types (vector, fts, hybrid)."""
async with HaikuRAG(db_path=temp_db_path) as client:
await client.create_document(
content="Machine learning is a subset of artificial intelligence",
uri="https://ai.example.com/ml.html",
title="ML Guide",
)
await client.create_document(
content="Deep learning uses neural networks",
uri="https://other.com/dl.html",
title="DL Guide",
)
# Test vector search with filter
results = await client.search(
"machine learning",
limit=5,
search_type="vector",
filter="uri LIKE '%ai.example%'",
)
assert len(results) > 0
for chunk, _ in results:
assert chunk.document_uri is not None
assert "ai.example" in chunk.document_uri
# Test FTS search with filter
results = await client.search(
"learning", limit=5, search_type="fts", filter="title = 'ML Guide'"
)
assert all(chunk.document_title == "ML Guide" for chunk, _ in results)
# Test hybrid search with filter (default)
results = await client.search(
"neural networks",
limit=5,
search_type="hybrid",
filter="uri LIKE '%other.com%'",
)
for chunk, _ in results:
assert chunk.document_uri is not None
assert "other.com" in chunk.document_uri