diff --git a/README.md b/README.md index e6851b5f..8368d909 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,8 @@ uv pip install haiku.rag # Add documents haiku-rag add "Your content here" -haiku-rag add-src document.pdf +haiku-rag add "Your content here" --meta author=alice --meta topic=notes +haiku-rag add-src document.pdf --meta source=manual # Search haiku-rag search "query" diff --git a/docs/cli.md b/docs/cli.md index 373e231c..4afc96aa 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -27,6 +27,9 @@ haiku-rag list From text: ```bash haiku-rag add "Your document content here" + +# Attach metadata (repeat --meta for multiple entries) +haiku-rag add "Your document content here" --meta author=alice --meta topic=notes ``` From file or URL: @@ -36,6 +39,10 @@ haiku-rag add-src https://example.com/article.html # Optionally set a human‑readable title stored in the DB schema haiku-rag add-src /mnt/data/doc1.pdf --title "Q3 Financial Report" + +# Optionally attach metadata (repeat --meta). Values use JSON parsing if possible: +# numbers, booleans, null, arrays/objects; otherwise kept as strings. +haiku-rag add-src /mnt/data/doc1.pdf --meta source=manual --meta page_count=12 --meta published=true ``` !!! note diff --git a/docs/index.md b/docs/index.md index 514fe5bf..648eb076 100644 --- a/docs/index.md +++ b/docs/index.md @@ -43,7 +43,8 @@ async with HaikuRAG("database.lancedb") as client: Or use the CLI: ```bash haiku-rag add "Your document content" -haiku-rag add-src /path/to/document.pdf --title "Q3 Financial Report" +haiku-rag add "Your document content" --meta author=alice +haiku-rag add-src /path/to/document.pdf --title "Q3 Financial Report" --meta source=manual haiku-rag search "query" haiku-rag ask "Who is the author of haiku.rag?" haiku-rag migrate old_database.sqlite # Migrate from SQLite diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 3f7c3c20..a2a690c5 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -144,17 +144,21 @@ class HaikuRAGApp: for doc in documents: self._rich_print_document(doc, truncate=True) - async def add_document_from_text(self, text: str): + async def add_document_from_text(self, text: str, metadata: dict | None = None): async with HaikuRAG(db_path=self.db_path) as self.client: - doc = await self.client.create_document(text) + doc = await self.client.create_document(text, metadata=metadata) self._rich_print_document(doc, truncate=True) self.console.print( f"[bold green]Document {doc.id} added successfully.[/bold green]" ) - async def add_document_from_source(self, source: str, title: str | None = None): + async def add_document_from_source( + 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(source, title=title) + doc = 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]" diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 517d7c71..ca8349fe 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -1,7 +1,9 @@ import asyncio +import json import warnings from importlib.metadata import version from pathlib import Path +from typing import Any import typer @@ -137,11 +139,41 @@ def list_documents( asyncio.run(app.list_documents()) +def _parse_meta_options(meta: list[str] | None) -> dict[str, Any]: + """Parse repeated --meta KEY=VALUE options into a dictionary. + + Raises a Typer error if any entry is malformed. + """ + result: dict[str, Any] = {} + if not meta: + return result + for item in meta: + if "=" not in item: + raise typer.BadParameter("--meta must be in KEY=VALUE format") + key, value = item.split("=", 1) + if not key: + raise typer.BadParameter("--meta key cannot be empty") + # Best-effort JSON coercion: numbers, booleans, null, arrays/objects + try: + parsed = json.loads(value) + result[key] = parsed + except Exception: + # Leave as string if not valid JSON literal + result[key] = value + return result + + @cli.command("add", help="Add a document from text input") def add_document_text( text: str = typer.Argument( help="The text content of the document to add", ), + meta: list[str] | None = typer.Option( + None, + "--meta", + help="Metadata entries as KEY=VALUE (repeatable)", + metavar="KEY=VALUE", + ), db: Path = typer.Option( Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", "--db", @@ -151,7 +183,8 @@ def add_document_text( from haiku.rag.app import HaikuRAGApp app = HaikuRAGApp(db_path=db) - asyncio.run(app.add_document_from_text(text=text)) + metadata = _parse_meta_options(meta) + 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") @@ -165,6 +198,12 @@ def add_document_src( "--title", help="Optional human-readable title to store with the document", ), + meta: list[str] | None = typer.Option( + None, + "--meta", + help="Metadata entries as KEY=VALUE (repeatable)", + metavar="KEY=VALUE", + ), db: Path = typer.Option( Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", "--db", @@ -174,7 +213,12 @@ def add_document_src( from haiku.rag.app import HaikuRAGApp app = HaikuRAGApp(db_path=db) - asyncio.run(app.add_document_from_source(source=source, title=title)) + metadata = _parse_meta_options(meta) + asyncio.run( + app.add_document_from_source( + source=source, title=title, metadata=metadata or None + ) + ) @cli.command("get", help="Get and display a document by its ID") diff --git a/src/haiku/rag/store/__init__.py b/src/haiku/rag/store/__init__.py index ad1e8338..362c9f66 100644 --- a/src/haiku/rag/store/__init__.py +++ b/src/haiku/rag/store/__init__.py @@ -1,4 +1,4 @@ from .engine import Store from .models import Chunk, Document -__all__ = ["Store", "Chunk", "Document"] \ No newline at end of file +__all__ = ["Store", "Chunk", "Document"] diff --git a/src/haiku/rag/store/models/__init__.py b/src/haiku/rag/store/models/__init__.py index 3e9d83be..f01bd304 100644 --- a/src/haiku/rag/store/models/__init__.py +++ b/src/haiku/rag/store/models/__init__.py @@ -1,4 +1,4 @@ from .chunk import Chunk from .document import Document -__all__ = ["Chunk", "Document"] \ No newline at end of file +__all__ = ["Chunk", "Document"] diff --git a/tests/test_app.py b/tests/test_app.py index b2d4a7f2..f596afa0 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -54,7 +54,10 @@ async def test_add_document_from_text(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.add_document_from_text("test document") - mock_client.create_document.assert_called_once_with("test document") + mock_client.create_document.assert_called_once() + args, kwargs = mock_client.create_document.call_args + assert args[0] == "test document" + assert kwargs.get("metadata") is None mock_rich_print.assert_called_once_with(mock_doc, truncate=True) mock_print.assert_called_once_with( "[bold green]Document 1 added successfully.[/bold green]" @@ -78,9 +81,11 @@ 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, title=None - ) + mock_client.create_document_from_source.assert_called_once() + args, kwargs = mock_client.create_document_from_source.call_args + assert args[0] == file_path + assert kwargs.get("title") is None + assert kwargs.get("metadata") is None mock_rich_print.assert_called_once_with(mock_doc, truncate=True) mock_print.assert_called_once_with( "[bold green]Document 1 added successfully.[/bold green]" diff --git a/tests/test_cli.py b/tests/test_cli.py index 6efb5eb2..19ea95a1 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -28,9 +28,10 @@ def test_add_document_text(): result = runner.invoke(cli, ["add", "test document"]) assert result.exit_code == 0 - mock_app_instance.add_document_from_text.assert_called_once_with( - text="test document" - ) + mock_app_instance.add_document_from_text.assert_called_once() + _, kwargs = mock_app_instance.add_document_from_text.call_args + assert kwargs.get("text") == "test document" + assert kwargs.get("metadata") is None def test_add_document_src(): @@ -57,8 +58,83 @@ def test_add_document_src_with_title(): mock_app_instance.add_document_from_source.assert_called_once() # Verify title is forwarded (inspect call kwargs) _, kwargs = mock_app_instance.add_document_from_source.call_args - assert kwargs.get("title") == "Nice Name" + assert kwargs.get("title") == "Nice Name" + assert kwargs.get("source") == "test.txt" + + +def test_add_document_text_with_meta(): + with patch("haiku.rag.app.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.add_document_from_text = AsyncMock() + mock_app.return_value = mock_app_instance + + result = runner.invoke( + cli, + [ + "add", + "some text", + "--meta", + "author=alice", + "--meta", + "topic=notes", + ], + ) + + assert result.exit_code == 0 + mock_app_instance.add_document_from_text.assert_called_once() + _, kwargs = mock_app_instance.add_document_from_text.call_args + assert kwargs.get("text") == "some text" + assert kwargs.get("metadata") == {"author": "alice", "topic": "notes"} + + +def test_add_document_src_with_meta(): + 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 + + result = runner.invoke( + cli, + [ + "add-src", + "test.txt", + "--meta", + "source=manual", + "--meta", + "lang=en", + ], + ) + + assert result.exit_code == 0 + mock_app_instance.add_document_from_source.assert_called_once() + _, kwargs = mock_app_instance.add_document_from_source.call_args assert kwargs.get("source") == "test.txt" + assert kwargs.get("metadata") == {"source": "manual", "lang": "en"} + + +def test_add_document_text_with_numeric_meta(): + with patch("haiku.rag.app.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.add_document_from_text = AsyncMock() + mock_app.return_value = mock_app_instance + + result = runner.invoke( + cli, + [ + "add", + "some text", + "--meta", + "version=3", + "--meta", + "published=true", + ], + ) + + assert result.exit_code == 0 + mock_app_instance.add_document_from_text.assert_called_once() + _, kwargs = mock_app_instance.add_document_from_text.call_args + assert kwargs.get("text") == "some text" + assert kwargs.get("metadata") == {"version": 3, "published": True} def test_get_document():