154 lines
6.0 KiB
Python
154 lines
6.0 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
|
|
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'
|
|
# 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)
|
|
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}
|