datarobot-genai 0.2.13__py3-none-any.whl → 0.2.20__py3-none-any.whl

datarobot_genai/drmcp/core/config.py

@@ -245,6 +245,30 @@ class MCPServerConfig(BaseSettings):
             os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
         )
 
+    # Gdrive tools
+    enable_gdrive_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
+            "ENABLE_GDRIVE_TOOLS",
+        ),
+        description="Enable/disable GDrive tools",
+    )
+    is_gdrive_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether GDrive OAuth provider is configured for GDrive integration",
+    )
+
+    @property
+    def is_gdrive_oauth_configured(self) -> bool:
+        return self.is_gdrive_oauth_provider_configured or bool(
+            os.getenv("GDRIVE_CLIENT_ID") and os.getenv("GDRIVE_CLIENT_SECRET")
+        )
+
     @field_validator(
         "otel_attributes",
         mode="before",

datarobot_genai/drmcp/core/tool_config.py

@@ -29,6 +29,7 @@ class ToolType(str, Enum):
     PREDICTIVE = "predictive"
     JIRA = "jira"
     CONFLUENCE = "confluence"
+    GDRIVE = "gdrive"
 
 
 class ToolConfig(TypedDict):
@@ -64,6 +65,13 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
         package_prefix="datarobot_genai.drmcp.tools.confluence",
         config_field_name="enable_confluence_tools",
     ),
+    ToolType.GDRIVE: ToolConfig(
+        name="gdrive",
+        oauth_check=lambda config: config.is_gdrive_oauth_configured,
+        directory="gdrive",
+        package_prefix="datarobot_genai.drmcp.tools.gdrive",
+        config_field_name="enable_gdrive_tools",
+    ),
 }
 
 

datarobot_genai/drmcp/core/utils.py

@@ -14,6 +14,7 @@
 import base64
 import uuid
 from typing import Any
+from urllib.parse import urlparse
 
 import boto3
 from fastmcp.resources import HttpResource
@@ -129,3 +130,9 @@ def format_response_as_tool_result(data: bytes, content_type: str, charset: str)
         }
 
     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/test_utils/elicitation_test_tool.py

@@ -0,0 +1,89 @@
+# 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.
+
+"""Test tool for elicitation testing.
+
+This module registers a test tool that can be used to test elicitation support.
+It should be imported in tests that need it.
+"""
+
+from fastmcp import Context
+from fastmcp.server.context import AcceptedElicitation
+from fastmcp.server.context import CancelledElicitation
+from fastmcp.server.context import DeclinedElicitation
+
+from datarobot_genai.drmcp.core.mcp_instance import mcp
+
+
+@mcp.tool(
+    name="get_user_greeting",
+    description=(
+        "Get a personalized greeting for a user. "
+        "Requires a username - if not provided, will request it via elicitation."
+    ),
+    tags={"test", "elicitation"},
+)
+async def get_user_greeting(ctx: Context, username: str | None = None) -> dict:
+    """
+    Get a personalized greeting for a user.
+
+    This tool demonstrates FastMCP's built-in elicitation by requiring a username parameter.
+    If username is not provided, it uses ctx.elicit() to request it from the user.
+
+    Args:
+        ctx: FastMCP context (automatically injected)
+        username: The username to greet. If None, elicitation will be triggered.
+
+    Returns
+    -------
+        Dictionary with greeting message or error if elicitation was declined/cancelled
+    """
+    if not username:
+        # Use elicitation to request username from the client
+        try:
+            result = await ctx.elicit(
+                message="Username is required to generate a personalized greeting",
+                response_type=str,
+            )
+
+            if isinstance(result, AcceptedElicitation):
+                username = result.data
+            elif isinstance(result, DeclinedElicitation):
+                return {
+                    "status": "error",
+                    "error": "Username declined by user",
+                    "message": "Cannot generate greeting without username",
+                }
+            elif isinstance(result, CancelledElicitation):
+                return {
+                    "status": "error",
+                    "error": "Operation cancelled",
+                    "message": "Greeting request was cancelled",
+                }
+        except Exception:
+            # Elicitation not supported by client - return graceful skip
+            return {
+                "status": "skipped",
+                "message": (
+                    "Elicitation not supported by client. "
+                    "Username parameter is required when client does not support elicitation."
+                ),
+                "elicitation_supported": False,
+            }
+
+    return {
+        "status": "success",
+        "message": f"Hello, {username}! Welcome to the DataRobot MCP server.",
+        "username": username,
+    }

datarobot_genai/drmcp/test_utils/integration_mcp_server.py

@@ -26,6 +26,13 @@ from typing import Any
 
 from datarobot_genai.drmcp import create_mcp_server
 
+# Import elicitation test tool to register it with the MCP server
+try:
+    from datarobot_genai.drmcp.test_utils import elicitation_test_tool  # noqa: F401
+except ImportError:
+    # Test utils not available (e.g., running in production)
+    pass
+
 # Import user components (will be used conditionally)
 try:
     from app.core.server_lifecycle import ServerLifecycle  # type: ignore  # noqa: F401

datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@@ -15,6 +15,7 @@ import asyncio
 import os
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
+from typing import Any
 
 import aiohttp
 from aiohttp import ClientSession as HttpClientSession
@@ -78,6 +79,7 @@ def get_headers() -> dict[str, str]:
 @asynccontextmanager
 async def ete_test_mcp_session(
     additional_headers: dict[str, str] | None = None,
+    elicitation_callback: Any | None = None,
 ) -> AsyncGenerator[ClientSession, None]:
     """Create an MCP session for each test.
 
@@ -85,6 +87,10 @@ async def ete_test_mcp_session(
     ----------
     additional_headers : dict[str, str], optional
         Additional headers to include in the MCP session (e.g., auth headers for testing).
+    elicitation_callback : callable, optional
+        Callback function to handle elicitation requests from the server.
+        The callback should have signature:
+        async def callback(context, params: ElicitRequestParams) -> ElicitResult
     """
     try:
         headers = get_headers()
@@ -96,7 +102,9 @@ async def ete_test_mcp_session(
             write_stream,
             _,
         ):
-            async with ClientSession(read_stream, write_stream) as session:
+            async with ClientSession(
+                read_stream, write_stream, elicitation_callback=elicitation_callback
+            ) as session:
                 await asyncio.wait_for(session.initialize(), timeout=5)
                 yield session
     except asyncio.TimeoutError:

datarobot_genai/drmcp/test_utils/mcp_utils_integration.py

@@ -17,6 +17,7 @@ import contextlib
 import os
 from collections.abc import AsyncGenerator
 from pathlib import Path
+from typing import Any
 
 from mcp import ClientSession
 from mcp.client.stdio import StdioServerParameters
@@ -34,7 +35,12 @@ def integration_test_mcp_server_params() -> StdioServerParameters:
         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",
+        # Disable all OTEL telemetry for integration tests
+        "OTEL_ENABLED": "false",
+        "OTEL_SDK_DISABLED": "true",
+        "OTEL_TRACES_EXPORTER": "none",
+        "OTEL_LOGS_EXPORTER": "none",
+        "OTEL_METRICS_EXPORTER": "none",
         "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP": os.environ.get(
             "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP"
         )
@@ -64,7 +70,9 @@ def integration_test_mcp_server_params() -> StdioServerParameters:
 
 @contextlib.asynccontextmanager
 async def integration_test_mcp_session(
-    server_params: StdioServerParameters | None = None, timeout: int = 30
+    server_params: StdioServerParameters | None = None,
+    timeout: int = 30,
+    elicitation_callback: Any | None = None,
 ) -> AsyncGenerator[ClientSession, None]:
     """
     Create and connect a client for the MCP server as a context manager.
@@ -72,6 +80,7 @@ async def integration_test_mcp_session(
     Args:
         server_params: Parameters for configuring the server connection
         timeout: Timeout
+        elicitation_callback: Optional callback for handling elicitation requests
 
     Yields
     ------
@@ -86,8 +95,12 @@ async def integration_test_mcp_session(
 
     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)
+            async with ClientSession(
+                read_stream, write_stream, elicitation_callback=elicitation_callback
+            ) as session:
+                init_result = await asyncio.wait_for(session.initialize(), timeout=timeout)
+                # Store the init result on the session for tests that need to inspect capabilities
+                session._init_result = init_result  # type: ignore[attr-defined]
                 yield session
 
     except asyncio.TimeoutError:

datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+from ast import literal_eval
 from typing import Any
 
 import openai
@@ -44,12 +45,39 @@ class LLMResponse:
 
 
 class LLMMCPClient:
-    """Client for interacting with LLMs via MCP."""
+    """
+    Client for interacting with LLMs via MCP.
 
-    def __init__(self, config: str):
-        """Initialize the LLM MCP client."""
+    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,
+    ):
+        """
+        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)
+        """
         # Parse config string to extract parameters
-        config_dict = eval(config) if isinstance(config, str) else config
+        if isinstance(config, str):
+            # Try JSON first (safer), fall back to literal_eval for Python dict strings
+            try:
+                config_dict = 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")
@@ -93,7 +121,21 @@ class LLMMCPClient:
     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."""
+        """
+        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
@@ -177,7 +219,26 @@ class LLMMCPClient:
     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."""
+        """
+        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)
 
@@ -191,8 +252,10 @@ class LLMMCPClient:
                 "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."
+                    "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},

datarobot_genai/drmcp/test_utils/test_interactive.py

@@ -0,0 +1,205 @@
+#!/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.
+
+"""Interactive MCP Client Test Script.
+
+This script allows you to test arbitrary commands with the MCP server
+using an LLM agent that can decide which tools to call.
+
+Supports elicitation - when tools require user input (like authentication tokens),
+the script will prompt you interactively.
+"""
+
+import asyncio
+import json
+import os
+import sys
+import traceback
+from pathlib import Path
+from typing import Any
+
+from dotenv import load_dotenv
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.context import RequestContext
+from mcp.types import ElicitRequestParams
+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
+
+# Re-export for backwards compatibility
+__all__ = ["LLMMCPClient", "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")
+        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")
+
+    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,
+        "save_llm_responses": False,
+    }
+
+    llm_client = LLMMCPClient(str(config))
+
+    # Get MCP server URL
+    mcp_server_url = get_dr_mcp_server_url()
+    if not mcp_server_url:
+        print("❌ Error: MCP server URL is not configured")
+        print("Please set DR_MCP_SERVER_URL environment variable or run: task test-interactive")
+        return
+
+    print(f"🔗 Connecting to MCP server at: {mcp_server_url}")
+
+    # Elicitation handler: prompt user for required values
+    async def elicitation_handler(
+        context: RequestContext[ClientSession, Any], params: ElicitRequestParams
+    ) -> ElicitResult:
+        print(f"\n📋 Elicitation Request: {params.message}")
+        if params.requestedSchema:
+            print(f"   Schema: {params.requestedSchema}")
+
+        while True:
+            try:
+                response = input("   Enter value (or 'decline'/'cancel'): ").strip()
+            except (EOFError, KeyboardInterrupt):
+                return ElicitResult(action="cancel")
+
+            if response.lower() == "decline":
+                return ElicitResult(action="decline")
+            if response.lower() == "cancel":
+                return ElicitResult(action="cancel")
+            if response:
+                return ElicitResult(action="accept", content={"value": response})
+            print("   Please enter a value or 'decline'/'cancel'")
+
+    try:
+        async with streamablehttp_client(
+            url=mcp_server_url,
+            headers=get_headers(),
+        ) as (read_stream, write_stream, _):
+            async with ClientSession(
+                read_stream,
+                write_stream,
+                elicitation_callback=elicitation_handler,
+            ) as session:
+                await session.initialize()
+
+                print("✅ Connected to MCP server!")
+                print("📋 Available tools:")
+
+                tools_result = await session.list_tools()
+                for i, tool in enumerate(tools_result.tools, 1):
+                    print(f"  {i}. {tool.name}: {tool.description}")
+
+                print("\n" + "=" * 60)
+                print("🎯 Interactive Testing Mode")
+                print("=" * 60)
+                print("Type your questions/commands. The AI will decide which tools to use.")
+                print("If a tool requires additional information, you will be prompted.")
+                print("Type 'quit' or 'exit' to stop.")
+                print()
+
+                while True:
+                    try:
+                        user_input = input("🤔 You: ").strip()
+
+                        if user_input.lower() in ["quit", "exit", "q"]:
+                            print("👋 Goodbye!")
+                            break
+
+                        if not user_input:
+                            continue
+                    except (EOFError, KeyboardInterrupt):
+                        print("\n👋 Goodbye!")
+                        break
+
+                    print("🤖 AI is thinking...")
+
+                    response = await llm_client.process_prompt_with_mcp_support(
+                        prompt=user_input,
+                        mcp_session=session,
+                    )
+
+                    print("\n🤖 AI Response:")
+                    print("-" * 40)
+                    print(response.content)
+
+                    if response.tool_calls:
+                        print("\n🔧 Tools Used:")
+                        for i, tool_call in enumerate(response.tool_calls, 1):
+                            print(f"  {i}. {tool_call.tool_name}")
+                            print(f"     Parameters: {tool_call.parameters}")
+                            print(f"     Reasoning: {tool_call.reasoning}")
+
+                            if i <= len(response.tool_results):
+                                result = response.tool_results[i - 1]
+                                try:
+                                    result_data = json.loads(result)
+                                    if result_data.get("status") == "error":
+                                        error_msg = result_data.get("error", "Unknown error")
+                                        print(f"     ❌ Error: {error_msg}")
+                                    elif result_data.get("status") == "success":
+                                        print("     ✅ Success")
+                                except json.JSONDecodeError:
+                                    if len(result) > 100:
+                                        print(f"     Result: {result[:100]}...")
+                                    else:
+                                        print(f"     Result: {result}")
+
+                    print("\n" + "=" * 60)
+    except Exception as e:
+        print(f"❌ Connection Error: {e}")
+        print(f"   Server URL: {mcp_server_url}")
+        traceback.print_exc()
+        return
+
+
+if __name__ == "__main__":
+    # Ensure we're in the right directory
+    if not Path("src").exists():
+        print("❌ Error: Please run this script from the project root")
+        sys.exit(1)
+
+    # Load environment variables from .env file
+    print("📄 Loading environment variables...")
+    load_dotenv()
+
+    print("🚀 Starting Interactive MCP Client Test")
+    print("Make sure the MCP server is running with: task drmcp-dev")
+    print()
+
+    asyncio.run(test_mcp_interactive())

datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -116,28 +116,30 @@ class ToolBaseE2E:
                     f"Should have called {test_expectations.tool_calls_expected[i].name} tool, but "
                     f"got: {tool_call.tool_name}"
                 )
-            assert tool_call.parameters == test_expectations.tool_calls_expected[i].parameters, (
-                f"Should have called {tool_call.tool_name} tool with the correct parameters, but "
-                f"got: {tool_call.parameters}"
-            )
-            if test_expectations.tool_calls_expected[i].result != SHOULD_NOT_BE_EMPTY:
-                expected_result = test_expectations.tool_calls_expected[i].result
-                if isinstance(expected_result, str):
-                    assert expected_result in response.tool_results[i], (
-                        f"Should have called {tool_call.tool_name} tool with the correct result, "
-                        f"but got: {response.tool_results[i]}"
-                    )
+                assert (
+                    tool_call.parameters == test_expectations.tool_calls_expected[i].parameters
+                ), (
+                    f"Should have called {tool_call.tool_name} tool with the correct parameters, "
+                    f"but got: {tool_call.parameters}"
+                )
+                if test_expectations.tool_calls_expected[i].result != SHOULD_NOT_BE_EMPTY:
+                    expected_result = test_expectations.tool_calls_expected[i].result
+                    if isinstance(expected_result, str):
+                        assert expected_result in response.tool_results[i], (
+                            f"Should have called {tool_call.tool_name} tool with the correct "
+                            f"result, but got: {response.tool_results[i]}"
+                        )
+                    else:
+                        actual_result = json.loads(response.tool_results[i])
+                        assert _check_dict_has_keys(expected_result, actual_result), (
+                            f"Should have called {tool_call.tool_name} tool with the correct "
+                            f"result structure, but got: {response.tool_results[i]}"
+                        )
                 else:
-                    actual_result = json.loads(response.tool_results[i])
-                    assert _check_dict_has_keys(expected_result, actual_result), (
-                        f"Should have called {tool_call.tool_name} tool with the correct result "
-                        f"structure, but got: {response.tool_results[i]}"
+                    assert len(response.tool_results[i]) > 0, (
+                        f"Should have called {tool_call.tool_name} tool with non-empty result, but "
+                        f"got: {response.tool_results[i]}"
                     )
-            else:
-                assert len(response.tool_results[i]) > 0, (
-                    f"Should have called {tool_call.tool_name} tool with non-empty result, but "
-                    f"got: {response.tool_results[i]}"
-                )
 
         # Verify LLM provided comprehensive response
         assert len(response.content) > 100, "LLM should provide detailed response"