Add monitor
This commit is contained in:
parent
0e4d886b3a
commit
5f604d31dd
4 changed files with 88 additions and 12 deletions
|
|
@ -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")
|
|
||||||
|
|
|
||||||
|
|
@ -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
24
src/haiku/rag/logging.py
Normal 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
54
src/haiku/rag/monitor.py
Normal 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
|
||||||
Loading…
Reference in a new issue