danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add-src in cli & create_document_from_source will recursively add directories

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-16 09:55:50 +03:00
parent 5763f89438
commit 1cab0339ad
No known key found for this signature in database
7 changed files with 138 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
pyproject.toml
View file

@ -63,7 +63,7 @@ dev = [
"mkdocs-material>=9.6.14",
"pydantic-evals>=1.0.8",
"pre-commit>=4.2.0",
"pyright>=1.1.405",
"pyright>=1.1.406",
"pytest>=8.4.2",
"pytest-asyncio>=1.2.0",
"pytest-cov>=7.0.0",

17
src/haiku/rag/app.py
View file

@ -160,13 +160,20 @@ class HaikuRAGApp:
self, source: str, title: str | None = None, metadata: dict | None = None
):
async with HaikuRAG(db_path=self.db_path) as self.client:
doc = await self.client.create_document_from_source(
result = await self.client.create_document_from_source(
source, title=title, metadata=metadata
)
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[bold green]Document {doc.id} added successfully.[/bold green]"
)
if isinstance(result, list):
for doc in result:
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[bold green]{len(result)} documents added successfully.[/bold green]"
)
else:
self._rich_print_document(result, truncate=True)
self.console.print(
f"[bold green]Document {result.id} added successfully.[/bold green]"
)
async def get_document(self, doc_id: str):
async with HaikuRAG(db_path=self.db_path) as self.client:

4
src/haiku/rag/cli.py
View file

@ -128,10 +128,10 @@ def add_document_text(
asyncio.run(app.add_document_from_text(text=text, metadata=metadata or None))
@cli.command("add-src", help="Add a document from a file path or URL")
@cli.command("add-src", help="Add a document from a file path, directory, or URL")
def add_document_src(
source: str = typer.Argument(
help="The file path or URL of the document to add",
help="The file path, directory, or URL of the document(s) to add",
),
title: str | None = typer.Option(
None,

51
src/haiku/rag/client.py
View file

@ -106,8 +106,8 @@ class HaikuRAG:
async def create_document_from_source(
self, source: str | Path, title: str | None = None, metadata: dict | None = None
) -> Document:
"""Create or update a document from a file path or URL.
) -> Document | list[Document]:
"""Create or update document(s) from a file path, directory, or URL.
Checks if a document with the same URI already exists:
- If MD5 is unchanged, returns existing document
@ -115,11 +115,13 @@ class HaikuRAG:
- If no document exists, creates a new one
Args:
source: File path (as string or Path) or URL to parse
source: File path, directory (as string or Path), or URL to parse
title: Optional title (only used for single files, not directories)
metadata: Optional metadata dictionary
Returns:
Document instance (created, updated, or existing)
Document instance (created, updated, or existing) for single files/URLs
List of Document instances for directories
Raises:
ValueError: If the file/URL cannot be parsed or doesn't exist
@ -142,6 +144,45 @@ class HaikuRAG:
else:
# Handle as regular file path
source_path = Path(source) if isinstance(source, str) else source
# Handle directories
if source_path.is_dir():
documents = []
supported_extensions = set(FileReader.extensions)
for file_path in source_path.rglob("*"):
if (
file_path.is_file()
and file_path.suffix.lower() in supported_extensions
):
doc = await self._create_document_from_file(
file_path, title=None, metadata=metadata
)
documents.append(doc)
return documents
# Handle single file
return await self._create_document_from_file(
source_path, title=title, metadata=metadata
)
async def _create_document_from_file(
self, source_path: Path, title: str | None = None, metadata: dict | None = None
) -> Document:
"""Create or update a document from a single file path.
Args:
source_path: Path to the file
title: Optional title
metadata: Optional metadata dictionary
Returns:
Document instance (created, updated, or existing)
Raises:
ValueError: If the file cannot be parsed or doesn't exist
"""
metadata = metadata or {}
if source_path.suffix.lower() not in FileReader.extensions:
raise ValueError(f"Unsupported file extension: {source_path.suffix}")
@ -592,6 +633,8 @@ class HaikuRAG:
new_doc = await self.create_document_from_source(
source=doc.uri, metadata=doc.metadata or {}
)
# URIs always point to single files/URLs, never directories
assert isinstance(new_doc, Document)
assert new_doc.id is not None, (
"New document ID should not be None"
)

23
tests/test_cli.py
View file

@ -372,3 +372,26 @@ def test_info():
assert result.exit_code == 0
mock_app_instance.info.assert_called_once()
def test_add_document_src_directory(tmp_path):
"""Test adding documents from a directory recursively."""
with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock()
mock_app_instance.add_document_from_source = AsyncMock()
mock_app.return_value = mock_app_instance
test_dir = tmp_path / "test_docs"
test_dir.mkdir()
(test_dir / "doc1.txt").write_text("doc1")
(test_dir / "doc2.md").write_text("doc2")
subdir = test_dir / "subdir"
subdir.mkdir()
(subdir / "doc3.pdf").write_text("doc3")
result = runner.invoke(cli, ["add-src", str(test_dir)])
assert result.exit_code == 0
mock_app_instance.add_document_from_source.assert_called_once()
call_args = mock_app_instance.add_document_from_source.call_args
assert call_args[1]["source"] == str(test_dir)

52
tests/test_client.py
View file

@ -9,6 +9,7 @@ from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
@pytest.mark.asyncio
@ -88,6 +89,7 @@ async def test_client_create_document_from_source(temp_db_path):
# Test create_document_from_source with Path
doc = await client.create_document_from_source(source=temp_path)
assert isinstance(doc, Document)
assert doc.id is not None
assert doc.content == test_content
@ -98,6 +100,7 @@ async def test_client_create_document_from_source(temp_db_path):
# Test create_document_from_source with string path
doc2 = await client.create_document_from_source(source=str(temp_path))
assert isinstance(doc2, Document)
assert doc2.id is not None
assert doc2.content == test_content
@ -118,6 +121,7 @@ async def test_client_create_document_from_source_with_title(temp_db_path):
doc = await client.create_document_from_source(
source=temp_path, title="My Doc"
)
assert isinstance(doc, Document)
assert doc.id is not None
assert doc.title == "My Doc"
@ -131,10 +135,12 @@ async def test_client_update_title_noop_behavior(temp_db_path):
temp_path.write_text("Original content")
doc1 = await client.create_document_from_source(temp_path, title="Title A")
assert isinstance(doc1, Document)
assert doc1.id is not None
# Re-add with same content but new title
doc2 = await client.create_document_from_source(temp_path, title="Title B")
assert isinstance(doc2, Document)
assert doc2.id == doc1.id
# Fetch and verify title updated
got = await client.get_document_by_id(doc1.id)
@ -169,6 +175,41 @@ async def test_client_create_document_from_source_nonexistent(temp_db_path):
await client.create_document_from_source(non_existent_path)
@pytest.mark.asyncio
async def test_client_create_document_from_directory(temp_db_path):
"""Test creating documents from a directory recursively."""
async with HaikuRAG(temp_db_path) as client:
with tempfile.TemporaryDirectory() as temp_dir:
test_dir = Path(temp_dir) / "test_docs"
test_dir.mkdir()
(test_dir / "doc1.txt").write_text("Content of doc1")
(test_dir / "doc2.md").write_text("# Content of doc2")
subdir = test_dir / "subdir"
subdir.mkdir()
(subdir / "doc3.py").write_text("print('hello')")
(test_dir / "unsupported.xyz").write_text("unsupported file")
result = await client.create_document_from_source(test_dir)
assert isinstance(result, list)
assert len(result) == 3
for doc in result:
assert doc.id is not None
assert doc.uri is not None
assert "md5" in doc.metadata
assert "contentType" in doc.metadata
uris = [doc.uri for doc in result if doc.uri]
assert any("doc1.txt" in uri for uri in uris)
assert any("doc2.md" in uri for uri in uris)
assert any("doc3.py" in uri for uri in uris)
assert not any("unsupported.xyz" in uri for uri in uris)
@pytest.mark.asyncio
async def test_client_create_document_from_url(temp_db_path):
"""Test creating a document from a URL."""
@ -183,6 +224,7 @@ async def test_client_create_document_from_url(temp_db_path):
doc = await client.create_document_from_source(
source="https://example.com/test.html", metadata={"source_type": "web"}
)
assert isinstance(doc, Document)
assert doc.id is not None
assert "Test Page" in doc.content
@ -212,6 +254,7 @@ async def test_client_create_document_from_url_with_different_content_types(
doc = await client.create_document_from_source(
"https://api.example.com/data.json"
)
assert isinstance(doc, Document)
assert doc.id is not None
assert "Test JSON" in doc.content
@ -230,6 +273,7 @@ async def test_client_create_document_from_url_with_different_content_types(
doc = await client.create_document_from_source(
"https://example.com/readme.txt"
)
assert isinstance(doc, Document)
assert doc.id is not None
assert doc.content == "This is plain text content from a URL."
@ -333,6 +377,7 @@ async def test_client_metadata_content_type_and_md5(temp_db_path):
temp_path.write_text(test_content)
doc = await client.create_document_from_source(temp_path)
assert isinstance(doc, Document)
assert doc.metadata["contentType"] == "text/plain"
assert doc.metadata["md5"] == expected_md5
@ -346,6 +391,7 @@ async def test_client_metadata_content_type_and_md5(temp_db_path):
url_doc = await client.create_document_from_source(
"https://example.com/test.txt"
)
assert isinstance(url_doc, Document)
assert url_doc.metadata["contentType"] == "text/plain"
assert url_doc.metadata["md5"] == expected_md5
@ -363,12 +409,14 @@ async def test_client_create_update_no_op_behavior(temp_db_path):
# First call - should create new document
doc1 = await client.create_document_from_source(temp_path)
assert isinstance(doc1, Document)
assert doc1.id is not None
assert doc1.content == test_content
original_id = doc1.id
# Second call with same content - should return existing document (no-op)
doc2 = await client.create_document_from_source(temp_path)
assert isinstance(doc2, Document)
assert doc2.id == original_id # Same document
assert doc2.content == test_content
@ -378,6 +426,7 @@ async def test_client_create_update_no_op_behavior(temp_db_path):
# Third call with changed content - should update existing document
doc3 = await client.create_document_from_source(temp_path)
assert isinstance(doc3, Document)
assert doc3.id == original_id # Same document ID
assert doc3.content == updated_content # Updated content
@ -404,11 +453,13 @@ async def test_client_url_create_update_no_op_behavior(temp_db_path):
with patch("httpx.AsyncClient.get", return_value=mock_response1):
# First call - should create new document
doc1 = await client.create_document_from_source(url)
assert isinstance(doc1, Document)
assert doc1.id is not None
original_id = doc1.id
# Second call with same content - should return existing document (no-op)
doc2 = await client.create_document_from_source(url)
assert isinstance(doc2, Document)
assert doc2.id == original_id # Same document
mock_response2 = AsyncMock()
@ -419,6 +470,7 @@ async def test_client_url_create_update_no_op_behavior(temp_db_path):
with patch("httpx.AsyncClient.get", return_value=mock_response2):
# Third call with changed content - should update existing document
doc3 = await client.create_document_from_source(url)
assert isinstance(doc3, Document)
assert doc3.id == original_id # Same document ID
assert doc3.content == updated_content.decode() # Updated content

2
uv.lock
View file

@ -1194,7 +1194,7 @@ dev = [
{ name = "mkdocs-material", specifier = ">=9.6.14" },
{ name = "pre-commit", specifier = ">=4.2.0" },
{ name = "pydantic-evals", specifier = ">=1.0.8" },
{ name = "pyright", specifier = ">=1.1.405" },
{ name = "pyright", specifier = ">=1.1.406" },
{ name = "pytest", specifier = ">=8.4.2" },
{ name = "pytest-asyncio", specifier = ">=1.2.0" },
{ name = "pytest-cov", specifier = ">=7.0.0" },