datarobot-genai 0.2.13__py3-none-any.whl → 0.2.20__py3-none-any.whl

datarobot_genai/drmcp/tools/clients/confluence.py

@@ -31,6 +31,14 @@ from .atlassian import get_atlassian_cloud_id
 logger = logging.getLogger(__name__)
 
 
+class ConfluenceError(Exception):
+    """Exception for Confluence API errors."""
+
+    def __init__(self, message: str, status_code: int | None = None) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+
+
 class ConfluencePage(BaseModel):
     """Pydantic model for Confluence page."""
 
@@ -51,6 +59,22 @@ class ConfluencePage(BaseModel):
         }
 
 
+class ConfluenceComment(BaseModel):
+    """Pydantic model for Confluence comment."""
+
+    comment_id: str = Field(..., description="The unique comment ID")
+    page_id: str = Field(..., description="The page ID where the comment was added")
+    body: str = Field(..., description="Comment content in storage format")
+
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the comment."""
+        return {
+            "comment_id": self.comment_id,
+            "page_id": self.page_id,
+            "body": self.body,
+        }
+
+
 class ConfluenceClient:
     """
     Client for interacting with Confluence API using OAuth access token.
@@ -133,7 +157,7 @@ class ConfluenceClient:
 
         Raises
         ------
-            ValueError: If page is not found
+            ConfluenceError: If page is not found
             httpx.HTTPStatusError: If the API request fails
         """
         cloud_id = await self._get_cloud_id()
@@ -142,7 +166,7 @@ class ConfluenceClient:
         response = await self._client.get(url, params={"expand": self.EXPAND_FIELDS})
 
         if response.status_code == HTTPStatus.NOT_FOUND:
-            raise ValueError(f"Page with ID '{page_id}' not found")
+            raise ConfluenceError(f"Page with ID '{page_id}' not found", status_code=404)
 
         response.raise_for_status()
         return self._parse_response(response.json())
@@ -161,7 +185,7 @@ class ConfluenceClient:
 
         Raises
         ------
-            ValueError: If the page is not found
+            ConfluenceError: If the page is not found
             httpx.HTTPStatusError: If the API request fails
         """
         cloud_id = await self._get_cloud_id()
@@ -181,10 +205,183 @@ class ConfluenceClient:
         results = data.get("results", [])
 
         if not results:
-            raise ValueError(f"Page with title '{title}' not found in space '{space_key}'")
+            raise ConfluenceError(
+                f"Page with title '{title}' not found in space '{space_key}'", status_code=404
+            )
 
         return self._parse_response(results[0])
 
+    def _extract_error_message(self, response: httpx.Response) -> str:
+        """Extract error message from Confluence API error response."""
+        try:
+            error_data = response.json()
+            # Confluence API returns errors in different formats
+            if "message" in error_data:
+                return error_data["message"]
+            if "errorMessages" in error_data and error_data["errorMessages"]:
+                return "; ".join(error_data["errorMessages"])
+            if "errors" in error_data:
+                errors = error_data["errors"]
+                if isinstance(errors, list):
+                    return "; ".join(str(e) for e in errors)
+                if isinstance(errors, dict):
+                    return "; ".join(f"{k}: {v}" for k, v in errors.items())
+        except Exception:
+            pass
+        return response.text or "Unknown error"
+
+    async def create_page(
+        self,
+        space_key: str,
+        title: str,
+        body_content: str,
+        parent_id: int | None = None,
+    ) -> ConfluencePage:
+        """
+        Create a new Confluence page in a specified space.
+
+        Args:
+            space_key: The key of the Confluence space where the page should live
+            title: The title of the new page
+            body_content: The content in Confluence Storage Format (XML) or raw text
+            parent_id: Optional ID of the parent page for creating a child page
+
+        Returns
+        -------
+            ConfluencePage with the created page data
+
+        Raises
+        ------
+            ConfluenceError: If space not found, parent page not found, duplicate title,
+                            permission denied, or invalid content
+            httpx.HTTPStatusError: If the API request fails with unexpected status
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/content"
+
+        payload: dict[str, Any] = {
+            "type": "page",
+            "title": title,
+            "space": {"key": space_key},
+            "body": {
+                "storage": {
+                    "value": body_content,
+                    "representation": "storage",
+                }
+            },
+        }
+
+        if parent_id is not None:
+            payload["ancestors"] = [{"id": parent_id}]
+
+        response = await self._client.post(url, json=payload)
+
+        if response.status_code == HTTPStatus.NOT_FOUND:
+            error_msg = self._extract_error_message(response)
+            if parent_id is not None and "ancestor" in error_msg.lower():
+                raise ConfluenceError(
+                    f"Parent page with ID '{parent_id}' not found", status_code=404
+                )
+            raise ConfluenceError(
+                f"Space '{space_key}' not found or resource unavailable: {error_msg}",
+                status_code=404,
+            )
+
+        if response.status_code == HTTPStatus.CONFLICT:
+            raise ConfluenceError(
+                f"A page with title '{title}' already exists in space '{space_key}'",
+                status_code=409,
+            )
+
+        if response.status_code == HTTPStatus.FORBIDDEN:
+            raise ConfluenceError(
+                f"Permission denied: you don't have access to create pages in space '{space_key}'",
+                status_code=403,
+            )
+
+        if response.status_code == HTTPStatus.BAD_REQUEST:
+            error_msg = self._extract_error_message(response)
+            raise ConfluenceError(f"Invalid request: {error_msg}", status_code=400)
+
+        if response.status_code == HTTPStatus.TOO_MANY_REQUESTS:
+            raise ConfluenceError("Rate limit exceeded. Please try again later.", status_code=429)
+
+        response.raise_for_status()
+
+        return self._parse_response(response.json())
+
+    def _parse_comment_response(self, data: dict, page_id: str) -> ConfluenceComment:
+        """Parse API response into ConfluenceComment."""
+        body_content = ""
+        body = data.get("body", {})
+        if isinstance(body, dict):
+            storage = body.get("storage", {})
+            if isinstance(storage, dict):
+                body_content = storage.get("value", "")
+
+        return ConfluenceComment(
+            comment_id=str(data.get("id", "")),
+            page_id=page_id,
+            body=body_content,
+        )
+
+    async def add_comment(self, page_id: str, comment_body: str) -> ConfluenceComment:
+        """
+        Add a comment to a Confluence page.
+
+        Args:
+            page_id: The numeric page ID where the comment will be added
+            comment_body: The text content of the comment
+
+        Returns
+        -------
+            ConfluenceComment with the created comment data
+
+        Raises
+        ------
+            ConfluenceError: If page not found, permission denied, or invalid content
+            httpx.HTTPStatusError: If the API request fails with unexpected status
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/content"
+
+        payload: dict[str, Any] = {
+            "type": "comment",
+            "container": {"id": page_id, "type": "page"},
+            "body": {
+                "storage": {
+                    "value": comment_body,
+                    "representation": "storage",
+                }
+            },
+        }
+
+        response = await self._client.post(url, json=payload)
+
+        if response.status_code == HTTPStatus.NOT_FOUND:
+            error_msg = self._extract_error_message(response)
+            raise ConfluenceError(
+                f"Page with ID '{page_id}' not found: {error_msg}",
+                status_code=404,
+            )
+
+        if response.status_code == HTTPStatus.FORBIDDEN:
+            raise ConfluenceError(
+                f"Permission denied: you don't have access to add comments to page '{page_id}'",
+                status_code=403,
+            )
+
+        if response.status_code == HTTPStatus.BAD_REQUEST:
+            error_msg = self._extract_error_message(response)
+            raise ConfluenceError(f"Invalid request: {error_msg}", status_code=400)
+
+        if response.status_code == HTTPStatus.TOO_MANY_REQUESTS:
+            raise ConfluenceError("Rate limit exceeded. Please try again later.", status_code=429)
+
+        response.raise_for_status()
+
+        return self._parse_comment_response(response.json(), page_id)
+
     async def __aenter__(self) -> "ConfluenceClient":
         """Async context manager entry."""
         return self

datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -0,0 +1,248 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Google Drive API Client and utilities for OAuth."""
+
+import logging
+from typing import Annotated
+from typing import Any
+
+import httpx
+from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
+from fastmcp.exceptions import ToolError
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+
+from datarobot_genai.drmcp.core.auth import get_access_token
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_FIELDS = "nextPageToken,files(id,name,size,mimeType,webViewLink,createdTime,modifiedTime)"
+DEFAULT_ORDER = "modifiedTime desc"
+MAX_PAGE_SIZE = 100
+LIMIT = 500
+
+
+async def get_gdrive_access_token() -> str | ToolError:
+    """
+    Get Google Drive OAuth access token with error handling.
+
+    Returns
+    -------
+        Access token string on success, ToolError on failure
+
+    Example:
+        ```python
+        token = await get_gdrive_access_token()
+        if isinstance(token, ToolError):
+            # Handle error
+            return token
+        # Use token
+        ```
+    """
+    try:
+        access_token = await get_access_token("google")
+        if not access_token:
+            logger.warning("Empty access token received")
+            return ToolError("Received empty access token. Please complete the OAuth flow.")
+        return access_token
+    except OAuthServiceClientErr as e:
+        logger.error(f"OAuth client error: {e}", exc_info=True)
+        return ToolError(
+            "Could not obtain access token for Google. Make sure the OAuth "
+            "permission was granted for the application to act on your behalf."
+        )
+    except Exception as e:
+        logger.error(f"Unexpected error obtaining access token: {e}", exc_info=True)
+        return ToolError("An unexpected error occurred while obtaining access token for Google.")
+
+
+class GoogleDriveError(Exception):
+    """Exception for Google Drive API errors."""
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+PrimitiveData = str | int | float | bool | None
+
+
+class GoogleDriveFile(BaseModel):
+    """Represents a file from Google Drive."""
+
+    id: str
+    name: str
+    mime_type: Annotated[str, Field(alias="mimeType")]
+    size: int | None = None
+    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
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    @classmethod
+    def from_api_response(cls, data: dict[str, Any]) -> "GoogleDriveFile":
+        """Create a GoogleDriveFile from API response data."""
+        return cls(
+            id=data.get("id", "Unknown"),
+            name=data.get("name", "Unknown"),
+            mime_type=data.get("mimeType", "Unknown"),
+            size=int(data["size"]) if data.get("size") else None,
+            web_view_link=data.get("webViewLink"),
+            created_time=data.get("createdTime"),
+            modified_time=data.get("modifiedTime"),
+        )
+
+
+class PaginatedResult(BaseModel):
+    """Result of a paginated API call."""
+
+    files: list[GoogleDriveFile]
+    next_page_token: str | None = None
+
+
+class GoogleDriveClient:
+    """Client for interacting with Google Drive API."""
+
+    def __init__(self, access_token: str) -> None:
+        self._client = httpx.AsyncClient(
+            base_url="https://www.googleapis.com/drive/v3/files",
+            headers={"Authorization": f"Bearer {access_token}"},
+            timeout=30.0,
+        )
+
+    async def list_files(
+        self,
+        page_size: int,
+        limit: int,
+        page_token: str | None = None,
+        query: str | None = None,
+    ) -> PaginatedResult:
+        """
+        List files from Google Drive.
+
+        It's public API for GoogleDriveClient.
+
+        Args:
+            page_size: Number of files to return per 1 gdrive api request.
+            limit: Maximum number of files to return.
+            page_token: Optional token (specific for gdrive api) allowing to query next page.
+            query: Optional query to filter results.
+                If not provided it'll list all authorized user files.
+                If the query doesn't contain operators (contains, =, etc.), it will be treated as
+                a name search: "name contains '{query}'".
+
+        Returns
+        -------
+            List of Google Drive files.
+        """
+        if page_size <= 0:
+            raise GoogleDriveError("Error: page size must be positive.")
+        if limit <= 0:
+            raise GoogleDriveError("Error: limit must be positive.")
+        if limit < page_size:
+            raise GoogleDriveError("Error: limit must be bigger than or equal to page size.")
+        if limit % page_size != 0:
+            raise GoogleDriveError("Error: limit must be multiplication of page size.")
+
+        page_size = min(page_size, MAX_PAGE_SIZE)
+        limit = min(limit, LIMIT)
+        fetched = 0
+
+        formatted_query = self._get_formatted_query(query)
+
+        files: list[GoogleDriveFile] = []
+
+        while fetched < limit:
+            data = await self._list_files(
+                page_size=page_size,
+                page_token=page_token,
+                query=formatted_query,
+            )
+            files.extend(data.files)
+            fetched += len(data.files)
+            page_token = data.next_page_token
+
+            if not page_token:
+                break
+
+        return PaginatedResult(files=files, next_page_token=page_token)
+
+    async def _list_files(
+        self,
+        page_size: int,
+        page_token: str | None = None,
+        query: str | None = None,
+    ) -> PaginatedResult:
+        """Fetch a page of files from Google Drive."""
+        params: dict[str, PrimitiveData] = {
+            "pageSize": page_size,
+            "fields": DEFAULT_FIELDS,
+            "orderBy": DEFAULT_ORDER,
+        }
+        if page_token:
+            params["pageToken"] = page_token
+        if query:
+            params["q"] = query
+
+        response = await self._client.get(url="/", params=params)
+        response.raise_for_status()
+        data = response.json()
+
+        files = [
+            GoogleDriveFile.from_api_response(file_data) for file_data in data.get("files", [])
+        ]
+        next_page_token = data.get("nextPageToken")
+        return PaginatedResult(files=files, next_page_token=next_page_token)
+
+    @staticmethod
+    def _get_formatted_query(query: str | None) -> str | None:
+        """Get formatted Google Drive API query.
+
+        Args:
+            query: Optional search query string (e.g., "name contains 'report'"").
+                If the query doesn't contain operators (contains, =, etc.), it will be treated as
+                a name search: "name contains '{query}'".
+
+        Returns
+        -------
+            Correctly formatted query (if provided)
+        """
+        if not query:
+            return None
+
+        # If query doesn't look like a formatted query (no operators), format it as a name search
+        # Check if query already has Google Drive API operators
+        has_operator = any(
+            op in query for op in [" contains ", "=", "!=", " in ", " and ", " or ", " not "]
+        )
+        formatted_query = query
+        if not has_operator and query.strip():
+            # Simple text search - format as name contains query
+            # Escape backslashes first, then single quotes for Google Drive API
+            escaped_query = query.replace("\\", "\\\\").replace("'", "\\'")
+            formatted_query = f"name contains '{escaped_query}'"
+            logger.debug(f"Auto-formatted query '{query}' to '{formatted_query}'")
+        return formatted_query
+
+    async def __aenter__(self) -> "GoogleDriveClient":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
+    ) -> None:
+        """Async context manager exit."""
+        await self._client.aclose()

datarobot_genai/drmcp/tools/clients/jira.py

@@ -25,6 +25,18 @@ from .atlassian import get_atlassian_cloud_id
 
 logger = logging.getLogger(__name__)
 
+RESPONSE_JIRA_ISSUE_FIELDS = {
+    "id",
+    "key",
+    "summary",
+    "status",
+    "reporter",
+    "assignee",
+    "created",
+    "updated",
+}
+RESPONSE_JIRA_ISSUE_FIELDS_STR = ",".join(RESPONSE_JIRA_ISSUE_FIELDS)
+
 
 class _IssuePerson(BaseModel):
     email_address: str = Field(alias="emailAddress")
@@ -107,10 +119,41 @@ class JiraClient:
         self._cloud_id = await get_atlassian_cloud_id(self._client, service_type="jira")
         return self._cloud_id
 
-    async def _get_full_url(self, url: str) -> str:
+    async def _get_full_url(self, path: str) -> str:
         """Return URL for Jira API."""
         cloud_id = await self._get_cloud_id()
-        return f"{ATLASSIAN_API_BASE}/ex/jira/{cloud_id}/rest/api/3/{url}"
+        return f"{ATLASSIAN_API_BASE}/ex/jira/{cloud_id}/rest/api/3/{path}"
+
+    async def search_jira_issues(self, jql_query: str, max_results: int) -> list[Issue]:
+        """
+        Search Jira issues using JQL (Jira Query Language).
+
+        Args:
+            jql_query: JQL Query
+            max_results: Maximum number of issues to return
+
+        Returns
+        -------
+            List of Jira issues
+
+        Raises
+        ------
+            httpx.HTTPStatusError: If the API request fails
+        """
+        url = await self._get_full_url("search/jql")
+        response = await self._client.post(
+            url,
+            json={
+                "jql": jql_query,
+                "fields": list(RESPONSE_JIRA_ISSUE_FIELDS),
+                "maxResults": max_results,
+            },
+        )
+
+        response.raise_for_status()
+        raw_issues = response.json().get("issues", [])
+        issues = [Issue(**issue) for issue in raw_issues]
+        return issues
 
     async def get_jira_issue(self, issue_key: str) -> Issue:
         """
@@ -128,9 +171,7 @@ class JiraClient:
             httpx.HTTPStatusError: If the API request fails
         """
         url = await self._get_full_url(f"issue/{issue_key}")
-        response = await self._client.get(
-            url, params={"fields": "id,key,summary,status,reporter,assignee,created,updated"}
-        )
+        response = await self._client.get(url, params={"fields": RESPONSE_JIRA_ISSUE_FIELDS_STR})
 
         if response.status_code == HTTPStatus.NOT_FOUND:
             raise ValueError(f"{issue_key} not found")
@@ -209,6 +250,79 @@ class JiraClient:
         jsoned = response.json()
         return jsoned["key"]
 
+    async def update_jira_issue(self, issue_key: str, fields: dict[str, Any]) -> list[str]:
+        """
+        Update Jira issue.
+
+        Args:
+            issue_key: The key of the Jira issue, e.g., 'PROJ-123'
+            fields: A dictionary of field names and their new values
+                e.g., {'description': 'New content'}
+
+        Returns
+        -------
+            List of updated fields
+
+        Raises
+        ------
+            httpx.HTTPStatusError: If the API request fails
+        """
+        url = await self._get_full_url(f"issue/{issue_key}")
+        payload = {"fields": fields}
+
+        response = await self._client.put(url, json=payload)
+
+        response.raise_for_status()
+        return list(fields.keys())
+
+    async def get_available_jira_transitions(self, issue_key: str) -> dict[str, str]:
+        """
+        Get Available Jira Transitions.
+
+        Args:
+            issue_key: The key of the Jira issue, e.g., 'PROJ-123'
+
+        Returns
+        -------
+            Dictionary where key is the transition name and value is the transition ID
+
+        Raises
+        ------
+            httpx.HTTPStatusError: If the API request fails
+        """
+        url = await self._get_full_url(f"issue/{issue_key}/transitions")
+        response = await self._client.get(url)
+        response.raise_for_status()
+        jsoned = response.json()
+        transitions = {
+            transition["name"]: transition["id"] for transition in jsoned.get("transitions", [])
+        }
+        return transitions
+
+    async def transition_jira_issue(self, issue_key: str, transition_id: str) -> None:
+        """
+        Transition Jira issue.
+
+        Args:
+            issue_key: The key of the Jira issue, e.g., 'PROJ-123'
+            transition_id: Id of target transitionm e.g. '123'.
+                Can be obtained from `get_available_jira_transitions`.
+
+        Returns
+        -------
+            Nothing
+
+        Raises
+        ------
+            httpx.HTTPStatusError: If the API request fails
+        """
+        url = await self._get_full_url(f"issue/{issue_key}")
+        payload = {"transition": {"id": transition_id}}
+
+        response = await self._client.post(url, json=payload)
+
+        response.raise_for_status()
+
     async def __aenter__(self) -> "JiraClient":
         """Async context manager entry."""
         return self