""" TAU2-bench data loader for tool-calling agent evaluation. This module provides data loading from tau2-bench, a benchmark for evaluating tool-calling agents in customer service domains (airline, retail, telecom). Setup Requirements: 1. Install tau2: pip install ace-framework[tau-bench] 2. Set TAU2_DATA_DIR environment variable to point to tau2 data directory 3. Download data from: https://github.com/sierra-research/tau2-bench """ from __future__ import annotations import logging import os from typing import Any, Dict, Iterator, List from ..base import DataLoader logger = logging.getLogger(__name__) class Tau2Loader(DataLoader): """ Data loader for TAU2-bench (τ²-bench) tasks. TAU2-bench evaluates tool-calling agents in customer service domains: - airline: Flight bookings, cancellations, seat changes - retail: Order management, returns, product inquiries - telecom: Account management, plan changes, billing Example: >>> loader = Tau2Loader() >>> for task in loader.load(domain="airline", task_split="base", limit=10): ... print(task["task_id"], task["instruction"]) Setup: 1. Install: pip install ace-framework[tau-bench] 2. Clone data: git clone https://github.com/sierra-research/tau2-bench 3. Set environment: export TAU2_DATA_DIR=/path/to/tau2-bench/data """ def supports_source(self, source: str) -> bool: """Check if this loader supports the given data source.""" return source == "tau2" def load(self, **kwargs) -> Iterator[Dict[str, Any]]: """ Load TAU2-bench tasks for a specific domain. Args: domain: Domain to load tasks from (airline, retail, telecom) task_split: Task split to use (base, human, gpt4o) - for airline/retail limit: Maximum number of tasks to load **kwargs: Additional arguments (unused) Yields: Dict containing task data: - task_id: Unique task identifier - instruction: Initial user instruction - tools: List of available tool definitions - user_llm: LLM model for user simulation - domain: Domain name - task_split: Split name - metadata: Additional task metadata Raises: ImportError: If tau2 is not installed ValueError: If data directory not configured or tasks cannot be loaded """ try: from tau2.registry import registry except ImportError: raise ImportError( "tau2 is required for TAU2 loader. " "Install with: pip install ace-framework[tau-bench]" ) # Check if data directory is configured data_dir = os.environ.get("TAU2_DATA_DIR") if not data_dir: raise ValueError( "TAU2_DATA_DIR environment variable not set. " "Please set it to point to the tau2-bench data directory. " "Clone data from: https://github.com/sierra-research/tau2-bench" ) domain = kwargs.get("domain", "airline") task_split = kwargs.get("task_split", "base") limit = kwargs.get("limit") user_llm = kwargs.get("user_llm", "gpt-4o-mini") # Get tasks for the domain using the registry try: tasks = self._get_tasks_for_domain(registry, domain, task_split) except FileNotFoundError as e: raise ValueError( f"Failed to load tasks for {domain}/{task_split}. " f"Ensure TAU2_DATA_DIR points to valid tau2 data directory. " f"Error: {e}" ) except Exception as e: raise ValueError(f"Failed to get tasks for {domain}/{task_split}: {e}") if not tasks: logger.warning(f"No tasks found for {domain}/{task_split}") return # Apply limit if specified if limit: tasks = tasks[:limit] # Yield each task for task in tasks: try: task_id = getattr(task, "id", str(id(task))) # Extract instruction from user_scenario instruction = self._extract_instruction(task) # Get tools from the domain environment tools = self._get_domain_tools(registry, domain) yield { "task_id": task_id, "instruction": instruction, "tools": tools, "user_llm": user_llm, "domain": domain, "task_split": task_split, "task": task, # Store the full task object for gym "metadata": { "task_id": task_id, "domain": domain, "task_split": task_split, "max_steps": 30, }, } except Exception as e: logger.warning(f"Failed to process task: {e}") continue def _extract_instruction(self, task) -> str: """Extract instruction text from a tau2 Task object.""" # Try user_scenario.instructions.reason_for_call first if hasattr(task, "user_scenario"): scenario = task.user_scenario if hasattr(scenario, "instructions"): instr = scenario.instructions if hasattr(instr, "reason_for_call") and instr.reason_for_call: return str(instr.reason_for_call) # Fallback to description if hasattr(task, "description") and task.description: return str(task.description) return "" def _get_tasks_for_domain( self, registry, domain: str, task_split: str ) -> List[Any]: """Get tasks for a domain using the appropriate registry method.""" # Get the task loader function for this domain tasks_loader = registry.get_tasks_loader(domain) # Load tasks with optional split if task_split and task_split != "base": # Check if domain supports splits splits_loader = registry.get_task_splits_loader(domain) if splits_loader: splits = splits_loader() if task_split in splits: # Filter tasks by split all_tasks = tasks_loader() split_ids = set(splits[task_split]) return [t for t in all_tasks if t.id in split_ids] # Default: load all tasks for domain return tasks_loader() def _get_domain_tools(self, registry, domain: str) -> List[Dict[str, Any]]: """Get available tools for a domain.""" try: env_constructor = registry.get_env_constructor(domain) env = env_constructor() # get_tools() returns list of Tool objects if hasattr(env, "get_tools"): tools = env.get_tools() if isinstance(tools, list): # Convert Tool objects to dicts return [ { "name": getattr(t, "name", str(t)), "description": getattr(t, "long_desc", ""), } for t in tools ] except Exception as e: logger.debug(f"Could not get tools for {domain}: {e}") return [] def get_domains(self) -> List[str]: """Get list of available domains.""" return ["airline", "retail", "telecom"] def get_task_splits(self) -> List[str]: """Get list of available task splits.""" return ["base", "human", "gpt4o"] def get_task_count(self, domain: str, task_split: str = "base") -> int: """Get number of tasks available for a domain/split combination.""" try: from tau2.registry import registry tasks = self._get_tasks_for_domain(registry, domain, task_split) return len(tasks) except ImportError: return 0 except Exception: return 0 def check_data_available(self) -> bool: """Check if tau2 data is available and configured.""" data_dir = os.environ.get("TAU2_DATA_DIR") if not data_dir: return False return os.path.isdir(data_dir)