datarobot-genai 0.2.17__tar.gz → 0.2.25__tar.gz

{datarobot_genai-0.2.17 → datarobot_genai-0.2.25}/PKG-INFO

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

{datarobot_genai-0.2.17 → datarobot_genai-0.2.25}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.17"
+version = "0.2.25"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -205,6 +205,9 @@ asyncio_mode = "auto"
 source = ["datarobot_genai"]
 omit = [
     "*/__init__.py",
+    "*/test_utils/*",
+    # nat requires Python >=3.11, can't be imported on 3.10 for coverage
+    "*/nat/*",
 ]
 
 [tool.coverage.report]

{datarobot_genai-0.2.17 → datarobot_genai-0.2.25}/src/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-0.2.17 → datarobot_genai-0.2.25}/src/datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -31,9 +31,6 @@ from .dynamic_prompts.register import register_prompts_from_datarobot_prompt_man
 from .dynamic_tools.deployment.register import register_tools_of_datarobot_deployments
 from .logging import MCPLogging
 from .mcp_instance import mcp
-from .mcp_server_tools import get_all_available_tags  # noqa # pylint: disable=unused-import
-from .mcp_server_tools import get_tool_info_by_name  # noqa # pylint: disable=unused-import
-from .mcp_server_tools import list_tools_by_tags  # noqa # pylint: disable=unused-import
 from .memory_management.manager import MemoryManager
 from .routes import register_routes
 from .routes_utils import prefix_mount_path

{datarobot_genai-0.2.17 → datarobot_genai-0.2.25}/src/datarobot_genai/drmcp/core/mcp_instance.py

@@ -16,17 +16,13 @@ import logging
 from collections.abc import Callable
 from functools import wraps
 from typing import Any
-from typing import overload
 
 from fastmcp import Context
 from fastmcp import FastMCP
 from fastmcp.exceptions import NotFoundError
 from fastmcp.prompts.prompt import Prompt
 from fastmcp.server.dependencies import get_context
-from fastmcp.tools import FunctionTool
 from fastmcp.tools import Tool
-from fastmcp.utilities.types import NotSet
-from fastmcp.utilities.types import NotSetT
 from mcp.types import AnyFunction
 from mcp.types import Tool as MCPTool
 from mcp.types import ToolAnnotations
@@ -120,86 +116,6 @@ class TaggedFastMCP(FastMCP):
                 "In stateless mode, clients will see changes on next request."
             )
 
-    @overload
-    def tool(
-        self,
-        name_or_fn: AnyFunction,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-    ) -> FunctionTool: ...
-
-    @overload
-    def tool(
-        self,
-        name_or_fn: str | None = None,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-    ) -> Callable[[AnyFunction], FunctionTool]: ...
-
-    def tool(
-        self,
-        name_or_fn: str | Callable[..., Any] | None = None,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-        **kwargs: Any,
-    ) -> Callable[[AnyFunction], FunctionTool] | FunctionTool:
-        """
-        Extend tool decorator that supports tags and other annotations, while remaining
-        signature-compatible with FastMCP.tool to avoid recursion issues with partials.
-        """
-        if isinstance(annotations, dict):
-            annotations = ToolAnnotations(**annotations)
-
-        # Ensure tags are available both via native fastmcp `tags` and inside annotations
-        if tags is not None:
-            tags_ = sorted(tags)
-            if annotations is None:
-                annotations = ToolAnnotations()  # type: ignore[call-arg]
-                annotations.tags = tags_  # type: ignore[attr-defined, union-attr]
-            else:
-                # At this point, annotations is ToolAnnotations (not dict)
-                assert isinstance(annotations, ToolAnnotations)
-                annotations.tags = tags_  # type: ignore[attr-defined]
-
-        return super().tool(
-            name_or_fn,
-            name=name,
-            title=title,
-            description=description,
-            tags=tags,
-            output_schema=output_schema
-            if output_schema is not None
-            else kwargs.get("output_schema"),
-            annotations=annotations,
-            exclude_args=exclude_args,
-            meta=meta,
-            enabled=enabled,
-        )
-
     async def list_tools(
         self, tags: list[str] | None = None, match_all: bool = False
     ) -> list[MCPTool]:
@@ -488,11 +404,10 @@ async def register_tools(
     # Apply dr_mcp_extras to the memory-aware function
     wrapped_fn = dr_mcp_extras()(memory_aware_fn)
 
-    # Create annotations with tags, deployment_id if provided
-    annotations = ToolAnnotations()  # type: ignore[call-arg]
-    if tags is not None:
-        annotations.tags = tags  # type: ignore[attr-defined]
+    # Create annotations only when additional metadata is required
+    annotations: ToolAnnotations | None = None  # type: ignore[assignment]
     if deployment_id is not None:
+        annotations = ToolAnnotations()  # type: ignore[call-arg]
         annotations.deployment_id = deployment_id  # type: ignore[attr-defined]
 
     tool = Tool.from_function(

{datarobot_genai-0.2.17 → datarobot_genai-0.2.25}/src/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-0.2.17 → datarobot_genai-0.2.25}/src/datarobot_genai/drmcp/core/tool_filter.py

@@ -41,7 +41,7 @@ def filter_tools_by_tags(
     filtered_tools = []
 
     for tool in tools:
-        tool_tags = getattr(tool.annotations, "tags", []) if tool.annotations else []
+        tool_tags = get_tool_tags(tool)
 
         if not tool_tags:
             continue
@@ -68,9 +68,18 @@ def get_tool_tags(tool: Tool | MCPTool) -> list[str]:
     -------
         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 []
 
 

{datarobot_genai-0.2.17 → datarobot_genai-0.2.25}/src/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-0.2.25/src/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-0.2.17 → datarobot_genai-0.2.25}/src/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-0.2.17 → datarobot_genai-0.2.25}/src/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-0.2.17 → datarobot_genai-0.2.25}/src/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-0.2.17 → datarobot_genai-0.2.25}/src/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},