SuperBizAgent-java/scripts/rebuild_hybrid_knowledge.py

#!/usr/bin/env python3
"""Rebuild knowledge into the configured Milvus collection (default: biz).

Clears:
  - milvus.collection (drop + recreate dense+BM25 schema)
  - MySQL api_document
  - in-memory L0 index

Then force-imports all markdown under server-side knowledge.base-path
(default: knowledge_base/).

Usage:
  # start Spring Boot first, then:
  python scripts/rebuild_hybrid_knowledge.py --confirm REBUILD

  python scripts/rebuild_hybrid_knowledge.py --base-url http://127.0.0.1:9900 --confirm REBUILD
"""

from __future__ import annotations

import argparse
import json
import sys
import urllib.error
import urllib.request
from typing import Any


DEFAULT_BASE_URL = "http://localhost:9900"


def http_json(method: str, url: str, timeout: float = 3600.0) -> tuple[int, Any]:
    req = urllib.request.Request(url=url, method=method.upper())
    req.add_header("Accept", "application/json")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
            status = getattr(resp, "status", 200)
            if not raw.strip():
                return status, None
            return status, json.loads(raw)
    except urllib.error.HTTPError as exc:
        raw = exc.read().decode("utf-8", errors="replace")
        body: Any
        try:
            body = json.loads(raw) if raw.strip() else None
        except json.JSONDecodeError:
            body = raw
        raise RuntimeError(f"HTTP {method} {url} failed status={exc.code}: {body}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"HTTP {method} {url} failed: {exc}") from exc


def pretty(obj: Any) -> str:
    return json.dumps(obj, ensure_ascii=False, indent=2)


def step(title: str) -> None:
    print()
    print(f"==> {title}")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Drop/recreate milvus.collection (default biz), clear MySQL api_document + L0, "
        "and reimport knowledge_base markdown into dense+BM25."
    )
    parser.add_argument(
        "--base-url",
        default=DEFAULT_BASE_URL,
        help=f"Service base URL (default: {DEFAULT_BASE_URL})",
    )
    parser.add_argument(
        "--confirm",
        required=True,
        choices=["REBUILD"],
        help="Must be REBUILD to execute destructive rebuild",
    )
    parser.add_argument(
        "--skip-stats",
        action="store_true",
        help="Skip before/after /api/knowledge/stats",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=7200.0,
        help="Rebuild request timeout seconds (default: 7200)",
    )
    args = parser.parse_args(argv)

    base_url = args.base_url.rstrip("/")

    print("Hybrid knowledge rebuild")
    print(f"  BaseUrl : {base_url}")
    print(f"  Confirm : {args.confirm}")
    print("  Source  : knowledge_base/ (server-side knowledge.base-path)")
    print()
    print("This will DESTROY data in:")
    print("  - Milvus collection milvus.collection (default: biz)")
    print("  - MySQL table api_document")
    print("  - In-memory L0 knowledge index")
    print("Then re-import all markdown under knowledge_base.")
    print()

    # 1) health
    step("Check service health")
    try:
        status, body = http_json("GET", f"{base_url}/milvus/health", timeout=30)
        print(f"  milvus health status={status}")
        print(pretty(body))
    except Exception as exc:  # noqa: BLE001 - ops script should continue on soft health failure
        print(f"  WARN: /milvus/health failed: {exc}")
        print("  Continue if app is up but milvus health endpoint has issues.")

    # 2) stats before
    if not args.skip_stats:
        step("Knowledge stats (before)")
        try:
            _, body = http_json("GET", f"{base_url}/api/knowledge/stats", timeout=30)
            print(pretty(body))
        except Exception as exc:  # noqa: BLE001
            print(f"  WARN: stats before failed: {exc}")

    # 3) rebuild
    step("POST /api/knowledge/rebuild-hybrid?confirm=REBUILD")
    rebuild_url = f"{base_url}/api/knowledge/rebuild-hybrid?confirm={args.confirm}"
    try:
        status, body = http_json("POST", rebuild_url, timeout=args.timeout)
    except RuntimeError as exc:
        print(str(exc))
        return 1

    print(f"  HTTP {status}")
    print(pretty(body))

    if not isinstance(body, dict):
        print("Unexpected rebuild response type", file=sys.stderr)
        return 1

    inserted = int(body.get("inserted") or 0)
    failed = int(body.get("failed") or 0)
    success = bool(body.get("success"))

    if not success:
        if inserted <= 0:
            print()
            print("Rebuild reported failure and inserted=0. Inspect details above.", file=sys.stderr)
            return 2
        print()
        print(f"Rebuild finished with failed={failed} inserted={inserted}. Review details.")
    else:
        print()
        print(f"Rebuild OK: inserted={inserted}, failed={failed}")

    # 4) stats after
    if not args.skip_stats:
        step("Knowledge stats (after)")
        try:
            _, body = http_json("GET", f"{base_url}/api/knowledge/stats", timeout=30)
            print(pretty(body))
        except Exception as exc:  # noqa: BLE001
            print(f"  WARN: stats after failed: {exc}")

    print()
    print("Done.")
    print("Next:")
    print("  1) Ensure application.yml has:")
    print("       milvus.collection: biz")
    print("       retrieval.search.mode: hybrid")
    print("  2) Smoke test lookup_knowledge / chat with a known doc query")
    return 0 if success or inserted > 0 else 2


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