vignette/scripts/run-alliance-calibration.py

#!/usr/bin/env python3
"""Run the Alliance agents on the versioned G0 gold scene pack.

The runner writes a raw, provenance-rich snapshot only.  Use
``compare-alliance-calibration.py`` to evaluate it against another version.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from uuid import NAMESPACE_URL, uuid5


REPO_ROOT = Path(__file__).resolve().parents[1]
API_ROOT = REPO_ROOT / "apps" / "api"
if str(API_ROOT) not in sys.path:
    sys.path.insert(0, str(API_ROOT))

from app.contracts.measurement import (  # noqa: E402
    ALLIANCE_DIMENSIONS,
    BenchmarkCase,
)
from app.engine_client import engine_client  # noqa: E402
from app.services.alliance_calibration import (  # noqa: E402
    AllianceCalibrationPrediction,
    AllianceCalibrationSnapshot,
)
from app.services.alliance_measurement import (  # noqa: E402
    TranscriptTurn,
    run_agent_assessment,
)


def _needed_perspectives(case: BenchmarkCase) -> tuple[str, ...]:
    perspectives = {
        expectation.perspective
        for expectation in case.expected
        if expectation.construct_key == "working_alliance"
        and expectation.dimension in ALLIANCE_DIMENSIONS
        and expectation.perspective
        in {"client_agent_report", "independent_observer"}
    }
    return tuple(sorted(perspectives))


async def _run(args: argparse.Namespace) -> int:
    payload = json.loads(args.benchmark.read_text(encoding="utf-8"))
    cases = tuple(BenchmarkCase.model_validate(item) for item in payload["cases"])
    predictions: list[AllianceCalibrationPrediction] = []
    run_models: set[str] = set()
    run_providers: set[str] = set()
    prompt_versions: set[str] = set()
    await engine_client.startup()
    try:
        for case in cases:
            turns = tuple(
                TranscriptTurn(
                    turn_id=uuid5(NAMESPACE_URL, f"{case.case_id}:turn:{index}"),
                    seq=index + 1,
                    speaker=turn.speaker,
                    text=turn.text,
                )
                for index, turn in enumerate(case.turns)
            )
            for perspective in _needed_perspectives(case):
                session_id = uuid5(
                    NAMESPACE_URL,
                    f"{args.run_id}:{case.case_id}:session",
                )
                pulse_id = uuid5(
                    NAMESPACE_URL,
                    f"{args.run_id}:{case.case_id}:{perspective}:pulse",
                )
                result = await run_agent_assessment(
                    pulse_id=pulse_id,
                    session_id=session_id,
                    checkpoint="post",
                    perspective=perspective,
                    turns=turns,
                    engine=engine_client,
                )
                run_models.add(result.model_run.model)
                run_providers.add(result.model_run.provider)
                prompt_versions.add(result.model_run.prompt_bundle_version)
                assessment = result.assessment.by_dimension() if result.assessment else {}
                for dimension in ALLIANCE_DIMENSIONS:
                    item = assessment.get(dimension)
                    predictions.append(
                        AllianceCalibrationPrediction(
                            case_id=case.case_id,
                            perspective=perspective,
                            dimension=dimension,
                            value=item.score if item is not None else None,
                            confidence=item.confidence if item is not None else None,
                            evidence_turn_indices=(
                                item.evidence_turn_indices if item is not None else ()
                            ),
                            status=result.measurement_status,
                            error_code=result.error_code,
                            provider=result.model_run.provider,
                            model=result.model_run.model,
                            prompt_bundle_version=result.model_run.prompt_bundle_version,
                            model_run_id=str(result.model_run.model_run_id),
                            attempt_count=len(result.all_model_runs),
                            prior_model_run_ids=tuple(
                                str(run.model_run_id)
                                for run in result.prior_model_runs
                            ),
                            prior_error_codes=tuple(
                                run.error_code or "unknown_error"
                                for run in result.prior_model_runs
                            ),
                            prompt_bundle_hash=result.model_run.prompt_bundle_hash,
                            input_evidence_hash=result.model_run.input_evidence_hash,
                        )
                    )
                print(
                    json.dumps(
                        {
                            "case_id": case.case_id,
                            "perspective": perspective,
                            "status": result.measurement_status,
                            "provider": result.model_run.provider,
                            "model": result.model_run.model,
                        },
                        ensure_ascii=False,
                    ),
                    flush=True,
                )
    finally:
        await engine_client.shutdown()

    snapshot = AllianceCalibrationSnapshot(
        run_id=args.run_id,
        provider=next(iter(run_providers)) if len(run_providers) == 1 else "mixed",
        model=next(iter(run_models)) if len(run_models) == 1 else "mixed",
        prompt_bundle_version=(
            next(iter(prompt_versions)) if len(prompt_versions) == 1 else "mixed"
        ),
        generated_at=datetime.now(timezone.utc),
        predictions=tuple(predictions),
    )
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(
        json.dumps(snapshot.model_dump(mode="json"), ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(str(args.out))
    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-id", required=True)
    parser.add_argument(
        "--benchmark",
        type=Path,
        default=API_ROOT / "app" / "data" / "outcome_alliance_benchmark_g0.v1.json",
    )
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    return asyncio.run(_run(args))


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