datarobot-genai 0.2.0__py3-none-any.whl

datarobot_genai/drmcp/core/utils.py

@@ -0,0 +1,131 @@
+# 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 base64
+import uuid
+from typing import Any
+
+import boto3
+from fastmcp.resources import HttpResource
+from fastmcp.tools.tool import ToolResult
+from pydantic import BaseModel
+
+from .constants import MAX_INLINE_SIZE
+from .mcp_instance import mcp
+
+
+def generate_presigned_url(bucket: str, key: str, expires_in: int = 2592000) -> str:
+    """
+    Generate a presigned S3 URL for the given bucket and key.
+    Args:
+        bucket (str): S3 bucket name.
+        key (str): S3 object key.
+        expires_in (int): Expiration in seconds (default 30 days).
+
+    Returns
+    -------
+        str: Presigned S3 URL for get_object.
+    """
+    s3 = boto3.client("s3")
+    result = s3.generate_presigned_url(
+        "get_object", Params={"Bucket": bucket, "Key": key}, ExpiresIn=expires_in
+    )
+    return str(result)
+
+
+class PredictionResponse(BaseModel):
+    type: str
+    data: str | None = None
+    resource_id: str | None = None
+    s3_url: str | None = None
+    show_explanations: bool | None = None
+
+
+def predictions_result_response(
+    df: Any, bucket: str, key: str, resource_name: str, show_explanations: bool = False
+) -> PredictionResponse:
+    csv_str = df.to_csv(index=False)
+    if len(csv_str.encode("utf-8")) < MAX_INLINE_SIZE:
+        return PredictionResponse(type="inline", data=csv_str, show_explanations=show_explanations)
+    else:
+        resource = save_df_to_s3_and_register_resource(df, bucket, key, resource_name)
+        return PredictionResponse(
+            type="resource",
+            resource_id=str(resource.uri),
+            s3_url=resource.url,
+            show_explanations=show_explanations,
+        )
+
+
+def save_df_to_s3_and_register_resource(
+    df: Any, bucket: str, key: str, resource_name: str, mime_type: str = "text/csv"
+) -> HttpResource:
+    """
+    Save a DataFrame to a temp CSV, upload to S3, register as a resource, and return the
+    presigned URL.
+    Args:
+        df (pd.DataFrame): DataFrame to save and upload.
+        bucket (str): S3 bucket name.
+        key (str): S3 object key.
+        resource_name (str): Name for the registered resource.
+        mime_type (str): MIME type for the resource (default 'text/csv').
+
+    Returns
+    -------
+        str: Presigned S3 URL for the uploaded file.
+    """
+    temp_csv = f"/tmp/{uuid.uuid4()}.csv"
+    df.to_csv(temp_csv, index=False)
+    s3 = boto3.client("s3")
+    s3.upload_file(temp_csv, bucket, key)
+    s3_url = generate_presigned_url(bucket, key)
+    resource = HttpResource(
+        uri="predictions://" + uuid.uuid4().hex,  # type: ignore[arg-type]
+        url=s3_url,
+        name=resource_name,
+        mime_type=mime_type,
+    )
+    mcp.add_resource(resource)
+    return resource
+
+
+def format_response_as_tool_result(data: bytes, content_type: str, charset: str) -> ToolResult:
+    """Format the deployment response into a ToolResult.
+
+    Using structured_content, to return as much information about
+    the response as possible, for LLMs to correctly interpret the
+    response.
+    """
+    charset = charset or "utf-8"
+    content_type = content_type.lower() if content_type else ""
+
+    if content_type.startswith("text/") or content_type == "application/json":
+        payload = {
+            "type": "text",
+            "mime_type": content_type,
+            "data": data.decode(charset),
+        }
+    elif content_type.startswith("image/"):
+        payload = {
+            "type": "image",
+            "mime_type": content_type,
+            "data_base64": base64.b64encode(data).decode(charset),
+        }
+    else:
+        payload = {
+            "type": "binary",
+            "mime_type": content_type,
+            "data_base64": base64.b64encode(data).decode(charset),
+        }
+
+    return ToolResult(structured_content=payload)

datarobot_genai/drmcp/server.py

@@ -0,0 +1,19 @@
+# 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 datarobot_genai.drmcp import create_mcp_server
+
+if __name__ == "__main__":
+    server = create_mcp_server()
+    server.run(show_banner=True)

datarobot_genai/drmcp/test_utils/__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/test_utils/integration_mcp_server.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+
+# 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.
+
+"""
+Integration test MCP server.
+
+This server works standalone (base tools only) or detects and loads
+user modules if they exist in the project structure.
+"""
+
+from pathlib import Path
+from typing import Any
+
+from datarobot_genai.drmcp import create_mcp_server
+
+# Import user components (will be used conditionally)
+try:
+    from app.core.server_lifecycle import ServerLifecycle  # type: ignore  # noqa: F401
+    from app.core.user_config import get_user_config  # type: ignore  # noqa: F401
+    from app.core.user_credentials import get_user_credentials  # type: ignore  # noqa: F401
+
+except ImportError:
+    # These imports will fail when running from library without user modules
+    pass
+
+
+def detect_user_modules() -> Any:
+    """
+    Detect if user modules exist in the project.
+
+    Returns
+    -------
+        Tuple of (config_factory, credentials_factory, lifecycle, module_paths) or None
+    """
+    # Try to find app directory
+    # When run from library: won't find it
+    # When run from project: will find it
+    current_dir = Path.cwd()
+
+    # Look for app in current directory or parent directories
+    for search_dir in [current_dir, current_dir.parent, current_dir.parent.parent]:
+        app_dir = search_dir / "app"
+        app_core_dir = app_dir / "core"
+        if app_core_dir.exists():
+            # Found user directory - load user modules
+            try:
+                module_paths = [
+                    (str(app_dir / "tools"), "app.tools"),
+                    (str(app_dir / "prompts"), "app.prompts"),
+                    (str(app_dir / "resources"), "app.resources"),
+                ]
+
+                return (
+                    get_user_config,
+                    get_user_credentials,
+                    ServerLifecycle(),
+                    module_paths,
+                )
+            except ImportError:
+                # User modules don't exist or can't be imported
+                pass
+
+    return None
+
+
+def main() -> None:
+    """Run the integration test MCP server."""
+    # Try to detect and load user modules
+    user_components = detect_user_modules()
+
+    if user_components:
+        # User modules found - create server with user extensions
+        config_factory, credentials_factory, lifecycle, module_paths = user_components
+        server = create_mcp_server(
+            config_factory=config_factory,
+            credentials_factory=credentials_factory,
+            lifecycle=lifecycle,
+            additional_module_paths=module_paths,
+            transport="stdio",
+        )
+    else:
+        # No user modules - create server with base tools only
+        server = create_mcp_server(transport="stdio")
+
+    server.run()
+
+
+if __name__ == "__main__":
+    main()

datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@@ -0,0 +1,96 @@
+# 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 os
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+
+from .utils import load_env
+
+load_env()
+
+
+def get_dr_mcp_server_url() -> str | None:
+    """Get DataRobot MCP server URL."""
+    return os.environ.get("DR_MCP_SERVER_URL")
+
+
+def get_openai_llm_client_config() -> dict[str, str]:
+    """Get OpenAI LLM client configuration."""
+    openai_api_key = os.environ.get("OPENAI_API_KEY")
+    openai_api_base = os.environ.get("OPENAI_API_BASE")
+    openai_api_deployment_id = os.environ.get("OPENAI_API_DEPLOYMENT_ID")
+    openai_api_version = os.environ.get("OPENAI_API_VERSION")
+    save_llm_responses = os.environ.get("SAVE_LLM_RESPONSES", "false").lower() == "true"
+
+    # Check for OpenAI configuration
+    if not openai_api_key:
+        raise ValueError("Missing required environment variable: OPENAI_API_KEY")
+    if (
+        openai_api_base and not openai_api_deployment_id
+    ):  # For Azure OpenAI, we need additional variables
+        raise ValueError("Missing required environment variable: OPENAI_API_DEPLOYMENT_ID")
+
+    config: dict[str, str] = {
+        "openai_api_key": openai_api_key,
+    }
+
+    if openai_api_base:
+        config["openai_api_base"] = openai_api_base
+    if openai_api_deployment_id:
+        config["openai_api_deployment_id"] = openai_api_deployment_id
+    if openai_api_version:
+        config["openai_api_version"] = openai_api_version
+    config["save_llm_responses"] = str(save_llm_responses)
+
+    return config
+
+
+def get_headers() -> dict[str, str]:
+    # When the MCP server is deployed in DataRobot, we have to include the API token in headers for
+    # authentication.
+    api_token = os.getenv("DATAROBOT_API_TOKEN")
+    headers = {"Authorization": f"Bearer {api_token}"}
+    return headers
+
+
+@asynccontextmanager
+async def ete_test_mcp_session(
+    additional_headers: dict[str, str] | None = None,
+) -> AsyncGenerator[ClientSession, None]:
+    """Create an MCP session for each test.
+
+    Parameters
+    ----------
+    additional_headers : dict[str, str], optional
+        Additional headers to include in the MCP session (e.g., auth headers for testing).
+    """
+    try:
+        headers = get_headers()
+        if additional_headers:
+            headers.update(additional_headers)
+
+        async with streamablehttp_client(url=get_dr_mcp_server_url(), headers=headers) as (
+            read_stream,
+            write_stream,
+            _,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                await asyncio.wait_for(session.initialize(), timeout=5)
+                yield session
+    except asyncio.TimeoutError:
+        raise TimeoutError(f"Check if the MCP server is running at {get_dr_mcp_server_url()}")

datarobot_genai/drmcp/test_utils/mcp_utils_integration.py

@@ -0,0 +1,94 @@
+# 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 contextlib
+import os
+from collections.abc import AsyncGenerator
+from pathlib import Path
+
+from mcp import ClientSession
+from mcp.client.stdio import StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from .utils import load_env
+
+load_env()
+
+
+def integration_test_mcp_server_params() -> StdioServerParameters:
+    env = {
+        "DATAROBOT_API_TOKEN": os.environ.get("DATAROBOT_API_TOKEN") or "test-token",
+        "DATAROBOT_ENDPOINT": os.environ.get("DATAROBOT_ENDPOINT")
+        or "https://test.datarobot.com/api/v2",
+        "MCP_SERVER_LOG_LEVEL": os.environ.get("MCP_SERVER_LOG_LEVEL") or "WARNING",
+        "APP_LOG_LEVEL": os.environ.get("APP_LOG_LEVEL") or "WARNING",
+        "OTEL_ENABLED": os.environ.get("OTEL_ENABLED") or "false",
+        "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP": os.environ.get(
+            "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP"
+        )
+        or "false",
+        "MCP_SERVER_REGISTER_DYNAMIC_PROMPTS_ON_STARTUP": os.environ.get(
+            "MCP_SERVER_REGISTER_DYNAMIC_PROMPTS_ON_STARTUP"
+        )
+        or "true",
+    }
+
+    script_dir = Path(__file__).resolve().parent
+    server_script = str(script_dir / "integration_mcp_server.py")
+    # Add src/ directory to Python path so datarobot_genai can be imported
+    src_dir = script_dir.parent.parent.parent
+
+    return StdioServerParameters(
+        command="uv",
+        args=["run", server_script],
+        env={
+            "PYTHONPATH": str(src_dir),
+            "MCP_SERVER_NAME": "integration",
+            "MCP_SERVER_PORT": "8081",
+            **env,
+        },
+    )
+
+
+@contextlib.asynccontextmanager
+async def integration_test_mcp_session(
+    server_params: StdioServerParameters | None = None, timeout: int = 30
+) -> AsyncGenerator[ClientSession, None]:
+    """
+    Create and connect a client for the MCP server as a context manager.
+
+    Args:
+        server_params: Parameters for configuring the server connection
+        timeout: Timeout
+
+    Yields
+    ------
+        ClientSession: Connected MCP client session
+
+    Raises
+    ------
+        ConnectionError: If session initialization fails
+        TimeoutError: If session initialization exceeds timeout
+    """
+    server_params = server_params or integration_test_mcp_server_params()
+
+    try:
+        async with stdio_client(server_params) as (read_stream, write_stream):
+            async with ClientSession(read_stream, write_stream) as session:
+                await asyncio.wait_for(session.initialize(), timeout=timeout)
+                yield session
+
+    except asyncio.TimeoutError:
+        raise TimeoutError(f"Session initialization timed out after {timeout} seconds")

datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py

@@ -0,0 +1,234 @@
+# 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 json
+from typing import Any
+
+import openai
+from mcp import ClientSession
+from mcp.types import CallToolResult
+from mcp.types import ListToolsResult
+from mcp.types import TextContent
+from openai.types.chat.chat_completion import ChatCompletion
+
+from .utils import save_response_to_file
+
+
+class ToolCall:
+    """Represents a tool call with its parameters and reasoning."""
+
+    def __init__(self, tool_name: str, parameters: dict[str, Any], reasoning: str):
+        self.tool_name = tool_name
+        self.parameters = parameters
+        self.reasoning = reasoning
+
+
+class LLMResponse:
+    """Represents an LLM response with content and tool calls."""
+
+    def __init__(self, content: str, tool_calls: list[ToolCall], tool_results: list[str]):
+        self.content = content
+        self.tool_calls = tool_calls
+        self.tool_results = tool_results
+
+
+class LLMMCPClient:
+    """Client for interacting with LLMs via MCP."""
+
+    def __init__(self, config: str):
+        """Initialize the LLM MCP client."""
+        # Parse config string to extract parameters
+        config_dict = eval(config) if isinstance(config, str) else config
+
+        openai_api_key = config_dict.get("openai_api_key")
+        openai_api_base = config_dict.get("openai_api_base")
+        openai_api_deployment_id = config_dict.get("openai_api_deployment_id")
+        model = config_dict.get("model", "gpt-3.5-turbo")
+        save_llm_responses = config_dict.get("save_llm_responses", True)
+
+        if openai_api_base and openai_api_deployment_id:
+            # Azure OpenAI
+            self.openai_client = openai.AzureOpenAI(
+                api_key=openai_api_key,
+                azure_endpoint=openai_api_base,
+                api_version=config_dict.get("openai_api_version", "2024-02-15-preview"),
+            )
+            self.model = openai_api_deployment_id
+        else:
+            # Regular OpenAI
+            self.openai_client = openai.OpenAI(api_key=openai_api_key)  # type: ignore[assignment]
+            self.model = model
+
+        self.save_llm_responses = save_llm_responses
+        self.available_tools: list[dict[str, Any]] = []
+        self.available_prompts: list[dict[str, Any]] = []
+        self.available_resources: list[dict[str, Any]] = []
+
+    async def _add_mcp_tool_to_available_tools(self, mcp_session: ClientSession) -> None:
+        """Add a tool to the available tools."""
+        tools_result: ListToolsResult = await mcp_session.list_tools()
+        self.available_tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": tool.inputSchema,
+                },
+            }
+            for tool in tools_result.tools
+        ]
+
+    async def _call_mcp_tool(
+        self, tool_name: str, parameters: dict[str, Any], mcp_session: ClientSession
+    ) -> str:
+        """Call an MCP tool and return the result as a string."""
+        result: CallToolResult = await mcp_session.call_tool(tool_name, parameters)
+        return (
+            result.content[0].text
+            if result.content and isinstance(result.content[0], TextContent)
+            else str(result.content)
+        )
+
+    async def _process_tool_calls(
+        self,
+        response: ChatCompletion,
+        messages: list[Any],
+        mcp_session: ClientSession,
+    ) -> tuple[list[ToolCall], list[str]]:
+        """Process tool calls from the response, and return the tool calls and tool results."""
+        tool_calls = []
+        tool_results = []
+
+        # If the response has tool calls, process them
+        if response.choices[0].message.tool_calls:
+            messages.append(response.choices[0].message)  # Add assistant's message with tool calls
+
+            for tool_call in response.choices[0].message.tool_calls:
+                tool_name = tool_call.function.name  # type: ignore[union-attr]
+                parameters = json.loads(tool_call.function.arguments)  # type: ignore[union-attr]
+
+                tool_calls.append(
+                    ToolCall(
+                        tool_name=tool_name,
+                        parameters=parameters,
+                        reasoning="Tool selected by LLM",
+                    )
+                )
+
+                try:
+                    result_text = await self._call_mcp_tool(tool_name, parameters, mcp_session)
+                    tool_results.append(result_text)
+
+                    # Add tool result to messages
+                    messages.append(
+                        {
+                            "role": "tool",
+                            "content": result_text,
+                            "tool_call_id": tool_call.id,
+                            "name": tool_name,
+                        }
+                    )
+                except Exception as e:
+                    error_msg = f"Error calling {tool_name}: {str(e)}"
+                    tool_results.append(error_msg)
+                    messages.append(
+                        {
+                            "role": "tool",
+                            "content": error_msg,
+                            "tool_call_id": tool_call.id,
+                            "name": tool_name,
+                        }
+                    )
+
+        return tool_calls, tool_results
+
+    async def _get_llm_response(
+        self, messages: list[dict[str, Any]], allow_tool_calls: bool = True
+    ) -> Any:
+        """Get a response from the LLM with optional tool calling capability."""
+        kwargs = {
+            "model": self.model,
+            "messages": messages,
+        }
+
+        if allow_tool_calls and self.available_tools:
+            kwargs["tools"] = self.available_tools
+            kwargs["tool_choice"] = "auto"
+
+        return self.openai_client.chat.completions.create(**kwargs)
+
+    async def process_prompt_with_mcp_support(
+        self, prompt: str, mcp_session: ClientSession, output_file_name: str = ""
+    ) -> LLMResponse:
+        """Process a prompt with MCP tool support."""
+        # Add MCP tools to available tools
+        await self._add_mcp_tool_to_available_tools(mcp_session)
+
+        if output_file_name:
+            print(f"Processing prompt for test: {output_file_name}")
+
+        # Initialize conversation
+        messages = [
+            {
+                "role": "system",
+                "content": (
+                    "You are a helpful AI assistant that can use tools to help users. "
+                    "If you need more information to provide a complete response, you can make "
+                    "multiple tool calls. When dealing with file paths, use them as raw paths "
+                    "without converting to file:// URLs."
+                ),
+            },
+            {"role": "user", "content": prompt},
+        ]
+
+        all_tool_calls = []
+        all_tool_results = []
+
+        while True:
+            # Get LLM response
+            response = await self._get_llm_response(messages)
+
+            # If no tool calls in response, this is the final response
+            if not response.choices[0].message.tool_calls:
+                final_response = response.choices[0].message.content
+                break
+
+            # Process tool calls
+            tool_calls, tool_results = await self._process_tool_calls(
+                response, messages, mcp_session
+            )
+            all_tool_calls.extend(tool_calls)
+            all_tool_results.extend(tool_results)
+
+            # Get another LLM response to see if we need more tool calls
+            response = await self._get_llm_response(messages, allow_tool_calls=True)
+
+            # If no more tool calls needed, this is the final response
+            if not response.choices[0].message.tool_calls:
+                final_response = response.choices[0].message.content
+                break
+
+        clean_content = final_response.replace("*", "").lower()
+
+        llm_response = LLMResponse(
+            content=clean_content,
+            tool_calls=all_tool_calls,
+            tool_results=all_tool_results,
+        )
+
+        if self.save_llm_responses:
+            save_response_to_file(llm_response, name=output_file_name)
+
+        return llm_response