Merge pull request #2 from ggozad/feat/monitor

File monitoring support
This commit is contained in:
Yiorgis Gozadinos 2025-06-20 10:30:40 +02:00 committed by GitHub
commit 3c98f246b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 270 additions and 29 deletions

View file

@ -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 - **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 - **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! - **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 ## Installation
@ -22,6 +25,13 @@ For other providers use:
## Configuration ## 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: 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: By default:
@ -77,7 +87,7 @@ haiku-rag search "machine learning"
# Search with custom options # Search with custom options
haiku-rag search "python programming" --limit 10 --k 100 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 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. 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: `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 - `list_documents` - List all documents with pagination
- `delete_document` - Delete documents by ID - `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 ## Using `haiku.rag` from python
### Managing documents ### Managing documents

View file

@ -6,6 +6,7 @@ authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" } license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" } readme = { file = "README.md", content-type = "text/markdown" }
requires-python = ">=3.10" requires-python = ">=3.10"
keywords = ["RAG", "sqlite", "sqlite-vec", "ml", "mcp"]
classifiers = [ classifiers = [
"Development Status :: 4 - Beta", "Development Status :: 4 - Beta",
"Environment :: Console", "Environment :: Console",

View file

@ -1,9 +1,13 @@
import asyncio
from pathlib import Path from pathlib import Path
from rich.console import Console from rich.console import Console
from rich.markdown import Markdown from rich.markdown import Markdown
from haiku.rag.client import HaikuRAG 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.chunk import Chunk
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
@ -88,20 +92,25 @@ class HaikuRAGApp:
self.console.print(content) self.console.print(content)
self.console.rule() self.console.rule()
def serve(self, transport: str | None = None): async def serve(self, transport: str | None = None):
"""Start the MCP server.""" """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) try:
if transport == "stdio":
if transport == "stdio": await server.run_stdio_async()
self.console.print("[green]Starting MCP server on stdio...[/green]") elif transport == "sse":
server.run("stdio") await server.run_sse_async("sse")
elif transport == "sse": else:
self.console.print( await server.run_http_async("streamable-http")
"[green]Starting MCP server with streamable HTTP...[/green]" except KeyboardInterrupt:
) pass
server.run("sse") finally:
else: monitor_task.cancel()
self.console.print("[green]Starting MCP server with HTTP...[/green]") try:
server.run("streamable-http") await monitor_task
except asyncio.CancelledError:
pass

View file

@ -146,7 +146,7 @@ def serve(
elif sse: elif sse:
transport = "sse" transport = "sse"
app.serve(transport=transport) event_loop.run_until_complete(app.serve(transport=transport))
if __name__ == "__main__": if __name__ == "__main__":

View file

@ -88,7 +88,7 @@ class HaikuRAG:
if not source_path.exists(): if not source_path.exists():
raise ValueError(f"File does not exist: {source_path}") 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() md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest()
# Check if document already exists # Check if document already exists

View file

@ -2,7 +2,7 @@ import os
from pathlib import Path from pathlib import Path
from dotenv import load_dotenv from dotenv import load_dotenv
from pydantic import BaseModel from pydantic import BaseModel, field_validator
from haiku.rag.utils import get_default_data_dir from haiku.rag.utils import get_default_data_dir
@ -13,6 +13,7 @@ class AppConfig(BaseModel):
ENV: str = "development" ENV: str = "development"
DEFAULT_DATA_DIR: Path = get_default_data_dir() DEFAULT_DATA_DIR: Path = get_default_data_dir()
MONITOR_DIRECTORIES: list[Path] = []
EMBEDDINGS_PROVIDER: str = "ollama" EMBEDDINGS_PROVIDER: str = "ollama"
EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_MODEL: str = "mxbai-embed-large"
@ -23,6 +24,17 @@ class AppConfig(BaseModel):
OLLAMA_BASE_URL: str = "http://localhost:11434" 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 # Expose Config object for app to import
Config = AppConfig.model_validate(os.environ) Config = AppConfig.model_validate(os.environ)

24
src/haiku/rag/logging.py Normal file
View file

@ -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

74
src/haiku/rag/monitor.py Normal file
View file

@ -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}")

View file

@ -98,7 +98,7 @@ async def test_client_create_document_from_source():
assert doc.id is not None assert doc.id is not None
assert doc.content == test_content 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 doc.metadata["source_type"] == "file"
assert "contentType" in doc.metadata assert "contentType" in doc.metadata
assert "md5" 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.id is not None
assert doc2.content == test_content 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 "contentType" in doc2.metadata
assert "md5" in doc2.metadata assert "md5" in doc2.metadata

99
tests/test_monitor.py Normal file
View file

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