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

datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -20,6 +20,7 @@ import logging
 import uuid
 from typing import Annotated
 from typing import Any
+from typing import Literal
 
 import httpx
 from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
@@ -33,7 +34,17 @@ from datarobot_genai.drmcp.core.auth import get_access_token
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_FIELDS = {"id", "name", "size", "mimeType", "webViewLink", "createdTime", "modifiedTime"}
+SUPPORTED_FIELDS = {
+    "id",
+    "name",
+    "size",
+    "mimeType",
+    "webViewLink",
+    "createdTime",
+    "modifiedTime",
+    "starred",
+    "trashed",
+}
 SUPPORTED_FIELDS_STR = ",".join(SUPPORTED_FIELDS)
 DEFAULT_FIELDS = f"nextPageToken,files({SUPPORTED_FIELDS_STR})"
 GOOGLE_DRIVE_FOLDER_MIME = "application/vnd.google-apps.folder"
@@ -119,6 +130,8 @@ class GoogleDriveFile(BaseModel):
     web_view_link: Annotated[str | None, Field(alias="webViewLink")] = None
     created_time: Annotated[str | None, Field(alias="createdTime")] = None
     modified_time: Annotated[str | None, Field(alias="modifiedTime")] = None
+    starred: bool | None = None
+    trashed: bool | None = None
 
     model_config = ConfigDict(populate_by_name=True)
 
@@ -133,8 +146,31 @@ class GoogleDriveFile(BaseModel):
             web_view_link=data.get("webViewLink"),
             created_time=data.get("createdTime"),
             modified_time=data.get("modifiedTime"),
+            starred=data.get("starred"),
+            trashed=data.get("trashed"),
         )
 
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the file."""
+        result: dict[str, Any] = {
+            "id": self.id,
+            "name": self.name,
+            "mimeType": self.mime_type,
+        }
+        if self.size is not None:
+            result["size"] = self.size
+        if self.web_view_link is not None:
+            result["webViewLink"] = self.web_view_link
+        if self.created_time is not None:
+            result["createdTime"] = self.created_time
+        if self.modified_time is not None:
+            result["modifiedTime"] = self.modified_time
+        if self.starred is not None:
+            result["starred"] = self.starred
+        if self.trashed is not None:
+            result["trashed"] = self.trashed
+        return result
+
 
 class PaginatedResult(BaseModel):
     """Result of a paginated API call."""
@@ -440,6 +476,66 @@ class GoogleDriveClient:
         response.raise_for_status()
         return GoogleDriveFile.from_api_response(response.json())
 
+    async def update_file_metadata(
+        self,
+        file_id: str,
+        new_name: str | None = None,
+        starred: bool | None = None,
+        trashed: bool | None = None,
+    ) -> GoogleDriveFile:
+        """Update file metadata in Google Drive.
+
+        Args:
+            file_id: The ID of the file to update.
+            new_name: A new name to rename the file. Must not be empty or whitespace.
+            starred: Set to True to star the file or False to unstar it.
+            trashed: Set to True to trash the file or False to restore it.
+
+        Returns
+        -------
+            GoogleDriveFile with updated metadata.
+
+        Raises
+        ------
+            GoogleDriveError: If no update fields are provided, file is not found,
+                             access is denied, or the request is invalid.
+        """
+        if new_name is None and starred is None and trashed is None:
+            raise GoogleDriveError(
+                "At least one of new_name, starred, or trashed must be provided."
+            )
+
+        if new_name is not None and not new_name.strip():
+            raise GoogleDriveError("new_name cannot be empty or whitespace.")
+
+        body: dict[str, Any] = {}
+        if new_name is not None:
+            body["name"] = new_name
+        if starred is not None:
+            body["starred"] = starred
+        if trashed is not None:
+            body["trashed"] = trashed
+
+        response = await self._client.patch(
+            f"/{file_id}",
+            json=body,
+            params={"fields": SUPPORTED_FIELDS_STR, "supportsAllDrives": "true"},
+        )
+
+        if response.status_code == 404:
+            raise GoogleDriveError(f"File with ID '{file_id}' not found.")
+        if response.status_code == 403:
+            raise GoogleDriveError(
+                f"Permission denied: you don't have permission to update file '{file_id}'."
+            )
+        if response.status_code == 400:
+            raise GoogleDriveError("Bad request: invalid parameters for file update.")
+        if response.status_code == 429:
+            raise GoogleDriveError("Rate limit exceeded. Please try again later.")
+
+        response.raise_for_status()
+        return GoogleDriveFile.from_api_response(response.json())
+
     async def _export_workspace_file(self, file_id: str, export_mime_type: str) -> str:
         """Export a Google Workspace file to the specified format.
 
@@ -726,6 +822,96 @@ class GoogleDriveClient:
             headers={"Content-Type": f"multipart/related; boundary={boundary}"},
         )
 
+    async def manage_access(
+        self,
+        *,
+        file_id: str,
+        action: Literal["add", "update", "remove"],
+        role: Literal["reader", "commenter", "writer", "fileOrganizer", "organizer", "owner"]
+        | None = None,
+        email_address: str | None = None,
+        permission_id: str | None = None,
+        transfer_ownership: bool = False,
+    ) -> str:
+        """Manage access permissions for a Google Drive file or folder.
+
+        Adds, updates, or removes sharing permissions on an existing Google Drive
+        file or folder using the Google Drive Permissions API.
+
+        This method supports granting access to users or groups, changing access
+        roles, and revoking permissions. Ownership transfer is supported for files
+        in "My Drive" when explicitly requested.
+
+        Args:
+            file_id: The ID of the Google Drive file or folder whose permissions
+                are being managed.
+            action: The permission operation to perform.
+            role: The access role to assign or update. Valid values include
+                Required for "add" and "update" actions.
+            email_address: The email address of the user or group to grant access to.
+                Required for the "add" action.
+            permission_id: The ID of the permission to update or remove.
+                Required for "update" and "remove" actions.
+            transfer_ownership: Whether to transfer ownership of the file.
+                Only applicable when action="update" and role="owner".
+
+        Returns
+        -------
+            Permission id.
+            For "add" its newly added permission.
+            For "update"/"remove" its previous permission.
+
+        Raises
+        ------
+            GoogleDriveError: If the permission operation fails (invalid arguments,
+                insufficient permissions, resource not found, ownership transfer
+                not allowed, rate limited, etc.).
+        """
+        if not file_id.strip():
+            raise GoogleDriveError("Argument validation error: 'file_id' cannot be empty.")
+
+        if action == "add" and not email_address:
+            raise GoogleDriveError("'email_address' is required for action 'add'.")
+
+        if action in ("update", "remove") and not permission_id:
+            raise GoogleDriveError("'permission_id' is required for action 'update' or 'remove'.")
+
+        if action != "remove" and not role:
+            raise GoogleDriveError("'role' is required for action 'add' or 'update'.")
+
+        if action == "add":
+            response = await self._client.post(
+                url=f"/{file_id}/permissions",
+                json={
+                    "type": "user",
+                    "role": role,
+                    "emailAddress": email_address,
+                },
+                params={"sendNotificationEmail": False, "supportsAllDrives": True},
+            )
+
+        elif action == "update":
+            response = await self._client.patch(
+                url=f"/{file_id}/permissions/{permission_id}",
+                json={"role": role},
+                params={"transferOwnership": transfer_ownership, "supportsAllDrives": True},
+            )
+
+        elif action == "remove":
+            response = await self._client.delete(url=f"/{file_id}/permissions/{permission_id}")
+
+        else:
+            raise GoogleDriveError(f"Invalid action '{action}'")
+
+        if response.status_code not in (200, 201, 204):
+            raise GoogleDriveError(f"Drive API error {response.status_code}: {response.text}")
+
+        if action == "add":
+            return response.json()["id"]
+
+        # Cannot be null here because of above validators
+        return permission_id  # type: ignore
+
     async def __aenter__(self) -> "GoogleDriveClient":
         """Async context manager entry."""
         return self

datarobot_genai/drmcp/tools/clients/microsoft_graph.py

@@ -16,6 +16,7 @@
 
 import logging
 from typing import Any
+from urllib.parse import quote
 
 import httpx
 from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
@@ -415,6 +416,131 @@ class MicrosoftGraphClient:
 
         return base_resource
 
+    async def get_personal_drive_id(self) -> str:
+        """Get the current user's personal OneDrive drive ID.
+
+        Returns
+        -------
+            The drive ID string for the user's personal OneDrive.
+
+        Raises
+        ------
+            MicrosoftGraphError: If the drive cannot be retrieved.
+        """
+        try:
+            response = await self._client.get(f"{GRAPH_API_BASE}/me/drive")
+            response.raise_for_status()
+            data = response.json()
+            return data["id"]
+        except httpx.HTTPStatusError as e:
+            status_code = e.response.status_code
+            if status_code == 401:
+                raise MicrosoftGraphError(
+                    "Authentication failed. Access token may be expired or invalid."
+                ) from e
+            if status_code == 403:
+                raise MicrosoftGraphError(
+                    "Permission denied: cannot access personal OneDrive. "
+                    "Requires Files.Read or Files.ReadWrite permission."
+                ) from e
+            raise MicrosoftGraphError(f"Failed to get personal OneDrive: HTTP {status_code}") from e
+
+    async def create_file(
+        self,
+        drive_id: str,
+        file_name: str,
+        content: str,
+        parent_folder_id: str = "root",
+        conflict_behavior: str = "rename",
+    ) -> MicrosoftGraphItem:
+        """Create a text file in a drive (SharePoint document library or OneDrive).
+
+        Uses Microsoft Graph's simple upload endpoint for files < 4MB.
+        Files are created as text/plain content.
+
+        Args:
+            drive_id: The ID of the drive (document library) where the file will be created.
+            file_name: The name of the file to create (e.g., 'report.txt').
+            content: The text content to store in the file.
+            parent_folder_id: ID of the parent folder. Defaults to "root" (drive root folder).
+            conflict_behavior: How to handle name conflicts. Options:
+                - "rename" (default): Auto-renames to 'filename (1).txt', etc.
+                - "fail": Returns 409 Conflict error
+                - "replace": Overwrites existing file
+
+        Returns
+        -------
+            MicrosoftGraphItem representing the created file.
+
+        Raises
+        ------
+            MicrosoftGraphError: If file creation fails.
+        """
+        if not drive_id or not drive_id.strip():
+            raise MicrosoftGraphError("drive_id cannot be empty")
+        if not file_name or not file_name.strip():
+            raise MicrosoftGraphError("file_name cannot be empty")
+
+        # URL encode the filename for path-based addressing
+        encoded_name = quote(file_name, safe="")
+
+        # Simple upload endpoint for files < 4MB
+        # Reference: https://learn.microsoft.com/en-us/graph/api/driveitem-put-content
+        upload_url = (
+            f"{GRAPH_API_BASE}/drives/{drive_id}/items/{parent_folder_id}:/{encoded_name}:/content"
+        )
+
+        try:
+            response = await self._client.put(
+                upload_url,
+                content=content.encode("utf-8"),
+                headers={"Content-Type": "text/plain"},
+                params={"@microsoft.graph.conflictBehavior": conflict_behavior},
+            )
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            raise self._handle_create_file_error(e, drive_id, file_name, parent_folder_id) from e
+
+        return MicrosoftGraphItem.from_api_response(response.json())
+
+    def _handle_create_file_error(
+        self,
+        error: httpx.HTTPStatusError,
+        drive_id: str,
+        file_name: str,
+        parent_folder_id: str,
+    ) -> MicrosoftGraphError:
+        """Handle HTTP errors for file creation and return appropriate MicrosoftGraphError."""
+        status_code = error.response.status_code
+        error_msg = f"Failed to create file: 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 creating file: {api_message}"
+            except Exception:
+                error_msg = "Bad request: invalid parameters for file creation."
+        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 create files in drive "
+                f"'{drive_id}'. Requires Files.ReadWrite.All permission."
+            )
+        elif status_code == 404:
+            error_msg = (
+                f"Parent folder '{parent_folder_id}' not found in drive '{drive_id}'."
+                if parent_folder_id != "root"
+                else f"Drive '{drive_id}' not found."
+            )
+        elif status_code == 409:
+            error_msg = f"File '{file_name}' already exists and conflict behavior is set to 'fail'."
+        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

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -16,6 +16,7 @@
 
 import logging
 from typing import Annotated
+from typing import Literal
 
 from fastmcp.exceptions import ToolError
 from fastmcp.tools.tool import ToolResult
@@ -33,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[
@@ -164,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})"
@@ -252,7 +254,6 @@ async def gdrive_create_file(
         logger.error(f"Unexpected error creating Google Drive file: {e}")
         raise ToolError(f"An unexpected error occurred while creating Google Drive file: {str(e)}")
 
-    # Build response message
     file_type = "folder" if mime_type == GOOGLE_DRIVE_FOLDER_MIME else "file"
     content_info = ""
     if initial_content and mime_type != GOOGLE_DRIVE_FOLDER_MIME:
@@ -260,11 +261,186 @@ async def gdrive_create_file(
 
     return ToolResult(
         content=f"Successfully created {file_type} '{created_file.name}'{content_info}.",
-        structured_content={
-            "id": created_file.id,
-            "name": created_file.name,
-            "mimeType": created_file.mime_type,
-            "webViewLink": created_file.web_view_link,
-            "createdTime": created_file.created_time,
-        },
+        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/tools.py

@@ -196,3 +196,82 @@ async def microsoft_graph_search_content(
             "count": n,
         },
     )
+
+
+@dr_mcp_tool(
+    tags={
+        "microsoft",
+        "graph api",
+        "sharepoint",
+        "onedrive",
+        "document library",
+        "create",
+        "file",
+        "write",
+    }
+)
+async def microsoft_create_file(
+    *,
+    file_name: Annotated[str, "The name of the file to create (e.g., 'report.txt')."],
+    content_text: Annotated[str, "The raw text content to be stored in the file."],
+    document_library_id: Annotated[
+        str | None,
+        "The ID of the document library (Drive). If not provided, saves to personal OneDrive.",
+    ] = None,
+    parent_folder_id: Annotated[
+        str | None,
+        "ID of the parent folder. Defaults to 'root' (root of the drive).",
+    ] = "root",
+) -> ToolResult | ToolError:
+    """
+    Create a new text file in SharePoint or OneDrive.
+
+    **Personal OneDrive:** Just provide file_name and content_text.
+    The file saves to your personal OneDrive root folder.
+
+    **SharePoint:** Provide document_library_id to save to a specific
+    SharePoint site. Get the ID from microsoft_graph_search_content
+    results ('documentLibraryId' field).
+
+    **Conflict Resolution:** If a file with the same name exists,
+    it will be automatically renamed (e.g., 'report (1).txt').
+    """
+    if not file_name or not file_name.strip():
+        raise ToolError("Error: file_name is required.")
+    if not content_text:
+        raise ToolError("Error: content_text is required.")
+
+    access_token = await get_microsoft_graph_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    folder_id = parent_folder_id if parent_folder_id else "root"
+
+    async with MicrosoftGraphClient(access_token=access_token) as client:
+        # Auto-fetch personal OneDrive if no library specified
+        if document_library_id is None:
+            drive_id = await client.get_personal_drive_id()
+            is_personal_onedrive = True
+        else:
+            drive_id = document_library_id
+            is_personal_onedrive = False
+
+        created_file = await client.create_file(
+            drive_id=drive_id,
+            file_name=file_name.strip(),
+            content=content_text,
+            parent_folder_id=folder_id,
+            conflict_behavior="rename",
+        )
+
+    return ToolResult(
+        content=f"File '{created_file.name}' created successfully.",
+        structured_content={
+            "file_name": created_file.name,
+            "destination": "onedrive" if is_personal_onedrive else "sharepoint",
+            "driveId": drive_id,
+            "id": created_file.id,
+            "webUrl": created_file.web_url,
+            "parentFolderId": created_file.parent_folder_id,
+        },
+    )