#!/usr/bin/env python3 """Run one real f1/f2 creation-knowledge pipeline sample to dry-run ingest.""" from __future__ import annotations import argparse import json from dataclasses import asdict, is_dataclass from datetime import datetime from pathlib import Path import sys from typing import Any ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from acquisition.queries.builder import ( # noqa: E402 QueryBuildOptions, TREES, build_creation_query_batch, persist_query_batch, ) from acquisition.repositories.postgres import PostgresAcquisitionRepository # noqa: E402 from acquisition.runner import DEFAULT_PLATFORMS, run_batch # noqa: E402 from core.config import CreationDbConfig, Settings # noqa: E402 from core.db_session import transaction # noqa: E402 from decode_content.repositories.postgres import PostgresDecodeRepository # noqa: E402 from decode_content.service import DecodeService # noqa: E402 from pipeline.decode_runner import run_decode_stage # noqa: E402 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, default=TREES) parser.add_argument("--per", type=int, default=1, help="Queries per active family") parser.add_argument("--batch-n", type=int, default=30) parser.add_argument("--seed", type=int, default=7) parser.add_argument( "--platform", action="append", choices=DEFAULT_PLATFORMS, help="Platform to run. Repeat for multiple platforms. Default: xiaohongshu.", ) parser.add_argument("--search-limit", type=int, default=1) parser.add_argument("--display-limit", type=int, default=1) parser.add_argument("--decode-limit", type=int, default=100) parser.add_argument("--name", default="") parser.add_argument("--frontend-base", default="http://127.0.0.1:5180/app/") return parser.parse_args(argv) def _now_key() -> str: return datetime.now().strftime("%Y%m%d-%H%M%S") def _model_dump(value: Any) -> dict[str, Any]: if hasattr(value, "model_dump"): return value.model_dump(mode="json") if is_dataclass(value): return asdict(value) if isinstance(value, dict): return value return dict(value) def _dry_ingest_payloads( repo: PostgresDecodeRepository, outputs: list[Any], ) -> list[dict[str, Any]]: records: list[dict[str, Any]] = [] for output in outputs: for draft in output.payload_drafts: if draft.id is None: continue repo.mark_payload_draft_ingested(draft.id) record = repo.save_ingest_record( payload_draft_id=draft.id, target_system="dry-run", target_id=str(draft.id), status="ingested", response_payload={ "dry_run": True, "note": "payload generated by singleton pipeline; external ingest API not called", "payload": draft.payload, }, ) records.append(_model_dump(record)) return records def main(argv: list[str] | None = None) -> int: args = parse_args(argv) settings = Settings.from_env(args.env_file) db_config = CreationDbConfig.from_env(args.env_file) platforms = tuple(args.platform or ("xiaohongshu",)) run_key_suffix = _now_key() name = args.name or f"creation-singleton-{run_key_suffix}" generated = build_creation_query_batch( settings, tree_path=args.tree_path, options=QueryBuildOptions( per=args.per, batch_n=args.batch_n, seed=args.seed, active_family_keys=("f1", "f2"), ), ) query_count = sum(len(family.get("items") or []) for family in generated.get("families") or []) kept_count = sum( 1 for family in generated.get("families") or [] for item in family.get("items") or [] if item.get("keep", True) ) with transaction(db_config) as conn: acquisition_repo = PostgresAcquisitionRepository(conn) batch, persisted_count = persist_query_batch( acquisition_repo, generated, name=name, source_type="generated", generation_method="creation_singleton_v1", target_platforms=list(platforms), ) with transaction(db_config) as conn: acquisition_repo = PostgresAcquisitionRepository(conn) acquisition = run_batch( acquisition_repo, batch_id=batch.id, settings=settings, platforms=platforms, search_limit=args.search_limit, display_limit=args.display_limit, classify=True, resume=False, skip_done=False, run_key=f"singleton-acquisition:{batch.id}:{run_key_suffix}", ) with transaction(db_config) as conn: acquisition_repo = PostgresAcquisitionRepository(conn) decode_repo = PostgresDecodeRepository(conn) decode_service = DecodeService(settings=settings, repository=decode_repo) decode = run_decode_stage( candidate_repo=acquisition_repo, decode_service=decode_service, run_id=acquisition.run_id, limit=args.decode_limit, ) ingest_records = _dry_ingest_payloads(decode_repo, decode.outputs) decoded_items = [str(output.item_id) for output in decode.outputs] first_item_id = decoded_items[0] if decoded_items else None summary = { "batch_id": str(batch.id), "run_id": str(acquisition.run_id), "active_family_keys": ["f1", "f2"], "query_count": query_count, "kept_count": kept_count, "persisted_queries": persisted_count, "platforms": list(platforms), "acquisition": _model_dump(acquisition), "decode": { "total": decode.total, "decoded": decode.decoded, "skipped": decode.skipped, "failed": decode.failed, "decoded_items": decoded_items, "payload_count": sum(len(output.payload_drafts) for output in decode.outputs), }, "dry_run_ingest_count": len(ingest_records), "detail_url": ( f"{args.frontend_base.rstrip('/')}/#/decode-item/{first_item_id}" if first_item_id else None ), } print(json.dumps(summary, ensure_ascii=False, indent=2, default=str)) return 0 if first_item_id else 2 if __name__ == "__main__": raise SystemExit(main())