datarobot-genai 0.2.39__py3-none-any.whl → 0.3.1__py3-none-any.whl

datarobot_genai/drmcp/core/mcp_instance.py

@@ -25,7 +25,6 @@ from fastmcp.prompts.prompt import Prompt
 from fastmcp.server.dependencies import get_context
 from fastmcp.tools import Tool
 from mcp.types import AnyFunction
-from mcp.types import Tool as MCPTool
 from mcp.types import ToolAnnotations
 from typing_extensions import Unpack
 
@@ -36,8 +35,6 @@ from .logging import log_execution
 from .memory_management.manager import MemoryManager
 from .memory_management.manager import get_memory_manager
 from .telemetry import trace_execution
-from .tool_filter import filter_tools_by_tags
-from .tool_filter import list_all_tags
 
 logger = logging.getLogger(__name__)
 
@@ -80,10 +77,8 @@ async def get_agent_and_storage_ids(
     return agent_id, storage_id
 
 
-class TaggedFastMCP(FastMCP):
-    """Extended FastMCP that supports tags, deployments and other annotations directly in the
-    tool decorator.
-    """
+class DataRobotMCP(FastMCP):
+    """Extended FastMCP that supports DataRobot specific features like deployments and prompts."""
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
@@ -118,45 +113,6 @@ class TaggedFastMCP(FastMCP):
                 "In stateless mode, clients will see changes on next request."
             )
 
-    async def list_tools(
-        self, tags: list[str] | None = None, match_all: bool = False
-    ) -> list[MCPTool]:
-        """
-        List all available tools, optionally filtered by tags.
-
-        Args:
-            tags: Optional list of tags to filter by. If None, returns all tools.
-            match_all: If True, tool must have all specified tags (AND logic).
-                      If False, tool must have at least one tag (OR logic).
-                      Only used when tags is provided.
-
-        Returns
-        -------
-            List of MCPTool objects that match the tag criteria.
-        """
-        # Get all tools from the parent class
-        all_tools = await super()._list_tools_mcp()
-
-        # If no tags specified, return all tools
-        if not tags:
-            return all_tools
-
-        # Filter tools by tags
-        filtered_tools = filter_tools_by_tags(list(all_tools), tags, match_all)
-
-        return filtered_tools  # type: ignore[return-value]
-
-    async def get_all_tags(self) -> list[str]:
-        """
-        Get all unique tags from all registered tools.
-
-        Returns
-        -------
-            List of all unique tags sorted alphabetically.
-        """
-        all_tools = await self._list_tools_mcp()
-        return list_all_tags(list(all_tools))
-
     async def get_deployment_mapping(self) -> dict[str, str]:
         """
         Get the list of deployment IDs for all registered dynamic tools.
@@ -279,10 +235,10 @@ class TaggedFastMCP(FastMCP):
             )
 
 
-# Create the tagged MCP instance
+# Create the DataRobot MCP instance
 mcp_server_configs: MCPServerConfig = get_config()
 
-mcp = TaggedFastMCP(
+mcp = DataRobotMCP(
     name=mcp_server_configs.mcp_server_name,
     on_duplicate_tools=mcp_server_configs.tool_registration_duplicate_behavior,
     on_duplicate_prompts=mcp_server_configs.prompt_registration_duplicate_behavior,
@@ -446,7 +402,7 @@ async def register_tools(
         await mcp.set_deployment_mapping(deployment_id, tool_name)
 
     # Verify tool is registered
-    tools = await mcp.list_tools()
+    tools = await mcp._list_tools_mcp()
     if not any(tool.name == tool_name for tool in tools):
         raise RuntimeError(f"Tool {tool_name} was not registered successfully")
     logger.info(f"Registered tools: {len(tools)}")

datarobot_genai/drmcp/core/routes.py

@@ -18,20 +18,28 @@ from botocore.exceptions import ClientError
 from starlette.requests import Request
 from starlette.responses import JSONResponse
 
+from datarobot_genai import __version__
+
+from .config import get_config
 from .dynamic_prompts.controllers import delete_registered_prompt_template
 from .dynamic_prompts.controllers import refresh_registered_prompt_template
 from .dynamic_prompts.controllers import register_prompt_from_prompt_template_id_and_version
 from .dynamic_tools.deployment.controllers import delete_registered_tool_deployment
 from .dynamic_tools.deployment.controllers import get_registered_tool_deployments
 from .dynamic_tools.deployment.controllers import register_tool_for_deployment_id
-from .mcp_instance import TaggedFastMCP
+from .mcp_instance import DataRobotMCP
 from .memory_management.manager import get_memory_manager
 from .routes_utils import prefix_mount_path
+from .tool_config import TOOL_CONFIGS
+from .tool_config import ToolType
+from .utils import get_prompt_tags
+from .utils import get_resource_tags
+from .utils import get_tool_tags
 
 logger = getLogger(__name__)
 
 
-def register_routes(mcp: TaggedFastMCP) -> None:
+def register_routes(mcp: DataRobotMCP) -> None:
     """Register all routes with the MCP server."""
 
     @mcp.custom_route(prefix_mount_path("/"), methods=["GET"])
@@ -44,26 +52,113 @@ def register_routes(mcp: TaggedFastMCP) -> None:
             },
         )
 
-    # Custom endpoint to get all tags
-    @mcp.custom_route(prefix_mount_path("/tags"), methods=["GET"])
-    async def handle_tags(_: Request) -> JSONResponse:
+    @mcp.custom_route(prefix_mount_path("/metadata"), methods=["GET"])
+    async def get_metadata(_: Request) -> JSONResponse:
+        """Get metadata about tools, prompts, resources, and system configuration."""
         try:
-            # TaggedFastMCP extends FastMCP with get_all_tags
-            tags = await mcp.get_all_tags()  # type: ignore[attr-defined]
+            # Get tools with tags
+            tools = await mcp._list_tools_mcp()
+            tools_metadata = [
+                {
+                    "name": tool.name,
+                    "tags": sorted(list(get_tool_tags(tool))),
+                }
+                for tool in tools
+            ]
+
+            # Get prompts with tags
+            prompts = await mcp._list_prompts_mcp()
+            prompts_metadata = [
+                {
+                    "name": prompt.name,
+                    "tags": sorted(list(get_prompt_tags(prompt))),
+                }
+                for prompt in prompts
+            ]
+
+            # Get resources with tags
+            resources = await mcp._list_resources_mcp()
+            resources_metadata = [
+                {
+                    "name": resource.name,
+                    "tags": sorted(list(get_resource_tags(resource))),
+                }
+                for resource in resources
+            ]
+
+            # Get safe configuration details
+            config = get_config()
+
+            # Build tool config status
+            tool_config_status = {}
+            for tool_type in ToolType:
+                tool_config = TOOL_CONFIGS[tool_type]
+                is_enabled = getattr(config.tool_config, tool_config["config_field_name"], False)
+                oauth_check_fn = tool_config["oauth_check"]
+                oauth_required = oauth_check_fn is not None
+                oauth_configured = None
+                if oauth_required and oauth_check_fn is not None:
+                    oauth_configured = oauth_check_fn(config)
+
+                tool_config_status[tool_type.value] = {
+                    "enabled": is_enabled,
+                    "oauth_required": oauth_required,
+                    "oauth_configured": oauth_configured,
+                }
+
+            # Safe config details (excluding sensitive information)
+            safe_config = {
+                "server": {
+                    "name": config.mcp_server_name,
+                    "port": config.mcp_server_port,
+                    "log_level": config.mcp_server_log_level,
+                    "app_log_level": config.app_log_level,
+                    "mount_path": config.mount_path,
+                    "drmcp_genai_version": __version__,
+                },
+                "features": {
+                    "register_dynamic_tools_on_startup": (
+                        config.mcp_server_register_dynamic_tools_on_startup
+                    ),
+                    "register_dynamic_prompts_on_startup": (
+                        config.mcp_server_register_dynamic_prompts_on_startup
+                    ),
+                    "tool_registration_allow_empty_schema": (
+                        config.tool_registration_allow_empty_schema
+                    ),
+                    "tool_registration_duplicate_behavior": (
+                        config.tool_registration_duplicate_behavior
+                    ),
+                    "prompt_registration_duplicate_behavior": (
+                        config.prompt_registration_duplicate_behavior
+                    ),
+                },
+                "tool_config": tool_config_status,
+            }
+
             return JSONResponse(
                 status_code=HTTPStatus.OK,
                 content={
-                    "tags": tags,
-                    "count": len(tags),
-                    "message": "All available tags retrieved successfully",
+                    "tools": {
+                        "items": tools_metadata,
+                        "count": len(tools_metadata),
+                    },
+                    "prompts": {
+                        "items": prompts_metadata,
+                        "count": len(prompts_metadata),
+                    },
+                    "resources": {
+                        "items": resources_metadata,
+                        "count": len(resources_metadata),
+                    },
+                    "config": safe_config,
                 },
             )
         except Exception as e:
+            logger.exception("Failed to retrieve metadata")
             return JSONResponse(
                 status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
-                content={
-                    "error": f"Failed to retrieve tags: {str(e)}",
-                },
+                content={"error": f"Failed to retrieve metadata: {str(e)}"},
             )
 
     memory_manager = get_memory_manager()

datarobot_genai/drmcp/core/tool_config.py

@@ -31,6 +31,8 @@ class ToolType(str, Enum):
     CONFLUENCE = "confluence"
     GDRIVE = "gdrive"
     MICROSOFT_GRAPH = "microsoft_graph"
+    PERPLEXITY = "perplexity"
+    TAVILY = "tavily"
 
 
 class ToolConfig(TypedDict):
@@ -80,6 +82,20 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
         package_prefix="datarobot_genai.drmcp.tools.microsoft_graph",
         config_field_name="enable_microsoft_graph_tools",
     ),
+    ToolType.PERPLEXITY: ToolConfig(
+        name="perplexity",
+        oauth_check=None,  # OAuth for Perplexity is not supported
+        directory="perplexity",
+        package_prefix="datarobot_genai.drmcp.tools.perplexity",
+        config_field_name="enable_perplexity_tools",
+    ),
+    ToolType.TAVILY: ToolConfig(
+        name="tavily",
+        oauth_check=None,
+        directory="tavily",
+        package_prefix="datarobot_genai.drmcp.tools.tavily",
+        config_field_name="enable_tavily_tools",
+    ),
 }
 
 

datarobot_genai/drmcp/core/utils.py

@@ -18,7 +18,11 @@ from urllib.parse import urlparse
 
 import boto3
 from fastmcp.resources import HttpResource
+from fastmcp.tools import Tool
 from fastmcp.tools.tool import ToolResult
+from mcp.types import Prompt as MCPPrompt
+from mcp.types import Resource as MCPResource
+from mcp.types import Tool as MCPTool
 from pydantic import BaseModel
 
 from .constants import MAX_INLINE_SIZE
@@ -136,3 +140,109 @@ def is_valid_url(url: str) -> bool:
     """Check if a URL is valid."""
     result = urlparse(url)
     return all([result.scheme, result.netloc])
+
+
+def get_prompt_tags(prompt: MCPPrompt) -> set[str]:
+    """
+    Extract tags from a prompt.
+
+    Args:
+        prompt: MCP protocol Prompt
+
+    Returns
+    -------
+        Set of tag strings, empty set if no tags found
+    """
+    # MCPPrompt has tags in meta._fastmcp.tags (as a list)
+    if not (prompt.meta and isinstance(prompt.meta, dict)):
+        return set()
+
+    fastmcp_meta = prompt.meta.get("_fastmcp")
+    if not (fastmcp_meta and isinstance(fastmcp_meta, dict)):
+        return set()
+
+    tags = fastmcp_meta.get("tags")
+    return set(tags) if tags else set()
+
+
+def get_resource_tags(resource: MCPResource) -> set[str]:
+    """
+    Extract tags from a resource.
+
+    Args:
+        resource: MCP protocol Resource
+
+    Returns
+    -------
+        Set of tag strings, empty set if no tags found
+    """
+    # MCPResource has tags in meta._fastmcp.tags (as a list)
+    if not (resource.meta and isinstance(resource.meta, dict)):
+        return set()
+
+    fastmcp_meta = resource.meta.get("_fastmcp")
+    if not (fastmcp_meta and isinstance(fastmcp_meta, dict)):
+        return set()
+
+    tags = fastmcp_meta.get("tags")
+    return set(tags) if tags else set()
+
+
+def get_tool_tags(tool: Tool | MCPTool) -> set[str]:
+    """
+    Extract tags from a tool, handling both FastMCP Tool and MCP protocol Tool types.
+
+    Args:
+        tool: Either a FastMCP Tool or MCP protocol Tool
+
+    Returns
+    -------
+        Set of tag strings, empty set if no tags found
+    """
+    if isinstance(tool, Tool):
+        # FastMCP Tool has tags directly as a set
+        return getattr(tool, "tags", None) or set()
+
+    # MCPTool has tags in meta._fastmcp.tags (as a list)
+    if not (tool.meta and isinstance(tool.meta, dict)):
+        return set()
+
+    fastmcp_meta = tool.meta.get("_fastmcp")
+    if not (fastmcp_meta and isinstance(fastmcp_meta, dict)):
+        return set()
+
+    tags = fastmcp_meta.get("tags")
+    return set(tags) if tags else set()
+
+
+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
+
+    # Convert tags to set for O(1) lookup instead of O(n)
+    tags_set = set(tags)
+
+    return [
+        tool
+        for tool in tools
+        if (tool_tags := get_tool_tags(tool))
+        and (tags_set.issubset(tool_tags) if match_all else tags_set & tool_tags)
+    ]

datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -39,6 +39,17 @@ class ETETestExpectations(BaseModel):
 SHOULD_NOT_BE_EMPTY = "SHOULD_NOT_BE_EMPTY"
 
 
+def _extract_content_when_structured_empty(tool_result: str) -> dict[str, str] | None:
+    r"""When Content is present but Structured content is empty, return {"error": content}."""
+    if "Content: " not in tool_result:
+        return None
+    content_part = tool_result.split("Content: ", 1)[1]
+    if "\nStructured content: " in content_part:
+        content_part = content_part.split("\nStructured content: ", 1)[0]
+    content_part = content_part.strip()
+    return {"error": content_part} if content_part else None
+
+
 def _extract_structured_content(tool_result: str) -> Any:
     r"""
     Extract and parse structured content from tool result string.
@@ -49,7 +60,9 @@ def _extract_structured_content(tool_result: str) -> Any:
     Structured content can be:
     1. A JSON object with a "result" key: {"result": "..."} or {"result": "{...}"}
     2. A direct JSON object: {"key": "value", ...}
-    3. Empty or missing
+    3. Empty or missing — when Content is present but Structured content is empty
+       (e.g. tool errors), returns {"error": content} so dict expectations can validate.
+    4. None if neither valid structured content nor Content is available
 
     Args:
         tool_result: The tool result string
@@ -58,33 +71,35 @@ def _extract_structured_content(tool_result: str) -> Any:
     -------
         Parsed structured content, or None if not available
     """
-    # Early returns for invalid inputs
-    if not tool_result or "Structured content: " not in tool_result:
-        return None
-
-    structured_part = tool_result.split("Structured content: ", 1)[1].strip()
-    # Parse JSON, return None on failure or empty structured_part
-    if not structured_part:
-        return None
-    try:
-        structured_data = json.loads(structured_part)
-    except json.JSONDecodeError:
-        return None
-
-    # If structured data has a "result" key, extract and parse that
-    if isinstance(structured_data, dict) and "result" in structured_data:
-        result_value = structured_data["result"]
-        # If result is a JSON string (starts with { or [), try to parse it
-        if isinstance(result_value, str) and result_value.strip().startswith(("{", "[")):
+    result: Any = None
+    if not tool_result:
+        pass
+    elif "Structured content: " in tool_result:
+        structured_part = tool_result.split("Structured content: ", 1)[1].strip()
+        if not structured_part:
+            result = _extract_content_when_structured_empty(tool_result)
+        else:
             try:
-                parsed_result = json.loads(result_value)
+                structured_data = json.loads(structured_part)
             except json.JSONDecodeError:
-                parsed_result = result_value  # Return string as-is if parsing fails
-            return parsed_result
-        return result_value  # Return result value directly
-
-    # If it's a direct JSON object (not wrapped in {"result": ...}), return it as-is
-    return structured_data
+                pass
+            else:
+                if isinstance(structured_data, dict) and "result" in structured_data:
+                    result_value = structured_data["result"]
+                    if isinstance(result_value, str) and result_value.strip().startswith(
+                        ("{", "[")
+                    ):
+                        try:
+                            result = json.loads(result_value)
+                        except json.JSONDecodeError:
+                            result = result_value
+                    else:
+                        result = result_value
+                else:
+                    result = structured_data
+    else:
+        result = _extract_content_when_structured_empty(tool_result)
+    return result
 
 
 def _check_dict_has_keys(

datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -351,6 +351,8 @@ class GoogleDriveClient:
             "pageSize": page_size,
             "fields": DEFAULT_FIELDS,
             "orderBy": DEFAULT_ORDER,
+            "supportsAllDrives": "true",
+            "includeItemsFromAllDrives": "true",
         }
         if page_token:
             params["pageToken"] = page_token

datarobot_genai/drmcp/tools/clients/microsoft_graph.py

@@ -586,6 +586,102 @@ class MicrosoftGraphClient:
                 f"Microsoft Graph API error {response.status_code}: {response.text}"
             )
 
+    async def update_item_metadata(
+        self,
+        item_id: str,
+        fields_to_update: dict[str, Any],
+        site_id: str | None = None,
+        list_id: str | None = None,
+        drive_id: str | None = None,
+    ) -> dict[str, Any]:
+        """Update metadata on a SharePoint list item or OneDrive/SharePoint drive item.
+
+        For SharePoint list items: Updates custom column values via the fields endpoint.
+        For drive items: Updates properties like name and description.
+
+        Args:
+            item_id: The ID of the item to update.
+            fields_to_update: Key-value pairs of metadata fields to modify.
+            site_id: For SharePoint list items - the site ID.
+            list_id: For SharePoint list items - the list ID.
+            drive_id: For OneDrive/drive items - the drive ID.
+
+        Returns
+        -------
+            The API response containing updated item data.
+
+        Raises
+        ------
+            MicrosoftGraphError: If validation fails or the API request fails.
+        """
+        if not item_id or not item_id.strip():
+            raise MicrosoftGraphError("item_id cannot be empty")
+        if not fields_to_update:
+            raise MicrosoftGraphError("fields_to_update cannot be empty")
+
+        # Determine the endpoint based on provided parameters
+        has_sharepoint_context = site_id is not None and list_id is not None
+        has_drive_context = drive_id is not None
+
+        if has_sharepoint_context and has_drive_context:
+            raise MicrosoftGraphError(
+                "Cannot specify both SharePoint (site_id + list_id) and OneDrive "
+                "(document_library_id) context. Choose one."
+            )
+
+        if not has_sharepoint_context and not has_drive_context:
+            raise MicrosoftGraphError(
+                "Must specify either SharePoint context (site_id + list_id) or "
+                "OneDrive context (document_library_id)."
+            )
+
+        try:
+            if has_sharepoint_context:
+                # PATCH /sites/{site-id}/lists/{list-id}/items/{item-id}/fields
+                url = f"{GRAPH_API_BASE}/sites/{site_id}/lists/{list_id}/items/{item_id}/fields"
+                response = await self._client.patch(url, json=fields_to_update)
+            else:
+                # Drive item: PATCH /drives/{drive-id}/items/{item-id}
+                url = f"{GRAPH_API_BASE}/drives/{drive_id}/items/{item_id}"
+                response = await self._client.patch(url, json=fields_to_update)
+
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPStatusError as e:
+            raise self._handle_update_metadata_error(e, item_id) from e
+
+    def _handle_update_metadata_error(
+        self,
+        error: httpx.HTTPStatusError,
+        item_id: str,
+    ) -> MicrosoftGraphError:
+        """Handle HTTP errors for metadata updates and return appropriate MicrosoftGraphError."""
+        status_code = error.response.status_code
+        error_msg = f"Failed to update metadata: HTTP {status_code}"
+
+        if status_code == 400:
+            try:
+                error_data = error.response.json()
+                api_message = error_data.get("error", {}).get("message", "Invalid request")
+                error_msg = f"Bad request updating metadata: {api_message}"
+            except Exception:
+                error_msg = "Bad request: invalid field names or values."
+        elif status_code == 401:
+            error_msg = "Authentication failed. Access token may be expired or invalid."
+        elif status_code == 403:
+            error_msg = (
+                f"Permission denied: you don't have permission to update item '{item_id}'. "
+                "Requires Sites.ReadWrite.All or Files.ReadWrite.All permission."
+            )
+        elif status_code == 404:
+            error_msg = f"Item '{item_id}' not found."
+        elif status_code == 409:
+            error_msg = f"Conflict: item '{item_id}' may have been modified concurrently."
+        elif status_code == 429:
+            error_msg = "Rate limit exceeded. Please try again later."
+
+        return MicrosoftGraphError(error_msg)
+
     async def __aenter__(self) -> "MicrosoftGraphClient":
         """Async context manager entry."""
         return self