datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/drmcp/core/config.py

@@ -0,0 +1,364 @@
+# 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 MCPToolConfig(BaseSettings):
+    """Tool configuration for MCP server."""
+
+    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",
+    )
+
+    enable_jira_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
+            "ENABLE_JIRA_TOOLS",
+        ),
+        description="Enable/disable Jira tools",
+    )
+
+    enable_confluence_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
+            "ENABLE_CONFLUENCE_TOOLS",
+        ),
+        description="Enable/disable Confluence tools",
+    )
+
+    enable_gdrive_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
+            "ENABLE_GDRIVE_TOOLS",
+        ),
+        description="Enable/disable GDrive tools",
+    )
+
+    enable_microsoft_graph_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_MICROSOFT_GRAPH_TOOLS",
+            "ENABLE_MICROSOFT_GRAPH_TOOLS",
+        ),
+        description="Enable/disable Sharepoint tools",
+    )
+
+    is_atlassian_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
+            "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Atlassian OAuth provider is configured for Atlassian integration",
+    )
+
+    @property
+    def is_atlassian_oauth_configured(self) -> bool:
+        """Check if Atlassian OAuth is configured via provider flag or environment variables."""
+        return self.is_atlassian_oauth_provider_configured or bool(
+            os.getenv("ATLASSIAN_CLIENT_ID") and os.getenv("ATLASSIAN_CLIENT_SECRET")
+        )
+
+    is_google_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Google OAuth provider is configured for Google integration",
+    )
+
+    @property
+    def is_google_oauth_configured(self) -> bool:
+        return self.is_google_oauth_provider_configured or bool(
+            os.getenv("GOOGLE_CLIENT_ID") and os.getenv("GOOGLE_CLIENT_SECRET")
+        )
+
+    is_microsoft_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
+            "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Microsoft OAuth provider is configured for Microsoft integration",
+    )
+
+    @property
+    def is_microsoft_oauth_configured(self) -> bool:
+        return self.is_microsoft_oauth_provider_configured or bool(
+            os.getenv("MICROSOFT_CLIENT_ID") and os.getenv("MICROSOFT_CLIENT_SECRET")
+        )
+
+    @field_validator(
+        "enable_predictive_tools",
+        "enable_jira_tools",
+        "enable_confluence_tools",
+        "enable_gdrive_tools",
+        "enable_microsoft_graph_tools",
+        "is_atlassian_oauth_provider_configured",
+        "is_google_oauth_provider_configured",
+        "is_microsoft_oauth_provider_configured",
+        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 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",
+    )
+
+    tool_config: MCPToolConfig = Field(
+        default_factory=MCPToolConfig,
+        description="Tool configuration",
+    )
+
+    @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",
+        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_"