haiku.rag/haiku_rag_slim/haiku/rag/app.py

import logging
from pathlib import Path
from typing import TYPE_CHECKING

from rich.console import Console
from rich.markdown import Markdown
from rich.markup import escape
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    SpinnerColumn,
    TaskID,
    TextColumn,
    TransferSpeedColumn,
)

from haiku.rag.client import HaikuRAG, RebuildMode
from haiku.rag.config import AppConfig, get_config
from haiku.rag.mcp import create_mcp_server
from haiku.rag.store.models.chunk import SearchType
from haiku.rag.store.models.document import Document

if TYPE_CHECKING:
    from haiku.rag.store.engine import Store
    from haiku.rag.store.models import SearchResult
from haiku.rag.config import redact_secrets
from haiku.rag.utils import format_bytes, format_citations_rich

logger = logging.getLogger(__name__)


class HaikuRAGApp:
    def __init__(
        self,
        db_path: Path,
        config: AppConfig | None = None,
        read_only: bool = False,
    ):
        self.db_path = db_path
        self.config = config if config is not None else get_config()
        self.read_only = read_only
        self.console = Console()

        from haiku.rag.store.engine import ConnectionMode

        self._is_local = ConnectionMode.from_config(self.config) == ConnectionMode.LOCAL
        self._display_path = self.db_path if self._is_local else self.config.lancedb.uri

    async def init(self):
        """Initialize a new database."""
        if self._is_local and self.db_path.exists():
            self.console.print(
                f"[yellow]Database already exists at {self.db_path}[/yellow]"
            )
            return

        # Create the database
        async with HaikuRAG(db_path=self.db_path, config=self.config, create=True):
            pass
        self.console.print(
            f"[bold green]Database initialized at {self._display_path}[/bold green]"
        )

    async def info(self):
        """Display read-only information about the database without modifying it."""

        from haiku.rag.store.info import gather_database_info

        # Basic: show path/URI
        self.console.print("[bold]haiku.rag database info[/bold]")
        self.console.print(
            f"  [repr.attrib_name]path[/repr.attrib_name]: {self._display_path}"
        )

        if self._is_local and not self.db_path.exists():
            self.console.print("[red]Database path does not exist.[/red]")
            return

        info = await gather_database_info(self.config, self.db_path)

        if not info.exists:
            self.console.print(
                "[red]Database is empty. Use 'haiku-rag init' to initialize.[/red]"
            )
            return

        self.console.print(
            f"  [repr.attrib_name]haiku.rag version (db)[/repr.attrib_name]: {info.stored_version}"
        )
        dim_part = (
            f"{info.embeddings.vector_dim}"
            if info.embeddings.vector_dim is not None
            else "unknown"
        )
        self.console.print(
            "  [repr.attrib_name]embeddings[/repr.attrib_name]: "
            f"{info.embeddings.provider}/{info.embeddings.name} (dim: {dim_part})"
        )

        tables = {t.name: t for t in info.tables}

        # Per-table row counts and sizes. Missing required tables are
        # reported as "absent" rather than raising.
        for name in ("documents", "document_meta", "chunks", "document_items"):
            entry = tables[name]
            if entry.exists:
                self.console.print(
                    f"  [repr.attrib_name]{name}[/repr.attrib_name]: {entry.num_rows} "
                    f"({format_bytes(entry.total_bytes)})"
                )
            else:
                self.console.print(
                    f"  [repr.attrib_name]{name}[/repr.attrib_name]: [yellow]absent[/yellow]"
                )

        # Vector index information
        if tables["chunks"].exists:
            num_chunks = tables["chunks"].num_rows
            if info.vector_index.exists:
                self.console.print(
                    "  [repr.attrib_name]vector index[/repr.attrib_name]: ✓ exists"
                )
                self.console.print(
                    f"  [repr.attrib_name]indexed chunks[/repr.attrib_name]: {info.vector_index.indexed_rows}"
                )
                if info.vector_index.unindexed_rows > 0:
                    self.console.print(
                        f"  [repr.attrib_name]unindexed chunks[/repr.attrib_name]: [yellow]{info.vector_index.unindexed_rows}[/yellow] "
                        "(consider running: haiku-rag create-index)"
                    )
                else:
                    self.console.print(
                        f"  [repr.attrib_name]unindexed chunks[/repr.attrib_name]: {info.vector_index.unindexed_rows}"
                    )
            else:
                if num_chunks >= 256:
                    self.console.print(
                        "  [repr.attrib_name]vector index[/repr.attrib_name]: [yellow]✗ not created[/yellow] "
                        "(run: haiku-rag create-index)"
                    )
                else:
                    self.console.print(
                        f"  [repr.attrib_name]vector index[/repr.attrib_name]: ✗ not created "
                        f"(need {256 - num_chunks} more chunks)"
                    )

        if tables["documents"].exists:
            self.console.print(
                f"  [repr.attrib_name]versions (documents)[/repr.attrib_name]: "
                f"{tables['documents'].num_versions}"
            )
        if tables["document_meta"].exists:
            self.console.print(
                f"  [repr.attrib_name]versions (document_meta)[/repr.attrib_name]: "
                f"{tables['document_meta'].num_versions}"
            )
        if tables["chunks"].exists:
            self.console.print(
                f"  [repr.attrib_name]versions (chunks)[/repr.attrib_name]: "
                f"{tables['chunks'].num_versions}"
            )

        # Migration status
        self.console.rule()
        if info.pending_migrations:
            self.console.print(
                f"[bold yellow]{len(info.pending_migrations)} migration(s) pending.[/bold yellow] "
                "Run [cyan]haiku-rag migrate[/cyan] to upgrade."
            )
            for step in info.pending_migrations:
                self.console.print(
                    f"  [yellow]→[/yellow] {step.version}: {step.description}"
                )
        else:
            self.console.print("[green]Database is up to date.[/green]")

        self.console.rule()
        self.console.print("[bold]Versions[/bold]")
        self.console.print(
            f"  [repr.attrib_name]haiku.rag[/repr.attrib_name]: {info.packages['haiku_rag']}"
        )
        self.console.print(
            f"  [repr.attrib_name]lancedb[/repr.attrib_name]: {info.packages['lancedb']}"
        )
        self.console.print(
            f"  [repr.attrib_name]docling[/repr.attrib_name]: {info.packages['docling']}"
        )
        self.console.print(
            f"  [repr.attrib_name]pydantic-ai[/repr.attrib_name]: {info.packages['pydantic_ai']}"
        )
        self.console.print(
            f"  [repr.attrib_name]docling-document schema[/repr.attrib_name]: {info.packages['docling_document_schema']}"
        )

    async def doctor(self, duplicates_out: Path | None = None) -> bool:
        """Run health checks and print a report. Returns True if any check failed."""
        import os
        from contextlib import nullcontext

        from haiku.rag.doctor import Severity, run_doctor

        self.console.print("[bold]haiku.rag doctor[/bold]")
        self.console.print(
            f"  [repr.attrib_name]path[/repr.attrib_name]: {self._display_path}"
        )

        if self._is_local and not self.db_path.exists():
            self.console.print("[red]Database path does not exist.[/red]")
            return True

        status = (
            self.console.status("Running checks") if self.console.is_terminal else None
        )

        def on_progress(label: str) -> None:
            if status is not None:
                status.update(f"{label}...")

        cm = status if status is not None else nullcontext()
        with cm:
            report = await run_doctor(
                self.config,
                self.db_path,
                dict(os.environ),
                duplicates_out=duplicates_out,
                on_progress=on_progress,
            )

        glyphs = {
            Severity.OK: "[green]✓[/green]",
            Severity.WARN: "[yellow]![/yellow]",
            Severity.FAIL: "[red]✗[/red]",
        }

        def render(result):
            self.console.print(f"{glyphs[result.severity]} {result.message}")
            for detail in result.details:
                self.console.print(f"    [dim]{detail}[/dim]")
            if result.remediation:
                self.console.print(f"    [dim]→ {result.remediation}[/dim]")

        database = [r for r in report.results if not r.name.startswith("provider:")]
        providers = [r for r in report.results if r.name.startswith("provider:")]

        self.console.rule("[bold]Database[/bold]")
        for result in database:
            render(result)
        if providers:
            self.console.rule("[bold]Providers[/bold]")
            for result in providers:
                render(result)

        self.console.rule()
        self.console.print(
            f"[green]{report.count(Severity.OK)} ok[/green], "
            f"[yellow]{report.count(Severity.WARN)} warning(s)[/yellow], "
            f"[red]{report.count(Severity.FAIL)} failure(s)[/red]"
        )
        if duplicates_out is not None:
            self.console.print(
                f"[dim]Duplicate-document groups written to {duplicates_out}[/dim]"
            )
        return report.failed

    async def history(self, table: str | None = None, limit: int | None = None):
        """Display version history for database tables.

        Args:
            table: Specific table to show history for (documents, chunks, settings).
                   If None, shows history for all tables.
            limit: Maximum number of versions to show per table.
        """
        from haiku.rag.store.engine import Store

        if self._is_local and not self.db_path.exists():
            self.console.print("[red]Database path does not exist.[/red]")
            return

        async with Store(
            self.db_path,
            config=self.config,
            skip_validation=True,
            read_only=True,
            skip_migration_check=True,
        ) as store:
            tables = [
                "documents",
                "document_meta",
                "chunks",
                "document_items",
                "settings",
            ]
            if table:
                if table not in tables:
                    self.console.print(
                        f"[red]Unknown table: {table}. Must be one of: {', '.join(tables)}[/red]"
                    )
                    return
                tables = [table]

            self.console.print("[bold]Version History[/bold]")

            try:
                tags = await store.list_tags()
            except Exception as exc:
                tags = {}
                self.console.print(
                    f"[yellow]Tag annotations unavailable: {escape(str(exc))}[/yellow]"
                )

            for table_name in tables:
                versions = await store.list_table_versions(table_name)

                # Sort by version descending (newest first)
                versions = sorted(versions, key=lambda v: v["version"], reverse=True)

                if limit:
                    versions = versions[:limit]

                version_tags: dict[int, list[str]] = {}
                for tag_name, info in tags.items():
                    tagged_version = info.tables.get(table_name)
                    if tagged_version is not None:
                        version_tags.setdefault(tagged_version, []).append(tag_name)

                self.console.print(f"\n[bold cyan]{table_name}[/bold cyan]")

                if not versions:
                    self.console.print("  [dim]No versions found[/dim]")
                    continue

                for v in versions:
                    version_num = v["version"]
                    timestamp = v["timestamp"]
                    suffix = ""
                    if version_num in version_tags:
                        names = ", ".join(
                            escape(n) for n in sorted(version_tags[version_num])
                        )
                        suffix = f"  [magenta]<- {names}[/magenta]"
                    self.console.print(
                        f"  [repr.attrib_name]v{version_num}[/repr.attrib_name]: {timestamp}{suffix}"
                    )

    def _tag_write_store(self) -> "Store":
        """Writable store for tag create/delete with normal validation and
        migration checks.

        A coordinated tag is only reliable when the database schema is
        current, and a writable open of a legacy database would create
        missing tables as a side effect.
        """
        from haiku.rag.store.engine import Store

        return Store(self.db_path, config=self.config, read_only=self.read_only)

    def _tag_read_store(self) -> "Store":
        """Read-only store for tag inspection; works on old or drifted DBs."""
        from haiku.rag.store.engine import Store

        return Store(
            self.db_path,
            config=self.config,
            skip_validation=True,
            skip_migration_check=True,
            read_only=True,
        )

    async def create_tag(self, name: str):
        """Tag the current version of every table."""
        if self._is_local and not self.db_path.exists():
            raise ValueError(f"Database path does not exist: {self.db_path}")
        async with self._tag_write_store() as store:
            await store.create_tag(name)
        self.console.print(f"[green]Created tag '{escape(name)}'[/green]")

    async def list_tags(self):
        """List database tags, flagging partial ones."""
        if self._is_local and not self.db_path.exists():
            raise ValueError(f"Database path does not exist: {self.db_path}")
        async with self._tag_read_store() as store:
            tags = await store.list_tags()

        if not tags:
            self.console.print("No tags")
            return

        self.console.print("[bold]Tags[/bold]")
        for name in sorted(tags):
            info = tags[name]
            versions = " ".join(f"{t}=v{v}" for t, v in info.tables.items())
            line = f"  [repr.attrib_name]{escape(name)}[/repr.attrib_name]: {versions}"
            if not info.complete:
                missing = ", ".join(info.missing_tables)
                line += f" [yellow](partial - missing: {missing})[/yellow]"
            self.console.print(line)

    async def delete_tag(self, name: str):
        """Delete a tag from every table that has it."""
        if self._is_local and not self.db_path.exists():
            raise ValueError(f"Database path does not exist: {self.db_path}")
        async with self._tag_write_store() as store:
            await store.delete_tag(name)
        self.console.print(f"[green]Deleted tag '{escape(name)}'[/green]")

    async def restore_tag(self, name: str):
        """Restore the database to a tagged state and report the outcome.

        The Store context exits before anything is printed; no high-level
        database access happens after the restore.

        Raises:
            ValueError: If the database path does not exist.
        """
        if self._is_local and not self.db_path.exists():
            raise ValueError(f"Database path does not exist: {self.db_path}")
        async with self._tag_write_store() as store:
            safety_tag = await store.restore_tag(name)
        self.console.print(f"[green]Restored database to tag '{escape(name)}'.[/green]")
        self.console.print(
            f"The previous state is preserved as '{escape(safety_tag)}'."
        )
        self.console.print(
            "The restored state is now live. Later historical versions remain "
            "until eligible for vacuum. Run [cyan]haiku-rag migrate[/cyan] if "
            "migration is required."
        )

    async def list_documents(self, filter: str | None = None):
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=True,
            skip_validation=True,
        ) as self.client:
            documents = await self.client.list_documents(filter=filter)
            for doc in documents:
                self._rich_print_document(doc, truncate=True)

    async def add_document_from_text(
        self, text: str, title: str | None = None, metadata: dict | None = None
    ):
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=self.read_only,
        ) as self.client:
            doc = await self.client.create_document(
                text, title=title, metadata=metadata
            )
            self._rich_print_document(doc, truncate=True)
            self.console.print(
                f"[bold green]Document {doc.id} added successfully.[/bold green]"
            )

    async def add_document_from_source(
        self, source: str, title: str | None = None, metadata: dict | None = None
    ):
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=self.read_only,
        ) as self.client:
            result = await self.client.create_document_from_source(
                source, title=title, metadata=metadata
            )
            if isinstance(result, list):
                for doc in result:
                    self._rich_print_document(doc, truncate=True)
                self.console.print(
                    f"[bold green]{len(result)} documents added successfully.[/bold green]"
                )
            else:
                self._rich_print_document(result, truncate=True)
                self.console.print(
                    f"[bold green]Document {result.id} added successfully.[/bold green]"
                )

    async def get_document(self, doc_id: str):
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=True,
            skip_validation=True,
        ) as self.client:
            doc = await self.client.get_document_by_id(doc_id)
            if doc is None:
                self.console.print(f"[red]Document with id {doc_id} not found.[/red]")
                return
            self._rich_print_document(doc, truncate=False)

    async def delete_document(self, doc_id: str):
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=self.read_only,
            skip_validation=True,
        ) as self.client:
            deleted = await self.client.delete_document(doc_id)
            if deleted:
                self.console.print(
                    f"[bold green]Document {doc_id} deleted successfully.[/bold green]"
                )
            else:
                self.console.print(
                    f"[yellow]Document with id {doc_id} not found.[/yellow]"
                )

    async def search(
        self,
        query: str | None = None,
        limit: int | None = None,
        filter: str | None = None,
        search_type: SearchType | None = None,
        image: Path | None = None,
    ):
        if query is None and image is None:
            self.console.print(
                "[red]Provide either a query argument or --image PATH.[/red]"
            )
            return
        if query is not None and image is not None:
            self.console.print("[red]Pass either a query or --image, not both.[/red]")
            return

        if query is None and search_type is not None:
            self.console.print("[red]Pass --search-type only for text queries[/red]")
            return

        search_input: str | bytes
        if image is not None:
            search_input = image.read_bytes()
        else:
            assert query is not None
            search_input = query

        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=True,
        ) as self.client:
            results = await self.client.search(
                search_input,
                limit=limit,
                filter=filter,
                search_type=search_type,
            )
            if not results:
                self.console.print("[yellow]No results found.[/yellow]")
                return
            for result in results:
                self._rich_print_search_result(result)

    async def visualize_chunk(self, chunk_id: str, expand: bool = True):
        """Display visual grounding images for a chunk."""
        from textual_image.renderable import Image as RichImage

        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=True,
            skip_validation=True,
        ) as self.client:
            chunk = await self.client.get_chunk_by_id(chunk_id)
            if not chunk:
                self.console.print(f"[red]Chunk with id {chunk_id} not found.[/red]")
                return

            images = await self.client.visualize_chunk(chunk, expand=expand)
            if not images:
                self.console.print(
                    "[yellow]No visual grounding available for this chunk.[/yellow]"
                )
                self.console.print(
                    "This may be because the document was converted without page images."
                )
                return

            self.console.print(f"[bold]Visual grounding for chunk {chunk_id}[/bold]")
            if chunk.document_uri:
                self.console.print(
                    f"[repr.attrib_name]document[/repr.attrib_name]: {chunk.document_uri}"
                )

            for i, img in enumerate(images):
                self.console.print(
                    f"\n[bold cyan]Page {i + 1}/{len(images)}[/bold cyan]"
                )
                self.console.print(RichImage(img))

    async def ask(
        self,
        question: str,
        filter: str | None = None,
        images: list[Path] | None = None,
    ):
        """Ask a question using the RAG system.

        Args:
            question: The question to ask
            filter: SQL WHERE clause to filter documents
            images: Paths of images to attach to the question
        """
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=True,
        ) as self.client:
            answer, citations = await self.client.ask(
                question,
                filter=filter,
                images=[path.read_bytes() for path in images] if images else None,
            )

            self.console.print(f"[bold blue]Question:[/bold blue] {question}")
            self.console.print()
            self.console.print("[bold green]Answer:[/bold green]")
            self.console.print(Markdown(answer))
            for renderable in await format_citations_rich(
                citations, client=self.client
            ):
                self.console.print(renderable)

    async def analyze(
        self,
        question: str,
        filter: str | None = None,
        images: list[Path] | None = None,
    ):
        """Answer a question using the analysis capability.

        Args:
            question: The question to answer
            filter: SQL WHERE clause to filter documents
            images: Paths of images to attach to the question
        """
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            read_only=True,
        ) as self.client:
            self.console.print(f"[bold blue]Question:[/bold blue] {question}")
            self.console.print()
            self.console.print(
                "[dim]Running analysis capability with code execution...[/dim]"
            )
            self.console.print()

            result = await self.client.analyze(
                question,
                filter=filter,
                images=[path.read_bytes() for path in images] if images else None,
            )

            self.console.print("[bold green]Answer:[/bold green]")
            self.console.print(Markdown(result.answer))
            for renderable in await format_citations_rich(
                result.citations, client=self.client
            ):
                self.console.print(renderable)

    async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            skip_validation=True,
            read_only=self.read_only,
        ) as client:
            if mode == RebuildMode.SET_EMBEDDER:
                async for _ in client.rebuild_database(mode=mode):
                    pass
                self.console.print(
                    "[bold green]Stored embedder settings updated.[/bold green]"
                )
                return

            documents = await client.list_documents()
            total_docs = len(documents)

            if total_docs == 0:
                self.console.print("[yellow]No documents found in database.[/yellow]")
                return

            mode_desc = {
                RebuildMode.FULL: "full rebuild",
                RebuildMode.RECHUNK: "rechunk",
                RebuildMode.EMBED_ONLY: "embed only",
                RebuildMode.TITLE_ONLY: "title only",
                RebuildMode.DESCRIPTIONS: "picture descriptions",
            }[mode]

            self.console.print(
                f"[bold cyan]Rebuilding database ({mode_desc}) with {total_docs} documents...[/bold cyan]"
            )
            with Progress() as progress:
                task = progress.add_task("Rebuilding...", total=total_docs)
                async for _ in client.rebuild_database(mode=mode):
                    progress.update(task, advance=1)

            self.console.print(
                "[bold green]Database rebuild completed successfully.[/bold green]"
            )

    async def vacuum(self):
        """Run database maintenance: optimize and cleanup table history."""
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            skip_validation=True,
            read_only=self.read_only,
        ) as client:
            await client.vacuum()
        self.console.print("[bold green]Vacuum completed successfully.[/bold green]")

    async def migrate(self) -> list[str]:
        """Run pending database migrations.

        Returns:
            List of descriptions of applied migrations.
        """
        from haiku.rag.store.engine import Store

        async with Store(
            self.db_path,
            config=self.config,
            skip_validation=True,
            skip_migration_check=True,
            read_only=self.read_only,
        ) as store:
            return await store.migrate()

    async def create_index(self):
        """Create vector index on the chunks table."""
        async with HaikuRAG(
            db_path=self.db_path,
            config=self.config,
            skip_validation=True,
            read_only=self.read_only,
        ) as client:
            row_count = await client.store.chunks_table.count_rows()
            self.console.print(f"Chunks in database: {row_count}")

            if row_count < 256:
                self.console.print(
                    f"[yellow]Warning: Need at least 256 chunks to create an index (have {row_count})[/yellow]"
                )
                return

            # Check if index already exists
            indices = await client.store.chunks_table.list_indices()
            has_vector_index = any("vector" in str(idx).lower() for idx in indices)

            if has_vector_index:
                self.console.print(
                    "[yellow]Rebuilding existing vector index...[/yellow]"
                )
            else:
                self.console.print("[bold]Creating vector index...[/bold]")

            await client.store._ensure_vector_index()
            self.console.print(
                "[bold green]Vector index created successfully.[/bold green]"
            )

    async def download_models(self):
        """Download Docling, HuggingFace tokenizer, and Ollama models per config."""
        from haiku.rag.client.downloads import download_models

        progress: Progress | None = None
        task_id: TaskID | None = None
        current_model = ""
        current_digest = ""

        async for event in download_models(self.config):
            if event.status == "start":
                self.console.print(
                    f"[bold blue]Downloading {event.model}...[/bold blue]"
                )
            elif event.status == "done":
                if progress:
                    progress.stop()
                    progress = None
                    task_id = None
                self.console.print(f"[green]✓[/green] {event.model}")
                current_model = ""
                current_digest = ""
            elif event.status == "pulling":
                self.console.print(f"[bold blue]Pulling {event.model}...[/bold blue]")
                current_model = event.model
                progress = Progress(
                    SpinnerColumn(),
                    TextColumn("[progress.description]{task.description}"),
                    BarColumn(),
                    DownloadColumn(),
                    TransferSpeedColumn(),
                    console=self.console,
                    transient=True,
                    auto_refresh=False,
                )
                progress.start()
                task_id = progress.add_task(event.model, total=None)
            elif event.status == "downloading" and progress and task_id is not None:
                if event.digest != current_digest:
                    current_digest = event.digest
                    short_digest = event.digest[:19] if event.digest else ""
                    progress.update(
                        task_id,
                        description=f"{current_model} ({short_digest})",
                        total=event.total,
                        completed=0,
                    )
                progress.update(task_id, completed=event.completed, refresh=True)
            elif progress and task_id is not None:
                progress.update(
                    task_id,
                    description=f"{current_model}: {event.status}",
                    refresh=True,
                )

    def show_settings(self):
        """Display current configuration settings."""
        self.console.print("[bold]haiku.rag configuration[/bold]")
        self.console.print()

        # redact_secrets walks the whole dump: masking only top-level names left
        # nested api keys, tokens and source passwords printed in full.
        for field_name, field_value in redact_secrets(self.config.model_dump()).items():
            self.console.print(
                f"  [repr.attrib_name]{field_name}[/repr.attrib_name]: {field_value}"
            )

    def _rich_print_document(self, doc: Document, truncate: bool = False):
        """Format a document for display."""
        if truncate:
            content = doc.content.splitlines()
            if len(content) > 3:
                content = content[:3] + ["\n…"]
            content = "\n".join(content)
            content = Markdown(content)
        else:
            content = Markdown(doc.content)
        title_part = (
            f" [repr.attrib_name]title[/repr.attrib_name]: {doc.title}"
            if doc.title
            else ""
        )
        self.console.print(
            f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} "
            f"[repr.attrib_name]uri[/repr.attrib_name]: {doc.uri}"
            + title_part
            + f" [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}"
        )
        self.console.print(
            f"[repr.attrib_name]created at[/repr.attrib_name]: {doc.created_at} [repr.attrib_name]updated at[/repr.attrib_name]: {doc.updated_at}"
        )
        self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
        self.console.print(content)
        self.console.rule()

    def _rich_print_search_result(self, result: "SearchResult"):
        """Format a search result for display."""
        content = Markdown(result.content)
        self.console.print(
            f"[repr.attrib_name]document_id[/repr.attrib_name]: {result.document_id} "
            f"[repr.attrib_name]chunk_id[/repr.attrib_name]: {result.chunk_id} "
            f"[repr.attrib_name]score[/repr.attrib_name]: {result.score:.4f}"
        )
        if result.document_uri:
            self.console.print(
                f"[repr.attrib_name]document uri[/repr.attrib_name]: {result.document_uri}"
            )
        if result.document_title:
            self.console.print("[repr.attrib_name]document title[/repr.attrib_name]:")
            self.console.print(result.document_title)
        if result.page_numbers:
            self.console.print("[repr.attrib_name]pages[/repr.attrib_name]:")
            self.console.print(", ".join(str(p) for p in result.page_numbers))
        if result.headings:
            self.console.print("[repr.attrib_name]headings[/repr.attrib_name]:")
            self.console.print(" > ".join(result.headings))
        self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
        self.console.print(content)
        self.console.rule()

    async def run_mcp(
        self,
        transport: str | None = None,
        host: str = "127.0.0.1",
        port: int = 8001,
    ):
        """Run the MCP server until interrupted."""
        async with HaikuRAG(
            self.db_path,
            config=self.config,
            read_only=self.read_only,
        ):
            server = create_mcp_server(
                self.db_path, config=self.config, read_only=self.read_only
            )
            try:
                if transport == "stdio":
                    await server.run_stdio_async()
                else:
                    logger.info(f"Starting MCP server on {host}:{port}")
                    await server.run_http_async(
                        transport="streamable-http", host=host, port=port
                    )
            except KeyboardInterrupt:
                pass