datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/drmcp/test_utils/clients/dr_gateway.py

@@ -0,0 +1,58 @@
+# Copyright 2026 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""DataRobot LLM Gateway MCP Client implementation."""
+
+import openai
+
+from .base import BaseLLMMCPClient
+
+
+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 (default: "gpt-4o-mini")
+                - save_llm_responses: Whether to save responses (default: True)
+        """
+        super().__init__(config)
+
+    def _create_llm_client(self, config_dict: dict) -> tuple[openai.OpenAI, str]:
+        """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", "gpt-4o-mini")
+
+        # 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

datarobot_genai/drmcp/test_utils/clients/openai.py

@@ -0,0 +1,68 @@
+# Copyright 2026 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""OpenAI LLM MCP Client implementation."""
+
+import openai
+
+from .base import BaseLLMMCPClient
+
+
+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 (default: "gpt-3.5-turbo")
+                - save_llm_responses: Whether to save responses (default: True)
+        """
+        super().__init__(config)
+
+    def _create_llm_client(
+        self, config_dict: dict
+    ) -> tuple[openai.OpenAI | openai.AzureOpenAI, str]:
+        """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", "gpt-3.5-turbo")
+
+        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", "2024-02-15-preview"),
+            )
+            return client, openai_api_deployment_id
+        else:
+            # Regular OpenAI
+            client = openai.OpenAI(api_key=openai_api_key)  # type: ignore[assignment]
+            return client, model

datarobot_genai/drmcp/test_utils/elicitation_test_tool.py

@@ -0,0 +1,89 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test tool for elicitation testing.
+
+This module registers a test tool that can be used to test elicitation support.
+It should be imported in tests that need it.
+"""
+
+from fastmcp import Context
+from fastmcp.server.context import AcceptedElicitation
+from fastmcp.server.context import CancelledElicitation
+from fastmcp.server.context import DeclinedElicitation
+
+from datarobot_genai.drmcp.core.mcp_instance import mcp
+
+
+@mcp.tool(
+    name="get_user_greeting",
+    description=(
+        "Get a personalized greeting for a user. "
+        "Requires a username - if not provided, will request it via elicitation."
+    ),
+    tags={"test", "elicitation"},
+)
+async def get_user_greeting(ctx: Context, username: str | None = None) -> dict:
+    """
+    Get a personalized greeting for a user.
+
+    This tool demonstrates FastMCP's built-in elicitation by requiring a username parameter.
+    If username is not provided, it uses ctx.elicit() to request it from the user.
+
+    Args:
+        ctx: FastMCP context (automatically injected)
+        username: The username to greet. If None, elicitation will be triggered.
+
+    Returns
+    -------
+        Dictionary with greeting message or error if elicitation was declined/cancelled
+    """
+    if not username:
+        # Use elicitation to request username from the client
+        try:
+            result = await ctx.elicit(
+                message="Username is required to generate a personalized greeting",
+                response_type=str,
+            )
+
+            if isinstance(result, AcceptedElicitation):
+                username = result.data
+            elif isinstance(result, DeclinedElicitation):
+                return {
+                    "status": "error",
+                    "error": "Username declined by user",
+                    "message": "Cannot generate greeting without username",
+                }
+            elif isinstance(result, CancelledElicitation):
+                return {
+                    "status": "error",
+                    "error": "Operation cancelled",
+                    "message": "Greeting request was cancelled",
+                }
+        except Exception:
+            # Elicitation not supported by client - return graceful skip
+            return {
+                "status": "skipped",
+                "message": (
+                    "Elicitation not supported by client. "
+                    "Username parameter is required when client does not support elicitation."
+                ),
+                "elicitation_supported": False,
+            }
+
+    return {
+        "status": "success",
+        "message": f"Hello, {username}! Welcome to the DataRobot MCP server.",
+        "username": username,
+    }

datarobot_genai/drmcp/test_utils/integration_mcp_server.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Integration test MCP server.
+
+This server works standalone (base tools only) or detects and loads
+user modules if they exist in the project structure.
+"""
+
+from pathlib import Path
+from typing import Any
+
+from datarobot_genai.drmcp import create_mcp_server
+
+# Import elicitation test tool to register it with the MCP server
+try:
+    from datarobot_genai.drmcp.test_utils import elicitation_test_tool  # noqa: F401
+except ImportError:
+    # Test utils not available (e.g., running in production)
+    pass
+
+# Import user components (will be used conditionally)
+try:
+    from app.core.server_lifecycle import ServerLifecycle  # type: ignore  # noqa: F401
+    from app.core.user_config import get_user_config  # type: ignore  # noqa: F401
+    from app.core.user_credentials import get_user_credentials  # type: ignore  # noqa: F401
+
+except ImportError:
+    # These imports will fail when running from library without user modules
+    pass
+
+
+def detect_user_modules() -> Any:
+    """
+    Detect if user modules exist in the project.
+
+    Returns
+    -------
+        Tuple of (config_factory, credentials_factory, lifecycle, module_paths) or None
+    """
+    # Try to find app directory
+    # When run from library: won't find it
+    # When run from project: will find it
+    current_dir = Path.cwd()
+
+    # Look for app in current directory or parent directories
+    for search_dir in [current_dir, current_dir.parent, current_dir.parent.parent]:
+        app_dir = search_dir / "app"
+        app_core_dir = app_dir / "core"
+        if app_core_dir.exists():
+            # Found user directory - load user modules
+            try:
+                module_paths = [
+                    (str(app_dir / "tools"), "app.tools"),
+                    (str(app_dir / "prompts"), "app.prompts"),
+                    (str(app_dir / "resources"), "app.resources"),
+                ]
+
+                return (
+                    get_user_config,
+                    get_user_credentials,
+                    ServerLifecycle(),
+                    module_paths,
+                )
+            except ImportError:
+                # User modules don't exist or can't be imported
+                pass
+
+    return None
+
+
+def main() -> None:
+    """Run the integration test MCP server."""
+    # Try to detect and load user modules
+    user_components = detect_user_modules()
+
+    if user_components:
+        # User modules found - create server with user extensions
+        config_factory, credentials_factory, lifecycle, module_paths = user_components
+        server = create_mcp_server(
+            config_factory=config_factory,
+            credentials_factory=credentials_factory,
+            lifecycle=lifecycle,
+            additional_module_paths=module_paths,
+            transport="stdio",
+        )
+    else:
+        # No user modules - create server with base tools only
+        server = create_mcp_server(transport="stdio")
+
+    server.run()
+
+
+if __name__ == "__main__":
+    main()

datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@@ -0,0 +1,133 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import asyncio
+import os
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import Any
+
+import aiohttp
+from aiohttp import ClientSession as HttpClientSession
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+
+from .utils import load_env
+
+load_env()
+
+
+def get_dr_mcp_server_url() -> str | None:
+    """Get DataRobot MCP server URL."""
+    return os.environ.get("DR_MCP_SERVER_URL")
+
+
+def get_dr_mcp_server_http_url() -> str | None:
+    """Get DataRobot MCP server http URL."""
+    return os.environ.get("DR_MCP_SERVER_HTTP_URL")
+
+
+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")
+
+    config: dict[str, str] = {
+        "openai_api_key": openai_api_key,
+    }
+
+    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)
+
+    return config
+
+
+def get_headers() -> dict[str, str]:
+    # When the MCP server is deployed in DataRobot, we have to include the API token in headers for
+    # authentication.
+    api_token = os.getenv("DATAROBOT_API_TOKEN")
+    headers = {"Authorization": f"Bearer {api_token}"}
+    return headers
+
+
+@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 asyncio.TimeoutError:
+        raise TimeoutError(f"Check if the MCP server is running at {get_dr_mcp_server_url()}")
+
+
+@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

datarobot_genai/drmcp/test_utils/mcp_utils_integration.py

@@ -0,0 +1,107 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import contextlib
+import os
+from collections.abc import AsyncGenerator
+from pathlib import Path
+from typing import Any
+
+from mcp import ClientSession
+from mcp.client.stdio import StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from .utils import load_env
+
+load_env()
+
+
+def integration_test_mcp_server_params() -> StdioServerParameters:
+    env = {
+        "DATAROBOT_API_TOKEN": os.environ.get("DATAROBOT_API_TOKEN") or "test-token",
+        "DATAROBOT_ENDPOINT": os.environ.get("DATAROBOT_ENDPOINT")
+        or "https://test.datarobot.com/api/v2",
+        "MCP_SERVER_LOG_LEVEL": os.environ.get("MCP_SERVER_LOG_LEVEL") or "WARNING",
+        "APP_LOG_LEVEL": os.environ.get("APP_LOG_LEVEL") or "WARNING",
+        # Disable all OTEL telemetry for integration tests
+        "OTEL_ENABLED": "false",
+        "OTEL_SDK_DISABLED": "true",
+        "OTEL_TRACES_EXPORTER": "none",
+        "OTEL_LOGS_EXPORTER": "none",
+        "OTEL_METRICS_EXPORTER": "none",
+        "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP": os.environ.get(
+            "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP"
+        )
+        or "false",
+        "MCP_SERVER_REGISTER_DYNAMIC_PROMPTS_ON_STARTUP": os.environ.get(
+            "MCP_SERVER_REGISTER_DYNAMIC_PROMPTS_ON_STARTUP"
+        )
+        or "true",
+    }
+
+    script_dir = Path(__file__).resolve().parent
+    server_script = str(script_dir / "integration_mcp_server.py")
+    # Add src/ directory to Python path so datarobot_genai can be imported
+    src_dir = script_dir.parent.parent.parent
+
+    return StdioServerParameters(
+        command="uv",
+        args=["run", server_script],
+        env={
+            "PYTHONPATH": str(src_dir),
+            "MCP_SERVER_NAME": "integration",
+            "MCP_SERVER_PORT": "8081",
+            **env,
+        },
+    )
+
+
+@contextlib.asynccontextmanager
+async def integration_test_mcp_session(
+    server_params: StdioServerParameters | None = None,
+    timeout: int = 30,
+    elicitation_callback: Any | None = None,
+) -> AsyncGenerator[ClientSession, None]:
+    """
+    Create and connect a client for the MCP server as a context manager.
+
+    Args:
+        server_params: Parameters for configuring the server connection
+        timeout: Timeout
+        elicitation_callback: Optional callback for handling elicitation requests
+
+    Yields
+    ------
+        ClientSession: Connected MCP client session
+
+    Raises
+    ------
+        ConnectionError: If session initialization fails
+        TimeoutError: If session initialization exceeds timeout
+    """
+    server_params = server_params or integration_test_mcp_server_params()
+
+    try:
+        async with stdio_client(server_params) as (read_stream, write_stream):
+            async with ClientSession(
+                read_stream, write_stream, elicitation_callback=elicitation_callback
+            ) as session:
+                init_result = await asyncio.wait_for(session.initialize(), timeout=timeout)
+                # Store the init result on the session for tests that need to inspect capabilities
+                session._init_result = init_result  # type: ignore[attr-defined]
+                yield session
+
+    except asyncio.TimeoutError:
+        raise TimeoutError(f"Session initialization timed out after {timeout} seconds")