Python: Fix tau2 tests by downloading data files directly from remote URLs (#960)

* Re-add skipped tests in lab-tau2

* resolve comments

* cache the files

---------

Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>

python/packages/lab/tau2/tests/purge_data.py

@@ -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()

python/packages/lab/tau2/tests/setup_data.py

@@ -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()

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