datarobot-genai 0.2.30__tar.gz → 0.2.32__tar.gz

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/PKG-INFO

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

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.30"
+version = "0.2.32"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/__init__.py

@@ -19,11 +19,11 @@ A reusable library for building Model Context Protocol (MCP) servers with DataRo
 """
 
 # Export main server components
+from datarobot_genai.drmcp.test_utils.clients.openai import OpenAILLMMCPClient
 from datarobot_genai.drmcp.test_utils.mcp_utils_ete import ete_test_mcp_session
 from datarobot_genai.drmcp.test_utils.mcp_utils_ete import get_dr_mcp_server_url
 from datarobot_genai.drmcp.test_utils.mcp_utils_ete import get_headers
 from datarobot_genai.drmcp.test_utils.mcp_utils_integration import integration_test_mcp_session
-from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMMCPClient
 from datarobot_genai.drmcp.test_utils.tool_base_ete import ETETestExpectations
 from datarobot_genai.drmcp.test_utils.tool_base_ete import ToolBaseE2E
 from datarobot_genai.drmcp.test_utils.tool_base_ete import ToolCallTestExpectations
@@ -70,7 +70,7 @@ __all__ = [
     "get_dr_mcp_server_url",
     "get_headers",
     "ete_test_mcp_session",
-    "LLMMCPClient",
+    "OpenAILLMMCPClient",
     "ETETestExpectations",
     "ToolBaseE2E",
     "ToolCallTestExpectations",

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/exceptions.py

@@ -19,7 +19,3 @@ class DynamicToolRegistrationError(Exception):
 
 class DynamicPromptRegistrationError(Exception):
     """Exception raised for errors in the dynamic prompt registration process."""
-
-
-class MCPError(Exception):
-    """Base class for MCP errors."""

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/logging.py

@@ -20,7 +20,7 @@ from collections.abc import Callable
 from typing import Any
 from typing import TypeVar
 
-from .exceptions import MCPError
+from fastmcp.exceptions import ToolError
 
 # Secret patterns to redact from logs
 SECRET_PATTERNS = [
@@ -93,6 +93,6 @@ def log_execution(func: F) -> F:
             return result
         except Exception as e:
             error_msg = _log_error(logger, func.__name__, e, args=args, kwargs=kwargs)
-            raise MCPError(error_msg)
+            raise ToolError(error_msg)
 
     return wrapper  # type: ignore[return-value]

datarobot_genai-0.2.32/src/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-0.2.30/src/datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py → datarobot_genai-0.2.32/src/datarobot_genai/drmcp/test_utils/clients/base.py

@@ -1,4 +1,4 @@
-# Copyright 2025 DataRobot, Inc.
+# 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.
@@ -12,7 +12,11 @@
 # 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
 
@@ -23,7 +27,7 @@ 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
+from datarobot_genai.drmcp.test_utils.utils import save_response_to_file
 
 
 class ToolCall:
@@ -44,9 +48,9 @@ class LLMResponse:
         self.tool_results = tool_results
 
 
-class LLMMCPClient:
+class BaseLLMMCPClient(ABC):
     """
-    Client for interacting with LLMs via MCP.
+    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.
@@ -54,54 +58,48 @@ class LLMMCPClient:
 
     def __init__(
         self,
-        config: str,
+        config: str | dict,
     ):
         """
         Initialize the LLM MCP client.
 
         Args:
-            config: Configuration string or dict with:
-                - openai_api_key: OpenAI API key
-                - openai_api_base: Optional Azure OpenAI endpoint
-                - openai_api_deployment_id: Optional Azure deployment ID
-                - openai_api_version: Optional Azure API version
-                - model: Model name (default: "gpt-3.5-turbo")
-                - save_llm_responses: Whether to save responses (default: True)
+            config: Configuration string or dict with provider-specific keys.
         """
-        # Parse config string to extract parameters
+        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:
-                config_dict = json.loads(config)
+                return json.loads(config)
             except json.JSONDecodeError:
                 # Fall back to literal_eval for Python dict literal strings
-                config_dict = literal_eval(config)
-        else:
-            config_dict = 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
+                return literal_eval(config)
+        return config
 
-        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]] = []
+    @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."""

datarobot_genai-0.2.32/src/datarobot_genai/drmcp/test_utils/clients/dr_gateway.py

@@ -0,0 +1,58 @@
+# 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.
+
+"""DataRobot LLM Gateway MCP Client implementation."""
+
+import openai
+
+from .base import BaseLLMMCPClient
+
+
+class DRLLMGatewayMCPClient(BaseLLMMCPClient):
+    """
+    Client for interacting with LLMs via MCP using DataRobot LLM Gateway.
+
+    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:
+                - datarobot_api_token: DataRobot API token
+                - datarobot_endpoint: DataRobot endpoint URL (default: "https://app.datarobot.com/api/v2")
+                - model: Model name (default: "gpt-4o-mini")
+                - 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 DataRobot LLM Gateway."""
+        datarobot_api_token = config_dict.get("datarobot_api_token")
+        datarobot_endpoint = config_dict.get(
+            "datarobot_endpoint", "https://app.datarobot.com/api/v2"
+        )
+        model = config_dict.get("model", "gpt-4o-mini")
+
+        # Build gateway URL: {endpoint}/genai/llmgw
+        gateway_url = datarobot_endpoint.rstrip("/") + "/genai/llmgw"
+
+        client = openai.OpenAI(api_key=datarobot_api_token, base_url=gateway_url)
+        return client, model

datarobot_genai-0.2.32/src/datarobot_genai/drmcp/test_utils/clients/openai.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.
+
+"""OpenAI LLM MCP Client implementation."""
+
+import openai
+
+from .base import BaseLLMMCPClient
+
+
+class OpenAILLMMCPClient(BaseLLMMCPClient):
+    """
+    Client for interacting with LLMs via MCP using OpenAI or Azure OpenAI.
+
+    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:
+                - openai_api_key: OpenAI API key
+                - openai_api_base: Optional Azure OpenAI endpoint
+                - openai_api_deployment_id: Optional Azure deployment ID
+                - openai_api_version: Optional Azure API version
+                - model: Model name (default: "gpt-3.5-turbo")
+                - save_llm_responses: Whether to save responses (default: True)
+        """
+        super().__init__(config)
+
+    def _create_llm_client(
+        self, config_dict: dict
+    ) -> tuple[openai.OpenAI | openai.AzureOpenAI, str]:
+        """Create the LLM client for OpenAI or Azure OpenAI."""
+        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")
+
+        if openai_api_base and openai_api_deployment_id:
+            # Azure 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"),
+            )
+            return client, openai_api_deployment_id
+        else:
+            # Regular OpenAI
+            client = openai.OpenAI(api_key=openai_api_key)  # type: ignore[assignment]
+            return client, model

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/test_interactive.py

@@ -40,40 +40,40 @@ from mcp.types import ElicitResult
 
 from datarobot_genai.drmcp import get_dr_mcp_server_url
 from datarobot_genai.drmcp import get_headers
-from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMMCPClient
-from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMResponse
-from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import ToolCall
+from datarobot_genai.drmcp.test_utils.clients.base import LLMResponse
+from datarobot_genai.drmcp.test_utils.clients.base import ToolCall
+from datarobot_genai.drmcp.test_utils.clients.dr_gateway import DRLLMGatewayMCPClient
 
 # Re-export for backwards compatibility
-__all__ = ["LLMMCPClient", "LLMResponse", "ToolCall", "test_mcp_interactive"]
+__all__ = ["DRLLMGatewayMCPClient", "LLMResponse", "ToolCall", "test_mcp_interactive"]
 
 
 async def test_mcp_interactive() -> None:
     """Test the MCP server interactively with LLM agent."""
     # Check for required environment variables
-    openai_api_key = os.environ.get("OPENAI_API_KEY")
-    if not openai_api_key:
-        print("❌ Error: OPENAI_API_KEY environment variable is required")
+    datarobot_api_token = os.environ.get("DATAROBOT_API_TOKEN")
+    if not datarobot_api_token:
+        print("❌ Error: DATAROBOT_API_TOKEN environment variable is required")
         print("Please set it in your .env file or export it")
         return
 
-    # Optional Azure OpenAI settings
-    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")
+    # Optional DataRobot settings
+    datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT")
+    model = os.environ.get("MODEL")
 
     print("🤖 Initializing LLM MCP Client...")
 
     # Initialize the LLM client with elicitation handler
     config = {
-        "openai_api_key": openai_api_key,
-        "openai_api_base": openai_api_base,
-        "openai_api_deployment_id": openai_api_deployment_id,
-        "openai_api_version": openai_api_version,
+        "datarobot_api_token": datarobot_api_token,
         "save_llm_responses": False,
     }
+    if datarobot_endpoint:
+        config["datarobot_endpoint"] = datarobot_endpoint
+    if model:
+        config["model"] = model
 
-    llm_client = LLMMCPClient(str(config))
+    llm_client = DRLLMGatewayMCPClient(str(config))
 
     # Get MCP server URL
     mcp_server_url = get_dr_mcp_server_url()

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -17,7 +17,7 @@ from typing import Any
 
 from pydantic import BaseModel
 
-from .openai_llm_mcp_client import LLMResponse
+from .clients.base import LLMResponse
 
 
 class ToolCallTestExpectations(BaseModel):

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/utils.py

@@ -19,7 +19,7 @@ from typing import TYPE_CHECKING
 from typing import Any
 
 if TYPE_CHECKING:
-    from .openai_llm_mcp_client import LLMResponse
+    from .clients.base import LLMResponse
 
 from dotenv import load_dotenv
 

{datarobot_genai-0.2.30 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -20,6 +20,7 @@ import logging
 import uuid
 from typing import Annotated
 from typing import Any
+from typing import Literal
 
 import httpx
 from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
@@ -821,6 +822,96 @@ class GoogleDriveClient:
             headers={"Content-Type": f"multipart/related; boundary={boundary}"},
         )
 
+    async def manage_access(
+        self,
+        *,
+        file_id: str,
+        action: Literal["add", "update", "remove"],
+        role: Literal["reader", "commenter", "writer", "fileOrganizer", "organizer", "owner"]
+        | None = None,
+        email_address: str | None = None,
+        permission_id: str | None = None,
+        transfer_ownership: bool = False,
+    ) -> str:
+        """Manage access permissions for a Google Drive file or folder.
+
+        Adds, updates, or removes sharing permissions on an existing Google Drive
+        file or folder using the Google Drive Permissions API.
+
+        This method supports granting access to users or groups, changing access
+        roles, and revoking permissions. Ownership transfer is supported for files
+        in "My Drive" when explicitly requested.
+
+        Args:
+            file_id: The ID of the Google Drive file or folder whose permissions
+                are being managed.
+            action: The permission operation to perform.
+            role: The access role to assign or update. Valid values include
+                Required for "add" and "update" actions.
+            email_address: The email address of the user or group to grant access to.
+                Required for the "add" action.
+            permission_id: The ID of the permission to update or remove.
+                Required for "update" and "remove" actions.
+            transfer_ownership: Whether to transfer ownership of the file.
+                Only applicable when action="update" and role="owner".
+
+        Returns
+        -------
+            Permission id.
+            For "add" its newly added permission.
+            For "update"/"remove" its previous permission.
+
+        Raises
+        ------
+            GoogleDriveError: If the permission operation fails (invalid arguments,
+                insufficient permissions, resource not found, ownership transfer
+                not allowed, rate limited, etc.).
+        """
+        if not file_id.strip():
+            raise GoogleDriveError("Argument validation error: 'file_id' cannot be empty.")
+
+        if action == "add" and not email_address:
+            raise GoogleDriveError("'email_address' is required for action 'add'.")
+
+        if action in ("update", "remove") and not permission_id:
+            raise GoogleDriveError("'permission_id' is required for action 'update' or 'remove'.")
+
+        if action != "remove" and not role:
+            raise GoogleDriveError("'role' is required for action 'add' or 'update'.")
+
+        if action == "add":
+            response = await self._client.post(
+                url=f"/{file_id}/permissions",
+                json={
+                    "type": "user",
+                    "role": role,
+                    "emailAddress": email_address,
+                },
+                params={"sendNotificationEmail": False, "supportsAllDrives": True},
+            )
+
+        elif action == "update":
+            response = await self._client.patch(
+                url=f"/{file_id}/permissions/{permission_id}",
+                json={"role": role},
+                params={"transferOwnership": transfer_ownership, "supportsAllDrives": True},
+            )
+
+        elif action == "remove":
+            response = await self._client.delete(url=f"/{file_id}/permissions/{permission_id}")
+
+        else:
+            raise GoogleDriveError(f"Invalid action '{action}'")
+
+        if response.status_code not in (200, 201, 204):
+            raise GoogleDriveError(f"Drive API error {response.status_code}: {response.text}")
+
+        if action == "add":
+            return response.json()["id"]
+
+        # Cannot be null here because of above validators
+        return permission_id  # type: ignore
+
     async def __aenter__(self) -> "GoogleDriveClient":
         """Async context manager entry."""
         return self