ai-agent-book/chapter2/prompt-engineering/tau_bench/envs/airline/env.py

# Copyright Sierra

from tau_bench.envs.airline.data import load_data
from tau_bench.envs.airline.rules import RULES
from tau_bench.envs.airline.tools import ALL_TOOLS
from tau_bench.envs.airline.wiki import WIKI
from tau_bench.envs.base import Env
from typing import Optional, Union
from tau_bench.envs.user import UserStrategy


class MockAirlineDomainEnv(Env):
    def __init__(
        self,
        user_strategy: Union[str, UserStrategy] = UserStrategy.LLM,
        user_model: str = "gpt-4o",
        user_provider: Optional[str] = None,
        task_split: str = "test",
        task_index: Optional[int] = None,
        user_seed: Optional[int] = None,
    ):
        match task_split:
            case "test":
                from tau_bench.envs.airline.tasks_test import TASKS as tasks
            case _:
                raise ValueError(f"Unknown task split: {task_split}")
        super().__init__(
            data_load_func=load_data,
            tools=ALL_TOOLS,
            tasks=tasks,
            wiki=WIKI,
            rules=RULES,
            user_strategy=user_strategy,
            user_model=user_model,
            user_provider=user_provider,
            task_index=task_index,
            user_seed=user_seed,
        )
        self.terminate_tools = ["transfer_to_human_agents"]