add extra_body to ModelProvider

Author: nabinchha

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/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/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}}