Merge pull request #2 from ggozad/feat/monitor

File monitoring support

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, OpenAI 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
 
@@ -22,6 +25,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:
@@ -77,7 +87,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
 ```
 
@@ -87,7 +97,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:
 
@@ -99,13 +128,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

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",

src/haiku/rag/app.py

@@ -1,9 +1,13 @@
+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.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
 from haiku.rag.store.models.document import Document
 
@@ -88,20 +92,25 @@ 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
+        async with HaikuRAG(self.db_path) as client:
+            monitor = FileWatcher(paths=Config.MONITOR_DIRECTORIES, client=client)
+            monitor_task = asyncio.create_task(monitor.observe())
+            server = create_mcp_server(self.db_path)
 
-        server = create_mcp_server(self.db_path)
-
-        if transport == "stdio":
-            self.console.print("[green]Starting MCP server on stdio...[/green]")
-            server.run("stdio")
-        elif transport == "sse":
-            self.console.print(
-                "[green]Starting MCP server with streamable HTTP...[/green]"
-            )
-            server.run("sse")
-        else:
-            self.console.print("[green]Starting MCP server with HTTP...[/green]")
-            server.run("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

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__":

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

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] = []
 
     EMBEDDINGS_PROVIDER: str = "ollama"
     EMBEDDINGS_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)

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

src/haiku/rag/monitor.py

@@ -0,0 +1,74 @@
+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
+
+logger = get_logger()
+
+
+class FileFilter(DefaultFilter):
+    def __init__(self, *, ignore_paths: list[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[Path], client: HaikuRAG):
+        self.paths = paths
+        self.client = client
+
+    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):
+        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)
+
+    async def _upsert_document(self, file: Path) -> Document | None:
+        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}")

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
 

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()