diff --git a/README.md b/README.md index fd62392b..d9d4557b 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, 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 diff --git a/pyproject.toml b/pyproject.toml index c4ac278e..939ca48f 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", diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 1972afae..6db14c83 100644 --- a/src/haiku/rag/app.py +++ b/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 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/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/src/haiku/rag/config.py b/src/haiku/rag/config.py index dd8546d9..26a29dfe 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] = [] 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) 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..97809e9e --- /dev/null +++ b/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}") 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 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()