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

datarobot_genai/drmcp/tools/clients/microsoft_graph.py

@@ -16,6 +16,8 @@
 
 import logging
 from typing import Any
+from typing import Literal
+from urllib.parse import quote
 
 import httpx
 from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
@@ -415,6 +417,175 @@ 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 share_item(
+        self,
+        file_id: str,
+        document_library_id: str,
+        recipient_emails: list[str],
+        role: Literal["read", "write"],
+        send_invitation: bool,
+    ) -> None:
+        """
+        Share sharepoint / ondrive item using Microsoft Graph API.
+        Under the hood all resources in sharepoint/onedrive
+            in MS Graph API are treated as 'driveItem'.
+
+        Args:
+            file_id: The ID of the file or folder to share.
+            document_library_id: The ID of the document library containing the item.
+            recipient_emails: A list of email addresses to invite.
+            role: The role to assign.
+            send_invitation: Flag determining if recipients should be notified
+
+        Returns
+        -------
+            None
+
+        Raises
+        ------
+            MicrosoftGraphError: If sharing fails
+        """
+        graph_url = f"{GRAPH_API_BASE}/drives/{document_library_id}/items/{file_id}/invite"
+
+        payload = {
+            "recipients": [{"email": email} for email in recipient_emails],
+            "requireSignIn": True,
+            "sendInvitation": send_invitation,
+            "roles": [role],
+        }
+
+        response = await self._client.post(url=graph_url, json=payload)
+
+        if response.status_code not in (200, 201):
+            raise MicrosoftGraphError(
+                f"Microsoft Graph API error {response.status_code}: {response.text}"
+            )
+
     async def __aenter__(self) -> "MicrosoftGraphClient":
         """Async context manager entry."""
         return self

datarobot_genai/drmcp/tools/confluence/tools.py

@@ -55,26 +55,16 @@ async def confluence_get_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            if page_id_or_title.isdigit():
-                page_response = await client.get_page_by_id(page_id_or_title)
-            else:
-                if not space_key:
-                    raise ToolError(
-                        "Argument validation error: "
-                        "'space_key' is required when identifying a page by title."
-                    )
-                page_response = await client.get_page_by_title(page_id_or_title, space_key)
-    except ConfluenceError as e:
-        logger.error(f"Confluence error getting page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error getting Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while getting Confluence page "
-            f"'{page_id_or_title}': {str(e)}"
-        )
+    async with ConfluenceClient(access_token) as client:
+        if page_id_or_title.isdigit():
+            page_response = await client.get_page_by_id(page_id_or_title)
+        else:
+            if not space_key:
+                raise ToolError(
+                    "Argument validation error: "
+                    "'space_key' is required when identifying a page by title."
+                )
+            page_response = await client.get_page_by_title(page_id_or_title, space_key)
 
     return ToolResult(
         content=f"Successfully retrieved page '{page_response.title}'.",
@@ -116,22 +106,12 @@ async def confluence_create_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            page_response = await client.create_page(
-                space_key=space_key,
-                title=title,
-                body_content=body_content,
-                parent_id=parent_id,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error creating page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error creating Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while creating Confluence page "
-            f"'{title}' in space '{space_key}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        page_response = await client.create_page(
+            space_key=space_key,
+            title=title,
+            body_content=body_content,
+            parent_id=parent_id,
         )
 
     return ToolResult(
@@ -164,19 +144,10 @@ async def confluence_add_comment(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            comment_response = await client.add_comment(
-                page_id=page_id,
-                comment_body=comment_body,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error adding comment: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error adding comment to Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while adding comment to page '{page_id}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        comment_response = await client.add_comment(
+            page_id=page_id,
+            comment_body=comment_body,
         )
 
     return ToolResult(
@@ -220,32 +191,24 @@ async def confluence_search(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            results = await client.search_confluence_content(
-                cql_query=cql_query, max_results=max_results
-            )
-
-            # If include_body is True, fetch full content for each page
-            if include_body and results:
-                data = []
-                for result in results:
-                    flat = result.as_flat_dict()
-                    try:
-                        page = await client.get_page_by_id(result.id)
-                        flat["body"] = page.body
-                    except ConfluenceError:
-                        flat["body"] = None  # Keep excerpt if page fetch fails
-                    data.append(flat)
-            else:
-                data = [result.as_flat_dict() for result in results]
-
-    except ConfluenceError as e:
-        logger.error(f"Confluence error searching content: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error searching Confluence content: {e}")
-        raise ToolError(f"An unexpected error occurred while searching Confluence: {str(e)}")
+    async with ConfluenceClient(access_token) as client:
+        results = await client.search_confluence_content(
+            cql_query=cql_query, max_results=max_results
+        )
+
+        # If include_body is True, fetch full content for each page
+        if include_body and results:
+            data = []
+            for result in results:
+                flat = result.as_flat_dict()
+                try:
+                    page = await client.get_page_by_id(result.id)
+                    flat["body"] = page.body
+                except ConfluenceError:
+                    flat["body"] = None  # Keep excerpt if page fetch fails
+                data.append(flat)
+        else:
+            data = [result.as_flat_dict() for result in results]
 
     n = len(results)
     return ToolResult(
@@ -296,20 +259,11 @@ async def confluence_update_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            page_response = await client.update_page(
-                page_id=page_id,
-                new_body_content=new_body_content,
-                version_number=version_number,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error updating page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error updating Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while updating Confluence page '{page_id}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        page_response = await client.update_page(
+            page_id=page_id,
+            new_body_content=new_body_content,
+            version_number=version_number,
         )
 
     return ToolResult(

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -28,7 +28,6 @@ from datarobot_genai.drmcp.tools.clients.gdrive import MAX_PAGE_SIZE
 from datarobot_genai.drmcp.tools.clients.gdrive import SUPPORTED_FIELDS
 from datarobot_genai.drmcp.tools.clients.gdrive import SUPPORTED_FIELDS_STR
 from datarobot_genai.drmcp.tools.clients.gdrive import GoogleDriveClient
-from datarobot_genai.drmcp.tools.clients.gdrive import GoogleDriveError
 from datarobot_genai.drmcp.tools.clients.gdrive import get_gdrive_access_token
 
 logger = logging.getLogger(__name__)
@@ -79,22 +78,15 @@ async def gdrive_find_contents(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with GoogleDriveClient(access_token) as client:
-            data = await client.list_files(
-                page_size=page_size,
-                page_token=page_token,
-                query=query,
-                limit=limit,
-                folder_id=folder_id,
-                recursive=recursive,
-            )
-    except GoogleDriveError as e:
-        logger.error(f"Google Drive error listing files: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error listing Google Drive files: {e}")
-        raise ToolError(f"An unexpected error occurred while listing Google Drive files: {str(e)}")
+    async with GoogleDriveClient(access_token) as client:
+        data = await client.list_files(
+            page_size=page_size,
+            page_token=page_token,
+            query=query,
+            limit=limit,
+            folder_id=folder_id,
+            recursive=recursive,
+        )
 
     filtered_fields = set(fields).intersection(SUPPORTED_FIELDS) if fields else SUPPORTED_FIELDS
     number_of_files = len(data.files)
@@ -155,17 +147,8 @@ async def gdrive_read_content(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with GoogleDriveClient(access_token) as client:
-            file_content = await client.read_file_content(file_id, target_format)
-    except GoogleDriveError as e:
-        logger.error(f"Google Drive error reading file content: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error reading Google Drive file content: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while reading Google Drive file content: {str(e)}"
-        )
+    async with GoogleDriveClient(access_token) as client:
+        file_content = await client.read_file_content(file_id, target_format)
 
     export_info = ""
     if file_content.was_exported:
@@ -239,20 +222,13 @@ async def gdrive_create_file(
     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)}")
+    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,
+        )
 
     file_type = "folder" if mime_type == GOOGLE_DRIVE_FOLDER_MIME else "file"
     content_info = ""
@@ -313,21 +289,12 @@ async def gdrive_update_metadata(
     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)}"
+    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,
         )
 
     changes: list[str] = []
@@ -408,23 +375,14 @@ async def gdrive_manage_access(
     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)}"
+    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,
         )
 
     # Build response

datarobot_genai/drmcp/tools/jira/tools.py

@@ -72,14 +72,8 @@ async def jira_get_issue(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with JiraClient(access_token) as client:
-            issue = await client.get_jira_issue(issue_key)
-    except Exception as e:
-        logger.error(f"Unexpected error while getting Jira issue: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while getting Jira issue '{issue_key}': {str(e)}"
-        )
+    async with JiraClient(access_token) as client:
+        issue = await client.get_jira_issue(issue_key)
 
     return ToolResult(
         content=f"Successfully retrieved details for issue '{issue_key}'.",
@@ -118,17 +112,13 @@ async def jira_create_issue(
                 f"Unexpected issue type `{issue_type}`. Possible values are {possible_issue_types}."
             )
 
-    try:
-        async with JiraClient(access_token) as client:
-            issue_key = await client.create_jira_issue(
-                project_key=project_key,
-                summary=summary,
-                issue_type_id=issue_type_id,
-                description=description,
-            )
-    except Exception as e:
-        logger.error(f"Unexpected error while creating Jira issue: {e}")
-        raise ToolError(f"An unexpected error occurred while creating Jira issue: {str(e)}")
+    async with JiraClient(access_token) as client:
+        issue_key = await client.create_jira_issue(
+            project_key=project_key,
+            summary=summary,
+            issue_type_id=issue_type_id,
+            description=description,
+        )
 
     return ToolResult(
         content=f"Successfully created issue '{issue_key}'.",
@@ -179,14 +169,10 @@ async def jira_update_issue(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with JiraClient(access_token) as client:
-            updated_fields = await client.update_jira_issue(
-                issue_key=issue_key, fields=fields_to_update
-            )
-    except Exception as e:
-        logger.error(f"Unexpected error while updating Jira issue: {e}")
-        raise ToolError(f"An unexpected error occurred while updating Jira issue: {str(e)}")
+    async with JiraClient(access_token) as client:
+        updated_fields = await client.update_jira_issue(
+            issue_key=issue_key, fields=fields_to_update
+        )
 
     updated_fields_str = ",".join(updated_fields)
     return ToolResult(
@@ -226,12 +212,8 @@ async def jira_transition_issue(
                 f"Possible values are {available_transitions_str}."
             )
 
-    try:
-        async with JiraClient(access_token) as client:
-            await client.transition_jira_issue(issue_key=issue_key, transition_id=transition_id)
-    except Exception as e:
-        logger.error(f"Unexpected error while transitioning Jira issue: {e}")
-        raise ToolError(f"An unexpected error occurred while transitioning Jira issue: {str(e)}")
+    async with JiraClient(access_token) as client:
+        await client.transition_jira_issue(issue_key=issue_key, transition_id=transition_id)
 
     return ToolResult(
         content=f"Successfully transitioned issue '{issue_key}' to status '{transition_name}'.",