make eval functions async

PiperOrigin-RevId: 756106627
This commit is contained in:
Xiang (Sean) Zhou
2025-05-07 19:52:47 -07:00
committed by Copybara-Service
parent cc1ef3f2ad
commit e7d9cf359a
11 changed files with 50 additions and 36 deletions

View File

@@ -32,7 +32,7 @@ class EvaluationGenerator:
"""Generates evaluation responses for agents."""
@staticmethod
def generate_responses(
async def generate_responses(
eval_dataset,
agent_module_path,
repeat_num=3,
@@ -107,7 +107,7 @@ class EvaluationGenerator:
)
@staticmethod
def _process_query_with_root_agent(
async def _process_query_with_root_agent(
data,
root_agent,
reset_func,
@@ -128,7 +128,7 @@ class EvaluationGenerator:
all_mock_tools.add(expected[EvalConstants.TOOL_NAME])
eval_data_copy = data.copy()
EvaluationGenerator.apply_before_tool_callback(
await EvaluationGenerator.apply_before_tool_callback(
root_agent,
lambda *args: EvaluationGenerator.before_tool_callback(
*args, eval_dataset=eval_data_copy
@@ -247,7 +247,7 @@ class EvaluationGenerator:
return None
@staticmethod
def apply_before_tool_callback(
async def apply_before_tool_callback(
agent: BaseAgent,
callback: BeforeToolCallback,
all_mock_tools: set[str],
@@ -265,6 +265,6 @@ class EvaluationGenerator:
# Apply recursively to subagents if they exist
for sub_agent in agent.sub_agents:
EvaluationGenerator.apply_before_tool_callback(
await EvaluationGenerator.apply_before_tool_callback(
sub_agent, callback, all_mock_tools
)