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

#!/usr/bin/env python3
"""Backfill pipeline trace ledger rows from existing acquisition runs.

Dry-run is the default. This script never fabricates historical LLM traces:
old runs only get structural events plus an explicit backfill note.
"""
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  # noqa: E402
from core.db_session import transaction  # noqa: E402
from pipeline.postgres import PostgresPipelineRepository  # noqa: E402
from pipeline.tracing import TraceContext  # 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 _acquisition_runs(conn, *, limit: int | None = None) -> list[dict]:
    sql = """
        SELECT ar.*,
               COUNT(DISTINCT aj.id)::int AS job_count,
               COUNT(DISTINCT ci.id)::int AS candidate_count
        FROM acquisition_runs ar
        LEFT JOIN acquisition_jobs aj ON aj.run_id = ar.id
        LEFT JOIN candidate_items ci ON ci.job_id = aj.id
        GROUP BY ar.id
        ORDER BY ar.created_at
    """
    if limit is not None:
        sql += " LIMIT %s"
        params = (limit,)
    else:
        params = ()
    return PostgresPipelineRepository(conn)._all(sql, params)


def _candidate_hits(conn, acquisition_run_id: UUID) -> list[dict]:
    return PostgresPipelineRepository(conn)._all(
        """
        SELECT
            aj.id AS acquisition_job_id,
            aj.query_id,
            aj.platform,
            ci.id AS item_id,
            ci.unique_key,
            ci.platform_item_id,
            ci.metadata,
            ci.source_payload
        FROM acquisition_jobs aj
        JOIN candidate_items ci ON ci.job_id = aj.id
        WHERE aj.run_id = %s
        ORDER BY aj.platform, ci.created_at
        """,
        (acquisition_run_id,),
    )


def backfill(*, env_file: str, apply: bool, limit: int | None = None) -> dict:
    db_config = CreationDbConfig.from_env(env_file)
    with transaction(db_config) as conn:
        repo = PostgresPipelineRepository(conn)
        runs = _acquisition_runs(conn, limit=limit)
        if not apply:
            return {
                "dry_run": True,
                "acquisition_run_count": len(runs),
                "candidate_hit_count": sum(int(row.get("candidate_count") or 0) for row in runs),
                "runs": [
                    {
                        "acquisition_run_id": str(row["id"]),
                        "run_key": row.get("run_key"),
                        "status": row.get("status"),
                        "job_count": row.get("job_count"),
                        "candidate_count": row.get("candidate_count"),
                    }
                    for row in runs
                ],
            }

        pipeline_count = 0
        hit_count = 0
        for row in runs:
            acquisition_run_id = row["id"]
            run = repo.create_pipeline_run(
                run_key=f"backfill:{acquisition_run_id}",
                batch_id=row.get("batch_id"),
                status=row.get("status") or "done",
                current_stage="ingest",
                metadata={
                    "source": "backfill_pipeline_traces",
                    "acquisition_run_id": str(acquisition_run_id),
                    "original_run_key": row.get("run_key"),
                },
            )
            pipeline_count += 1
            context = TraceContext(
                pipeline_run_id=run.id,
                acquisition_run_id=acquisition_run_id,
                stage="search",
            )
            repo.append_event(
                context=context,
                stage="query",
                event_type="backfilled_without_llm_trace",
                status="done",
                target_table="acquisition_runs",
                target_id=acquisition_run_id,
                payload={
                    "reason": "historical raw LLM requests and responses were not stored",
                    "job_count": row.get("job_count"),
                    "candidate_count": row.get("candidate_count"),
                },
            )
            for hit in _candidate_hits(conn, acquisition_run_id):
                metadata = hit.get("metadata") or {}
                repo.record_candidate_hit(
                    context=context.child(
                        acquisition_job_id=hit.get("acquisition_job_id"),
                        query_id=hit.get("query_id"),
                        item_id=hit.get("item_id"),
                        platform=hit.get("platform"),
                    ),
                    item_id=hit.get("item_id"),
                    platform=hit.get("platform"),
                    unique_key=hit.get("unique_key"),
                    platform_item_id=hit.get("platform_item_id"),
                    search_provider=metadata.get("search_provider"),
                    detail_provider=metadata.get("detail_provider"),
                    attempt_index=metadata.get("retry_attempt_index"),
                    page_index=metadata.get("page_index"),
                    page_rank=metadata.get("page_rank"),
                    candidate_rank=metadata.get("candidate_rank"),
                    source_cursor=metadata.get("source_cursor"),
                    is_duplicate_hit=metadata.get("acquisition_match_status") == "existing",
                    raw_candidate=metadata.get("matched_candidate") or {},
                    metadata={"source": "backfill_pipeline_traces"},
                )
                hit_count += 1
        return {
            "dry_run": False,
            "pipeline_run_count": pipeline_count,
            "candidate_hit_count": hit_count,
        }


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("--apply", action="store_true", help="Actually write trace rows. Default is dry-run.")
    parser.add_argument("--limit", type=int, default=None)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    print(json.dumps(
        backfill(env_file=args.env_file, apply=args.apply, limit=args.limit),
        ensure_ascii=False,
        default=str,
        indent=2,
    ))
    return 0


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