134 lines
4.9 KiB
Python
134 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from ..models import RolloutTrace
|
|
from ..scoring.state import acp_events_to_state, read_acp_events
|
|
|
|
|
|
VARIANT_TO_COMPILE_TYPE = {
|
|
"model_skill": "model_compile",
|
|
"ori_skill": "ori",
|
|
}
|
|
|
|
|
|
def event_to_state(trajectory_path: Path) -> tuple[list[dict[str, str]], bool]:
|
|
events = read_acp_events(trajectory_path)
|
|
skill_invoked = any(
|
|
event.get("type") == "tool_call"
|
|
and any(
|
|
str(event.get(field, "")).strip().lower() == "skill"
|
|
for field in ("title", "kind")
|
|
)
|
|
for event in events
|
|
)
|
|
return acp_events_to_state(events), skill_invoked
|
|
|
|
|
|
def trajectory_for_test(test_dir: Path) -> Path | None:
|
|
candidates = sorted(test_dir.rglob("acp_trajectory.jsonl"))
|
|
if not candidates:
|
|
return None
|
|
canonical = [path for path in candidates if "trajectory" in path.parts]
|
|
return canonical[0] if canonical else candidates[0]
|
|
|
|
|
|
def _result_for_trajectory(trajectory_path: Path) -> dict[str, Any]:
|
|
result_path = trajectory_path.parent.parent / "result.json"
|
|
if not result_path.is_file():
|
|
raise ValueError(f"missing structured BenchFlow result: {result_path}")
|
|
value = json.loads(result_path.read_text(encoding="utf-8"))
|
|
if not isinstance(value, dict):
|
|
raise ValueError(f"BenchFlow result must be a JSON object: {result_path}")
|
|
return value
|
|
|
|
|
|
def load_benchflow_trace(
|
|
test_dir: Path,
|
|
task_name: str,
|
|
compile_type: str,
|
|
) -> RolloutTrace:
|
|
trajectory = trajectory_for_test(test_dir)
|
|
if trajectory is None:
|
|
raise ValueError(f"missing acp_trajectory.jsonl under {test_dir}")
|
|
result = _result_for_trajectory(trajectory)
|
|
agent_timeout = result.get("agent_timeout_info")
|
|
idle_timeout = result.get("idle_timeout_info")
|
|
timeout_info = (
|
|
agent_timeout
|
|
if isinstance(agent_timeout, dict)
|
|
else idle_timeout if isinstance(idle_timeout, dict) else None
|
|
)
|
|
timed_out = timeout_info is not None
|
|
metadata: dict[str, Any] = {
|
|
"source": str(test_dir.resolve()),
|
|
"termination": "timeout" if timed_out else "completed",
|
|
}
|
|
timing = result.get("timing") if isinstance(result.get("timing"), dict) else {}
|
|
execution_seconds = timing.get("agent_execution")
|
|
if execution_seconds is None and isinstance(timeout_info, dict):
|
|
execution_seconds = timeout_info.get(
|
|
"wall_clock_elapsed_sec", timeout_info.get("timeout_sec")
|
|
)
|
|
if isinstance(execution_seconds, (int, float)):
|
|
metadata["agent_execution_seconds"] = float(execution_seconds)
|
|
if isinstance(result.get("n_tool_calls"), int):
|
|
metadata["tool_calls"] = result["n_tool_calls"]
|
|
if timed_out:
|
|
metadata["timeout_reason"] = timeout_info.get("reason")
|
|
metadata["timeout_seconds"] = timeout_info.get(
|
|
"timeout_sec", timeout_info.get("idle_timeout_sec")
|
|
)
|
|
metadata["partial_trajectory"] = bool(result.get("partial_trajectory", False))
|
|
metadata["error_category"] = result.get("error_category")
|
|
state, skill_invoked = event_to_state(trajectory)
|
|
return RolloutTrace(
|
|
trace_id=f"{task_name}/{compile_type}/{test_dir.name}",
|
|
task_name=task_name,
|
|
compile_type=compile_type,
|
|
test_name=test_dir.name,
|
|
state=state,
|
|
skill_invoked=skill_invoked,
|
|
timed_out=timed_out,
|
|
metadata=metadata,
|
|
)
|
|
|
|
|
|
def _variant_dirs(input_path: Path) -> list[tuple[Path, str, str]]:
|
|
if input_path.name in VARIANT_TO_COMPILE_TYPE:
|
|
return [(
|
|
input_path,
|
|
input_path.parent.name,
|
|
VARIANT_TO_COMPILE_TYPE[input_path.name],
|
|
)]
|
|
direct_variants = [
|
|
(input_path / variant_name, input_path.name, compile_type)
|
|
for variant_name, compile_type in VARIANT_TO_COMPILE_TYPE.items()
|
|
if (input_path / variant_name).is_dir()
|
|
]
|
|
if direct_variants:
|
|
return direct_variants
|
|
variants = []
|
|
for task_dir in sorted(path for path in input_path.iterdir() if path.is_dir()):
|
|
for variant_name, compile_type in VARIANT_TO_COMPILE_TYPE.items():
|
|
variant_dir = task_dir / variant_name
|
|
if variant_dir.is_dir():
|
|
variants.append((variant_dir, task_dir.name, compile_type))
|
|
return variants
|
|
|
|
|
|
def load_benchflow_traces(input_path: Path) -> list[RolloutTrace]:
|
|
input_path = input_path.resolve()
|
|
if not input_path.is_dir():
|
|
raise ValueError(f"BenchFlow input does not exist: {input_path}")
|
|
variants = _variant_dirs(input_path)
|
|
if not variants:
|
|
raise ValueError(f"no model_skill or ori_skill directories under {input_path}")
|
|
traces = []
|
|
for variant_dir, task_name, compile_type in variants:
|
|
for test_dir in sorted(path for path in variant_dir.glob("test-*") if path.is_dir()):
|
|
traces.append(load_benchflow_trace(test_dir, task_name, compile_type))
|
|
return traces
|