Skip to content

datarobot_genai.drmcp.test_utils.clients.openai

openai

OpenAI LLM MCP Client implementation.

OpenAILLMMCPClient

Bases: BaseLLMMCPClient

Client for interacting with LLMs via MCP using OpenAI or Azure OpenAI.

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/openai.py
class OpenAILLMMCPClient(BaseLLMMCPClient):
    """
    Client for interacting with LLMs via MCP using OpenAI or Azure OpenAI.

    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:
                - openai_api_key: OpenAI API key
                - openai_api_base: Optional Azure OpenAI endpoint
                - openai_api_deployment_id: Optional Azure deployment ID
                - openai_api_version: Optional Azure API version
                - model: Model name (**required** for non-Azure; not needed for Azure
                    when deployment ID is set)
                - 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 | openai.AzureOpenAI, str | None]:
        """Create the LLM client for OpenAI or Azure OpenAI."""
        openai_api_key = config_dict.get("openai_api_key")
        openai_api_base = config_dict.get("openai_api_base")
        openai_api_deployment_id = config_dict.get("openai_api_deployment_id")
        model = config_dict.get("model")

        if openai_api_base and openai_api_deployment_id:
            # Azure OpenAI
            client = openai.AzureOpenAI(
                api_key=openai_api_key,
                azure_endpoint=openai_api_base,
                api_version=config_dict.get("openai_api_version"),
            )
            return client, openai_api_deployment_id
        else:
            # Regular OpenAI
            client = openai.OpenAI(api_key=openai_api_key)  # type: ignore[assignment]
            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: - openai_api_key: OpenAI API key - openai_api_base: Optional Azure OpenAI endpoint - openai_api_deployment_id: Optional Azure deployment ID - openai_api_version: Optional Azure API version - model: Model name (required for non-Azure; not needed for Azure when deployment ID is set) - save_llm_responses: Whether to save responses (default: True) - temperature: (optional float, default: None)

required
Source code in datarobot_genai/drmcp/test_utils/clients/openai.py
def __init__(
    self,
    config: str | dict,
):
    """
    Initialize the LLM MCP client.

    Args:
        config: Configuration string or dict with:
            - openai_api_key: OpenAI API key
            - openai_api_base: Optional Azure OpenAI endpoint
            - openai_api_deployment_id: Optional Azure deployment ID
            - openai_api_version: Optional Azure API version
            - model: Model name (**required** for non-Azure; not needed for Azure
                when deployment ID is set)
            - save_llm_responses: Whether to save responses (default: True)
            - temperature: (optional float, default: None)
    """
    super().__init__(config)