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

datarobot_genai/drmcp/core/config.py

@@ -245,6 +245,30 @@ class MCPServerConfig(BaseSettings):
             os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
         )
 
+    # Gdrive tools
+    enable_gdrive_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
+            "ENABLE_GDRIVE_TOOLS",
+        ),
+        description="Enable/disable GDrive tools",
+    )
+    is_gdrive_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether GDrive OAuth provider is configured for GDrive integration",
+    )
+
+    @property
+    def is_gdrive_oauth_configured(self) -> bool:
+        return self.is_gdrive_oauth_provider_configured or bool(
+            os.getenv("GDRIVE_CLIENT_ID") and os.getenv("GDRIVE_CLIENT_SECRET")
+        )
+
     @field_validator(
         "otel_attributes",
         mode="before",

datarobot_genai/drmcp/core/tool_config.py

@@ -29,6 +29,7 @@ class ToolType(str, Enum):
     PREDICTIVE = "predictive"
     JIRA = "jira"
     CONFLUENCE = "confluence"
+    GDRIVE = "gdrive"
 
 
 class ToolConfig(TypedDict):
@@ -64,6 +65,13 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
         package_prefix="datarobot_genai.drmcp.tools.confluence",
         config_field_name="enable_confluence_tools",
     ),
+    ToolType.GDRIVE: ToolConfig(
+        name="gdrive",
+        oauth_check=lambda config: config.is_gdrive_oauth_configured,
+        directory="gdrive",
+        package_prefix="datarobot_genai.drmcp.tools.gdrive",
+        config_field_name="enable_gdrive_tools",
+    ),
 }
 
 

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/gdrive/tools.py

@@ -0,0 +1,88 @@
+# 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 MCP tools for interacting with Google Drive API."""
+
+import logging
+from typing import Annotated
+
+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 LIMIT
+from datarobot_genai.drmcp.tools.clients.gdrive import MAX_PAGE_SIZE
+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__)
+
+
+@dr_mcp_tool(tags={"google", "gdrive", "list", "files"})
+async def google_drive_list_files(
+    *,
+    page_size: Annotated[
+        int, f"Maximum number of files to return per page (max {MAX_PAGE_SIZE})."
+    ] = 10,
+    limit: Annotated[int, f"Total maximum number of files to return (max {LIMIT})."] = 50,
+    page_token: Annotated[
+        str | None, "The token for the next page of results, retrieved from a previous call."
+    ] = None,
+    query: Annotated[
+        str | None, "Optional filter to narrow results (e.g., 'trashed = false')."
+    ] = None,
+) -> ToolResult | ToolError:
+    """
+    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.
+
+    Limit must be bigger than or equal to page size and it must be multiplication of page size.
+    Ex.
+        page size = 10 limit = 50
+        page size = 3 limit = 3
+        page size = 12 limit = 36
+    """
+    access_token = await get_gdrive_access_token()
+    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
+            )
+    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)}")
+
+    number_of_files = len(data.files)
+    next_page_info = (
+        f"Next page token needed to fetch more data: {data.next_page_token}"
+        if data.next_page_token
+        else "There're no more pages."
+    )
+    return ToolResult(
+        content=f"Successfully listed {number_of_files} files. {next_page_info}",
+        structured_content={
+            "files": [
+                file.model_dump(by_alias=True, include={"id", "name"}) for file in data.files
+            ],
+            "count": number_of_files,
+            "nextPageToken": data.next_page_token,
+        },
+    )

{datarobot_genai-0.2.19.dist-info → datarobot_genai-0.2.20.dist-info}/METADATA

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

{datarobot_genai-0.2.19.dist-info → datarobot_genai-0.2.20.dist-info}/RECORD

@@ -27,7 +27,7 @@ datarobot_genai/drmcp/server.py,sha256=KE4kjS5f9bfdYftG14HBHrfvxDfCD4pwCXePfvl1O
 datarobot_genai/drmcp/core/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
 datarobot_genai/drmcp/core/auth.py,sha256=E-5wrGbBFEBlD5377g6Exddrc7HsazamwX8tWr2RLXY,5815
 datarobot_genai/drmcp/core/clients.py,sha256=y-yG8617LbmiZ_L7FWfMrk4WjIekyr76u_Q80aLqGpI,5524
-datarobot_genai/drmcp/core/config.py,sha256=0Y0OjbLq-buP66U9tiXMyqrX9wWIn557vKvkBvjM_cM,11797
+datarobot_genai/drmcp/core/config.py,sha256=69QDsVVSvjzv1uIHOjtQGzdg7_Ic4sA3vLA6ZbZJ1Ok,12674
 datarobot_genai/drmcp/core/config_utils.py,sha256=U-aieWw7MyP03cGDFIp97JH99ZUfr3vD9uuTzBzxn7w,6428
 datarobot_genai/drmcp/core/constants.py,sha256=lUwoW_PTrbaBGqRJifKqCn3EoFacoEgdO-CpoFVrUoU,739
 datarobot_genai/drmcp/core/credentials.py,sha256=PYEUDNMVw1BoMzZKLkPVTypNkVevEPtmk3scKnE-zYg,6706
@@ -41,7 +41,7 @@ datarobot_genai/drmcp/core/routes.py,sha256=dqE2M0UzAyyN9vQjlyTjYW4rpju3LT039po5
 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
 datarobot_genai/drmcp/core/telemetry.py,sha256=NEkSTC1w6uQgtukLHI-sWvR4EMgInysgATcvfQ5CplM,15378
-datarobot_genai/drmcp/core/tool_config.py,sha256=5OkC3e-vekZtdqg-DbwcyadGSrDxmZqSDey2YyGVn1M,2978
+datarobot_genai/drmcp/core/tool_config.py,sha256=5JCWO70ZH-K-34yS7vYJG2nl4i9UO_q_W9NCoWSXXno,3271
 datarobot_genai/drmcp/core/tool_filter.py,sha256=tLOcG50QBvS48cOVHM6OqoODYiiS6KeM_F-2diaHkW0,2858
 datarobot_genai/drmcp/core/utils.py,sha256=EvfpqKZ3tECMoxpIQ_tA_3rOgy6KJEYKC0lWZo_Daag,4517
 datarobot_genai/drmcp/core/dynamic_prompts/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
@@ -79,10 +79,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=gDzy8t5t3b1mwEr-CuZ5BwXXQ52AXke8J_Ra7i_8T1g,13692
+datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=QmNTmJdSqYO5Y5Vnp3roNZiNNJeocBVjF9UcSzcjgRY,8635
 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=jSF7yXGFqqlMcavkRIY4HbMxb7tCeunA2ST41wa2vGI,7219
+datarobot_genai/drmcp/tools/gdrive/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=wmCUSaCWqepdlOIApA8tZ-grPYSV7wZKoer6uRy26Qg,3459
 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
@@ -108,9 +111,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.19.dist-info/METADATA,sha256=rywu2LteAnHoLxfq84xckJSU10XX_sKMqli7pqsuCgg,6301
-datarobot_genai-0.2.19.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.19.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.19.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.19.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.19.dist-info/RECORD,,
+datarobot_genai-0.2.20.dist-info/METADATA,sha256=iWXFNkplo7YmQoa5bCbVPhhv-KzDIHsotkbKf-bEbEk,6301
+datarobot_genai-0.2.20.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.20.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.20.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.20.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.20.dist-info/RECORD,,