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

Handle title in client & CLI

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-22 11:08:08 +03:00
parent 312d29333c
commit 4077e3ecb2
No known key found for this signature in database
7 changed files with 143 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -39,9 +39,9 @@ class HaikuRAGApp:
f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
)
async def add_document_from_source(self, source: str):
async def add_document_from_source(self, source: str, title: str | None = None):
async with HaikuRAG(db_path=self.db_path) as self.client:
doc = await self.client.create_document_from_source(source)
doc = await self.client.create_document_from_source(source, title=title)
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"

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

@ -160,6 +160,11 @@ def add_document_src(
help="The file path or URL of the document to add",
autocompletion=complete_local_paths,
),
title: str | None = typer.Option(
None,
"--title",
help="Optional human-readable title to store with the document",
),
db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db",
@ -169,7 +174,7 @@ def add_document_src(
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.add_document_from_source(source=source))
asyncio.run(app.add_document_from_source(source=source, title=title))
@cli.command("get", help="Get and display a document by its ID")

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

@ -50,6 +50,7 @@ class HaikuRAG:
self,
docling_document,
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
chunks: list[Chunk] | None = None,
) -> Document:
@ -58,6 +59,7 @@ class HaikuRAG:
document = Document(
content=content,
uri=uri,
title=title,
metadata=metadata or {},
)
return await self.document_repository._create_with_docling(
@ -68,6 +70,7 @@ class HaikuRAG:
self,
content: str,
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
chunks: list[Chunk] | None = None,
) -> Document:
@ -88,6 +91,7 @@ class HaikuRAG:
document = Document(
content=content,
uri=uri,
title=title,
metadata=metadata or {},
)
return await self.document_repository._create_with_docling(
@ -95,7 +99,7 @@ class HaikuRAG:
)
async def create_document_from_source(
self, source: str | Path, metadata: dict = {}
self, source: str | Path, title: str | None = None, metadata: dict | None = None
) -> Document:
"""Create or update a document from a file path or URL.
@ -116,11 +120,16 @@ class HaikuRAG:
httpx.RequestError: If URL request fails
"""
# Normalize metadata
metadata = metadata or {}
# Check if it's a URL
source_str = str(source)
parsed_url = urlparse(source_str)
if parsed_url.scheme in ("http", "https"):
return await self._create_or_update_document_from_url(source_str, metadata)
return await self._create_or_update_document_from_url(
source_str, title=title, metadata=metadata
)
elif parsed_url.scheme == "file":
# Handle file:// URI by converting to path
source_path = Path(parsed_url.path)
@ -136,37 +145,51 @@ class HaikuRAG:
uri = source_path.absolute().as_uri()
md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest()
# Check if document already exists
existing_doc = await self.get_document_by_uri(uri)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
# MD5 unchanged, return existing document
return existing_doc
docling_document = FileReader.parse_file(source_path)
# Get content type from file extension
# Get content type from file extension (do before early return)
content_type, _ = mimetypes.guess_type(str(source_path))
if not content_type:
content_type = "application/octet-stream"
# Merge metadata with contentType and md5
metadata.update({"contentType": content_type, "md5": md5_hash})
# Check if document already exists
existing_doc = await self.get_document_by_uri(uri)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
# MD5 unchanged; update title/metadata if provided
updated = False
if title is not None and title != existing_doc.title:
existing_doc.title = title
updated = True
if metadata:
existing_doc.metadata = {**(existing_doc.metadata or {}), **metadata}
updated = True
if updated:
return await self.document_repository.update(existing_doc)
return existing_doc
# Parse file only when content changed or new document
docling_document = FileReader.parse_file(source_path)
if existing_doc:
# Update existing document
existing_doc.content = docling_document.export_to_markdown()
existing_doc.metadata = metadata
if title is not None:
existing_doc.title = title
return await self.document_repository._update_with_docling(
existing_doc, docling_document
)
else:
# Create new document using DoclingDocument
return await self._create_document_with_docling(
docling_document=docling_document, uri=uri, metadata=metadata
docling_document=docling_document,
uri=uri,
title=title,
metadata=metadata,
)
async def _create_or_update_document_from_url(
self, url: str, metadata: dict = {}
self, url: str, title: str | None = None, metadata: dict | None = None
) -> Document:
"""Create or update a document from a URL by downloading and parsing the content.
@ -186,20 +209,35 @@ class HaikuRAG:
ValueError: If the content cannot be parsed
httpx.RequestError: If URL request fails
"""
metadata = metadata or {}
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status()
md5_hash = hashlib.md5(response.content).hexdigest()
# Get content type early (used for potential no-op update)
content_type = response.headers.get("content-type", "").lower()
# Check if document already exists
existing_doc = await self.get_document_by_uri(url)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
# MD5 unchanged, return existing document
# MD5 unchanged; update title/metadata if provided
updated = False
if title is not None and title != existing_doc.title:
existing_doc.title = title
updated = True
metadata.update({"contentType": content_type, "md5": md5_hash})
if metadata:
existing_doc.metadata = {
**(existing_doc.metadata or {}),
**metadata,
}
updated = True
if updated:
return await self.document_repository.update(existing_doc)
return existing_doc
# Get content type to determine file extension
content_type = response.headers.get("content-type", "").lower()
file_extension = self._get_extension_from_content_type_or_url(
url, content_type
)
@ -226,12 +264,17 @@ class HaikuRAG:
if existing_doc:
existing_doc.content = docling_document.export_to_markdown()
existing_doc.metadata = metadata
if title is not None:
existing_doc.title = title
return await self.document_repository._update_with_docling(
existing_doc, docling_document
)
else:
return await self._create_document_with_docling(
docling_document=docling_document, uri=url, metadata=metadata
docling_document=docling_document,
uri=url,
title=title,
metadata=metadata,
)
def _get_extension_from_content_type_or_url(
@ -522,7 +565,7 @@ class HaikuRAG:
# Try to re-create from source (this creates the document with chunks)
new_doc = await self.create_document_from_source(
doc.uri, doc.metadata or {}
source=doc.uri, metadata=doc.metadata or {}
)
assert new_doc.id is not None, "New document ID should not be None"

24
src/haiku/rag/mcp.py
View file

@ -17,6 +17,7 @@ class DocumentResult(BaseModel):
id: str | None
content: str
uri: str | None = None
title: str | None = None
metadata: dict[str, Any] = {}
created_at: str
updated_at: str
@ -28,13 +29,15 @@ def create_mcp_server(db_path: Path) -> FastMCP:
@mcp.tool()
async def add_document_from_file(
file_path: str, metadata: dict[str, Any] | None = None
file_path: str,
metadata: dict[str, Any] | None = None,
title: str | None = None,
) -> str | None:
"""Add a document to the RAG system from a file path."""
try:
async with HaikuRAG(db_path) as rag:
document = await rag.create_document_from_source(
Path(file_path), metadata or {}
Path(file_path), title=title, metadata=metadata or {}
)
return document.id
except Exception:
@ -42,24 +45,31 @@ def create_mcp_server(db_path: Path) -> FastMCP:
@mcp.tool()
async def add_document_from_url(
url: str, metadata: dict[str, Any] | None = None
url: str, metadata: dict[str, Any] | None = None, title: str | None = None
) -> str | None:
"""Add a document to the RAG system from a URL."""
try:
async with HaikuRAG(db_path) as rag:
document = await rag.create_document_from_source(url, metadata or {})
document = await rag.create_document_from_source(
url, title=title, metadata=metadata or {}
)
return document.id
except Exception:
return None
@mcp.tool()
async def add_document_from_text(
content: str, uri: str | None = None, metadata: dict[str, Any] | None = None
content: str,
uri: str | None = None,
metadata: dict[str, Any] | None = None,
title: str | None = None,
) -> str | None:
"""Add a document to the RAG system from text content."""
try:
async with HaikuRAG(db_path) as rag:
document = await rag.create_document(content, uri, metadata or {})
document = await rag.create_document(
content, uri, title=title, metadata=metadata or {}
)
return document.id
except Exception:
return None
@ -102,6 +112,7 @@ def create_mcp_server(db_path: Path) -> FastMCP:
id=document.id,
content=document.content,
uri=document.uri,
title=document.title,
metadata=document.metadata,
created_at=str(document.created_at),
updated_at=str(document.updated_at),
@ -123,6 +134,7 @@ def create_mcp_server(db_path: Path) -> FastMCP:
id=doc.id,
content=doc.content,
uri=doc.uri,
title=doc.title,
metadata=doc.metadata,
created_at=str(doc.created_at),
updated_at=str(doc.updated_at),

4
tests/test_app.py
View file

@ -78,7 +78,9 @@ async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch):
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
await app.add_document_from_source(file_path)
mock_client.create_document_from_source.assert_called_once_with(file_path)
mock_client.create_document_from_source.assert_called_once_with(
file_path, title=None
)
mock_rich_print.assert_called_once_with(mock_doc, truncate=True)
mock_print.assert_called_once_with(
"[b]Document with id [cyan]1[/cyan] added successfully.[/b]"

15
tests/test_cli.py
View file

@ -45,6 +45,21 @@ def test_add_document_src():
mock_app_instance.add_document_from_source.assert_called_once()
def test_add_document_src_with_title():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock()
mock_app_instance.add_document_from_source = AsyncMock()
mock_app.return_value = mock_app_instance
result = runner.invoke(cli, ["add-src", "test.txt", "--title", "Nice Name"])
assert result.exit_code == 0
mock_app_instance.add_document_from_source.assert_called_once()
# Verify title is forwarded
await_args = mock_app_instance.add_document_from_source.await_args
assert await_args.kwargs.get("title") == "Nice Name"
def test_get_document():
with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock()

36
tests/test_client.py
View file

@ -105,6 +105,42 @@ async def test_client_create_document_from_source(temp_db_path):
assert "md5" in doc2.metadata
@pytest.mark.asyncio
async def test_client_create_document_from_source_with_title(temp_db_path):
"""Test creating a document from a file source with a title."""
async with HaikuRAG(temp_db_path) as client:
with tempfile.TemporaryDirectory() as temp_dir:
test_content = "This is test content from a file."
temp_path = Path(temp_dir) / "test_title.txt"
temp_path.write_text(test_content)
doc = await client.create_document_from_source(
source=temp_path, title="My Doc"
)
assert doc.id is not None
assert doc.title == "My Doc"
@pytest.mark.asyncio
async def test_client_update_title_noop_behavior(temp_db_path):
"""When content is unchanged, updating title should update document without re-chunking."""
async with HaikuRAG(temp_db_path) as client:
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test_update_title.txt"
temp_path.write_text("Original content")
doc1 = await client.create_document_from_source(temp_path, title="Title A")
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 doc2.id == doc1.id
# Fetch and verify title updated
got = await client.get_document_by_id(doc1.id)
assert got is not None
assert got.title == "Title B"
@pytest.mark.asyncio
async def test_client_create_document_from_source_unsupported(temp_db_path):
"""Test creating a document from an unsupported file type."""