feat: add extra_body to ModelProvider

src/data_designer/config/models.py

@@ -221,7 +221,8 @@ class ModelProvider(ConfigBase):
     name: str
     endpoint: str
     provider_type: str = "openai"
-    api_key: str | None = None
+    api_key: Optional[str] = None
+    extra_body: Optional[dict[str, Any]] = None
 
 
 def load_model_configs(model_configs: Union[list[ModelConfig], str, Path]) -> list[ModelConfig]:

src/data_designer/config/utils/info.py

@@ -12,9 +12,9 @@
 
 
 class InfoType(str, Enum):
-    SAMPLERS = "SAMPLERS"
-    MODEL_CONFIGS = "MODEL_CONFIGS"
-    MODEL_PROVIDERS = "MODEL_PROVIDERS"
+    SAMPLERS = "samplers"
+    MODEL_CONFIGS = "model_configs"
+    MODEL_PROVIDERS = "model_providers"
 
 
 ConfigBuilderInfoType = Literal[InfoType.SAMPLERS, InfoType.MODEL_CONFIGS]

src/data_designer/engine/models/facade.py

@@ -11,7 +11,7 @@
 from litellm.types.router import DeploymentTypedDict, LiteLLM_Params
 from litellm.types.utils import ModelResponse
 
-from data_designer.config.models import ModelConfig
+from data_designer.config.models import ModelConfig, ModelProvider
 from data_designer.engine.model_provider import ModelProviderRegistry
 from data_designer.engine.models.errors import (
     GenerationValidationFailureError,
@@ -45,9 +45,13 @@ def __init__(
     def model_name(self) -> str:
         return self._model_config.model
 
+    @property
+    def model_provider(self) -> ModelProvider:
+        return self._model_provider_registry.get_provider(self._model_config.provider)
+
     @property
     def model_provider_name(self) -> str:
-        return self._model_provider_registry.get_provider(self._model_config.provider).name
+        return self.model_provider.name
 
     @property
     def model_alias(self) -> str:
@@ -63,6 +67,8 @@ def completion(self, messages: list[dict[str, str]], skip_usage_tracking: bool =
             extra={"model": self.model_name, "messages": messages, "sensitive": True},
         )
         response = None
+        if self.model_provider.extra_body:
+            kwargs["extra_body"] = {**kwargs.get("extra_body", {}), **self.model_provider.extra_body}
         try:
             response = self._router.completion(self.model_name, messages, **kwargs)
             logger.debug(

tests/config/utils/test_info.py

@@ -38,7 +38,7 @@ def test_config_builder_unsupported_info_type(stub_model_configs):
     info = ConfigBuilderInfo(model_configs=stub_model_configs)
     with pytest.raises(
         ValueError,
-        match="Unsupported info_type: 'unsupported_type'. ConfigBuilderInfo only supports 'SAMPLERS' and 'MODEL_CONFIGS'.",
+        match="Unsupported info_type: 'unsupported_type'. ConfigBuilderInfo only supports 'samplers' and 'model_configs'.",
     ):
         info.display("unsupported_type")
 
@@ -53,6 +53,6 @@ def test_interface_model_providers_info(mock_display_model_providers_table, stub
 def test_interface_unsupported_info_type(stub_model_providers):
     info = InterfaceInfo(model_providers=stub_model_providers)
     with pytest.raises(
-        ValueError, match="Unsupported info_type: 'unsupported_type'. InterfaceInfo only supports 'MODEL_PROVIDERS'."
+        ValueError, match="Unsupported info_type: 'unsupported_type'. InterfaceInfo only supports 'model_providers'."
     ):
         info.display("unsupported_type")

tests/engine/models/test_facade.py

@@ -148,3 +148,27 @@ def mock_completion(model_name, messages, **kwargs):
 
     assert result == stub_expected_response
     assert captured_kwargs == kwargs
+
+
+@patch("data_designer.engine.models.facade.CustomRouter.completion", autospec=True)
+def test_completion_with_extra_body(mock_router_completion, stub_model_facade):
+    messages = [{"role": "user", "content": "test"}]
+
+    # completion call has no extra body argument and provider has no extra body
+    _ = stub_model_facade.completion(messages)
+    assert len(mock_router_completion.call_args) == 2
+    assert mock_router_completion.call_args[0][1] == "stub-model-text"
+    assert mock_router_completion.call_args[0][2] == messages
+
+    # completion call has no extra body argument and provider has extra body.
+    # Should pull extra body from model provider
+    custom_extra_body = {"some_custom_key": "some_custom_value"}
+    stub_model_facade.model_provider.extra_body = custom_extra_body
+    _ = stub_model_facade.completion(messages)
+    assert mock_router_completion.call_args[1] == {"extra_body": custom_extra_body}
+
+    # completion call has extra body argument and provider has extra body.
+    # Should merge the two with provider extra body taking precedence
+    completion_extra_body = {"some_completion_key": "some_completion_value", "some_custom_key": "some_different_value"}
+    _ = stub_model_facade.completion(messages, extra_body=completion_extra_body)
+    assert mock_router_completion.call_args[1] == {"extra_body": {**completion_extra_body, **custom_extra_body}}