Merge pull request #125 from ggozad/feat/filtering

Filtered searches

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?"
 

docs/cli.md

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

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:

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

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")

src/haiku/rag/client.py

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

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