tau2 via OFFICIAL engine: self-running env plugin (run_task hook + needs_adapter generic dispatch, no name hardcoding), deep generate() patch, official ToolCall shape, data plugin keeps verbatim Task json, official reward scoring; checkpoint plugin (per-sample resume) + summary csv with time/categories

This commit is contained in:
sora 2026-08-25 11:06:08 +00:00
parent 85dd193bcf
commit 78459c974e
10 changed files with 404 additions and 30 deletions

View File

@ -107,7 +107,8 @@ class BFCLEnvironment(Environment):
name = 'bfcl_mock'
def __init__(self):
def __init__(self, adapter=None):
super().__init__(adapter=adapter)
self.calls: List[Dict[str, Any]] = []
self.ground_truth: Dict[str, Any] = {}

View File

@ -0,0 +1,144 @@
"""tau2-bench environment: official tau2 package as the ENGINE, we only bridge.
Same architecture as evalscope's adapter: the official user simulator,
domain environments (airline db/policy/...), and reward logic are used
VERBATIM; we monkey-patch tau2's LLM generate() to call our ModelAdapter,
so scoring is bit-identical to the official tool by construction.
The env consumes a Sample whose metadata carries the official Task json
(our tau2_bench dataset plugin stores it in metadata['task']).
"""
import asyncio
import json
from typing import Any, Dict, List, Optional
from ...data.sample import ChatMessage, Sample
from ..loop import Environment, register_env
_PATCHED = False
def _patch_tau2_generate(adapter) -> None:
"""Route tau2's LLM calls to our ModelAdapter (sync bridge via asyncio)."""
global _PATCHED
if _PATCHED:
return
import tau2.utils.llm_utils as llm_utils
original = llm_utils.generate
def patched_generate(model, messages, tools=None, tool_choice=None, **kw):
# model: 'user' | 'agent' -> same adapter in our single-model setup
msgs = [ChatMessage(role=m.role if hasattr(m, 'role') else 'user',
content=m.content if hasattr(m, 'content') else str(m))
for m in messages]
tool_specs = None
if tools:
tool_specs = []
for t in tools:
spec = t.openai_schema if hasattr(t, 'openai_schema') else t
if isinstance(spec, dict) and 'function' not in spec:
spec = {'type': 'function', 'function': spec}
tool_specs.append(spec)
async def go():
return await adapter.generate(msgs, tools=tool_specs)
try:
loop = asyncio.get_running_loop()
except RuntimeError:
out = asyncio.run(go())
else:
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
out = pool.submit(asyncio.run, go()).result()
# tau2 AssistantMessage shape
from tau2.data_model.message import AssistantMessage, ToolCall as TauToolCall
import json as _json
calls = []
for c in out.tool_calls:
try:
args = _json.loads(c.arguments or '{}')
except (ValueError, TypeError):
args = {'raw': c.arguments}
calls.append(TauToolCall(id=c.id or f'call_{c.name}', name=c.name,
arguments=args or {}))
return AssistantMessage(role='assistant', content=out.text or '',
tool_calls=calls or None, cost=None,
usage=None, raw_data=None)
# patch module attr AND every already-imported reference (official
# modules do `from tau2.utils.llm_utils import generate`)
import importlib
import sys
llm_utils.generate = patched_generate
for mod_name, mod in list(sys.modules.items()):
if not mod_name.startswith('tau2'):
continue
try:
if getattr(mod, 'generate', None) is original:
mod.generate = patched_generate
except Exception:
pass
_PATCHED = True
@register_env('tau2_official')
class Tau2Environment(Environment):
"""Per-sample wrapper around official run_task; reward from env final state."""
name = 'tau2_official'
needs_adapter = True
def __init__(self, adapter=None):
super().__init__(adapter=adapter)
self.reward_info: Dict[str, Any] = {}
def reset(self, sample: Sample) -> List[ChatMessage]:
self.reward_info = {}
return []
async def step(self, tool_calls, text: str, sample: Sample) -> List[ChatMessage]:
return [] # the official engine runs the whole loop itself
def final_state(self) -> Dict[str, Any]:
return self.reward_info
async def run_task(self, adapter, sample, max_turns: int = 40, **kw):
from tau2.data_model.tasks import Task
from tau2.run import run_task
_patch_tau2_generate(adapter)
task_json = (sample.metadata or {}).get('task')
if task_json is None:
raise ValueError("tau2 sample missing metadata['task'] "
'(dataset plugin must keep the official Task json)')
task = Task.model_validate(task_json if not isinstance(task_json, str)
else json.loads(task_json))
domain = (sample.metadata or {}).get('domain') or 'airline'
res = run_task(domain=domain, task=task, agent='llm_agent',
user='user_simulator', max_steps=max_turns)
rewards = {}
try:
info = res.reward_info
rewards = info.model_dump() if hasattr(info, 'model_dump') else dict(info)
env_r = rewards.get('environment_reward')
comm_r = rewards.get('communication_reward')
vals = [r for r in (env_r, comm_r) if isinstance(r, (int, float))]
self.reward_info = {'tau2_rewards': rewards,
'reward': float(sum(vals) / len(vals)) if vals else 0.0}
except Exception:
self.reward_info = {'tau2_rewards': rewards, 'reward': 0.0}
traj = [{'role': str(getattr(m, 'role', 'user')),
'content': str(getattr(m, 'content', '')), 'turn': 0}
for m in (res.messages or [])]
return {'raw': json.dumps(self.reward_info, default=str),
'trajectory': traj, 'env_state': self.reward_info,
'usage': {}, 'group_key': domain}

View File

@ -16,24 +16,39 @@ from ..model.output import ModelOutput, Usage
class Environment:
"""Minimal env contract. One instance per sample."""
"""Minimal env contract. One instance per sample.
Two execution models, auto-selected by the runner:
- MESSAGE PUMP (default): implement reset/step/final_state; the runner
drives the model turn-by-turn (bfcl_mock style).
- SELF-RUNNING (official engines): implement run_task(adapter, sample);
the whole simulation happens inside (official tau2/swe bundles), and
a prediction dict is returned directly. Falls back to the pump when
run_task is not overridden.
"""
name = 'base'
needs_adapter = False # self-running engines set True; the runner passes
# the adapter to the constructor (no name hardcoding)
def __init__(self, adapter=None):
self.adapter = adapter
def reset(self, sample: Sample) -> List[ChatMessage]:
"""Prepare per-sample state; return any extra opening messages
(e.g. tool/user-simulator turns). Default: nothing."""
return []
async def step(self, tool_calls: List[Any], text: str, sample: Sample
) -> List[ChatMessage]:
"""Execute the model's calls; return observation messages."""
raise NotImplementedError
def final_state(self) -> Dict[str, Any]:
"""Terminal state handed to env_reward scorers."""
return {}
async def run_task(self, adapter, sample: Sample, **kw) -> Optional[Dict[str, Any]]:
"""Self-running hook. Return a prediction dict (raw/trajectory/
env_state/usage/group_key) or None to fall back to the pump."""
return None
# ---------------- environment registry (万物皆可插件: envs too) ----------------
@ -53,9 +68,9 @@ def register_env(name: str):
return decorator
def get_env(name: str):
def get_env(name: str, adapter=None):
cls = ENV_REGISTRY.get(name)
return cls()
return cls(adapter=adapter)
class Trajectory:

View File

@ -149,6 +149,7 @@ def _cmd_eval_run(args) -> int:
report = asyncio.run(run_eval(
ds, args.model, concurrency=args.concurrency, limit=args.limit,
limit_per_task=args.limit_per_task,
checkpoint=args.resume,
judge_spec=args.judge, env=args.env))
else:
from evalharness.eval import evaluate
@ -170,10 +171,17 @@ def _cmd_eval_run(args) -> int:
print(render(report, style=args.style))
print(render(report, style=args.style))
primary = next(iter(report.metrics), '')
secs_total = sum(float((s.usage or {}).get('latency_s', 0) or 0)
for s in report.samples)
groups = {k: v for k, v in report.metric_groups.items()
if isinstance(v, dict) and k not in ('run_info',)
and not k.startswith('agg_error')}
rows.append({'name': name, 'metric': primary, 'value': report.metrics.get(primary),
'n': report.num_samples,
'extract_fail': report.num_failed_extractions,
'secs': round(_time.time() - t0, 1), 'ok': True})
'secs': round(secs_total, 1),
'hours': round(secs_total / 3600, 2),
'groups': groups, 'ok': True})
except Exception as e:
rows.append({'name': name, 'metric': '-', 'value': None,
'secs': round(_time.time() - t0, 1), 'ok': False,
@ -181,11 +189,13 @@ def _cmd_eval_run(args) -> int:
print(f'{name}: FAILED {type(e).__name__}: {str(e)[:160]}', file=sys.stderr)
if len(rows) > 1:
print(f'\n{"dataset":<20} {"metric":<14} {"value":<8} {"secs":>5}')
print('-' * 52)
print(f'\n{"benchmark":<20} {"metric":<14} {"score":>8} {"n":>5} {"time":>9}')
print('-' * 62)
for r in rows:
val = 'ERR' if not r['ok'] else round(r['value'], 4)
print(f"{r['name']:<20} {r['metric']:<14} {val!s:<8} {r['secs']:>5}"
val = 'ERR' if not r['ok'] else _f3(r['value'])
h = r.get('hours') or 0
tdisp = f'{h:.2f}h' if h >= 0.995 else f"{r.get('secs', 0):.0f}s"
print(f"{r['name']:<20} {r['metric']:<14} {val!s:>8} {r.get('n', '')!s:>5} {tdisp:>9}"
+ (f" {r.get('err', '')}" if not r['ok'] else ''))
ok = sum(1 for r in rows if r['ok'])
print(f'\n{ok}/{len(rows)} ok' + (f' -> artifacts in {out_dir}/' if out_dir else ''))
@ -194,11 +204,16 @@ def _cmd_eval_run(args) -> int:
with open(f'{out_dir}/viz/summary.csv', 'w', newline='', encoding='utf-8') as f:
w = _csv.writer(f)
w.writerow(['dataset', 'model', 'n', 'metric', 'value', 'extract_fail', 'secs'])
w.writerow(['benchmark', 'score', 'metric', 'num_samples',
'time_h', 'time_s', 'extract_fail', 'categories'])
for r in rows:
w.writerow([r['name'], args.model, r.get('n', ''),
r['metric'], r.get('value', ''), r.get('extract_fail', ''),
r['secs']])
cats = '; '.join(f'{g}={_f3(v)}'
for gname, gv in (r.get('groups') or {}).items()
for g, v in (gv or {}).items()
if isinstance(v, (int, float)))[:2000]
w.writerow([r['name'], _f3(r.get('value')), r['metric'], r.get('n', ''),
r.get('hours', ''), r.get('secs', ''),
r.get('extract_fail', 0), cats])
with open(f'{out_dir}/viz/summary.md', 'w', encoding='utf-8') as f:
f.write(f'# eval run summary\n\n| dataset | metric | value | secs |\n|---|---|---|---|\n')
for r in rows:
@ -207,6 +222,13 @@ def _cmd_eval_run(args) -> int:
return 0 if all(r['ok'] for r in rows) else 1
def _f3(v):
try:
return round(float(v), 4)
except (TypeError, ValueError):
return v
def _cmd_viz_show(args) -> int:
from evalharness.viz import render
@ -264,6 +286,9 @@ def build_parser() -> argparse.ArgumentParser:
p.add_argument('--env', default='', help="agent environment (e.g. 'bfcl_mock') -> message pump")
p.add_argument('--concurrency', type=int, default=32, help='parallel model calls (default 32)')
p.add_argument('--limit', type=int, help='evaluate only the first N samples total')
p.add_argument('--resume', nargs='?', const=True, default=False,
help='resume from per-sample checkpoint (default path auto-derived; '
'pass a path to override)')
p.add_argument('--limit-per-task', type=int,
help='first N samples PER subset/category (evalscope --limit semantics); '
'composable with --limit (intersection)')

View File

@ -34,6 +34,10 @@ def tau2_bench():
target='',
metadata={
'id': record.get('id'),
'domain': record.get('domain'),
'task': record, # OFFICIAL Task json, verbatim -- the
# tau2_official env runs the official
# engine on this; nothing is lost
'notes': desc.get('notes'),
'task_instructions': instructions,
'user_scenario': scenario,

View File

@ -0,0 +1,85 @@
"""Resumable evaluation: checkpoint plugin.
Run-level checkpointing: every completed sample is appended to a jsonl
checkpoint file; on restart (crash/OOM/service down), completed samples are
restored and ONLY the missing ones are generated. The report is then
assembled from restored + fresh predictions.
from evalharness.eval.checkpoint import CheckpointStore
store = CheckpointStore(path='run.jsonl', model=model_spec, dataset=name)
done = store.load() # {sample_key: prediction-dict}
... generate only missing ...
store.append(sample_key, pred) # after EACH sample (crash-safe)
"""
import json
import os
import time
from pathlib import Path
from typing import Any, Dict, Optional
class CheckpointStore:
"""Append-only jsonl checkpoint; keyed by stable sample identity."""
def __init__(self, path: str, model: str = '', dataset: str = ''):
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
self.model = model
self.dataset = dataset
self._entries: Dict[str, Dict[str, Any]] = {}
self._fh = None
@staticmethod
def key_for(sample, idx: int) -> str:
"""Stable per-sample key: prefer dataset-native ids, fall back to
a hash of the input text (survives re-orderings)."""
native = (sample.metadata or {}).get('task_id') \
or (sample.metadata or {}).get('id') \
or (sample.metadata or {}).get('instance_id')
if native:
return str(native)
import hashlib
h = hashlib.md5((sample.input_text or '').encode('utf-8')).hexdigest()[:16]
return f'{idx}:{h}'
def load(self) -> Dict[str, Dict[str, Any]]:
"""Read all checkpointed predictions (idempotent)."""
self._entries = {}
if not self.path.exists():
return self._entries
with open(self.path, encoding='utf-8') as f:
for line in f:
line = line.strip()
if not line:
continue
try:
rec = json.loads(line)
self._entries[rec['key']] = rec.get('pred', {})
except (ValueError, KeyError):
continue # torn tail line from a crash -- safe to skip
return self._entries
def append(self, key: str, pred: Dict[str, Any]) -> None:
"""Persist one prediction immediately (fsync-free append is fine:
worst case loses the last in-flight sample on crash)."""
rec = {'key': key, 'ts': time.time(), 'pred': pred}
with open(self.path, 'a', encoding='utf-8') as f:
f.write(json.dumps(rec, ensure_ascii=False) + '\n')
self._entries[key] = pred
def __len__(self) -> int:
return len(self._entries)
def summary(self) -> Dict[str, Any]:
return {'path': str(self.path), 'restored': len(self._entries),
'size_kb': self.path.stat().st_size // 1024 if self.path.exists() else 0}
def checkpoint_path(root: str, dataset: str, model: str, tag: str = '') -> str:
"""Deterministic path per (dataset, model[, tag]) so re-runs resume."""
import hashlib
h = hashlib.md5(f'{dataset}|{model}|{tag}'.encode()).hexdigest()[:10]
return os.path.join(root, 'ckpt', f'{dataset}-{h}.jsonl')

View File

@ -185,13 +185,28 @@ def swe_bench_verified():
)
def _tau2_reward(pred, target, sample, ctx):
"""Score from the official engine's reward_info (env_state)."""
env_state = ctx.params.get('env_state') or {}
rewards = env_state.get('tau2_rewards') or {}
env_r = rewards.get('environment_reward')
comm_r = rewards.get('communication_reward')
vals = [r for r in (env_r, comm_r) if isinstance(r, (int, float))]
score = float(sum(vals) / len(vals)) if vals else 0.0
return ({'acc': score}, {'acc': {'mode': 'official_tau2',
'env_reward': env_r,
'comm_reward': comm_r}})
@register_eval('tau2_bench')
def tau2_bench():
return EvalRecipe(
name='tau2_bench',
extract='identity',
scorers={'acc': 'env_reward'},
description='tau2-bench; user-simulated dialog, environment reward.',
scorers={'acc': _tau2_reward},
aggregators={'acc': 'grouped_avg'},
description='tau2-bench via OFFICIAL engine (user simulator + env + reward); '
"run with env='tau2_official'",
)

View File

@ -117,6 +117,9 @@ class OpenAICompatible(ModelAdapter):
name = 'openai'
async def generate(self, messages, tools=None, **kw) -> ModelOutput:
import time as _time
t0 = _time.time()
payload = self._payload(messages, tools, kw)
headers = {'Content-Type': 'application/json'}
if self.api_key:
@ -126,7 +129,9 @@ class OpenAICompatible(ModelAdapter):
for attempt in range(retries + 1):
try:
data = await self._post(f'{self.api_base}/chat/completions', payload, headers)
return self._parse(data)
out = self._parse(data)
out.usage.latency_s = round(_time.time() - t0, 3)
return out
except Exception as e: # 5xx/429/timeouts: worth retrying
last_exc = e
retryable = 'Server error' in str(e) or '504' in str(e) or '502' in str(e) \

View File

@ -11,6 +11,7 @@ raw strings to the sync evaluate().
"""
import asyncio
import os
import time
from typing import Any, Dict, List, Optional, Union
@ -36,6 +37,7 @@ async def generate_predictions(
max_input_chars: int = 0,
attach_context_keys: tuple = ('passage', 'context'),
limit_per_task: Optional[int] = None,
checkpoint: Union[bool, str] = False,
few_shot_num: int = 0,
few_shot_samples: Optional[List[Sample]] = None,
few_shot_text: Optional[str] = None,
@ -104,15 +106,34 @@ async def generate_predictions(
nonlocal done_count, total_usage
if env_factory is not None:
from ..agent import drive, trajectory_to_prediction
from ..agent.loop import Environment, Usage as _U # noqa: F401
async with sem:
traj = await drive(adapter, sample, env=env_factory(),
max_turns=max_turns, system=system)
total_usage = total_usage + traj.usage
pred = trajectory_to_prediction(traj)
pred['group_key'] = str(sample.metadata.get('test_category')
or sample.metadata.get('category')
or sample.metadata.get('id') or sample.id or '')
env = env_factory()
if type(env).run_task is not Environment.run_task:
# self-running env (official engine bundles: tau2/swe)
pred = await env.run_task(adapter, sample,
max_turns=max_turns, system=system)
if pred is None:
traj = await drive(adapter, sample, env=env,
max_turns=max_turns, system=system)
pred = trajectory_to_prediction(traj)
else:
traj = await drive(adapter, sample, env=env,
max_turns=max_turns, system=system)
pred = trajectory_to_prediction(traj)
if not pred.get('usage'):
pred['usage'] = traj.usage.model_dump() if 'traj' in dir() else {}
pred.setdefault('group_key', str(sample.metadata.get('test_category')
or sample.metadata.get('category')
or sample.metadata.get('domain')
or sample.metadata.get('id') or sample.id or ''))
u = pred.get('usage') or {}
total_usage = total_usage + Usage(
input_tokens=int(u.get('input_tokens', 0) or 0),
output_tokens=int(u.get('output_tokens', 0) or 0),
total_tokens=int(u.get('total_tokens', 0) or 0),
latency_s=float(u.get('latency_s', 0) or 0))
done_count += 1
_progress(progress, done_count, len(samples), t0, total_usage)
return pred
@ -143,7 +164,48 @@ async def generate_predictions(
return {'raw': text, 'usage': out.usage.model_dump()}
work = _apply_limits(samples, limit, limit_per_task)
preds = list(await asyncio.gather(*(one(s) for s in work)))
# checkpointing: restore completed samples, generate only the rest
ckpt_store = None
if checkpoint:
from ..eval.checkpoint import CheckpointStore, checkpoint_path
ckpt = checkpoint if isinstance(checkpoint, str) else None
if ckpt is None:
from ..eval.checkpoint import checkpoint_path as _cp
ckpt = _cp('/data/evalharness' if False else os.path.expanduser('~/.cache/evalharness'),
getattr(spec, 'name', 'adhoc') if spec is not None else 'adhoc',
model_spec)
ckpt_store = CheckpointStore(ckpt, model=adapter.model or str(adapter))
restored = ckpt_store.load()
else:
restored = {}
keys = []
pending = []
preds_by_key: Dict[str, Dict[str, Any]] = {}
for i, s in enumerate(work):
k = CheckpointStore.key_for(s, i) if ckpt_store else str(i)
keys.append(k)
if k in restored:
preds_by_key[k] = restored[k]
else:
pending.append((i, s))
if ckpt_store is not None and restored:
print(f'checkpoint: restored {len(restored)} predictions '
f'({len(pending)} to generate) -> {ckpt_store.path}', flush=True)
async def run_one(i_s):
i, s = i_s
pred = await one(s)
if ckpt_store is not None:
ckpt_store.append(keys[i], pred)
return i, pred
fresh = await asyncio.gather(*(run_one((i, s)) for i, s in pending))
for i, pred in fresh:
preds_by_key[keys[i]] = pred
preds = [preds_by_key[k] for k in keys]
usages = [p.get('usage', {}) for p in preds]
return preds, usages, total_usage
@ -191,6 +253,7 @@ async def run_eval(
max_turns: int = 8,
max_input_chars: int = 0,
limit_per_task: Optional[int] = None,
checkpoint: Union[bool, str] = False,
few_shot_num: int = -1,
prompt_style: str = 'strict_letter',
) -> EvalReport:
@ -229,11 +292,15 @@ async def run_eval(
env_factory = None
if env:
from ..agent.loop import ENV_REGISTRY, get_env
from ..agent import ENV_REGISTRY, get_env
if env not in ENV_REGISTRY:
raise KeyError(f'unknown env {env!r}; available: {", ".join(ENV_REGISTRY.names())}')
env_factory = lambda: get_env(env) # noqa: E731
probe = get_env(env)
if getattr(probe, 'needs_adapter', False):
env_factory = lambda: get_env(env, adapter=adapter) # noqa: E731
else:
env_factory = lambda: get_env(env) # noqa: E731
# paper-faithful few-shot exemplars: official hand-written hooks first
# (bbh CoT), else the dataset's own dev/train split
@ -268,6 +335,7 @@ async def run_eval(
gen_kwargs=gen_kwargs, env_factory=env_factory,
system=system, max_turns=max_turns, max_input_chars=max_input_chars,
limit_per_task=limit_per_task,
checkpoint=checkpoint,
few_shot_num=few_shot_num,
few_shot_samples=few_shot_samples, few_shot_text=few_shot_text,
prompt_style=prompt_style)

View File

@ -25,6 +25,18 @@ def text_table(target: Union[EvalReport, List[EvalReport]], opts: Dict) -> str:
out.append('=' * max(len(head), 40))
out.append(head)
out.append('=' * max(len(head), 40))
# run stats: duration, tokens, cost (from run_info + usage aggregates)
info = rep.metric_groups.get('run_info', {}) or {}
total_tokens = info.get('gen_total_tokens', 0)
tok_s = ''
if total_tokens:
tok_s = f' tokens={total_tokens}'
secs = 0.0
for s in rep.samples:
secs += float((s.usage or {}).get('latency_s', 0) or 0)
dur = f' time={secs / 3600:.2f}h' if secs >= 3600 else (f' time={secs:.0f}s' if secs else '')
if dur or tok_s:
out.append(f'n_samples={rep.num_samples}{dur}{tok_s}')
if rep.num_failed_extractions:
warn = (f'!! {rep.num_failed_extractions}/{rep.num_samples} extractions failed '
f'({_pct(rep.metrics.get("extraction_failure_rate", 0))}) -- check recipe/model fit')