datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/crewai/mcp.py

@@ -0,0 +1,59 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# 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.
+
+"""
+MCP integration for CrewAI using MCPServerAdapter.
+
+This module provides MCP server connection management for CrewAI agents.
+"""
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+from crewai_tools import MCPServerAdapter
+
+from datarobot_genai.core.mcp.common import MCPConfig
+
+
+@contextmanager
+def mcp_tools_context(
+    authorization_context: dict[str, Any] | None = None,
+    forwarded_headers: dict[str, str] | None = None,
+) -> Generator[list[Any], None, None]:
+    """Context manager for MCP tools that handles connection lifecycle."""
+    config = MCPConfig(
+        authorization_context=authorization_context,
+        forwarded_headers=forwarded_headers,
+    )
+    # If no MCP server configured, return empty tools list
+    if not config.server_config:
+        print("No MCP server configured, using empty tools list", flush=True)
+        yield []
+        return
+
+    print(f"Connecting to MCP server: {config.server_config['url']}", flush=True)
+
+    # Use MCPServerAdapter as context manager with the server config
+    try:
+        with MCPServerAdapter(config.server_config) as tools:
+            print(
+                f"Successfully connected to MCP server, got {len(tools)} tools",
+                flush=True,
+            )
+            yield tools
+    except Exception as exc:
+        # Gracefully degrade when connection fails or adapter initialization raises
+        print(f"Failed to connect to MCP server: {exc}", flush=True)
+        yield []

datarobot_genai/drmcp/__init__.py

@@ -0,0 +1,78 @@
+# 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.
+
+"""
+DataRobot MCP Server Library.
+
+A reusable library for building Model Context Protocol (MCP) servers with DataRobot integration.
+"""
+
+# Export main server components
+from datarobot_genai.drmcp.test_utils.clients.openai import OpenAILLMMCPClient
+from datarobot_genai.drmcp.test_utils.mcp_utils_ete import ete_test_mcp_session
+from datarobot_genai.drmcp.test_utils.mcp_utils_ete import get_dr_mcp_server_url
+from datarobot_genai.drmcp.test_utils.mcp_utils_ete import get_headers
+from datarobot_genai.drmcp.test_utils.mcp_utils_integration import integration_test_mcp_session
+from datarobot_genai.drmcp.test_utils.tool_base_ete import ETETestExpectations
+from datarobot_genai.drmcp.test_utils.tool_base_ete import ToolBaseE2E
+from datarobot_genai.drmcp.test_utils.tool_base_ete import ToolCallTestExpectations
+
+from .core.clients import get_sdk_client
+from .core.config import MCPServerConfig
+from .core.config import get_config
+from .core.config_utils import extract_datarobot_credential_runtime_param_payload
+from .core.config_utils import extract_datarobot_dict_runtime_param_payload
+from .core.config_utils import extract_datarobot_runtime_param_payload
+from .core.constants import RUNTIME_PARAM_ENV_VAR_NAME_PREFIX
+from .core.credentials import MCPServerCredentials
+from .core.credentials import get_credentials
+from .core.dr_mcp_server import BaseServerLifecycle
+from .core.dr_mcp_server import DataRobotMCPServer
+from .core.dr_mcp_server import create_mcp_server
+from .core.logging import MCPLogging
+from .core.mcp_instance import dr_mcp_tool
+from .core.mcp_instance import register_tools
+
+__all__ = [
+    # Main server
+    "DataRobotMCPServer",
+    "create_mcp_server",
+    "BaseServerLifecycle",
+    # Configuration
+    "get_config",
+    "MCPServerConfig",
+    # Credentials
+    "get_credentials",
+    "MCPServerCredentials",
+    # Constants
+    "RUNTIME_PARAM_ENV_VAR_NAME_PREFIX",
+    # User extensibility
+    "get_sdk_client",
+    "dr_mcp_tool",
+    "register_tools",
+    # Utilities
+    "MCPLogging",
+    "extract_datarobot_runtime_param_payload",
+    "extract_datarobot_dict_runtime_param_payload",
+    "extract_datarobot_credential_runtime_param_payload",
+    # Test utilities
+    "get_dr_mcp_server_url",
+    "get_headers",
+    "ete_test_mcp_session",
+    "OpenAILLMMCPClient",
+    "ETETestExpectations",
+    "ToolBaseE2E",
+    "ToolCallTestExpectations",
+    "integration_test_mcp_session",
+]

datarobot_genai/drmcp/core/__init__.py

@@ -0,0 +1,13 @@
+# 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.

datarobot_genai/drmcp/core/auth.py

@@ -0,0 +1,165 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# 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.
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from datarobot.auth.session import AuthCtx
+from datarobot.models.genai.agent.auth import ToolAuth
+from fastmcp.server.dependencies import get_context
+from fastmcp.server.dependencies import get_http_headers
+from fastmcp.server.middleware import CallNext
+from fastmcp.server.middleware import Middleware
+from fastmcp.server.middleware import MiddlewareContext
+from fastmcp.tools.tool import ToolResult
+
+from datarobot_genai.core.utils.auth import AsyncOAuthTokenProvider
+from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
+
+logger = logging.getLogger(__name__)
+
+
+AUTH_CTX_KEY = "authorization_context"
+
+
+class OAuthMiddleWare(Middleware):
+    """Middleware that parses `x-datarobot-authorization-context` for tool calls.
+
+    The header is expected to be a JWT-encoded token representing an authentication
+    context compatible with :class:`datarobot.auth.session.AuthCtx`.
+
+    Attributes
+    ----------
+    auth_handler : AuthContextHeaderHandler
+        Handler for encoding/decoding JWT tokens containing auth context.
+    """
+
+    def __init__(self, auth_handler: AuthContextHeaderHandler | None = None) -> None:
+        self.auth_handler = auth_handler or AuthContextHeaderHandler()
+
+    async def on_call_tool(
+        self, context: MiddlewareContext, call_next: CallNext[Any, ToolResult]
+    ) -> ToolResult:
+        """Parse header and attach an AuthCtx to the context before running the tool.
+
+        Parameters
+        ----------
+        context : MiddlewareContext
+            The middleware context that will be passed to the tool.
+        call_next : CallNext[Any, ToolResult]
+            The next handler in the middleware chain.
+
+        Returns
+        -------
+        ToolResult
+            The result from the tool execution.
+        """
+        auth_context = self._extract_auth_context()
+        if not auth_context:
+            logger.debug("No valid authorization context extracted from request headers.")
+
+        if context.fastmcp_context is not None:
+            context.fastmcp_context.set_state(AUTH_CTX_KEY, auth_context)
+            logger.debug("Authorization context attached to state.")
+
+        return await call_next(context)
+
+    def _extract_auth_context(self) -> AuthCtx | None:
+        """Extract and validate authentication context from request headers.
+
+        Returns
+        -------
+        Optional[AuthCtx]
+            The validated authentication context, or None if extraction fails.
+        """
+        try:
+            headers = get_http_headers()
+            return self.auth_handler.get_context(headers)
+        except (ValueError, KeyError, TypeError) as exc:
+            logger.warning("Failed to extract auth context from headers: %s", exc, exc_info=True)
+            return None
+        except Exception as exc:
+            logger.error("Unexpected error extracting auth context: %s", exc, exc_info=True)
+            return None
+
+
+async def must_get_auth_context() -> AuthCtx:
+    """Retrieve the AuthCtx from the current request context or raise error.
+
+    Raises
+    ------
+    RuntimeError
+        If no authorization context is found in the request.
+
+    Returns
+    -------
+    AuthCtx
+        The authorization context associated with the current request.
+    """
+    context = get_context()
+
+    auth_ctx = context.get_state(AUTH_CTX_KEY)
+    if not auth_ctx:
+        raise RuntimeError("Could not retrieve authorization context from FastMCP context state.")
+
+    return auth_ctx
+
+
+async def get_access_token(provider_type: str | None = None) -> str:
+    """Retrieve access token from the DataRobot OAuth Provider Service.
+
+    OAuth access tokens can be retrieved only for providers where the user completed
+    the OAuth flow and granted consent.
+
+    Note:
+        *   Currently, only On-Behalf-Of (OBO) tokens are supported, which allow tools to
+            act on behalf of the authenticated user, after the user has granted his consent.
+
+    Parameters
+    ----------
+    provider_type : str, optional
+        The name of the OAuth provider. It should match the name of the provider configured
+        during provider setup. If no value is provided and only one OAuth provider exists, that
+        provider will be used. If multiple providers exist and none is specified, an error will be
+        raised.
+
+    Returns
+    -------
+    The oauth access token.
+    """
+    auth_ctx = await must_get_auth_context()
+    logger.debug("Retrieved authorization context")
+
+    oauth_token_provider = AsyncOAuthTokenProvider(auth_ctx)
+    oauth_access_token = await oauth_token_provider.get_token(
+        auth_type=ToolAuth.OBO,
+        provider_type=provider_type,
+    )
+    return oauth_access_token
+
+
+def initialize_oauth_middleware(mcp: Any) -> None:
+    """Initialize and register OAuth middleware with the MCP server.
+
+    Parameters
+    ----------
+    mcp : FastMCP
+        The FastMCP server instance to register the middleware with.
+    secret_key : Optional[str]
+        Secret key for JWT validation. If None, uses the value from config.
+    """
+    mcp.add_middleware(OAuthMiddleWare())
+    logger.info("OAuth middleware registered successfully")

datarobot_genai/drmcp/core/clients.py

@@ -0,0 +1,180 @@
+# 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 logging
+from typing import Any
+from typing import cast
+
+import datarobot as dr
+from datarobot.context import Context as DRContext
+from datarobot.rest import RESTClientObject
+from fastmcp.server.dependencies import get_http_headers
+
+from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
+from datarobot_genai.core.utils.auth import DRAppCtx
+
+from .credentials import get_credentials
+
+logger = logging.getLogger(__name__)
+
+# Header names to check for authorization tokens (in order of preference)
+HEADER_TOKEN_CANDIDATE_NAMES = [
+    "authorization",
+    "x-datarobot-api-token",
+    "x-datarobot-api-key",
+]
+
+
+def _extract_token_from_headers(headers: dict[str, str]) -> str | None:
+    """
+    Extract a Bearer token from headers by checking multiple header name candidates.
+
+    Args:
+        headers: Dictionary of headers (keys should be lowercase)
+
+    Returns
+    -------
+        The extracted token string, or None if not found
+    """
+    for candidate_name in HEADER_TOKEN_CANDIDATE_NAMES:
+        auth_header = headers.get(candidate_name)
+        if not auth_header:
+            continue
+
+        if not isinstance(auth_header, str):
+            continue
+
+        # Handle Bearer token format
+        bearer_prefix = "bearer "
+        if auth_header.lower().startswith(bearer_prefix):
+            token = auth_header[len(bearer_prefix) :].strip()
+        else:
+            # Assume it's a plain token
+            token = auth_header.strip()
+
+        if token:
+            return token
+
+    return None
+
+
+def _extract_token_from_auth_context(headers: dict[str, str]) -> str | None:
+    """
+    Extract API token from authorization context metadata as a fallback.
+
+    Args:
+        headers: Dictionary of headers (keys should be lowercase)
+
+    Returns
+    -------
+        The extracted API key from auth context metadata, or None if not found
+    """
+    try:
+        auth_handler = AuthContextHeaderHandler()
+
+        auth_ctx = auth_handler.get_context(headers)
+        if not auth_ctx or not auth_ctx.metadata:
+            return None
+
+        metadata = auth_ctx.metadata
+        if not isinstance(metadata, dict):
+            return None
+
+        dr_ctx: DRAppCtx = DRAppCtx(**metadata.get("dr_ctx", {}))
+        if dr_ctx.api_key:
+            logger.debug("Extracted token from auth context")
+            return dr_ctx.api_key
+
+        return None
+
+    except Exception as e:
+        logger.debug(f"Failed to get token from auth context: {e}")
+        return None
+
+
+def extract_token_from_headers(headers: dict[str, str]) -> str | None:
+    """
+    Extract a token from headers with multiple fallback strategies.
+
+    This function attempts to extract a token in the following order:
+    1. From standard authorization headers (Bearer token, x-datarobot-api-token, etc.)
+    2. From authorization context metadata (dr_ctx.api_key)
+
+    Args:
+        headers: Dictionary of headers (keys should be lowercase)
+
+    Returns
+    -------
+        The extracted token string, or None if not found
+    """
+    if token := _extract_token_from_headers(headers):
+        return token
+
+    if token := _extract_token_from_auth_context(headers):
+        return token
+
+    return None
+
+
+def get_sdk_client() -> Any:
+    """
+    Get a DataRobot SDK client, using the user's Bearer token from the request.
+
+    This function attempts to extract the Bearer token from the HTTP request headers
+    with fallback strategies:
+    1. Standard authorization headers (Bearer token, x-datarobot-api-token, etc.)
+    2. Authorization context metadata (dr_ctx.api_key)
+    3. Application credentials as final fallback
+    """
+    token = None
+
+    try:
+        headers = get_http_headers()
+        if headers:
+            token = extract_token_from_headers(headers)
+            if token:
+                logger.debug("Using API token found in HTTP headers")
+    except Exception:
+        # No HTTP context e.g. stdio transport
+        logger.warning(
+            "Could not get HTTP headers, falling back to application credentials", exc_info=True
+        )
+
+    credentials = get_credentials()
+
+    # Fallback: Use application token
+    if not token:
+        token = credentials.datarobot.application_api_token
+        logger.debug("Using application API token from credentials")
+
+    dr.Client(token=token, endpoint=credentials.datarobot.endpoint)
+    # The trafaret setting up a use case in the context, seem to mess up the tool calls
+    DRContext.use_case = None
+    return dr
+
+
+def get_api_client() -> RESTClientObject:
+    """Get a DataRobot SDK api client using application credentials."""
+    dr = get_sdk_client()
+
+    return cast(RESTClientObject, dr.client.get_client())
+
+
+def get_s3_bucket_info() -> dict[str, str]:
+    """Get S3 bucket configuration."""
+    credentials = get_credentials()
+    return {
+        "bucket": credentials.aws_predictions_s3_bucket,
+        "prefix": credentials.aws_predictions_s3_prefix,
+    }