datarobot-genai 0.2.24__py3-none-any.whl → 0.2.29__py3-none-any.whl

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -21,6 +21,7 @@ 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 GOOGLE_DRIVE_FOLDER_MIME
 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
@@ -60,7 +61,7 @@ async def gdrive_find_contents(
         "Optional list of metadata fields to include. Ex. id, name, mimeType. "
         f"Default = {SUPPORTED_FIELDS_STR}",
     ] = None,
-) -> ToolResult | ToolError:
+) -> ToolResult:
     """
     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.
@@ -121,7 +122,7 @@ async def gdrive_read_content(
         "(e.g., 'text/markdown' for Docs, 'text/csv' for Sheets). "
         "If not specified, uses sensible defaults. Has no effect on regular files.",
     ] = None,
-) -> ToolResult | ToolError:
+) -> ToolResult:
     """
     Retrieve the content of a specific file by its ID. Google Workspace files are
     automatically exported to LLM-readable formats (Push-Down).
@@ -175,3 +176,95 @@ async def gdrive_read_content(
         ),
         structured_content=file_content.as_flat_dict(),
     )
+
+
+@dr_mcp_tool(tags={"google", "gdrive", "create", "write", "file", "folder"}, enabled=False)
+async def gdrive_create_file(
+    *,
+    name: Annotated[str, "The name for the new file or folder."],
+    mime_type: Annotated[
+        str,
+        "The MIME type of the file (e.g., 'text/plain', "
+        "'application/vnd.google-apps.document', 'application/vnd.google-apps.folder').",
+    ],
+    parent_id: Annotated[
+        str | None, "The ID of the parent folder where the file should be created."
+    ] = None,
+    initial_content: Annotated[
+        str | None, "Text content to populate the new file, if applicable."
+    ] = None,
+) -> ToolResult:
+    """
+    Create a new file or folder in Google Drive.
+
+    This tool is essential for an AI agent to generate new output (like reports or
+    documentation) directly into the Drive structure.
+
+    Usage:
+        - Create empty file: gdrive_create_file(name="report.txt", mime_type="text/plain")
+        - Create Google Doc: gdrive_create_file(
+            name="My Report",
+            mime_type="application/vnd.google-apps.document",
+            initial_content="# Report Title"
+          )
+        - Create folder: gdrive_create_file(
+            name="Reports",
+            mime_type="application/vnd.google-apps.folder"
+          )
+        - Create in subfolder: gdrive_create_file(
+            name="file.txt",
+            mime_type="text/plain",
+            parent_id="folder_id_here",
+            initial_content="File content"
+          )
+
+    Supported MIME types:
+        - text/plain: Plain text file
+        - application/vnd.google-apps.document: Google Doc (content auto-converted)
+        - application/vnd.google-apps.spreadsheet: Google Sheet (CSV content works best)
+        - application/vnd.google-apps.folder: Folder (initial_content is ignored)
+
+    Note: For Google Workspace files, the Drive API automatically converts plain text
+    content to the appropriate format.
+    """
+    if not name or not name.strip():
+        raise ToolError("Argument validation error: 'name' cannot be empty.")
+
+    if not mime_type or not mime_type.strip():
+        raise ToolError("Argument validation error: 'mime_type' cannot be empty.")
+
+    access_token = await get_gdrive_access_token()
+    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)}")
+
+    # Build response message
+    file_type = "folder" if mime_type == GOOGLE_DRIVE_FOLDER_MIME else "file"
+    content_info = ""
+    if initial_content and mime_type != GOOGLE_DRIVE_FOLDER_MIME:
+        content_info = " with initial content"
+
+    return ToolResult(
+        content=f"Successfully created {file_type} '{created_file.name}'{content_info}.",
+        structured_content={
+            "id": created_file.id,
+            "name": created_file.name,
+            "mimeType": created_file.mime_type,
+            "webViewLink": created_file.web_view_link,
+            "createdTime": created_file.created_time,
+        },
+    )

datarobot_genai/drmcp/tools/microsoft_graph/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 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.

datarobot_genai/drmcp/tools/microsoft_graph/tools.py

@@ -0,0 +1,198 @@
+# Copyright 2026 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.
+
+"""Microsoft Graph MCP tools for searching SharePoint and OneDrive content."""
+
+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.microsoft_graph import MicrosoftGraphClient
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import MicrosoftGraphError
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import get_microsoft_graph_access_token
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import validate_site_url
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(
+    tags={
+        "microsoft",
+        "graph api",
+        "sharepoint",
+        "drive",
+        "list",
+        "search",
+        "files",
+        "find",
+        "contents",
+    }
+)
+async def microsoft_graph_search_content(
+    *,
+    search_query: Annotated[str, "The search string to find files, folders, or list items."],
+    site_url: Annotated[
+        str | None,
+        "Optional SharePoint site URL to scope the search "
+        "(e.g., https://tenant.sharepoint.com/sites/sitename). "
+        "If not provided, searches across all accessible sites.",
+    ] = None,
+    site_id: Annotated[
+        str | None,
+        "Optional ID of the site to scope the search. If provided, takes precedence over site_url.",
+    ] = None,
+    from_offset: Annotated[
+        int,
+        "The zero-based index of the first result to return. Use this for pagination. "
+        "Default: 0 (start from the beginning). To get the next page, increment by the size "
+        "value (e.g., first page: from=0 size=250, second page: from=250 size=250, "
+        "third page: from=500 size=250).",
+    ] = 0,
+    size: Annotated[
+        int,
+        "Maximum number of results to return in this request. Default is 250, max is 250. "
+        "The LLM should control pagination by making multiple calls with different 'from' values.",
+    ] = 250,
+    entity_types: Annotated[
+        list[str] | None,
+        "Optional list of entity types to search. Valid values: 'driveItem', 'listItem', "
+        "'site', 'list', 'drive'. Default: ['driveItem', 'listItem']. "
+        "Multiple types can be specified.",
+    ] = None,
+    filters: Annotated[
+        list[str] | None,
+        "Optional list of KQL filter expressions to refine search results "
+        "(e.g., ['fileType:docx', 'size>1000']).",
+    ] = None,
+    include_hidden_content: Annotated[
+        bool,
+        "Whether to include hidden content in search results. Only works with delegated "
+        "permissions, not application permissions. Default: False.",
+    ] = False,
+    region: Annotated[
+        str | None,
+        "Optional region code for application permissions (e.g., 'NAM', 'EUR', 'APC'). "
+        "Required when using application permissions to search SharePoint content in "
+        "specific regions.",
+    ] = None,
+) -> ToolResult | ToolError:
+    """
+    Search for SharePoint and OneDrive content using Microsoft Graph Search API.
+
+    Search Scope:
+    - When site_url or site_id is provided: searches within the specified SharePoint site
+    - When neither is provided: searches across all accessible SharePoint sites and OneDrive
+
+    Supported Entity Types:
+    - driveItem: Files and folders in document libraries and OneDrive
+    - listItem: Items in SharePoint lists
+    - site: SharePoint sites
+    - list: SharePoint lists
+    - drive: Document libraries/drives
+
+    Filtering:
+    - Filters use KQL (Keyword Query Language) syntax
+    - Multiple filters are combined with AND operators
+    - Examples: ['fileType:docx', 'size>1000', 'lastModifiedTime>2024-01-01']
+    - Filters are applied in addition to the search query
+
+    Pagination:
+    - Controlled via from_offset (zero-based index) and size parameters
+    - Maximum size per request: 250 results
+    - To paginate: increment from_offset by size value for each subsequent page
+    - Example pagination sequence:
+      * Page 1: from_offset=0, size=250 (returns results 0-249)
+      * Page 2: from_offset=250, size=250 (returns results 250-499)
+      * Page 3: from_offset=500, size=250 (returns results 500-749)
+
+    API Reference:
+    - Endpoint: POST /search/query
+    - Documentation: https://learn.microsoft.com/en-us/graph/api/search-query
+    - Search concepts: https://learn.microsoft.com/en-us/graph/search-concept-files
+
+    Permissions:
+    - Requires Sites.Read.All or Sites.Search.All permission
+    - include_hidden_content only works with delegated permissions
+    - region parameter is required for application permissions in multi-region environments
+    """
+    if not search_query:
+        raise ToolError("Argument validation error: 'search_query' cannot be empty.")
+
+    # Validate site_url if provided
+    if site_url:
+        validation_error = validate_site_url(site_url)
+        if validation_error:
+            raise ToolError(validation_error)
+
+    access_token = await get_microsoft_graph_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with MicrosoftGraphClient(access_token=access_token, site_url=site_url) as client:
+            items = await client.search_content(
+                search_query=search_query,
+                site_id=site_id,
+                from_offset=from_offset,
+                size=size,
+                entity_types=entity_types,
+                filters=filters,
+                include_hidden_content=include_hidden_content,
+                region=region,
+            )
+    except MicrosoftGraphError as e:
+        logger.error(f"Microsoft Graph error searching content: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error searching Microsoft Graph content: {e}", exc_info=True)
+        raise ToolError(
+            f"An unexpected error occurred while searching Microsoft Graph content: {str(e)}"
+        )
+
+    results = []
+    for item in items:
+        result_dict = {
+            "id": item.id,  # Unique ID of the file, folder, or list item
+            "name": item.name,
+            "webUrl": item.web_url,
+            "size": item.size,
+            "createdDateTime": item.created_datetime,
+            "lastModifiedDateTime": item.last_modified_datetime,
+            "isFolder": item.is_folder,
+            "mimeType": item.mime_type,
+            # Document library/drive ID (driveId in Microsoft Graph API)
+            "documentLibraryId": item.drive_id,
+            "parentFolderId": item.parent_folder_id,  # Parent folder ID
+        }
+        results.append(result_dict)
+
+    n = len(results)
+    return ToolResult(
+        content=(
+            f"Successfully searched Microsoft Graph and retrieved {n} result(s) for "
+            f"'{search_query}' (from={from_offset}, size={size})."
+        ),
+        structured_content={
+            "query": search_query,
+            "siteUrl": site_url,
+            "siteId": site_id,
+            "from": from_offset,
+            "size": size,
+            "results": results,
+            "count": n,
+        },
+    )

datarobot_genai/drmcp/tools/predictive/data.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import logging
 import os
 from typing import Annotated
@@ -28,6 +29,7 @@ logger = logging.getLogger(__name__)
 
 @dr_mcp_tool(tags={"predictive", "data", "write", "upload", "catalog"})
 async def upload_dataset_to_ai_catalog(
+    *,
     file_path: Annotated[str, "The path to the dataset file to upload."] | None = None,
     file_url: Annotated[str, "The URL to the dataset file to upload."] | None = None,
 ) -> ToolError | ToolResult:
@@ -80,11 +82,17 @@ async def list_ai_catalog_items() -> ToolResult:
             structured_content={"datasets": []},
         )
 
+    datasets_dict = {ds.id: ds.name for ds in datasets}
+    datasets_count = len(datasets)
+
     return ToolResult(
-        content=f"Found {len(datasets)} AI Catalog items.",
+        content=(
+            f"Found {datasets_count} AI Catalog items, here are the details:\n"
+            f"{json.dumps(datasets_dict, indent=2)}"
+        ),
         structured_content={
-            "datasets": [{"id": ds.id, "name": ds.name} for ds in datasets],
-            "count": len(datasets),
+            "datasets": datasets_dict,
+            "count": datasets_count,
         },
     )
 

datarobot_genai/drmcp/tools/predictive/project.py

@@ -14,6 +14,10 @@
 
 import json
 import logging
+from typing import Annotated
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
 
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
@@ -21,35 +25,39 @@ from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
 logger = logging.getLogger(__name__)
 
 
-@dr_mcp_tool(tags={"project", "management", "list"})
-async def list_projects() -> str:
-    """
-    List all DataRobot projects for the authenticated user.
-
-    Returns
-    -------
-        A string summary of the user's DataRobot projects.
-    """
+@dr_mcp_tool(tags={"predictive", "project", "read", "management", "list"})
+async def list_projects() -> ToolResult:
+    """List all DataRobot projects for the authenticated user."""
     client = get_sdk_client()
     projects = client.Project.list()
-    if not projects:
-        return "No projects found."
-    return "\n".join(f"{p.id}: {p.project_name}" for p in projects)
+    projects = {p.id: p.project_name for p in projects}
 
+    return ToolResult(
+        content=(
+            json.dumps(projects, indent=2)
+            if projects
+            else json.dumps({"message": "No projects found."}, indent=2)
+        ),
+        structured_content=projects,
+    )
 
-@dr_mcp_tool(tags={"project", "data", "info"})
-async def get_project_dataset_by_name(project_id: str, dataset_name: str) -> str:
-    """
-    Get a dataset ID by name for a given project.
 
-    Args:
-        project_id: The ID of the DataRobot project.
-        dataset_name: The name of the dataset to find (e.g., 'training', 'holdout').
+@dr_mcp_tool(tags={"predictive", "project", "read", "data", "info"})
+async def get_project_dataset_by_name(
+    *,
+    project_id: Annotated[str, "The ID of the DataRobot project."] | None = None,
+    dataset_name: Annotated[str, "The name of the dataset to find (e.g., 'training', 'holdout')."]
+    | None = None,
+) -> ToolError | ToolResult:
+    """Get a dataset ID by name for a given project.
 
-    Returns
-    -------
-        The dataset ID and the dataset type (source or prediction) as a string, or an error message.
+    The dataset ID and the dataset type (source or prediction) as a string, or an error message.
     """
+    if not project_id:
+        return ToolError("Project ID is required.")
+    if not dataset_name:
+        return ToolError("Dataset name is required.")
+
     client = get_sdk_client()
     project = client.Project.get(project_id)
     all_datasets = []
@@ -61,12 +69,22 @@ async def get_project_dataset_by_name(project_id: str, dataset_name: str) -> str
         all_datasets.extend([{"type": "prediction", "dataset": ds} for ds in prediction_datasets])
     for ds in all_datasets:
         if dataset_name.lower() in ds["dataset"].name.lower():
-            return json.dumps(
-                {
+            return ToolResult(
+                content=(
+                    json.dumps(
+                        {
+                            "dataset_id": ds["dataset"].id,
+                            "dataset_type": ds["type"],
+                        },
+                        indent=2,
+                    )
+                ),
+                structured_content={
                     "dataset_id": ds["dataset"].id,
                     "dataset_type": ds["type"],
-                    "ui_panel": ["dataset"],
                 },
-                indent=2,
             )
-    return f"Dataset with name containing '{dataset_name}' not found in project {project_id}."
+    return ToolResult(
+        content=f"Dataset with name containing '{dataset_name}' not found in project {project_id}.",
+        structured_content={},
+    )

datarobot_genai/drmcp/tools/predictive/training.py

@@ -617,6 +617,7 @@ async def get_model_feature_impact(
 
 @dr_mcp_tool(tags={"predictive", "training", "read", "model", "evaluation"})
 async def get_model_lift_chart(
+    *,
     project_id: Annotated[str, "The ID of the DataRobot project"] | None = None,
     model_id: Annotated[str, "The ID of the model to analyze"] | None = None,
     source: Annotated[