datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/drmcp/core/dynamic_prompts/controllers.py

@@ -0,0 +1,130 @@
+# 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.prompts.prompt import Prompt
+
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template_version
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template_versions
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_templates
+from datarobot_genai.drmcp.core.dynamic_prompts.register import (
+    register_prompt_from_datarobot_prompt_management,
+)
+from datarobot_genai.drmcp.core.exceptions import DynamicPromptRegistrationError
+from datarobot_genai.drmcp.core.mcp_instance import mcp
+
+logger = logging.getLogger(__name__)
+
+
+async def register_prompt_from_prompt_template_id_and_version(
+    prompt_template_id: str, prompt_template_version_id: str | None
+) -> Prompt:
+    """Register a Prompt for a specific prompt template ID and version.
+
+    Args:
+        prompt_template_id: The ID of the DataRobot prompt template.
+        prompt_template_version_id: Optional ID of the DataRobot prompt template version.
+            If not provided latest will be used
+
+    Raises
+    ------
+        DynamicPromptRegistrationError: If registration fails at any step.
+
+    Returns
+    -------
+        The registered Prompt instance.
+    """
+    prompt_template = get_datarobot_prompt_template(prompt_template_id)
+
+    if not prompt_template:
+        raise DynamicPromptRegistrationError("Registration failed. Could not find prompt template.")
+
+    if not prompt_template_version_id:
+        return await register_prompt_from_datarobot_prompt_management(
+            prompt_template=prompt_template
+        )
+
+    prompt_template_version = get_datarobot_prompt_template_version(
+        prompt_template_id, prompt_template_version_id
+    )
+
+    if not prompt_template_version:
+        raise DynamicPromptRegistrationError(
+            "Registration failed. Could not find prompt template version."
+        )
+
+    return await register_prompt_from_datarobot_prompt_management(
+        prompt_template=prompt_template, prompt_template_version=prompt_template_version
+    )
+
+
+async def delete_registered_prompt_template(prompt_template_id: str) -> bool:
+    """Delete the prompt registered for the prompt template id in the MCP instance."""
+    prompt_templates_mappings = await mcp.get_prompt_mapping()
+    if prompt_template_id not in prompt_templates_mappings:
+        logger.debug(f"No prompt registered for prompt template id {prompt_template_id}")
+        return False
+
+    prompt_template_version_id, prompt_name = prompt_templates_mappings[prompt_template_id]
+    await mcp.remove_prompt_mapping(prompt_template_id, prompt_template_version_id)
+    logger.info(
+        f"Deleted prompt name {prompt_name} for prompt template id {prompt_template_id}, "
+        f"version {prompt_template_version_id}"
+    )
+    return True
+
+
+async def refresh_registered_prompt_template() -> None:
+    """Refresh all registered prompt templates in the MCP instance."""
+    prompt_templates = get_datarobot_prompt_templates()
+    prompt_templates_ids = {p.id for p in prompt_templates}
+    prompt_templates_versions = get_datarobot_prompt_template_versions(list(prompt_templates_ids))
+
+    mcp_prompt_templates_mappings = await mcp.get_prompt_mapping()
+
+    for prompt_template in prompt_templates:
+        prompt_template_versions = prompt_templates_versions.get(prompt_template.id)
+        if not prompt_template_versions:
+            continue
+
+        latest_version = max(prompt_template_versions, key=lambda v: v.version)
+
+        if prompt_template.id not in mcp_prompt_templates_mappings:
+            # New prompt template -> add
+            await register_prompt_from_datarobot_prompt_management(
+                prompt_template=prompt_template, prompt_template_version=latest_version
+            )
+            continue
+
+        mcp_prompt_template_version, mcp_prompt = mcp_prompt_templates_mappings[prompt_template.id]
+
+        if mcp_prompt_template_version != latest_version:
+            # Current version saved in MCP is not the latest one => update it
+            await register_prompt_from_datarobot_prompt_management(
+                prompt_template=prompt_template, prompt_template_version=latest_version
+            )
+            continue
+
+        # Else => mcp_prompt_template_version == latest_version
+        # For now it means nothing changed as there's no possibility to edit promp template version.
+
+    for mcp_prompt_template_id, (
+        mcp_prompt_template_version_id,
+        _,
+    ) in mcp_prompt_templates_mappings.items():
+        if mcp_prompt_template_id not in prompt_templates_ids:
+            # We need to also delete prompt templates that are
+            await mcp.remove_prompt_mapping(mcp_prompt_template_id, mcp_prompt_template_version_id)

datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py

@@ -0,0 +1,70 @@
+# 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.
+from collections import defaultdict
+
+import datarobot as dr
+
+from datarobot_genai.drmcp.core.clients import get_api_client
+
+
+def get_datarobot_prompt_templates() -> list[dr.genai.PromptTemplate]:
+    try:
+        return dr.genai.PromptTemplate.list()
+    except Exception:
+        return []
+
+
+def get_datarobot_prompt_template_versions(
+    prompt_template_ids: list[str],
+) -> dict[str, list[dr.genai.PromptTemplateVersion]]:
+    # Still missing in SDK
+    prompt_template_versions_data = dr.utils.pagination.unpaginate(
+        initial_url="genai/promptTemplates/versions/",
+        initial_params={
+            "promptTemplateIds": prompt_template_ids,
+        },
+        client=get_api_client(),
+    )
+    prompt_template_versions = defaultdict(list)
+    for prompt_template_version in prompt_template_versions_data:
+        prompt_template_versions[prompt_template_version["promptTemplateId"]].append(
+            dr.genai.PromptTemplateVersion(
+                id=prompt_template_version["id"],
+                prompt_template_id=prompt_template_version["promptTemplateId"],
+                prompt_text=prompt_template_version["promptText"],
+                commit_comment=prompt_template_version["commitComment"],
+                version=prompt_template_version["version"],
+                variables=prompt_template_version["variables"],
+                creation_date=prompt_template_version["creationDate"],
+                creation_user_id=prompt_template_version["creationUserId"],
+                user_name=prompt_template_version["userName"],
+            )
+        )
+    return prompt_template_versions
+
+
+def get_datarobot_prompt_template(prompt_template_id: str) -> dr.genai.PromptTemplate | None:
+    try:
+        return dr.genai.PromptTemplate.get(prompt_template_id)
+    except Exception:
+        return None
+
+
+def get_datarobot_prompt_template_version(
+    prompt_template_id: str, prompt_template_version_id: str
+) -> dr.genai.PromptTemplateVersion | None:
+    try:
+        return dr.genai.PromptTemplateVersion.get(prompt_template_id, prompt_template_version_id)
+    except Exception:
+        return None

datarobot_genai/drmcp/core/dynamic_prompts/register.py

@@ -0,0 +1,205 @@
+# 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 keyword
+import logging
+import re
+from collections.abc import Callable
+from inspect import Parameter
+from inspect import Signature
+
+import datarobot as dr
+from fastmcp.prompts.prompt import Prompt
+from pydantic import Field
+
+from datarobot_genai.drmcp.core.exceptions import DynamicPromptRegistrationError
+from datarobot_genai.drmcp.core.mcp_instance import register_prompt
+
+from .dr_lib import get_datarobot_prompt_template_versions
+from .dr_lib import get_datarobot_prompt_templates
+
+logger = logging.getLogger(__name__)
+
+
+async def register_prompts_from_datarobot_prompt_management() -> None:
+    """Register prompts from DataRobot Prompt Management."""
+    prompts = get_datarobot_prompt_templates()
+    logger.info(f"Found {len(prompts)} prompts in Prompts Management.")
+    all_prompts_versions = get_datarobot_prompt_template_versions(
+        prompt_template_ids=list({prompt.id for prompt in prompts})
+    )
+
+    # Try to register each prompt, continue on failure
+    for prompt in prompts:
+        prompt_versions = all_prompts_versions.get(prompt.id)
+        if not prompt_versions:
+            logger.warning(f"Prompt template id {prompt.id} has no versions.")
+            continue
+
+        latest_version = max(prompt_versions, key=lambda v: v.version)
+
+        try:
+            await register_prompt_from_datarobot_prompt_management(prompt, latest_version)
+        except DynamicPromptRegistrationError:
+            pass
+
+
+async def register_prompt_from_datarobot_prompt_management(
+    prompt_template: dr.genai.PromptTemplate,
+    prompt_template_version: dr.genai.PromptTemplateVersion | None = None,
+) -> Prompt:
+    """Register a single prompt.
+
+    Args:
+        prompt_template: The prompt within DataRobot Prompt Management.
+        prompt_template_version: Optional prompt version within DataRobot Prompt Management.
+            If not provided -- latest version will be used
+
+    Raises
+    ------
+        DynamicPromptRegistrationError: If registration fails at any step.
+
+    Returns
+    -------
+        The registered Prompt instance.
+    """
+    if not prompt_template_version:
+        prompt_template_version_to_register = prompt_template.get_latest_version()
+
+        if prompt_template_version_to_register is None:
+            logger.info(
+                f"No latest version in Prompts Management for prompt id: {prompt_template.id}"
+            )
+            raise DynamicPromptRegistrationError
+
+    else:
+        prompt_template_version_to_register = prompt_template_version
+
+    logger.info(
+        f"Found prompt: id: {prompt_template.id}, "
+        f"name: {prompt_template.name}, "
+        f"prompt version id: {prompt_template_version_to_register.id}, "
+        f"version: {prompt_template_version_to_register.version}."
+    )
+
+    try:
+        valid_fn_name = to_valid_mcp_prompt_name(prompt_template.name)
+    except ValueError as e:
+        raise DynamicPromptRegistrationError from e
+
+    prompt_fn = make_prompt_function(
+        name=valid_fn_name,
+        description=prompt_template.description,
+        prompt_text=prompt_template_version_to_register.prompt_text,
+        variables=prompt_template_version_to_register.variables,
+    )
+
+    try:
+        # Register using generic external tool registration with the config
+        return await register_prompt(
+            fn=prompt_fn,
+            name=prompt_template.name,
+            description=prompt_template.description,
+            meta={
+                "prompt_template_id": prompt_template.id,
+                "prompt_template_version_id": prompt_template_version_to_register.id,
+            },
+            prompt_template=(prompt_template.id, prompt_template_version_to_register.id),
+        )
+
+    except Exception as exc:
+        logger.error(f"Skipping prompt {prompt_template.id}. Registration failed: {exc}")
+        raise DynamicPromptRegistrationError(
+            "Registration failed. Could not create prompt."
+        ) from exc
+
+
+def _escape_non_ascii(s: str) -> str:
+    out = []
+    for ch in s:
+        # If its space -> change to underscore
+        if ch.isspace():
+            out.append("_")
+        # ASCII letter, digit or underscore -> keep
+        elif ch.isascii() and (ch.isalnum() or ch == "_"):
+            out.append(ch)
+        # Everything else -> encode as 'xHEX'
+        else:
+            out.append(f"x{ord(ch):x}")
+    return "".join(out)
+
+
+def to_valid_mcp_prompt_name(s: str) -> str:
+    """Convert an arbitrary string into a valid MCP prompt name."""
+    # If its ONLY numbers return "prompt_[number]"
+    if s.isdigit():
+        return f"prompt_{s}"
+
+    # First, ASCII-transliterate using hex escape for non-ASCII
+    if not s.isascii():
+        # whole string non-ascii? -> escape and prefix with prompt_
+        encoded = _escape_non_ascii(s)
+        return f"prompt_{encoded}"
+
+    # Replace any sequence of invalid characters with '_'
+    s = re.sub(r"[^0-9a-zA-Z_]+", "_", s)
+
+    # Remove leading characters that are not letters or underscores (can't start with a digit or _)
+    s = re.sub(r"^[^a-zA-Z]+", "", s)
+
+    # Remove following _
+    s = re.sub(r"_+$", "", s)
+
+    # If string is empty after cleaning, raise error
+    if not s:
+        raise ValueError(f"Cannot convert {s} to valid MCP prompt name.")
+
+    # Make sure it's a valid identifier and not a reserved keyword
+    if keyword.iskeyword(s) or not s.isidentifier():
+        s = f"{s}_prompt"
+
+    return s
+
+
+def make_prompt_function(
+    name: str, description: str, prompt_text: str, variables: list[dr.genai.Variable]
+) -> Callable:
+    params = []
+    for v in variables:
+        if keyword.iskeyword(v.name):
+            raise ValueError(f"Variable name '{v.name}' is invalid.")
+
+        try:
+            param = Parameter(
+                name=v.name,
+                kind=Parameter.POSITIONAL_OR_KEYWORD,
+                default=Field(description=v.description),
+            )
+        except ValueError as e:
+            raise ValueError(f"Variable name '{v.name}' is invalid.") from e
+
+        params.append(param)
+
+    async def template_function(**kwargs) -> str:  # type: ignore
+        prompt_text_correct = prompt_text.replace("{{", "{").replace("}}", "}")
+        try:
+            return prompt_text_correct.format(**kwargs)
+        except KeyError as exc:
+            raise ValueError(f"Missing variable {exc.args[0]} for prompt '{name}'") from exc
+
+    # Apply metadata
+    template_function.__name__ = name
+    template_function.__doc__ = description
+    template_function.__signature__ = Signature(params)  # type: ignore
+
+    return template_function

datarobot_genai/drmcp/core/dynamic_prompts/utils.py

@@ -0,0 +1,33 @@
+# 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.
+from uuid import uuid4
+
+from fastmcp import FastMCP
+from fastmcp.exceptions import NotFoundError
+
+_SUFFIX_LENGTH: int = 4
+
+
+async def get_prompt_name_no_duplicate(mcp: FastMCP, prompt_name: str) -> str:
+    """Handle prompt name duplicate.
+
+    We're working optimistic here -- we're keeping default names unless there's collision
+    """
+    try:
+        prompt = await mcp.get_prompt(prompt_name)
+    except NotFoundError:
+        return prompt_name
+
+    prompt_name_suffix = str(uuid4())[:_SUFFIX_LENGTH]
+    return f"{prompt.name} ({prompt_name_suffix})"

datarobot_genai/drmcp/core/dynamic_tools/__init__.py

@@ -0,0 +1,14 @@
+# 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/dynamic_tools/deployment/adapters/__init__.py

@@ -0,0 +1,14 @@
+# 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/dynamic_tools/deployment/adapters/base.py

@@ -0,0 +1,72 @@
+# 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.
+
+from abc import ABC
+from abc import abstractmethod
+from typing import Any
+from typing import Literal
+
+
+class MetadataBase(ABC):
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """The name of the tool, for the LLM to identify and call the tool."""
+        pass
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """The description of the tool, for the LLM to understand its purpose
+        and all additional instructions and context about the tool, which can
+        help the LLM to better utilize the tool in the right context.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def method(self) -> Literal["GET", "POST", "PATCH", "PUT", "DELETE"]:
+        """HTTP method to use when calling the tool, e.g. `POST` or `GET` etc."""
+        pass
+
+    @property
+    @abstractmethod
+    def endpoint(self) -> str:
+        """The endpoint path of the tool, e.g. `/weather/{city}/forecast`."""
+        pass
+
+    @property
+    @abstractmethod
+    def input_schema(self) -> dict[str, Any]:
+        """The JSON schema defining the input parameters for the tool.
+
+        Structure:
+            {
+                "type": "object",
+                "properties": {
+                    "path_params": {...},    # Optional: path parameter schemas
+                    "query_params": {...},   # Optional: query parameter schemas
+                    "data": {...},           # Optional: form/body data schema
+                    "json": {...}            # Optional: JSON body schema
+                },
+                "required": [...]            # Optional: required properties
+            }
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def headers(self) -> dict[str, str]:
+        """Optional HTTP headers to include when calling the tool."""
+        pass

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

@@ -0,0 +1,82 @@
+# 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.
+from typing import Any
+from typing import Literal
+from typing import cast
+
+from .base import MetadataBase
+
+
+class Metadata(MetadataBase):
+    """Default adapter for external deployment metadata."""
+
+    def __init__(self, metadata: dict[str, Any]) -> None:
+        self.metadata = metadata
+
+    @property
+    def name(self) -> str:
+        return str(self.metadata.get("name", ""))
+
+    @property
+    def description(self) -> str:
+        return str(self.metadata.get("description", ""))
+
+    @property
+    def base_url(self) -> str:
+        base_url = self.metadata.get("base_url")
+        if not base_url or not isinstance(base_url, str):
+            raise ValueError(
+                "Deployment missing required 'base_url' field in /info/ metadata. "
+                "This field is required for MCP Server to route tool requests correctly."
+            )
+        return str(base_url)
+
+    @property
+    def endpoint(self) -> str:
+        endpoint = self.metadata.get("endpoint")
+        if not endpoint or not isinstance(endpoint, str):
+            raise ValueError(
+                "Deployment missing required 'endpoint' field in /info/ metadata. "
+                "This field is required for MCP Server to route tool requests correctly."
+            )
+        return str(endpoint)
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        input_schema = self.metadata.get("input_schema")
+        if not input_schema or not isinstance(input_schema, dict):
+            raise ValueError(
+                "Deployment missing required 'inputSchema' field in /info/ metadata. "
+                "This field is required for MCP Server to route tool requests correctly."
+            )
+        return dict(input_schema)
+
+    @property
+    def method(self) -> Literal["GET", "POST", "PATCH", "PUT", "DELETE"]:
+        method = self.metadata.get("method", "").upper()
+        if not method or not isinstance(method, str):
+            raise ValueError(
+                "Deployment missing required 'method' field in /info/ metadata. "
+                "This field is required for MCP Server to route tool requests correctly."
+            )
+        if method not in ("GET", "POST", "PATCH", "PUT", "DELETE"):
+            raise ValueError(f"Deployment metadata is invalid, unsupported `method`: {method}.")
+        return cast(Literal["GET", "POST", "PATCH", "PUT", "DELETE"], method)
+
+    @property
+    def headers(self) -> dict[str, str]:
+        headers = self.metadata.get("headers", {})
+        if not isinstance(headers, dict):
+            raise ValueError("Deployment metadata 'headers' field must be a dictionary.")
+        return dict(headers)