datarobot-genai 0.2.0__py3-none-any.whl

datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py

@@ -0,0 +1,128 @@
+# 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
+from dataclasses import dataclass
+
+import datarobot as dr
+
+from datarobot_genai.drmcp.core.clients import get_api_client
+
+# Needed SDK version (3.10.0) is not published yet. We'll reimplement simplified version of it.
+# get_datarobot_prompt_templates = dr.genai.PromptTemplate.list()
+# DrPrompt = dr.genai.PromptTemplate
+# DrPromptVersion = dr.genai.PromptTemplateVersion
+# DrVariable = dr.genai.Variable
+
+
+@dataclass
+class DrVariable:
+    name: str
+    description: str
+
+
+@dataclass
+class DrPromptVersion:
+    id: str
+    prompt_template_id: str
+    version: int
+    prompt_text: str
+    variables: list[DrVariable]
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "DrPromptVersion":
+        variables = [
+            DrVariable(name=v["name"], description=v["description"]) for v in d["variables"]
+        ]
+        return cls(
+            id=d["id"],
+            prompt_template_id=d["promptTemplateId"],
+            version=d["version"],
+            prompt_text=d["promptText"],
+            variables=variables,
+        )
+
+
+@dataclass
+class DrPrompt:
+    id: str
+    name: str
+    description: str
+
+    def get_latest_version(self) -> DrPromptVersion | None:
+        all_prompt_template_versions = get_datarobot_prompt_template_versions([self.id])
+        prompt_template_versions = all_prompt_template_versions.get(self.id)
+
+        if not prompt_template_versions:
+            return None
+        latest_version = max(prompt_template_versions, key=lambda v: v.version)
+        return latest_version
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "DrPrompt":
+        return cls(id=d["id"], name=d["name"], description=d["description"])
+
+
+def get_datarobot_prompt_templates() -> list[DrPrompt]:
+    prompt_templates_data = dr.utils.pagination.unpaginate(
+        initial_url="genai/promptTemplates/", initial_params={}, client=get_api_client()
+    )
+
+    return [DrPrompt.from_dict(prompt_template) for prompt_template in prompt_templates_data]
+
+
+def get_datarobot_prompt_template_versions(
+    prompt_template_ids: list[str],
+) -> dict[str, list[DrPromptVersion]]:
+    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(
+            DrPromptVersion.from_dict(prompt_template_version)
+        )
+    return prompt_template_versions
+
+
+def get_datarobot_prompt_template(prompt_template_id: str) -> DrPrompt | None:
+    api_client = get_api_client()
+    try:
+        prompt_template_response = api_client.get(
+            f"genai/promptTemplates/{prompt_template_id}/", join_endpoint=True
+        )
+        prompt_template_json = prompt_template_response.json()
+    except Exception:
+        return None
+
+    return DrPrompt.from_dict(prompt_template_json)
+
+
+def get_datarobot_prompt_template_version(
+    prompt_template_id: str, prompt_template_version_id: str
+) -> DrPromptVersion | None:
+    api_client = get_api_client()
+    try:
+        prompt_template_version_response = api_client.get(
+            f"genai/promptTemplates/{prompt_template_id}/versions/{prompt_template_version_id}/",
+            join_endpoint=True,
+        )
+        prompt_template_version_json = prompt_template_version_response.json()
+    except Exception:
+        return None
+
+    return DrPromptVersion.from_dict(prompt_template_version_json)

datarobot_genai/drmcp/core/dynamic_prompts/register.py

@@ -0,0 +1,206 @@
+# 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
+
+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 DrPrompt
+from .dr_lib import DrPromptVersion
+from .dr_lib import DrVariable
+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: DrPrompt, prompt_template_version: DrPromptVersion | 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[DrVariable]
+) -> 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)