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

Merge pull request #125 from ggozad/feat/filtering

Browse source 
Filtered searches
This commit is contained in:
Yiorgis Gozadinos 2025-10-30 14:28:42 +02:00 committed by GitHub
commit 5c7c79397c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 303 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
README.md
View file

@ -38,6 +38,9 @@ haiku-rag add-src document.pdf --meta source=manual
# Search
haiku-rag search "query"
# Search with filters
haiku-rag search "query" --filter "uri LIKE '%.pdf' AND title LIKE '%paper%'"
# Ask questions
haiku-rag ask "Who is the author of haiku.rag?"

12
docs/cli.md
View file

@ -86,6 +86,18 @@ With options:
haiku-rag search "python programming" --limit 10
```
With filters (filter by document properties):
```bash
# Filter by URI pattern
haiku-rag search "neural networks" --filter "uri LIKE '%arxiv%'"
# Filter by exact title
haiku-rag search "transformers" --filter "title = 'Deep Learning Guide'"
# Combine multiple conditions
haiku-rag search "AI" --filter "uri LIKE '%.pdf' AND title LIKE '%paper%'"
```
## Question Answering
Ask questions about your documents:

42
docs/python.md
View file

@ -168,6 +168,48 @@ for chunk, relevance_score in results:
print(f"Document metadata: {chunk.document_meta}")
```
### Filtering Search Results
Filter search results to only include chunks from documents matching specific criteria:
```python
# Filter by document URI pattern
results = await client.search(
query="machine learning",
limit=5,
filter="uri LIKE '%arxiv%'"
)
# Filter by exact document title
results = await client.search(
query="neural networks",
limit=5,
filter="title = 'Deep Learning Guide'"
)
# Combine multiple filter conditions
results = await client.search(
query="AI research",
limit=5,
filter="uri LIKE '%.pdf' AND title LIKE '%paper%'"
)
# Filter with any search type
results = await client.search(
query="transformers",
limit=5,
search_type="vector",
filter="uri LIKE '%huggingface%'"
)
```
**Note:** Filters apply to document properties only. Available columns for filtering:
- `id` - Document ID
- `uri` - Document URI/URL
- `title` - Document title (if set)
- `created_at`, `updated_at` - Timestamps
- `metadata` - Document metadata (as string, use LIKE for pattern matching)
### Expanding Search Context
Expand search results with adjacent chunks for more complete context:

4
src/haiku/rag/app.py
View file

@ -192,9 +192,9 @@ class HaikuRAGApp:
f"[yellow]Document with id {doc_id} not found.[/yellow]"
)
async def search(self, query: str, limit: int = 5):
async def search(self, query: str, limit: int = 5, filter: str | None = None):
async with HaikuRAG(db_path=self.db_path) as self.client:
results = await self.client.search(query, limit=limit)
results = await self.client.search(query, limit=limit, filter=filter)
if not results:
self.console.print("[yellow]No results found.[/yellow]")
return

8
src/haiku/rag/cli.py
View file

@ -221,6 +221,12 @@ def search(
"-l",
help="Maximum number of results to return",
),
filter: str | None = typer.Option(
None,
"--filter",
"-f",
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
),
db: Path = typer.Option(
Config.storage.data_dir / "haiku.rag.lancedb",
"--db",
@ -230,7 +236,7 @@ def search(
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.search(query=query, limit=limit))
asyncio.run(app.search(query=query, limit=limit, filter=filter))
@cli.command("ask", help="Ask a question using the QA agent")

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

@ -423,7 +423,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.
@ -431,6 +435,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.
@ -440,12 +445,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