datarobot-genai 0.1.67__py3-none-any.whl → 0.1.69__py3-none-any.whl

datarobot_genai/core/mcp/common.py

@@ -13,16 +13,20 @@
 # limitations under the License.
 
 import json
+import logging
 import re
+from http import HTTPStatus
 from typing import Any
 from typing import Literal
-from urllib.parse import urlparse
 
+import requests
 from datarobot.core.config import DataRobotAppFrameworkBaseSettings
 from pydantic import field_validator
 
 from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
 
+logger = logging.getLogger(__name__)
+
 
 class MCPConfig(DataRobotAppFrameworkBaseSettings):
     """Configuration for MCP server connection.
@@ -39,6 +43,7 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
     datarobot_api_token: str | None = None
     authorization_context: dict[str, Any] | None = None
     forwarded_headers: dict[str, str] | None = None
+    mcp_server_port: int | None = None
 
     _auth_context_handler: AuthContextHeaderHandler | None = None
     _server_config: dict[str, Any] | None = None
@@ -49,17 +54,14 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
         if value is None:
             return None
 
-        if not isinstance(value, str):
-            msg = "external_mcp_headers must be a JSON string"
-            raise TypeError(msg)
-
         candidate = value.strip()
 
         try:
             json.loads(candidate)
-        except json.JSONDecodeError as exc:
+        except json.JSONDecodeError:
             msg = "external_mcp_headers must be valid JSON"
-            raise ValueError(msg) from exc
+            logger.warning(msg)
+            return None
 
         return candidate
 
@@ -69,15 +71,12 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
         if value is None:
             return None
 
-        if not isinstance(value, str):
-            msg = "mcp_deployment_id must be a string"
-            raise TypeError(msg)
-
         candidate = value.strip()
 
         if not re.fullmatch(r"[0-9a-fA-F]{24}", candidate):
             msg = "mcp_deployment_id must be a valid 24-character hex ID"
-            raise ValueError(msg)
+            logger.warning(msg)
+            return None
 
         return candidate
 
@@ -112,6 +111,45 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
             # Authorization context not available (e.g., in tests)
             return {}
 
+    def _build_authenticated_headers(self) -> dict[str, str]:
+        """Build headers for authenticated requests.
+
+        Returns
+        -------
+            Dictionary containing forwarded headers (if available) and authentication headers.
+        """
+        headers: dict[str, str] = {}
+        if self.forwarded_headers:
+            headers.update(self.forwarded_headers)
+        headers.update(self._authorization_bearer_header())
+        headers.update(self._authorization_context_header())
+        return headers
+
+    def _check_localhost_server(self, url: str, timeout: float = 2.0) -> bool:
+        """Check if MCP server is running on localhost.
+
+        Parameters
+        ----------
+        url : str
+            The URL to check.
+        timeout : float, optional
+            Request timeout in seconds (default: 2.0).
+
+        Returns
+        -------
+        bool
+            True if server is running and responding with OK status, False otherwise.
+        """
+        try:
+            response = requests.get(url, timeout=timeout)
+            return (
+                response.status_code == HTTPStatus.OK
+                and response.json().get("message") == "DataRobot MCP Server is running"
+            )
+        except requests.RequestException as e:
+            logger.debug(f"Failed to connect to MCP server at {url}: {e}")
+            return False
+
     def _build_server_config(self) -> dict[str, Any] | None:
         """
         Get MCP server configuration.
@@ -121,34 +159,6 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
             Server configuration dict with url, transport, and optional headers,
             or None if not configured.
         """
-        if self.external_mcp_url:
-            # External MCP URL - no authentication needed
-            headers: dict[str, str] = {}
-
-            # Forward headers for localhost connections
-            if self.forwarded_headers:
-                try:
-                    parsed_url = urlparse(self.external_mcp_url)
-                    hostname = parsed_url.hostname or ""
-                    # Check if hostname is localhost or 127.0.0.1
-                    if hostname in ("localhost", "127.0.0.1", "::1"):
-                        headers.update(self.forwarded_headers)
-                except Exception:
-                    # If URL parsing fails, fall back to simple string check
-                    if "localhost" in self.external_mcp_url or "127.0.0.1" in self.external_mcp_url:
-                        headers.update(self.forwarded_headers)
-
-            # Merge external headers if provided
-            if self.external_mcp_headers:
-                external_headers = json.loads(self.external_mcp_headers)
-                headers.update(external_headers)
-
-            return {
-                "url": self.external_mcp_url.rstrip("/"),
-                "transport": self.external_mcp_transport,
-                "headers": headers,
-            }
-
         if self.mcp_deployment_id:
             # DataRobot deployment ID - requires authentication
             if self.datarobot_endpoint is None:
@@ -165,15 +175,9 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
                 base_url = f"{base_url}/api/v2"
 
             url = f"{base_url}/deployments/{self.mcp_deployment_id}/directAccess/mcp"
+            headers = self._build_authenticated_headers()
 
-            # Start with forwarded headers if available
-            headers = {}
-            if self.forwarded_headers:
-                headers.update(self.forwarded_headers)
-
-            # Add authentication headers
-            headers.update(self._authorization_bearer_header())
-            headers.update(self._authorization_context_header())
+            logger.info(f"Using DataRobot hosted MCP deployment: {url}")
 
             return {
                 "url": url,
@@ -181,4 +185,34 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
                 "headers": headers,
             }
 
+        if self.external_mcp_url:
+            # External MCP URL - no authentication needed
+            headers = {}
+
+            # Merge external headers if provided
+            if self.external_mcp_headers:
+                external_headers = json.loads(self.external_mcp_headers)
+                headers.update(external_headers)
+
+            logger.info(f"Using external MCP URL: {self.external_mcp_url}")
+
+            return {
+                "url": self.external_mcp_url.rstrip("/"),
+                "transport": self.external_mcp_transport,
+                "headers": headers,
+            }
+
+        # No MCP configuration found, setup localhost if running locally
+        if self.mcp_server_port:
+            url = f"http://localhost:{self.mcp_server_port}"
+            if self._check_localhost_server(url):
+                headers = self._build_authenticated_headers()
+                logger.info(f"Using localhost MCP server: {url}")
+                return {
+                    "url": f"{url}/mcp",
+                    "transport": "streamable-http",
+                    "headers": headers,
+                }
+            logger.warning(f"MCP server is not running or not responding at {url}")
+
         return None

datarobot_genai/drmcp/core/dynamic_prompts/controllers.py

@@ -18,6 +18,8 @@ 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,
 )
@@ -83,3 +85,46 @@ async def delete_registered_prompt_template(prompt_template_id: str) -> bool:
         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/routes.py

@@ -19,6 +19,7 @@ from starlette.requests import Request
 from starlette.responses import JSONResponse
 
 from .dynamic_prompts.controllers import delete_registered_prompt_template
+from .dynamic_prompts.controllers import refresh_registered_prompt_template
 from .dynamic_prompts.controllers import register_prompt_from_prompt_template_id_and_version
 from .dynamic_tools.deployment.controllers import delete_registered_tool_deployment
 from .dynamic_tools.deployment.controllers import get_registered_tool_deployments
@@ -418,6 +419,18 @@ def register_routes(mcp: TaggedFastMCP) -> None:
             )
         except Exception as e:
             return JSONResponse(
-                status_code=HTTPStatus.BAD_REQUEST,
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
                 content={"error": f"Failed to add prompt template: {str(e)}"},
             )
+
+    @mcp.custom_route(prefix_mount_path("/registeredPrompts"), methods=["PUT"])
+    async def refresh_prompt_templates(_: Request) -> JSONResponse:
+        """Refresh prompt templates."""
+        try:
+            await refresh_registered_prompt_template()
+            return JSONResponse(status_code=HTTPStatus.NO_CONTENT, content=None)
+        except Exception as e:
+            return JSONResponse(
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                content={"error": f"Failed to refresh prompt templates: {str(e)}"},
+            )

{datarobot_genai-0.1.67.dist-info → datarobot_genai-0.1.69.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.67
+Version: 0.1.69
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.

{datarobot_genai-0.1.67.dist-info → datarobot_genai-0.1.69.dist-info}/RECORD

@@ -13,7 +13,7 @@ datarobot_genai/core/cli/__init__.py,sha256=B93Yb6VavoZpatrh8ltCL6YglIfR5FHgytXb
 datarobot_genai/core/cli/agent_environment.py,sha256=BJzQoiDvZF5gW4mFE71U0yeg-l72C--kxiE-fv6W194,1662
 datarobot_genai/core/cli/agent_kernel.py,sha256=3XX58DQ6XPpWB_tn5m3iGb3XTfhZf5X3W9tc6ADieU4,7790
 datarobot_genai/core/mcp/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/core/mcp/common.py,sha256=Y24Ij8e4BwOnxojnz2BJRIFcs_Y7EfbORLrN6txowpY,6833
+datarobot_genai/core/mcp/common.py,sha256=Y8SjuquUODKEfI7T9X-QuTMKdIlpCWFI1b3xs6tmHFA,7812
 datarobot_genai/core/utils/__init__.py,sha256=VxtRUz6iwb04eFQQy0zqTNXLAkYpPXcJxVoKV0nOdXk,59
 datarobot_genai/core/utils/auth.py,sha256=Xo1PxVr6oMgtMHkmHdS02klDKK1cyDpjGvIMF4Tx0Lo,7874
 datarobot_genai/core/utils/urls.py,sha256=tk0t13duDEPcmwz2OnS4vwEdatruiuX8lnxMMhSaJik,2289
@@ -37,14 +37,14 @@ datarobot_genai/drmcp/core/exceptions.py,sha256=eqsGI-lxybgvWL5w4BFhbm3XzH1eU5te
 datarobot_genai/drmcp/core/logging.py,sha256=Y_hig4eBWiXGaVV7B_3wBcaYVRNH4ydptbEQhrP9-mY,3414
 datarobot_genai/drmcp/core/mcp_instance.py,sha256=wMsP39xqTmNBYqd49olEQb5UHTSsxj6BOIoIElorRB0,19235
 datarobot_genai/drmcp/core/mcp_server_tools.py,sha256=odNZKozfx0VV38SLZHw9lY0C0JM_JnRI06W3BBXnyE4,4278
-datarobot_genai/drmcp/core/routes.py,sha256=zW4Ukr9f_odF-_oXq_WlTaCH5P4mb7gBSFp67IiuAoo,17208
+datarobot_genai/drmcp/core/routes.py,sha256=nrgzYkiWWcwSy0BUJ-k5AuXdUHNrnGRudCo4AhtYvlY,17856
 datarobot_genai/drmcp/core/routes_utils.py,sha256=vSseXWlplMSnRgoJgtP_rHxWSAVYcx_tpTv4lyTpQoc,944
 datarobot_genai/drmcp/core/server_life_cycle.py,sha256=WKGJWGxalvqxupzJ2y67Kklc_9PgpZT0uyjlv_sr5wc,3419
 datarobot_genai/drmcp/core/telemetry.py,sha256=NEkSTC1w6uQgtukLHI-sWvR4EMgInysgATcvfQ5CplM,15378
 datarobot_genai/drmcp/core/tool_filter.py,sha256=tLOcG50QBvS48cOVHM6OqoODYiiS6KeM_F-2diaHkW0,2858
 datarobot_genai/drmcp/core/utils.py,sha256=dSjrayWVcnC5GxQcvOIOSHaoEymPIVtG_s2ZBMlmSOw,4336
 datarobot_genai/drmcp/core/dynamic_prompts/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
-datarobot_genai/drmcp/core/dynamic_prompts/controllers.py,sha256=vCMYxwYjNDadRxSlRk6p8pHK6h_2K-PkbBTTW_lqBJ0,3318
+datarobot_genai/drmcp/core/dynamic_prompts/controllers.py,sha256=AGJlKqgHRO0Kd7Gl-Ulw9KYBgzjTTFXWBvOUF-SuKUI,5454
 datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py,sha256=IEdD2Gqm4SfUdiXJB99RiWxkN6frGaxJ2SfATetMM3c,4243
 datarobot_genai/drmcp/core/dynamic_prompts/register.py,sha256=5AEh1m8GX-gPZHUdiE1VATt7IKJQk-eThcxh01sWn0I,7204
 datarobot_genai/drmcp/core/dynamic_prompts/utils.py,sha256=BZ3792AgfvYlwL0_J0MzQfGecyEA5_OKUMynEZYzCds,1136
@@ -93,9 +93,9 @@ datarobot_genai/nat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 datarobot_genai/nat/agent.py,sha256=siBLDWAff2-JwZ8Q3iNpM_e4_IoSwG9IvY0hyEjNenw,10292
 datarobot_genai/nat/datarobot_llm_clients.py,sha256=IZq_kooUL8QyDTkpEreszLQk9vCzg6-FbTjIkXR9wc0,7203
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=lOVaL_0Fl6-7GFYl3HmfqttqKpKt-2w8o92P3T7B6cU,3683
-datarobot_genai-0.1.67.dist-info/METADATA,sha256=hfSJaoB-QbjSH8wEMUAYxqzmdyudEz8IYdqG53tRdb4,5918
-datarobot_genai-0.1.67.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.1.67.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
-datarobot_genai-0.1.67.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.1.67.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.1.67.dist-info/RECORD,,
+datarobot_genai-0.1.69.dist-info/METADATA,sha256=ANM5ghmbD_jGo8UHS87Y63jcMZlZY4Pnm7suIWRdtsg,5918
+datarobot_genai-0.1.69.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.1.69.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
+datarobot_genai-0.1.69.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.1.69.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.1.69.dist-info/RECORD,,