datarobot_genai.drmcp.test_utils.mcp_utils_ete

mcp_utils_ete

get_dr_mcp_server_url

get_dr_mcp_server_url() -> str | None

Get DataRobot MCP server URL.

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

def get_dr_mcp_server_url() -> str | None:
    """Get DataRobot MCP server URL."""
    return os.environ.get("DR_MCP_SERVER_URL")

get_dr_mcp_server_http_url

get_dr_mcp_server_http_url() -> str | None

Get DataRobot MCP server http URL.

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

def get_dr_mcp_server_http_url() -> str | None:
    """Get DataRobot MCP server http URL."""
    return os.environ.get("DR_MCP_SERVER_HTTP_URL")

get_openai_llm_client_config

get_openai_llm_client_config() -> dict[str, str]

Get OpenAI LLM client configuration.

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

def get_openai_llm_client_config() -> dict[str, str]:
    """Get OpenAI LLM client configuration."""
    openai_api_key = os.environ.get("OPENAI_API_KEY")
    openai_api_base = os.environ.get("OPENAI_API_BASE")
    openai_api_deployment_id = os.environ.get("OPENAI_API_DEPLOYMENT_ID")
    openai_api_version = os.environ.get("OPENAI_API_VERSION")
    save_llm_responses = os.environ.get("SAVE_LLM_RESPONSES", "false").lower() == "true"

    # Check for OpenAI configuration
    if not openai_api_key:
        raise ValueError("Missing required environment variable: OPENAI_API_KEY")
    if (
        openai_api_base and not openai_api_deployment_id
    ):  # For Azure OpenAI, we need additional variables
        raise ValueError("Missing required environment variable: OPENAI_API_DEPLOYMENT_ID")

    openai_model = os.environ.get("OPENAI_MODEL")
    if not openai_api_deployment_id and not openai_model:
        raise ValueError("Missing required environment variable: OPENAI_MODEL")

    config: dict[str, str] = {
        "openai_api_key": openai_api_key,
    }

    if openai_model:
        config["model"] = openai_model
    if openai_api_base:
        config["openai_api_base"] = openai_api_base
    if openai_api_deployment_id:
        config["openai_api_deployment_id"] = openai_api_deployment_id
    if openai_api_version:
        config["openai_api_version"] = openai_api_version
    config["save_llm_responses"] = str(save_llm_responses)

    llm_temperature = os.environ.get("LLM_TEMPERATURE")
    if llm_temperature is not None:
        config["temperature"] = llm_temperature

    return config

get_dr_llm_gateway_client_config

get_dr_llm_gateway_client_config() -> dict[str, str]

Get DataRobot LLM Gateway client configuration.

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

def get_dr_llm_gateway_client_config() -> dict[str, str]:
    """Get DataRobot LLM Gateway client configuration."""
    datarobot_api_token = os.environ.get("DATAROBOT_API_TOKEN")
    datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT")
    save_llm_responses = os.environ.get("SAVE_LLM_RESPONSES", "false").lower() == "true"

    if not datarobot_api_token:
        raise ValueError("Missing required environment variable: DATAROBOT_API_TOKEN")

    dr_llm_gateway_model = os.environ.get("DR_LLM_GATEWAY_MODEL")
    if not dr_llm_gateway_model:
        raise ValueError("Missing required environment variable: DR_LLM_GATEWAY_MODEL")

    config: dict[str, str] = {
        "datarobot_api_token": datarobot_api_token,
        "save_llm_responses": str(save_llm_responses),
    }

    config["model"] = dr_llm_gateway_model
    if datarobot_endpoint:
        config["datarobot_endpoint"] = datarobot_endpoint

    llm_temperature = os.environ.get("LLM_TEMPERATURE")
    if llm_temperature is not None:
        config["temperature"] = llm_temperature

    return config

ete_test_mcp_session async

ete_test_mcp_session(additional_headers: dict[str, str] | None = None, elicitation_callback: Any | None = None) -> AsyncGenerator[ClientSession, None]

Create an MCP session for each test.

Parameters

additional_headers : dict[str, str], optional Additional headers to include in the MCP session (e.g., auth headers for testing). elicitation_callback : callable, optional Callback function to handle elicitation requests from the server. The callback should have signature: async def callback(context, params: ElicitRequestParams) -> ElicitResult

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@asynccontextmanager
async def ete_test_mcp_session(
    additional_headers: dict[str, str] | None = None,
    elicitation_callback: Any | None = None,
) -> AsyncGenerator[ClientSession, None]:
    """Create an MCP session for each test.

    Parameters
    ----------
    additional_headers : dict[str, str], optional
        Additional headers to include in the MCP session (e.g., auth headers for testing).
    elicitation_callback : callable, optional
        Callback function to handle elicitation requests from the server.
        The callback should have signature:
        async def callback(context, params: ElicitRequestParams) -> ElicitResult
    """
    try:
        headers = get_headers()
        if additional_headers:
            headers.update(additional_headers)

        async with streamablehttp_client(url=get_dr_mcp_server_url(), headers=headers) as (
            read_stream,
            write_stream,
            _,
        ):
            async with ClientSession(
                read_stream, write_stream, elicitation_callback=elicitation_callback
            ) as session:
                await asyncio.wait_for(session.initialize(), timeout=5)
                yield session
    except TimeoutError:
        raise TimeoutError(f"Check if the MCP server is running at {get_dr_mcp_server_url()}")

ete_test_http_session async

ete_test_http_session(additional_headers: dict[str, str] | None = None) -> AsyncGenerator[HttpClientSession, None]

Create an HTTP session for each test that can connect to MCP custom http routes.

Parameters

additional_headers : dict[str, str], optional Additional headers to include in the HTTP session (e.g., auth headers for testing).

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@asynccontextmanager
async def ete_test_http_session(
    additional_headers: dict[str, str] | None = None,
) -> AsyncGenerator[HttpClientSession, None]:
    """Create an HTTP session for each test that can connect to MCP custom http routes.

    Parameters
    ----------
    additional_headers : dict[str, str], optional
        Additional headers to include in the HTTP session (e.g., auth headers for testing).
    """
    headers = get_headers()
    if additional_headers:
        headers.update(additional_headers)

    async with ete_test_mcp_session(additional_headers=additional_headers):
        async with aiohttp.ClientSession(
            base_url=get_dr_mcp_server_http_url(), headers=headers
        ) as client:
            yield client