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

#!/usr/bin/env python3
"""Run formal acquisition, decode every coarse-hit item, and build payload drafts."""
from __future__ import annotations

import argparse
import json
import logging
from dataclasses import asdict, is_dataclass
from uuid import UUID

from acquisition.repositories.postgres import PostgresAcquisitionRepository
from acquisition.runner import DEFAULT_PLATFORMS, run_batch
from core.config import CreationDbConfig, IngestApiConfig, Settings
from core.db_session import transaction
from decode_content.ingest import KnowledgeIngestClient, ingest_payload_draft
from decode_content.repositories.postgres import PostgresDecodeRepository
from decode_content.service import DecodeService
from pipeline.decode_runner import run_decode_stage
from pipeline.postgres import PostgresPipelineRepository
from pipeline.tracing import TraceContext, TraceWriter, new_trace_writer

logger = logging.getLogger(__name__)


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 _ingest_payloads(
    repo: PostgresDecodeRepository,
    outputs: list,
    *,
    dry_run: bool,
    env_file: str,
    trace_writer: TraceWriter | None = None,
    trace_context: TraceContext | None = None,
) -> list[dict]:
    records: list[dict] = []
    client = None if dry_run else KnowledgeIngestClient(IngestApiConfig.from_env(env_file))
    for output in outputs:
        for draft in output.payload_drafts:
            if draft.id is None:
                continue
            context = (trace_context or TraceContext()).child(
                item_id=output.item_id,
                payload_draft_id=draft.id,
                stage="ingest",
            )
            if trace_writer:
                trace_writer.event(
                    context=context,
                    stage="ingest",
                    event_type="ingest_started",
                    target_table="payload_drafts",
                    target_id=draft.id,
                    payload={"dry_run": dry_run},
                )
            record = ingest_payload_draft(repo, draft, dry_run=dry_run, client=client)
            if trace_writer:
                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 _ensure_pipeline_run(
    db_config: CreationDbConfig,
    *,
    batch_id: UUID,
    run_key: str | None,
    platforms: tuple[str, ...],
    args: argparse.Namespace,
) -> UUID | None:
    try:
        with transaction(db_config) as conn:
            run = PostgresPipelineRepository(conn).create_pipeline_run(
                run_key=run_key or f"creation-pipeline:{batch_id}",
                batch_id=batch_id,
                status="running",
                current_stage="search",
                config={
                    "platforms": list(platforms),
                    "search_limit": args.search_limit,
                    "display_limit": args.display_limit,
                    "decode_limit": args.decode_limit,
                    "resume": not args.no_resume,
                    "skip_done": not args.no_skip_done,
                    "real_ingest": args.real_ingest,
                },
                metadata={"entrypoint": "scripts/run_creation_pipeline.py"},
            )
            return run.id
    except Exception as exc:
        logger.warning("pipeline run trace setup skipped: %s", exc)
        return None


def _mark_pipeline_run(
    db_config: CreationDbConfig,
    run_id: UUID | None,
    *,
    status: str,
    current_stage: str | None = None,
    error_message: str | None = None,
    metadata: dict | None = None,
) -> None:
    if run_id is None:
        return
    try:
        with transaction(db_config) as conn:
            PostgresPipelineRepository(conn).mark_pipeline_run_status(
                run_id,
                status=status,
                current_stage=current_stage,
                error_message=error_message,
                metadata=metadata,
            )
    except Exception as exc:
        logger.warning("pipeline run trace status skipped: %s", exc)


def _start_pipeline_job(
    db_config: CreationDbConfig,
    pipeline_run_id: UUID | None,
    *,
    stage: str,
    target_table: str | None = None,
    target_id: UUID | None = None,
    metadata: dict | None = None,
) -> UUID | None:
    if pipeline_run_id is None:
        return None
    try:
        with transaction(db_config) as conn:
            job = PostgresPipelineRepository(conn).save_pipeline_job(
                run_id=pipeline_run_id,
                stage=stage,
                target_table=target_table,
                target_id=target_id,
                status="running",
                metadata=metadata,
            )
            return job.id
    except Exception as exc:
        logger.warning("pipeline job trace setup skipped: %s", exc)
        return None


def _mark_pipeline_job(
    db_config: CreationDbConfig,
    job_id: UUID | None,
    *,
    status: str,
    error_message: str | None = None,
    metadata: dict | None = None,
) -> None:
    if job_id is None:
        return
    try:
        with transaction(db_config) as conn:
            PostgresPipelineRepository(conn).mark_job_status(
                job_id,
                status=status,
                error_message=error_message,
                metadata=metadata,
            )
    except Exception as exc:
        logger.warning("pipeline job trace status skipped: %s", exc)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--batch-id", required=True, help="Formal query batch UUID")
    parser.add_argument(
        "--platform",
        action="append",
        choices=DEFAULT_PLATFORMS,
        help="Platform to run. Repeat to run multiple platforms. Default: all.",
    )
    parser.add_argument("--search-limit", type=int, default=10)
    parser.add_argument("--display-limit", type=int, default=5)
    parser.add_argument("--decode-limit", type=int, default=100)
    parser.add_argument("--run-key")
    parser.add_argument("--env-file", default=".env")
    parser.add_argument("--no-resume", action="store_true")
    parser.add_argument("--no-skip-done", action="store_true")
    parser.add_argument("--no-dry-ingest-record", action="store_true")
    parser.add_argument("--real-ingest", action="store_true", help="Call CK_INGEST_API_URL instead of dry-run records.")
    return parser.parse_args(argv)


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 DEFAULT_PLATFORMS)
    batch_id = UUID(args.batch_id)
    pipeline_run_id = _ensure_pipeline_run(
        db_config,
        batch_id=batch_id,
        run_key=args.run_key,
        platforms=platforms,
        args=args,
    )
    trace_writer = new_trace_writer(db_config, env_file=args.env_file)
    search_job_id = _start_pipeline_job(
        db_config,
        pipeline_run_id,
        stage="search",
        target_table="query_batches",
        target_id=batch_id,
        metadata={"platforms": list(platforms)},
    )
    trace_context = TraceContext(pipeline_run_id=pipeline_run_id, pipeline_job_id=search_job_id, stage="search")

    try:
        trace_writer.event(
            context=trace_context,
            stage="search",
            event_type="pipeline_started",
            status="running",
            payload={"batch_id": str(batch_id), "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=not args.no_resume,
                skip_done=not args.no_skip_done,
                run_key=args.run_key,
                trace_writer=trace_writer,
                trace_context=trace_context,
            )
        _mark_pipeline_job(
            db_config,
            search_job_id,
            status="partial" if acquisition.failed else "done",
            metadata={"acquisition_run_id": str(acquisition.run_id), "failed": acquisition.failed},
        )

        decode_job_id = _start_pipeline_job(
            db_config,
            pipeline_run_id,
            stage="decode",
            target_table="acquisition_runs",
            target_id=acquisition.run_id,
            metadata={"decode_limit": args.decode_limit},
        )
        decode_context = trace_context.child(
            pipeline_job_id=decode_job_id,
            acquisition_run_id=acquisition.run_id,
            stage="decode",
        )
        _mark_pipeline_run(db_config, pipeline_run_id, status="running", current_stage="decode")
        with transaction(db_config) as conn:
            acquisition_repo = PostgresAcquisitionRepository(conn)
            decode_repo = PostgresDecodeRepository(conn)
            decode_service = DecodeService(
                settings=settings,
                repository=decode_repo,
                trace_writer=trace_writer,
                trace_context=decode_context,
            )
            decode = run_decode_stage(
                candidate_repo=acquisition_repo,
                decode_service=decode_service,
                run_id=acquisition.run_id,
                limit=args.decode_limit,
                trace_writer=trace_writer,
                trace_context=decode_context,
            )
            _mark_pipeline_job(
                db_config,
                decode_job_id,
                status="partial" if decode.failed else "done",
                metadata={"decode_total": decode.total, "decode_failed": decode.failed},
            )
            ingest_job_id = _start_pipeline_job(
                db_config,
                pipeline_run_id,
                stage="ingest",
                target_table="acquisition_runs",
                target_id=acquisition.run_id,
                metadata={"real_ingest": args.real_ingest},
            )
            ingest_context = decode_context.child(pipeline_job_id=ingest_job_id, stage="ingest")
            _mark_pipeline_run(db_config, pipeline_run_id, status="running", current_stage="ingest")
            ingest_records = [] if args.no_dry_ingest_record else _ingest_payloads(
                decode_repo,
                decode.outputs,
                dry_run=not args.real_ingest,
                env_file=args.env_file,
                trace_writer=trace_writer,
                trace_context=ingest_context,
            )
            ingest_failed = sum(1 for record in ingest_records if record.get("status") == "failed")
            _mark_pipeline_job(
                db_config,
                ingest_job_id,
                status="partial" if ingest_failed else "done",
                metadata={"ingest_record_count": len(ingest_records), "failed": ingest_failed},
            )
        final_status = "partial" if acquisition.failed or decode.failed else "done"
        _mark_pipeline_run(
            db_config,
            pipeline_run_id,
            status=final_status,
            current_stage="ingest",
            metadata={
                "acquisition_run_id": str(acquisition.run_id),
                "decode_total": decode.total,
                "decode_failed": decode.failed,
                "ingest_record_count": len(ingest_records),
            },
        )
        trace_writer.event(
            context=trace_context.child(acquisition_run_id=acquisition.run_id),
            stage="ingest",
            event_type="pipeline_finished",
            status=final_status,
            payload={"decode": _model_dump(decode), "ingest_record_count": len(ingest_records)},
        )
    except Exception as exc:
        _mark_pipeline_run(
            db_config,
            pipeline_run_id,
            status="failed",
            error_message=str(exc),
        )
        trace_writer.event(
            context=trace_context,
            stage="ingest",
            event_type="pipeline_failed",
            status="failed",
            severity="error",
            error_message=str(exc),
        )
        raise

    print(json.dumps(
        {
            "batch_id": str(batch_id),
            "pipeline_run_id": str(pipeline_run_id) if pipeline_run_id else None,
            "run_id": str(acquisition.run_id),
            "platforms": list(platforms),
            "acquisition": _model_dump(acquisition),
            "decode": _model_dump(decode),
            "ingest_records": ingest_records,
            "dry_ingest_records": ingest_records if not args.real_ingest else [],
        },
        ensure_ascii=False,
        default=str,
        indent=2,
    ))
    return 0


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