From d7c9b41a2c5b39c5f29c7ffcd548693fbf92219f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 18 Jun 2025 15:30:57 +0200 Subject: [PATCH] Console app --- README.md | 84 ++++++++++++++++--------------- pyproject.toml | 5 ++ src/haiku/rag/app.py | 89 +++++++++++++++++++++++++++++++++ src/haiku/rag/cli.py | 115 +++++++++++++++++++++++++++++++++++++++++++ uv.lock | 62 +++++++++++++++++++++++ 5 files changed, 314 insertions(+), 41 deletions(-) create mode 100644 src/haiku/rag/app.py create mode 100644 src/haiku/rag/cli.py diff --git a/README.md b/README.md index a3ce1e53..d7a33989 100644 --- a/README.md +++ b/README.md @@ -4,11 +4,9 @@ A SQLite-based Retrieval-Augmented Generation (RAG) system built for efficient d ## Features - **Local SQLite**: No need to run additional servers -- **Support for various embedding providers**: You can use Ollama, VoyageAI, OpenAI or add your own -- **Vector Embeddings**: Uses sqlite-vec for efficient similarity search -- **Hybrid Search**: Full-text search (FTS5) combined with vector embeddings using Reciprocal Rank Fusion -- **Multi-format Support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, audio and more -- **Web Content**: Direct URL ingestion with automatic content type detection +- **Support for various embedding providers**: You can use Ollama, VoyageAI or add your own +- **Hybrid Search**: Vector search using `sqlite-vec` combined with full-text search `FTS5`, using Reciprocal Rank Fusion +- **Multi-format Support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, audio and more. Or add a url! ## Installation @@ -40,7 +38,45 @@ EMBEDDING_MODEL="voyage-3.5" # or any other model EMBEDDING_VECTOR_DIM=1024 ``` -## Quick Start +## Command Line Interface + +`haiku.rag` includes a CLI application for managing documents and performing searches from the command line: + +### Available Commands + +```bash +# List all documents +haiku-rag list + +# Add document from text +haiku-rag add "Your document content here" + +# Add document from file or URL +haiku-rag add-src /path/to/document.pdf +haiku-rag add-src https://example.com/article.html + +# Get and display a specific document +haiku-rag get 1 + +# Delete a document by ID +haiku-rag delete 1 + +# Search documents +haiku-rag search "machine learning" + +# Search with custom options +haiku-rag search "python programming" --limit 10 --k 100 +``` + +All commands support the `--db` option to specify a custom database path. Run +```bash +haiku-rag command -h +``` +to see additional parameters for a command. + +## Using `haiku.rag` from python + +### Managing documents ```python from pathlib import Path @@ -82,23 +118,9 @@ async with HaikuRAG("path/to/database.db") as client: print(f"Content: {chunk.content}") print(f"Document ID: {chunk.document_id}") print("---") - - -# Or use without the context manager. -client = HaikuRAG(":memory:") -try: - # ... operations ... -finally: - client.close() ``` -## Search Functionality - -`haiku.rag` provides hybrid search combining vector similarity and full-text search: -1. **Vector Search**: Uses embeddings to find semantically similar content -2. **Full-text Search**: Uses SQLite FTS5 for exact keyword matching -3. **Hybrid Ranking**: Combines both using Reciprocal Rank Fusion (RRF) -4. **Chunked Results**: Returns relevant document chunks with scores +## Searching documents ```python async with HaikuRAG("database.db") as client: @@ -115,23 +137,3 @@ async with HaikuRAG("database.db") as client: print(f"Content: {chunk.content}") print(f"From document: {chunk.document_id}") ``` - - -## Supported File Formats - -`haiku.rag` supports 40+ file formats through MarkItDown: - -- **Documents**: PDF, DOCX, PPTX, XLSX -- **Web**: HTML, XML -- **Text**: TXT, MD, CSV, JSON, YAML -- **Code**: PY, JS, TS, C, CPP, JAVA, GO, RS, and more -- **Media**: MP3, WAV (transcription) - -## Contributing - -1. Fork the repository -2. Create a feature branch -3. Add tests for new functionality -4. Ensure all tests pass: `pytest` -5. Run type checking & linting with `pyright` & `ruff check` -6. Submit a pull request diff --git a/pyproject.toml b/pyproject.toml index 76f5a9e0..d1116ff2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,14 +12,19 @@ dependencies = [ "ollama>=0.5.1", "pydantic>=2.11.7", "python-dotenv>=1.1.0", + "rich>=14.0.0", "sqlite-vec>=0.1.6", "tiktoken>=0.9.0", + "typer>=0.16.0", "watchfiles>=1.1.0", ] [project.optional-dependencies] voyageai = ["voyageai>=0.3.2"] +[project.scripts] +haiku-rag = "haiku.rag.cli:cli" + [build-system] requires = ["hatchling"] build-backend = "hatchling.build" diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py new file mode 100644 index 00000000..6629981f --- /dev/null +++ b/src/haiku/rag/app.py @@ -0,0 +1,89 @@ +from pathlib import Path + +from rich.console import Console +from rich.markdown import Markdown + +from haiku.rag.client import HaikuRAG +from haiku.rag.store.models.chunk import Chunk +from haiku.rag.store.models.document import Document + + +class HaikuRAGApp: + def __init__(self, db_path: Path): + self.db_path = db_path + self.console = Console() + + async def list_documents(self): + async with HaikuRAG(db_path=self.db_path) as self.client: + documents = await self.client.list_documents() + for doc in documents: + self._rich_print_document(doc, truncate=True) + + async def add_document_from_text(self, text: str): + async with HaikuRAG(db_path=self.db_path) as self.client: + doc = await self.client.create_document(text) + self._rich_print_document(doc, truncate=True) + self.console.print( + f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]" + ) + + async def add_document_from_source(self, file_path: Path): + async with HaikuRAG(db_path=self.db_path) as self.client: + doc = await self.client.create_document_from_source(file_path) + self._rich_print_document(doc, truncate=True) + self.console.print( + f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]" + ) + + async def get_document(self, doc_id: int): + async with HaikuRAG(db_path=self.db_path) as self.client: + doc = await self.client.get_document_by_id(doc_id) + if doc is None: + self.console.print(f"[red]Document with id {doc_id} not found.[/red]") + return + self._rich_print_document(doc, truncate=False) + + async def delete_document(self, doc_id: int): + async with HaikuRAG(db_path=self.db_path) as self.client: + await self.client.delete_document(doc_id) + self.console.print(f"[b]Document {doc_id} deleted successfully.[/b]") + + async def search(self, query: str, limit: int = 5, k: int = 60): + async with HaikuRAG(db_path=self.db_path) as self.client: + results = await self.client.search(query, limit=limit, k=k) + if not results: + self.console.print("[red]No results found.[/red]") + return + for chunk, score in results: + self._rich_print_search_result(chunk, score) + + def _rich_print_document(self, doc: Document, truncate: bool = False): + """Format a document for display.""" + if truncate: + content = doc.content.splitlines() + if len(content) > 3: + content = content[:3] + ["\n…"] + content = "\n".join(content) + content = Markdown(content) + else: + content = Markdown(doc.content) + self.console.print( + f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} [repr.attrib_name]uri[/repr.attrib_name]: {doc.uri} [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}" + ) + self.console.print( + f"[repr.attrib_name]created at[/repr.attrib_name]: {doc.created_at} [repr.attrib_name]updated at[/repr.attrib_name]: {doc.updated_at}" + ) + self.console.print("[repr.attrib_name]content[/repr.attrib_name]:") + self.console.print(content) + self.console.rule() + + def _rich_print_search_result(self, chunk: Chunk, score: float): + """Format a search result chunk for display.""" + content = Markdown(chunk.content) + self.console.print( + f"[repr.attrib_name]document_id[/repr.attrib_name]: {chunk.document_id} " + f"[repr.attrib_name]score[/repr.attrib_name]: {score:.4f}" + ) + self.console.print("[repr.attrib_name]content[/repr.attrib_name]:") + self.console.print(content) + self.console.rule() diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py new file mode 100644 index 00000000..fe48199e --- /dev/null +++ b/src/haiku/rag/cli.py @@ -0,0 +1,115 @@ +import asyncio +from pathlib import Path + +import typer + +from haiku.rag.app import HaikuRAGApp +from haiku.rag.utils import get_default_data_dir + +cli = typer.Typer( + context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True +) + +event_loop = asyncio.get_event_loop() + + +@cli.command("list", help="List all stored documents") +def list_documents( + db: Path = typer.Option( + get_default_data_dir() / "haiku.rag.sqlite", + "--db", + help="The path to the sqlite db to use", + ), +): + app = HaikuRAGApp(db_path=db) + event_loop.run_until_complete(app.list_documents()) + + +@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", + ), + db: Path = typer.Option( + get_default_data_dir() / "haiku.rag.sqlite", + "--db", + help="The path to the sqlite db to use", + ), +): + app = HaikuRAGApp(db_path=db) + event_loop.run_until_complete(app.add_document_from_text(text=text)) + + +@cli.command("add-src", help="Add a document from a file path or URL") +def add_document_src( + file_path: Path = typer.Argument( + help="The file path or URL of the document to add", + ), + db: Path = typer.Option( + get_default_data_dir() / "haiku.rag.sqlite", + "--db", + help="The path to the sqlite db to use", + ), +): + app = HaikuRAGApp(db_path=db) + event_loop.run_until_complete(app.add_document_from_source(file_path=file_path)) + + +@cli.command("get", help="Get and display a document by its ID") +def get_document( + doc_id: int = typer.Argument( + help="The ID of the document to get", + ), + db: Path = typer.Option( + get_default_data_dir() / "haiku.rag.sqlite", + "--db", + help="The path to the sqlite db to use", + ), +): + app = HaikuRAGApp(db_path=db) + event_loop.run_until_complete(app.get_document(doc_id=doc_id)) + + +@cli.command("delete", help="Delete a document by its ID") +def delete_document( + doc_id: int = typer.Argument( + help="The ID of the document to delete", + ), + db: Path = typer.Option( + get_default_data_dir() / "haiku.rag.sqlite", + "--db", + help="The path to the sqlite db to use", + ), +): + app = HaikuRAGApp(db_path=db) + event_loop.run_until_complete(app.delete_document(doc_id=doc_id)) + + +@cli.command("search", help="Search for documents by a query") +def search( + query: str = typer.Argument( + help="The search query to use", + ), + limit: int = typer.Option( + 5, + "--limit", + "-l", + help="Maximum number of results to return", + ), + k: int = typer.Option( + 60, + "--k", + help="Reciprocal Rank Fusion k parameter", + ), + db: Path = typer.Option( + get_default_data_dir() / "haiku.rag.sqlite", + "--db", + help="The path to the sqlite db to use", + ), +): + app = HaikuRAGApp(db_path=db) + event_loop.run_until_complete(app.search(query=query, limit=limit, k=k)) + + +if __name__ == "__main__": + cli() diff --git a/uv.lock b/uv.lock index d63b0215..7267a11b 100644 --- a/uv.lock +++ b/uv.lock @@ -474,8 +474,10 @@ dependencies = [ { name = "ollama" }, { name = "pydantic" }, { name = "python-dotenv" }, + { name = "rich" }, { name = "sqlite-vec" }, { name = "tiktoken" }, + { name = "typer" }, { name = "watchfiles" }, ] @@ -502,8 +504,10 @@ requires-dist = [ { name = "ollama", specifier = ">=0.5.1" }, { name = "pydantic", specifier = ">=2.11.7" }, { name = "python-dotenv", specifier = ">=1.1.0" }, + { name = "rich", specifier = ">=14.0.0" }, { name = "sqlite-vec", specifier = ">=0.1.6" }, { name = "tiktoken", specifier = ">=0.9.0" }, + { name = "typer", specifier = ">=0.16.0" }, { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" }, { name = "watchfiles", specifier = ">=1.1.0" }, ] @@ -676,6 +680,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/be/0c/3153f159b78e368ac473a00e955d69d976e4b69740ed07c76c9f72a161b8/mammoth-1.9.1-py2.py3-none-any.whl", hash = "sha256:f0569bd640cee6c77a07e7c75c5dc10d745dc4dc95d530cfcbb0a5d9536d636c", size = 52991, upload-time = "2025-05-28T19:17:54.62Z" }, ] +[[package]] +name = "markdown-it-py" +version = "3.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/38/71/3b932df36c1a044d397a1f92d1cf91ee0a503d91e470cbd670aa66b07ed0/markdown-it-py-3.0.0.tar.gz", hash = "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb", size = 74596, upload-time = "2023-06-03T06:41:14.443Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/42/d7/1ec15b46af6af88f19b8e5ffea08fa375d433c998b8a7639e76935c14f1f/markdown_it_py-3.0.0-py3-none-any.whl", hash = "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1", size = 87528, upload-time = "2023-06-03T06:41:11.019Z" }, +] + [[package]] name = "markdownify" version = "1.1.0" @@ -726,6 +742,15 @@ xlsx = [ { name = "pandas" }, ] +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + [[package]] name = "mpmath" version = "1.3.0" @@ -1307,6 +1332,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7c/e4/56027c4a6b4ae70ca9de302488c5ca95ad4a39e190093d6c1a8ace08341b/requests-2.32.4-py3-none-any.whl", hash = "sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c", size = 64847, upload-time = "2025-06-09T16:43:05.728Z" }, ] +[[package]] +name = "rich" +version = "14.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/53/830aa4c3066a8ab0ae9a9955976fb770fe9c6102117c8ec4ab3ea62d89e8/rich-14.0.0.tar.gz", hash = "sha256:82f1bc23a6a21ebca4ae0c45af9bdbc492ed20231dcb63f297d6d1021a9d5725", size = 224078, upload-time = "2025-03-30T14:15:14.23Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/9b/63f4c7ebc259242c89b3acafdb37b41d1185c07ff0011164674e9076b491/rich-14.0.0-py3-none-any.whl", hash = "sha256:1c9491e1951aac09caffd42f448ee3d04e58923ffe14993f6e83068dc395d7e0", size = 243229, upload-time = "2025-03-30T14:15:12.283Z" }, +] + [[package]] name = "ruff" version = "0.11.13" @@ -1332,6 +1370,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/bf/b273dd11673fed8a6bd46032c0ea2a04b2ac9bfa9c628756a5856ba113b0/ruff-0.11.13-py3-none-win_arm64.whl", hash = "sha256:b4385285e9179d608ff1d2fb9922062663c658605819a6876d8beef0c30b7f3b", size = 10683928, upload-time = "2025-06-05T21:00:13.758Z" }, ] +[[package]] +name = "shellingham" +version = "1.5.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, +] + [[package]] name = "six" version = "1.17.0" @@ -1483,6 +1530,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" }, ] +[[package]] +name = "typer" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "rich" }, + { name = "shellingham" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/8c/7d682431efca5fd290017663ea4588bf6f2c6aad085c7f108c5dbc316e70/typer-0.16.0.tar.gz", hash = "sha256:af377ffaee1dbe37ae9440cb4e8f11686ea5ce4e9bae01b84ae7c63b87f1dd3b", size = 102625, upload-time = "2025-05-26T14:30:31.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/76/42/3efaf858001d2c2913de7f354563e3a3a2f0decae3efe98427125a8f441e/typer-0.16.0-py3-none-any.whl", hash = "sha256:1f79bed11d4d02d4310e3c1b7ba594183bcedb0ac73b27a9e5f28f6fb5b98855", size = 46317, upload-time = "2025-05-26T14:30:30.523Z" }, +] + [[package]] name = "typing-extensions" version = "4.14.0"