fix a few more

Author: Benny Chen

eval_protocol/benchmarks/test_livebench_data_analysis.py

@@ -424,8 +424,8 @@ def _extract_gt(row: EvaluationRow) -> Dict[str, Any]:
 
 @evaluation_test(
     completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/gpt-oss-120b"}],
-    # Provide a flat list per run (Sequence[InputMessagesParam]) to match signature
-    input_messages=[[m for m in r.messages] for r in _CTA_ROWS],
+    # Wrap dataset messages in an extra list to match Sequence[list[InputMessagesParam]]
+    input_messages=[[[m for m in r.messages] for r in _CTA_ROWS]],
     rollout_processor_kwargs={"extra_body": {"reasoning_effort": "low"}},
     rollout_processor=SingleTurnRolloutProcessor(),
     aggregation_method="mean",
@@ -468,7 +468,7 @@ def test_livebench_cta_pointwise(row: EvaluationRow) -> EvaluationRow:
 
 @evaluation_test(
     completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/gpt-oss-120b"}],
-    input_messages=[[m for m in r.messages] for r in _TABLEJOIN_ROWS],
+    input_messages=[[[m for m in r.messages] for r in _TABLEJOIN_ROWS]],
     rollout_processor_kwargs={"extra_body": {"reasoning_effort": "low"}},
     rollout_processor=LiveBenchGroundTruthRolloutProcessor(_TABLEJOIN_ROWS),
     aggregation_method="mean",
@@ -511,7 +511,7 @@ def test_livebench_tablejoin_pointwise(row: EvaluationRow) -> EvaluationRow:
 
 @evaluation_test(
     completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/gpt-oss-120b"}],
-    input_messages=[[m for m in r.messages] for r in _TABLEREFORMAT_ROWS],
+    input_messages=[[[m for m in r.messages] for r in _TABLEREFORMAT_ROWS]],
     rollout_processor_kwargs={"extra_body": {"reasoning_effort": "low"}},
     rollout_processor=LiveBenchGroundTruthRolloutProcessor(_TABLEREFORMAT_ROWS),
     aggregation_method="mean",

eval_protocol/benchmarks/test_tau_bench_airline.py

@@ -147,7 +147,7 @@ def test_tau_bench_airline_evaluation(row: EvaluationRow) -> EvaluationRow:
     messages = row.messages
 
     # Get evaluation criteria and user_simulation from input_metadata.dataset_info
-    dataset_info = row.input_metadata.dataset_info if row.input_metadata else {}
+    dataset_info = (row.input_metadata.dataset_info or {}) if row.input_metadata else {}
     evaluation_criteria = dataset_info.get("evaluation_criteria", {})
 
     nl_assertions = evaluation_criteria.get("nl_assertions", [])

eval_protocol/benchmarks/test_tau_bench_retail.py

@@ -137,7 +137,7 @@ def test_tau_bench_retail_evaluation(row: EvaluationRow) -> EvaluationRow:
     messages = row.messages
 
     # Get evaluation criteria and user_simulation from input_metadata.dataset_info
-    dataset_info = row.input_metadata.dataset_info if row.input_metadata else {}
+    dataset_info = (row.input_metadata.dataset_info or {}) if row.input_metadata else {}
     evaluation_criteria = dataset_info.get("evaluation_criteria", {})
 
     nl_assertions = evaluation_criteria.get("nl_assertions", [])

eval_protocol/execution/pipeline.py

@@ -87,6 +87,7 @@ async def _discover_tools_for_sample(self, sample_id: str, mcp_backend_ref: str)
 
         try:
             backend_requests = [{"backend_name_ref": mcp_backend_ref, "num_instances": 1}]
+            assert self.mcp_intermediary_client is not None
             init_response = await self.mcp_intermediary_client.initialize_session(backend_requests)
 
             if init_response.get("error"):
@@ -109,6 +110,7 @@ async def _discover_tools_for_sample(self, sample_id: str, mcp_backend_ref: str)
                     current_instance_id = inst_info_dict.get("instance_id")
                     if not current_instance_id:
                         continue
+                    assert self.mcp_intermediary_client is not None
                     list_tools_result = await self.mcp_intermediary_client.list_backend_tools(
                         rk_session_id=rk_session_id,
                         instance_id=current_instance_id,
@@ -130,6 +132,7 @@ async def _discover_tools_for_sample(self, sample_id: str, mcp_backend_ref: str)
             if rk_session_id and self.mcp_intermediary_client:
                 logger.info(f"Sample {sample_id}: Cleaning up tool discovery session '{rk_session_id}'.")
                 try:
+                    assert self.mcp_intermediary_client is not None
                     await self.mcp_intermediary_client.cleanup_session(rk_session_id)
                 except Exception as e_cl:
                     logger.error(
@@ -276,6 +279,7 @@ async def _execute_mcp_agent_rollout(
 
         try:
             backend_requests = [{"backend_name_ref": mcp_backend_ref, "num_instances": 1}]
+            assert self.mcp_intermediary_client is not None
             init_response = await self.mcp_intermediary_client.initialize_session(backend_requests)
             if init_response.get("error"):
                 raise RuntimeError(
@@ -331,6 +335,7 @@ async def _execute_mcp_agent_rollout(
                             if not isinstance(tool_args_dict, dict):
                                 raise ValueError("Args not dict")
 
+                            assert self.mcp_intermediary_client is not None
                             exec_result = await self.mcp_intermediary_client.call_backend_tool(
                                 rk_session_id=rk_session_id,
                                 instance_id=primary_instance_id_for_agent_actions,
@@ -405,6 +410,7 @@ async def _execute_mcp_agent_rollout(
             state_capture_tool = self.cfg.agent.get("state_capture_tool")
             if state_capture_tool:
                 state_capture_args = dict(self.cfg.agent.get("state_capture_args", OmegaConf.create({})))
+                assert self.mcp_intermediary_client is not None
                 final_filesystem_state_from_mcp = await self.mcp_intermediary_client.call_backend_tool(
                     rk_session_id=rk_session_id,
                     instance_id=primary_instance_id_for_agent_actions,
@@ -432,6 +438,7 @@ async def _execute_mcp_agent_rollout(
             }
         finally:
             if rk_session_id and self.mcp_intermediary_client:
+                assert self.mcp_intermediary_client is not None
                 await self.mcp_intermediary_client.cleanup_session(rk_session_id)
 
     async def _process_single_sample(

eval_protocol/integrations/braintrust.py

@@ -48,7 +48,7 @@ def scorer(input_val: Any, output: Any, expected: Any) -> float:
         ground_truth = None
         if expected is not None:
             ground_truth = [Message(role="assistant", content=str(expected))]
-        result = reward_fn(messages=messages, ground_truth=ground_truth)
+        result = reward_fn(messages, ground_truth)
         return float(result.score)
 
     return scorer

eval_protocol/mcp/execution/manager.py

@@ -281,8 +281,15 @@ def extract_text_content(msg_dict):
                         # Generate user response using the simulator
                         # Pass the assistant message content to drive the simulated user's next response
                         last_assistant = user_simulator_messages[-1]
+                        # Convert last assistant message into a valid user input message for simulator
+                        from vendor.tau2.data_model.message import UserMessage as TauUserMessage
+
+                        converted_user_prompt = (
+                            last_assistant.content if getattr(last_assistant, "content", None) else ""
+                        )
+                        converted_message = TauUserMessage(role="user", content=converted_user_prompt)
                         user_message, user_simulator_state = await user_simulator.generate_next_message(
-                            last_assistant,
+                            converted_message,
                             user_simulator_state,
                         )
                         user_content = user_message.content if user_message.content else ""

eval_protocol/mcp/simulation_server.py

@@ -288,7 +288,7 @@ def _discover_and_register_resources(self):
         if discovered_resources:
 
             @self.app.read_resource()
-            async def read_resource(uri: str):
+            async def read_resource(uri: AnyUrl):
                 # Get the current request context
                 ctx = self.app.request_context
 

eval_protocol/pytest/default_langchain_rollout_processor.py

@@ -56,9 +56,17 @@ def __init__(self, content: str):
 
             # Resolve the appropriate async invoke function
             if hasattr(target, "graph") and hasattr(target.graph, "ainvoke"):
-                invoke_fn = target.graph.ainvoke
+
+                async def _invoke_graph(payload):
+                    return await target.graph.ainvoke(payload)  # type: ignore[attr-defined]
+
+                invoke_fn = _invoke_graph
             elif hasattr(target, "ainvoke"):
-                invoke_fn = target.ainvoke
+
+                async def _invoke_direct(payload):
+                    return await target.ainvoke(payload)  # type: ignore[attr-defined]
+
+                invoke_fn = _invoke_direct
             elif callable(target):
 
                 async def _invoke_wrapper(payload):

eval_protocol/rewards/function_calling.py

@@ -451,7 +451,7 @@ def schema_jaccard_reward(
         DeprecationWarning,
         stacklevel=2,
     )
-    return exact_tool_match_reward(messages=messages, ground_truth=ground_truth, **kwargs)
+    return exact_tool_match_reward(messages, ground_truth, **kwargs)
 
 
 @reward_function
@@ -493,7 +493,7 @@ def llm_judge_reward(
         DeprecationWarning,
         stacklevel=2,
     )
-    return exact_tool_match_reward(messages=messages, ground_truth=ground_truth, **kwargs)
+    return exact_tool_match_reward(messages, ground_truth, **kwargs)
 
 
 @reward_function
@@ -537,7 +537,7 @@ def composite_function_call_reward(
         DeprecationWarning,
         stacklevel=2,
     )
-    return exact_tool_match_reward(messages=messages, ground_truth=ground_truth, **kwargs)
+    return exact_tool_match_reward(messages, ground_truth, **kwargs)
 
 
 # JSON schema reward functions have been moved to json_schema.py module

eval_protocol/rewards/json_schema.py

@@ -290,8 +290,8 @@ def json_schema_reward_with_llm_judge(
     normalized_weights = {k: v / total_weight for k, v in weights.items()}
 
     schema_result = json_schema_reward(
-        messages=messages,
-        ground_truth=ground_truth,
+        messages,
+        ground_truth,
         json_content=json_content,
         expected_schema=expected_schema,
         **kwargs,