Initial commit

2026-09-04 14:58:42 +08:00
commit 439cad87d9
4601 changed files with 29440 additions and 0 deletions
@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from ..models import RolloutTrace
+from ..storage import atomic_write_jsonl
+from .agentrm import AgentRM
+from .pre_score import PreScorer, adjust_agent_score
+from .identity import composite_id
+
+
+class TraceScorer:
+    """Resolve pre-scores and AgentRM scores into one effective score per trace."""
+
+    def __init__(
+        self,
+        pre_scorer: PreScorer,
+        agentrm: AgentRM | None = None,
+    ):
+        self.pre_scorer = pre_scorer
+        self.agentrm = agentrm if agentrm is not None else AgentRM()
+
+    def score_all(
+        self,
+        traces: list[RolloutTrace],
+        final_score_dir: Path,
+        agentrm_dir: Path | None = None,
+    ) -> dict[str, float]:
+        final_score_dir.mkdir(parents=True, exist_ok=True)
+        agentrm_dir = agentrm_dir or final_score_dir
+        agentrm_dir.mkdir(parents=True, exist_ok=True)
+        results = self.pre_scorer.score_all(traces)
+        pre_scores = {result.trace_id: result for result in results}
+        if set(pre_scores) != {trace.trace_id for trace in traces}:
+            raise ValueError("pre-scorer returned incomplete or duplicate results")
+        atomic_write_jsonl(
+            final_score_dir / "pre_scores.jsonl",
+            [pre_scores[trace.trace_id].to_dict() for trace in traces],
+        )
+        agentrm_traces = [
+            trace for trace in traces if pre_scores[trace.trace_id].route == "agentrm"
+        ]
+        requests = [trace.agentrm_request() for trace in agentrm_traces]
+        local_scores = agentrm_dir / "agentrm_scores.jsonl"
+        found = self.agentrm.score_requests(requests)
+        atomic_write_jsonl(local_scores, found)
+
+        by_key = {composite_id(row): float(row["score"]) for row in found}
+        resolved = []
+        scores: dict[str, float] = {}
+        for trace in traces:
+            pre_score = pre_scores[trace.trace_id]
+            if pre_score.route == "fixed_score":
+                if pre_score.score is None:
+                    raise ValueError(f"fixed pre-score is missing for {trace.trace_id}")
+                score = pre_score.score
+                source = "pre_score"
+            else:
+                score = adjust_agent_score(by_key[trace.key.as_tuple()], pre_score)
+                source = "agentrm"
+            scores[trace.trace_id] = score
+            resolved.append({
+                "task_name": trace.task_name,
+                "compile_type": trace.compile_type,
+                "test_name": trace.test_name,
+                "trace_id": trace.trace_id,
+                "effective_score": score,
+                "score_source": source,
+                "pre_score_reason": pre_score.reason,
+                "agentrm_score": by_key.get(trace.key.as_tuple()),
+                "efficiency_cost": pre_score.efficiency_cost,
+                "efficiency_penalty": pre_score.efficiency_penalty,
+            })
+        atomic_write_jsonl(final_score_dir / "effective_scores.jsonl", resolved)
+        return scores