Merge pull request #67 from ggozad/chore/fast-startup

Late imports for faster startup. Remove obsolete SSE MCP transport.
This commit is contained in:
Yiorgis Gozadinos 2025-09-22 12:42:45 +03:00 committed by GitHub
commit 2150d6a670
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 69 additions and 81 deletions

View file

@ -111,9 +111,6 @@ haiku-rag serve
# stdio transport # stdio transport
haiku-rag serve --stdio haiku-rag serve --stdio
# SSE transport
haiku-rag serve --sse
``` ```
## Settings ## Settings

View file

@ -19,7 +19,7 @@ The MCP server exposes `haiku.rag` as MCP tools for compatible MCP clients.
## Starting MCP Server ## Starting MCP Server
The MCP server starts automatically with the serve command and supports Streamable HTTP, stdio and SSE transports: The MCP server starts automatically with the serve command and supports Streamable HTTP and stdio transports:
```bash ```bash
# Default streamable HTTP transport # Default streamable HTTP transport
@ -27,7 +27,4 @@ haiku-rag serve
# stdio transport (for Claude Desktop) # stdio transport (for Claude Desktop)
haiku-rag serve --stdio haiku-rag serve --stdio
# SSE transport
haiku-rag serve --sse
``` ```

View file

@ -11,7 +11,6 @@ haiku-rag serve
Transport options: Transport options:
- Default - Streamable HTTP transport - Default - Streamable HTTP transport
- `--stdio` - Standard input/output transport - `--stdio` - Standard input/output transport
- `--sse` - Server-sent events transport
## File Monitoring ## File Monitoring

View file

@ -289,8 +289,6 @@ class HaikuRAGApp:
try: try:
if transport == "stdio": if transport == "stdio":
await server.run_stdio_async() await server.run_stdio_async()
elif transport == "sse":
await server.run_sse_async()
else: else:
await server.run_http_async(transport="streamable-http") await server.run_http_async(transport="streamable-http")
except KeyboardInterrupt: except KeyboardInterrupt:

View file

@ -3,28 +3,16 @@ import warnings
from importlib.metadata import version from importlib.metadata import version
from pathlib import Path from pathlib import Path
import logfire
import typer import typer
from rich.console import Console
from haiku.rag.app import HaikuRAGApp
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.logging import configure_cli_logging from haiku.rag.logging import configure_cli_logging
from haiku.rag.migration import migrate_sqlite_to_lancedb
from haiku.rag.utils import is_up_to_date from haiku.rag.utils import is_up_to_date
if Config.ENV == "development":
logfire.configure(send_to_logfire="if-token-present")
logfire.instrument_pydantic_ai()
else:
warnings.filterwarnings("ignore")
cli = typer.Typer( cli = typer.Typer(
context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
) )
console = Console()
def complete_document_ids(ctx: typer.Context, incomplete: str): def complete_document_ids(ctx: typer.Context, incomplete: str):
"""Autocomplete document IDs from the selected DB.""" """Autocomplete document IDs from the selected DB."""
@ -89,16 +77,16 @@ async def check_version():
"""Check if haiku.rag is up to date and show warning if not.""" """Check if haiku.rag is up to date and show warning if not."""
up_to_date, current_version, latest_version = await is_up_to_date() up_to_date, current_version, latest_version = await is_up_to_date()
if not up_to_date: if not up_to_date:
console.print( typer.echo(
f"[yellow]Warning: haiku.rag is outdated. Current: {current_version}, Latest: {latest_version}[/yellow]" f"Warning: haiku.rag is outdated. Current: {current_version}, Latest: {latest_version}",
) )
console.print("[yellow]Please update.[/yellow]") typer.echo("Please update.")
def version_callback(value: bool): def version_callback(value: bool):
if value: if value:
v = version("haiku.rag") v = version("haiku.rag")
console.print(f"haiku.rag version {v}") typer.echo(f"haiku.rag version {v}")
raise typer.Exit() raise typer.Exit()
@ -113,10 +101,26 @@ def main(
), ),
): ):
"""haiku.rag CLI - Vector database RAG system""" """haiku.rag CLI - Vector database RAG system"""
# Ensure only haiku.rag logs are emitted in CLI context # Configure logging minimally for CLI context
configure_cli_logging() if Config.ENV == "development":
# Lazy import logfire only in development
try:
import logfire # type: ignore
logfire.configure(send_to_logfire="if-token-present")
logfire.instrument_pydantic_ai()
except Exception:
pass
else:
configure_cli_logging()
warnings.filterwarnings("ignore")
# Run version check before any command # Run version check before any command
asyncio.run(check_version()) try:
asyncio.run(check_version())
except Exception:
# Do not block CLI on version check issues
pass
@cli.command("list", help="List all stored documents") @cli.command("list", help="List all stored documents")
@ -127,6 +131,8 @@ def list_documents(
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.list_documents()) asyncio.run(app.list_documents())
@ -142,6 +148,8 @@ def add_document_text(
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.add_document_from_text(text=text)) asyncio.run(app.add_document_from_text(text=text))
@ -158,6 +166,8 @@ def add_document_src(
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.add_document_from_source(source=source)) asyncio.run(app.add_document_from_source(source=source))
@ -174,6 +184,8 @@ def get_document(
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.get_document(doc_id=doc_id)) asyncio.run(app.get_document(doc_id=doc_id))
@ -190,6 +202,8 @@ def delete_document(
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.delete_document(doc_id=doc_id)) asyncio.run(app.delete_document(doc_id=doc_id))
@ -215,6 +229,8 @@ def search(
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.search(query=query, limit=limit)) asyncio.run(app.search(query=query, limit=limit))
@ -235,6 +251,8 @@ def ask(
help="Include citations in the response", help="Include citations in the response",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.ask(question=question, cite=cite)) asyncio.run(app.ask(question=question, cite=cite))
@ -271,6 +289,8 @@ def research(
help="Show verbose progress output", help="Show verbose progress output",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run( asyncio.run(
app.research( app.research(
@ -285,6 +305,8 @@ def research(
@cli.command("settings", help="Display current configuration settings") @cli.command("settings", help="Display current configuration settings")
def settings(): def settings():
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=Path()) # Don't need actual DB for settings app = HaikuRAGApp(db_path=Path()) # Don't need actual DB for settings
app.show_settings() app.show_settings()
@ -300,6 +322,8 @@ def rebuild(
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.rebuild()) asyncio.run(app.rebuild())
@ -312,6 +336,8 @@ def vacuum(
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
): ):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.vacuum()) asyncio.run(app.vacuum())
@ -330,24 +356,15 @@ def serve(
"--stdio", "--stdio",
help="Run MCP server on stdio Transport", help="Run MCP server on stdio Transport",
), ),
sse: bool = typer.Option(
False,
"--sse",
help="Run MCP server on SSE transport",
),
) -> None: ) -> None:
"""Start the MCP server.""" """Start the MCP server."""
if stdio and sse: from haiku.rag.app import HaikuRAGApp
console.print("[red]Error: Cannot use both --stdio and --http options[/red]")
raise typer.Exit(1)
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
transport = None transport = None
if stdio: if stdio:
transport = "stdio" transport = "stdio"
elif sse:
transport = "sse"
asyncio.run(app.serve(transport=transport)) asyncio.run(app.serve(transport=transport))
@ -361,6 +378,9 @@ def migrate(
# Generate LanceDB path in same parent directory # Generate LanceDB path in same parent directory
lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb") lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
# Lazy import to avoid heavy deps on simple invocations
from haiku.rag.migration import migrate_sqlite_to_lancedb
success = asyncio.run(migrate_sqlite_to_lancedb(sqlite_path, lancedb_path)) success = asyncio.run(migrate_sqlite_to_lancedb(sqlite_path, lancedb_path))
if not success: if not success:

View file

@ -9,10 +9,6 @@ from io import BytesIO
from pathlib import Path from pathlib import Path
from types import ModuleType from types import ModuleType
import httpx
from docling.document_converter import DocumentConverter
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.io import DocumentStream
from packaging.version import Version, parse from packaging.version import Version, parse
@ -82,6 +78,9 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
the running version and the latest version. the running version and the latest version.
""" """
# Lazy import to avoid pulling httpx (and its deps) on module import
import httpx
async with httpx.AsyncClient() as client: async with httpx.AsyncClient() as client:
running_version = parse(metadata.version("haiku.rag")) running_version = parse(metadata.version("haiku.rag"))
try: try:
@ -94,7 +93,7 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
return running_version >= pypi_version, running_version, pypi_version return running_version >= pypi_version, running_version, pypi_version
def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocument: def text_to_docling_document(text: str, name: str = "content.md"):
"""Convert text content to a DoclingDocument. """Convert text content to a DoclingDocument.
Args: Args:
@ -104,6 +103,10 @@ def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocu
Returns: Returns:
A DoclingDocument created from the text content. A DoclingDocument created from the text content.
""" """
# Lazy import docling deps to keep import-time light
from docling.document_converter import DocumentConverter # type: ignore
from docling_core.types.io import DocumentStream # type: ignore
bytes_io = BytesIO(text.encode("utf-8")) bytes_io = BytesIO(text.encode("utf-8"))
doc_stream = DocumentStream(name=name, stream=bytes_io) doc_stream = DocumentStream(name=name, stream=bytes_io)
converter = DocumentConverter() converter = DocumentConverter()

View file

@ -172,7 +172,7 @@ async def test_search_no_results(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["stdio", "sse", "http", None]) @pytest.mark.parametrize("transport", ["stdio", "http", None])
async def test_serve(app: HaikuRAGApp, monkeypatch, transport): async def test_serve(app: HaikuRAGApp, monkeypatch, transport):
"""Test the serve method with different transports.""" """Test the serve method with different transports."""
mock_server = AsyncMock() mock_server = AsyncMock()
@ -199,8 +199,6 @@ async def test_serve(app: HaikuRAGApp, monkeypatch, transport):
if transport == "stdio": if transport == "stdio":
mock_server.run_stdio_async.assert_called_once() mock_server.run_stdio_async.assert_called_once()
elif transport == "sse":
mock_server.run_sse_async.assert_called_once()
else: else:
mock_server.run_http_async.assert_called_once_with(transport="streamable-http") mock_server.run_http_async.assert_called_once_with(transport="streamable-http")

View file

@ -8,7 +8,7 @@ runner = CliRunner()
def test_list_documents(): def test_list_documents():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.list_documents = AsyncMock() mock_app_instance.list_documents = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -20,7 +20,7 @@ def test_list_documents():
def test_add_document_text(): def test_add_document_text():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.add_document_from_text = AsyncMock() mock_app_instance.add_document_from_text = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -34,7 +34,7 @@ def test_add_document_text():
def test_add_document_src(): def test_add_document_src():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.add_document_from_source = AsyncMock() mock_app_instance.add_document_from_source = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -46,7 +46,7 @@ def test_add_document_src():
def test_get_document(): def test_get_document():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.get_document = AsyncMock() mock_app_instance.get_document = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -58,7 +58,7 @@ def test_get_document():
def test_delete_document(): def test_delete_document():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.delete_document = AsyncMock() mock_app_instance.delete_document = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -70,7 +70,7 @@ def test_delete_document():
def test_search(): def test_search():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.search = AsyncMock() mock_app_instance.search = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -82,7 +82,7 @@ def test_search():
def test_serve(): def test_serve():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.serve = AsyncMock() mock_app_instance.serve = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -94,7 +94,7 @@ def test_serve():
def test_serve_stdio(): def test_serve_stdio():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.serve = AsyncMock() mock_app_instance.serve = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -105,32 +105,8 @@ def test_serve_stdio():
mock_app_instance.serve.assert_called_once_with(transport="stdio") mock_app_instance.serve.assert_called_once_with(transport="stdio")
def test_serve_sse():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock()
mock_app_instance.serve = AsyncMock()
mock_app.return_value = mock_app_instance
result = runner.invoke(cli, ["serve", "--sse"])
assert result.exit_code == 0
mock_app_instance.serve.assert_called_once_with(transport="sse")
def test_serve_stdio_and_sse():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock()
mock_app_instance.serve = AsyncMock()
mock_app.return_value = mock_app_instance
result = runner.invoke(cli, ["serve", "--stdio", "--sse"])
assert result.exit_code == 1
assert "Error: Cannot use both --stdio and --http options" in result.stdout
def test_ask(): def test_ask():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.ask = AsyncMock() mock_app_instance.ask = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance
@ -144,7 +120,7 @@ def test_ask():
def test_ask_with_cite(): def test_ask_with_cite():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock() mock_app_instance = MagicMock()
mock_app_instance.ask = AsyncMock() mock_app_instance.ask = AsyncMock()
mock_app.return_value = mock_app_instance mock_app.return_value = mock_app_instance