DRY completion params and make it a dict

Author: Dylan Huang

eval_protocol/adapters/huggingface.py

@@ -4,11 +4,11 @@
 to EvaluationRow format for use in evaluation pipelines.
 """
 
-from typing import Any, Dict, Iterator, List, Optional
-from datetime import datetime
 import logging
+from datetime import datetime
+from typing import Any, Dict, Iterator, List, Optional
 
-from eval_protocol.models import EvaluationRow, Message, InputMetadata, CompletionParams
+from eval_protocol.models import EvaluationRow, InputMetadata, Message
 
 logger = logging.getLogger(__name__)
 
@@ -277,20 +277,20 @@ def _create_input_metadata(self, trace: Any, observations: List[Any]) -> InputMe
             InputMetadata object
         """
         # Extract completion parameters from observations
-        completion_params = CompletionParams()
+        completion_params = {}
 
         # Look for model parameters in observations
         for obs in observations:
             if hasattr(obs, "model") and obs.model:
-                completion_params.model = obs.model
+                completion_params["model"] = obs.model
             if hasattr(obs, "model_parameters") and obs.model_parameters:
                 params = obs.model_parameters
                 if "temperature" in params:
-                    completion_params.temperature = params["temperature"]
+                    completion_params["temperature"] = params["temperature"]
                 if "max_tokens" in params:
-                    completion_params.max_tokens = params["max_tokens"]
+                    completion_params["max_tokens"] = params["max_tokens"]
                 if "top_p" in params:
-                    completion_params.top_p = params["top_p"]
+                    completion_params["top_p"] = params["top_p"]
                 break
 
         # Create dataset info from trace metadata

eval_protocol/adapters/langfuse.py

@@ -11,7 +11,7 @@
 from typing import Any, Dict, List
 
 from eval_protocol.benchmarks.registry import export_benchmark
-from eval_protocol.models import CompletionParams, EvaluateResult, EvaluationRow, InputMetadata, Message
+from eval_protocol.models import EvaluateResult, EvaluationRow, InputMetadata, Message
 from eval_protocol.pytest import evaluation_test
 from eval_protocol.pytest.default_mcp_gym_rollout_processor import default_mcp_gym_rollout_processor
 from vendor.tau2.data_model.message import (

eval_protocol/benchmarks/suites/tau_bench_retail.py

@@ -20,7 +20,7 @@
 from vendor.tau2.data_model.message import AssistantMessage, UserMessage
 from vendor.tau2.user.user_simulator import UserSimulator
 
-from ...models import CompletionParams, EvaluationRow, InputMetadata, Message
+from ...models import EvaluationRow, InputMetadata, Message
 from ...types import MCPSession, MCPToolCall, TerminationReason, Trajectory
 
 if TYPE_CHECKING:
@@ -128,12 +128,12 @@ async def _execute_with_semaphore(idx):
                 evaluation_row.messages = messages
                 evaluation_row.tools = shared_tool_schema
                 evaluation_row.usage = CompletionUsage(**trajectory.usage)
-                evaluation_row.input_metadata.completion_params = CompletionParams(
-                    model=policy.model_id,
-                    temperature=getattr(policy, "temperature", None),
-                    max_tokens=getattr(policy, "max_tokens", None),
-                    max_tool_calls=getattr(policy, "max_tools_per_turn", None),
-                )
+                evaluation_row.input_metadata.completion_params = {
+                    "model": policy.model_id,
+                    "temperature": getattr(policy, "temperature", None),
+                    "max_tokens": getattr(policy, "max_tokens", None),
+                    "max_tool_calls": getattr(policy, "max_tools_per_turn", None),
+                }
 
                 if trajectory.terminated:
                     if trajectory.termination_reason == TerminationReason.ERROR:

eval_protocol/mcp/execution/manager.py

@@ -1,6 +1,6 @@
 import os
 from datetime import datetime
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
 
 from openai.types import CompletionUsage
 from openai.types.chat.chat_completion_message import (
@@ -178,16 +178,18 @@ def __iter__(self):
         return iter(self.__fields__.keys())  # Changed to __fields__
 
 
-class CompletionParams(BaseModel):
-    """Configuration for the language model used in the session."""
+CompletionParams = Dict[str, Any]
+"""
+Common set of completion parameters that most model providers support in their
+API. Set total=False to allow extra fields since LiteLLM + providers have their
+own set of parameters. The following parameters are common fields that are
+populated.
 
-    model: str = Field(..., description="Model identifier (e.g., 'gpt-4.1', 'fireworks/llama')")
-    temperature: Optional[float] = Field(None, description="Temperature setting for model generation")
-    max_tokens: Optional[int] = Field(None, description="Maximum tokens to generate")
-    max_tool_calls: Optional[int] = Field(None, description="Maximum tool calls per turn")
-
-    # there might be model or provider specific parameters that you want to pass that should be preserved
-    model_config = ConfigDict(extra="allow")
+model: str
+temperature: Optional[float]
+max_tokens: Optional[int]
+top_p: Optional[float]
+"""
 
 
 class InputMetadata(BaseModel):
@@ -196,7 +198,7 @@ class InputMetadata(BaseModel):
     model_config = ConfigDict(extra="allow")
 
     row_id: Optional[str] = Field(default_factory=generate_id, description="Unique string to ID the row")
-    completion_params: Optional[CompletionParams] = Field(None, description="Completion endpoint parameters used")
+    completion_params: CompletionParams = Field(..., description="Completion endpoint parameters used")
     dataset_info: Optional[Dict[str, Any]] = Field(
         None, description="Dataset row details: seed, system_prompt, environment_context, etc"
     )

eval_protocol/models.py

@@ -41,24 +41,20 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
 
         messages_payload = [{"role": m.role, "content": m.content} for m in row.messages]
 
-        request_params = {
-            "model": config.completion_params.model,
-            "messages": messages_payload,
-            **config.completion_params,
-        }
+        request_params = {"messages": messages_payload, **config.completion_params}
         # Ensure caching is disabled only for this request (review feedback)
         request_params["cache"] = {"no-cache": True}
         # Single-level reasoning effort: expect `reasoning_effort` only
         effort_val = None
-        if isinstance(config.completion_params, dict):
-            if "reasoning_effort" in config.completion_params:
-                effort_val = str(config.completion_params["reasoning_effort"])  # flat shape
-            elif (
-                isinstance(config.completion_params.get("extra_body"), dict)
-                and "reasoning_effort" in config.completion_params["extra_body"]
-            ):
-                # Accept if user passed it directly inside extra_body
-                effort_val = str(config.completion_params["extra_body"]["reasoning_effort"])  # already in extra_body
+
+        if "reasoning_effort" in config.completion_params:
+            effort_val = str(config.completion_params["reasoning_effort"])  # flat shape
+        elif (
+            isinstance(config.completion_params.get("extra_body"), dict)
+            and "reasoning_effort" in config.completion_params["extra_body"]
+        ):
+            # Accept if user passed it directly inside extra_body
+            effort_val = str(config.completion_params["extra_body"]["reasoning_effort"])  # already in extra_body
 
         if effort_val:
             # Always under extra_body so LiteLLM forwards to provider-specific param set

eval_protocol/pytest/default_single_turn_rollout_process.py

@@ -26,7 +26,6 @@
 from eval_protocol.pytest.default_dataset_adapter import default_dataset_adapter
 from eval_protocol.pytest.default_no_op_rollout_process import default_no_op_rollout_processor
 from eval_protocol.pytest.types import (
-    CompletionsParams,
     Dataset,
     DatasetPathParam,
     EvaluationInputParam,
@@ -52,7 +51,7 @@
 
 def evaluation_test(  # noqa: C901
     *,
-    completion_params: List[CompletionsParams],
+    completion_params: List[CompletionParams],
     input_messages: Optional[List[InputMessagesParam]] = None,
     input_dataset: Optional[List[DatasetPathParam]] = None,
     dataset_adapter: Callable[[List[Dict[str, Any]]], Dataset] = default_dataset_adapter,
@@ -240,7 +239,7 @@ def generate_combinations():
                         datasets = [[input_dataset]]  # type: ignore
             else:
                 datasets = [None]
-            cps: List[Optional[CompletionsParams]] = completion_params if completion_params is not None else [None]  # type: ignore
+            cps: List[Optional[CompletionParams]] = completion_params if completion_params is not None else [None]  # type: ignore
             # 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).
@@ -348,7 +347,16 @@ def _log_eval_error(
                     else:
                         raise ValueError("No input dataset or input messages provided")
 
-                    completions_params = kwargs.get("completion_params") or {}
+                    if "completion_params" not in kwargs or not kwargs["completion_params"]:
+                        raise ValueError(
+                            "No completion parameters provided. Please provide a completion parameters object."
+                        )
+                    completion_params = kwargs["completion_params"]
+                    if "model" not in completion_params or not completion_params["model"]:
+                        raise ValueError(
+                            "No model provided. Please provide a model in the completion parameters object."
+                        )
+
                     # Optional global overrides via environment for ad-hoc experimentation
                     # EP_INPUT_PARAMS_JSON can contain a JSON object that will be deep-merged
                     # into input_params (e.g., '{"temperature":0,"extra_body":{"reasoning":{"effort":"low"}}}').
@@ -359,7 +367,7 @@ def _log_eval_error(
                         if _env_override:
                             override_obj = _json.loads(_env_override)
                             if isinstance(override_obj, dict):
-                                completions_params = _deep_update_dict(dict(completions_params), override_obj)
+                                completion_params = _deep_update_dict(dict(completion_params), override_obj)
                     except Exception:
                         pass
 
@@ -374,11 +382,6 @@ def _log_eval_error(
                         passed=None,
                     )
 
-                    # Populate completion_params in input_metadata for all rows and initialize eval_metadata BEFORE rollouts
-                    completion_params = CompletionParams(
-                        **completions_params,
-                    )
-
                     for row in data:
                         if row.input_metadata is None:
                             row.input_metadata = InputMetadata()
@@ -398,13 +401,13 @@ def _log_eval_error(
 
                     # Prepare rollout processor config once; we will generate fresh outputs per run
                     config = RolloutProcessorConfig(
-                        completion_params=CompletionParams(**completions_params),
+                        completion_params=completion_params,
                         mcp_config_path=mcp_config_path or "",
                         max_concurrent_rollouts=max_concurrent_rollouts,
                         server_script_path=server_script_path,
                         steps=steps,
                         logger=active_logger,
-                        kwargs=rollout_processor_kwargs,
+                        kwargs=rollout_processor_kwargs or {},
                     )
 
                     for i in range(num_runs):
@@ -611,7 +614,7 @@ def _extract_effort_tag(params: dict) -> str | None:
                                 return None
 
                             model_slug = _sanitize_filename(model_used)
-                            effort_tag = _extract_effort_tag(completions_params) or ""
+                            effort_tag = _extract_effort_tag(completion_params) or ""
                             effort_suffix = f"__effort-{_sanitize_filename(effort_tag)}" if effort_tag else ""
                             base_name = f"{suite_name}__{model_slug}{effort_suffix}__{mode}__runs{num_runs}.json"
 
@@ -788,7 +791,7 @@ def __ep_run_direct(
                     input_messages=cfg.get("input_messages"),
                     input_dataset=cfg.get("input_dataset"),
                     dataset_adapter=cfg.get("dataset_adapter"),
-                    completions_params=rip,
+                    completion_params=rip,
                     rollout_processor=cfg.get("rollout_processor"),
                     aggregation_method=cfg.get("aggregation_method"),
                     passed_threshold=cfg.get("passed_threshold"),
@@ -818,7 +821,7 @@ def run_evaluation_test_direct(
     input_messages: Optional[List[InputMessagesParam]] = None,
     input_dataset: Optional[List[DatasetPathParam]] = None,
     dataset_adapter: Callable[[List[Dict[str, Any]]], Dataset] = default_dataset_adapter,
-    completions_params: Optional[CompletionsParams] = None,
+    completion_params: Optional[CompletionParams] = None,
     rollout_processor: RolloutProcessor = default_no_op_rollout_processor,
     rollout_processor_kwargs: Optional[RolloutProcessorInputParam] = None,
     aggregation_method: AggregationMethod = "mean",
@@ -885,7 +888,7 @@ def _deep_update_dict(base: dict, override: dict) -> dict:
         raise ValueError("No input dataset or input messages provided")
 
     # Build input params and apply env JSON override
-    completion_params: Dict[str, Any] = completions_params or {}
+    completion_params: Dict[str, Any] = completion_params or {}
     try:
         import json as _json
 
@@ -911,7 +914,7 @@ def _deep_update_dict(base: dict, override: dict) -> dict:
     for row in data:
         if row.input_metadata is None:
             row.input_metadata = InputMetadata()
-        row.input_metadata.completion_params = CompletionParams(**completion_params)
+        row.input_metadata.completion_params = completion_params
         if row.input_metadata.session_data is None:
             row.input_metadata.session_data = {}
         row.input_metadata.session_data["mode"] = mode
@@ -925,7 +928,7 @@ def _deep_update_dict(base: dict, override: dict) -> dict:
         max_concurrent_rollouts=max_concurrent_rollouts,
         server_script_path=server_script_path,
         steps=steps,
-        kwargs=rollout_processor_kwargs,
+        kwargs=rollout_processor_kwargs or {},
     )
 
     all_results: List[EvaluationRow] = []

eval_protocol/pytest/evaluation_test.py

@@ -8,11 +8,10 @@
 from eval_protocol.dataset_logger import default_logger
 from eval_protocol.dataset_logger.dataset_logger import DatasetLogger
 
-from ..models import EvaluationRow, Message
+from ..models import CompletionParams, EvaluationRow, Message
 
 ModelParam = str  # gpt-4o, gpt-4o-mini, accounts/fireworks/models/llama-3.1-8b-instruct
 DatasetPathParam = str
-CompletionsParams = Dict[str, Any]
 InputMessagesParam = List[Message]
 EvaluationInputParam = Dict[str, Any]
 RolloutProcessorInputParam = Dict[str, Any]
@@ -41,7 +40,7 @@
 
 @dataclass
 class RolloutProcessorConfig:
-    completion_params: CompletionsParams  # input parameters for inference
+    completion_params: CompletionParams  # input parameters for inference
     mcp_config_path: str
     server_script_path: Optional[str] = (
         None  # TODO: change from server_script_path to mcp_config_path for agent rollout processor

eval_protocol/pytest/types.py

@@ -7,7 +7,7 @@
 
 from typing import Any, Dict, List
 
-from eval_protocol.models import CompletionParams, EvaluateResult, EvaluationRow, InputMetadata, Message, MetricResult
+from eval_protocol.models import EvaluateResult, EvaluationRow, InputMetadata, Message, MetricResult
 from eval_protocol.pytest import evaluation_test
 from eval_protocol.pytest.default_mcp_gym_rollout_processor import default_mcp_gym_rollout_processor
 

tests/pytest/test_frozen_lake.py

@@ -7,7 +7,7 @@
 
 from typing import Any, Dict, List
 
-from eval_protocol.models import CompletionParams, EvaluateResult, EvaluationRow, InputMetadata, Message
+from eval_protocol.models import EvaluateResult, EvaluationRow, InputMetadata, Message
 from eval_protocol.pytest import evaluation_test
 from eval_protocol.pytest.default_mcp_gym_rollout_processor import default_mcp_gym_rollout_processor