datarobot-genai 0.2.26__py3-none-any.whl → 0.2.34__py3-none-any.whl

datarobot_genai/drmcp/tools/predictive/model.py

@@ -14,9 +14,12 @@
 
 import json
 import logging
+from typing import Annotated
 from typing import Any
 
 from datarobot.models.model import Model
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
 
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
@@ -50,33 +53,25 @@ class ModelEncoder(json.JSONEncoder):
         return super().default(obj)
 
 
-@dr_mcp_tool(tags={"model", "management", "info"})
-async def get_best_model(project_id: str, metric: str | None = None) -> str:
-    """
-    Get the best model for a DataRobot project, optionally by a specific metric.
+@dr_mcp_tool(tags={"predictive", "model", "read", "management", "info"})
+async def get_best_model(
+    *,
+    project_id: Annotated[str, "The DataRobot project ID"] | None = None,
+    metric: Annotated[str, "The metric to use for best model selection (e.g., 'AUC', 'LogLoss')"]
+    | None = None,
+) -> ToolError | ToolResult:
+    """Get the best model for a DataRobot project, optionally by a specific metric."""
+    if not project_id:
+        raise ToolError("Project ID must be provided")
 
-    Args:
-        project_id: The ID of the DataRobot project.
-        metric: (Optional) The metric to use for best model selection (e.g., 'AUC', 'LogLoss').
-
-    Returns
-    -------
-        A formatted string describing the best model.
-
-    Raises
-    ------
-        Exception: If project not found or no models exist in the project.
-    """
     client = get_sdk_client()
     project = client.Project.get(project_id)
     if not project:
-        logger.error(f"Project with ID {project_id} not found")
-        raise Exception(f"Project with ID {project_id} not found.")
+        raise ToolError(f"Project with ID {project_id} not found.")
 
     leaderboard = project.get_models()
     if not leaderboard:
-        logger.info(f"No models found for project {project_id}")
-        raise Exception("No models found for this project.")
+        raise ToolError("No models found for this project.")
 
     if metric:
         reverse_sort = metric.upper() in [
@@ -98,51 +93,91 @@ async def get_best_model(project_id: str, metric: str | None = None) -> str:
     best_model = leaderboard[0]
     logger.info(f"Found best model {best_model.id} for project {project_id}")
 
-    # Format the response as a human-readable string
     metric_info = ""
+    metric_value = None
+
     if metric and best_model.metrics and metric in best_model.metrics:
         metric_value = best_model.metrics[metric].get("validation")
         if metric_value is not None:
             metric_info = f" with {metric}: {metric_value:.2f}"
 
-    return f"Best model: {best_model.model_type}{metric_info}"
-
-
-@dr_mcp_tool(tags={"model", "prediction", "scoring"})
-async def score_dataset_with_model(project_id: str, model_id: str, dataset_url: str) -> str:
-    """
-    Score a dataset using a specific DataRobot model.
+    # Include full metrics in the response
+    best_model_dict = model_to_dict(best_model)
+    best_model_dict["metric"] = metric
+    best_model_dict["metric_value"] = metric_value
+
+    # Format metrics for human-readable content
+    metrics_text = ""
+    if best_model.metrics:
+        metrics_list = []
+        for metric_name, metric_data in best_model.metrics.items():
+            if isinstance(metric_data, dict) and "validation" in metric_data:
+                val = metric_data["validation"]
+                if val is not None:
+                    metrics_list.append(f"{metric_name}: {val:.4f}")
+        if metrics_list:
+            metrics_text = "\nPerformance metrics:\n" + "\n".join(f"  - {m}" for m in metrics_list)
+
+    return ToolResult(
+        content=f"Best model: {best_model.model_type}{metric_info}{metrics_text}",
+        structured_content={
+            "project_id": project_id,
+            "best_model": best_model_dict,
+        },
+    )
+
+
+@dr_mcp_tool(tags={"predictive", "model", "read", "scoring", "dataset"})
+async def score_dataset_with_model(
+    *,
+    project_id: Annotated[str, "The DataRobot project ID"] | None = None,
+    model_id: Annotated[str, "The DataRobot model ID"] | None = None,
+    dataset_url: Annotated[str, "The dataset URL"] | None = None,
+) -> ToolError | ToolResult:
+    """Score a dataset using a specific DataRobot model."""
+    if not project_id:
+        raise ToolError("Project ID must be provided")
+    if not model_id:
+        raise ToolError("Model ID must be provided")
+    if not dataset_url:
+        raise ToolError("Dataset URL must be provided")
 
-    Args:
-        project_id: The ID of the DataRobot project.
-        model_id: The ID of the DataRobot model to use for scoring.
-        dataset_url: The URL to the dataset to score (must be accessible to DataRobot).
-
-    Returns
-    -------
-        A string summary of the scoring job or a meaningful error message.
-    """
     client = get_sdk_client()
     project = client.Project.get(project_id)
     model = client.Model.get(project, model_id)
     job = model.score(dataset_url)
-    logger.info(f"Started scoring job {job.id} for model {model_id}")
-    return f"Scoring job started: {job.id}"
-
 
-@dr_mcp_tool(tags={"model", "management", "list"})
-async def list_models(project_id: str) -> str:
-    """
-    List all models in a project.
+    return ToolResult(
+        content=f"Scoring job started: {job.id}",
+        structured_content={
+            "scoring_job_id": job.id,
+            "project_id": project_id,
+            "model_id": model_id,
+            "dataset_url": dataset_url,
+        },
+    )
+
+
+@dr_mcp_tool(tags={"predictive", "model", "read", "management", "list"})
+async def list_models(
+    *,
+    project_id: Annotated[str, "The DataRobot project ID"] | None = None,
+) -> ToolError | ToolResult:
+    """List all models in a project."""
+    if not project_id:
+        raise ToolError("Project ID must be provided")
 
-    Args:
-        project_id: The ID of the DataRobot project.
-
-    Returns
-    -------
-        A string summary of the models in the project.
-    """
     client = get_sdk_client()
     project = client.Project.get(project_id)
     models = project.get_models()
-    return json.dumps(models, indent=2, cls=ModelEncoder)
+
+    return ToolResult(
+        content=(
+            f"Found {len(models)} models in project {project_id}, here are the details:\n"
+            f"{json.dumps(models, indent=2, cls=ModelEncoder)}"
+        ),
+        structured_content={
+            "project_id": project_id,
+            "models": [model_to_dict(model) for model in models],
+        },
+    )

datarobot_genai/drmcp/tools/predictive/project.py

@@ -54,9 +54,9 @@ async def get_project_dataset_by_name(
     The dataset ID and the dataset type (source or prediction) as a string, or an error message.
     """
     if not project_id:
-        return ToolError("Project ID is required.")
+        raise ToolError("Project ID is required.")
     if not dataset_name:
-        return ToolError("Dataset name is required.")
+        raise ToolError("Dataset name is required.")
 
     client = get_sdk_client()
     project = client.Project.get(project_id)

datarobot_genai/drmcp/tools/predictive/training.py

@@ -63,7 +63,7 @@ async def analyze_dataset(
 ) -> ToolError | ToolResult:
     """Analyze a dataset to understand its structure and potential use cases."""
     if not dataset_id:
-        return ToolError("Dataset ID must be provided")
+        raise ToolError("Dataset ID must be provided")
 
     client = get_sdk_client()
     dataset = client.Dataset.get(dataset_id)
@@ -116,7 +116,7 @@ async def suggest_use_cases(
 ) -> ToolError | ToolResult:
     """Analyze a dataset and suggest potential machine learning use cases."""
     if not dataset_id:
-        return ToolError("Dataset ID must be provided")
+        raise ToolError("Dataset ID must be provided")
 
     client = get_sdk_client()
     dataset = client.Dataset.get(dataset_id)
@@ -148,7 +148,7 @@ async def get_exploratory_insights(
 ) -> ToolError | ToolResult:
     """Generate exploratory data insights for a dataset."""
     if not dataset_id:
-        return ToolError("Dataset ID must be provided")
+        raise ToolError("Dataset ID must be provided")
 
     client = get_sdk_client()
     dataset = client.Dataset.get(dataset_id)
@@ -481,9 +481,9 @@ async def start_autopilot(
 
     if not project_id:
         if not dataset_url and not dataset_id:
-            return ToolError("Either dataset_url or dataset_id must be provided")
+            raise ToolError("Either dataset_url or dataset_id must be provided")
         if dataset_url and dataset_id:
-            return ToolError("Please provide either dataset_url or dataset_id, not both")
+            raise ToolError("Please provide either dataset_url or dataset_id, not both")
 
         if dataset_url:
             dataset = client.Dataset.create_from_url(dataset_url)
@@ -497,7 +497,7 @@ async def start_autopilot(
         project = client.Project.get(project_id)
 
     if not target:
-        return ToolError("Target variable must be specified")
+        raise ToolError("Target variable must be specified")
 
     try:
         # Start modeling
@@ -517,7 +517,7 @@ async def start_autopilot(
         )
 
     except Exception as e:
-        return ToolError(
+        raise ToolError(
             content=json.dumps(
                 {
                     "error": f"Failed to start Autopilot: {str(e)}",
@@ -546,9 +546,9 @@ async def get_model_roc_curve(
 ) -> ToolError | ToolResult:
     """Get detailed ROC curve for a specific model."""
     if not project_id:
-        return ToolError("Project ID must be provided")
+        raise ToolError("Project ID must be provided")
     if not model_id:
-        return ToolError("Model ID must be provided")
+        raise ToolError("Model ID must be provided")
 
     client = get_sdk_client()
     project = client.Project.get(project_id)
@@ -587,7 +587,7 @@ async def get_model_roc_curve(
             structured_content={"data": roc_data},
         )
     except Exception as e:
-        return ToolError(f"Failed to get ROC curve: {str(e)}")
+        raise ToolError(f"Failed to get ROC curve: {str(e)}")
 
 
 @dr_mcp_tool(tags={"predictive", "training", "read", "model", "evaluation"})
@@ -598,9 +598,9 @@ async def get_model_feature_impact(
 ) -> ToolError | ToolResult:
     """Get detailed feature impact for a specific model."""
     if not project_id:
-        return ToolError("Project ID must be provided")
+        raise ToolError("Project ID must be provided")
     if not model_id:
-        return ToolError("Model ID must be provided")
+        raise ToolError("Model ID must be provided")
 
     client = get_sdk_client()
     project = client.Project.get(project_id)
@@ -617,6 +617,7 @@ async def get_model_feature_impact(
 
 @dr_mcp_tool(tags={"predictive", "training", "read", "model", "evaluation"})
 async def get_model_lift_chart(
+    *,
     project_id: Annotated[str, "The ID of the DataRobot project"] | None = None,
     model_id: Annotated[str, "The ID of the model to analyze"] | None = None,
     source: Annotated[
@@ -630,9 +631,9 @@ async def get_model_lift_chart(
 ) -> ToolError | ToolResult:
     """Get detailed lift chart for a specific model."""
     if not project_id:
-        return ToolError("Project ID must be provided")
+        raise ToolError("Project ID must be provided")
     if not model_id:
-        return ToolError("Model ID must be provided")
+        raise ToolError("Model ID must be provided")
 
     client = get_sdk_client()
     project = client.Project.get(project_id)

datarobot_genai/nat/datarobot_llm_clients.py

@@ -12,22 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import TypeVar
 
-from crewai import LLM
-from langchain_openai import ChatOpenAI
-from llama_index.core.base.llms.types import LLMMetadata
-from llama_index.llms.litellm import LiteLLM
 from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_llm_client
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.retry_mixin import RetryMixin
-from nat.plugins.langchain.llm import (
-    _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
-)
 from nat.utils.exception_handlers.automatic_retries import patch_with_retry
 
 from ..nat.datarobot_llm_providers import DataRobotLLMComponentModelConfig
@@ -35,6 +31,11 @@ from ..nat.datarobot_llm_providers import DataRobotLLMDeploymentModelConfig
 from ..nat.datarobot_llm_providers import DataRobotLLMGatewayModelConfig
 from ..nat.datarobot_llm_providers import DataRobotNIMModelConfig
 
+if TYPE_CHECKING:
+    from crewai import LLM
+    from langchain_openai import ChatOpenAI
+    from llama_index.llms.litellm import LiteLLM
+
 ModelType = TypeVar("ModelType")
 
 
@@ -50,42 +51,53 @@ def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) ->
     return client
 
 
-class DataRobotChatOpenAI(ChatOpenAI):
-    def _get_request_payload(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> dict:
-        # We need to default to include_usage=True for streaming but we get 400 response
-        # if stream_options is present for a non-streaming call.
-        payload = super()._get_request_payload(*args, **kwargs)
-        if not payload.get("stream"):
-            payload.pop("stream_options", None)
-        return payload
-
-
-class DataRobotLiteLLM(LiteLLM):  # type: ignore[misc]
-    """DataRobotLiteLLM is a small LiteLLM wrapper class that makes all LiteLLM endpoints
-    compatible with the LlamaIndex library.
-    """
-
-    @property
-    def metadata(self) -> LLMMetadata:
-        """Returns the metadata for the LLM.
-
-        This is required to enable the is_chat_model and is_function_calling_model, which are
-        mandatory for LlamaIndex agents. By default, LlamaIndex assumes these are false unless each
-        individual model config in LiteLLM explicitly sets them to true. To use custom LLM
-        endpoints with LlamaIndex agents, you must override this method to return the appropriate
-        metadata.
+def _create_datarobot_chat_openai(config: dict[str, Any]) -> Any:
+    from langchain_openai import ChatOpenAI  # noqa: PLC0415
+
+    class DataRobotChatOpenAI(ChatOpenAI):
+        def _get_request_payload(  # type: ignore[override]
+            self,
+            *args: Any,
+            **kwargs: Any,
+        ) -> dict:
+            # We need to default to include_usage=True for streaming but we get 400 response
+            # if stream_options is present for a non-streaming call.
+            payload = super()._get_request_payload(*args, **kwargs)
+            if not payload.get("stream"):
+                payload.pop("stream_options", None)
+            return payload
+
+    return DataRobotChatOpenAI(**config)
+
+
+def _create_datarobot_litellm(config: dict[str, Any]) -> Any:
+    from llama_index.core.base.llms.types import LLMMetadata  # noqa: PLC0415
+    from llama_index.llms.litellm import LiteLLM  # noqa: PLC0415
+
+    class DataRobotLiteLLM(LiteLLM):  # type: ignore[misc]
+        """DataRobotLiteLLM is a small LiteLLM wrapper class that makes all LiteLLM endpoints
+        compatible with the LlamaIndex library.
         """
-        return LLMMetadata(
-            context_window=128000,
-            num_output=self.max_tokens or -1,
-            is_chat_model=True,
-            is_function_calling_model=True,
-            model_name=self.model,
-        )
+
+        @property
+        def metadata(self) -> LLMMetadata:
+            """Returns the metadata for the LLM.
+
+            This is required to enable the is_chat_model and is_function_calling_model, which are
+            mandatory for LlamaIndex agents. By default, LlamaIndex assumes these are false unless
+            each individual model config in LiteLLM explicitly sets them to true. To use custom LLM
+            endpoints with LlamaIndex agents, you must override this method to return the
+            appropriate metadata.
+            """
+            return LLMMetadata(
+                context_window=128000,
+                num_output=self.max_tokens or -1,
+                is_chat_model=True,
+                is_function_calling_model=True,
+                model_name=self.model,
+            )
+
+    return DataRobotLiteLLM(**config)
 
 
 @register_llm_client(
@@ -94,11 +106,15 @@ class DataRobotLiteLLM(LiteLLM):  # type: ignore[misc]
 async def datarobot_llm_gateway_langchain(
     llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
 ) -> AsyncGenerator[ChatOpenAI]:
+    from nat.plugins.langchain.llm import (  # noqa: PLC0415
+        _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+    )
+
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
     config["base_url"] = config["base_url"] + "/genai/llmgw"
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
-    client = DataRobotChatOpenAI(**config)
+    client = _create_datarobot_chat_openai(config)
     yield langchain_patch_llm_based_on_config(client, config)
 
 
@@ -108,6 +124,8 @@ async def datarobot_llm_gateway_langchain(
 async def datarobot_llm_gateway_crewai(
     llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
 ) -> AsyncGenerator[LLM]:
+    from crewai import LLM  # noqa: PLC0415
+
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
@@ -121,12 +139,12 @@ async def datarobot_llm_gateway_crewai(
 )
 async def datarobot_llm_gateway_llamaindex(
     llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
-) -> AsyncGenerator[LLM]:
+) -> AsyncGenerator[LiteLLM]:
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url").removesuffix("/api/v2")
-    client = DataRobotLiteLLM(**config)
+    client = _create_datarobot_litellm(config)
     yield _patch_llm_based_on_config(client, config)
 
 
@@ -136,6 +154,10 @@ async def datarobot_llm_gateway_llamaindex(
 async def datarobot_llm_deployment_langchain(
     llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
 ) -> AsyncGenerator[ChatOpenAI]:
+    from nat.plugins.langchain.llm import (  # noqa: PLC0415
+        _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+    )
+
     config = llm_config.model_dump(
         exclude={"type", "thinking"},
         by_alias=True,
@@ -143,7 +165,7 @@ async def datarobot_llm_deployment_langchain(
     )
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
-    client = DataRobotChatOpenAI(**config)
+    client = _create_datarobot_chat_openai(config)
     yield langchain_patch_llm_based_on_config(client, config)
 
 
@@ -153,6 +175,8 @@ async def datarobot_llm_deployment_langchain(
 async def datarobot_llm_deployment_crewai(
     llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
 ) -> AsyncGenerator[LLM]:
+    from crewai import LLM  # noqa: PLC0415
+
     config = llm_config.model_dump(
         exclude={"type", "thinking"},
         by_alias=True,
@@ -170,7 +194,7 @@ async def datarobot_llm_deployment_crewai(
 )
 async def datarobot_llm_deployment_llamaindex(
     llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
-) -> AsyncGenerator[LLM]:
+) -> AsyncGenerator[LiteLLM]:
     config = llm_config.model_dump(
         exclude={"type", "thinking"},
         by_alias=True,
@@ -179,7 +203,7 @@ async def datarobot_llm_deployment_llamaindex(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
-    client = DataRobotLiteLLM(**config)
+    client = _create_datarobot_litellm(config)
     yield _patch_llm_based_on_config(client, config)
 
 
@@ -187,6 +211,10 @@ async def datarobot_llm_deployment_llamaindex(
 async def datarobot_nim_langchain(
     llm_config: DataRobotNIMModelConfig, builder: Builder
 ) -> AsyncGenerator[ChatOpenAI]:
+    from nat.plugins.langchain.llm import (  # noqa: PLC0415
+        _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+    )
+
     config = llm_config.model_dump(
         exclude={"type", "thinking"},
         by_alias=True,
@@ -194,7 +222,7 @@ async def datarobot_nim_langchain(
     )
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
-    client = DataRobotChatOpenAI(**config)
+    client = _create_datarobot_chat_openai(config)
     yield langchain_patch_llm_based_on_config(client, config)
 
 
@@ -202,6 +230,8 @@ async def datarobot_nim_langchain(
 async def datarobot_nim_crewai(
     llm_config: DataRobotNIMModelConfig, builder: Builder
 ) -> AsyncGenerator[LLM]:
+    from crewai import LLM  # noqa: PLC0415
+
     config = llm_config.model_dump(
         exclude={"type", "thinking", "max_retries"},
         by_alias=True,
@@ -217,7 +247,7 @@ async def datarobot_nim_crewai(
 @register_llm_client(config_type=DataRobotNIMModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
 async def datarobot_nim_llamaindex(
     llm_config: DataRobotNIMModelConfig, builder: Builder
-) -> AsyncGenerator[LLM]:
+) -> AsyncGenerator[LiteLLM]:
     config = llm_config.model_dump(
         exclude={"type", "thinking"},
         by_alias=True,
@@ -226,7 +256,7 @@ async def datarobot_nim_llamaindex(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
-    client = DataRobotLiteLLM(**config)
+    client = _create_datarobot_litellm(config)
     yield _patch_llm_based_on_config(client, config)
 
 
@@ -236,13 +266,17 @@ async def datarobot_nim_llamaindex(
 async def datarobot_llm_component_langchain(
     llm_config: DataRobotLLMComponentModelConfig, builder: Builder
 ) -> AsyncGenerator[ChatOpenAI]:
+    from nat.plugins.langchain.llm import (  # noqa: PLC0415
+        _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+    )
+
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
     if config["use_datarobot_llm_gateway"]:
         config["base_url"] = config["base_url"] + "/genai/llmgw"
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
     config.pop("use_datarobot_llm_gateway")
-    client = DataRobotChatOpenAI(**config)
+    client = _create_datarobot_chat_openai(config)
     yield langchain_patch_llm_based_on_config(client, config)
 
 
@@ -252,6 +286,8 @@ async def datarobot_llm_component_langchain(
 async def datarobot_llm_component_crewai(
     llm_config: DataRobotLLMComponentModelConfig, builder: Builder
 ) -> AsyncGenerator[LLM]:
+    from crewai import LLM  # noqa: PLC0415
+
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
@@ -269,7 +305,7 @@ async def datarobot_llm_component_crewai(
 )
 async def datarobot_llm_component_llamaindex(
     llm_config: DataRobotLLMComponentModelConfig, builder: Builder
-) -> AsyncGenerator[LLM]:
+) -> AsyncGenerator[LiteLLM]:
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
@@ -278,5 +314,5 @@ async def datarobot_llm_component_llamaindex(
     else:
         config["api_base"] = config.pop("base_url") + "/chat/completions"
     config.pop("use_datarobot_llm_gateway")
-    client = DataRobotLiteLLM(**config)
+    client = _create_datarobot_litellm(config)
     yield _patch_llm_based_on_config(client, config)