brain_cli.py

#!/usr/bin/env python
"""brain_cli.py — CLI for interacting with the Brain KB API.

Usage:
    python brain_cli.py status
    python brain_cli.py ingest --text "..." --space personal --title "..."
    python brain_cli.py query "question here" [--space personal] [--top_k 5]
    python brain_cli.py list [--space personal]
    python brain_cli.py ingest-batch facts.json
"""

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request

_PORT = os.getenv("SCP_DASHBOARD_PORT", "9001")
BASE = os.getenv("SCP_BRAIN_URL", f"http://localhost:{_PORT}/api/brain")
_QDRANT_PORT = os.getenv("SCP_QDRANT_PORT", "6333")


def _post(path: str, body: dict) -> dict:
    data = json.dumps(body, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        f"{BASE}{path}", data=data,
        headers={"Content-Type": "application/json; charset=utf-8"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return json.loads(r.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        print(f"HTTP {e.code}: {e.read().decode()}", file=sys.stderr)
        sys.exit(1)
    except urllib.error.URLError as e:
        print(f"Connection error: {e.reason}\n→ Is the dashboard running? (manage.ps1 start)", file=sys.stderr)
        sys.exit(1)


def _get(path: str) -> dict:
    req = urllib.request.Request(f"{BASE}{path}", method="GET")
    try:
        with urllib.request.urlopen(req, timeout=10) as r:
            return json.loads(r.read().decode("utf-8"))
    except urllib.error.URLError as e:
        print(f"Connection error: {e.reason}", file=sys.stderr)
        sys.exit(1)


def cmd_status(_args):
    r = _get("/status")
    print(f"Status:  {r.get('status')}")
    print(f"Points:  {r.get('points_count')}")
    print(f"Embed:   {r.get('embed_model')}")
    print(f"LLM:     {r.get('generate_model')}")


def cmd_ingest(args):
    tags = [t.strip() for t in args.tags.split(",")] if args.tags else []
    body = {
        "text":   args.text,
        "space":  args.space,
        "title":  args.title or "",
        "tags":   tags,
        "source": args.source or "brain-cli",
    }
    r = _post("/ingest", body)
    print(f"✓ Ingested — doc_id: {r.get('doc_id')}  chunks: {r.get('chunks')}")


def cmd_query(args):
    contexts = [c.strip() for c in args.space.split(",")] if args.space else ["*"]
    body = {
        "question": args.question,
        "contexts": contexts,
        "top_k":    args.top_k,
        "generate": args.generate,
    }
    r = _post("/query", body)
    chunks = r.get("chunks_found", 0)
    print(f"\n🧠 Brain — {chunks} chunk(s) found\n{'─'*60}")

    if r.get("answer"):
        print(f"\n{r['answer']}\n")

    sources = r.get("sources") or []
    if sources:
        print("Sources:")
        for s in sources:
            print(f"  [{s['space']}] {s['title']}  score={s['score']}")
            print(f"    {s['preview'][:120]}...")
            print()
    elif not r.get("answer"):
        print("(no results)")


def cmd_list(args):
    # Usar Qdrant scroll directo
    qdrant = f"http://localhost:{_QDRANT_PORT}/collections/brain/points/scroll"
    body: dict = {"limit": 100, "with_payload": True, "with_vector": False}
    if args.space:
        body["filter"] = {"must": [{"key": "space", "match": {"value": args.space}}]}
    data = json.dumps(body).encode()
    req = urllib.request.Request(qdrant, data=data,
                                  headers={"Content-Type": "application/json"},
                                  method="POST")
    with urllib.request.urlopen(req, timeout=10) as r:
        result = json.loads(r.read())

    points = result["result"]["points"]
    from itertools import groupby
    points.sort(key=lambda p: p["payload"].get("space", ""))
    for space, group in groupby(points, key=lambda p: p["payload"].get("space", "")):
        items = list(group)
        print(f"\n=== {space} ({len(items)}) ===")
        for p in sorted(items, key=lambda x: x["payload"].get("title", "")):
            title = p["payload"].get("title", "—")
            text  = p["payload"].get("text", "")[:80]
            print(f"  · {title}")
            print(f"    {text}...")


def cmd_ingest_batch(args):
    """Ingestar múltiples hechos desde un archivo JSON."""
    with open(args.file, encoding="utf-8") as f:
        items = json.load(f)
    ok = 0
    for item in items:
        r = _post("/ingest", item)
        ok += 1
        print(f"OK  {item.get('title', '?')[:60]}")
        time.sleep(0.25)
    print(f"\nTotal: {ok}/{len(items)}")


def main():
    p = argparse.ArgumentParser(description="Brain KB CLI")
    sub = p.add_subparsers(dest="cmd")

    # status
    sub.add_parser("status")

    # ingest
    pi = sub.add_parser("ingest")
    pi.add_argument("--text",   required=True)
    pi.add_argument("--space",  default="general")
    pi.add_argument("--title",  default="")
    pi.add_argument("--tags",   default="", help="comma-separated")
    pi.add_argument("--source", default="brain-cli")

    # ingest-batch
    pb = sub.add_parser("ingest-batch")
    pb.add_argument("file", help="JSON file with list of ingest objects")

    # query
    pq = sub.add_parser("query")
    pq.add_argument("question")
    pq.add_argument("--space",    default="",  help="comma-separated spaces, empty=all")
    pq.add_argument("--top_k",   type=int, default=5)
    pq.add_argument("--generate", action="store_true", default=False,
                    help="Use LLM to generate answer (requires Ollama running)")

    # list
    pl = sub.add_parser("list")
    pl.add_argument("--space", default="")

    args = p.parse_args()
    {"status": cmd_status, "ingest": cmd_ingest, "ingest-batch": cmd_ingest_batch,
     "query": cmd_query, "list": cmd_list}.get(args.cmd, lambda _: p.print_help())(args)


if __name__ == "__main__":
    main()