Create-knowledge-find-decode/scripts/build_creation_demo.py

#!/usr/bin/env python3
"""Build formal creation-query batches.

Default mode only prints a summary. Use --export-json for a local review file
or --persist to write query_batches/queries into the formal PostgreSQL store.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from acquisition.queries.builder import (
    QueryBuildOptions,
    TREES,
    build_creation_query_batch,
    persist_query_batch,
)
from acquisition.repositories.postgres import PostgresAcquisitionRepository
from core.config import CreationDbConfig, Settings
from core.db_session import transaction


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--env-file", default=".env")
    parser.add_argument("--tree-path", type=Path)
    parser.add_argument("--per", type=int, default=30)
    parser.add_argument("--batch-n", type=int, default=30)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--dry", action="store_true", help="Skip LLM query filtering")
    parser.add_argument("--export-json", type=Path, help="Optional local review export")
    parser.add_argument("--persist", action="store_true", help="Persist to formal PG")
    parser.add_argument("--name", default="creation-demo")
    return parser.parse_args(argv)


def _summary(generated: dict) -> dict:
    families = generated.get("families") or []
    total = sum(len(f.get("items") or []) for f in families)
    kept = sum(
        1
        for family in families
        for item in family.get("items") or []
        if item.get("keep", True)
    )
    return {
        "family_count": len(families),
        "query_count": total,
        "kept_count": kept,
        "metadata": generated.get("metadata") or {},
    }


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    settings = Settings.from_env(args.env_file)
    generated = build_creation_query_batch(
        settings,
        tree_path=args.tree_path or TREES,
        options=QueryBuildOptions(
            per=args.per,
            batch_n=args.batch_n,
            seed=args.seed,
            dry=args.dry,
        ),
    )
    summary = _summary(generated)

    if args.export_json:
        args.export_json.parent.mkdir(parents=True, exist_ok=True)
        args.export_json.write_text(
            json.dumps(generated, ensure_ascii=False, indent=1),
            encoding="utf-8",
        )
        summary["export_json"] = str(args.export_json)

    if args.persist:
        db_config = CreationDbConfig.from_env(args.env_file)
        with transaction(db_config) as conn:
            repo = PostgresAcquisitionRepository(conn)
            batch, count = persist_query_batch(
                repo,
                generated,
                name=args.name,
            )
        summary["batch_id"] = str(batch.id)
        summary["persisted_queries"] = count

    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


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