cli.py

"""The ``engram`` command-line interface.

Four commands, mirroring the Python API:

    engram init      configure models (local Ollama or your own API key)
    engram ingest    load files into the local store
    engram ask       ask a question, get a grounded + cited answer
    engram doctor    show every component and whether any data can leave the box

Designed to be driven by a person *or* a coding agent: ``init`` auto-detects a
running Ollama and your environment and writes ``engram.toml`` without
prompting, so an agent can set everything up unattended. We never ship API
keys — you point Engram at your own local models or your own API key.
"""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path

_LOCAL_EMBED_PREFERENCE = ("nomic-embed-text", "mxbai-embed-large", "all-minilm")
_DEFAULT_LOCAL_EMBED = "nomic-embed-text"
_DEFAULT_LOCAL_READER = "llama3.1:8b"


def _ollama_models() -> list[str]:
    """Best-effort list of locally pulled Ollama model names (empty if none)."""
    try:
        import ollama

        data = ollama.list()
        models = data.get("models", []) if isinstance(data, dict) else getattr(data, "models", [])
        names = []
        for m in models:
            name = (
                (m.get("model") or m.get("name"))
                if isinstance(m, dict)
                else getattr(m, "model", None)
            )
            if name:
                names.append(name)
        if names:
            return names
    except Exception:
        pass
    try:
        out = subprocess.run(
            ["ollama", "list"], capture_output=True, text=True, timeout=5, check=False
        )
        if out.returncode == 0:
            return [ln.split()[0] for ln in out.stdout.strip().splitlines()[1:] if ln.strip()]
    except Exception:
        pass
    return []


def _pick_local_models() -> tuple[str, str]:
    """Choose (embedder, reader) from what Ollama has pulled, with fallbacks."""
    models = _ollama_models()
    embed = next(
        (m for pref in _LOCAL_EMBED_PREFERENCE for m in models if m.startswith(pref)),
        _DEFAULT_LOCAL_EMBED,
    )
    reader = next(
        (m for m in models if not any(m.startswith(p) for p in _LOCAL_EMBED_PREFERENCE)),
        _DEFAULT_LOCAL_READER,
    )
    return embed, reader


def _render_toml(*, embedder: str, reader: str, rerank_mode: str, store: str, profile: str) -> str:
    return (
        "# engram.toml — generated by `engram init`. Hand-edit anytime.\n\n"
        "[models]\n"
        f'embedder = "{embedder}"\n'
        f'reader   = "{reader}"\n\n'
        "[rerank]\n"
        f'mode  = "{rerank_mode}"          # local | cohere | off\n'
        'model = "bge-reranker-base"\n\n'
        "[store]\n"
        f'path = "{store}"          # LMDB + vector index live here\n\n'
        "[runtime]\n"
        f'profile = "{profile}"          # local hard-fails on any network egress\n'
    )


def cmd_init(args: argparse.Namespace) -> int:
    import os

    out_path = Path(args.output or "engram.toml")
    if out_path.exists() and not args.force:
        print(
            f"engram: {out_path} already exists. Re-run with --force to overwrite.", file=sys.stderr
        )
        return 1

    if args.local:
        profile = "local"
    elif args.cloud:
        profile = "cloud"
    else:
        # Auto: prefer local if Ollama is reachable; else cloud if a key exists.
        profile = "local" if _ollama_models() or not os.environ.get("OPENAI_API_KEY") else "cloud"

    if profile == "local":
        embedder, reader = _pick_local_models()
        embedder = f"ollama/{embedder}"
        reader = f"ollama/{reader}"
        rerank_mode = "local"
    else:
        embedder = "openai/text-embedding-3-small"
        reader = "openai/gpt-4o-mini"
        rerank_mode = "cohere"

    toml = _render_toml(
        embedder=embedder, reader=reader, rerank_mode=rerank_mode, store=args.store, profile=profile
    )
    out_path.write_text(toml)
    print(f"engram: wrote {out_path}  (profile={profile})")
    print(f"  embedder  {embedder}")
    print(f"  reader    {reader}")
    print(f"  reranker  {rerank_mode}")
    if profile == "local":
        need = {embedder.split("/", 1)[1], reader.split("/", 1)[1]}
        have = set(_ollama_models())
        missing = [m for m in need if not any(h.startswith(m) for h in have)]
        if missing:
            print("\n  Next: pull the models you don't have yet:")
            for m in missing:
                print(f"    ollama pull {m}")
    else:
        print("\n  Cloud profile uses your own API keys (OPENAI_API_KEY, AWS creds for rerank).")
    print("\n  Verify with:  engram doctor")
    return 0


def cmd_ingest(args: argparse.Namespace) -> int:
    from engram import Engram

    eg = Engram()
    enriched = eg.ingest(args.paths, build_graph=args.graph)
    store = eg._config.store_path if eg._config else "?"
    print(
        f"engram: ingested {len(enriched)} chunk(s) into {store}"
        + (" (with KG)" if args.graph else "")
    )
    return 0


def cmd_ask(args: argparse.Namespace) -> int:
    from engram import Engram

    eg = Engram()
    ans = eg.ask(args.question, mode=args.mode)
    if args.json:
        import json
        from dataclasses import asdict

        print(json.dumps(asdict(ans), indent=2))
        return 0
    print(ans.text)
    if ans.citations:
        print("\nSources:")
        for c in ans.citations:
            loc = f" p.{c.page}" if c.page is not None else ""
            print(f"  [{c.document}] {c.source}{loc}")
    print(f"\n(route: {ans.route} · {ans.latency_ms / 1000:.1f}s)")
    return 0


def cmd_doctor(args: argparse.Namespace) -> int:
    from engram.config import ConfigError, load_config

    try:
        cfg = load_config()
    except ConfigError as exc:
        print(f"engram doctor: configuration error\n{exc}", file=sys.stderr)
        return 1
    print(f"engram doctor  (config: {cfg.source}, profile: {cfg.profile})\n")
    for component, value, endpoint in cfg.manifest():
        print(f"  {component:<10} {value:<34} → {endpoint}")
    if cfg.egress_possible:
        print("\n  egress: POSSIBLE — some component may send text off the machine.")
    else:
        print("\n  egress: none — nothing can leave this machine.")
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="engram",
        description="Grounded, cited answers from your own documents — local or cloud.",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p_init = sub.add_parser("init", help="Configure models and write engram.toml")
    p_init.add_argument("--local", action="store_true", help="Fully-local (Ollama) config")
    p_init.add_argument("--cloud", action="store_true", help="Cloud (OpenAI + Cohere) config")
    p_init.add_argument("--store", default="./.engram", help="Where the local index lives")
    p_init.add_argument("--output", "-o", help="Config file path (default: engram.toml)")
    p_init.add_argument("--force", action="store_true", help="Overwrite an existing config")
    p_init.set_defaults(func=cmd_init)

    p_ing = sub.add_parser("ingest", help="Load files (PDF/txt/md/docx) into the store")
    p_ing.add_argument("paths", nargs="+", help="Files, globs, or directories")
    p_ing.add_argument(
        "--graph", action="store_true", help="Also build the knowledge graph (slower)"
    )
    p_ing.set_defaults(func=cmd_ingest)

    p_ask = sub.add_parser("ask", help="Ask a question against the ingested corpus")
    p_ask.add_argument("question", help="Your question, quoted")
    p_ask.add_argument(
        "--mode",
        choices=["hybrid", "kg", "auto"],
        default=None,
        help="Retrieval mode (default: hybrid)",
    )
    p_ask.add_argument("--json", action="store_true", help="Emit the full Answer as JSON")
    p_ask.set_defaults(func=cmd_ask)

    p_doc = sub.add_parser("doctor", help="Show components + whether any data can leave the box")
    p_doc.set_defaults(func=cmd_doctor)

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    return int(args.func(args))


if __name__ == "__main__":
    raise SystemExit(main())