knowledge_agent/cli.py

"""
cli.py — Command-line interface for the Knowledge Agent
=======================================================
Usage examples:
    python cli.py index                          # index ./documents/
    python cli.py index --dir ~/my_docs          # index a custom folder
    python cli.py index --force                  # force re-index everything
    python cli.py ask "What is our pricing?"     # one-shot Q&A
    python cli.py chat                           # interactive chat loop
    python cli.py list                           # list indexed documents
    python cli.py stats                          # show DB stats
    python cli.py delete report.pdf              # remove a document
"""

import argparse
import sys
import os
from typing import List, Dict

# ── Pretty printing helpers ───────────────────────────────────────────────────
try:
    from rich.console import Console
    from rich.markdown import Markdown
    from rich.table import Table
    from rich.panel import Panel
    from rich import print as rprint
    RICH = True
    console = Console()
except ImportError:
    RICH = False
    console = None

from knowledge_api import (
    query_knowledge,
    index_docs,
    index_single,
    get_knowledge_stats,
    list_indexed_docs,
    delete_doc,
)


def print_answer(result: Dict):
    answer  = result.get("answer", "")
    sources = result.get("sources", [])
    chunks  = result.get("chunks_used", 0)

    if RICH:
        console.print(Panel(
            Markdown(answer),
            title="[bold cyan]Answer[/bold cyan]",
            border_style="cyan",
        ))
        if sources:
            console.print(f"[dim]Sources ({chunks} chunks): {', '.join(sources)}[/dim]")
    else:
        print("\n" + "="*60)
        print("ANSWER:")
        print(answer)
        if sources:
            print(f"\nSources ({chunks} chunks): {', '.join(sources)}")
        print("="*60 + "\n")


def cmd_index(args):
    docs_dir = args.dir or "./documents"
    print(f"[Knowledge Agent] Indexing documents in: {docs_dir}")
    index_docs(docs_dir=docs_dir, force=args.force)


def cmd_ask(args):
    question = " ".join(args.question)
    if RICH:
        console.print(f"\n[bold]Question:[/bold] {question}")
    else:
        print(f"\nQuestion: {question}")

    result = query_knowledge(question, top_k=args.top_k)
    print_answer(result)


def cmd_chat(args):
    """Interactive REPL — keeps asking questions until the user exits."""
    if RICH:
        console.print(Panel(
            "[bold cyan]Knowledge Agent — Interactive Chat[/bold cyan]\n"
            "Ask anything about your documents.\n"
            "Type [yellow]exit[/yellow] or [yellow]quit[/yellow] to stop.",
            border_style="cyan",
        ))
    else:
        print("\n" + "="*60)
        print(" Knowledge Agent — Interactive Chat")
        print(" Type 'exit' to quit.")
        print("="*60 + "\n")

    while True:
        try:
            if RICH:
                question = console.input("[bold green]You:[/bold green] ").strip()
            else:
                question = input("You: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nGoodbye!")
            break

        if not question:
            continue
        if question.lower() in ("exit", "quit", "q", "bye"):
            print("Goodbye!")
            break

        # Special commands in chat mode
        if question.lower() == "/stats":
            cmd_stats(args)
            continue
        if question.lower() == "/list":
            cmd_list(args)
            continue
        if question.lower().startswith("/index "):
            path = question[7:].strip()
            index_single(path)
            continue

        result = query_knowledge(question, top_k=args.top_k if hasattr(args, "top_k") else 5)
        print_answer(result)


def cmd_list(args):
    docs = list_indexed_docs()
    if not docs:
        print("No documents indexed yet.")
        return

    if RICH:
        table = Table(title="Indexed Documents", border_style="cyan")
        table.add_column("File", style="bold")
        table.add_column("Status")
        table.add_column("Path", style="dim")
        for d in docs:
            status = "[green]✓ exists[/green]" if d["exists"] else "[red]✗ missing[/red]"
            table.add_row(d["name"], status, d["path"])
        console.print(table)
    else:
        print(f"\n{'File':<40} {'Status':<12} Path")
        print("-" * 80)
        for d in docs:
            status = "✓ exists" if d["exists"] else "✗ missing"
            print(f"{d['name']:<40} {status:<12} {d['path']}")


def cmd_stats(args):
    stats = get_knowledge_stats()
    if RICH:
        console.print(Panel(
            f"[bold]Total chunks:[/bold]   {stats['total_chunks']}\n"
            f"[bold]Indexed files:[/bold]  {stats['indexed_files']}\n"
            f"[bold]Embed model:[/bold]    {stats['embed_model']}\n"
            f"[bold]ChromaDB dir:[/bold]   {stats['chroma_dir']}",
            title="[cyan]Knowledge Base Stats[/cyan]",
            border_style="cyan",
        ))
    else:
        print("\nKnowledge Base Stats:")
        for k, v in stats.items():
            print(f"  {k}: {v}")


def cmd_delete(args):
    filename = args.filename
    confirm = input(f"Delete all chunks for '{filename}'? [y/N] ").strip().lower()
    if confirm == "y":
        delete_doc(filename)
        print(f"Deleted: {filename}")
    else:
        print("Cancelled.")


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

def main():
    parser = argparse.ArgumentParser(
        prog="knowledge",
        description="Personal Knowledge Agent — RAG over your documents",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # index
    p_index = subparsers.add_parser("index", help="Index documents into ChromaDB")
    p_index.add_argument("--dir",   default=None,  help="Documents directory (default: ./documents)")
    p_index.add_argument("--force", action="store_true", help="Force re-index all files")
    p_index.set_defaults(func=cmd_index)

    # ask
    p_ask = subparsers.add_parser("ask", help="Ask a one-shot question")
    p_ask.add_argument("question", nargs="+", help="Your question (in quotes or multiple words)")
    p_ask.add_argument("--top-k", dest="top_k", type=int, default=5,
                       help="Number of chunks to retrieve (default: 5)")
    p_ask.set_defaults(func=cmd_ask)

    # chat
    p_chat = subparsers.add_parser("chat", help="Interactive chat loop")
    p_chat.add_argument("--top-k", dest="top_k", type=int, default=5)
    p_chat.set_defaults(func=cmd_chat)

    # list
    p_list = subparsers.add_parser("list", help="List indexed documents")
    p_list.set_defaults(func=cmd_list)

    # stats
    p_stats = subparsers.add_parser("stats", help="Show knowledge base statistics")
    p_stats.set_defaults(func=cmd_stats)

    # delete
    p_del = subparsers.add_parser("delete", help="Remove a document from the index")
    p_del.add_argument("filename", help="Filename to delete (e.g. report.pdf)")
    p_del.set_defaults(func=cmd_delete)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()