| """Persisted "repair library" — the model's accumulated knowledge of |
| known breakage -> repair pairs. Curated from successful rollouts during |
| training. Loaded at inference time as a few-shot prefix when the agent |
| recognises a familiar error class. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import asdict, dataclass, field |
| from pathlib import Path |
| from typing import Any, Optional |
|
|
|
|
| @dataclass |
| class RepairExample: |
| primitive_type: str |
| breakage_params: dict[str, Any] |
| error_signature: str |
| repair_diff: str |
| visible_reward: float |
| held_out: dict[str, float] |
| task_id: str = "" |
|
|
| def signature_key(self) -> str: |
| return f"{self.primitive_type}::{self.error_signature[:80]}" |
|
|
|
|
| @dataclass |
| class RepairLibrary: |
| examples: list[RepairExample] = field(default_factory=list) |
|
|
| def add(self, example: RepairExample) -> None: |
| self.examples.append(example) |
|
|
| def best_match(self, primitive_type: str, error_text: str) -> Optional[RepairExample]: |
| """Return the highest-reward example whose primitive_type matches and |
| whose error text overlaps.""" |
| candidates = [ |
| e for e in self.examples if e.primitive_type == primitive_type |
| ] |
| if not candidates: |
| return None |
| scored = sorted( |
| candidates, |
| key=lambda e: ( |
| _ngram_overlap(e.error_signature, error_text), |
| e.visible_reward, |
| ), |
| reverse=True, |
| ) |
| return scored[0] if scored else None |
|
|
| def to_dict(self) -> dict: |
| return { |
| "version": "1", |
| "examples": [asdict(e) for e in self.examples], |
| "size": len(self.examples), |
| "by_primitive": _count_by_primitive(self.examples), |
| } |
|
|
| def save(self, path: str | Path) -> None: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps(self.to_dict(), indent=2), encoding="utf-8") |
|
|
| @classmethod |
| def load(cls, path: str | Path) -> "RepairLibrary": |
| data = json.loads(Path(path).read_text(encoding="utf-8")) |
| examples = [RepairExample(**e) for e in data.get("examples", [])] |
| return cls(examples=examples) |
|
|
|
|
| def _ngram_overlap(a: str, b: str, n: int = 3) -> float: |
| if not a or not b: |
| return 0.0 |
|
|
| def grams(text: str) -> set[str]: |
| text = text.lower() |
| return {text[i : i + n] for i in range(len(text) - n + 1)} |
|
|
| ga, gb = grams(a), grams(b) |
| if not ga or not gb: |
| return 0.0 |
| return len(ga & gb) / max(1, len(ga | gb)) |
|
|
|
|
| def _count_by_primitive(examples: list[RepairExample]) -> dict[str, int]: |
| counts: dict[str, int] = {} |
| for e in examples: |
| counts[e.primitive_type] = counts.get(e.primitive_type, 0) + 1 |
| return counts |
|
|
|
|
| def curate_from_rollouts( |
| rollout_results: list, |
| min_reward: float = 0.6, |
| min_held_out_clean: float = 0.5, |
| ) -> RepairLibrary: |
| """Build a RepairLibrary from a list of rollout dicts/RolloutResults.""" |
| lib = RepairLibrary() |
| for r in rollout_results: |
| get = r.get if isinstance(r, dict) else lambda k, default=None: getattr(r, k, default) |
| if float(get("visible_reward", 0.0) or 0.0) < min_reward: |
| continue |
| if float(get("held_out_breakdown", {}).get("executed_cleanly", 0.0)) < min_held_out_clean: |
| continue |
| lib.add( |
| RepairExample( |
| primitive_type=str(get("primitive_type", "unknown")), |
| breakage_params=dict(get("info", {}).get("breakage_spec", {}).get("params", {})) |
| if isinstance(get("info", {}), dict) |
| else {}, |
| error_signature=str(get("error_trace", "") or "")[:160], |
| repair_diff=str(get("repair_completion", "") or get("info", {}).get("repair_diff", ""))[:2000], |
| visible_reward=float(get("visible_reward", 0.0) or 0.0), |
| held_out=dict(get("held_out_breakdown", {}) or {}), |
| task_id=str(get("task_id", "")), |
| ) |
| ) |
| return lib |
|
|
