datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/drmcp/core/tool_filter.py

@@ -0,0 +1,117 @@
+# 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 fastmcp.tools import Tool
+from mcp.types import Tool as MCPTool
+
+
+def filter_tools_by_tags(
+    tools: list[Tool | MCPTool],
+    tags: list[str] | None = None,
+    match_all: bool = False,
+) -> list[Tool | MCPTool]:
+    """
+    Filter tools by tags.
+
+    Args:
+        tools: List of tools to filter
+        tags: List of tags to filter by. If None, returns all tools
+        match_all: If True, tool must have all specified tags. If False, tool must have at least
+            one tag.
+
+    Returns
+    -------
+        List of tools that match the tag criteria
+    """
+    if not tags:
+        return tools
+
+    filtered_tools = []
+
+    for tool in tools:
+        tool_tags = get_tool_tags(tool)
+
+        if not tool_tags:
+            continue
+
+        if match_all:
+            # Tool must have all specified tags
+            if all(tag in tool_tags for tag in tags):
+                filtered_tools.append(tool)
+        elif any(tag in tool_tags for tag in tags):
+            # Tool must have at least one specified tag
+            filtered_tools.append(tool)
+
+    return filtered_tools
+
+
+def get_tool_tags(tool: Tool | MCPTool) -> list[str]:
+    """
+    Get tags for a specific tool.
+
+    Args:
+        tool: The tool to get tags for
+
+    Returns
+    -------
+        List of tags for the tool
+    """
+    # Primary: native FastMCP meta location
+    if hasattr(tool, "meta") and getattr(tool, "meta"):
+        fastmcp_meta = tool.meta.get("_fastmcp", {})
+        meta_tags = fastmcp_meta.get("tags", [])
+        if isinstance(meta_tags, list):
+            return meta_tags
+
+    # Fallback: annotations.tags (for compatibility during transition)
+    if tool.annotations and hasattr(tool.annotations, "tags"):
+        tags = getattr(tool.annotations, "tags", [])
+        return tags if isinstance(tags, list) else []
+
+    return []
+
+
+def list_all_tags(tools: list[Tool | MCPTool]) -> list[str]:
+    """
+    Get all unique tags from a list of tools.
+
+    Args:
+        tools: List of tools to extract tags from
+
+    Returns
+    -------
+        List of unique tags
+    """
+    all_tags = set()
+    for tool in tools:
+        tool_tags = get_tool_tags(tool)
+        all_tags.update(tool_tags)
+
+    return sorted(list(all_tags))
+
+
+def get_tools_by_tag(tools: list[Tool | MCPTool], tag: str) -> list[Tool | MCPTool]:
+    """
+    Get all tools that have a specific tag.
+
+    Args:
+        tools: List of tools to search
+        tag: The tag to search for
+
+    Returns
+    -------
+        List of tools with the specified tag
+    """
+    return filter_tools_by_tags(tools, [tag])

datarobot_genai/drmcp/core/utils.py

@@ -0,0 +1,138 @@
+# 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
+from urllib.parse import urlparse
+
+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)
+
+
+def is_valid_url(url: str) -> bool:
+    """Check if a URL is valid."""
+    result = urlparse(url)
+    return all([result.scheme, result.netloc])

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/clients/anthropic.py

@@ -0,0 +1,68 @@
+# Copyright 2026 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.
+
+"""Anthropic LLM MCP Client implementation (example).
+
+This is an example implementation showing how easy it is to add a new LLM provider.
+Anthropic's API is OpenAI-compatible, so we can use the OpenAI SDK with their endpoint.
+"""
+
+import openai
+
+from .base import BaseLLMMCPClient
+
+
+class AnthropicMCPClient(BaseLLMMCPClient):
+    """
+    Client for interacting with LLMs via MCP using Anthropic Claude.
+
+    Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
+    Tools using FastMCP's built-in elicitation will work automatically.
+
+    Example:
+        ```python
+        config = {
+            "anthropic_api_key": "sk-ant-...",
+            "model": "claude-3-5-sonnet-20241022",
+        }
+        client = AnthropicMCPClient(str(config))
+        ```
+    """
+
+    def __init__(
+        self,
+        config: str | dict,
+    ):
+        """
+        Initialize the LLM MCP client.
+
+        Args:
+            config: Configuration string or dict with:
+                - anthropic_api_key: Anthropic API key
+                - model: Model name (default: "claude-3-5-sonnet-20241022")
+                - save_llm_responses: Whether to save responses (default: True)
+        """
+        super().__init__(config)
+
+    def _create_llm_client(self, config_dict: dict) -> tuple[openai.OpenAI, str]:
+        """Create the LLM client for Anthropic (OpenAI-compatible endpoint)."""
+        anthropic_api_key = config_dict.get("anthropic_api_key")
+        model = config_dict.get("model", "claude-3-5-sonnet-20241022")
+
+        # Anthropic provides an OpenAI-compatible endpoint
+        client = openai.OpenAI(
+            api_key=anthropic_api_key,
+            base_url="https://api.anthropic.com/v1",
+        )
+        return client, model

datarobot_genai/drmcp/test_utils/clients/base.py

@@ -0,0 +1,300 @@
+# Copyright 2026 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.
+
+"""Base classes for LLM MCP clients."""
+
+import json
+from abc import ABC
+from abc import abstractmethod
+from ast import literal_eval
+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 BaseLLMMCPClient(ABC):
+    """
+    Base class for LLM MCP clients.
+
+    Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
+    Tools using FastMCP's built-in elicitation will work automatically.
+    """
+
+    def __init__(
+        self,
+        config: str | dict,
+    ):
+        """
+        Initialize the LLM MCP client.
+
+        Args:
+            config: Configuration string or dict with provider-specific keys.
+        """
+        config_dict = self._parse_config(config)
+        self.openai_client, self.model = self._create_llm_client(config_dict)
+        self.save_llm_responses = config_dict.get("save_llm_responses", True)
+        self.available_tools: list[dict[str, Any]] = []
+        self.available_prompts: list[dict[str, Any]] = []
+        self.available_resources: list[dict[str, Any]] = []
+
+    @staticmethod
+    def _parse_config(config: str | dict) -> dict:
+        """Parse config string to dict."""
+        if isinstance(config, str):
+            # Try JSON first (safer), fall back to literal_eval for Python dict strings
+            try:
+                return json.loads(config)
+            except json.JSONDecodeError:
+                # Fall back to literal_eval for Python dict literal strings
+                return literal_eval(config)
+        return config
+
+    @abstractmethod
+    def _create_llm_client(
+        self, config_dict: dict
+    ) -> tuple[openai.OpenAI | openai.AzureOpenAI, str]:
+        """
+        Create the LLM client.
+
+        Args:
+            config_dict: Parsed configuration dictionary
+
+        Returns
+        -------
+            Tuple of (LLM client instance, model name)
+        """
+        pass
+
+    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.
+
+        Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
+        Tools using FastMCP's built-in elicitation will work automatically.
+
+        Args:
+            tool_name: Name of the tool to call
+            parameters: Parameters to pass to the tool
+            mcp_session: MCP client session
+
+        Returns
+        -------
+            Result text from the tool call
+        """
+        result: CallToolResult = await mcp_session.call_tool(tool_name, parameters)
+        content = (
+            result.content[0].text
+            if result.content and isinstance(result.content[0], TextContent)
+            else str(result.content)
+        )
+        if result.structuredContent is not None:
+            structured_content = json.dumps(result.structuredContent)
+        else:
+            structured_content = ""
+        return f"Content: {content}\nStructured content: {structured_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 and elicitation handling.
+
+        This method:
+        1. Adds MCP tools to available tools
+        2. Sends prompt to LLM
+        3. Processes tool calls
+        4. Continues until LLM provides final response
+
+        Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
+
+        Args:
+            prompt: User prompt
+            mcp_session: MCP client session
+            output_file_name: Optional file name to save response
+
+        Returns
+        -------
+            LLMResponse with content, tool calls, and tool results
+        """
+        # 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 or ask the user for more info, but prefer tool calls "
+                    "when possible. "
+                    "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