Skip to content

datarobot_genai.drmcp.test_utils.clients.dr_gateway

dr_gateway

DataRobot LLM Gateway MCP Client implementation.

DRLLMGatewayMCPClient

Bases: BaseLLMMCPClient

Client for interacting with LLMs via MCP using DataRobot LLM Gateway.

Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit(). Tools using FastMCP's built-in elicitation will work automatically.

Source code in datarobot_genai/drmcp/test_utils/clients/dr_gateway.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class DRLLMGatewayMCPClient(BaseLLMMCPClient):
    """
    Client for interacting with LLMs via MCP using DataRobot LLM Gateway.

    Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
    Tools using FastMCP's built-in elicitation will work automatically.
    """

    def __init__(
        self,
        config: str | dict,
    ):
        """
        Initialize the LLM MCP client.

        Args:
            config: Configuration string or dict with:
                - datarobot_api_token: DataRobot API token
                - datarobot_endpoint: DataRobot endpoint URL (default: "https://app.datarobot.com/api/v2")
                - model: Model name (**required**)
                - save_llm_responses: Whether to save responses (default: True)
                - temperature: (optional float, default: None)
        """
        super().__init__(config)

    def _create_llm_client(self, config_dict: dict) -> tuple[openai.OpenAI, str | None]:
        """Create the LLM client for DataRobot LLM Gateway."""
        datarobot_api_token = config_dict.get("datarobot_api_token")
        datarobot_endpoint = config_dict.get(
            "datarobot_endpoint", "https://app.datarobot.com/api/v2"
        )
        model = config_dict.get("model")

        # Build gateway URL: {endpoint}/genai/llmgw
        gateway_url = datarobot_endpoint.rstrip("/") + "/genai/llmgw"

        client = openai.OpenAI(api_key=datarobot_api_token, base_url=gateway_url)
        return client, model

__init__ 

__init__(config: str | dict)

Initialize the LLM MCP client.

Parameters:

Name Type Description Default
config str | dict

Configuration string or dict with: - datarobot_api_token: DataRobot API token - datarobot_endpoint: DataRobot endpoint URL (default: "https://app.datarobot.com/api/v2") - model: Model name (required) - save_llm_responses: Whether to save responses (default: True) - temperature: (optional float, default: None)

 required
Source code in datarobot_genai/drmcp/test_utils/clients/dr_gateway.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(
    self,
    config: str | dict,
):
    """
    Initialize the LLM MCP client.

    Args:
        config: Configuration string or dict with:
            - datarobot_api_token: DataRobot API token
            - datarobot_endpoint: DataRobot endpoint URL (default: "https://app.datarobot.com/api/v2")
            - model: Model name (**required**)
            - save_llm_responses: Whether to save responses (default: True)
            - temperature: (optional float, default: None)
    """
    super().__init__(config)