"""Audit the frozen E08 design and generate its balanced reliability cells.""" from __future__ import annotations from collections import Counter from hashlib import sha256 import json from pathlib import Path from typing import Any from agent_harness.repository import GitSnapshot from agent_harness.specs import ( load_agent_systems, load_embeddings, load_experiments, load_harnesses, load_models, load_repositories, load_task_split, load_tasks, validate_configuration_tree, ) from agent_harness.study2_experiment import tokenizer_for REPEAT_TREATMENTS = ("H000", "H003", "H007", "H011", "A001", "A002") MODEL_IDS = ("M002", "M003") REPOSITORY_IDS = ("R001", "R002", "R003") def validation_record(root: Path, task_id: str) -> dict[str, Any]: candidates = ( root / "tasks" / "validation" / "study2" / f"{task_id}.json", root / "tasks" / "validation" / f"{task_id}.json", ) path = next((item for item in candidates if item.exists()), None) if path is None: raise RuntimeError(f"missing hidden-test validation record for {task_id}") value = json.loads(path.read_text(encoding="utf-8")) valid = value.get("valid_end_to_end") if valid is None: valid = value.get("validation", {}).get("valid") if valid is not True: raise RuntimeError(f"task validation is not successful for {task_id}: {path}") return value def repository_metrics( root: Path, repository_id: str, models: dict[str, Any], ) -> dict[str, Any]: repository = load_repositories(root)[repository_id] snapshot = GitSnapshot(root / repository.local_path) files = tuple(snapshot.iter_files(repository.pinned_head, repository.source_suffixes)) text = "\n".join(f"FILE: {item.path}\n{item.text}" for item in files) return { "repository_id": repository_id, "name": repository.name, "language": repository.language, "pinned_head": repository.pinned_head, "source_files": len(files), "source_lines": sum(item.text.count("\n") + 1 for item in files), "source_bytes": sum(len(item.text.encode("utf-8")) for item in files), "full_source_tokens": { model_id: tokenizer_for(models[model_id]).count(text) for model_id in MODEL_IDS }, } def balanced_repeat_cells(root: Path, split: tuple[str, ...]) -> list[dict[str, str]]: tasks = load_tasks(root) repositories = load_repositories(root) by_repository = { repository_id: [ task_id for task_id in split if tasks[task_id].repository_url == repositories[repository_id].repository_url ] for repository_id in REPOSITORY_IDS } strata = [ (repository_id, model_id) for repository_id in REPOSITORY_IDS for model_id in MODEL_IDS ] cells: list[dict[str, str]] = [] for treatment_index, treatment_id in enumerate(REPEAT_TREATMENTS): excluded = {treatment_index % 6, (treatment_index + 1) % 6} for stratum_index, (repository_id, model_id) in enumerate(strata): if stratum_index in excluded: continue task_id = min( by_repository[repository_id], key=lambda value: sha256( f"E08-repeat|{treatment_id}|{repository_id}|{model_id}|{value}".encode() ).hexdigest(), ) cells.append( { "task_id": task_id, "treatment_id": treatment_id, "model_id": model_id, } ) if len(cells) != 24 or len({tuple(item.values()) for item in cells}) != 24: raise RuntimeError("reliability selection did not produce 24 unique cells") treatment_counts = Counter(item["treatment_id"] for item in cells) stratum_counts = Counter( (next( repository_id for repository_id in REPOSITORY_IDS if tasks[item["task_id"]].repository_url == repositories[repository_id].repository_url ), item["model_id"]) for item in cells ) if set(treatment_counts.values()) != {4} or set(stratum_counts.values()) != {4}: raise RuntimeError( f"reliability balance failed: treatments={treatment_counts}, strata={stratum_counts}" ) return cells def audit(root: Path, write: bool) -> dict[str, Any]: errors, warnings = validate_configuration_tree(root) if errors or warnings: raise RuntimeError(f"configuration audit failed: errors={errors}, warnings={warnings}") experiment = load_experiments(root)["E08"] models = load_models(root) tasks = load_tasks(root) repositories = load_repositories(root) harnesses = load_harnesses(root) systems = load_agent_systems(root) embeddings = load_embeddings(root) split = load_task_split(root / "tasks" / "splits" / "study2_confirmatory.txt") if len(split) != 60: raise RuntimeError(f"Study 2 split has {len(split)} tasks instead of 60") for task_id in split: validation_record(root, task_id) task = tasks[task_id] for relative in (task.gold_patch, task.test_patch): if not (root / "tasks" / relative).is_file(): raise RuntimeError(f"missing frozen patch for {task_id}: {relative}") repository_counts = Counter( next( repository_id for repository_id, repository in repositories.items() if repository.repository_url == tasks[task_id].repository_url ) for task_id in split ) if repository_counts != Counter({"R001": 20, "R002": 20, "R003": 20}): raise RuntimeError(f"unbalanced repository task counts: {repository_counts}") metrics = [ repository_metrics(root, repository_id, models) for repository_id in REPOSITORY_IDS ] if any( min(item["full_source_tokens"].values()) <= experiment.context_budgets[0] for item in metrics ): raise RuntimeError("at least one repository fits inside the Study 2 context cap") repeat_cells = balanced_repeat_cells(root, split) result = { "schema_version": 1, "experiment_id": experiment.experiment_id, "task_count": len(split), "repository_task_counts": dict(sorted(repository_counts.items())), "language_task_counts": dict( sorted(Counter(tasks[item].language for item in split).items()) ), "multi_file_tasks": sum(len(tasks[item].gold_files) > 1 for item in split), "component_treatments": { item: harnesses[item].config_hash for item in experiment.harness_ids }, "controlled_systems": { item: systems[item].config_hash for item in experiment.agent_system_ids }, "models": {item: models[item].config_hash for item in experiment.model_ids}, "embedding": { "embedding_id": experiment.embedding_id, "config_hash": embeddings[experiment.embedding_id].config_hash, }, "main_generation_profile": { "temperature": 0.0, "top_p": 1.0, "seeds": list(experiment.seeds), }, "reliability_generation_profile": { "temperature": 0.2, "top_p": 1.0, "seeds": [0, 1, 2], }, "repositories": metrics, "main_cells": experiment.cells_per_task() * len(split), "repeat_base_cells": len(repeat_cells), "repeat_additional_cells": len(repeat_cells) * 3, "planned_live_cells": experiment.cells_per_task() * len(split) + len(repeat_cells) * 3, "repeat_cells": repeat_cells, } if write: reliability_dir = root / "configs" / "reliability" reliability_dir.mkdir(parents=True, exist_ok=True) (reliability_dir / "E08_repeat_cells.json").write_text( json.dumps( { "schema_version": 1, "experiment_id": "E08", "description": ( "Twenty-four balanced non-oracle E08 cells repeated at " "temperature 0.2 with seeds 0, 1, and 2." ), "temperature": 0.2, "top_p": 1.0, "seeds": [0, 1, 2], "cells": repeat_cells, }, indent=2, sort_keys=True, ) + "\n", encoding="utf-8", ) (root / "docs" / "STUDY2_DESIGN_AUDIT.json").write_text( json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" ) return result def main() -> None: import argparse parser = argparse.ArgumentParser() parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1]) parser.add_argument("--write", action="store_true") arguments = parser.parse_args() result = audit(arguments.root.resolve(), arguments.write) print(json.dumps(result, indent=2, sort_keys=True)) if __name__ == "__main__": main()