scripts/build_kg.py

"""
Build Knowledge Graph Script.

Reads persona corpus, extracts (subject, relation, object) triples via Qwen3-32B,
and serializes the resulting NetworkX graph.

Usage:
    python scripts/build_kg.py --persona alex_rivera
    python scripts/build_kg.py --persona alex_rivera --corpus-dir data/personas/alex_rivera
"""

from __future__ import annotations

import argparse
import asyncio
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Build knowledge graph for a persona from corpus text."
    )
    parser.add_argument(
        "--persona", required=True,
        help="Persona ID (e.g., alex_rivera)."
    )
    parser.add_argument(
        "--corpus-dir", default=None,
        help="Path to persona corpus directory. Defaults to data/personas/{persona}."
    )
    parser.add_argument(
        "--config", default="memorybridge/config/settings.yaml",
        help="Path to settings.yaml."
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    corpus_dir = args.corpus_dir or f"data/personas/{args.persona}"

    from memorybridge.core.models import ModelRegistry
    from memorybridge.memory.kg_builder import KGBuilder

    registry = ModelRegistry(args.config)
    builder = KGBuilder(registry)

    print(f"Building KG for: {args.persona}")
    print(f"Corpus directory: {corpus_dir}")
    print(f"Output: data/knowledge_graphs/{args.persona}.gpickle")

    builder.build(corpus_dir, args.persona)
    print("Done.")


if __name__ == "__main__":
    main()