Add monitor

This commit is contained in:
Yiorgis Gozadinos 2025-06-20 07:44:03 +02:00
parent 0e4d886b3a
commit 5f604d31dd
No known key found for this signature in database
4 changed files with 88 additions and 12 deletions

View file

@ -1,9 +1,12 @@
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.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 +91,15 @@ 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 monitor = FileWatcher(paths=[])
asyncio.create_task(monitor.observe())
server = create_mcp_server(self.db_path) server = create_mcp_server(self.db_path)
if transport == "stdio": if transport == "stdio":
self.console.print("[green]Starting MCP server on stdio...[/green]") await server.run_stdio_async()
server.run("stdio")
elif transport == "sse": elif transport == "sse":
self.console.print( await server.run_sse_async("sse")
"[green]Starting MCP server with streamable HTTP...[/green]"
)
server.run("sse")
else: else:
self.console.print("[green]Starting MCP server with HTTP...[/green]") await server.run_http_async("streamable-http")
server.run("streamable-http")

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

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

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

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