datarobot-genai 0.2.29__tar.gz → 0.2.30__tar.gz

{datarobot_genai-0.2.29 → datarobot_genai-0.2.30}/PKG-INFO

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

{datarobot_genai-0.2.29 → datarobot_genai-0.2.30}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.29"
+version = "0.2.30"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"

{datarobot_genai-0.2.29 → datarobot_genai-0.2.30}/src/datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -33,7 +33,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 +129,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 +145,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 +475,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.
 

{datarobot_genai-0.2.29 → datarobot_genai-0.2.30}/src/datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -164,7 +164,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 +251,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 +258,90 @@ 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(),
     )