datarobot-genai 0.2.25__py3-none-any.whl → 0.2.27__py3-none-any.whl

datarobot_genai/drmcp/core/mcp_instance.py

@@ -16,6 +16,7 @@ import logging
 from collections.abc import Callable
 from functools import wraps
 from typing import Any
+from typing import TypedDict
 
 from fastmcp import Context
 from fastmcp import FastMCP
@@ -26,6 +27,7 @@ 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
 
 from .config import MCPServerConfig
 from .config import get_config
@@ -287,16 +289,37 @@ mcp = TaggedFastMCP(
 )
 
 
+class ToolKwargs(TypedDict, total=False):
+    """Keyword arguments passed through to FastMCP's mcp.tool() decorator.
+
+    All parameters are optional and forwarded directly to FastMCP tool registration.
+    See FastMCP documentation for full details on each parameter.
+    """
+
+    name: str | None
+    title: str | None
+    description: str | None
+    icons: list[Any] | None
+    tags: set[str] | None
+    output_schema: dict[str, Any] | None
+    annotations: Any | None
+    exclude_args: list[str] | None
+    meta: dict[str, Any] | None
+    enabled: bool | None
+
+
 def dr_core_mcp_tool(
-    name: str | None = None,
-    description: str | None = None,
-    tags: set[str] | None = None,
+    **kwargs: Unpack[ToolKwargs],
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
-    """Combine decorator that includes mcp.tool() and dr_mcp_extras()."""
+    """Combine decorator that includes mcp.tool() and dr_mcp_extras().
+
+    All keyword arguments are passed through to FastMCP's mcp.tool() decorator.
+    See ToolKwargs for available parameters.
+    """
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         instrumented = dr_mcp_extras()(func)
-        mcp.tool(name=name, description=description, tags=tags)(instrumented)
+        mcp.tool(**kwargs)(instrumented)
         return instrumented
 
     return decorator
@@ -329,27 +352,23 @@ async def memory_aware_wrapper(func: Callable[..., Any], *args: Any, **kwargs: A
 
 
 def dr_mcp_tool(
-    name: str | None = None,
-    description: str | None = None,
-    tags: set[str] | None = None,
+    **kwargs: Unpack[ToolKwargs],
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     """Combine decorator that includes mcp.tool(), dr_mcp_extras(), and capture memory ids from
     the request headers if they exist.
 
-    Args:
-        name: Tool name
-        description: Tool description
-        tags: Optional set of tags to apply to the tool
+    All keyword arguments are passed through to FastMCP's mcp.tool() decorator.
+    See ToolKwargs for available parameters.
     """
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         @wraps(func)
-        async def wrapper(*args: Any, **kwargs: Any) -> Any:
-            return await memory_aware_wrapper(func, *args, **kwargs)
+        async def wrapper(*args: Any, **inner_kwargs: Any) -> Any:
+            return await memory_aware_wrapper(func, *args, **inner_kwargs)
 
         # Apply the MCP decorators
         instrumented = dr_mcp_extras()(wrapper)
-        mcp.tool(name=name, description=description, tags=tags)(instrumented)
+        mcp.tool(**kwargs)(instrumented)
         return instrumented
 
     return decorator

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/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,
+        },
+    )

{datarobot_genai-0.2.25.dist-info → datarobot_genai-0.2.27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.25
+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.25.dist-info → datarobot_genai-0.2.27.dist-info}/RECORD

@@ -35,7 +35,7 @@ datarobot_genai/drmcp/core/dr_mcp_server.py,sha256=czcjbwhZAeW9EtG_Bys0GARPOuQul
 datarobot_genai/drmcp/core/dr_mcp_server_logo.py,sha256=hib-nfR1SNTW6CnpFsFCkL9H_OMwa4YYyinV7VNOuLk,4708
 datarobot_genai/drmcp/core/exceptions.py,sha256=eqsGI-lxybgvWL5w4BFhbm3XzH1eU5tetwjnhJxelpc,905
 datarobot_genai/drmcp/core/logging.py,sha256=Y_hig4eBWiXGaVV7B_3wBcaYVRNH4ydptbEQhrP9-mY,3414
-datarobot_genai/drmcp/core/mcp_instance.py,sha256=cbhy9vSw7LMZWPAYkpCyT0V55wfxyRLkzOlgwKh1T2o,17794
+datarobot_genai/drmcp/core/mcp_instance.py,sha256=nt4gOlAQklMcqmohRIKovYcyhgLdb08NHMo28DBYmOk,18362
 datarobot_genai/drmcp/core/routes.py,sha256=dqE2M0UzAyyN9vQjlyTjYW4rpju3LT039po5weuO__I,17936
 datarobot_genai/drmcp/core/routes_utils.py,sha256=vSseXWlplMSnRgoJgtP_rHxWSAVYcx_tpTv4lyTpQoc,944
 datarobot_genai/drmcp/core/server_life_cycle.py,sha256=WKGJWGxalvqxupzJ2y67Kklc_9PgpZT0uyjlv_sr5wc,3419
@@ -78,13 +78,13 @@ datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosH
 datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCEYD6OmFx_LuywHCxgXsA,6472
 datarobot_genai/drmcp/tools/clients/confluence.py,sha256=h_G0By_kDnJeWDT_d-IREsaZ5-0xB5GoLXOqblYP5MA,20706
-datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=e28XwX0C8E3nql85-_NbUEMB-4s0lsQ2f5spj9BgsgM,21455
+datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=8GztWTdpJ7Ir3NIvIoOHPzDscoR1Ui7Ct2IiKmuUzIc,26012
 datarobot_genai/drmcp/tools/clients/jira.py,sha256=Rm91JAyrNIqxu66-9rU1YqoRXVnWbEy-Ahvy6f6HlVg,9823
 datarobot_genai/drmcp/tools/clients/s3.py,sha256=GmwzvurFdNfvxOooA8g5S4osRysHYU0S9ypg_177Glg,953
 datarobot_genai/drmcp/tools/confluence/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/confluence/tools.py,sha256=_-ws65WLK8KZP_mKkf4yJ7ZunR8qdyoiMwHQX47MSMw,12362
 datarobot_genai/drmcp/tools/gdrive/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=EvoEr3AEI-xRldwCTZHiQfBRHQfLtgHuojTx8mXhlU4,7074
+datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=G8LlnGEINZqV83Q-b3ZliWkDjouhbozDam3w6GfA7s0,10711
 datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/jira/tools.py,sha256=dfkqTU2HH-7n44hX80ODFacKq0p0LOchFcZtIIKFNMM,9687
 datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
@@ -110,9 +110,9 @@ datarobot_genai/nat/datarobot_llm_clients.py,sha256=Yu208Ed_p_4P3HdpuM7fYnKcXtim
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
 datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
 datarobot_genai/nat/helpers.py,sha256=Q7E3ADZdtFfS8E6OQPyw2wgA6laQ58N3bhLj5CBWwJs,3265
-datarobot_genai-0.2.25.dist-info/METADATA,sha256=C2_0ev_hz7YV8hiWuMVjoS5O9fz62WZ_oN3DuYDj6rA,6301
-datarobot_genai-0.2.25.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.25.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.25.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.25.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.25.dist-info/RECORD,,
+datarobot_genai-0.2.27.dist-info/METADATA,sha256=XfhoUrhYzD4u2Glt9HIR0AUEzlj170ct0ZkcSujc_W4,6301
+datarobot_genai-0.2.27.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.27.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.27.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.27.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.27.dist-info/RECORD,,