update

Author: xzrderek

eval_protocol/pytest/default_agent_rollout_processor.py

@@ -22,7 +22,7 @@
 from openai.types import CompletionUsage
 from eval_protocol.pytest.rollout_processor import RolloutProcessor
 from eval_protocol.pytest.types import Dataset, RolloutProcessorConfig
-from eval_protocol.pytest.evaluation_test_utils import normalize_fireworks_model_for_litellm
+from eval_protocol.pytest.utils import normalize_fireworks_model_for_litellm
 from pydantic import BaseModel
 from typing import Optional
 

eval_protocol/pytest/default_klavis_sandbox_rollout_processor.py

@@ -11,7 +11,7 @@
 from eval_protocol.models import EvaluationRow
 from eval_protocol.pytest.rollout_processor import RolloutProcessor
 from eval_protocol.pytest.types import RolloutProcessorConfig
-from eval_protocol.pytest.evaluation_test_utils import normalize_fireworks_model_for_litellm
+from eval_protocol.pytest.utils import normalize_fireworks_model_for_litellm
 
 from eval_protocol.pytest.default_agent_rollout_processor import Agent
 from klavis import Klavis

eval_protocol/pytest/default_mcp_gym_rollout_processor.py

@@ -14,7 +14,7 @@
 from eval_protocol.models import EvaluationRow
 from eval_protocol.pytest.rollout_processor import RolloutProcessor
 from eval_protocol.pytest.types import RolloutProcessorConfig, ServerMode
-from eval_protocol.pytest.evaluation_test_utils import normalize_fireworks_model_for_litellm
+from eval_protocol.pytest.utils import normalize_fireworks_model_for_litellm
 
 
 class MCPServerManager:

eval_protocol/pytest/default_single_turn_rollout_process.py

@@ -17,7 +17,7 @@
 from openai.types import CompletionUsage
 from eval_protocol.pytest.rollout_processor import RolloutProcessor
 from eval_protocol.pytest.types import RolloutProcessorConfig
-from eval_protocol.pytest.evaluation_test_utils import normalize_fireworks_model_for_litellm
+from eval_protocol.pytest.utils import normalize_fireworks_model_for_litellm
 
 logger = logging.getLogger(__name__)
 

eval_protocol/pytest/evaluation_test_utils.py

@@ -619,22 +619,3 @@ def build_rollout_processor_config(
         server_script_path=None,
         kwargs=rollout_processor_kwargs,
     )
-
-
-def normalize_fireworks_model_for_litellm(completion_params: CompletionParams | None) -> CompletionParams | None:
-    """Fireworks model names like 'accounts/<org>/models/<model>' need the fireworks_ai/
-    prefix when routing through LiteLLM. This function adds the prefix if missing.
-    """
-    if completion_params is None:
-        return None
-
-    model = completion_params.get("model")
-    if (
-        model
-        and isinstance(model, str)
-        and not model.startswith("fireworks_ai/")
-        and re.match(r"^accounts/[^/]+/models/.+", model)
-    ):
-        completion_params = completion_params.copy()
-        completion_params["model"] = f"fireworks_ai/{model}"
-    return completion_params

eval_protocol/pytest/openenv_rollout_processor.py

@@ -24,7 +24,7 @@
 from eval_protocol.models import EvaluationRow, Message
 from eval_protocol.pytest.rollout_processor import RolloutProcessor
 from eval_protocol.pytest.types import RolloutProcessorConfig
-from eval_protocol.pytest.evaluation_test_utils import normalize_fireworks_model_for_litellm
+from eval_protocol.pytest.utils import normalize_fireworks_model_for_litellm
 
 logger = logging.getLogger(__name__)
 

eval_protocol/pytest/tracing_utils.py

@@ -12,7 +12,7 @@
 from eval_protocol.utils.evaluation_row_utils import filter_longest_conversation
 from eval_protocol.types.remote_rollout_processor import DataLoaderConfig, RolloutMetadata, InitRequest
 from eval_protocol.pytest.types import RolloutProcessorConfig
-from eval_protocol.pytest.evaluation_test_utils import normalize_fireworks_model_for_litellm
+from eval_protocol.pytest.utils import normalize_fireworks_model_for_litellm
 
 
 def default_fireworks_output_data_loader(config: DataLoaderConfig) -> DynamicDataLoader:

eval_protocol/pytest/utils.py

@@ -0,0 +1,25 @@
+"""Utility functions for model name handling."""
+
+import re
+from typing import Optional
+
+from eval_protocol.models import CompletionParams
+
+
+def normalize_fireworks_model_for_litellm(completion_params: CompletionParams | None) -> CompletionParams | None:
+    """Fireworks model names like 'accounts/<org>/models/<model>' or 'accounts/<org>/deployments/<model>'
+    need the fireworks_ai/ prefix when routing through LiteLLM. This function adds the prefix if missing.
+    """
+    if completion_params is None:
+        return None
+
+    model = completion_params.get("model")
+    if (
+        model
+        and isinstance(model, str)
+        and not model.startswith("fireworks_ai/")
+        and re.match(r"^accounts/[^/]+/(models|deployments)/.+", model)
+    ):
+        completion_params = completion_params.copy()
+        completion_params["model"] = f"fireworks_ai/{model}"
+    return completion_params