datarobot-genai 0.2.26__py3-none-any.whl → 0.2.34__py3-none-any.whl

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -16,11 +16,13 @@
 
 import logging
 from typing import Annotated
+from typing import Literal
 
 from fastmcp.exceptions import ToolError
 from fastmcp.tools.tool import ToolResult
 
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.tools.clients.gdrive import GOOGLE_DRIVE_FOLDER_MIME
 from datarobot_genai.drmcp.tools.clients.gdrive import LIMIT
 from datarobot_genai.drmcp.tools.clients.gdrive import MAX_PAGE_SIZE
 from datarobot_genai.drmcp.tools.clients.gdrive import SUPPORTED_FIELDS
@@ -32,7 +34,9 @@ from datarobot_genai.drmcp.tools.clients.gdrive import get_gdrive_access_token
 logger = logging.getLogger(__name__)
 
 
-@dr_mcp_tool(tags={"google", "gdrive", "list", "search", "files", "find", "contents"})
+@dr_mcp_tool(
+    tags={"google", "gdrive", "list", "search", "files", "find", "contents"}, enabled=False
+)
 async def gdrive_find_contents(
     *,
     page_size: Annotated[
@@ -60,7 +64,7 @@ async def gdrive_find_contents(
         "Optional list of metadata fields to include. Ex. id, name, mimeType. "
         f"Default = {SUPPORTED_FIELDS_STR}",
     ] = None,
-) -> ToolResult | ToolError:
+) -> ToolResult:
     """
     Search or list files in the user's Google Drive with pagination and filtering support.
     Use this tool to discover file names and IDs for use with other tools.
@@ -121,7 +125,7 @@ async def gdrive_read_content(
         "(e.g., 'text/markdown' for Docs, 'text/csv' for Sheets). "
         "If not specified, uses sensible defaults. Has no effect on regular files.",
     ] = None,
-) -> ToolResult | ToolError:
+) -> ToolResult:
     """
     Retrieve the content of a specific file by its ID. Google Workspace files are
     automatically exported to LLM-readable formats (Push-Down).
@@ -163,7 +167,6 @@ async def gdrive_read_content(
             f"An unexpected error occurred while reading Google Drive file content: {str(e)}"
         )
 
-    # Provide helpful context about the conversion
     export_info = ""
     if file_content.was_exported:
         export_info = f" (exported from {file_content.original_mime_type})"
@@ -175,3 +178,269 @@ async def gdrive_read_content(
         ),
         structured_content=file_content.as_flat_dict(),
     )
+
+
+@dr_mcp_tool(tags={"google", "gdrive", "create", "write", "file", "folder"}, enabled=False)
+async def gdrive_create_file(
+    *,
+    name: Annotated[str, "The name for the new file or folder."],
+    mime_type: Annotated[
+        str,
+        "The MIME type of the file (e.g., 'text/plain', "
+        "'application/vnd.google-apps.document', 'application/vnd.google-apps.folder').",
+    ],
+    parent_id: Annotated[
+        str | None, "The ID of the parent folder where the file should be created."
+    ] = None,
+    initial_content: Annotated[
+        str | None, "Text content to populate the new file, if applicable."
+    ] = None,
+) -> ToolResult:
+    """
+    Create a new file or folder in Google Drive.
+
+    This tool is essential for an AI agent to generate new output (like reports or
+    documentation) directly into the Drive structure.
+
+    Usage:
+        - Create empty file: gdrive_create_file(name="report.txt", mime_type="text/plain")
+        - Create Google Doc: gdrive_create_file(
+            name="My Report",
+            mime_type="application/vnd.google-apps.document",
+            initial_content="# Report Title"
+          )
+        - Create folder: gdrive_create_file(
+            name="Reports",
+            mime_type="application/vnd.google-apps.folder"
+          )
+        - Create in subfolder: gdrive_create_file(
+            name="file.txt",
+            mime_type="text/plain",
+            parent_id="folder_id_here",
+            initial_content="File content"
+          )
+
+    Supported MIME types:
+        - text/plain: Plain text file
+        - application/vnd.google-apps.document: Google Doc (content auto-converted)
+        - application/vnd.google-apps.spreadsheet: Google Sheet (CSV content works best)
+        - application/vnd.google-apps.folder: Folder (initial_content is ignored)
+
+    Note: For Google Workspace files, the Drive API automatically converts plain text
+    content to the appropriate format.
+    """
+    if not name or not name.strip():
+        raise ToolError("Argument validation error: 'name' cannot be empty.")
+
+    if not mime_type or not mime_type.strip():
+        raise ToolError("Argument validation error: 'mime_type' cannot be empty.")
+
+    access_token = await get_gdrive_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with GoogleDriveClient(access_token) as client:
+            created_file = await client.create_file(
+                name=name,
+                mime_type=mime_type,
+                parent_id=parent_id,
+                initial_content=initial_content,
+            )
+    except GoogleDriveError as e:
+        logger.error(f"Google Drive error creating file: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error creating Google Drive file: {e}")
+        raise ToolError(f"An unexpected error occurred while creating Google Drive file: {str(e)}")
+
+    file_type = "folder" if mime_type == GOOGLE_DRIVE_FOLDER_MIME else "file"
+    content_info = ""
+    if initial_content and mime_type != GOOGLE_DRIVE_FOLDER_MIME:
+        content_info = " with initial content"
+
+    return ToolResult(
+        content=f"Successfully created {file_type} '{created_file.name}'{content_info}.",
+        structured_content=created_file.as_flat_dict(),
+    )
+
+
+@dr_mcp_tool(
+    tags={"google", "gdrive", "update", "metadata", "rename", "star", "trash"}, enabled=False
+)
+async def gdrive_update_metadata(
+    *,
+    file_id: Annotated[str, "The ID of the file or folder to update."],
+    new_name: Annotated[str | None, "A new name to rename the file."] = None,
+    starred: Annotated[bool | None, "Set to True to star the file or False to unstar it."] = None,
+    trash: Annotated[bool | None, "Set to True to trash the file or False to restore it."] = None,
+) -> ToolResult:
+    """
+    Update non-content metadata fields of a Google Drive file or folder.
+
+    This tool allows you to:
+    - Rename files and folders by setting new_name
+    - Star or unstar files (per-user preference) with starred
+    - Move files to trash or restore them with trash
+
+    Usage:
+        - Rename: gdrive_update_metadata(file_id="1ABC...", new_name="New Name.txt")
+        - Star: gdrive_update_metadata(file_id="1ABC...", starred=True)
+        - Unstar: gdrive_update_metadata(file_id="1ABC...", starred=False)
+        - Trash: gdrive_update_metadata(file_id="1ABC...", trash=True)
+        - Restore: gdrive_update_metadata(file_id="1ABC...", trash=False)
+        - Multiple: gdrive_update_metadata(file_id="1ABC...", new_name="New.txt", starred=True)
+
+    Note:
+        - At least one of new_name, starred, or trash must be provided.
+        - Starring is per-user: starring a shared file only affects your view.
+        - Trashing a folder trashes all contents recursively.
+        - Trashing requires permissions (owner for My Drive, organizer for Shared Drives).
+    """
+    if not file_id or not file_id.strip():
+        raise ToolError("Argument validation error: 'file_id' cannot be empty.")
+
+    if new_name is None and starred is None and trash is None:
+        raise ToolError(
+            "Argument validation error: at least one of 'new_name', 'starred', or 'trash' "
+            "must be provided."
+        )
+
+    if new_name is not None and not new_name.strip():
+        raise ToolError("Argument validation error: 'new_name' cannot be empty or whitespace.")
+
+    access_token = await get_gdrive_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with GoogleDriveClient(access_token) as client:
+            updated_file = await client.update_file_metadata(
+                file_id=file_id,
+                new_name=new_name,
+                starred=starred,
+                trashed=trash,
+            )
+    except GoogleDriveError as e:
+        logger.error(f"Google Drive error updating file metadata: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error updating Google Drive file metadata: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while updating Google Drive file metadata: {str(e)}"
+        )
+
+    changes: list[str] = []
+    if new_name is not None:
+        changes.append(f"renamed to '{new_name}'")
+    if starred is True:
+        changes.append("starred")
+    elif starred is False:
+        changes.append("unstarred")
+    if trash is True:
+        changes.append("moved to trash")
+    elif trash is False:
+        changes.append("restored from trash")
+
+    changes_description = ", ".join(changes)
+
+    return ToolResult(
+        content=f"Successfully updated file '{updated_file.name}': {changes_description}.",
+        structured_content=updated_file.as_flat_dict(),
+    )
+
+
+@dr_mcp_tool(tags={"google", "gdrive", "manage", "access", "acl"})
+async def gdrive_manage_access(
+    *,
+    file_id: Annotated[str, "The ID of the file or folder."],
+    action: Annotated[Literal["add", "update", "remove"], "The operation to perform."],
+    role: Annotated[
+        Literal["reader", "commenter", "writer", "fileOrganizer", "organizer", "owner"] | None,
+        "The access level.",
+    ] = None,
+    email_address: Annotated[
+        str | None, "The email of the user or group (required for 'add')."
+    ] = None,
+    permission_id: Annotated[
+        str | None, "The specific permission ID (required for 'update' or 'remove')."
+    ] = None,
+    transfer_ownership: Annotated[
+        bool, "Whether to transfer ownership (only for 'update' to 'owner' role)."
+    ] = False,
+) -> ToolResult:
+    """
+    Consolidated tool for sharing files and managing permissions.
+    Pushes all logic to the Google Drive API permissions resource (create, update, delete).
+
+    Usage:
+        - Add role: gdrive_manage_access(
+            file_id="SomeFileId",
+            action="add",
+            role="reader",
+            email_address="dummy@user.com"
+          )
+        - Update role: gdrive_manage_access(
+            file_id="SomeFileId",
+            action="update",
+            role="reader",
+            permission_id="SomePermissionId"
+          )
+        - Remove permission: gdrive_manage_access(
+            file_id="SomeFileId",
+            action="remove",
+            permission_id="SomePermissionId"
+          )
+    """
+    if not file_id or not file_id.strip():
+        raise ToolError("Argument validation error: 'file_id' cannot be empty.")
+
+    if action == "add" and not email_address:
+        raise ToolError("'email_address' is required for action 'add'.")
+
+    if action in ("update", "remove") and not permission_id:
+        raise ToolError("'permission_id' is required for action 'update' or 'remove'.")
+
+    if action != "remove" and not role:
+        raise ToolError("'role' is required for action 'add' or 'update'.")
+
+    access_token = await get_gdrive_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with GoogleDriveClient(access_token) as client:
+            permission_id = await client.manage_access(
+                file_id=file_id,
+                action=action,
+                role=role,
+                email_address=email_address,
+                permission_id=permission_id,
+                transfer_ownership=transfer_ownership,
+            )
+    except GoogleDriveError as e:
+        logger.error(f"Google Drive permission operation failed: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error changing permissions for Google Drive file {file_id}: {e}")
+        raise ToolError(
+            f"Unexpected error changing permissions for Google Drive file {file_id}: {str(e)}"
+        )
+
+    # Build response
+    structured_content = {"affectedFileId": file_id}
+    if action == "add":
+        content = (
+            f"Successfully added role '{role}' for '{email_address}' for gdrive file '{file_id}'. "
+            f"New permission id '{permission_id}'."
+        )
+        structured_content["newPermissionId"] = permission_id
+    elif action == "update":
+        content = (
+            f"Successfully updated role '{role}' (permission '{permission_id}') "
+            f"for gdrive file '{file_id}'."
+        )
+    else:  # action == "remove":
+        content = f"Successfully removed permission '{permission_id}' for gdrive file '{file_id}'."
+
+    return ToolResult(content=content, structured_content=structured_content)

datarobot_genai/drmcp/tools/microsoft_graph/__init__.py

@@ -0,0 +1,13 @@
+# 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_genai/drmcp/tools/microsoft_graph/tools.py

@@ -0,0 +1,198 @@
+# 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.
+
+"""Microsoft Graph MCP tools for searching SharePoint and OneDrive content."""
+
+import logging
+from typing import Annotated
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
+
+from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import MicrosoftGraphClient
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import MicrosoftGraphError
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import get_microsoft_graph_access_token
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import validate_site_url
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(
+    tags={
+        "microsoft",
+        "graph api",
+        "sharepoint",
+        "drive",
+        "list",
+        "search",
+        "files",
+        "find",
+        "contents",
+    }
+)
+async def microsoft_graph_search_content(
+    *,
+    search_query: Annotated[str, "The search string to find files, folders, or list items."],
+    site_url: Annotated[
+        str | None,
+        "Optional SharePoint site URL to scope the search "
+        "(e.g., https://tenant.sharepoint.com/sites/sitename). "
+        "If not provided, searches across all accessible sites.",
+    ] = None,
+    site_id: Annotated[
+        str | None,
+        "Optional ID of the site to scope the search. If provided, takes precedence over site_url.",
+    ] = None,
+    from_offset: Annotated[
+        int,
+        "The zero-based index of the first result to return. Use this for pagination. "
+        "Default: 0 (start from the beginning). To get the next page, increment by the size "
+        "value (e.g., first page: from=0 size=250, second page: from=250 size=250, "
+        "third page: from=500 size=250).",
+    ] = 0,
+    size: Annotated[
+        int,
+        "Maximum number of results to return in this request. Default is 250, max is 250. "
+        "The LLM should control pagination by making multiple calls with different 'from' values.",
+    ] = 250,
+    entity_types: Annotated[
+        list[str] | None,
+        "Optional list of entity types to search. Valid values: 'driveItem', 'listItem', "
+        "'site', 'list', 'drive'. Default: ['driveItem', 'listItem']. "
+        "Multiple types can be specified.",
+    ] = None,
+    filters: Annotated[
+        list[str] | None,
+        "Optional list of KQL filter expressions to refine search results "
+        "(e.g., ['fileType:docx', 'size>1000']).",
+    ] = None,
+    include_hidden_content: Annotated[
+        bool,
+        "Whether to include hidden content in search results. Only works with delegated "
+        "permissions, not application permissions. Default: False.",
+    ] = False,
+    region: Annotated[
+        str | None,
+        "Optional region code for application permissions (e.g., 'NAM', 'EUR', 'APC'). "
+        "Required when using application permissions to search SharePoint content in "
+        "specific regions.",
+    ] = None,
+) -> ToolResult | ToolError:
+    """
+    Search for SharePoint and OneDrive content using Microsoft Graph Search API.
+
+    Search Scope:
+    - When site_url or site_id is provided: searches within the specified SharePoint site
+    - When neither is provided: searches across all accessible SharePoint sites and OneDrive
+
+    Supported Entity Types:
+    - driveItem: Files and folders in document libraries and OneDrive
+    - listItem: Items in SharePoint lists
+    - site: SharePoint sites
+    - list: SharePoint lists
+    - drive: Document libraries/drives
+
+    Filtering:
+    - Filters use KQL (Keyword Query Language) syntax
+    - Multiple filters are combined with AND operators
+    - Examples: ['fileType:docx', 'size>1000', 'lastModifiedTime>2024-01-01']
+    - Filters are applied in addition to the search query
+
+    Pagination:
+    - Controlled via from_offset (zero-based index) and size parameters
+    - Maximum size per request: 250 results
+    - To paginate: increment from_offset by size value for each subsequent page
+    - Example pagination sequence:
+      * Page 1: from_offset=0, size=250 (returns results 0-249)
+      * Page 2: from_offset=250, size=250 (returns results 250-499)
+      * Page 3: from_offset=500, size=250 (returns results 500-749)
+
+    API Reference:
+    - Endpoint: POST /search/query
+    - Documentation: https://learn.microsoft.com/en-us/graph/api/search-query
+    - Search concepts: https://learn.microsoft.com/en-us/graph/search-concept-files
+
+    Permissions:
+    - Requires Sites.Read.All or Sites.Search.All permission
+    - include_hidden_content only works with delegated permissions
+    - region parameter is required for application permissions in multi-region environments
+    """
+    if not search_query:
+        raise ToolError("Argument validation error: 'search_query' cannot be empty.")
+
+    # Validate site_url if provided
+    if site_url:
+        validation_error = validate_site_url(site_url)
+        if validation_error:
+            raise ToolError(validation_error)
+
+    access_token = await get_microsoft_graph_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with MicrosoftGraphClient(access_token=access_token, site_url=site_url) as client:
+            items = await client.search_content(
+                search_query=search_query,
+                site_id=site_id,
+                from_offset=from_offset,
+                size=size,
+                entity_types=entity_types,
+                filters=filters,
+                include_hidden_content=include_hidden_content,
+                region=region,
+            )
+    except MicrosoftGraphError as e:
+        logger.error(f"Microsoft Graph error searching content: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error searching Microsoft Graph content: {e}", exc_info=True)
+        raise ToolError(
+            f"An unexpected error occurred while searching Microsoft Graph content: {str(e)}"
+        )
+
+    results = []
+    for item in items:
+        result_dict = {
+            "id": item.id,  # Unique ID of the file, folder, or list item
+            "name": item.name,
+            "webUrl": item.web_url,
+            "size": item.size,
+            "createdDateTime": item.created_datetime,
+            "lastModifiedDateTime": item.last_modified_datetime,
+            "isFolder": item.is_folder,
+            "mimeType": item.mime_type,
+            # Document library/drive ID (driveId in Microsoft Graph API)
+            "documentLibraryId": item.drive_id,
+            "parentFolderId": item.parent_folder_id,  # Parent folder ID
+        }
+        results.append(result_dict)
+
+    n = len(results)
+    return ToolResult(
+        content=(
+            f"Successfully searched Microsoft Graph and retrieved {n} result(s) for "
+            f"'{search_query}' (from={from_offset}, size={size})."
+        ),
+        structured_content={
+            "query": search_query,
+            "siteUrl": site_url,
+            "siteId": site_id,
+            "from": from_offset,
+            "size": size,
+            "results": results,
+            "count": n,
+        },
+    )

datarobot_genai/drmcp/tools/predictive/data.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import logging
 import os
 from typing import Annotated
@@ -28,14 +29,15 @@ logger = logging.getLogger(__name__)
 
 @dr_mcp_tool(tags={"predictive", "data", "write", "upload", "catalog"})
 async def upload_dataset_to_ai_catalog(
+    *,
     file_path: Annotated[str, "The path to the dataset file to upload."] | None = None,
     file_url: Annotated[str, "The URL to the dataset file to upload."] | None = None,
 ) -> ToolError | ToolResult:
     """Upload a dataset to the DataRobot AI Catalog / Data Registry."""
     if not file_path and not file_url:
-        return ToolError("Either file_path or file_url must be provided.")
+        raise ToolError("Either file_path or file_url must be provided.")
     if file_path and file_url:
-        return ToolError("Please provide either file_path or file_url, not both.")
+        raise ToolError("Please provide either file_path or file_url, not both.")
 
     # Get client
     client = get_sdk_client()
@@ -45,17 +47,17 @@ async def upload_dataset_to_ai_catalog(
         # Does file exist?
         if not os.path.exists(file_path):
             logger.error("File not found: %s", file_path)
-            return ToolError(f"File not found: {file_path}")
+            raise ToolError(f"File not found: {file_path}")
         catalog_item = client.Dataset.create_from_file(file_path)
     else:
         # Does URL exist?
         if file_url is None or not is_valid_url(file_url):
             logger.error("Invalid file URL: %s", file_url)
-            return ToolError(f"Invalid file URL: {file_url}")
+            raise ToolError(f"Invalid file URL: {file_url}")
         catalog_item = client.Dataset.create_from_url(file_url)
 
     if not catalog_item:
-        return ToolError("Failed to upload dataset.")
+        raise ToolError("Failed to upload dataset.")
 
     return ToolResult(
         content=f"Successfully uploaded dataset: {catalog_item.id}",
@@ -80,11 +82,17 @@ async def list_ai_catalog_items() -> ToolResult:
             structured_content={"datasets": []},
         )
 
+    datasets_dict = {ds.id: ds.name for ds in datasets}
+    datasets_count = len(datasets)
+
     return ToolResult(
-        content=f"Found {len(datasets)} AI Catalog items.",
+        content=(
+            f"Found {datasets_count} AI Catalog items, here are the details:\n"
+            f"{json.dumps(datasets_dict, indent=2)}"
+        ),
         structured_content={
-            "datasets": [{"id": ds.id, "name": ds.name} for ds in datasets],
-            "count": len(datasets),
+            "datasets": datasets_dict,
+            "count": datasets_count,
         },
     )