15b1574857
Function-level silencers kept losing the race to whichever 'import tau2.*' fired first in the process. The filter + TAU2_DATA_DIR now install when the env registry loads the plugin, strictly before any tau2 import can happen anywhere. Co-Authored-By: Claude <noreply@anthropic.com>
272 lines
11 KiB
Python
272 lines
11 KiB
Python
"""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
|
||
import os
|
||
import sys
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from ...data.sample import ChatMessage, Sample
|
||
from ..loop import Environment, register_env
|
||
|
||
# ---- module-level, order-independent tau2 silencing --------------------
|
||
# Every function-level silencer proved fragile to import order (banners
|
||
# fire at 'import tau2.*' wherever it first happens in the process).
|
||
# This block runs when the ENV_REGISTRY loads this plugin -- strictly
|
||
# before any tau2 import anywhere -- and is therefore the one that counts.
|
||
import os as _os
|
||
import sys as _sys
|
||
|
||
_os.environ.setdefault('TAU2_DATA_DIR',
|
||
'/data1/sora/evalscope/tools/tau2-bench/data')
|
||
try:
|
||
from loguru import logger as _lg
|
||
|
||
_lg.remove()
|
||
_lg.add(_sys.stderr, level='ERROR')
|
||
except Exception:
|
||
pass
|
||
# ------------------------------------------------------------------------
|
||
|
||
_PATCHED = False
|
||
|
||
|
||
_TAU2_GIT = 'git+https://github.com/sierra-research/tau2-bench'
|
||
# 本机常见的 tau2-bench checkout 位置（es 仓库 tools/ 下），按序探测
|
||
_TAU2_LOCAL_CANDIDATES = [
|
||
'/data1/sora/evalscope/tools/tau2-bench',
|
||
os.path.expanduser('~/tau2-bench',
|
||
) if hasattr(os.path, 'expanduser') else '',
|
||
]
|
||
|
||
|
||
def _silence_tau2_logs():
|
||
"""tau2's loguru chatter: DEBUG dumps every full message, WARNING-level
|
||
banners fire at import (.env note, data-dir note, litellm cost-map
|
||
timeout). ERROR+ only -- the terminal stays ours."""
|
||
try:
|
||
# kill the data-dir banner at the SOURCE (right directory set)
|
||
os.environ.setdefault(
|
||
'TAU2_DATA_DIR',
|
||
'/data1/sora/evalscope/tools/tau2-bench/data')
|
||
from loguru import logger
|
||
|
||
logger.remove()
|
||
logger.add(sys.stderr, level='ERROR')
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _ensure_tau2_engine():
|
||
# SILENCE FIRST: tau2's modules emit loguru INFO at import time
|
||
# ('Registering default components...', data-dir banner); any tau2
|
||
# import below triggers them before later silencing can help
|
||
_silence_tau2_logs()
|
||
"""Auto-install the REAL tau2-bench engine when missing.
|
||
|
||
PyPI's 'tau2' is an unrelated physics package -- never 'pip install
|
||
tau2'. We try: (1) already importable, (2) a local checkout
|
||
(editable), (3) the GitHub URL, then print the manual fallback."""
|
||
try:
|
||
import tau2.data_model.message # noqa: F401
|
||
return
|
||
except ImportError:
|
||
pass
|
||
import importlib
|
||
import subprocess
|
||
import sys
|
||
|
||
cands = [p for p in _TAU2_LOCAL_CANDIDATES if p and os.path.isdir(p)]
|
||
attempts = [('local checkout', ['-e', p]) for p in cands] + \
|
||
[('github', [_TAU2_GIT])]
|
||
for src, arg in attempts:
|
||
print(f'· tau2 engine missing -- installing from {src} ...', flush=True)
|
||
r = subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', *arg],
|
||
capture_output=True, text=True, timeout=600)
|
||
if r.returncode == 0:
|
||
try:
|
||
importlib.invalidate_caches()
|
||
import tau2.data_model.message # noqa: F401
|
||
print(f'✓ tau2 engine installed from {src}', flush=True)
|
||
return
|
||
except ImportError:
|
||
continue
|
||
raise RuntimeError(
|
||
'tau2-bench 引擎不可用且自动安装失败。手动安装（任选一）：\n'
|
||
f' A. pip install git+{_TAU2_GIT}\n'
|
||
' B. 本机源码: pip install -e /path/to/tau2-bench\n'
|
||
' C. 内网无出口时：在有网的机器 '
|
||
'`pip download tau2-bench --no-deps -d pkg/`？注意 PyPI 的 tau2 是'
|
||
'无关物理库，必须用 GitHub 源或本机源码！')
|
||
|
||
|
||
def _patch_tau2_generate(adapter, user_adapter=None, gen_kwargs=None) -> None:
|
||
"""user_adapter: separate model for the USER simulator (es production
|
||
parity: strong user model like DeepSeek while the agent under test stays
|
||
Qwen). None = single-model setup (agent doubles as user)."""
|
||
"""Route tau2's LLM calls to our ModelAdapter (sync bridge via asyncio)."""
|
||
global _PATCHED
|
||
if _PATCHED:
|
||
return
|
||
|
||
_ensure_tau2_engine()
|
||
_silence_tau2_logs()
|
||
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' — route the user simulator to its own
|
||
# adapter when one is configured (strong-user parity mode)
|
||
target = user_adapter if (user_adapter is not None and str(model) in ('user', 'llm_user')) else adapter
|
||
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():
|
||
# es parity: agent/user gen params (temp 0, max_tokens 16k) --
|
||
# without these the adapter defaults (4096 tokens, server temp)
|
||
# truncate long action sequences and add sampling noise
|
||
return await target.generate(msgs, tools=tool_specs, **(gen_kwargs or {}))
|
||
|
||
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
|
||
|
||
# strip thinking leakage: hybrid backends inline reasoning in
|
||
# content as '...scenario reasoning...</think>visible reply' -- the
|
||
# USER simulator's leak hands the agent the scenario's secret
|
||
# instructions (task goal, what to disclose), inflating rewards
|
||
text = out.text or ''
|
||
if '</think>' in text:
|
||
text = text.rsplit('</think>', 1)[-1].strip()
|
||
|
||
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=text,
|
||
tool_calls=calls or None,
|
||
is_final_chunk=not calls, 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, user_adapter=None, gen_kwargs=None, **kw):
|
||
# the runner's generic default (8 turns) is for single-tool benches;
|
||
# a tau2 phone call needs greet -> verify -> find -> policy -> act ->
|
||
# confirm (official runs use 40+). Floor it, never cap below it.
|
||
max_turns = max(int(max_turns or 0), 40)
|
||
from tau2.data_model.tasks import Task
|
||
from tau2.run import run_task
|
||
|
||
_patch_tau2_generate(adapter, user_adapter, gen_kwargs)
|
||
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'
|
||
# the official engine is SYNCHRONOUS and calls our adapter back via a
|
||
# private event loop in a worker thread; run it off the main loop so it
|
||
# never blocks the runner's other benches
|
||
import concurrent.futures
|
||
|
||
loop = asyncio.get_running_loop()
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
||
res = await loop.run_in_executor(
|
||
pool, lambda: 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)
|
||
# tau2 official reward_info fields: reward (composite), db_check,
|
||
# action_checks, ... -- 'reward' is THE score es reports too
|
||
r = rewards.get('reward')
|
||
if r is None:
|
||
vals = [v for v in (rewards.get('environment_reward'),
|
||
rewards.get('communication_reward'))
|
||
if isinstance(v, (int, float))]
|
||
r = float(sum(vals) / len(vals)) if vals else 0.0
|
||
self.reward_info = {'tau2_rewards': rewards,
|
||
'reward': float(r)}
|
||
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}
|