diff --git a/python/packages/lab/tau2/tests/purge_data.py b/python/packages/lab/tau2/tests/purge_data.py deleted file mode 100644 index 612caabf07..0000000000 --- a/python/packages/lab/tau2/tests/purge_data.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import shutil -from pathlib import Path - - -def purge_tau2_data(): - """Purge tau2 data directory if it exists.""" - - data_dir = Path.cwd() / "data" - - if data_dir.exists(): - shutil.rmtree(data_dir) - print(f"Data directory at {data_dir} has been purged.") - else: - print("Data directory not found. Skipping purge.") - - -if __name__ == "__main__": - purge_tau2_data() diff --git a/python/packages/lab/tau2/tests/setup_data.py b/python/packages/lab/tau2/tests/setup_data.py deleted file mode 100644 index 70be4c0aa1..0000000000 --- a/python/packages/lab/tau2/tests/setup_data.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import shutil -import subprocess -from pathlib import Path - - -def setup_tau2_data(): - """Set up tau2 data directory by cloning repository if needed.""" - - # Get project directory (parent of tests directory) - data_dir = Path.cwd() / "data" - print(data_dir) - - print("Setting up tau2 data directory...") - - # Check if data directory already exists - if data_dir.exists(): - print(f"Data directory already exists at {data_dir}") - else: - print("Data directory not found. Cloning tau2-bench repository...") - - try: - # Clone the repository - print("Cloning https://github.com/sierra-research/tau2-bench.git...") - subprocess.run( - ["git", "clone", "https://github.com/sierra-research/tau2-bench.git"], - check=True, - capture_output=True, - text=True, - ) - - # Move data directory - print("Moving data directory...") - tau2_bench_dir = Path.cwd() / "tau2-bench" - tau2_data_dir = tau2_bench_dir / "data" - - if tau2_data_dir.exists(): - shutil.move(str(tau2_data_dir), str(data_dir)) - else: - raise FileNotFoundError(f"Data directory not found in cloned repository: {tau2_data_dir}") - - # Clean up cloned repository - print("Cleaning up cloned repository...") - shutil.rmtree(tau2_bench_dir) - - print("Data directory setup completed successfully!") - - except subprocess.CalledProcessError as e: - print(f"ERROR: Failed to clone repository: {e}") - raise - except Exception as e: - print(f"ERROR: Failed to set up data directory: {e}") - raise - - print(f"TAU2_DATA_DIR should be set to: {data_dir}") - - return str(data_dir) - - -if __name__ == "__main__": - setup_tau2_data() diff --git a/python/packages/lab/tau2/tests/test_tau2_utils.py b/python/packages/lab/tau2/tests/test_tau2_utils.py index e0c856d5a2..21e8622da8 100644 --- a/python/packages/lab/tau2/tests/test_tau2_utils.py +++ b/python/packages/lab/tau2/tests/test_tau2_utils.py @@ -2,6 +2,9 @@ """Tests for tau2 utils module.""" +import urllib.request +from pathlib import Path + import pytest from agent_framework._tools import AIFunction from agent_framework._types import ChatMessage, FunctionCallContent, FunctionResultContent, Role, TextContent @@ -10,27 +13,48 @@ from agent_framework_lab_tau2._tau2_utils import ( convert_tau2_tool_to_ai_function, ) from tau2.data_model.message import AssistantMessage, SystemMessage, ToolCall, ToolMessage, UserMessage - -# Try to import get_environment and handle missing data files -try: - from tau2.domains.airline.environment import get_environment - - # Try to initialize the environment to check if data files are available - try: - get_environment() - TAU2_DATA_AVAILABLE = True - except FileNotFoundError: - TAU2_DATA_AVAILABLE = False -except ImportError: - TAU2_DATA_AVAILABLE = False +from tau2.domains.airline.data_model import FlightDB +from tau2.domains.airline.tools import AirlineTools +from tau2.environment.environment import Environment -@pytest.mark.skipif(not TAU2_DATA_AVAILABLE, reason="tau2 data files not available") -def test_convert_tau2_tool_to_ai_function_basic(): +@pytest.fixture(scope="session") +def tau2_airline_environment() -> Environment: + airline_db_remote_path = "https://raw.githubusercontent.com/sierra-research/tau2-bench/5ba9e3e56db57c5e4114bf7f901291f09b2c5619/data/tau2/domains/airline/db.json" + airline_policy_remote_path = "https://raw.githubusercontent.com/sierra-research/tau2-bench/5ba9e3e56db57c5e4114bf7f901291f09b2c5619/data/tau2/domains/airline/policy.md" + + # Create cache directory + cache_dir = Path(__file__).parent / "data" + cache_dir.mkdir(exist_ok=True) + + # Define cache file paths + db_cache_path = cache_dir / "airline_db.json" + policy_cache_path = cache_dir / "airline_policy.md" + + # Download files only if they don't exist in cache + if not db_cache_path.exists(): + urllib.request.urlretrieve(airline_db_remote_path, db_cache_path) + + if not policy_cache_path.exists(): + urllib.request.urlretrieve(airline_policy_remote_path, policy_cache_path) + + # Load data from cached files + db = FlightDB.load(str(db_cache_path)) + tools = AirlineTools(db) + with open(policy_cache_path) as fp: + policy = fp.read() + + yield Environment( + domain_name="airline", + policy=policy, + tools=tools, + ) + + +def test_convert_tau2_tool_to_ai_function_basic(tau2_airline_environment): """Test basic conversion from tau2 tool to AIFunction.""" # Get real tools from tau2 environment - env = get_environment() - tools = env.get_tools() + tools = tau2_airline_environment.get_tools() # Use the first available tool for testing assert len(tools) > 0, "No tools available in environment" @@ -49,12 +73,10 @@ def test_convert_tau2_tool_to_ai_function_basic(): assert callable(ai_function.func) -@pytest.mark.skipif(not TAU2_DATA_AVAILABLE, reason="tau2 data files not available") -def test_convert_tau2_tool_to_ai_function_multiple_tools(): +def test_convert_tau2_tool_to_ai_function_multiple_tools(tau2_airline_environment): """Test conversion with multiple tau2 tools.""" # Get real tools from tau2 environment - env = get_environment() - tools = env.get_tools() + tools = tau2_airline_environment.get_tools() # Convert multiple tools ai_functions = [convert_tau2_tool_to_ai_function(tool) for tool in tools[:3]] # Test first 3 tools