From 7d57fd6c97769a08e88981d5d6d212af467db79d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 30 Oct 2025 13:03:23 +0200 Subject: [PATCH] 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