datarobot-genai 0.2.21__py3-none-any.whl → 0.2.22__py3-none-any.whl

datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -29,7 +29,10 @@ 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)"
+SUPPORTED_FIELDS = {"id", "name", "size", "mimeType", "webViewLink", "createdTime", "modifiedTime"}
+SUPPORTED_FIELDS_STR = ",".join(SUPPORTED_FIELDS)
+DEFAULT_FIELDS = f"nextPageToken,files({SUPPORTED_FIELDS_STR})"
+GOOGLE_DRIVE_FOLDER_MIME = "application/vnd.google-apps.folder"
 DEFAULT_ORDER = "modifiedTime desc"
 MAX_PAGE_SIZE = 100
 LIMIT = 500
@@ -129,6 +132,8 @@ class GoogleDriveClient:
         limit: int,
         page_token: str | None = None,
         query: str | None = None,
+        folder_id: str | None = None,
+        recursive: bool = False,
     ) -> PaginatedResult:
         """
         List files from Google Drive.
@@ -143,6 +148,10 @@ class GoogleDriveClient:
                 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}'".
+            folder_id: The ID of a specific folder to list or search within.
+                If omitted, searches the entire Drive.
+            recursive: If True, searches all subfolders.
+                If False and folder_id is provided, only lists immediate children.
 
         Returns
         -------
@@ -159,26 +168,85 @@ class GoogleDriveClient:
 
         page_size = min(page_size, MAX_PAGE_SIZE)
         limit = min(limit, LIMIT)
-        fetched = 0
+        formatted_query = self._build_query(query, folder_id)
+
+        if not recursive or not folder_id:
+            files, next_token = await self._fetch_paginated(
+                page_size=page_size,
+                limit=limit,
+                page_token=page_token,
+                query=formatted_query,
+            )
+            return PaginatedResult(files=files, next_page_token=next_token)
 
-        formatted_query = self._get_formatted_query(query)
+        files = await self._fetch_recursive(
+            root_folder_id=folder_id,
+            base_query=query,
+            page_size=page_size,
+            limit=limit,
+        )
 
+        return PaginatedResult(files=files, next_page_token=page_token)
+
+    async def _fetch_paginated(
+        self,
+        page_size: int,
+        limit: int,
+        page_token: str | None,
+        query: str | None,
+    ) -> tuple[list[GoogleDriveFile], str | None]:
+        fetched = 0
         files: list[GoogleDriveFile] = []
+        next_page_token = page_token
 
         while fetched < limit:
             data = await self._list_files(
                 page_size=page_size,
-                page_token=page_token,
-                query=formatted_query,
+                page_token=next_page_token,
+                query=query,
             )
+
             files.extend(data.files)
             fetched += len(data.files)
-            page_token = data.next_page_token
+            next_page_token = data.next_page_token
 
-            if not page_token:
+            if not next_page_token:
                 break
 
-        return PaginatedResult(files=files, next_page_token=page_token)
+        return files, next_page_token
+
+    async def _fetch_recursive(
+        self,
+        root_folder_id: str,
+        base_query: str | None,
+        page_size: int,
+        limit: int,
+    ) -> list[GoogleDriveFile]:
+        collected: list[GoogleDriveFile] = []
+        folders_to_visit: list[str] = [root_folder_id]
+
+        while folders_to_visit and len(collected) < limit:
+            current_folder = folders_to_visit.pop(0)
+
+            query = self._build_query(base_query, current_folder)
+
+            files, _ = await self._fetch_paginated(
+                page_size=page_size,
+                limit=limit - len(collected),
+                page_token=None,
+                query=query,
+            )
+
+            for file in files:
+                collected.append(file)
+
+                if file.mime_type == GOOGLE_DRIVE_FOLDER_MIME:
+                    folders_to_visit.append(file.id)
+
+                if len(collected) >= limit:
+                    break
+
+        return collected
 
     async def _list_files(
         self,
@@ -207,6 +275,45 @@ class GoogleDriveClient:
         next_page_token = data.get("nextPageToken")
         return PaginatedResult(files=files, next_page_token=next_page_token)
 
+    def _build_query(self, query: str | None, folder_id: str | None) -> str | None:
+        """Build 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}'".
+            folder_id: Optional folder id.
+                If provided it'll narrow query to search/list only in given folder.
+
+        Returns
+        -------
+            Correctly builded query (if provided)
+        """
+        base_query = self._get_formatted_query(query)
+
+        if base_query:
+            # Case #1 -- Some query provided and contains in parents (gdrive "folder id")
+            if "in parents" in base_query and folder_id:
+                logger.debug(
+                    "In-parents (parent folder) already used in query. "
+                    "Omiting folder_id argument. "
+                    f"Query: {base_query} | FolderId: {folder_id}"
+                )
+                return base_query
+            # Case #2 -- Some query provided without "in parents" and folder id provided.
+            elif folder_id:
+                return f"{base_query} and '{folder_id}' in parents"
+            # Case #3 -- Query provided without "in parents" and no folder id.
+            else:
+                return base_query
+
+        # Case #4 -- Base query is null but folder id provided
+        if folder_id:
+            return f"'{folder_id}' in parents"
+
+        # Case #5 -- Neither query not folder provided
+        return None
+
     @staticmethod
     def _get_formatted_query(query: str | None) -> str | None:
         """Get formatted Google Drive API query.

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -23,6 +23,8 @@ 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 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
@@ -30,8 +32,8 @@ 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(
+@dr_mcp_tool(tags={"google", "gdrive", "list", "search", "files", "find", "contents"})
+async def gdrive_find_contents(
     *,
     page_size: Annotated[
         int, f"Maximum number of files to return per page (max {MAX_PAGE_SIZE})."
@@ -43,9 +45,24 @@ async def google_drive_list_files(
     query: Annotated[
         str | None, "Optional filter to narrow results (e.g., 'trashed = false')."
     ] = None,
+    folder_id: Annotated[
+        str | None,
+        "The ID of a specific folder to list or search within. "
+        "If omitted, searches the entire Drive.",
+    ] = None,
+    recursive: Annotated[
+        bool,
+        "If True, searches all subfolders. "
+        "If False and folder_id is provided, only lists immediate children.",
+    ] = False,
+    fields: Annotated[
+        list[str] | None,
+        "Optional list of metadata fields to include. Ex. id, name, mimeType. "
+        f"Default = {SUPPORTED_FIELDS_STR}",
+    ] = None,
 ) -> ToolResult | ToolError:
     """
-    List files in the user's Google Drive with pagination and filtering support.
+    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.
 
     Limit must be bigger than or equal to page size and it must be multiplication of page size.
@@ -61,7 +78,12 @@ async def google_drive_list_files(
     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
+                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}")
@@ -70,6 +92,7 @@ async def google_drive_list_files(
         logger.error(f"Unexpected error listing Google Drive files: {e}")
         raise ToolError(f"An unexpected error occurred while listing Google Drive files: {str(e)}")
 
+    filtered_fields = set(fields).intersection(SUPPORTED_FIELDS) if fields else SUPPORTED_FIELDS
     number_of_files = len(data.files)
     next_page_info = (
         f"Next page token needed to fetch more data: {data.next_page_token}"
@@ -80,7 +103,7 @@ async def google_drive_list_files(
         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
+                file.model_dump(by_alias=True, include=filtered_fields) for file in data.files
             ],
             "count": number_of_files,
             "nextPageToken": data.next_page_token,

{datarobot_genai-0.2.21.dist-info → datarobot_genai-0.2.22.dist-info}/METADATA

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

{datarobot_genai-0.2.21.dist-info → datarobot_genai-0.2.22.dist-info}/RECORD

@@ -79,13 +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=YS5XsKd-jK5Yg0rgwOcC76v9e8fDJgUZIW5B9kcq5B0,17101
-datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=QmNTmJdSqYO5Y5Vnp3roNZiNNJeocBVjF9UcSzcjgRY,8635
+datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=gRtWWCENHcmLepKQbS7qsF4R6vbQQK1Ru-EqfeUbldY,12550
 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=ySwABe8osAzky3BO3lRaF6UHnXQgaurkmvM0iHFfL30,9849
 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/gdrive/tools.py,sha256=BP5tcpciuijakmXTjEgS6CySg5TUBAmlKYPkTgpVZbc,4406
 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
@@ -111,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.21.dist-info/METADATA,sha256=wdbafq2Z77RRyxy7nG-kYCUU5pfCbsDbBTH7QdENnRU,6301
-datarobot_genai-0.2.21.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.21.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.21.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.21.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.21.dist-info/RECORD,,
+datarobot_genai-0.2.22.dist-info/METADATA,sha256=Jc7FEPYKQq7DI3dpk2qbMSYJolHyrBeO-A_9i7IsFFM,6301
+datarobot_genai-0.2.22.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.22.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.22.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.22.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.22.dist-info/RECORD,,