datarobot-genai 0.2.0__py3-none-any.whl

datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -0,0 +1,316 @@
+# 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 asyncio
+import glob
+import importlib
+import logging
+import os
+from collections.abc import Callable
+from typing import Any
+
+from fastmcp import FastMCP
+from starlette.middleware import Middleware
+
+from .auth import initialize_oauth_middleware
+from .config import get_config
+from .credentials import get_credentials
+from .dr_mcp_server_logo import log_server_custom_banner
+from .dynamic_prompts.register import register_prompts_from_datarobot_prompt_management
+from .dynamic_tools.deployment.register import register_tools_of_datarobot_deployments
+from .logging import MCPLogging
+from .mcp_instance import mcp
+from .mcp_server_tools import get_all_available_tags  # noqa # pylint: disable=unused-import
+from .mcp_server_tools import get_tool_info_by_name  # noqa # pylint: disable=unused-import
+from .mcp_server_tools import list_tools_by_tags  # noqa # pylint: disable=unused-import
+from .memory_management.manager import MemoryManager
+from .routes import register_routes
+from .routes_utils import prefix_mount_path
+from .server_life_cycle import BaseServerLifecycle
+from .telemetry import OtelASGIMiddleware
+from .telemetry import initialize_telemetry
+
+
+def _import_modules_from_dir(
+    directory: str, package_prefix: str, module_name: str | None = None
+) -> None:
+    """Dynamically import all modules from a directory."""
+    if not os.path.exists(directory):
+        return
+    if module_name:
+        module_name = f"{package_prefix}.{module_name}"
+        try:
+            importlib.import_module(module_name)
+        except ImportError as e:
+            logging.warning(f"Failed to import module {module_name}: {e}")
+    else:
+        for file in glob.glob(os.path.join(directory, "*.py")):
+            if os.path.basename(file) != "__init__.py":
+                module_name = f"{package_prefix}.{os.path.splitext(os.path.basename(file))[0]}"
+                try:
+                    importlib.import_module(module_name)
+                except ImportError as e:
+                    logging.warning(f"Failed to import module {module_name}: {e}")
+
+
+class DataRobotMCPServer:
+    """
+    DataRobot MCP server implementation using FastMCP framework.
+
+    This server can be extended by providing custom configuration, credentials,
+    and lifecycle handlers.
+    """
+
+    def __init__(
+        self,
+        mcp: FastMCP,
+        transport: str = "streamable-http",
+        config_factory: Callable[[], Any] | None = None,
+        credentials_factory: Callable[[], Any] | None = None,
+        lifecycle: BaseServerLifecycle | None = None,
+        additional_module_paths: list[tuple[str, str]] | None = None,
+    ):
+        """
+        Initialize the server.
+
+        Args:
+            mcp: FastMCP instance
+            transport: Transport type ("streamable-http" or "stdio")
+            config_factory: Optional factory function for user config
+            credentials_factory: Optional factory function for user credentials
+            lifecycle: Optional lifecycle handler (defaults to BaseServerLifecycle())
+            additional_module_paths: Optional list of (directory, package_prefix) tuples for
+                loading additional modules
+        """
+        # Initialize config and logging
+        self._config = get_config()
+        MCPLogging(self._config.app_log_level)
+        self._logger = logging.getLogger(self.__class__.__name__)
+        self._logger.info(f"Config initialized: {self._config}")
+
+        # Initialize credentials
+        self._credentials = get_credentials()
+        self._logger.info("Credentials initialized")
+
+        self._user_config = config_factory() if config_factory else None
+        self._logger.info(f"User config initialized: {self._user_config}")
+        self._user_credentials = credentials_factory() if credentials_factory else None
+        self._logger.info("User credentials initialized")
+
+        # Initialize lifecycle
+        self._lifecycle = lifecycle if lifecycle else BaseServerLifecycle()
+        self._logger.info("Lifecycle initialized")
+
+        self._mcp = mcp
+        self._mcp_transport = transport
+
+        # Initialize telemetry
+        initialize_telemetry(mcp)
+
+        # Initialize OAuth middleware
+        initialize_oauth_middleware(mcp)
+
+        # Initialize memory manager if AWS credentials are available
+        self._memory_manager: MemoryManager | None = None
+        if self._config.enable_memory_management:
+            if self._credentials.has_aws_credentials():
+                self._logger.info("Initializing memory manager")
+                try:
+                    self._memory_manager = MemoryManager.get_instance()
+                except Exception as e:
+                    self._logger.error(f"Error initializing memory manager: {e}")
+                    self._logger.info("Skipping memory manager initialization")
+                    self._memory_manager = None
+            else:
+                self._logger.info(
+                    "No AWS credentials found, skipping memory manager initialization"
+                )
+
+        # 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",
+            )
+
+        # Load memory management tools if available
+        if self._memory_manager:
+            _import_modules_from_dir(
+                directory=os.path.join(base_dir, "core", "memory_management"),
+                package_prefix="datarobot_genai.drmcp.core.memory_management",
+                module_name="memory_tools",
+            )
+
+        # Load additional recipe user modules if provided
+        if additional_module_paths:
+            for directory, package_prefix in additional_module_paths:
+                self._logger.info(f"Loading additional modules from {directory}")
+                _import_modules_from_dir(directory, package_prefix)
+
+        # Register HTTP routes if using streamable-http transport
+        if transport == "streamable-http":
+            register_routes(self._mcp)
+
+    def run(self, show_banner: bool = False) -> None:
+        """Run the DataRobot MCP server synchronously."""
+        try:
+            # Validate configuration
+            if not self._credentials.has_datarobot_credentials():
+                self._logger.error("DataRobot credentials not configured")
+                raise ValueError("Missing required DataRobot credentials")
+
+            if self._config.mcp_server_register_dynamic_tools_on_startup:
+                self._logger.info("Registering dynamic tools from deployments...")
+                asyncio.run(register_tools_of_datarobot_deployments())
+
+            if self._config.mcp_server_register_dynamic_prompts_on_startup:
+                self._logger.info("Registering dynamic prompts from prompt management...")
+                asyncio.run(register_prompts_from_datarobot_prompt_management())
+
+            # 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())
+            resources = asyncio.run(self._mcp._list_resources_mcp())
+
+            tools_count = len(tools)
+            prompts_count = len(prompts)
+            resources_count = len(resources)
+
+            self._logger.info(f"Registered tools: {tools_count}")
+            for tool in tools:
+                self._logger.info(f" > {tool.name}")
+            self._logger.info(f"Registered prompts: {prompts_count}")
+            for prompt in prompts:
+                self._logger.info(f" > {prompt.name}")
+            self._logger.info(f"Registered resources: {resources_count}")
+            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)
+
+            async def run_server(show_banner: bool = show_banner) -> None:
+                # Start server in background based on transport type
+
+                if show_banner:
+                    log_server_custom_banner(
+                        self._mcp,
+                        self._mcp_transport,
+                        port=self._config.mcp_server_port,
+                        tools_count=tools_count,
+                        prompts_count=prompts_count,
+                        resources_count=resources_count,
+                    )
+
+                if self._mcp_transport == "stdio":
+                    server_task = asyncio.create_task(self._mcp.run_stdio_async(show_banner=False))
+                elif self._mcp_transport == "streamable-http":
+                    server_task = asyncio.create_task(
+                        self._mcp.run_http_async(
+                            transport="http",
+                            middleware=[Middleware(OtelASGIMiddleware)],
+                            show_banner=False,
+                            port=self._config.mcp_server_port,
+                            log_level=self._config.mcp_server_log_level,
+                            host=self._config.mcp_server_host,
+                            stateless_http=True,
+                            path=prefix_mount_path("/mcp"),
+                        )
+                    )
+                else:
+                    raise ValueError(f"Unsupported transport: {self._mcp_transport}")
+
+                # Give the server a moment to initialize
+                await asyncio.sleep(1)
+
+                # Execute post-server start actions
+                await self._lifecycle.post_server_start(self._mcp)
+
+                # Wait for server to complete
+                await server_task
+
+            # Start the server
+            self._logger.info("Starting MCP server...")
+            try:
+                loop.run_until_complete(run_server(show_banner=show_banner))
+            except KeyboardInterrupt:
+                self._logger.info("Server interrupted by user")
+            finally:
+                # Execute pre-shutdown actions
+                self._logger.info("Shutting down server...")
+                loop.run_until_complete(self._lifecycle.pre_server_shutdown(self._mcp))
+                loop.close()
+
+        except Exception as e:
+            self._logger.error(f"Server error: {e}")
+            raise
+
+
+def create_mcp_server(
+    config_factory: Callable[[], Any] | None = None,
+    credentials_factory: Callable[[], Any] | None = None,
+    lifecycle: BaseServerLifecycle | None = None,
+    additional_module_paths: list[tuple[str, str]] | None = None,
+    transport: str = "streamable-http",
+) -> DataRobotMCPServer:
+    """
+    Create a DataRobot MCP server.
+
+    Args:
+        config_factory: Optional factory function for user config
+        credentials_factory: Optional factory function for user credentials
+        lifecycle: Optional lifecycle handler
+        additional_module_paths: Optional list of (directory, package_prefix) tuples
+        transport: Transport type ("streamable-http" or "stdio")
+
+    Returns
+    -------
+        Configured DataRobotMCPServer instance
+
+    Example:
+        ```python
+        # Basic usage with defaults
+        server = create_mcp_server()
+        server.run()
+
+        # With custom configuration
+        from myapp.config import get_my_config
+        from myapp.lifecycle import MyLifecycle
+
+        server = create_mcp_server(
+            config_factory=get_my_config,
+            lifecycle=MyLifecycle(),
+            additional_module_paths=[
+                ("/path/to/my/tools", "myapp.tools"),
+                ("/path/to/my/prompts", "myapp.prompts"),
+            ]
+        )
+        server.run()
+        ```
+    """
+    # Use the global mcp instance that tools are registered with
+
+    return DataRobotMCPServer(
+        mcp=mcp,
+        transport=transport,
+        config_factory=config_factory,
+        credentials_factory=credentials_factory,
+        lifecycle=lifecycle,
+        additional_module_paths=additional_module_paths,
+    )

datarobot_genai/drmcp/core/dr_mcp_server_logo.py

@@ -0,0 +1,136 @@
+# 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 fastmcp import FastMCP
+from rich.align import Align
+from rich.console import Console
+from rich.console import Group
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+from datarobot_genai import __version__ as datarobot_genai_version
+
+
+# Green color #81FBA5 = RGB(129, 251, 165)
+def _apply_green(text: str) -> str:
+    """Apply green color #81FBA5 to all characters in the text."""
+    # Apply ANSI escape code for RGB color #81FBA5 (129, 251, 165)
+    green_start = "\x1b[38;2;129;251;165m"
+    green_end = "\x1b[39m"
+    # Wrap the entire text (except trailing newline) with green color codes
+    lines = text.split("\n")
+    colored_lines = [green_start + line + green_end if line else "" for line in lines]
+    return "\n".join(colored_lines)
+
+
+DR_LOGO_ASCII = _apply_green(r"""
+ ____        _        ____       _           _   
+|  _ \  __ _| |_ __ _|  _ \ ___ | |__   ___ | |_ 
+| | | |/ _` | __/ _` | |_) / _ \| '_ \ / _ \| __|
+| |_| | (_| | || (_| |  _ < (_) | |_) | (_) | |_ 
+|____/ \__,_|\__\__,_|_| \_\___/|_.__/ \___/ \__|
+""")
+
+
+def log_server_custom_banner(
+    server: FastMCP[Any],
+    transport: str,
+    *,
+    host: str | None = None,
+    port: int | None = None,
+    path: str | None = None,
+    tools_count: int | None = None,
+    prompts_count: int | None = None,
+    resources_count: int | None = None,
+) -> None:
+    """
+    Create and log a formatted banner with server information and logo.
+
+    Args:
+        transport: The transport protocol being used
+        server_name: Optional server name to display
+        host: Host address (for HTTP transports)
+        port: Port number (for HTTP transports)
+        path: Server path (for HTTP transports)
+        tools_count: Number of tools registered
+        prompts_count: Number of prompts registered
+        resources_count: Number of resources registered
+    """
+    # Create the logo text
+    # Use Text with no_wrap and markup disabled to preserve ANSI escape codes
+    logo_text = Text.from_ansi(DR_LOGO_ASCII, no_wrap=True)
+
+    # Create the main title
+    title_text = Text(f"DataRobot MCP Server {datarobot_genai_version}", style="dim green")
+    stats_text = Text(
+        f"{tools_count} tools, {prompts_count} prompts, {resources_count} resources",
+        style="bold green",
+    )
+
+    # Create the information table
+    info_table = Table.grid(padding=(0, 1))
+    info_table.add_column(style="bold", justify="center")  # Emoji column
+    info_table.add_column(style="cyan", justify="left")  # Label column
+    info_table.add_column(style="dim", justify="left")  # Value column
+
+    match transport:
+        case "http" | "streamable-http":
+            display_transport = "HTTP"
+        case "sse":
+            display_transport = "SSE"
+        case "stdio":
+            display_transport = "STDIO"
+
+    info_table.add_row("🖥", "Server name:", Text(server.name + "\n", style="bold blue"))
+    info_table.add_row("📦", "Transport:", display_transport)
+    info_table.add_row("🌐", "MCP port:", str(port))
+
+    # Show connection info based on transport
+    if transport in ("http", "streamable-http", "sse") and host and port:
+        server_url = f"http://{host}:{port}"
+        if path:
+            server_url += f"/{path.lstrip('/')}"
+        info_table.add_row("🔗", "Server URL:", server_url)
+
+    # Add documentation link
+    info_table.add_row("", "", "")
+    info_table.add_row("📚", "Docs:", "https://docs.datarobot.com")
+    info_table.add_row("🚀", "Hosting:", "https://datarobot.com")
+
+    # Create panel with logo, title, and information using Group
+    panel_content = Group(
+        Align.center(logo_text),
+        "",
+        Align.center(title_text),
+        Align.center(stats_text),
+        "",
+        "",
+        Align.center(info_table),
+    )
+
+    panel = Panel(
+        panel_content,
+        border_style="dim",
+        padding=(1, 4),
+        # expand=False,
+        width=80,  # Set max width for the pane
+    )
+
+    console = Console(stderr=True)
+    # Center the panel itself
+    console.print(Group("\n", Align.center(panel), "\n"))

datarobot_genai/drmcp/core/dynamic_prompts/__init__.py

@@ -0,0 +1,13 @@
+# 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_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)