diff --git a/docs/cli.md b/docs/cli.md index 2c5b543a..c1c650e4 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -22,6 +22,18 @@ The `haiku-rag` CLI provides complete document management functionality. haiku-rag list ``` +Filter documents by properties: +```bash +# Filter by URI pattern +haiku-rag list --filter "uri LIKE '%arxiv%'" + +# Filter by exact title +haiku-rag list --filter "title = 'My Document'" + +# Combine multiple conditions +haiku-rag list --filter "uri LIKE '%.pdf' AND title LIKE '%paper%'" +``` + ### Add Documents From text: diff --git a/docs/mcp.md b/docs/mcp.md index ec2d02c9..1fe06f51 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -10,7 +10,7 @@ The MCP server exposes `haiku.rag` as MCP tools for compatible MCP clients. - `add_document_from_url` - Add documents from URLs - `add_document_from_text` - Add documents from raw text content - `get_document` - Retrieve specific documents by ID -- `list_documents` - List all documents with pagination +- `list_documents` - List all documents with pagination and optional filtering - `delete_document` - Delete documents by ID ### Search diff --git a/docs/python.md b/docs/python.md index bf1bdbce..cd2c58db 100644 --- a/docs/python.md +++ b/docs/python.md @@ -84,6 +84,21 @@ List all documents: docs = await client.list_documents(limit=10, offset=0) ``` +Filter documents by properties: +```python +# Filter by URI pattern +docs = await client.list_documents(filter="uri LIKE '%arxiv%'") + +# Filter by exact title +docs = await client.list_documents(filter="title = 'My Document'") + +# Combine multiple conditions +docs = await client.list_documents( + limit=10, + filter="uri LIKE '%.pdf' AND title LIKE '%paper%'" +) +``` + ### Updating Documents ```python diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index d5926312..8e95ae7c 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -139,9 +139,9 @@ class HaikuRAGApp: f" [repr.attrib_name]docling[/repr.attrib_name]: {docling_version}" ) - async def list_documents(self): + async def list_documents(self, filter: str | None = None): async with HaikuRAG(db_path=self.db_path) as self.client: - documents = await self.client.list_documents() + documents = await self.client.list_documents(filter=filter) for doc in documents: self._rich_print_document(doc, truncate=True) diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index f1df1978..1d654717 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -92,11 +92,17 @@ def list_documents( "--db", help="Path to the LanceDB database file", ), + filter: str | None = typer.Option( + None, + "--filter", + "-f", + help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", + ), ): from haiku.rag.app import HaikuRAGApp app = HaikuRAGApp(db_path=db) - asyncio.run(app.list_documents()) + asyncio.run(app.list_documents(filter=filter)) def _parse_meta_options(meta: list[str] | None) -> dict[str, Any]: diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 2cc76b00..5c399e6a 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -409,18 +409,24 @@ class HaikuRAG: return await self.document_repository.delete(document_id) async def list_documents( - self, limit: int | None = None, offset: int | None = None + self, + limit: int | None = None, + offset: int | None = None, + filter: str | None = None, ) -> list[Document]: - """List all documents with optional pagination. + """List all documents with optional pagination and filtering. Args: limit: Maximum number of documents to return. offset: Number of documents to skip. + filter: Optional SQL WHERE clause to filter documents. Returns: - List of Document instances. + List of Document instances matching the criteria. """ - return await self.document_repository.list_all(limit=limit, offset=offset) + return await self.document_repository.list_all( + limit=limit, offset=offset, filter=filter + ) async def search( self, diff --git a/src/haiku/rag/mcp.py b/src/haiku/rag/mcp.py index 7248d744..61754759 100644 --- a/src/haiku/rag/mcp.py +++ b/src/haiku/rag/mcp.py @@ -130,12 +130,23 @@ def create_mcp_server(db_path: Path) -> FastMCP: @mcp.tool() async def list_documents( - limit: int | None = None, offset: int | None = None + limit: int | None = None, + offset: int | None = None, + filter: str | None = None, ) -> list[DocumentResult]: - """List all documents with optional pagination.""" + """List all documents with optional pagination and filtering. + + Args: + limit: Maximum number of documents to return. + offset: Number of documents to skip. + filter: Optional SQL WHERE clause to filter documents. + + Returns: + List of DocumentResult instances matching the criteria. + """ try: async with HaikuRAG(db_path) as rag: - documents = await rag.list_documents(limit, offset) + documents = await rag.list_documents(limit, offset, filter) return [ DocumentResult( diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index 70643023..74199dc0 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/src/haiku/rag/store/repositories/document.py @@ -123,11 +123,25 @@ class DocumentRepository: return True async def list_all( - self, limit: int | None = None, offset: int | None = None + self, + limit: int | None = None, + offset: int | None = None, + filter: str | None = None, ) -> list[Document]: - """List all documents with optional pagination.""" + """List all documents with optional pagination and filtering. + + Args: + limit: Maximum number of documents to return. + offset: Number of documents to skip. + filter: Optional SQL WHERE clause to filter documents. + + Returns: + List of Document instances matching the criteria. + """ query = self.store.documents_table.search() + if filter is not None: + query = query.where(filter) if offset is not None: query = query.offset(offset) if limit is not None: diff --git a/tests/test_document.py b/tests/test_document.py index 61051505..46dd9c50 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -105,3 +105,46 @@ async def test_document_repository_crud(qa_corpus: Dataset, temp_db_path): assert retrieved_document is None store.close() + + +@pytest.mark.asyncio +async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path): + """Test listing documents with filter clause.""" + store = Store(temp_db_path) + doc_repo = DocumentRepository(store) + + first_doc = qa_corpus[0] + document_text = first_doc["document_extracted"] + + doc1 = Document( + content=document_text, + uri="https://example.com/doc1.txt", + metadata={"source": "test", "category": "A"}, + ) + doc2 = Document( + content=document_text, + uri="https://arxiv.org/paper.pdf", + metadata={"source": "test", "category": "B"}, + ) + doc3 = Document( + content=document_text, + uri="https://example.com/doc3.txt", + metadata={"source": "test", "category": "A"}, + ) + + created_doc1 = await doc_repo.create(doc1) + created_doc2 = await doc_repo.create(doc2) + created_doc3 = await doc_repo.create(doc3) + + all_documents = await doc_repo.list_all() + assert len(all_documents) == 3 + + arxiv_documents = await doc_repo.list_all(filter="uri LIKE '%arxiv%'") + assert len(arxiv_documents) == 1 + assert arxiv_documents[0].id == created_doc2.id + + example_documents = await doc_repo.list_all(filter="uri LIKE '%example.com%'") + assert len(example_documents) == 2 + assert {doc.id for doc in example_documents} == {created_doc1.id, created_doc3.id} + + store.close() diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 02b1b275..e6a303c6 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -215,7 +215,7 @@ async def test_mcp_list_documents(): assert len(result) == 2 assert result[0].id == "doc1" assert result[1].id == "doc2" - mock_rag.list_documents.assert_called_once_with(10, 0) + mock_rag.list_documents.assert_called_once_with(10, 0, None) @pytest.mark.asyncio