fix input_message data type

Author: Dylan Huang

eval_protocol/benchmarks/test_gpqa.py

@@ -2,7 +2,6 @@
 import csv
 import io
 import re
-from typing import List
 
 import requests
 
@@ -20,12 +19,12 @@
 )
 
 
-def _load_gpqa_messages_from_csv() -> List[List[Message]]:
+def _load_gpqa_messages_from_csv() -> list[list[list[Message]]]:
     url = "https://openaipublic.blob.core.windows.net/simple-evals/gpqa_diamond.csv"
     resp = requests.get(url, timeout=60)
     resp.raise_for_status()
 
-    messages_list: List[List[Message]] = []
+    messages_list: list[list[Message]] = []
     reader = csv.DictReader(io.StringIO(resp.text))
     for ex in reader:
         q = str(ex.get("Question", ""))
@@ -45,7 +44,7 @@ def _load_gpqa_messages_from_csv() -> List[List[Message]]:
         )
     if not messages_list:
         raise RuntimeError("Failed to load GPQA messages: no rows found from source")
-    return messages_list
+    return [messages_list]
 
 
 def _extract_abcd_letter(text: str) -> str | None:
@@ -58,7 +57,7 @@ def _extract_abcd_letter(text: str) -> str | None:
 _GPQA_INPUT_MESSAGES = _load_gpqa_messages_from_csv()
 
 
-def _strip_gt_messages(msgs: List[Message]) -> List[Message]:
+def _strip_gt_messages(msgs: list[Message]) -> list[Message]:
     return [m for m in msgs if not (m.role == "system" and (m.content or "").startswith("__GT__:"))]
 
 
@@ -69,9 +68,9 @@ def __init__(self):
         super().__init__()
         self.single_turn_processor = SingleTurnRolloutProcessor()
 
-    def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[asyncio.Task[EvaluationRow]]:
+    def __call__(self, rows: list[EvaluationRow], config: RolloutProcessorConfig) -> list[asyncio.Task[EvaluationRow]]:
         """Preprocess rows and delegate to SingleTurnRolloutProcessor."""
-        processed: List[EvaluationRow] = []
+        processed: list[EvaluationRow] = []
 
         for r in rows:
             gt_tokens = [

eval_protocol/benchmarks/test_livebench_data_analysis.py

@@ -409,7 +409,7 @@ def _extract_gt(row: EvaluationRow) -> Dict[str, Any]:
 
 @evaluation_test(
     completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/gpt-oss-120b"}],
-    input_messages=[[m for m in r.messages] for r in _CTA_ROWS],
+    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",
@@ -451,7 +451,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",

eval_protocol/pytest/evaluation_test.py

@@ -62,7 +62,7 @@
 def evaluation_test(
     *,
     completion_params: Sequence[CompletionParams | None] | None = None,
-    input_messages: Sequence[InputMessagesParam | None] | None = None,
+    input_messages: Sequence[list[InputMessagesParam] | None] | None = None,
     input_dataset: Sequence[DatasetPathParam] | None = None,
     input_rows: Sequence[list[EvaluationRow]] | None = None,
     dataset_adapter: Callable[[list[dict[str, Any]]], Dataset] = default_dataset_adapter,  # pyright: ignore[reportExplicitAny]
@@ -232,7 +232,7 @@ def _log_eval_error(status: Status, rows: list[EvaluationRow] | None, passed: bo
                     elif "input_messages" in kwargs and kwargs["input_messages"] is not None:
                         # Support either a single row (List[Message]) or many rows (List[List[Message]])
                         im = kwargs["input_messages"]
-                        data = [EvaluationRow(messages=im)]
+                        data = [EvaluationRow(messages=dataset_messages) for dataset_messages in im]
                     elif "input_rows" in kwargs and kwargs["input_rows"] is not None:
                         # Use pre-constructed EvaluationRow objects directly
                         data = kwargs["input_rows"]

eval_protocol/pytest/generate_parameter_combinations.py

@@ -18,7 +18,7 @@
 Either a single completion params object or None.
 """
 
-InputMessagesKwarg = InputMessagesParam | None
+InputMessagesKwarg = list[InputMessagesParam] | None
 InputRowsKwarg = Dataset | None
 EvaluationTestKwargs = EvaluationInputParam | None
 
@@ -47,7 +47,7 @@ class ParameterizedTestKwargs(TypedDict):
 def generate_parameter_combinations(
     input_dataset: Sequence[DatasetPathParam] | None,
     completion_params: Sequence[CompletionParams | None],
-    input_messages: Sequence[InputMessagesParam | None] | None,
+    input_messages: Sequence[list[InputMessagesParam] | None] | None,
     input_rows: Sequence[list[EvaluationRow] | None] | None,
     evaluation_test_kwargs: Sequence[EvaluationInputParam | None] | None,
     max_dataset_rows: int | None,
@@ -83,11 +83,15 @@ def generate_parameter_combinations(
     # Apply EP_MAX_DATASET_ROWS to input_messages, but do NOT parameterize over
     # each row. Instead, pass the entire sliced list through in a single test run
     # so summaries aggregate all rows together (AIME-style behavior).
-    messages: Sequence[InputMessagesParam | None] = [None]
+    messages: Sequence[list[InputMessagesParam] | None] = [None]
     if input_messages is not None:
         effective_max_rows = parse_ep_max_rows(max_dataset_rows)
         if effective_max_rows is not None:
-            sliced_messages: Sequence[InputMessagesParam | None] = input_messages[:effective_max_rows]
+            sliced_messages: Sequence[list[InputMessagesParam] | None] = [
+                dataset_messages[:effective_max_rows]
+                for dataset_messages in input_messages
+                if dataset_messages is not None
+            ]
         else:
             sliced_messages = input_messages
         # Wrap as a single parameter payload

eval_protocol/pytest/parameterize.py

@@ -18,7 +18,7 @@ def pytest_parametrize(
     combinations: list[CombinationTuple],
     input_dataset: Sequence[DatasetPathParam] | None,
     completion_params: Sequence[CompletionParams | None] | None,
-    input_messages: Sequence[InputMessagesParam | None] | None,
+    input_messages: Sequence[list[InputMessagesParam] | None] | None,
     input_rows: Sequence[list[EvaluationRow]] | None,
     evaluation_test_kwargs: Sequence[EvaluationInputParam | None] | None,
 ) -> PytestParametrizeArgs:

eval_protocol/pytest/rollout_processor.py

@@ -1,6 +1,5 @@
 import asyncio
 from abc import ABC, abstractmethod
-from typing import List
 
 from eval_protocol.models import EvaluationRow
 from eval_protocol.pytest.types import RolloutProcessorConfig
@@ -12,7 +11,7 @@ class RolloutProcessor(ABC):
     """
 
     @abstractmethod
-    def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[asyncio.Task[EvaluationRow]]:
+    def __call__(self, rows: list[EvaluationRow], config: RolloutProcessorConfig) -> list[asyncio.Task[EvaluationRow]]:
         """Process evaluation rows and return async tasks. Must be implemented by subclasses."""
         pass
 

examples/healthbench/tests/test_evaluation.py

@@ -1,6 +1,3 @@
-import json
-from typing import Dict, List
-
 from eval_protocol.models import EvaluateResult, EvaluationRow, Message, MetricResult
 from eval_protocol.pytest.default_single_turn_rollout_process import (
     SingleTurnRolloutProcessor,
@@ -34,13 +31,15 @@
     },
 ]
 
-_HB_INPUT_MESSAGES: List[List[Message]] = []
-_HB_RUBRICS_MAP: Dict[str, List[Dict]] = {}
+_HB_INPUT_MESSAGES: list[list[list[Message]]] = []
+_HB_RUBRICS_MAP: dict[str, list[dict]] = {}
 for s in _HB_SAMPLES:
     _HB_INPUT_MESSAGES.append(
         [
-            Message(role="system", content=SYSTEM_PROMPT),
-            Message(role="user", content=s["prompt_text"]),
+            [
+                Message(role="system", content=SYSTEM_PROMPT),
+                Message(role="user", content=s["prompt_text"]),
+            ]
         ]
     )
     _HB_RUBRICS_MAP[s["prompt_text"]] = s["rubrics"]

tests/chinook/test_pydantic_chinook.py

@@ -21,7 +21,7 @@
 
 @pytest.mark.asyncio
 @evaluation_test(
-    input_messages=[[Message(role="user", content="What is the total number of tracks in the database?")]],
+    input_messages=[[[Message(role="user", content="What is the total number of tracks in the database?")]]],
     completion_params=[
         {
             "model": {

tests/pytest/test_get_metadata.py

@@ -7,11 +7,13 @@
 @evaluation_test(
     input_messages=[
         [
-            Message(role="user", content="What is the capital of France?"),
-        ],
-        [
-            Message(role="user", content="What is the capital of the moon?"),
-        ],
+            [
+                Message(role="user", content="What is the capital of France?"),
+            ],
+            [
+                Message(role="user", content="What is the capital of the moon?"),
+            ],
+        ]
     ],
     completion_params=[{"model": "accounts/fireworks/models/kimi-k2-instruct"}] * 2,
     mode="groupwise",

tests/pytest/test_pydantic_agent.py

@@ -11,7 +11,7 @@
 
 @pytest.mark.asyncio
 @evaluation_test(
-    input_messages=[[Message(role="user", content="Hello, how are you?")]],
+    input_messages=[[[Message(role="user", content="Hello, how are you?")]]],
     completion_params=[
         {"model": "accounts/fireworks/models/gpt-oss-120b", "provider": "fireworks"},
     ],