Refactor Agent class to accept EvaluationRow directly in the constructor, simplifying initialization. Update _call_model method to return a Message object with structured response data. Modify default_agent_rollout_processor to utilize the new constructor and streamline dataset population. Enhance evaluation_test to initialize eval_metadata for each row before running rollouts, ensuring consistent metadata handling.

Author: Dylan Huang

eval_protocol/pytest/default_agent_rollout_processor.py

@@ -20,9 +20,9 @@ class Agent:
     A really simple agent that calls the model until no more tool calls are needed.
     """
 
-    def __init__(self, model: str, initial_messages: list[Message], config_path: str):
+    def __init__(self, model: str, row: EvaluationRow, config_path: str):
         self.model = model
-        self.evaluation_row: EvaluationRow = EvaluationRow(messages=initial_messages)
+        self.evaluation_row: EvaluationRow = row
         self._policy = LiteLLMPolicy(model_id=model)
         self.mcp_client = MCPMultiClient(config_path=config_path) if config_path else None
         self.tools: Union[List[ChatCompletionToolParam], NotGiven] = NOT_GIVEN
@@ -82,13 +82,15 @@ async def call_agent(self) -> str:
             return await self.call_agent()
         return message["content"]
 
-    async def _call_model(
-        self, messages: list[Message], tools: Optional[list[ChatCompletionToolParam]]
-    ) -> ChatCompletionMessage:
+    async def _call_model(self, messages: list[Message], tools: Optional[list[ChatCompletionToolParam]]) -> Message:
         messages = [message.model_dump() if hasattr(message, "model_dump") else message for message in messages]
         tools = [{"function": tool["function"].model_dump(), "type": "function"} for tool in tools] if tools else []
         response = await self._policy._make_llm_call(messages=messages, tools=tools)
-        return response["choices"][0]["message"]
+        return Message(
+            role=response["choices"][0]["message"]["role"],
+            content=response["choices"][0]["message"]["content"],
+            tool_calls=response["choices"][0]["message"]["tool_calls"],
+        )
 
     async def _execute_tool_call(
         self, tool_call_id: str, tool_name: str, tool_args_dict: dict
@@ -114,12 +116,10 @@ async def default_agent_rollout_processor(
 ) -> List[EvaluationRow]:
     dataset: Dataset = []
     for row in rows:
-        agent = Agent(model=config.model, initial_messages=row.messages, config_path=config.mcp_config_path)
+        agent = Agent(model=config.model, row=row, config_path=config.mcp_config_path)
         await agent.setup()
         await agent.call_agent()
-        dataset.append(
-            EvaluationRow(ground_truth=row.ground_truth, **agent.evaluation_row.model_dump(exclude_none=True))
-        )
+        dataset.append(agent.evaluation_row)
         if agent.mcp_client:
             await agent.mcp_client.cleanup()
     return dataset

eval_protocol/pytest/evaluation_test.py

@@ -206,34 +206,49 @@ def wrapper_body(**kwargs):
                 else:
                     raise ValueError("No input dataset or input messages provided")
 
-                input_dataset: List[EvaluationRow] = []
                 input_params = kwargs.get("input_params") or {}
-                config = RolloutProcessorConfig(
-                    model=model_name,
-                    input_params=input_params,
-                    mcp_config_path=mcp_config_path or "",
-                    max_concurrent_rollouts=max_concurrent_rollouts,
-                    server_script_path=server_script_path,
-                    steps=steps,
+
+                # Create eval metadata with test function info and current commit hash
+                eval_metadata = EvalMetadata(
+                    name=test_func.__name__,
+                    description=test_func.__doc__,
+                    version=versioneer.get_version(),
+                    status="running",
+                    num_runs=num_runs,
+                    aggregation_method=aggregation_method,
+                    threshold_of_success=threshold_of_success,
+                    passed=None,
                 )
-                input_dataset = execute_function(rollout_processor, rows=data, config=config)
 
-                # Populate completion_params in input_metadata for all rows
+                # Populate completion_params in input_metadata for all rows and initialize eval_metadata BEFORE rollouts
                 completion_params = CompletionParams(
                     model=model_name,
                     temperature=input_params.get("temperature"),
                     max_tokens=input_params.get("max_tokens"),
                     max_tool_calls=input_params.get("max_tool_calls"),
                 )
 
-                for row in input_dataset:
+                for row in data:
                     if row.input_metadata is None:
                         row.input_metadata = InputMetadata()
                     row.input_metadata.completion_params = completion_params
                     # Add mode to session_data
                     if row.input_metadata.session_data is None:
                         row.input_metadata.session_data = {}
                     row.input_metadata.session_data["mode"] = mode
+                    # Initialize eval_metadata for each row
+                    row.eval_metadata = eval_metadata
+
+                # Now run the rollout processor with metadata-initialized data
+                config = RolloutProcessorConfig(
+                    model=model_name,
+                    input_params=input_params,
+                    mcp_config_path=mcp_config_path or "",
+                    max_concurrent_rollouts=max_concurrent_rollouts,
+                    server_script_path=server_script_path,
+                    steps=steps,
+                )
+                input_dataset = execute_function(rollout_processor, rows=data, config=config)
 
                 all_results: List[EvaluationRow] = []
                 for _ in range(num_runs):
@@ -283,21 +298,11 @@ def wrapper_body(**kwargs):
                 if threshold_of_success is not None:
                     passed = agg_score >= threshold_of_success
 
-                # Create eval metadata with test function info and current commit hash
-                eval_metadata = EvalMetadata(
-                    name=test_func.__name__,
-                    description=test_func.__doc__,
-                    version=versioneer.get_version(),
-                    status="finished",
-                    num_runs=num_runs,
-                    aggregation_method=aggregation_method,
-                    threshold_of_success=threshold_of_success,
-                    passed=passed,
-                )
-
-                # Add metadata to all results before logging
+                # Update eval metadata status and passed field for all results
                 for r in all_results:
-                    r.eval_metadata = eval_metadata
+                    if r.eval_metadata is not None:
+                        r.eval_metadata.status = "finished"
+                        r.eval_metadata.passed = passed
                     default_logger.log(r)
 
                 # Check threshold after logging