datarobot-genai 0.2.0__tar.gz → 0.2.8__tar.gz

{datarobot_genai-0.2.0 → datarobot_genai-0.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.0
+Version: 0.2.8
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -11,13 +11,15 @@ Requires-Python: <3.13,>=3.10
 Requires-Dist: ag-ui-protocol<0.2.0,>=0.1.9
 Requires-Dist: datarobot-drum<2.0.0,>=1.17.5
 Requires-Dist: datarobot-predict<2.0.0,>=1.13.2
-Requires-Dist: datarobot<4.0.0,>=3.9.1
+Requires-Dist: datarobot<4.0.0,>=3.10.0
 Requires-Dist: openai<2.0.0,>=1.76.2
 Requires-Dist: opentelemetry-instrumentation-aiohttp-client<1.0.0,>=0.43b0
 Requires-Dist: opentelemetry-instrumentation-httpx<1.0.0,>=0.43b0
 Requires-Dist: opentelemetry-instrumentation-openai<1.0.0,>=0.40.5
 Requires-Dist: opentelemetry-instrumentation-requests<1.0.0,>=0.43b0
+Requires-Dist: opentelemetry-instrumentation-threading<1.0.0,>=0.43b0
 Requires-Dist: pandas<3.0.0,>=2.2.3
+Requires-Dist: pyarrow==20.0.0
 Requires-Dist: pyjwt<3.0.0,>=2.10.1
 Requires-Dist: pypdf<7.0.0,>=6.1.3
 Requires-Dist: ragas<0.4.0,>=0.3.8
@@ -57,9 +59,11 @@ Requires-Dist: llama-index<0.14.0,>=0.13.6; extra == 'llamaindex'
 Requires-Dist: opentelemetry-instrumentation-llamaindex<1.0.0,>=0.40.5; extra == 'llamaindex'
 Requires-Dist: pypdf<7.0.0,>=6.0.0; extra == 'llamaindex'
 Provides-Extra: nat
+Requires-Dist: anyio==4.11.0; extra == 'nat'
 Requires-Dist: crewai>=1.1.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: llama-index-llms-litellm<0.7.0,>=0.4.1; extra == 'nat'
 Requires-Dist: nvidia-nat-langchain==1.3.0; (python_version >= '3.11') and extra == 'nat'
+Requires-Dist: nvidia-nat-mcp==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat-opentelemetry==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'nat'

{datarobot_genai-0.2.0 → datarobot_genai-0.2.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.0"
+version = "0.2.8"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -12,7 +12,7 @@ license = { text = "Apache-2.0" }
 authors = [{ name = "DataRobot, Inc." }]
 dependencies = [
   "requests>=2.32.4,<3.0.0",
-  "datarobot>=3.9.1,<4.0.0",
+  "datarobot>=3.10.0,<4.0.0",
   "datarobot-predict>=1.13.2,<2.0.0",
   "openai>=1.76.2,<2.0.0",
   "pandas>=2.2.3,<3.0.0",
@@ -24,7 +24,9 @@ dependencies = [
   "opentelemetry-instrumentation-aiohttp-client>=0.43b0,<1.0.0",
   "opentelemetry-instrumentation-httpx>=0.43b0,<1.0.0",
   "opentelemetry-instrumentation-openai>=0.40.5,<1.0.0",
+  "opentelemetry-instrumentation-threading>=0.43b0,<1.0.0",
   "ag-ui-protocol>=0.1.9,<0.2.0",
+  "pyarrow==20.0.0",
 ]
 
 [project.urls]
@@ -33,6 +35,8 @@ Homepage = "https://github.com/datarobot-oss/datarobot-genai"
 [project.entry-points.'nat.plugins']
 datarobot_llm_providers = "datarobot_genai.nat.datarobot_llm_providers"
 datarobot_llm_clients = "datarobot_genai.nat.datarobot_llm_clients"
+datarobot_mcp_client = "datarobot_genai.nat.datarobot_mcp_client"
+datarobot_auth_provider = "datarobot_genai.nat.datarobot_auth_provider"
 
 [project.optional-dependencies]
 crewai = [
@@ -61,11 +65,13 @@ nat = [
   "nvidia-nat==1.3.0; python_version >= '3.11'",
   "nvidia-nat-opentelemetry==1.3.0; python_version >= '3.11'",
   "nvidia-nat-langchain==1.3.0; python_version >= '3.11'",
+  "nvidia-nat-mcp==1.3.0; python_version >= '3.11'",
   "crewai>=1.1.0; python_version >= '3.11'",
   "llama-index-llms-litellm>=0.4.1,<0.7.0",  # Need this to support datarobot-llm plugin
   "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
   "opentelemetry-instrumentation-llamaindex>=0.40.5,<1.0.0",
   "opentelemetry-instrumentation-langchain>=0.40.5,<1.0.0",
+  "anyio==4.11.0",
 ]
 pydanticai = [
   "logfire>=4.6.0,<5.0.0",
@@ -105,7 +111,7 @@ include = [
 [tool.uv]
 package = true
 override-dependencies = [
-  "ragas>=0.3.8",
+  "ragas>=0.3.8,<0.4.0",
 ]
 exclude-dependencies = [
   "uv",
@@ -122,6 +128,7 @@ dev = [
   "pytest-asyncio>=0.23.5",
   "responses>=0.25.8",
   "ipdb>=0.13.13",
+  "aioresponses>=0.7.8",
 ]
 
 

{datarobot_genai-0.2.0 → datarobot_genai-0.2.8}/src/datarobot_genai/core/utils/auth.py

@@ -12,17 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+import os
 import warnings
 from typing import Any
+from typing import Protocol
 
+import aiohttp
 import jwt
-from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotAsyncOAuthClient
+from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotOAuthClient
+from datarobot.auth.exceptions import OAuthProviderNotFound
+from datarobot.auth.exceptions import OAuthValidationErr
 from datarobot.auth.identity import Identity
-from datarobot.auth.oauth import AsyncOAuthComponent
+from datarobot.auth.oauth import OAuthToken
 from datarobot.auth.session import AuthCtx
 from datarobot.core.config import DataRobotAppFrameworkBaseSettings
 from datarobot.models.genai.agent.auth import ToolAuth
 from datarobot.models.genai.agent.auth import get_authorization_context
+from mypyc.ir.ops import Sequence
 from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
@@ -183,52 +189,203 @@ class AuthContextHeaderHandler:
             return None
 
 
+# --- OAuth Token Provider Implementation ---
+
+
+class TokenRetriever(Protocol):
+    """Protocol for OAuth token retrievers."""
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter identities to only those valid for this retriever implementation."""
+        ...
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Refresh the access token for the given identity ID.
+
+        Parameters
+        ----------
+        identity_id : str
+            The provider identity ID to refresh the token for.
+
+        Returns
+        -------
+        OAuthToken
+            The refreshed OAuth token.
+        """
+        ...
+
+
+class DatarobotTokenRetriever:
+    """Retrieves OAuth tokens using the DataRobot platform."""
+
+    def __init__(self) -> None:
+        self._client = DatarobotOAuthClient()
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter oauth2 identities to only those with provider_identity_id.
+
+        The `provider_identity_id` is required in order to identify the provider
+        and retrieve the access token from DataRobot OAuth Providers service.
+        """
+        return [i for i in identities if i.type == "oauth2" and i.provider_identity_id]
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Refresh the access token using DataRobot's OAuth client."""
+        return await self._client.refresh_access_token(provider_id=identity.provider_identity_id)
+
+
+class AuthlibTokenRetriever:
+    """Retrieves OAuth tokens from a generic Authlib-based endpoint."""
+
+    def __init__(self, application_endpoint: str) -> None:
+        if not application_endpoint:
+            raise ValueError("AuthlibTokenRetriever requires 'application_endpoint'.")
+        self.application_endpoint = application_endpoint.rstrip("/")
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter identities to only OAuth2 identities."""
+        return [i for i in identities if i.type == "oauth2" and i.provider_identity_id is None]
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Retrieve an OAuth token via an HTTP POST request.
+
+        Parameters
+        ----------
+        identity : Identity
+            The identity to retrieve the token for.
+
+        Returns
+        -------
+        OAuthToken
+            The retrieved OAuth token.
+        """
+        api_token = os.environ.get("DATAROBOT_API_TOKEN")
+        if not api_token:
+            raise ValueError("DATAROBOT_API_TOKEN environment variable is required but not set.")
+
+        token_url = f"{self.application_endpoint}/oauth/token/"
+        headers = {"Authorization": f"Bearer {api_token}"}
+        payload = {"identity_id": identity.id}
+        timeout = aiohttp.ClientTimeout(total=30)
+
+        try:
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(token_url, headers=headers, json=payload) as response:
+                    response.raise_for_status()
+                    data = await response.json()
+                    logger.debug(f"Retrieved access token from {token_url}")
+                    return OAuthToken(**data)
+        except aiohttp.ClientError as e:
+            logger.error(f"Error retrieving token from {token_url}: {e}")
+            raise
+
+
+class OAuthConfig(BaseModel):
+    """Configuration extracted from AuthCtx metadata for OAuth operations."""
+
+    implementation: str = "datarobot"
+    application_endpoint: str | None = None
+
+    @classmethod
+    def from_auth_ctx(cls, auth_ctx: AuthCtx) -> "OAuthConfig":
+        metadata = auth_ctx.metadata or {}
+        return cls(
+            implementation=metadata.get("oauth_implementation", "datarobot"),
+            application_endpoint=metadata.get("application_endpoint"),
+        )
+
+
+def create_token_retriever(config: OAuthConfig) -> TokenRetriever:
+    """Create a token retriever based on the OAuth configuration.
+
+    Parameters
+    ----------
+    config : OAuthConfig
+        The OAuth configuration specifying implementation type and endpoints.
+
+    Returns
+    -------
+    TokenRetriever
+        The configured token retriever instance.
+    """
+    if config.implementation == "datarobot":
+        return DatarobotTokenRetriever()
+
+    if config.implementation == "authlib":
+        if not config.application_endpoint:
+            raise ValueError("Required 'application_endpoint' not found in metadata.")
+        return AuthlibTokenRetriever(config.application_endpoint)
+
+    raise ValueError(
+        f"Unsupported OAuth implementation: '{config.implementation}'. "
+        f"Supported values: datarobot, authlib."
+    )
+
+
 class AsyncOAuthTokenProvider:
-    """Manages OAuth access tokens using generic OAuth client."""
+    """Provides OAuth tokens for authorized users.
+
+    This class manages OAuth token retrieval for users with multiple identity providers.
+    It uses either DataRobot or Authlib as the OAuth token storage and refresh backend
+    based on the auth context metadata.
+    """
 
     def __init__(self, auth_ctx: AuthCtx) -> None:
+        """Initialize the provider with an authorization context.
+
+        Parameters
+        ----------
+        auth_ctx : AuthCtx
+            The authorization context containing user identities and metadata.
+        """
         self.auth_ctx = auth_ctx
-        self.oauth_client = self._create_oauth_client()
+        config = OAuthConfig.from_auth_ctx(auth_ctx)
+        self._retriever = create_token_retriever(config)
 
     def _get_identity(self, provider_type: str | None) -> Identity:
-        """Retrieve the appropriate identity from the authentication context."""
-        identities = [x for x in self.auth_ctx.identities if x.provider_identity_id is not None]
-
-        if not identities:
-            raise ValueError("No identities found in authorization context.")
+        """Get identity from auth context, filtered by provider_type if specified."""
+        oauth_identities = self._retriever.filter_identities(self.auth_ctx.identities)
+        if not oauth_identities:
+            raise OAuthProviderNotFound("No OAuth provider found.")
 
         if provider_type is None:
-            if len(identities) > 1:
-                raise ValueError(
-                    "Multiple identities found. Please specify 'provider_type' parameter."
+            if len(oauth_identities) > 1:
+                raise OAuthValidationErr(
+                    "Multiple OAuth providers found. Specify 'provider_type' parameter."
                 )
-            return identities[0]
-
-        identity = next((id for id in identities if id.provider_type == provider_type), None)
+            return oauth_identities[0]
 
+        identity = next((i for i in oauth_identities if i.provider_type == provider_type), None)
         if identity is None:
-            raise ValueError(f"No identity found for provider '{provider_type}'.")
-
+            raise OAuthValidationErr(f"No identity found for provider '{provider_type}'.")
         return identity
 
     async def get_token(self, auth_type: ToolAuth, provider_type: str | None = None) -> str:
-        """Get OAuth access token using the specified method."""
+        """Get an OAuth access token for the specified auth type and provider.
+
+        Parameters
+        ----------
+        auth_type : ToolAuth
+            Authentication type (only OBO is supported).
+        provider_type : str, optional
+            The specific provider to use (e.g., 'google'). Required if multiple
+            identities are available.
+
+        Returns
+        -------
+        str
+            The retrieved OAuth access token.
+
+        Raises
+        ------
+        ValueError
+            If the auth type is unsupported or if a suitable identity cannot be found.
+        """
         if auth_type != ToolAuth.OBO:
             raise ValueError(
-                f"Unsupported auth type: {auth_type}. Only {ToolAuth.OBO} is supported."
+                f"Unsupported auth type: {auth_type}. Only OBO (on-behalf-of) is supported."
             )
 
         identity = self._get_identity(provider_type)
-        token_data = await self.oauth_client.refresh_access_token(
-            identity_id=identity.provider_identity_id
-        )
+        token_data = await self._retriever.refresh_access_token(identity)
         return token_data.access_token
-
-    def _create_oauth_client(self) -> AsyncOAuthComponent:
-        """Create either DataRobot or Authlib OAuth client based on
-        authorization context.
-
-        Note: at the moment, only DataRobot OAuth client is supported.
-        """
-        logger.debug("Using DataRobot OAuth client")
-        return DatarobotAsyncOAuthClient()

{datarobot_genai-0.2.0 → datarobot_genai-0.2.8}/src/datarobot_genai/drmcp/core/config.py

@@ -197,6 +197,54 @@ class MCPServerConfig(BaseSettings):
         description="Enable/disable predictive tools",
     )
 
+    # Jira 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",
+    )
+    is_jira_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
+            "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Jira OAuth provider is configured for Jira integration",
+    )
+
+    @property
+    def is_jira_oauth_configured(self) -> bool:
+        return self.is_jira_oauth_provider_configured or bool(
+            os.getenv("JIRA_CLIENT_ID") and os.getenv("JIRA_CLIENT_SECRET")
+        )
+
+    # Confluence 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",
+    )
+    is_confluence_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Confluence OAuth provider is configured for Confluence integration",
+    )
+
+    @property
+    def is_confluence_oauth_configured(self) -> bool:
+        return self.is_confluence_oauth_provider_configured or bool(
+            os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
+        )
+
     @field_validator(
         "otel_attributes",
         mode="before",
@@ -220,6 +268,10 @@ class MCPServerConfig(BaseSettings):
         "tool_registration_duplicate_behavior",
         "mcp_server_register_dynamic_prompts_on_startup",
         "enable_predictive_tools",
+        "enable_jira_tools",
+        "is_jira_oauth_provider_configured",
+        "enable_confluence_tools",
+        "is_confluence_oauth_provider_configured",
         mode="before",
     )
     @classmethod

{datarobot_genai-0.2.0 → datarobot_genai-0.2.8}/src/datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -40,6 +40,8 @@ from .routes_utils import prefix_mount_path
 from .server_life_cycle import BaseServerLifecycle
 from .telemetry import OtelASGIMiddleware
 from .telemetry import initialize_telemetry
+from .tool_config import TOOL_CONFIGS
+from .tool_config import is_tool_enabled
 
 
 def _import_modules_from_dir(
@@ -115,6 +117,9 @@ class DataRobotMCPServer:
         self._mcp = mcp
         self._mcp_transport = transport
 
+        # Configure MCP server capabilities
+        self._configure_mcp_capabilities()
+
         # Initialize telemetry
         initialize_telemetry(mcp)
 
@@ -139,11 +144,12 @@ class DataRobotMCPServer:
 
         # Load static tools modules
         base_dir = os.path.dirname(os.path.dirname(__file__))
-        if self._config.enable_predictive_tools:
-            _import_modules_from_dir(
-                os.path.join(base_dir, "tools", "predictive"),
-                "datarobot_genai.drmcp.tools.predictive",
-            )
+        for tool_type, tool_config in TOOL_CONFIGS.items():
+            if is_tool_enabled(tool_type, self._config):
+                _import_modules_from_dir(
+                    os.path.join(base_dir, "tools", tool_config["directory"]),
+                    tool_config["package_prefix"],
+                )
 
         # Load memory management tools if available
         if self._memory_manager:
@@ -163,6 +169,37 @@ class DataRobotMCPServer:
         if transport == "streamable-http":
             register_routes(self._mcp)
 
+    def _configure_mcp_capabilities(self) -> None:
+        """Configure MCP capabilities that FastMCP doesn't expose directly.
+
+        See: https://github.com/modelcontextprotocol/python-sdk/issues/1126
+        """
+        server = self._mcp._mcp_server
+
+        # Declare prompts_changed capability  (capabilities.prompts.listChanged: true)
+        server.notification_options.prompts_changed = True
+
+        # Declare experimental capabilities ( experimental.dynamic_prompts: true)
+        server.experimental_capabilities = {"dynamic_prompts": {"enabled": True}}
+
+        # Patch to include experimental_capabilities (FastMCP doesn't expose this)
+        original = server.create_initialization_options
+
+        def patched(
+            notification_options: Any = None,
+            experimental_capabilities: dict[str, dict[str, Any]] | None = None,
+            **kwargs: Any,
+        ) -> Any:
+            if experimental_capabilities is None:
+                experimental_capabilities = getattr(server, "experimental_capabilities", None)
+            return original(
+                notification_options=notification_options,
+                experimental_capabilities=experimental_capabilities,
+                **kwargs,
+            )
+
+        server.create_initialization_options = patched
+
     def run(self, show_banner: bool = False) -> None:
         """Run the DataRobot MCP server synchronously."""
         try:
@@ -179,6 +216,9 @@ class DataRobotMCPServer:
                 self._logger.info("Registering dynamic prompts from prompt management...")
                 asyncio.run(register_prompts_from_datarobot_prompt_management())
 
+            # Execute pre-server start actions
+            asyncio.run(self._lifecycle.pre_server_start(self._mcp))
+
             # List registered tools, prompts, and resources before starting server
             tools = asyncio.run(self._mcp._list_tools_mcp())
             prompts = asyncio.run(self._mcp._list_prompts_mcp())
@@ -198,9 +238,6 @@ class DataRobotMCPServer:
             for resource in resources:
                 self._logger.info(f" > {resource.name}")
 
-            # Execute pre-server start actions
-            asyncio.run(self._lifecycle.pre_server_start(self._mcp))
-
             # Create event loop for async operations
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)

datarobot_genai-0.2.8/src/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-0.2.0 → datarobot_genai-0.2.8}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py

@@ -18,15 +18,13 @@ 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 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
 
@@ -57,7 +55,8 @@ async def register_prompts_from_datarobot_prompt_management() -> None:
 
 
 async def register_prompt_from_datarobot_prompt_management(
-    prompt_template: DrPrompt, prompt_template_version: DrPromptVersion | None = None
+    prompt_template: dr.genai.PromptTemplate,
+    prompt_template_version: dr.genai.PromptTemplateVersion | None = None,
 ) -> Prompt:
     """Register a single prompt.
 
@@ -173,7 +172,7 @@ def to_valid_mcp_prompt_name(s: str) -> str:
 
 
 def make_prompt_function(
-    name: str, description: str, prompt_text: str, variables: list[DrVariable]
+    name: str, description: str, prompt_text: str, variables: list[dr.genai.Variable]
 ) -> Callable:
     params = []
     for v in variables: