From 5f604d31dd545da5e2fb0c3ef36e7071e88d8cb9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 20 Jun 2025 07:44:03 +0200 Subject: [PATCH 1/7] Add monitor --- src/haiku/rag/app.py | 20 +++++++-------- src/haiku/rag/cli.py | 2 +- src/haiku/rag/logging.py | 24 ++++++++++++++++++ src/haiku/rag/monitor.py | 54 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 12 deletions(-) create mode 100644 src/haiku/rag/logging.py create mode 100644 src/haiku/rag/monitor.py diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 1972afae..b8e3c6d0 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -1,9 +1,12 @@ +import asyncio from pathlib import Path from rich.console import Console from rich.markdown import Markdown from haiku.rag.client import HaikuRAG +from haiku.rag.mcp import create_mcp_server +from haiku.rag.monitor import FileWatcher from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document @@ -88,20 +91,15 @@ class HaikuRAGApp: self.console.print(content) self.console.rule() - def serve(self, transport: str | None = None): + async def serve(self, transport: str | None = None): """Start the MCP server.""" - from haiku.rag.mcp import create_mcp_server - + monitor = FileWatcher(paths=[]) + asyncio.create_task(monitor.observe()) server = create_mcp_server(self.db_path) if transport == "stdio": - self.console.print("[green]Starting MCP server on stdio...[/green]") - server.run("stdio") + await server.run_stdio_async() elif transport == "sse": - self.console.print( - "[green]Starting MCP server with streamable HTTP...[/green]" - ) - server.run("sse") + await server.run_sse_async("sse") else: - self.console.print("[green]Starting MCP server with HTTP...[/green]") - server.run("streamable-http") + await server.run_http_async("streamable-http") diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 21b5d4d1..71e2c8b9 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -146,7 +146,7 @@ def serve( elif sse: transport = "sse" - app.serve(transport=transport) + event_loop.run_until_complete(app.serve(transport=transport)) if __name__ == "__main__": diff --git a/src/haiku/rag/logging.py b/src/haiku/rag/logging.py new file mode 100644 index 00000000..33197b8d --- /dev/null +++ b/src/haiku/rag/logging.py @@ -0,0 +1,24 @@ +import logging + +from rich.console import Console +from rich.logging import RichHandler + + +def get_logger() -> logging.Logger: + logger = logging.getLogger("haiku.rag") + + handler = RichHandler( + console=Console(stderr=True), + rich_tracebacks=True, + ) + formatter = logging.Formatter("%(message)s") + handler.setFormatter(formatter) + + logger.setLevel("INFO") + + # Remove any existing handlers to avoid duplicates on reconfiguration + for hdlr in logger.handlers[:]: + logger.removeHandler(hdlr) + + logger.addHandler(handler) + return logger diff --git a/src/haiku/rag/monitor.py b/src/haiku/rag/monitor.py new file mode 100644 index 00000000..dccdb764 --- /dev/null +++ b/src/haiku/rag/monitor.py @@ -0,0 +1,54 @@ +from pathlib import Path + +from watchfiles import Change, DefaultFilter, awatch + +from haiku.rag.logging import get_logger +from haiku.rag.reader import FileReader +from haiku.rag.store.models.document import Document + +logger = get_logger() + + +class FileFilter(DefaultFilter): + def __init__(self, *, ignore_paths: list[str | Path] | None = None) -> None: + self.extensions = tuple(FileReader.extensions) + super().__init__(ignore_paths=ignore_paths) + + def __call__(self, change: "Change", path: str) -> bool: + return path.endswith(self.extensions) and super().__call__(change, path) + + +class FileWatcher: + def __init__(self, paths: list[str | Path]): + self.paths = paths + + async def observe(self): + logger.info(f"Watching files in {self.paths}") + filter = FileFilter() + await self.refresh() + + async for changes in awatch(*self.paths, watch_filter=filter): + await self.handler(changes) + + async def handler(self, changes: set[tuple[Change, str]]): + for change, path in changes: + if change == Change.added or change == Change.modified: + await self._upsert_document(Path(path)) + elif change == Change.deleted: + await self._delete_document(Path(path)) + + async def refresh(self, paths: list[str | Path] | None = None): + if paths is None: + paths = self.paths + for path in paths: + for f in Path(path).rglob("**/*"): + if f.is_file() and f.suffix in FileReader.extensions: + await self._upsert_document(f) + + async def _delete_document(self, file: Path): + logger.info(f"Deleting document from {file}") + pass + + async def _upsert_document(self, file: Path) -> Document | None: + logger.info(f"Updating document from {file}") + pass From 8132a44994d51882aad87cfbfbd29de9cd062ba9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 20 Jun 2025 09:04:07 +0200 Subject: [PATCH 2/7] Use MONITOR_DIRECTORIES for a list of dirs to monitor --- src/haiku/rag/app.py | 3 ++- src/haiku/rag/config.py | 14 +++++++++++++- src/haiku/rag/monitor.py | 10 ++++------ 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index b8e3c6d0..371b571f 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -5,6 +5,7 @@ from rich.console import Console from rich.markdown import Markdown from haiku.rag.client import HaikuRAG +from haiku.rag.config import Config from haiku.rag.mcp import create_mcp_server from haiku.rag.monitor import FileWatcher from haiku.rag.store.models.chunk import Chunk @@ -93,7 +94,7 @@ class HaikuRAGApp: async def serve(self, transport: str | None = None): """Start the MCP server.""" - monitor = FileWatcher(paths=[]) + monitor = FileWatcher(paths=Config.MONITOR_DIRECTORIES) asyncio.create_task(monitor.observe()) server = create_mcp_server(self.db_path) diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index ffbe3054..6483a092 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -2,7 +2,7 @@ import os from pathlib import Path from dotenv import load_dotenv -from pydantic import BaseModel +from pydantic import BaseModel, field_validator from haiku.rag.utils import get_default_data_dir @@ -13,6 +13,7 @@ class AppConfig(BaseModel): ENV: str = "development" DEFAULT_DATA_DIR: Path = get_default_data_dir() + MONITOR_DIRECTORIES: list[Path] = [] EMBEDDING_PROVIDER: str = "ollama" EMBEDDING_MODEL: str = "mxbai-embed-large" @@ -23,6 +24,17 @@ class AppConfig(BaseModel): OLLAMA_BASE_URL: str = "http://localhost:11434" + @field_validator("MONITOR_DIRECTORIES", mode="before") + @classmethod + def parse_monitor_directories(cls, v): + if isinstance(v, str): + if not v.strip(): + return [] + return [ + Path(path.strip()).absolute() for path in v.split(",") if path.strip() + ] + return v + # Expose Config object for app to import Config = AppConfig.model_validate(os.environ) diff --git a/src/haiku/rag/monitor.py b/src/haiku/rag/monitor.py index dccdb764..be12f9d9 100644 --- a/src/haiku/rag/monitor.py +++ b/src/haiku/rag/monitor.py @@ -10,7 +10,7 @@ logger = get_logger() class FileFilter(DefaultFilter): - def __init__(self, *, ignore_paths: list[str | Path] | None = None) -> None: + def __init__(self, *, ignore_paths: list[Path] | None = None) -> None: self.extensions = tuple(FileReader.extensions) super().__init__(ignore_paths=ignore_paths) @@ -19,7 +19,7 @@ class FileFilter(DefaultFilter): class FileWatcher: - def __init__(self, paths: list[str | Path]): + def __init__(self, paths: list[Path]): self.paths = paths async def observe(self): @@ -37,10 +37,8 @@ class FileWatcher: elif change == Change.deleted: await self._delete_document(Path(path)) - async def refresh(self, paths: list[str | Path] | None = None): - if paths is None: - paths = self.paths - for path in paths: + async def refresh(self): + for path in self.paths: for f in Path(path).rglob("**/*"): if f.is_file() and f.suffix in FileReader.extensions: await self._upsert_document(f) From ee57bc0a82639b44874ef72b661146c664ba68ef Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 20 Jun 2025 09:27:01 +0200 Subject: [PATCH 3/7] Use path.as_uri() instead of path.resolve() --- src/haiku/rag/client.py | 2 +- tests/test_client.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index d25ddb2b..920f262a 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -88,7 +88,7 @@ class HaikuRAG: if not source_path.exists(): raise ValueError(f"File does not exist: {source_path}") - uri = str(source_path.resolve()) + uri = source_path.as_uri() md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest() # Check if document already exists diff --git a/tests/test_client.py b/tests/test_client.py index 77430e24..086facc0 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -98,7 +98,7 @@ async def test_client_create_document_from_source(): assert doc.id is not None assert doc.content == test_content - assert doc.uri == str(temp_path.resolve()) + assert doc.uri == temp_path.as_uri() assert doc.metadata["source_type"] == "file" assert "contentType" in doc.metadata assert "md5" in doc.metadata @@ -109,7 +109,7 @@ async def test_client_create_document_from_source(): assert doc2.id is not None assert doc2.content == test_content - assert doc2.uri == str(temp_path.resolve()) + assert doc2.uri == temp_path.as_uri() assert "contentType" in doc2.metadata assert "md5" in doc2.metadata From 51e7f990a099e6f367f3da0d6679c70785df4d8f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 20 Jun 2025 09:28:44 +0200 Subject: [PATCH 4/7] Implement upsert and delete in monitor --- src/haiku/rag/app.py | 19 ++++++++++--------- src/haiku/rag/monitor.py | 36 +++++++++++++++++++++++++++++------- 2 files changed, 39 insertions(+), 16 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 371b571f..3589c003 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -94,13 +94,14 @@ class HaikuRAGApp: async def serve(self, transport: str | None = None): """Start the MCP server.""" - monitor = FileWatcher(paths=Config.MONITOR_DIRECTORIES) - asyncio.create_task(monitor.observe()) - server = create_mcp_server(self.db_path) + async with HaikuRAG(self.db_path) as client: + monitor = FileWatcher(paths=Config.MONITOR_DIRECTORIES, client=client) + asyncio.create_task(monitor.observe()) + server = create_mcp_server(self.db_path) - if transport == "stdio": - await server.run_stdio_async() - elif transport == "sse": - await server.run_sse_async("sse") - else: - await server.run_http_async("streamable-http") + if transport == "stdio": + await server.run_stdio_async() + elif transport == "sse": + await server.run_sse_async("sse") + else: + await server.run_http_async("streamable-http") diff --git a/src/haiku/rag/monitor.py b/src/haiku/rag/monitor.py index be12f9d9..97809e9e 100644 --- a/src/haiku/rag/monitor.py +++ b/src/haiku/rag/monitor.py @@ -2,6 +2,7 @@ from pathlib import Path from watchfiles import Change, DefaultFilter, awatch +from haiku.rag.client import HaikuRAG from haiku.rag.logging import get_logger from haiku.rag.reader import FileReader from haiku.rag.store.models.document import Document @@ -19,8 +20,9 @@ class FileFilter(DefaultFilter): class FileWatcher: - def __init__(self, paths: list[Path]): + def __init__(self, paths: list[Path], client: HaikuRAG): self.paths = paths + self.client = client async def observe(self): logger.info(f"Watching files in {self.paths}") @@ -43,10 +45,30 @@ class FileWatcher: if f.is_file() and f.suffix in FileReader.extensions: await self._upsert_document(f) - async def _delete_document(self, file: Path): - logger.info(f"Deleting document from {file}") - pass - async def _upsert_document(self, file: Path) -> Document | None: - logger.info(f"Updating document from {file}") - pass + try: + uri = file.as_uri() + existing_doc = await self.client.get_document_by_uri(uri) + print(uri) + if existing_doc: + doc = await self.client.create_document_from_source(str(file)) + logger.info(f"Updated document {existing_doc.id} from {file}") + return doc + else: + doc = await self.client.create_document_from_source(str(file)) + logger.info(f"Created new document {doc.id} from {file}") + return doc + except Exception as e: + logger.error(f"Failed to upsert document from {file}: {e}") + return None + + async def _delete_document(self, file: Path): + try: + uri = file.as_uri() + existing_doc = await self.client.get_document_by_uri(uri) + + if existing_doc and existing_doc.id: + await self.client.delete_document(existing_doc.id) + logger.info(f"Deleted document {existing_doc.id} for {file}") + except Exception as e: + logger.error(f"Failed to delete document for {file}: {e}") From 27f65adda7470897c9731bf7695e369156265bf8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 20 Jun 2025 09:35:20 +0200 Subject: [PATCH 5/7] Test FileWatcher --- tests/test_monitor.py | 99 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100644 tests/test_monitor.py diff --git a/tests/test_monitor.py b/tests/test_monitor.py new file mode 100644 index 00000000..ac909631 --- /dev/null +++ b/tests/test_monitor.py @@ -0,0 +1,99 @@ +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from haiku.rag.client import HaikuRAG +from haiku.rag.monitor import FileWatcher +from haiku.rag.store.models.document import Document + + +@pytest.mark.asyncio +async def test_file_watcher_upsert_document(): + """Test FileWatcher._upsert_document method.""" + + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("Test content for file watcher") + temp_path = Path(f.name) + + try: + mock_client = AsyncMock(spec=HaikuRAG) + mock_doc = Document(id=1, content="Test content", uri=temp_path.as_uri()) + mock_client.create_document_from_source.return_value = mock_doc + mock_client.get_document_by_uri.return_value = None # No existing document + + watcher = FileWatcher(paths=[temp_path.parent], client=mock_client) + + result = await watcher._upsert_document(temp_path) + + assert result is not None + assert result.id == 1 + mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri()) + mock_client.create_document_from_source.assert_called_once_with(str(temp_path)) + + finally: + temp_path.unlink(missing_ok=True) + + +@pytest.mark.asyncio +async def test_file_watcher_upsert_existing_document(): + """Test FileWatcher._upsert_document with existing document.""" + + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("Test content for file watcher") + temp_path = Path(f.name) + + try: + mock_client = AsyncMock(spec=HaikuRAG) + existing_doc = Document(id=1, content="Old content", uri=temp_path.as_uri()) + updated_doc = Document(id=1, content="Updated content", uri=temp_path.as_uri()) + + mock_client.get_document_by_uri.return_value = existing_doc + mock_client.create_document_from_source.return_value = updated_doc + + watcher = FileWatcher(paths=[temp_path.parent], client=mock_client) + + result = await watcher._upsert_document(temp_path) + + assert result is not None + assert result.content == "Updated content" + mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri()) + mock_client.create_document_from_source.assert_called_once_with(str(temp_path)) + + finally: + temp_path.unlink(missing_ok=True) + + +@pytest.mark.asyncio +async def test_file_watcher_delete_document(): + """Test FileWatcher._delete_document method.""" + temp_path = Path("/tmp/test_file.txt") + + mock_client = AsyncMock(spec=HaikuRAG) + existing_doc = Document(id=1, content="Content to delete", uri=temp_path.as_uri()) + mock_client.get_document_by_uri.return_value = existing_doc + mock_client.delete_document.return_value = True + + watcher = FileWatcher(paths=[temp_path.parent], client=mock_client) + + await watcher._delete_document(temp_path) + + mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri()) + mock_client.delete_document.assert_called_once_with(1) + + +@pytest.mark.asyncio +async def test_file_watcher_delete_nonexistent_document(): + """Test FileWatcher._delete_document with non-existent document.""" + temp_path = Path("/tmp/nonexistent_file.txt") + + mock_client = AsyncMock(spec=HaikuRAG) + mock_client.get_document_by_uri.return_value = None + + watcher = FileWatcher(paths=[temp_path.parent], client=mock_client) + + await watcher._delete_document(temp_path) + + mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri()) + mock_client.delete_document.assert_not_called() From 18e1e7959811ea0fca1140b025c1019c00da0619 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 20 Jun 2025 10:03:00 +0200 Subject: [PATCH 6/7] Update docs --- README.md | 40 +++++++++++++++++++++++++++++++--------- pyproject.toml | 1 + 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 9c73285b..554d4d61 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,9 @@ A SQLite-based Retrieval-Augmented Generation (RAG) system built for efficient d - **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! +- **File monitoring** when run as a server automatically indexing your files +- **MCP server** Exposes functionality as MCP tools. +- **Python client** Call `haiku.rag` from your own python applications. ## Installation @@ -21,6 +24,13 @@ For other providers use: ## Configuration +You can set the directories to monitor using the `MONITOR_DIRECTORIES` environment variable (as comma separated values) : + +```bash +# Monitor single directory +export MONITOR_DIRECTORIES="/path/to/documents,/another_path/to/documents" +``` + If you want to use an alternative embeddings provider (Ollama being the default) you will need to set the provider details through environment variables: By default: @@ -67,7 +77,7 @@ haiku-rag search "machine learning" # Search with custom options haiku-rag search "python programming" --limit 10 --k 100 -# Start MCP server (default HTTP transport) +# Start file monitoring & MCP server (default HTTP transport) haiku-rag serve # --stdio for stdio transport or --sse for SSE transport ``` @@ -77,7 +87,26 @@ haiku-rag command -h ``` to see additional parameters for a command. -## MCP Server +## File Monitoring & MCP server + +You can start the server (using Streamble HTTP, stdio or SSE transports) with: + +```bash +# Start with default HTTP transport +haiku-rag serve # --stdio for stdio transport or --sse for SSE transport +``` + +You need to have set the `MONITOR_DIRECTORIES` environment variable for monitoring to take place. + +### File monitoring + +`haiku.rag` can watch directories for changes and automatically update the document store: + +- **Startup**: Scan all monitored directories and add any new files +- **File Added/Modified**: Automatically parse and add/update the document in the database +- **File Deleted**: Remove the corresponding document from the database + +### MCP Server `haiku.rag` includes a Model Context Protocol (MCP) server that exposes RAG functionality as tools for AI assistants like Claude Desktop. The MCP server provides the following tools: @@ -89,13 +118,6 @@ to see additional parameters for a command. - `list_documents` - List all documents with pagination - `delete_document` - Delete documents by ID -You can start the server (using Streamble HTTP, stdio or SSE transports) with: - -```bash -# Start with default HTTP transport -haiku-rag serve # --stdio for stdio transport or --sse for SSE transport -``` - ## Using `haiku.rag` from python ### Managing documents diff --git a/pyproject.toml b/pyproject.toml index 9fc68928..68dd2c0a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,7 @@ authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.10" +keywords = ["RAG", "sqlite", "sqlite-vec", "ml", "mcp"] classifiers = [ "Development Status :: 4 - Beta", "Environment :: Console", From 03b2b20ae2d3704861bf4e6a54519eb27a4de09d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 20 Jun 2025 10:18:07 +0200 Subject: [PATCH 7/7] Handle keyboard interrupt --- src/haiku/rag/app.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 3589c003..6db14c83 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -96,12 +96,21 @@ class HaikuRAGApp: """Start the MCP server.""" async with HaikuRAG(self.db_path) as client: monitor = FileWatcher(paths=Config.MONITOR_DIRECTORIES, client=client) - asyncio.create_task(monitor.observe()) + monitor_task = asyncio.create_task(monitor.observe()) server = create_mcp_server(self.db_path) - if transport == "stdio": - await server.run_stdio_async() - elif transport == "sse": - await server.run_sse_async("sse") - else: - await server.run_http_async("streamable-http") + try: + if transport == "stdio": + await server.run_stdio_async() + elif transport == "sse": + await server.run_sse_async("sse") + else: + await server.run_http_async("streamable-http") + except KeyboardInterrupt: + pass + finally: + monitor_task.cancel() + try: + await monitor_task + except asyncio.CancelledError: + pass