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

datarobot_genai/drmcp/tools/clients/confluence.py

@@ -30,6 +30,9 @@ from .atlassian import get_atlassian_cloud_id
 
 logger = logging.getLogger(__name__)
 
+# Search expand fields for CQL search - content.space gives us space.key directly
+SEARCH_EXPAND_FIELDS = "content.space"
+
 
 class ConfluenceError(Exception):
     """Exception for Confluence API errors."""
@@ -75,6 +78,32 @@ class ConfluenceComment(BaseModel):
         }
 
 
+class ContentSearchResult(BaseModel):
+    """Pydantic model for Confluence search result item."""
+
+    id: str
+    title: str
+    type: str
+    space_key: str = ""
+    space_name: str = ""
+    excerpt: str = ""
+    last_modified: str | None = None
+    url: str = ""
+
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the search result."""
+        return {
+            "id": self.id,
+            "title": self.title,
+            "type": self.type,
+            "spaceKey": self.space_key,
+            "spaceName": self.space_name,
+            "excerpt": self.excerpt,
+            "lastModified": self.last_modified,
+            "url": self.url,
+        }
+
+
 class ConfluenceClient:
     """
     Client for interacting with Confluence API using OAuth access token.
@@ -382,6 +411,76 @@ class ConfluenceClient:
 
         return self._parse_comment_response(response.json(), page_id)
 
+    async def search_confluence_content(
+        self, cql_query: str, max_results: int
+    ) -> list[ContentSearchResult]:
+        """
+        Search Confluence content using CQL (Confluence Query Language).
+
+        Args:
+            cql_query: CQL Query
+            max_results: Maximum number of results to return
+
+        Returns
+        -------
+            List of Confluence content search results
+
+        Raises
+        ------
+            ConfluenceError: If the API request fails (400, 403, 429)
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/search"
+
+        response = await self._client.get(
+            url,
+            params={
+                "cql": cql_query,
+                "limit": max_results,
+                "expand": SEARCH_EXPAND_FIELDS,
+            },
+        )
+
+        if response.status_code == HTTPStatus.BAD_REQUEST:
+            error_msg = self._extract_error_message(response)
+            raise ConfluenceError(f"Invalid CQL query: {error_msg}", status_code=400)
+
+        if response.status_code == HTTPStatus.FORBIDDEN:
+            raise ConfluenceError(
+                "Permission denied: you don't have access to search this content",
+                status_code=403,
+            )
+
+        if response.status_code == HTTPStatus.TOO_MANY_REQUESTS:
+            raise ConfluenceError("Rate limit exceeded. Please try again later.", status_code=429)
+
+        response.raise_for_status()
+        raw_results = response.json().get("results", [])
+        results = [ContentSearchResult(**self._parse_search_item(item)) for item in raw_results]
+        return results
+
+    def _parse_search_item(self, item: dict) -> dict:
+        """Parse raw search API response item into model-compatible dict."""
+        content = item.get("content", item)
+        links = content.get("_links", {})
+        base_url = links.get("base", "")
+        webui = links.get("webui", "")
+        url = f"{base_url}{webui}" if base_url and webui else webui
+
+        # Get space from content.space (requires expand=content.space)
+        content_space = content.get("space", {})
+
+        return {
+            "id": str(content.get("id", "")),
+            "title": content.get("title", ""),
+            "type": content.get("type", "page"),
+            "space_key": content_space.get("key", ""),
+            "space_name": content_space.get("name", ""),
+            "excerpt": item.get("excerpt", ""),
+            "last_modified": item.get("lastModified"),
+            "url": url,
+        }
+
     async def __aenter__(self) -> "ConfluenceClient":
         """Async context manager entry."""
         return self

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

@@ -186,3 +186,69 @@ async def confluence_add_comment(
             "page_id": page_id,
         },
     )
+
+
+@dr_mcp_tool(tags={"confluence", "search", "content"})
+async def confluence_search(
+    *,
+    cql_query: Annotated[
+        str,
+        "The CQL (Confluence Query Language) string used to filter content, "
+        "e.g., 'type=page and space=DOC'.",
+    ],
+    max_results: Annotated[int, "Maximum number of content items to return. Default is 10."] = 10,
+    include_body: Annotated[
+        bool,
+        "If True, fetch full page body content for each result (slower, "
+        "makes additional API calls). Default is False, which returns only excerpts.",
+    ] = False,
+) -> ToolResult:
+    """
+    Search Confluence pages and content efficiently using a CQL query string.
+    This pushes the search logic to the Confluence API (Push-Down).
+
+    Refer to Confluence documentation for advanced searching using CQL:
+    https://developer.atlassian.com/cloud/confluence/advanced-searching-using-cql/
+    """
+    if not cql_query:
+        raise ToolError("Argument validation error: 'cql_query' cannot be empty.")
+
+    if max_results < 1 or max_results > 100:
+        raise ToolError("Argument validation error: 'max_results' must be between 1 and 100.")
+
+    access_token = await get_atlassian_access_token()
+    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)}")
+
+    n = len(results)
+    return ToolResult(
+        content=f"Successfully executed CQL query and retrieved {n} result(s).",
+        structured_content={"data": data, "count": n},
+    )

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.20.dist-info → datarobot_genai-0.2.22.dist-info}/METADATA

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

@@ -78,14 +78,14 @@ datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdY
 datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 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/confluence.py,sha256=YS5XsKd-jK5Yg0rgwOcC76v9e8fDJgUZIW5B9kcq5B0,17101
+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=jSF7yXGFqqlMcavkRIY4HbMxb7tCeunA2ST41wa2vGI,7219
+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.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,,
+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,,