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

#!/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.ingest import ingest_payload_draft  # 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
            record = ingest_payload_draft(repo, draft, dry_run=True)
            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())