merge

Author: Shrey Modi

eval_protocol/pytest/github_action_rollout_processor.py

@@ -2,7 +2,7 @@
 import os
 import time
 from typing import Any, Callable, Dict, List, Optional
-
+import json
 import requests
 from datetime import datetime, timezone, timedelta
 from eval_protocol.models import EvaluationRow, Status
@@ -87,10 +87,14 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
 
             def _dispatch_workflow():
                 url = f"https://api.github.com/repos/{self.owner}/{self.repo}/actions/workflows/{self.workflow_id}/dispatches"
+
+                model = init_request.completion_params.get("model")
+                if not model:
+                    raise ValueError("model is required in completion_params")
                 payload = {
                     "ref": self.ref,
                     "inputs": {
-                        "model": init_request.model,
+                        "model": model,
                         "metadata": init_request.metadata.model_dump_json(),
                         "model_base_url": init_request.model_base_url,
                     },

eval_protocol/pytest/tracing_utils.py

@@ -80,21 +80,25 @@ def build_init_request(
         row_id=row.input_metadata.row_id,
     )
 
-    # Extract model
-    model: Optional[str] = None
+    # Build completion_params from row and config
+    completion_params_dict: Dict[str, Any] = {}
+    
+    # Start with config-level completion_params
+    if config.completion_params and isinstance(config.completion_params, dict):
+        completion_params_dict.update(config.completion_params)
+    
+    # Override with row-specific completion_params
     if row.input_metadata and row.input_metadata.completion_params:
-        model = row.input_metadata.completion_params.get("model")
-    if model is None and config.completion_params:
-        model = config.completion_params.get("model")
-    if model is None:
-        raise ValueError("Model must be provided in row.input_metadata.completion_params or config.completion_params")
-
+        row_cp = row.input_metadata.completion_params
+        if isinstance(row_cp, dict):
+            completion_params_dict.update(row_cp)
+    
+    # Validate model is present
+    if not completion_params_dict.get("model"):
+        raise ValueError("Model must be provided in completion_params")
+    
     # Extract base_url from completion_params
-    completion_params_base_url: Optional[str] = None
-    if row.input_metadata and row.input_metadata.completion_params:
-        completion_params_base_url = row.input_metadata.completion_params.get("base_url")
-    if completion_params_base_url is None and config.completion_params:
-        completion_params_base_url = config.completion_params.get("base_url")
+    completion_params_base_url: Optional[str] = completion_params_dict.get("base_url")
 
     # Strip non-OpenAI fields from messages
     allowed_message_fields = {"role", "content", "tool_calls", "tool_call_id", "name"}
@@ -124,7 +128,7 @@ def build_init_request(
         final_model_base_url = build_fireworks_tracing_url(model_base_url, meta, completion_params_base_url)
 
     return InitRequest(
-        model=model,
+        completion_params=completion_params_dict,
         messages=clean_messages,
         tools=row.tools,
         metadata=meta,

tests/github_actions/rollout_worker.py

@@ -18,12 +18,20 @@ def main():
 
     # Required arguments from workflow inputs
     parser.add_argument("--model", required=True, help="Model to use")
+    parser.add_argument("--completion-params", required=False, help="JSON completion params (optional)")
     parser.add_argument("--metadata", required=True, help="JSON serialized metadata object")
     parser.add_argument("--model-base-url", required=True, help="Base URL for the model API")
 
     args = parser.parse_args()
 
     # Parse the metadata
+    completion_params = {}
+    if args.completion_params:
+        try:
+            completion_params = json.loads(args.completion_params)
+        except Exception as e:
+            print(f"⚠️  Failed to parse completion_params: {e}")
+
     try:
         metadata = json.loads(args.metadata)
     except Exception as e:
@@ -51,6 +59,9 @@ def main():
     try:
         completion_kwargs = {"model": args.model, "messages": messages}
 
+        if completion_params.get("model_kwargs"):
+            completion_kwargs.update(completion_params["model_kwargs"])
+
         client = OpenAI(base_url=args.model_base_url, api_key=os.environ.get("FIREWORKS_API_KEY"))
 
         print("📡 Calling OpenAI completion...")

tests/remote_server/remote_server.py

@@ -35,7 +35,7 @@ def _worker():
         try:
             if not req.messages:
                 raise ValueError("messages is required")
-
+            
             model = req.completion_params.get("model")
             if not model:
                 raise ValueError("model is required in completion_params")
@@ -50,10 +50,12 @@ def _worker():
                 model_kwargs = req.completion_params["model_kwargs"]
                 if isinstance(model_kwargs, dict):
                     completion_kwargs.update(model_kwargs)
-                    
+
             if req.tools:
                 completion_kwargs["tools"] = req.tools
 
+            logger.info(f"Final completion_kwargs: {completion_kwargs}")
+            
             client = OpenAI(base_url=req.model_base_url, api_key=os.environ.get("FIREWORKS_API_KEY"))
 
             logger.info(f"Sending completion request to model {model}")

tests/remote_server/remote_server_multi_turn.py

@@ -31,6 +31,10 @@ def _worker():
         try:
             if not req.messages:
                 raise ValueError("messages is required")
+            
+            model = req.completion_params.get("model")
+            if not model:
+                raise ValueError("model is required in completion_params")
 
             client = OpenAI(base_url=req.model_base_url, api_key=os.environ.get("FIREWORKS_API_KEY"))
 

tests/remote_server/test_remote_fireworks.py

@@ -58,7 +58,9 @@ def rows() -> List[EvaluationRow]:
 
 
 @pytest.mark.skipif(os.environ.get("CI") == "true", reason="Only run this test locally (skipped in CI)")
-@pytest.mark.parametrize("completion_params", [{"model": "fireworks_ai/accounts/fireworks/models/gpt-oss-120b"}])
+@pytest.mark.parametrize("completion_params", [{"model": "fireworks_ai/accounts/fireworks/models/gpt-oss-120b", 
+"model_kwargs": {"temperature": 0.5}
+}])
 @evaluation_test(
     data_loaders=DynamicDataLoader(
         generators=[rows],
@@ -82,5 +84,6 @@ async def test_remote_rollout_and_fetch_fireworks(row: EvaluationRow) -> Evaluat
     assert row.execution_metadata.rollout_id in ROLLOUT_IDS, (
         f"Row rollout_id {row.execution_metadata.rollout_id} should be in tracked rollout_ids: {ROLLOUT_IDS}"
     )
-
+    assert row.input_metadata.completion_params["model_kwargs"] == {"temperature": 0.5}, "Row should have correct model_kwargs"
+    
     return row