datarobot-genai 0.2.0__py3-none-any.whl

datarobot_genai/drmcp/core/config.py

@@ -0,0 +1,250 @@
+# 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 os
+from typing import Any
+from typing import Literal
+from urllib.parse import urlparse
+from urllib.parse import urlunparse
+
+from fastmcp.settings import DuplicateBehavior
+from pydantic import AliasChoices
+from pydantic import Field
+from pydantic import field_validator
+from pydantic_settings import BaseSettings
+from pydantic_settings import SettingsConfigDict
+
+from .config_utils import extract_datarobot_dict_runtime_param_payload
+from .config_utils import extract_datarobot_runtime_param_payload
+from .constants import DEFAULT_DATAROBOT_ENDPOINT
+from .constants import RUNTIME_PARAM_ENV_VAR_NAME_PREFIX
+
+
+class MCPServerConfig(BaseSettings):
+    """MCP Server configuration using pydantic settings."""
+
+    mcp_server_name: str = Field(
+        default="datarobot-mcp-server",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_NAME",
+            "MCP_SERVER_NAME",
+        ),
+        description="Name of the MCP server",
+    )
+    mcp_server_port: int = Field(
+        default=8080,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_PORT",
+            "MCP_SERVER_PORT",
+        ),
+        description="Port number for the MCP server",
+    )
+    mcp_server_log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = Field(
+        default="WARNING",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_LOG_LEVEL",
+            "MCP_SERVER_LOG_LEVEL",
+        ),
+        description="Log level for the MCP server",
+    )
+    mcp_server_host: str = Field(
+        default="0.0.0.0",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_HOST",
+            "MCP_SERVER_HOST",
+        ),
+        description="Host address for the MCP server",
+    )
+    app_log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = Field(
+        default="INFO",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "APP_LOG_LEVEL",
+            "APP_LOG_LEVEL",
+        ),
+        description="App log level",
+    )
+    # When the server is run in a custom model, it is important to mount all routes under the
+    # prefix provided in the URL_PREFIX
+    mount_path: str = Field(default="/", alias="URL_PREFIX")
+
+    @staticmethod
+    def _get_default_otel_endpoint() -> str:
+        """Get the default OpenTelemetry endpoint e.g. https://app.datarobot.com/otel."""
+        parsed_url = urlparse(os.environ.get("DATAROBOT_ENDPOINT", DEFAULT_DATAROBOT_ENDPOINT))
+        stripped_url = (parsed_url.scheme, parsed_url.netloc, "otel", "", "", "")
+        return urlunparse(stripped_url)
+
+    otel_collector_base_url: str = Field(
+        default=_get_default_otel_endpoint(),
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "OTEL_COLLECTOR_BASE_URL",
+            "OTEL_COLLECTOR_BASE_URL",
+        ),
+        description="Base URL for the OpenTelemetry collector",
+    )
+    otel_entity_id: str = Field(
+        default="",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "OTEL_ENTITY_ID",
+            "OTEL_ENTITY_ID",
+        ),
+        description="Entity ID for tracing",
+    )
+    otel_attributes: dict[str, Any] = Field(
+        default={},
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "OTEL_ATTRIBUTES",
+            "OTEL_ATTRIBUTES",
+        ),
+        description="Attributes for tracing (as JSON string)",
+    )
+    otel_enabled: bool = Field(
+        default=True,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "OTEL_ENABLED",
+            "OTEL_ENABLED",
+        ),
+        description="Enable/disable OpenTelemetry",
+    )
+    otel_enabled_http_instrumentors: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "OTEL_ENABLED_HTTP_INSTRUMENTORS",
+            "OTEL_ENABLED_HTTP_INSTRUMENTORS",
+        ),
+        description="Enable/disable HTTP instrumentors",
+    )
+    mcp_server_register_dynamic_tools_on_startup: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP",
+            "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP",
+        ),
+        description="Register dynamic tools on startup. When enabled, the MCP server will "
+        "automatically register all DataRobot tool deployments as MCP tools during startup.",
+    )
+    tool_registration_allow_empty_schema: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_TOOL_REGISTRATION_ALLOW_EMPTY_SCHEMA",
+            "MCP_SERVER_TOOL_REGISTRATION_ALLOW_EMPTY_SCHEMA",
+        ),
+        description="Allow registration of tools with no input parameters. When enabled, "
+        "tools can be registered with empty schemas for static endpoints that don't require any "
+        "inputs. "
+        "Disabled by default, as this is not typical use case and can hide potential issues with "
+        "schema.",
+    )
+    tool_registration_duplicate_behavior: DuplicateBehavior = Field(
+        default="warn",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_TOOL_REGISTRATION_DUPLICATE_BEHAVIOR",
+            "MCP_SERVER_TOOL_REGISTRATION_DUPLICATE_BEHAVIOR",
+        ),
+        description="Behavior when a tool with the same name already exists in the MCP server. "
+        " - 'warn': will log a warning and replace the existing tool. "
+        " - 'replace': will replace the existing tool without a warning. "
+        " - 'error': will raise an error and prevent registration. "
+        " - 'ignore': will skip registration of the new tool.",
+    )
+    mcp_server_register_dynamic_prompts_on_startup: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_REGISTER_DYNAMIC_PROMPTS_ON_STARTUP",
+            "MCP_SERVER_REGISTER_DYNAMIC_PROMPTS_ON_STARTUP",
+        ),
+        description="Register dynamic prompts on startup. When enabled, the MCP server will "
+        "automatically register all prompts from DataRobot Prompt Management "
+        "as MCP prompts during startup.",
+    )
+    prompt_registration_duplicate_behavior: DuplicateBehavior = Field(
+        default="warn",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "MCP_SERVER_PROMPT_REGISTRATION_DUPLICATE_BEHAVIOR",
+            "MCP_SERVER_PROMPT_REGISTRATION_DUPLICATE_BEHAVIOR",
+        ),
+        description="Behavior when a prompt with the same name already exists in the MCP server. "
+        " - 'warn': will log a warning and replace the existing tool. "
+        " - 'replace': will replace the existing tool without a warning. "
+        " - 'error': will raise an error and prevent registration. "
+        " - 'ignore': will skip registration of the new tool.",
+    )
+    enable_memory_management: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_MEMORY_MANAGEMENT",
+            "ENABLE_MEMORY_MANAGEMENT",
+        ),
+        description="Enable/disable memory management",
+    )
+    enable_predictive_tools: bool = Field(
+        default=True,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
+            "ENABLE_PREDICTIVE_TOOLS",
+        ),
+        description="Enable/disable predictive tools",
+    )
+
+    @field_validator(
+        "otel_attributes",
+        mode="before",
+    )
+    @classmethod
+    def validate_dict_runtime_params(cls, v: Any) -> Any:
+        """Validate dict runtime parameters."""
+        return extract_datarobot_dict_runtime_param_payload(v)
+
+    @field_validator(
+        "mcp_server_name",
+        "mcp_server_log_level",
+        "app_log_level",
+        "otel_collector_base_url",
+        "otel_entity_id",
+        "otel_enabled",
+        "otel_enabled_http_instrumentors",
+        "enable_memory_management",
+        "tool_registration_allow_empty_schema",
+        "mcp_server_register_dynamic_tools_on_startup",
+        "tool_registration_duplicate_behavior",
+        "mcp_server_register_dynamic_prompts_on_startup",
+        "enable_predictive_tools",
+        mode="before",
+    )
+    @classmethod
+    def validate_runtime_params(cls, v: Any) -> Any:
+        """Validate runtime parameters."""
+        return extract_datarobot_runtime_param_payload(v)
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        case_sensitive=False,
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+
+# Global configuration instance
+_config: MCPServerConfig | None = None
+
+
+def get_config() -> MCPServerConfig:
+    """Get the global configuration instance."""
+    # Use a local variable to avoid global statement warning
+    config = _config
+    if config is None:
+        config = MCPServerConfig()
+        # Update the global variable
+        globals()["_config"] = config
+    return config

datarobot_genai/drmcp/core/config_utils.py

@@ -0,0 +1,174 @@
+# 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 json
+from typing import Any
+
+from pydantic_core import PydanticUseDefault
+
+
+def extract_datarobot_runtime_param_payload(v: Any) -> Any:
+    """Extract payload from DataRobot runtime parameter JSON format.
+
+    DataRobot runtime parameters come in the format:
+    {"type":"string","payload":"value"} or {"type":"boolean","payload":false}
+
+    If payload is None, raises PydanticUseDefault so Pydantic uses the field default.
+
+    This function extracts the payload value for simple types (strings, booleans, numbers).
+    For dict types with nested JSON, use extract_datarobot_dict_runtime_param_payload instead.
+
+    Args:
+        v: The input value (may be a raw value, JSON string, or DataRobot runtime param format)
+
+    Returns
+    -------
+        The extracted payload value
+
+    Raises
+    ------
+        PydanticUseDefault: When payload is None, signaling Pydantic to use field default
+    """
+    # If it's a string, try to parse as JSON
+    if isinstance(v, str):
+        # Handle Python-style boolean strings (True/False) by converting to lowercase
+        v_normalized = v.lower() if v.lower() in ("true", "false") else v
+
+        try:
+            parsed = json.loads(v_normalized)
+            if isinstance(parsed, dict) and "payload" in parsed:
+                payload = parsed["payload"]
+                if payload is not None:
+                    return payload
+                # If payload is None, use field default
+                raise PydanticUseDefault
+            # If it's a plain JSON value (boolean, number, etc.), return it
+            # This handles cases like "true", "false", "123", etc.
+            return parsed
+        except (json.JSONDecodeError, ValueError):
+            pass
+    return v
+
+
+def extract_datarobot_dict_runtime_param_payload(v: Any) -> Any:  # noqa: PLR0911
+    r"""Extract and parse dict from DataRobot runtime parameter JSON format.
+
+    DataRobot runtime parameters for dict fields come in the format:
+    {"type":"string","payload":"{\\"key\\":\\"value\\"}"}
+
+    The payload itself is a JSON string that needs to be parsed into a dict.
+    If payload is None, raises PydanticUseDefault so Pydantic uses the field default.
+
+    Args:
+        v: The input value (may be a dict, JSON string, or DataRobot runtime param format)
+
+    Returns
+    -------
+        The extracted and parsed dict
+
+    Raises
+    ------
+        PydanticUseDefault: When payload is None, signaling Pydantic to use field default
+    """
+    # If it's already a dict, check if it's in DataRobot format
+    if isinstance(v, dict):
+        # If it has "payload" key, it's in DataRobot format - extract it
+        if "payload" in v:
+            payload = v["payload"]
+            if payload is None:
+                raise PydanticUseDefault
+            # If payload is a string, parse it as JSON
+            if isinstance(payload, str):
+                try:
+                    return json.loads(payload)
+                except (json.JSONDecodeError, ValueError):
+                    return {}
+            return payload
+        # Otherwise, return as-is (already a plain dict)
+        return v
+
+    # If it's a string, try to parse as JSON
+    if isinstance(v, str):
+        try:
+            parsed = json.loads(v)
+            if isinstance(parsed, dict) and "payload" in parsed:
+                payload = parsed["payload"]
+                if payload is not None:
+                    # If payload is a string, parse it as JSON
+                    if isinstance(payload, str):
+                        try:
+                            return json.loads(payload)
+                        except (json.JSONDecodeError, ValueError):
+                            return {}
+                    return payload
+                raise PydanticUseDefault
+            # If it's already a dict, return it
+            if isinstance(parsed, dict):
+                return parsed
+        except (json.JSONDecodeError, ValueError):
+            pass
+
+    # Fallback: return empty dict
+    return {}
+
+
+def extract_datarobot_credential_runtime_param_payload(v: Any) -> Any:
+    """Extract credential payload from DataRobot runtime parameter JSON format.
+
+    DataRobot credential runtime parameters come in the format:
+    {"type":"credential","payload":{...}} where the payload is the full credential object.
+
+    For credential types, the entire dict is preserved.
+    If payload is None, raises PydanticUseDefault so Pydantic uses the field default.
+
+    Args:
+        v: The input value (may be a dict, JSON string, or DataRobot runtime param format)
+
+    Returns
+    -------
+        The extracted credential dict
+
+    Raises
+    ------
+        PydanticUseDefault: When payload is None, signaling Pydantic to use field default
+    """
+    # If it's already a dict (credential object or wrapped payload)
+    if isinstance(v, dict):
+        # If it's a wrapped payload, extract it
+        if "payload" in v:
+            payload = v["payload"]
+            if payload is not None:
+                return payload
+            # If payload is None, use field default
+            raise PydanticUseDefault
+        # Otherwise return the dict as-is (it's the credential object)
+        return v
+
+    # If it's a string, try to parse as JSON
+    if isinstance(v, str):
+        try:
+            parsed = json.loads(v)
+            if isinstance(parsed, dict):
+                # If it's a wrapped payload, extract it
+                if "payload" in parsed:
+                    payload = parsed["payload"]
+                    if payload is not None:
+                        return payload
+                    # If payload is None, use field default
+                    raise PydanticUseDefault
+                # Otherwise return the dict as-is
+                return parsed
+        except (json.JSONDecodeError, ValueError):
+            pass
+    return v

datarobot_genai/drmcp/core/constants.py

@@ -0,0 +1,18 @@
+# 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.
+
+
+MAX_INLINE_SIZE = 1024 * 1024  # 1MB
+DEFAULT_DATAROBOT_ENDPOINT = "https://app.datarobot.com/api/v2"
+RUNTIME_PARAM_ENV_VAR_NAME_PREFIX = "MLOPS_RUNTIME_PARAM_"

datarobot_genai/drmcp/core/credentials.py

@@ -0,0 +1,190 @@
+# 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 pydantic import AliasChoices
+from pydantic import Field
+from pydantic import field_validator
+from pydantic_settings import BaseSettings
+from pydantic_settings import SettingsConfigDict
+
+from .config_utils import extract_datarobot_credential_runtime_param_payload
+from .config_utils import extract_datarobot_runtime_param_payload
+from .constants import DEFAULT_DATAROBOT_ENDPOINT
+from .constants import RUNTIME_PARAM_ENV_VAR_NAME_PREFIX
+
+
+class DataRobotCredentials(BaseSettings):
+    """DataRobot API credentials."""
+
+    application_api_token: str = Field(
+        default="",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "DATAROBOT_API_TOKEN",
+            "DATAROBOT_API_TOKEN",
+        ),
+        description="DataRobot API token",
+    )
+    endpoint: str = Field(
+        default=DEFAULT_DATAROBOT_ENDPOINT,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "DATAROBOT_ENDPOINT",
+            "DATAROBOT_ENDPOINT",
+        ),
+        description="DataRobot API endpoint",
+    )
+
+    @field_validator(
+        "application_api_token",
+        "endpoint",
+        mode="before",
+    )
+    @classmethod
+    def validate_runtime_params(cls, v: Any) -> Any:
+        """Validate runtime parameters."""
+        return extract_datarobot_runtime_param_payload(v)
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        case_sensitive=False,
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+
+class MCPServerCredentials(BaseSettings):
+    """Application credentials combining DataRobot and AWS credentials."""
+
+    datarobot: DataRobotCredentials = Field(default_factory=DataRobotCredentials)
+
+    # AWS Credentials - loaded from DataRobot credential object via aws_credential runtime parameter
+    aws_credential: dict[str, Any] | None = Field(
+        default=None,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "AWS_CREDENTIAL",
+            "AWS_CREDENTIAL",
+        ),
+        description="DataRobot AWS Credential object (contains awsAccessKeyId, "
+        "awsSecretAccessKey, awsSessionToken)",
+    )
+    # AWS credentials are also available as direct environment variables for local development
+    aws_access_key_id: str | None = Field(
+        default=None,
+        alias="AWS_ACCESS_KEY_ID",
+        description="AWS Access Key ID (direct, for local use)",
+    )
+    aws_secret_access_key: str | None = Field(
+        default=None,
+        alias="AWS_SECRET_ACCESS_KEY",
+        description="AWS Secret Access Key (direct, for local use)",
+    )
+    aws_session_token: str | None = Field(
+        default=None,
+        alias="AWS_SESSION_TOKEN",
+        description="AWS Session Token (direct, for local use)",
+    )
+
+    aws_predictions_s3_bucket: str = Field(
+        default="datarobot-rd",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "AWS_PREDICTIONS_S3_BUCKET",
+            "AWS_PREDICTIONS_S3_BUCKET",
+        ),
+        description="S3 bucket name",
+    )
+    aws_predictions_s3_prefix: str = Field(
+        default="dev/mcp-temp-storage/predictions/",
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "AWS_PREDICTIONS_S3_PREFIX",
+            "AWS_PREDICTIONS_S3_PREFIX",
+        ),
+        description="S3 key prefix",
+    )
+
+    @field_validator(
+        "aws_credential",
+        "aws_predictions_s3_bucket",
+        "aws_predictions_s3_prefix",
+        mode="before",
+    )
+    @classmethod
+    def validate_credential_runtime_params(cls, v: Any) -> Any:
+        """Validate credential runtime parameters."""
+        return extract_datarobot_credential_runtime_param_payload(v)
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        case_sensitive=False,
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+    def has_aws_credentials(self) -> bool:
+        """Check if AWS credentials are configured (either direct or via credential object)."""
+        return bool((self.aws_access_key_id and self.aws_secret_access_key) or self.aws_credential)
+
+    def has_datarobot_credentials(self) -> bool:
+        """Check if DataRobot credentials are configured."""
+        return bool(self.datarobot.application_api_token)
+
+    def get_aws_credentials(self) -> tuple[str | None, str | None, str | None]:
+        """Get AWS credentials (access_key_id, secret_access_key, session_token).
+
+        If aws_credential dict is set, extracts credentials from it.
+        Otherwise, returns the direct environment variable values.
+
+        Returns
+        -------
+            Tuple of (access_key_id, secret_access_key, session_token)
+        """
+        # If credentials are provided directly (local development), use them
+        if self.aws_access_key_id and self.aws_secret_access_key:
+            return (
+                self.aws_access_key_id,
+                self.aws_secret_access_key,
+                self.aws_session_token,
+            )
+
+        # If credential object is provided, extract keys from it
+        if self.aws_credential:
+            try:
+                return (
+                    self.aws_credential.get("awsAccessKeyId"),
+                    self.aws_credential.get("awsSecretAccessKey"),
+                    self.aws_credential.get("awsSessionToken"),
+                )
+            except Exception as e:
+                logging.getLogger(__name__).warning(
+                    f"Failed to extract AWS credentials from credential object: {e}"
+                )
+                return (None, None, None)
+
+        return (None, None, None)
+
+
+# Global credentials instance
+_credentials: MCPServerCredentials | None = None
+
+
+def get_credentials() -> MCPServerCredentials:
+    """Get the global credentials instance."""
+    # Use a local variable to avoid global statement warning
+    credentials = _credentials
+    if credentials is None:
+        credentials = MCPServerCredentials()
+        # Update the global variable
+        globals()["_credentials"] = credentials
+    return credentials