Adding in DB, Action, Communicate Checks for Tau (#80)

* working

* changing tests

* updating llm usage

* bug with accessing msg

* temp

* not finished yet

* adding tau2 checks

* removing erroneous tau2 subfolder

* remove workflows folder

* fix test

Author: xzrderek

eval_protocol/mcp/execution/base_policy.py

@@ -34,6 +34,7 @@ def __init__(
         temperature: float = 0.2,
         max_tokens: int = 4096,
         max_tools_per_turn: Optional[int] = None,
+        base_url: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -53,6 +54,7 @@ def __init__(
         self.temperature = temperature
         self.max_tokens = max_tokens
         self.max_tools_per_turn = max_tools_per_turn
+        self.base_url = base_url
 
         # Initialize conversation state tracking for proper OpenAI trajectories
         self.initialized = False

eval_protocol/mcp/execution/manager.py

@@ -253,7 +253,7 @@ async def _execute_rollout(
             user_simulator = UserSimulator(
                 instructions=dataset_row.user_simulation.get("system_prompt"),
                 llm=dataset_row.user_simulation.get("llm", "gpt-4.1"),
-                llm_args=dataset_row.user_simulation.get("llm_args", {"temperature": 0.7}),
+                llm_args=dataset_row.user_simulation.get("llm_args", {"temperature": 0.0}),
             )
 
             # Get initial messages in tau2-bench format for user simulator

eval_protocol/mcp/execution/policy.py

@@ -36,6 +36,7 @@ def __init__(
         temperature: float = 0.2,
         max_tokens: int = 4096,
         max_tools_per_turn: Optional[int] = None,
+        base_url: Optional[str] = None,
         # LiteLLM-specific parameters
         use_caching: bool = True,
         cache_type: Literal["memory", "redis", "dual", "s3", "disk"] = "memory",
@@ -58,7 +59,7 @@ def __init__(
             num_retries: Number of retries for failed requests
             retry_strategy: Retry strategy (literal: "exponential_backoff_retry", "constant_retry")
         """
-        super().__init__(model_id, temperature, max_tokens, max_tools_per_turn, **kwargs)
+        super().__init__(model_id, temperature, max_tokens, max_tools_per_turn, base_url, **kwargs)
 
         self.num_retries = num_retries
         self.retry_strategy = retry_strategy
@@ -162,6 +163,7 @@ async def _make_llm_call(self, messages: List[Dict], tools: List[Dict]) -> Dict:
             "caching": True,
             "num_retries": self.num_retries,
             "retry_strategy": self.retry_strategy,
+            "base_url": self.base_url,
         }
 
         # Add tools if provided
@@ -266,10 +268,19 @@ def __init__(self, model_id: str, **kwargs):
         super().__init__(model_id=f"fireworks_ai/{model_id}", **kwargs)
 
 
+class LocalPolicy(LiteLLMPolicy):
+    """Local policy using LiteLLM for local model endpoints."""
+
+    def __init__(self, model_id: str, base_url: str, **kwargs):
+        """Initialize LocalPolicy for local model endpoints."""
+        super().__init__(model_id=model_id, base_url=base_url, **kwargs)
+
+
 # Export the policies
 __all__ = [
     "LiteLLMPolicy",
     "OpenAIPolicy",
     "AnthropicPolicy",
     "FireworksPolicy",
+    "LocalPolicy",
 ]

examples/tau2_mcp/tests/datasets/airline.json

@@ -6,11 +6,15 @@
     "user_simulation": {
       "enabled": true,
       "llm": "gpt-4.1",
-      "system_prompt": "Domain: airline\nReason for call: You want to cancel reservation EHGLP3.\nKnown info: You are Emma Kim (user id emma_kim_9957).\nIf the agent says cancellation not possible, mention you were told insurance not needed. You don't want to cancel if no refund."
+      "system_prompt": "Instructions:\n\tDomain: airline\nReason for call: You want to cancel reservation EHGLP3.\nKnown info: You are Emma Kim (user id emma_kim_9957).\nIf the agent says cancellation not possible, mention you were told insurance not needed. You don't want to cancel if no refund."
     },
-    "assertions": [
-      "Agent should refuse to proceed with the cancellation."
-    ]
+    "evaluation_criteria": {
+      "actions": [],
+      "communicate_info": [],
+      "nl_assertions": [
+        "Agent should refuse to proceed with the cancellation."
+      ]
+    }
   },
   {
     "id": "airline_task_2",
@@ -19,11 +23,32 @@
     "user_simulation": {
       "enabled": true,
       "llm": "gpt-4.1",
-      "system_prompt": "Domain: airline\nReason for call: You recently spoke on the phone with a customer support representative that told you that a service agent will be able to help you cancel your reservation.\nKnown info: You are Raj Sanchez. Your user id is raj_sanchez_7340.\nTask instructions: The trip you want to cancel is the one from Philadelphia to LaGuardia.\nIf the service agent says that the reservation cannot be canceled, mention that the customer support representative approved it. You don't want to go ahead with the cancellation if you don't get a refund."
+      "system_prompt": "Instructions:\n\tDomain: airline\nReason for call: You recently spoke on the phone with a customer support representative that told you that a service agent will be able to help you cancel your reservation.\nKnown info: You are Raj Sanchez. Your user id is raj_sanchez_7340.\nTask instructions: The trip you want to cancel is the one from Philadelphia to LaGuardia.\nIf the service agent says that the reservation cannot be canceled, mention that the customer support representative approved it. You don't want to go ahead with the cancellation if you don't get a refund."
     },
-    "assertions": [
-      "Agent should not approve the cancellation."
-    ]
+    "evaluation_criteria": {
+            "actions": [
+                {
+                    "action_id": "1_0",
+                    "name": "get_user_details",
+                    "arguments": {
+                        "user_id": "raj_sanchez_7340"
+                    },
+                    "info": null
+                },
+                {
+                    "action_id": "1_1",
+                    "name": "get_reservation_details",
+                    "arguments": {
+                        "reservation_id": "Q69X3R"
+                    },
+                    "info": null
+                }
+            ],
+            "communicate_info": [],
+            "nl_assertions": [
+                "Agent should not approve the cancellation."
+            ]
+        }
   },
   {
     "id": "airline_task_3",
@@ -32,13 +57,51 @@
     "user_simulation": {
       "enabled": true,
       "llm": "gpt-4.1",
-      "system_prompt": "Domain: airline\nReason for call: First, try to book a flight from sf to ny. You will have 3 passengers. Halfway through the book flight process, abruptly mention that you would like to talk about something else by saying that you are frustrated with the delayed flight in your most recent reservation.\nKnown info: You are Noah Muller. Your user id is noah_muller_9847."
+      "system_prompt": "Instructions:\n\tDomain: airline\nReason for call: First, try to book a flight from sf to ny. You will have 3 passengers. Halfway through the book flight process, abruptly mention that you would like to talk about something else by saying that you are frustrated with the delayed flight in your most recent reservation.\nKnown info: You are Noah Muller. Your user id is noah_muller_9847."
     },
-    "assertions": [
-      "Agent should not offer compensation unless the user asks for it.",
-      "Agent should check that the flight was indeed delayed.",
-      "Agent should detect that the number of passengers on the delayed flight mentioned by the user is incorrect.",
-      "Agent should offer a certificate of $50."
-    ]
+    "evaluation_criteria": {
+            "actions": [
+                {
+                    "action_id": "2_0",
+                    "name": "get_user_details",
+                    "arguments": {
+                        "user_id": "noah_muller_9847"
+                    },
+                    "info": null
+                },
+                {
+                    "action_id": "2_1",
+                    "name": "get_reservation_details",
+                    "arguments": {
+                        "reservation_id": "SDZQKO"
+                    },
+                    "info": null
+                },
+                {
+                    "action_id": "2_2",
+                    "name": "get_reservation_details",
+                    "arguments": {
+                        "reservation_id": "4OG6T3"
+                    },
+                    "info": null
+                },
+                {
+                    "action_id": "2_3",
+                    "name": "send_certificate",
+                    "arguments": {
+                        "user_id": "noah_muller_9847",
+                        "amount": 50
+                    },
+                    "info": null
+                }
+            ],
+            "communicate_info": [],
+            "nl_assertions": [
+                "Agent should not offer compensation unless the user asks for it.",
+                "Agent should check that the flight was indeed delayed.",
+                "Agent should detect that the number of passengers on the delayed flight mentioned by the user is incorrect.",
+                "Agent should offer a certificate of $50."
+            ]
+        }
   }
 ]