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

#!/usr/bin/env python3
"""Ingest existing payload drafts for one acquisition run.

This script does not search, classify, or decode. It only reads payload_drafts
already produced for candidate items under the given run_id.
"""
from __future__ import annotations

import argparse
import json
import sys
from dataclasses import asdict, is_dataclass
from pathlib import Path
from uuid import UUID

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from core.config import CreationDbConfig, IngestApiConfig  # noqa: E402
from core.db_session import transaction  # noqa: E402
from decode_content.ingest import KnowledgeIngestClient, ingest_payload_draft  # noqa: E402
from decode_content.models import PayloadDraft  # noqa: E402
from decode_content.repositories.postgres import PostgresDecodeRepository  # noqa: E402
from pipeline.postgres import PostgresPipelineRepository  # noqa: E402
from pipeline.tracing import TraceContext, new_trace_writer  # noqa: E402


def _model_dump(value):
    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 list_payload_drafts_for_run(
    repo: PostgresDecodeRepository,
    *,
    run_id: UUID,
    include_real_ingested: bool = False,
    limit: int | None = None,
) -> list[PayloadDraft]:
    skip_real_ingested = "" if include_real_ingested else """
      AND NOT EXISTS (
        SELECT 1 FROM ingest_records ir
        WHERE ir.payload_draft_id = pd.id
          AND ir.target_system = 'knowledge-ingest-api'
          AND ir.status = 'ingested'
      )
    """
    limit_sql = "LIMIT %s" if limit is not None else ""
    params: tuple = (run_id, limit) if limit is not None else (run_id,)
    rows = repo._all(
        f"""
        SELECT pd.*
        FROM payload_drafts pd
        JOIN candidate_items ci ON ci.id = pd.item_id
        JOIN acquisition_jobs aj ON aj.id = ci.job_id
        WHERE aj.run_id = %s
        {skip_real_ingested}
        ORDER BY ci.created_at, pd.created_at, pd.id
        {limit_sql}
        """,
        params,
    )
    return [PayloadDraft.model_validate(row) for row in rows]


def count_run_payloads(repo: PostgresDecodeRepository, *, run_id: UUID) -> dict[str, int]:
    row = repo._one(
        """
        SELECT
          COUNT(DISTINCT pd.id)::int AS total_payloads,
          COUNT(DISTINCT pd.id) FILTER (
            WHERE EXISTS (
              SELECT 1 FROM ingest_records ir
              WHERE ir.payload_draft_id = pd.id
                AND ir.target_system = 'knowledge-ingest-api'
                AND ir.status = 'ingested'
            )
          )::int AS already_real_ingested
        FROM payload_drafts pd
        JOIN candidate_items ci ON ci.id = pd.item_id
        JOIN acquisition_jobs aj ON aj.id = ci.job_id
        WHERE aj.run_id = %s
        """,
        (run_id,),
    )
    return {
        "total_payloads": int(row.get("total_payloads") or 0),
        "already_real_ingested": int(row.get("already_real_ingested") or 0),
    }


def ingest_existing_payloads(
    repo: PostgresDecodeRepository,
    drafts: list[PayloadDraft],
    *,
    dry_run: bool,
    env_file: str,
    trace_writer=None,
    trace_context: TraceContext | None = None,
) -> list[dict]:
    client = None
    if not dry_run:
        config = IngestApiConfig.from_env(env_file)
        if not config.url:
            raise RuntimeError("CK_INGEST_API_URL 未配置，真实 ingest 已关闭")
        client = KnowledgeIngestClient(config)

    records: list[dict] = []
    for draft in drafts:
        if draft.id is None:
            continue
        context = (trace_context or TraceContext(stage="ingest")).child(
            item_id=draft.item_id,
            payload_draft_id=draft.id,
            stage="ingest",
        )
        if trace_writer is not None:
            trace_writer.event(
                context=context,
                stage="ingest",
                event_type="ingest_started",
                status="running",
                target_table="payload_drafts",
                target_id=draft.id,
                payload={"dry_run": dry_run, "entrypoint": "scripts/ingest_payloads.py"},
            )
        record = ingest_payload_draft(repo, draft, dry_run=dry_run, client=client)
        if trace_writer is not None:
            trace_writer.event(
                context=context.child(ingest_record_id=record.id),
                stage="ingest",
                event_type="ingest_recorded",
                status=record.status,
                target_table="ingest_records",
                target_id=record.id,
                payload={"dry_run": dry_run, "target_system": record.target_system},
                error_message=record.error_message,
            )
        records.append(_model_dump(record))
    return records


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-id", required=True, help="acquisition_runs.id whose existing payloads should be ingested")
    parser.add_argument("--pipeline-run-id", help="Optional pipeline_runs.id for trace events")
    parser.add_argument("--env-file", default=".env")
    parser.add_argument("--real", action="store_true", help="Call CK_INGEST_API_URL. Default is dry-run.")
    parser.add_argument("--limit", type=int, default=None, help="Optional max payload drafts to process")
    parser.add_argument(
        "--include-real-ingested",
        action="store_true",
        help="Also process drafts that already have successful knowledge-ingest-api records.",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    run_id = UUID(args.run_id)
    db_config = CreationDbConfig.from_env(args.env_file)
    pipeline_run_id = UUID(args.pipeline_run_id) if args.pipeline_run_id else None
    if pipeline_run_id is None:
        try:
            with transaction(db_config) as conn:
                pipeline_run_id = PostgresPipelineRepository(conn).get_pipeline_run_by_acquisition_run(run_id).id
        except Exception:
            pipeline_run_id = None
    trace_writer = new_trace_writer(db_config, env_file=args.env_file)
    trace_context = TraceContext(
        pipeline_run_id=pipeline_run_id,
        acquisition_run_id=run_id,
        stage="ingest",
    )

    with transaction(db_config) as conn:
        repo = PostgresDecodeRepository(conn)
        counts = count_run_payloads(repo, run_id=run_id)
        drafts = list_payload_drafts_for_run(
            repo,
            run_id=run_id,
            include_real_ingested=args.include_real_ingested,
            limit=args.limit,
        )
        records = ingest_existing_payloads(
            repo,
            drafts,
            dry_run=not args.real,
            env_file=args.env_file,
            trace_writer=trace_writer,
            trace_context=trace_context,
        )

    ingested_count = sum(1 for record in records if record.get("status") == "ingested")
    failed_count = sum(1 for record in records if record.get("status") == "failed")
    skipped_already_real = (
        0
        if args.include_real_ingested
        else max(0, counts["already_real_ingested"])
    )
    print(json.dumps(
        {
            "run_id": str(run_id),
            "pipeline_run_id": str(pipeline_run_id) if pipeline_run_id else None,
            "dry_run": not args.real,
            "total_payloads": counts["total_payloads"],
            "already_real_ingested": counts["already_real_ingested"],
            "selected_payloads": len(drafts),
            "skipped_already_real_ingested": skipped_already_real,
            "ingested_count": ingested_count,
            "failed_count": failed_count,
            "records": records,
        },
        ensure_ascii=False,
        default=str,
        indent=2,
    ))
    return 0


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