Add filtering to document listing in CLI/client
This commit is contained in:
parent
8fd5d169a8
commit
1f22f74e1a
10 changed files with 121 additions and 14 deletions
12
docs/cli.md
12
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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue