haiku.rag/haiku_rag_slim/haiku/rag/cli.py

import asyncio
import json
import sys
import warnings
from datetime import datetime
from importlib.metadata import version
from pathlib import Path
from typing import Any

import typer
from dotenv import find_dotenv, load_dotenv

# Load environment variables from .env file for API keys and service URLs.
# Env loading needs to be before config import; usecwd=True searches from cwd, not this .py file's location
load_dotenv(find_dotenv(usecwd=True))

from haiku.rag.app import HaikuRAGApp  # noqa: E402
from haiku.rag.config import (  # noqa: E402
    AppConfig,
    find_config_file,
    get_config,
    load_yaml_config,
    set_config,
)
from haiku.rag.logging import configure_cli_logging  # noqa: E402
from haiku.rag.store.exceptions import (  # noqa: E402
    MigrationRequiredError,
    ReadOnlyError,
)
from haiku.rag.utils import is_up_to_date  # noqa: E402

_cli = typer.Typer(
    context_settings={"help_option_names": ["-h", "--help"]},
    no_args_is_help=True,
    pretty_exceptions_show_locals=False,
)


def cli():
    try:
        _cli()
    except (MigrationRequiredError, ReadOnlyError) as e:
        typer.echo(f"Error: {e}", err=True)
        sys.exit(1)


# Module-level flags set by callback
_read_only: bool = False
_before: datetime | None = None


def create_app(db: Path | None = None) -> HaikuRAGApp:  # pragma: no cover
    """Create HaikuRAGApp with loaded config and resolved database path.

    Args:
        db: Optional database path. If None, uses path from config.

    Returns:
        HaikuRAGApp instance with proper config and db path.
    """
    config = get_config()
    db_path = db if db else config.storage.data_dir / "haiku.rag.lancedb"
    return HaikuRAGApp(
        db_path=db_path, config=config, read_only=_read_only, before=_before
    )


async def check_version():  # pragma: no cover
    """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()
    if not up_to_date:
        typer.echo(
            f"Warning: haiku.rag is outdated. Current: {current_version}, Latest: {latest_version}",
        )
        typer.echo("Please update.")


def version_callback(value: bool):  # pragma: no cover
    if value:
        v = version("haiku.rag-slim")
        typer.echo(f"haiku.rag version {v}")
        raise typer.Exit()


@_cli.callback()
def main(
    _version: bool = typer.Option(
        False,
        "-v",
        "--version",
        callback=version_callback,
        help="Show version and exit",
    ),
    config: Path | None = typer.Option(
        None,
        "--config",
        help="Path to YAML configuration file",
    ),
    read_only: bool = typer.Option(
        False,
        "--read-only",
        help="Open database in read-only mode",
    ),
    before: str | None = typer.Option(
        None,
        "--before",
        help="Query database as it existed before this datetime (implies --read-only). "
        "Accepts ISO 8601 format (e.g., 2025-01-15T14:30:00) or date (e.g., 2025-01-15)",
    ),
):
    """haiku.rag CLI - Vector database RAG system"""
    global _read_only, _before
    _read_only = read_only

    # Parse and store before datetime
    if before is not None:  # pragma: no cover
        from haiku.rag.utils import parse_datetime, to_utc

        try:
            _before = to_utc(parse_datetime(before))
        except ValueError as e:
            typer.echo(f"Error: {e}")
            raise typer.Exit(1)
    else:
        _before = None
    # Load config from --config, local folder, or default directory
    config_path = find_config_file(cli_path=config)
    if config_path:
        yaml_data = load_yaml_config(config_path)
        loaded_config = AppConfig.model_validate(yaml_data)
        set_config(loaded_config)

    # Configure logging for CLI context
    configure_cli_logging()

    # Configure logfire (only sends data if token is present)
    try:
        import logfire

        is_production = get_config().environment != "development"
        logfire.configure(
            send_to_logfire="if-token-present",
            console=False if is_production else None,
        )
        logfire.instrument_pydantic_ai()
    except Exception:  # pragma: no cover
        pass

    if get_config().environment != "development":
        # Suppress warnings in production
        warnings.filterwarnings("ignore")

    # Run version check before any command
    try:
        asyncio.run(check_version())
    except Exception:  # pragma: no cover
        # Do not block CLI on version check issues
        pass


@_cli.command("list", help="List all stored documents")
def list_documents(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
    filter: str | None = typer.Option(
        None,
        "--filter",
        "-f",
        help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
    ),
):
    app = create_app(db)
    asyncio.run(app.list_documents(filter=filter))


def _parse_meta_options(meta: list[str] | None) -> dict[str, Any]:
    """Parse repeated --meta KEY=VALUE options into a dictionary.

    Raises a Typer error if any entry is malformed.
    """
    result: dict[str, Any] = {}
    if not meta:
        return result
    for item in meta:
        if "=" not in item:
            raise typer.BadParameter("--meta must be in KEY=VALUE format")
        key, value = item.split("=", 1)
        if not key:
            raise typer.BadParameter("--meta key cannot be empty")
        # Best-effort JSON coercion: numbers, booleans, null, arrays/objects
        try:
            parsed = json.loads(value)
            result[key] = parsed
        except Exception:
            # Leave as string if not valid JSON literal
            result[key] = value
    return result


@_cli.command("add", help="Add a document from text input")
def add_document_text(  # pragma: no cover
    text: str = typer.Argument(
        help="The text content of the document to add",
    ),
    title: str | None = typer.Option(
        None,
        "--title",
        help="Optional title for the document",
    ),
    meta: list[str] | None = typer.Option(
        None,
        "--meta",
        help="Metadata entries as KEY=VALUE (repeatable)",
        metavar="KEY=VALUE",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    metadata = _parse_meta_options(meta)
    asyncio.run(
        app.add_document_from_text(text=text, title=title, metadata=metadata or None)
    )


@_cli.command("add-src", help="Add a document from a file path, directory, or URL")
def add_document_src(  # pragma: no cover
    source: str = typer.Argument(
        help="The file path, directory, or URL of the document(s) to add",
    ),
    title: str | None = typer.Option(
        None,
        "--title",
        help="Optional human-readable title to store with the document",
    ),
    meta: list[str] | None = typer.Option(
        None,
        "--meta",
        help="Metadata entries as KEY=VALUE (repeatable)",
        metavar="KEY=VALUE",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    metadata = _parse_meta_options(meta)
    asyncio.run(
        app.add_document_from_source(
            source=source, title=title, metadata=metadata or None
        )
    )


@_cli.command("get", help="Get and display a document by its ID")
def get_document(  # pragma: no cover
    doc_id: str = typer.Argument(
        help="The ID of the document to get",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    asyncio.run(app.get_document(doc_id=doc_id))


@_cli.command("delete", help="Delete a document by its ID")
def delete_document(  # pragma: no cover
    doc_id: str = typer.Argument(
        help="The ID of the document to delete",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    asyncio.run(app.delete_document(doc_id=doc_id))


# Add alias `rm` for delete
_cli.command("rm", help="Alias for delete: remove a document by its ID")(
    delete_document
)


@_cli.command("search", help="Search for documents by a query")
def search(  # pragma: no cover
    query: str | None = typer.Argument(
        None,
        help="The search query (omit when using --image)",
    ),
    limit: int | None = typer.Option(
        None,
        "--limit",
        "-l",
        help="Maximum number of results to return (default: config search.default_limit)",
    ),
    filter: str | None = typer.Option(
        None,
        "--filter",
        "-f",
        help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
    ),
    image: Path | None = typer.Option(
        None,
        "--image",
        help="Path to an image file to use as the query (requires a multimodal embedder)",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    asyncio.run(app.search(query=query, limit=limit, filter=filter, image=image))


@_cli.command("visualize", help="Show visual grounding for a chunk")
def visualize(  # pragma: no cover
    chunk_id: str = typer.Argument(
        help="The ID of the chunk to visualize",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    asyncio.run(app.visualize_chunk(chunk_id=chunk_id))


@_cli.command("ask", help="Ask a question using the QA agent")
def ask(  # pragma: no cover
    question: str = typer.Argument(
        help="The question to ask",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
    cite: bool = typer.Option(
        False,
        "--cite",
        help="Include citations in the response",
    ),
    filter: str | None = typer.Option(
        None,
        "--filter",
        "-f",
        help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
    ),
):
    app = create_app(db)
    asyncio.run(
        app.ask(
            question=question,
            cite=cite,
            filter=filter,
        )
    )


@_cli.command("analyze", help="Answer questions using code execution (analysis agent)")
def analyze(  # pragma: no cover
    question: str = typer.Argument(
        help="The question to answer",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
    document: str | None = typer.Option(
        None,
        "--document",
        "-d",
        help="Document ID or title to pre-load for analysis",
    ),
    filter: str | None = typer.Option(
        None,
        "--filter",
        "-f",
        help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
    ),
):
    app = create_app(db)
    asyncio.run(
        app.analyze(
            question=question,
            document=document,
            filter=filter,
        )
    )


@_cli.command("research", help="Run multi-agent research and output a concise report")
def research(  # pragma: no cover
    question: str = typer.Argument(..., help="The research question to investigate"),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
    filter: str | None = typer.Option(
        None,
        "--filter",
        "-f",
        help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
    ),
):
    app = create_app(db)
    asyncio.run(app.research(question=question, filter=filter))


@_cli.command("settings", help="Display current configuration settings")
def settings():  # pragma: no cover
    config = get_config()
    app = HaikuRAGApp(db_path=Path(), config=config)
    app.show_settings()


@_cli.command("init-config", help="Generate a YAML configuration file")
def init_config(  # pragma: no cover
    output: Path = typer.Argument(
        Path("haiku.rag.yaml"),
        help="Output path for the config file",
    ),
):
    """Generate a YAML configuration file with defaults."""
    import yaml

    from haiku.rag.config.loader import generate_default_config

    if output.exists():
        typer.echo(
            f"Error: {output} already exists. Remove it first or choose a different path."
        )
        raise typer.Exit(1)

    config_data = generate_default_config()

    # Write YAML with comments
    with open(output, "w") as f:
        f.write("# haiku.rag configuration file\n")
        f.write(
            "# See https://ggozad.github.io/haiku.rag/configuration/ for details\n\n"
        )
        yaml.dump(config_data, f, default_flow_style=False, sort_keys=False)

    typer.echo(f"Configuration file created: {output}")
    typer.echo("Edit the file to customize your settings.")


@_cli.command(
    "rebuild",
    help="Rebuild the database by deleting all chunks and re-indexing all documents",
)
def rebuild(
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
    embed_only: bool = typer.Option(
        False,
        "--embed-only",
        help="Only regenerate embeddings, keep existing chunks",
    ),
    rechunk: bool = typer.Option(
        False,
        "--rechunk",
        help="Re-chunk from existing content without accessing source files",
    ),
    title_only: bool = typer.Option(
        False,
        "--title-only",
        help="Only generate titles for documents without one",
    ),
    descriptions: bool = typer.Option(
        False,
        "--descriptions",
        help=(
            "Run the VLM over already-stored picture bytes, patch descriptions "
            "into the docling blob, then re-chunk + re-embed. Skips the docling "
            "parse entirely. Requires processing.pictures='description'."
        ),
    ),
):
    from haiku.rag.client import RebuildMode

    exclusive = sum([embed_only, rechunk, title_only, descriptions])
    if exclusive > 1:
        typer.echo(
            "Error: --embed-only, --rechunk, --title-only, and --descriptions "
            "are mutually exclusive"
        )
        raise typer.Exit(1)

    if embed_only:  # pragma: no cover
        mode = RebuildMode.EMBED_ONLY
    elif rechunk:  # pragma: no cover
        mode = RebuildMode.RECHUNK
    elif title_only:  # pragma: no cover
        mode = RebuildMode.TITLE_ONLY
    elif descriptions:  # pragma: no cover
        mode = RebuildMode.DESCRIPTIONS
    else:  # pragma: no cover
        mode = RebuildMode.FULL

    app = create_app(db)  # pragma: no cover
    asyncio.run(app.rebuild(mode=mode))  # pragma: no cover


@_cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage")
def vacuum(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    asyncio.run(app.vacuum())


@_cli.command("migrate", help="Run pending database migrations")
def migrate(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    try:
        applied = asyncio.run(app.migrate())
        if applied:
            typer.echo(f"Applied {len(applied)} migration(s):")
            for desc in applied:
                typer.echo(f"  - {desc}")
            typer.echo("Migration completed successfully.")
        else:
            typer.echo("No migrations pending. Database is up to date.")
    except Exception as e:
        typer.echo(f"Migration failed: {e}")
        raise typer.Exit(1)


@_cli.command(
    "create-index", help="Create vector index for efficient similarity search"
)
def create_index(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    asyncio.run(app.create_index())


@_cli.command("init", help="Initialize a new database")
def init_db(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    asyncio.run(app.init())


@_cli.command("info", help="Show database info")
def info(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    app = create_app(db)
    asyncio.run(app.info())


@_cli.command("history", help="Show version history for database tables")
def history(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
    table: str | None = typer.Option(
        None,
        "--table",
        "-t",
        help="Specific table to show history for (documents, chunks, settings)",
    ),
    limit: int | None = typer.Option(
        None,
        "--limit",
        "-l",
        help="Maximum number of versions to show per table",
    ),
):
    app = create_app(db)
    asyncio.run(app.history(table=table, limit=limit))


@_cli.command("download-models", help="Download Docling and Ollama models per config")
def download_models_cmd():  # pragma: no cover
    app = HaikuRAGApp(db_path=Path(), config=get_config())
    try:
        asyncio.run(app.download_models())
    except Exception as e:
        typer.echo(f"Error downloading models: {e}")
        raise typer.Exit(1)


@_cli.command("inspect", help="Launch interactive TUI to inspect database contents")
def inspect(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
):
    """Launch the inspector TUI for browsing documents and chunks."""
    try:
        from haiku.rag.inspector import run_inspector
    except ImportError as e:
        typer.echo(f"Error: {e}", err=True)
        raise typer.Exit(1) from e

    db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb"
    run_inspector(db_path, read_only=_read_only, before=_before)


@_cli.command("chat", help="Launch interactive chat TUI for conversational RAG")
def chat(  # pragma: no cover
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
    model: str | None = typer.Option(
        None,
        "--model",
        help="Model to use for the chat (e.g. openai:gpt-4o)",
    ),
    skill: list[str] | None = typer.Option(
        None,
        "--skill",
        "-s",
        help="Skills to enable: rag, analysis (can repeat, default: rag)",
    ),
):
    """Launch the chat TUI for conversational RAG."""
    from haiku.rag.chat import run_chat

    db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb"
    skills = skill if skill else ["rag"]

    run_chat(
        db_path,
        read_only=_read_only,
        before=_before,
        model=model,
        skills=skills,
    )


@_cli.command(
    "serve",
    help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
)
def serve(
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database file",
    ),
    monitor: bool = typer.Option(
        False,
        "--monitor",
        help="Enable file monitoring",
    ),
    mcp: bool = typer.Option(
        False,
        "--mcp",
        help="Enable MCP server",
    ),
    stdio: bool = typer.Option(
        False,
        "--stdio",
        help="Run MCP server on stdio Transport (requires --mcp)",
    ),
    mcp_port: int = typer.Option(
        8001,
        "--mcp-port",
        help="Port to bind MCP server to (ignored with --stdio)",
    ),
) -> None:
    """Start the server with selected services."""
    # Require at least one service flag
    if not (monitor or mcp):
        typer.echo(
            "Error: At least one service flag (--monitor or --mcp) must be specified"
        )
        raise typer.Exit(1)

    if stdio and not mcp:
        typer.echo("Error: --stdio requires --mcp")
        raise typer.Exit(1)

    app = create_app(db)  # pragma: no cover

    transport = "stdio" if stdio else None  # pragma: no cover

    asyncio.run(  # pragma: no cover
        app.serve(
            enable_monitor=monitor,
            enable_mcp=mcp,
            mcp_transport=transport,
            mcp_port=mcp_port,
        )
    )


@_cli.command(
    "create-skill",
    help="Generate a standalone skill package with an embedded or remote database",
)
def create_skill_cmd(  # pragma: no cover
    name: str = typer.Option(
        ...,
        "--name",
        help="Skill name (lowercase alphanumeric and hyphens)",
    ),
    db: Path | None = typer.Option(
        None,
        "--db",
        help="Path to the LanceDB database to embed (omit for remote storage)",
    ),
    description: str | None = typer.Option(
        None,
        "--description",
        help="Skill description (default: standard RAG description)",
    ),
    tools: str = typer.Option(
        "all",
        "--tools",
        help="Comma-separated tool names, or 'all'",
    ),
    preamble: str | None = typer.Option(
        None,
        "--preamble",
        help="Custom preamble for the skill instructions",
    ),
    config_file: Path | None = typer.Option(
        None,
        "--config-file",
        help="Path to haiku.rag.yaml to embed in the skill",
    ),
    output: Path = typer.Option(
        Path("."),
        "--output",
        "-o",
        help="Output directory for the generated package",
    ),
):
    """Generate a standalone haiku.skills package with an embedded database."""
    from haiku.rag.skill_generator import (
        AVAILABLE_TOOLS,
        DEFAULT_DESCRIPTION,
        generate_skill,
    )

    if description is None:
        description = DEFAULT_DESCRIPTION

    if tools.strip().lower() == "all":
        tool_names = sorted(AVAILABLE_TOOLS)
    else:
        tool_names = [t.strip() for t in tools.split(",") if t.strip()]

    try:
        result = generate_skill(
            db_path=db,
            output_dir=output,
            name=name,
            description=description,
            tool_names=tool_names,
            config_path=config_file,
            preamble=preamble,
        )
        typer.echo(f"Skill generated: {result}")
    except ValueError as e:
        typer.echo(f"Error: {e}", err=True)
        raise typer.Exit(1)


if __name__ == "__main__":  # pragma: no cover
    cli()