cli.py

#!/usr/bin/env python3
"""
cli.py
──────
Command-line interface for the Local RAG Knowledge Base.

Usage examples
--------------
# Ingest a directory
python cli.py ingest ./my_notes

# Ask a question
python cli.py query "What are the main ideas in my ML notes?"

# Summarise a topic
python cli.py summarise "deep learning optimisers"

# Generate insights
python cli.py insights "transformer architectures"

# Show stats
python cli.py stats
"""

import argparse
import logging
import sys
import textwrap
from pathlib import Path

# ── Bootstrap path ─────────────────────────────────────────────────────────────
sys.path.insert(0, str(Path(__file__).parent))

from src.rag_pipeline import RAGPipeline   # noqa: E402

logging.basicConfig(
    level=logging.WARNING,
    format="%(levelname)s | %(name)s | %(message)s",
)

BANNER = """
╔══════════════════════════════════════════════════════╗
║       🧠  Local RAG Knowledge Base  CLI              ║
╚══════════════════════════════════════════════════════╝
"""

# ANSI colours (no extra deps)
GREEN  = "\033[92m"
YELLOW = "\033[93m"
CYAN   = "\033[96m"
BOLD   = "\033[1m"
RESET  = "\033[0m"


def _separator(char: str = "─", width: int = 60) -> str:
    return char * width


def _print_response(resp) -> None:
    print(f"\n{BOLD}Answer{RESET}")
    print(_separator())
    print(textwrap.fill(resp.answer, width=80))
    print()

    if resp.sources:
        print(f"{CYAN}Sources:{RESET}")
        for src in resp.source_files:
            print(f"  • {src}")

    if resp.refinement and resp.refinement.was_refined:
        print(f"\n{YELLOW}Query was refined over {resp.refinement.iterations} iterations.{RESET}")
        print(f"  Final query used: \"{resp.refinement.final_query}\"")


# ── Sub-commands ───────────────────────────────────────────────────────────────


def cmd_ingest(args, pipeline: RAGPipeline) -> None:
    path = Path(args.path)
    print(f"Ingesting: {path.resolve()} …")

    if path.is_dir():
        count = pipeline.ingest_directory(path, recursive=not args.no_recursive)
    elif path.is_file():
        count = pipeline.ingest_file(path)
    else:
        print(f"[ERROR] Path not found: {path}", file=sys.stderr)
        sys.exit(1)

    print(f"{GREEN}✓ Ingested {count} chunks successfully.{RESET}")


def cmd_query(args, pipeline: RAGPipeline) -> None:
    question = " ".join(args.question)
    print(f"\n{BOLD}Question:{RESET} {question}")
    print("Searching knowledge base …\n")

    if not pipeline.llm.is_available():
        print(
            f"{YELLOW}⚠  Ollama not reachable at {pipeline.llm.base_url}.{RESET}\n"
            "   Start Ollama and run `ollama pull llama3`."
        )
        sys.exit(1)

    if args.stream:
        gen = pipeline.query(question, stream=True, refine=not args.no_refine,
                             search_mode=args.mode)
        print(f"\n{BOLD}Answer{RESET}\n{_separator()}")
        for token in gen:
            print(token, end="", flush=True)
        print()
    else:
        resp = pipeline.query(question, stream=False, refine=not args.no_refine,
                               search_mode=args.mode)
        _print_response(resp)


def cmd_summarise(args, pipeline: RAGPipeline) -> None:
    topic = " ".join(args.topic)
    print(f"Summarising topic: {topic} …\n")
    resp = pipeline.summarise(topic, search_mode=args.mode)
    _print_response(resp)


def cmd_insights(args, pipeline: RAGPipeline) -> None:
    topic = " ".join(args.topic) if args.topic else ""
    label = topic or "general knowledge base"
    print(f"Generating insights on: {label} …\n")
    resp = pipeline.generate_insights(topic, search_mode=args.mode)
    _print_response(resp)


def cmd_stats(_, pipeline: RAGPipeline) -> None:
    stats = pipeline.stats()
    print(f"\n{BOLD}Knowledge Base Stats{RESET}")
    print(_separator())
    print(f"  Total chunks    : {stats['total_chunks']}")
    print(f"  Unique sources  : {stats['unique_sources']}")
    print(f"  Embedding model : {stats['embedding_model']}")
    print(f"  LLM model       : {stats['ollama_model']}")
    print(f"  LLM available   : {'Yes ✓' if stats['llm_available'] else 'No ✗'}")
    if stats["sources"]:
        print(f"\n  {CYAN}Ingested files:{RESET}")
        for src in stats["sources"]:
            print(f"    • {Path(src).name}")


def cmd_interactive(_, pipeline: RAGPipeline) -> None:
    """Simple REPL loop."""
    print(BANNER)
    print("Type your question and press Enter.  Commands: /quit, /stats, /mode")
    mode = "hybrid"

    while True:
        try:
            user_input = input(f"\n{BOLD}[{mode}]>{RESET} ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye!")
            break

        if not user_input:
            continue

        if user_input.lower() in ("/quit", "/exit", "q"):
            print("Bye!")
            break

        if user_input == "/stats":
            cmd_stats(None, pipeline)
            continue

        if user_input.startswith("/mode "):
            new_mode = user_input.split(" ", 1)[1].strip()
            if new_mode in ("hybrid", "semantic", "bm25"):
                mode = new_mode
                print(f"Search mode set to: {mode}")
            else:
                print("Valid modes: hybrid | semantic | bm25")
            continue

        if not pipeline.llm.is_available():
            print(f"{YELLOW}Ollama not available. Please start it first.{RESET}")
            continue

        resp = pipeline.query(user_input, stream=False, refine=True, search_mode=mode)
        _print_response(resp)


# ── Argument parser ────────────────────────────────────────────────────────────


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="rag-kb",
        description="Local RAG-Powered Personal Knowledge Base",
    )
    parser.add_argument(
        "--db", default=None, help="Path to ChromaDB directory (overrides config)"
    )
    parser.add_argument(
        "--model", default=None, help="Ollama model name (overrides config)"
    )

    sub = parser.add_subparsers(dest="command", required=True)

    # ingest
    p_ingest = sub.add_parser("ingest", help="Ingest documents")
    p_ingest.add_argument("path", help="File or directory to ingest")
    p_ingest.add_argument("--no-recursive", action="store_true",
                          help="Don't recurse into subdirectories")

    # query
    p_query = sub.add_parser("query", help="Ask a question")
    p_query.add_argument("question", nargs="+", help="Your question")
    p_query.add_argument("--stream", action="store_true", help="Stream the response")
    p_query.add_argument("--no-refine", action="store_true",
                         help="Disable query refinement loop")
    p_query.add_argument("--mode", choices=["hybrid", "semantic", "bm25"],
                         default="hybrid", help="Retrieval mode")

    # summarise
    p_sum = sub.add_parser("summarise", help="Summarise a topic")
    p_sum.add_argument("topic", nargs="+", help="Topic to summarise")
    p_sum.add_argument("--mode", choices=["hybrid", "semantic", "bm25"],
                       default="hybrid")

    # insights
    p_ins = sub.add_parser("insights", help="Generate insights")
    p_ins.add_argument("topic", nargs="*", help="Optional topic focus")
    p_ins.add_argument("--mode", choices=["hybrid", "semantic", "bm25"],
                       default="hybrid")

    # stats
    sub.add_parser("stats", help="Show knowledge base statistics")

    # interactive / REPL
    sub.add_parser("repl", help="Interactive REPL mode")

    return parser


# ── Entry point ────────────────────────────────────────────────────────────────


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    kwargs = {}
    if args.db:
        kwargs["db_path"] = args.db
    if args.model:
        kwargs["ollama_model"] = args.model

    pipeline = RAGPipeline(**kwargs)

    dispatch = {
        "ingest":    cmd_ingest,
        "query":     cmd_query,
        "summarise": cmd_summarise,
        "insights":  cmd_insights,
        "stats":     cmd_stats,
        "repl":      cmd_interactive,
    }
    dispatch[args.command](args, pipeline)


if __name__ == "__main__":
    main()