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

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

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(
+            results = self.store.chunks_table.search(query, query_type="fts")
-                limit
+
-            )
+            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")
+                results.vector(query_embedding)
-                .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 = []

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