datarobot-genai 0.2.29__py3-none-any.whl → 0.2.37__py3-none-any.whl

datarobot_genai/core/cli/agent_kernel.py

@@ -107,7 +107,10 @@ class AgentKernel:
         stream: bool = False,
         config: Any | None = None,
     ) -> ChatCompletion | Stream[ChatCompletionChunk]:
-        chat_api_url = config.agent_endpoint if config else self.base_url
+        if config is not None:
+            chat_api_url = f"http://localhost:{config.local_dev_port}"
+        else:
+            chat_api_url = self.base_url
         print(chat_api_url)
 
         return self._do_chat_completion(chat_api_url, user_prompt, completion_json, stream=stream)

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/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/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/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/{openai_llm_mcp_client.py → 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/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/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/drmcp/test_utils/mcp_utils_ete.py

@@ -68,6 +68,26 @@ def get_openai_llm_client_config() -> dict[str, str]:
     return config
 
 
+def get_dr_llm_gateway_client_config() -> dict[str, str]:
+    """Get DataRobot LLM Gateway client configuration."""
+    datarobot_api_token = os.environ.get("DATAROBOT_API_TOKEN")
+    datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT")
+    save_llm_responses = os.environ.get("SAVE_LLM_RESPONSES", "false").lower() == "true"
+
+    if not datarobot_api_token:
+        raise ValueError("Missing required environment variable: DATAROBOT_API_TOKEN")
+
+    config: dict[str, str] = {
+        "datarobot_api_token": datarobot_api_token,
+        "save_llm_responses": str(save_llm_responses),
+    }
+
+    if datarobot_endpoint:
+        config["datarobot_endpoint"] = datarobot_endpoint
+
+    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.

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/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/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