datarobot-genai 0.2.26__tar.gz → 0.2.27__tar.gz

{datarobot_genai-0.2.26 → datarobot_genai-0.2.27}/PKG-INFO

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

{datarobot_genai-0.2.26 → datarobot_genai-0.2.27}/pyproject.toml

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

{datarobot_genai-0.2.26 → datarobot_genai-0.2.27}/src/datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -15,7 +15,9 @@
 """Google Drive API Client and utilities for OAuth."""
 
 import io
+import json
 import logging
+import uuid
 from typing import Annotated
 from typing import Any
 
@@ -45,6 +47,12 @@ GOOGLE_WORKSPACE_EXPORT_MIMES: dict[str, str] = {
     "application/vnd.google-apps.presentation": "text/plain",
 }
 
+# MIME type mappings for content conversion during upload to Google Workspace formats
+UPLOAD_CONTENT_TYPES: dict[str, str] = {
+    "application/vnd.google-apps.document": "text/plain",
+    "application/vnd.google-apps.spreadsheet": "text/csv",
+}
+
 BINARY_MIME_PREFIXES = (
     "image/",
     "audio/",
@@ -599,6 +607,125 @@ class GoogleDriveClient:
             web_view_link=file_metadata.web_view_link,
         )
 
+    async def create_file(
+        self,
+        name: str,
+        mime_type: str,
+        parent_id: str | None = None,
+        initial_content: str | None = None,
+    ) -> GoogleDriveFile:
+        """Create a new file or folder in Google Drive.
+
+        Creates a new file with the specified name and MIME type. Optionally places
+        it in a specific folder and populates it with initial content.
+
+        For Google Workspace files (Docs, Sheets), the Drive API automatically
+        converts plain text content to the appropriate format.
+
+        Args:
+            name: The name for the new file or folder.
+            mime_type: The MIME type of the file (e.g., 'text/plain',
+                'application/vnd.google-apps.document',
+                'application/vnd.google-apps.folder').
+            parent_id: Optional ID of the parent folder. If not specified,
+                the file is created in the root of the user's Drive.
+            initial_content: Optional text content to populate the file.
+                Ignored for folders.
+
+        Returns
+        -------
+            GoogleDriveFile with the created file's metadata.
+
+        Raises
+        ------
+            GoogleDriveError: If file creation fails (permission denied,
+                parent not found, rate limited, etc.).
+        """
+        metadata: dict[str, Any] = {
+            "name": name,
+            "mimeType": mime_type,
+        }
+        if parent_id:
+            metadata["parents"] = [parent_id]
+
+        if mime_type == GOOGLE_DRIVE_FOLDER_MIME or not initial_content:
+            response = await self._client.post(
+                "/",
+                json=metadata,
+                params={"fields": SUPPORTED_FIELDS_STR, "supportsAllDrives": "true"},
+            )
+        else:
+            response = await self._create_file_with_content(
+                metadata=metadata,
+                content=initial_content,
+                target_mime_type=mime_type,
+            )
+
+        if response.status_code == 404:
+            raise GoogleDriveError(
+                f"Parent folder with ID '{parent_id}' not found."
+                if parent_id
+                else "Resource not found."
+            )
+        if response.status_code == 403:
+            raise GoogleDriveError(
+                "Permission denied: you don't have permission to create files in this location."
+            )
+        if response.status_code == 400:
+            raise GoogleDriveError(
+                f"Bad request: invalid parameters for file creation. "
+                f"Check that the MIME type '{mime_type}' is valid."
+            )
+        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 _create_file_with_content(
+        self,
+        metadata: dict[str, Any],
+        content: str,
+        target_mime_type: str,
+    ) -> httpx.Response:
+        """Create a file with content using multipart upload.
+
+        Args:
+            metadata: File metadata dictionary.
+            content: Text content for the file.
+            target_mime_type: The target MIME type for the file.
+
+        Returns
+        -------
+            The HTTP response from the upload.
+        """
+        content_type = UPLOAD_CONTENT_TYPES.get(target_mime_type, "text/plain")
+        boundary = f"===gdrive_boundary_{uuid.uuid4().hex}==="
+        body_parts = [
+            f"--{boundary}",
+            "Content-Type: application/json; charset=UTF-8",
+            "",
+            json.dumps(metadata),
+            f"--{boundary}",
+            f"Content-Type: {content_type}",
+            "",
+            content,
+            f"--{boundary}--",
+        ]
+        body = "\r\n".join(body_parts)
+
+        upload_url = "https://www.googleapis.com/upload/drive/v3/files"
+        return await self._client.post(
+            upload_url,
+            content=body.encode("utf-8"),
+            params={
+                "uploadType": "multipart",
+                "fields": SUPPORTED_FIELDS_STR,
+                "supportsAllDrives": "true",
+            },
+            headers={"Content-Type": f"multipart/related; boundary={boundary}"},
+        )
+
     async def __aenter__(self) -> "GoogleDriveClient":
         """Async context manager entry."""
         return self

{datarobot_genai-0.2.26 → datarobot_genai-0.2.27}/src/datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -21,6 +21,7 @@ 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
@@ -60,7 +61,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 +122,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).
@@ -175,3 +176,95 @@ 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)}")
+
+    # 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:
+        content_info = " with initial content"
+
+    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,
+        },
+    )