content-find-agent-new/aigc_push_registry/cli.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

from aigc_push_registry.registry import AigcPushRegistry, import_live_logs


def main() -> None:
    parser = argparse.ArgumentParser(description="AIGC push registry")
    parser.add_argument("--env-file", default=".env", help="Project env file for CONTENT_SUPPLY_DB_*")
    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("init-db", help="Create AIGC push registry tables")

    import_parser = sub.add_parser("import-live-logs", help="Import only live AIGC push logs into DB")
    import_parser.add_argument("--log-dir", default="e2e_logs")

    upsert_parser = sub.add_parser("upsert", help="Upsert one manual pushed record JSON")
    upsert_parser.add_argument("--file", required=True)

    list_parser = sub.add_parser("list", help="List pushed records from DB")
    list_parser.add_argument("--platform")
    list_parser.add_argument("--demand-content-id", type=int)
    list_parser.add_argument("--run-id")
    list_parser.add_argument("--limit", type=int, default=200)

    args = parser.parse_args()
    registry = AigcPushRegistry.from_env(args.env_file)

    if args.command == "init-db":
        registry.init_db()
        print("created content_agent_aigc_push_* tables")
        return
    if args.command == "import-live-logs":
        registry.init_db()
        imported = import_live_logs(registry, Path(args.log_dir))
        print(json.dumps({"imported": len(imported), "push_record_ids": [row["push_record_id"] for row in imported]}, ensure_ascii=False, indent=2))
        return
    if args.command == "upsert":
        registry.init_db()
        payload = json.loads(Path(args.file).read_text(encoding="utf-8"))
        registry.upsert_record(payload)
        print(json.dumps({"upserted": payload.get("push_record_id")}, ensure_ascii=False))
        return
    if args.command == "list":
        payload = registry.list_records(
            platform=args.platform,
            demand_content_id=args.demand_content_id,
            run_id=args.run_id,
            limit=args.limit,
        )
        print(json.dumps(payload, ensure_ascii=False, indent=2))
        return


if __name__ == "__main__":
    main()