From 7d57fd6c97769a08e88981d5d6d212af467db79d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 30 Oct 2025 13:03:23 +0200 Subject: [PATCH 1/3] 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 --- src/haiku/rag/client.py | 11 +- src/haiku/rag/store/repositories/chunk.py | 55 ++++++- tests/test_filter.py | 181 ++++++++++++++++++++++ 3 files changed, 237 insertions(+), 10 deletions(-) create mode 100644 tests/test_filter.py diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index b9d4a1b5..76e639a7 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -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 diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index b92c96de..397f4117 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -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 = [] diff --git a/tests/test_filter.py b/tests/test_filter.py new file mode 100644 index 00000000..8e8664a8 --- /dev/null +++ b/tests/test_filter.py @@ -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 From f3565afdf773f8e6e8889649d2e4a4a93394debe Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 30 Oct 2025 13:04:25 +0200 Subject: [PATCH 2/3] Cli accepts filter --- src/haiku/rag/app.py | 4 ++-- src/haiku/rag/cli.py | 8 +++++++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 51e63751..d5926312 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -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 diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 1d778487..1cc04154 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -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") From f0f844d21d6b2de7cc3f44e752c22a68b7f3f849 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 30 Oct 2025 13:07:13 +0200 Subject: [PATCH 3/3] Update documentation --- README.md | 3 +++ docs/cli.md | 12 ++++++++++++ docs/python.md | 42 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 57 insertions(+) diff --git a/README.md b/README.md index da0f2dff..c2ab5903 100644 --- a/README.md +++ b/README.md @@ -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?" diff --git a/docs/cli.md b/docs/cli.md index 2b2e88ec..bddfdf31 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -78,6 +78,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: diff --git a/docs/python.md b/docs/python.md index d462b625..3bd93d00 100644 --- a/docs/python.md +++ b/docs/python.md @@ -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: