datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/drum.py

@@ -0,0 +1,238 @@
+# 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 enum import Enum
+from functools import lru_cache
+from pathlib import Path
+from typing import Any
+from typing import Literal
+
+from .base import MetadataBase
+
+# Path to the schemas directory
+SCHEMAS_DIR = Path(__file__).parent.parent / "schemas"
+
+
+class DrumTargetType(str, Enum):
+    BINARY = "binary"
+    REGRESSION = "regression"
+    ANOMALY = "anomaly"
+    UNSTRUCTURED = "unstructured"
+    MULTICLASS = "multiclass"
+    TEXT_GENERATION = "textgeneration"
+    GEO_POINT = "geopoint"
+    VECTOR_DATABASE = "vectordatabase"
+    AGENTIC_WORKFLOW = "agenticworkflow"
+
+    @classmethod
+    def prediction_types(cls) -> set["DrumTargetType"]:
+        """Get the set of DRUM target types that correspond to structured predictions."""
+        return {
+            DrumTargetType.BINARY,
+            DrumTargetType.REGRESSION,
+            DrumTargetType.ANOMALY,
+            DrumTargetType.MULTICLASS,
+            DrumTargetType.TEXT_GENERATION,
+            DrumTargetType.GEO_POINT,
+            DrumTargetType.VECTOR_DATABASE,
+        }
+
+
+@lru_cache(maxsize=1)
+def _get_prediction_fallback_schema() -> dict[str, Any]:
+    """Get the default prediction input schema for DRUM deployments."""
+    schema_path = SCHEMAS_DIR / "drum_prediction_fallback_schema.json"
+    with open(schema_path) as f:
+        schema: dict[str, Any] = json.load(f)
+        return schema
+
+
+@lru_cache(maxsize=1)
+def _get_agentic_fallback_schema() -> dict[str, Any]:
+    """Get the default agentic workflow input schema for DRUM deployments."""
+    schema_path = SCHEMAS_DIR / "drum_agentic_fallback_schema.json"
+    with open(schema_path) as f:
+        schema: dict[str, Any] = json.load(f)
+        return schema
+
+
+def get_default_schema(target_type: str) -> dict[str, Any]:
+    """Get the default input schema for a given DRUM target type, when
+    the deployment does not provide one. This fallback mechanism is here
+    to lower the friction of using DRUM deployments with MCP, for more
+    advanced use cases it is recommended to provide a custom input and
+    expose it via model-metadata.yaml inputSchema parameter.
+
+    Args:
+        target_type: The target type of the DRUM deployment.
+
+    Returns
+    -------
+        A dictionary representing the default input schema wrapped in HTTP request structure.
+    """
+    if target_type == DrumTargetType.AGENTIC_WORKFLOW:
+        return _get_agentic_fallback_schema()
+
+    if target_type in DrumTargetType.prediction_types():
+        return _get_prediction_fallback_schema()
+
+    return {}
+
+
+def is_drum(metadata: dict[str, Any]) -> bool:
+    """Check if the deployment is a DRUM deployment.
+
+    DRUM deployments are identified by the presence of both drum_server
+    and drum_version fields in the metadata response.
+
+    Args:
+        metadata: The response retrieved from the custom model /info/ route.
+
+    Returns
+    -------
+        True if this is a DRUM deployment, False otherwise.
+    """
+    drum_server = metadata.get("drum_server")
+    drum_version = metadata.get("drum_version")
+    return bool(drum_server or drum_version)
+
+
+class DrumMetadataAdapter(MetadataBase):
+    """Adapter for DRUM deployment metadata."""
+
+    def __init__(self, metadata: dict[str, Any]):
+        """Initialize adapter with validated metadata.
+
+        Args:
+            metadata: Dictionary containing at minimum a 'target_type' key.
+
+        Note:
+            Use class methods `from_deployment_metadata()` or `from_target_type()`
+            for construction instead of calling this directly.
+        """
+        self.metadata = metadata
+        self._validate_tool_support()
+
+    def _validate_tool_support(self) -> None:
+        """Validate that DRUM deployments are supported in the current environment.
+
+        Raises
+        ------
+            ValueError: If DRUM deployments are not supported.
+        """
+        if self.target_type not in list(DrumTargetType):
+            raise ValueError(
+                f"The deployment target_type: {self.target_type} "
+                f"is not supported, to be registered as MCP Tool."
+            )
+
+    @classmethod
+    def from_deployment_metadata(cls, metadata: dict[str, Any]) -> "DrumMetadataAdapter":
+        """Create adapter from full deployment metadata.
+
+        Args:
+            metadata: The response retrieved from the custom model /info/ route.
+
+        Returns
+        -------
+            DrumMetadataAdapter instance.
+
+        Raises
+        ------
+            ValueError: If metadata is not from a DRUM deployment.
+        """
+        if not is_drum(metadata):
+            raise ValueError("Provided metadata is not from a DRUM deployment.")
+        return cls(metadata)
+
+    @classmethod
+    def from_target_type(cls, target_type: str) -> "DrumMetadataAdapter":
+        """Create adapter from target type only.
+
+        Used for testing/minimal setup when broader set of information
+        from metadata built from model-metadata.yaml information is
+        not available i.e. datarobot predictive models.
+
+        Args:
+            target_type: The DRUM target type (e.g., 'binary', 'regression').
+
+        Returns
+        -------
+            DrumMetadataAdapter instance with minimal metadata.
+        """
+        return cls({"target_type": target_type.lower()})
+
+    @property
+    def target_type(self) -> str:
+        return str(self.metadata["target_type"])
+
+    @property
+    def name(self) -> str:
+        return str(self.model_metadata.get("name", ""))
+
+    @property
+    def description(self) -> str:
+        return str(self.model_metadata.get("description", ""))
+
+    @property
+    def endpoint(self) -> str:
+        """Return the appropriate endpoint for the DRUM target type."""
+        predictions_endpoint = "/predictions"
+
+        endpoint_map: dict[str, str] = {
+            DrumTargetType.BINARY: predictions_endpoint,
+            DrumTargetType.REGRESSION: predictions_endpoint,
+            DrumTargetType.ANOMALY: predictions_endpoint,
+            DrumTargetType.MULTICLASS: predictions_endpoint,
+            DrumTargetType.TEXT_GENERATION: predictions_endpoint,
+            DrumTargetType.GEO_POINT: predictions_endpoint,
+            DrumTargetType.UNSTRUCTURED: "/predictionsUnstructured",
+            DrumTargetType.VECTOR_DATABASE: predictions_endpoint,
+            DrumTargetType.AGENTIC_WORKFLOW: "/chat/completions",
+        }
+
+        return endpoint_map[self.target_type]
+
+    @property
+    def model_metadata(self) -> dict[str, Any]:
+        result = self.metadata.get("model_metadata", {})
+        return dict(result)
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        input_schema = self.model_metadata.get("input_schema", get_default_schema(self.target_type))
+
+        if not input_schema or not isinstance(input_schema, dict):
+            raise ValueError(
+                "DRUM deployment is missing a valid input schema. Please make "
+                "sure the model-metadata.yaml file includes `inputSchema` "
+                "definition and that custom model is using datarobot-drum in "
+                "version v1.17.2 or later."
+            )
+        return dict(input_schema)
+
+    @property
+    def method(self) -> Literal["GET", "POST", "PATCH", "PUT", "DELETE"]:
+        return "POST"
+
+    @property
+    def headers(self) -> dict[str, str]:
+        """Return HTTP headers required for this deployment type."""
+        if self.target_type in DrumTargetType.prediction_types():
+            # structured predictions send data as CSV bytes, which
+            # requires an explicit Content-Type header since aiohttp
+            # won't set it automatically
+            return {"Content-Type": "text/csv"}
+
+        return {}

datarobot_genai/drmcp/core/dynamic_tools/deployment/config.py

@@ -0,0 +1,228 @@
+# 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.
+
+"""Configuration assembly for DataRobot deployment tools.
+
+This module is responsible for creating complete ExternalToolRegistrationConfig
+objects from DataRobot deployments. It handles all aspects of configuration
+including metadata fetching, URL construction, authentication, and schema assembly.
+"""
+
+import re
+from urllib.parse import urljoin
+
+import datarobot as dr
+
+from datarobot_genai.drmcp.core.clients import get_api_client
+from datarobot_genai.drmcp.core.dynamic_tools.deployment.adapters.base import MetadataBase
+from datarobot_genai.drmcp.core.dynamic_tools.deployment.metadata import get_mcp_tool_metadata
+from datarobot_genai.drmcp.core.dynamic_tools.register import ExternalToolRegistrationConfig
+
+
+def create_deployment_tool_config(
+    deployment: dr.Deployment,
+) -> ExternalToolRegistrationConfig:
+    """Create an ExternalToolRegistrationConfig from deployment.
+
+    This is the main public API of this module. It gathers all the information
+    needed to register a DataRobot deployment as an external tool, handling both
+    DRUM and standard deployments. It fetches metadata, extracts deployment-specific
+    configuration, and assembles everything into a complete registration config.
+
+    Args:
+        deployment: The DataRobot deployment object.
+
+    Returns
+    -------
+        ExternalToolRegistrationConfig with all parameters needed for registration.
+
+    Raises
+    ------
+        Exception: If metadata fetching or config assembly fails.
+    """
+    # Fetch and validate metadata from the deployment
+    metadata = get_mcp_tool_metadata(deployment)
+
+    # Get deployment-specific infrastructure configuration
+    base_url = _get_deployment_base_url(deployment)
+    auth_headers = _get_deployment_auth_headers(deployment)
+
+    # Merge metadata headers with deployment headers
+    merged_headers = {**auth_headers, **metadata.headers}
+
+    # Build endpoint path
+    endpoint = metadata.endpoint.lstrip("/")
+
+    # Generate tool name and description
+    tool_name = _get_tool_name(deployment, metadata)
+    tool_description = _get_tool_description(deployment, metadata)
+
+    return ExternalToolRegistrationConfig(
+        name=tool_name,
+        title=deployment.label,
+        description=tool_description,
+        method=metadata.method,
+        base_url=base_url,
+        endpoint=endpoint,
+        headers=merged_headers,
+        input_schema=metadata.input_schema,
+        tags=set(),  # Add missing tags parameter
+    )
+
+
+def _is_serverless_deployment(deployment: dr.Deployment) -> bool:
+    """Check if deployment is serverless."""
+    if not deployment.prediction_environment:
+        return False
+    return deployment.prediction_environment.get("platform") == "datarobotServerless"
+
+
+def _get_deployment_base_url(deployment: dr.Deployment) -> str:
+    """Get base URL for deployment prediction server.
+
+    Args:
+        deployment: DataRobot deployment instance
+
+    Returns
+    -------
+        Formatted deployment URL including deployment ID path
+
+    Raises
+    ------
+        ValueError: If prediction server cannot be determined
+    """
+    api_client = get_api_client()
+
+    # Determine base URL based on deployment type
+    if _is_serverless_deployment(deployment):
+        base_url = api_client.endpoint
+    elif "datarobot-nginx" in api_client.endpoint:
+        # On-prem/ST SAAS environments
+        base_url = "http://datarobot-prediction-server:80/predApi/v1.0"
+    else:
+        # Regular prediction server
+        pred_server = deployment.default_prediction_server
+        if not pred_server:
+            raise ValueError(f"Deployment {deployment.id} has no default prediction server")
+
+        url = pred_server["url"]
+        if not url:
+            raise ValueError(f"Deployment {deployment.id} prediction server has no URL")
+        base_url = f"{url}/predApi/v1.0"
+
+    merged_url = urljoin(base_url.rstrip("/") + "/", f"deployments/{deployment.id}/")
+    return merged_url
+
+
+def _get_deployment_auth_headers(deployment: dr.Deployment) -> dict[str, str]:
+    """Get authentication headers for deployment.
+
+    Args:
+        deployment: DataRobot deployment instance
+
+    Returns
+    -------
+        Dictionary of authentication headers
+    """
+    headers = {"Authorization": f"Bearer {get_api_client().token}"}
+
+    # For non-serverless deployments, include datarobot-key
+    if not _is_serverless_deployment(deployment):
+        pred_server = deployment.default_prediction_server
+        if pred_server:
+            dr_key = pred_server.get("datarobot-key")
+            if dr_key:
+                headers["datarobot-key"] = dr_key
+
+    return headers
+
+
+def _get_tool_name(deployment: dr.Deployment, metadata: MetadataBase) -> str:
+    """Generate tool name from deployment and metadata."""
+    tool_name = deployment.label or metadata.name or f"deployment_{deployment.id}"
+    return _convert_tool_string(tool_name)
+
+
+def _get_additional_prediction_instructions(deployment_id: str) -> str:
+    """Generate additional instructions for scoring prediction models, to make tool usage more
+    reliable.
+    """
+    return f"""
+    
+Follow these steps in order:
+1. Get deployment info: Call tools with the deployment_id="{deployment_id}" to learn about
+   features and requirements.
+2. Retrieve features: Use `get_deployment_features` to see all required and optional features
+   with their importance scores.
+3. Prepare data: Use `generate_prediction_data_template` to create the correctly structured
+   CSV format.
+4. Consider feature importance: For high-importance features, always provide values (infer or
+   ask). Low-importance features can be left blank.
+5. Validate: Run `validate_prediction_data` before submission to catch errors early.
+6. Time series note: Ensure `datetime_column` and `series_id_columns` are properly formatted
+   if applicable.
+
+Parameter details and format requirements are specified in the input schema below."""
+
+
+def _get_tool_description(deployment: dr.Deployment, metadata: MetadataBase) -> str:
+    """Generate tool description from deployment and metadata.
+
+    Args:
+        deployment: The DataRobot deployment object.
+        metadata: The metadata adapter containing tool information.
+
+    Returns
+    -------
+        Complete tool description, optionally enhanced with workflow instructions
+        for prediction endpoints.
+    """
+    base_description = deployment.description or metadata.description
+
+    if metadata.endpoint.endswith("predictions"):
+        additional_instructions = _get_additional_prediction_instructions(deployment.id)
+        return f"{base_description}{additional_instructions}"
+
+    return base_description
+
+
+def _convert_tool_string(text: str | None) -> str:
+    """Convert a string to a valid tool name format.
+
+    Removes brackets, replaces spaces/hyphens with underscores, removes special
+    characters, converts to lowercase, and cleans up multiple underscores.
+    """
+    if not text:
+        return ""
+
+    # Remove anything within brackets (including the brackets)
+    text = re.sub(r"\[.*?\]", "", text)
+
+    # Replace spaces with underscores
+    text = text.replace(" ", "_")
+    text = text.replace("-", "_")
+
+    # Remove all non-alphanumeric characters except underscores
+    text = re.sub(r"[^a-zA-Z0-9_]", "", text)
+
+    # Convert to lowercase
+    text = text.lower()
+
+    # Clean up any multiple underscores that might result
+    text = re.sub(r"_+", "_", text)
+
+    # Remove leading/trailing underscores
+    text = text.strip("_")
+
+    return text

datarobot_genai/drmcp/core/dynamic_tools/deployment/controllers.py

@@ -0,0 +1,63 @@
+# 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 fastmcp.tools.tool import Tool
+
+from datarobot_genai.drmcp.core.clients import get_sdk_client
+from datarobot_genai.drmcp.core.dynamic_tools.deployment.register import (
+    register_tool_of_datarobot_deployment,
+)
+from datarobot_genai.drmcp.core.mcp_instance import mcp
+
+logger = logging.getLogger(__name__)
+
+
+async def register_tool_for_deployment_id(deployment_id: str) -> Tool:
+    """Register a tool for a specific deployment ID.
+
+    Args:
+        deployment_id: The ID of the DataRobot deployment to register as a tool.
+
+    Raises
+    ------
+        DynamicToolRegistrationError: If registration fails at any step.
+
+    Returns
+    -------
+        The registered Tool instance.
+    """
+    deployment = get_sdk_client().Deployment.get(deployment_id)
+    registered_tool = await register_tool_of_datarobot_deployment(deployment)
+    return registered_tool
+
+
+async def get_registered_tool_deployments() -> dict[str, str]:
+    """Get the tool registered for the deployment in the MCP instance."""
+    deployments = await mcp.get_deployment_mapping()
+    return deployments
+
+
+async def delete_registered_tool_deployment(deployment_id: str) -> bool:
+    """Delete the tool registered for the deployment in the MCP instance."""
+    deployments = await mcp.get_deployment_mapping()
+    if deployment_id not in deployments:
+        logger.debug(f"No tool registered for deployment {deployment_id}")
+        return False
+
+    tool_name = deployments[deployment_id]
+    await mcp.remove_deployment_mapping(deployment_id)
+    logger.info(f"Deleted tool {tool_name} for deployment {deployment_id}")
+    return True

datarobot_genai/drmcp/core/dynamic_tools/deployment/metadata.py

@@ -0,0 +1,162 @@
+# 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
+
+import datarobot as dr
+from datarobot.utils import from_api
+
+from datarobot_genai.drmcp.core.clients import get_api_client
+from datarobot_genai.drmcp.core.dynamic_tools.deployment.adapters.base import MetadataBase
+from datarobot_genai.drmcp.core.dynamic_tools.deployment.adapters.default import Metadata
+from datarobot_genai.drmcp.core.dynamic_tools.deployment.adapters.drum import DrumMetadataAdapter
+from datarobot_genai.drmcp.core.dynamic_tools.deployment.adapters.drum import DrumTargetType
+from datarobot_genai.drmcp.core.dynamic_tools.deployment.adapters.drum import is_drum
+
+logger = logging.getLogger(__name__)
+
+
+def _normalize_api_response(response: Any) -> dict[str, Any]:
+    """Normalize API response to a dictionary.
+
+    The API can return either a list or a dict. This function ensures
+    we always get a dict, taking the first element if it's a list.
+
+    Args:
+        response: The raw API response object.
+
+    Returns
+    -------
+        Normalized dictionary representation of the response.
+    """
+    data = from_api(response.json())
+    if isinstance(data, list):
+        return data[0] if data else {}
+    if isinstance(data, dict):
+        return data
+    return {}
+
+
+def _get_model_attribute(model: dict[str, Any], key: str, default: str = "") -> str:
+    """Safely extract a string attribute from a model dictionary.
+
+    Args:
+        model: The deployment model dictionary.
+        key: The attribute key to retrieve.
+        default: Default value if key is not found.
+
+    Returns
+    -------
+        The attribute value as a lowercase string.
+    """
+    value = model.get(key, default)
+    return str(value).lower() if value else default
+
+
+def _is_datarobot_structured_prediction(deployment: dr.Deployment) -> str | None:
+    """Check if deployment is a DataRobot structured prediction model.
+
+    DataRobot native predictive models with structured predictions don't support
+    metadata fetching via API, so they need special handling.
+
+    Args:
+        deployment: The DataRobot deployment object.
+
+    Returns
+    -------
+        Target type string if it's a DataRobot structured prediction, None otherwise.
+    """
+    if deployment.model is None:
+        return None
+
+    # deployment.model is a TypedDict at type-check time, but dict at runtime
+    # Cast it to Dict[str, Any] to allow dynamic key access
+    model_dict: dict[str, Any] = deployment.model  # type: ignore[assignment]
+    target_type = _get_model_attribute(model_dict, "target_type")
+    build_env = _get_model_attribute(model_dict, "build_environment_type")
+
+    if not target_type or not build_env:
+        return None
+
+    if build_env == "datarobot" and target_type in DrumTargetType.prediction_types():
+        return target_type
+
+    return None
+
+
+def _fetch_deployment_metadata(deployment: dr.Deployment) -> dict[str, Any]:
+    """Fetch metadata from deployment's directAccess/info endpoint.
+
+    Args:
+        deployment: The DataRobot deployment object.
+
+    Returns
+    -------
+        Normalized metadata dictionary.
+
+    Raises
+    ------
+        RuntimeError: If API call fails.
+    """
+    api_client = get_api_client()
+
+    try:
+        response = api_client.get(url=f"deployments/{deployment.id}/directAccess/info/")
+        response.raise_for_status()
+        return _normalize_api_response(response)
+    except Exception as exc:
+        logger.error(f"Failed to fetch metadata for deployment {deployment.id}: {exc}")
+        raise RuntimeError(f"Could not retrieve metadata for deployment {deployment.id}") from exc
+
+
+def get_mcp_tool_metadata(deployment: dr.Deployment) -> MetadataBase:
+    """Fetch and validate metadata for a given deployment.
+
+    This method retrieves deployment metadata from the /directAccess/info/
+    endpoint and validates it contains the required fields for tool registration.
+    It uses a universal approach where DRUM deployments are treated as a
+    special case and converted to a standard format.
+
+    The returned metadata must contain at minimum:
+    - input_schema: JSON schema describing the tool's input parameters
+    - endpoint: The deployment endpoint path to call
+
+    Args:
+        deployment: The DataRobot deployment object.
+
+    Raises
+    ------
+        RuntimeError: If metadata fetching fails.
+        ValueError: If the deployment metadata is missing or invalid.
+
+    Returns
+    -------
+        MetadataBase adapter instance to expose required
+        metadata properties for tool registration in a standard way.
+    """
+    # Check if this is a DataRobot native structured prediction model
+    # These don't support metadata API, so we create minimal metadata
+    target_type = _is_datarobot_structured_prediction(deployment)
+    if target_type:
+        return DrumMetadataAdapter.from_target_type(target_type)
+
+    # Fetch metadata from the deployment's info endpoint
+    metadata = _fetch_deployment_metadata(deployment)
+
+    # Return appropriate adapter based on metadata type
+    if is_drum(metadata):
+        return DrumMetadataAdapter.from_deployment_metadata(metadata)
+
+    return Metadata(metadata)