datarobot-genai 0.2.37__py3-none-any.whl → 0.2.39__py3-none-any.whl

datarobot_genai/drmcp/tools/clients/microsoft_graph.py

@@ -16,6 +16,7 @@
 
 import logging
 from typing import Any
+from typing import Literal
 from urllib.parse import quote
 
 import httpx
@@ -541,6 +542,50 @@ class MicrosoftGraphClient:
 
         return MicrosoftGraphError(error_msg)
 
+    async def share_item(
+        self,
+        file_id: str,
+        document_library_id: str,
+        recipient_emails: list[str],
+        role: Literal["read", "write"],
+        send_invitation: bool,
+    ) -> None:
+        """
+        Share sharepoint / ondrive item using Microsoft Graph API.
+        Under the hood all resources in sharepoint/onedrive
+            in MS Graph API are treated as 'driveItem'.
+
+        Args:
+            file_id: The ID of the file or folder to share.
+            document_library_id: The ID of the document library containing the item.
+            recipient_emails: A list of email addresses to invite.
+            role: The role to assign.
+            send_invitation: Flag determining if recipients should be notified
+
+        Returns
+        -------
+            None
+
+        Raises
+        ------
+            MicrosoftGraphError: If sharing fails
+        """
+        graph_url = f"{GRAPH_API_BASE}/drives/{document_library_id}/items/{file_id}/invite"
+
+        payload = {
+            "recipients": [{"email": email} for email in recipient_emails],
+            "requireSignIn": True,
+            "sendInvitation": send_invitation,
+            "roles": [role],
+        }
+
+        response = await self._client.post(url=graph_url, json=payload)
+
+        if response.status_code not in (200, 201):
+            raise MicrosoftGraphError(
+                f"Microsoft Graph API error {response.status_code}: {response.text}"
+            )
+
     async def __aenter__(self) -> "MicrosoftGraphClient":
         """Async context manager entry."""
         return self

datarobot_genai/drmcp/tools/confluence/tools.py

@@ -55,26 +55,16 @@ async def confluence_get_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            if page_id_or_title.isdigit():
-                page_response = await client.get_page_by_id(page_id_or_title)
-            else:
-                if not space_key:
-                    raise ToolError(
-                        "Argument validation error: "
-                        "'space_key' is required when identifying a page by title."
-                    )
-                page_response = await client.get_page_by_title(page_id_or_title, space_key)
-    except ConfluenceError as e:
-        logger.error(f"Confluence error getting page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error getting Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while getting Confluence page "
-            f"'{page_id_or_title}': {str(e)}"
-        )
+    async with ConfluenceClient(access_token) as client:
+        if page_id_or_title.isdigit():
+            page_response = await client.get_page_by_id(page_id_or_title)
+        else:
+            if not space_key:
+                raise ToolError(
+                    "Argument validation error: "
+                    "'space_key' is required when identifying a page by title."
+                )
+            page_response = await client.get_page_by_title(page_id_or_title, space_key)
 
     return ToolResult(
         content=f"Successfully retrieved page '{page_response.title}'.",
@@ -116,22 +106,12 @@ async def confluence_create_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            page_response = await client.create_page(
-                space_key=space_key,
-                title=title,
-                body_content=body_content,
-                parent_id=parent_id,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error creating page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error creating Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while creating Confluence page "
-            f"'{title}' in space '{space_key}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        page_response = await client.create_page(
+            space_key=space_key,
+            title=title,
+            body_content=body_content,
+            parent_id=parent_id,
         )
 
     return ToolResult(
@@ -164,19 +144,10 @@ async def confluence_add_comment(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            comment_response = await client.add_comment(
-                page_id=page_id,
-                comment_body=comment_body,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error adding comment: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error adding comment to Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while adding comment to page '{page_id}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        comment_response = await client.add_comment(
+            page_id=page_id,
+            comment_body=comment_body,
         )
 
     return ToolResult(
@@ -220,32 +191,24 @@ async def confluence_search(
     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)}")
+    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]
 
     n = len(results)
     return ToolResult(
@@ -296,20 +259,11 @@ async def confluence_update_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            page_response = await client.update_page(
-                page_id=page_id,
-                new_body_content=new_body_content,
-                version_number=version_number,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error updating page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error updating Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while updating Confluence page '{page_id}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        page_response = await client.update_page(
+            page_id=page_id,
+            new_body_content=new_body_content,
+            version_number=version_number,
         )
 
     return ToolResult(

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -28,7 +28,6 @@ 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
 
 logger = logging.getLogger(__name__)
@@ -79,22 +78,15 @@ async def gdrive_find_contents(
     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,
-                folder_id=folder_id,
-                recursive=recursive,
-            )
-    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)}")
+    async with GoogleDriveClient(access_token) as client:
+        data = await client.list_files(
+            page_size=page_size,
+            page_token=page_token,
+            query=query,
+            limit=limit,
+            folder_id=folder_id,
+            recursive=recursive,
+        )
 
     filtered_fields = set(fields).intersection(SUPPORTED_FIELDS) if fields else SUPPORTED_FIELDS
     number_of_files = len(data.files)
@@ -155,17 +147,8 @@ async def gdrive_read_content(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with GoogleDriveClient(access_token) as client:
-            file_content = await client.read_file_content(file_id, target_format)
-    except GoogleDriveError as e:
-        logger.error(f"Google Drive error reading file content: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error reading Google Drive file content: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while reading Google Drive file content: {str(e)}"
-        )
+    async with GoogleDriveClient(access_token) as client:
+        file_content = await client.read_file_content(file_id, target_format)
 
     export_info = ""
     if file_content.was_exported:
@@ -239,20 +222,13 @@ async def gdrive_create_file(
     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)}")
+    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,
+        )
 
     file_type = "folder" if mime_type == GOOGLE_DRIVE_FOLDER_MIME else "file"
     content_info = ""
@@ -313,21 +289,12 @@ async def gdrive_update_metadata(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with GoogleDriveClient(access_token) as client:
-            updated_file = await client.update_file_metadata(
-                file_id=file_id,
-                new_name=new_name,
-                starred=starred,
-                trashed=trash,
-            )
-    except GoogleDriveError as e:
-        logger.error(f"Google Drive error updating file metadata: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error updating Google Drive file metadata: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while updating Google Drive file metadata: {str(e)}"
+    async with GoogleDriveClient(access_token) as client:
+        updated_file = await client.update_file_metadata(
+            file_id=file_id,
+            new_name=new_name,
+            starred=starred,
+            trashed=trash,
         )
 
     changes: list[str] = []
@@ -408,23 +375,14 @@ async def gdrive_manage_access(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with GoogleDriveClient(access_token) as client:
-            permission_id = await client.manage_access(
-                file_id=file_id,
-                action=action,
-                role=role,
-                email_address=email_address,
-                permission_id=permission_id,
-                transfer_ownership=transfer_ownership,
-            )
-    except GoogleDriveError as e:
-        logger.error(f"Google Drive permission operation failed: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error changing permissions for Google Drive file {file_id}: {e}")
-        raise ToolError(
-            f"Unexpected error changing permissions for Google Drive file {file_id}: {str(e)}"
+    async with GoogleDriveClient(access_token) as client:
+        permission_id = await client.manage_access(
+            file_id=file_id,
+            action=action,
+            role=role,
+            email_address=email_address,
+            permission_id=permission_id,
+            transfer_ownership=transfer_ownership,
         )
 
     # Build response

datarobot_genai/drmcp/tools/jira/tools.py

@@ -72,14 +72,8 @@ async def jira_get_issue(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with JiraClient(access_token) as client:
-            issue = await client.get_jira_issue(issue_key)
-    except Exception as e:
-        logger.error(f"Unexpected error while getting Jira issue: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while getting Jira issue '{issue_key}': {str(e)}"
-        )
+    async with JiraClient(access_token) as client:
+        issue = await client.get_jira_issue(issue_key)
 
     return ToolResult(
         content=f"Successfully retrieved details for issue '{issue_key}'.",
@@ -118,17 +112,13 @@ async def jira_create_issue(
                 f"Unexpected issue type `{issue_type}`. Possible values are {possible_issue_types}."
             )
 
-    try:
-        async with JiraClient(access_token) as client:
-            issue_key = await client.create_jira_issue(
-                project_key=project_key,
-                summary=summary,
-                issue_type_id=issue_type_id,
-                description=description,
-            )
-    except Exception as e:
-        logger.error(f"Unexpected error while creating Jira issue: {e}")
-        raise ToolError(f"An unexpected error occurred while creating Jira issue: {str(e)}")
+    async with JiraClient(access_token) as client:
+        issue_key = await client.create_jira_issue(
+            project_key=project_key,
+            summary=summary,
+            issue_type_id=issue_type_id,
+            description=description,
+        )
 
     return ToolResult(
         content=f"Successfully created issue '{issue_key}'.",
@@ -179,14 +169,10 @@ async def jira_update_issue(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with JiraClient(access_token) as client:
-            updated_fields = await client.update_jira_issue(
-                issue_key=issue_key, fields=fields_to_update
-            )
-    except Exception as e:
-        logger.error(f"Unexpected error while updating Jira issue: {e}")
-        raise ToolError(f"An unexpected error occurred while updating Jira issue: {str(e)}")
+    async with JiraClient(access_token) as client:
+        updated_fields = await client.update_jira_issue(
+            issue_key=issue_key, fields=fields_to_update
+        )
 
     updated_fields_str = ",".join(updated_fields)
     return ToolResult(
@@ -226,12 +212,8 @@ async def jira_transition_issue(
                 f"Possible values are {available_transitions_str}."
             )
 
-    try:
-        async with JiraClient(access_token) as client:
-            await client.transition_jira_issue(issue_key=issue_key, transition_id=transition_id)
-    except Exception as e:
-        logger.error(f"Unexpected error while transitioning Jira issue: {e}")
-        raise ToolError(f"An unexpected error occurred while transitioning Jira issue: {str(e)}")
+    async with JiraClient(access_token) as client:
+        await client.transition_jira_issue(issue_key=issue_key, transition_id=transition_id)
 
     return ToolResult(
         content=f"Successfully transitioned issue '{issue_key}' to status '{transition_name}'.",

datarobot_genai/drmcp/tools/microsoft_graph/tools.py

@@ -16,13 +16,13 @@
 
 import logging
 from typing import Annotated
+from typing import Literal
 
 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
 
@@ -142,25 +142,16 @@ async def microsoft_graph_search_content(
     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)}"
+    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,
         )
 
     results = []
@@ -198,6 +189,68 @@ async def microsoft_graph_search_content(
     )
 
 
+@dr_mcp_tool(tags={"microsoft", "graph api", "sharepoint", "onedrive", "share"}, enabled=False)
+async def microsoft_graph_share_item(
+    *,
+    file_id: Annotated[str, "The ID of the file or folder to share."],
+    document_library_id: Annotated[str, "The ID of the document library containing the item."],
+    recipient_emails: Annotated[list[str], "A list of email addresses to invite."],
+    role: Annotated[Literal["read", "write"], "The role to assign: 'read' or 'write'."] = "read",
+    send_invitation: Annotated[
+        bool, "Flag determining if recipients should be notified. Default False"
+    ] = False,
+) -> ToolResult | ToolError:
+    """
+    Share a SharePoint or Onedrive file or folder with one or more users.
+    It works with internal users or existing guest users in the
+    tenant. It does NOT create new guest accounts and does NOT use the tenant-level
+    /invitations endpoint.
+
+    Microsoft Graph API is treating OneDrive and SharePoint resources as driveItem.
+
+    API Reference:
+    - DriveItem Resource Type: https://learn.microsoft.com/en-us/graph/api/resources/driveitem
+    - API Documentation: https://learn.microsoft.com/en-us/graph/api/driveitem-invite
+    """
+    if not file_id or not file_id.strip():
+        raise ToolError("Argument validation error: 'file_id' cannot be empty.")
+
+    if not document_library_id or not document_library_id.strip():
+        raise ToolError("Argument validation error: 'document_library_id' cannot be empty.")
+
+    if not recipient_emails:
+        raise ToolError("Argument validation error: you must provide at least one 'recipient'.")
+
+    access_token = await get_microsoft_graph_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    async with MicrosoftGraphClient(access_token=access_token) as client:
+        await client.share_item(
+            file_id=file_id,
+            document_library_id=document_library_id,
+            recipient_emails=recipient_emails,
+            role=role,
+            send_invitation=send_invitation,
+        )
+
+    n = len(recipient_emails)
+    return ToolResult(
+        content=(
+            f"Successfully shared file {file_id} "
+            f"from document library {document_library_id} "
+            f"with {n} recipients with '{role}' role."
+        ),
+        structured_content={
+            "fileId": file_id,
+            "documentLibraryId": document_library_id,
+            "recipientEmails": recipient_emails,
+            "n": n,
+            "role": role,
+        },
+    )
+
+
 @dr_mcp_tool(
     tags={
         "microsoft",

datarobot_genai/drmcp/tools/predictive/deployment_info.py

@@ -19,9 +19,13 @@ import json
 import logging
 from datetime import datetime
 from datetime import timedelta
+from typing import Annotated
 from typing import Any
 
 import pandas as pd
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
+from mcp.types import TextContent
 
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
@@ -29,40 +33,18 @@ from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
 logger = logging.getLogger(__name__)
 
 
-@dr_mcp_tool(tags={"deployment", "info", "metadata"})
-async def get_deployment_info(deployment_id: str) -> str:
+@dr_mcp_tool(tags={"predictive", "deployment", "read", "info", "metadata"})
+async def get_deployment_info(
+    *,
+    deployment_id: Annotated[str, "The ID of the DataRobot deployment"] | None = None,
+) -> ToolError | ToolResult:
     """
     Retrieve information about the deployment, including the list of
     features needed to make predictions on this deployment.
-
-    Args:
-        deployment_id: The ID of the DataRobot deployment
-
-    Returns
-    -------
-        JSON string containing model and feature information including:
-        For datarobot native models will return model information for custom models
-        this will likely just return features and total_features values.
-
-        - model_type: Type of model
-        - target: Name of the target feature
-        - target_type: Type of the target feature
-        - features: List of features with their importance and type
-        - total_features: Total number of features
-        - time_series_config: Time series configuration if applicable
-
-            for features:
-            - feature_name: Name of the feature
-            - ``name`` : str, feature name
-            - ``feature_type`` : str, feature type
-            - ``importance`` : float, numeric measure of the relationship strength between
-                the feature and target (independent of model or other features)
-            - ``date_format`` : str or None, the date format string for how this feature was
-                interpreted, null if not a date feature, compatible with
-                https://docs.python.org/2/library/time.html#time.strftime.
-            - ``known_in_advance`` : bool, whether the feature was selected as known in advance in
-                a time series model, false for non-time series models.
     """
+    if not deployment_id:
+        raise ToolError("Deployment ID must be provided")
+
     client = get_sdk_client()
     deployment = client.Deployment.get(deployment_id)
 
@@ -112,40 +94,34 @@ async def get_deployment_info(deployment_id: str) -> str:
             "series_id_columns": partition.multiseries_id_columns or [],
         }
 
-    return json.dumps(result, indent=2)
+    return ToolResult(
+        content=json.dumps(result, indent=2),
+        structured_content=result,
+    )
 
 
-@dr_mcp_tool(tags={"deployment", "template", "data"})
-async def generate_prediction_data_template(deployment_id: str, n_rows: int = 1) -> str:
-    """
-    Generate a template CSV with the correct structure for making predictions.
-
-    This creates a template with:
-    - All required feature columns in the correct order
-    - Sample values based on feature types
-    - Comments explaining each feature
-    - When using this tool, always consider feature importance. For features with high importance,
-      try to infer or ask for a reasonable value, using frequent values or domain knowledge if
-      available. For less important features, you may leave them blank.
-    - If frequent values are available for a feature, they will be used as sample values;
-      otherwise, blank fields will be used.
-      Please note that using frequent values in your predictions data can influence the prediction,
-      think of it as sending in the average value for the feature. If you don't want this effect on
-      your predictions leave the field blank you in predictions dataset.
-
-    Args:
-        deployment_id: The ID of the DataRobot deployment
-        n_rows: Number of template rows to generate (default 1)
-
-    Returns
-    -------
-        CSV template string with sample data ready for predictions
-    """
+@dr_mcp_tool(tags={"predictive", "deployment", "read", "template", "data"})
+async def generate_prediction_data_template(
+    *,
+    deployment_id: Annotated[str, "The ID of the DataRobot deployment"] | None = None,
+    n_rows: Annotated[int, "Number of template rows to generate"] = 1,
+) -> ToolError | ToolResult:
+    """Generate a template CSV with the correct structure for making predictions."""
+    if not deployment_id:
+        raise ToolError("Deployment ID must be provided")
+    if n_rows is None or n_rows <= 0:
+        n_rows = 1
+
     # Get feature information
-    features_json = await get_deployment_features(deployment_id)
+    features_result = await get_deployment_features(deployment_id=deployment_id)
     # Add error handling for empty or error responses
+    # Extract text content from ToolResult
+    if features_result.content and isinstance(features_result.content[0], TextContent):
+        features_json = features_result.content[0].text
+    else:
+        features_json = str(features_result.content)
     if not features_json or features_json.strip().startswith("Error"):
-        return f"Error: {features_json}"
+        raise ToolError(f"Error with feature information: {features_json}")
     features_info = json.loads(features_json)
 
     # Create template data
@@ -218,49 +194,55 @@ async def generate_prediction_data_template(deployment_id: str, n_rows: int = 1)
     result += f"# Total Features: {features_info['total_features']}\n"
     result += df.to_csv(index=False)
 
-    return str(result)
+    # Build structured content with template data and metadata
+    structured_content = {
+        "deployment_id": deployment_id,
+        "model_type": features_info["model_type"],
+        "target": features_info["target"],
+        "target_type": features_info["target_type"],
+        "total_features": features_info["total_features"],
+        "template_data": df.to_dict("records"),  # Convert DataFrame to list of dicts
+    }
+
+    if "time_series_config" in features_info:
+        structured_content["time_series_config"] = features_info["time_series_config"]
+
+    return ToolResult(
+        content=str(result),
+        structured_content=structured_content,
+    )
 
 
-@dr_mcp_tool(tags={"deployment", "validation", "data"})
+@dr_mcp_tool(tags={"predictive", "deployment", "read", "validation", "data"})
 async def validate_prediction_data(
-    deployment_id: str,
-    file_path: str | None = None,
-    csv_string: str | None = None,
-) -> str:
-    """
-    Validate if a CSV file is suitable for making predictions with a deployment.
-
-    Checks:
-    - All required features are present
-    - Feature types match expectations
-    - Missing values (null, empty string, or blank fields) are allowed and will not cause errors
-    - No critical issues that would prevent predictions
-
-    Args:
-        deployment_id: The ID of the DataRobot deployment
-        file_path: Path to the CSV file to validate (optional if csv_string is provided)
-        csv_string: CSV data as a string (optional, used if file_path is not provided)
-
-    Returns
-    -------
-        Validation report including any errors, warnings, and suggestions
-    """
+    *,
+    deployment_id: Annotated[str, "The ID of the DataRobot deployment"] | None = None,
+    file_path: Annotated[
+        str, "Path to the CSV file to validate (optional if csv_string is provided)"
+    ]
+    | None = None,
+    csv_string: Annotated[str, "CSV data as a string (optional, used if file_path is not provided)"]
+    | None = None,
+) -> ToolError | ToolResult:
+    """Validate if a CSV file is suitable for making predictions with a deployment."""
     # Load the data
     if csv_string is not None:
         df = pd.read_csv(io.StringIO(csv_string))
     elif file_path is not None:
         df = pd.read_csv(file_path)
     else:
-        return json.dumps(
-            {
-                "status": "error",
-                "error": "Must provide either file_path or csv_string.",
-            },
-            indent=2,
-        )
+        raise ToolError("Must provide either file_path or csv_string.")
+
+    if not deployment_id:
+        raise ToolError("Deployment ID must be provided")
 
     # Get deployment features
-    features_json = await get_deployment_features(deployment_id)
+    features_result = await get_deployment_features(deployment_id=deployment_id)
+    # Extract text content from ToolResult
+    if features_result.content and isinstance(features_result.content[0], TextContent):
+        features_json = features_result.content[0].text
+    else:
+        features_json = str(features_result.content)
     features_info = json.loads(features_json)
 
     validation_report: dict[str, Any] = {
@@ -359,22 +341,29 @@ async def validate_prediction_data(
         "model_type": features_info["model_type"],
     }
 
-    return json.dumps(validation_report, indent=2)
-
-
-@dr_mcp_tool(tags={"deployment", "features", "info"})
-async def get_deployment_features(deployment_id: str) -> str:
-    """
-    Retrieve only the features list for a deployment, as JSON string.
-    Args:
-        deployment_id: The ID of the DataRobot deployment
-    Returns:
-        JSON string containing only the features list and time series config if present.
-    """
-    info_json = await get_deployment_info(deployment_id)
+    return ToolResult(
+        content=json.dumps(validation_report, indent=2),
+        structured_content=validation_report,
+    )
+
+
+@dr_mcp_tool(tags={"predictive", "deployment", "read", "features", "info"})
+async def get_deployment_features(
+    *,
+    deployment_id: Annotated[str, "The ID of the DataRobot deployment"] | None = None,
+) -> ToolError | ToolResult:
+    """Retrieve only the features list for a deployment, as JSON string."""
+    if not deployment_id:
+        raise ToolError("Deployment ID must be provided")
+
+    info_result = await get_deployment_info(deployment_id=deployment_id)
+    # Extract text content from ToolResult
+    if info_result.content and isinstance(info_result.content[0], TextContent):
+        info_json = info_result.content[0].text
+    else:
+        info_json = str(info_result.content)
     if not info_json.strip().startswith("{"):
-        # Return a default error JSON
-        return json.dumps({"features": [], "total_features": 0, "error": info_json}, indent=2)
+        raise ToolError(f"Error with deployment info: {info_json}")
     info = json.loads(info_json)
     # Only keep features, time_series_config, and total_features
     result = {
@@ -389,4 +378,8 @@ async def get_deployment_features(deployment_id: str) -> str:
         result["target"] = info["target"]
     if "target_type" in info:
         result["target_type"] = info["target_type"]
-    return json.dumps(result, indent=2)
+
+    return ToolResult(
+        content=json.dumps(result, indent=2),
+        structured_content=result,
+    )

{datarobot_genai-0.2.37.dist-info → datarobot_genai-0.2.39.dist-info}/METADATA

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

{datarobot_genai-0.2.37.dist-info → datarobot_genai-0.2.39.dist-info}/RECORD

@@ -84,20 +84,20 @@ datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCE
 datarobot_genai/drmcp/tools/clients/confluence.py,sha256=h_G0By_kDnJeWDT_d-IREsaZ5-0xB5GoLXOqblYP5MA,20706
 datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=RK4IISpYb99aK6WgDthesDoglaZxwGpG_PPAAe6xsVM,33064
 datarobot_genai/drmcp/tools/clients/jira.py,sha256=Rm91JAyrNIqxu66-9rU1YqoRXVnWbEy-Ahvy6f6HlVg,9823
-datarobot_genai/drmcp/tools/clients/microsoft_graph.py,sha256=-g0EhaBVElKbujaO2cHdgc86hwFEkkyEyZVAw8pq7yM,24468
+datarobot_genai/drmcp/tools/clients/microsoft_graph.py,sha256=xh3J3QvBpsk-jO5tPIikn2AR1q5Fp3Z3_zdfb-C6UKE,25970
 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=_-ws65WLK8KZP_mKkf4yJ7ZunR8qdyoiMwHQX47MSMw,12362
+datarobot_genai/drmcp/tools/confluence/tools.py,sha256=tbZxpSkMqFWSz8HxCnjFuJ0JL06RD5t-B1alxxKnMl4,10314
 datarobot_genai/drmcp/tools/gdrive/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=7bNrp7E3opKwsBDYfLIOsOGfPXW-Ae9KvcimEzetR0A,17631
+datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=38XxGXQHZcoX06HQDFf0-BFmpNxNG4wdvijaoqIM4ho,15572
 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/jira/tools.py,sha256=qEQyyri2bPYpphDR9SVxon1_gTmBO2bPuP9nOLLYdk0,8775
 datarobot_genai/drmcp/tools/microsoft_graph/__init__.py,sha256=CuOaMt1AJo7cHx_GuhO3s_aqxZas_wlDsoBorBsvbeU,577
-datarobot_genai/drmcp/tools/microsoft_graph/tools.py,sha256=cNctozv_4lRC5Kva3D2j4taZfeQHDE6LTAjcmeQXwWA,10446
+datarobot_genai/drmcp/tools/microsoft_graph/tools.py,sha256=iQnmj-AEIVpGBVjlg2ogwDeHdQlTexdTQzX-DYQbFfQ,12479
 datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
 datarobot_genai/drmcp/tools/predictive/data.py,sha256=VbGs8ERP8vNFtTTryGhI61JItNVaJsx1gxpRX1ZFZcg,4626
 datarobot_genai/drmcp/tools/predictive/deployment.py,sha256=Pc6lz9V2JOw3Ufw-SsGAhMKf6-YhvbjGoNLRFOIcSSY,3670
-datarobot_genai/drmcp/tools/predictive/deployment_info.py,sha256=BGEF_dmbxOBJR0n1Tt9TO2-iNTQSBTr-oQUyaxLZ0ZI,15297
+datarobot_genai/drmcp/tools/predictive/deployment_info.py,sha256=I9YxznndDOq0H2QgIjkb5O5sX8S6GTbYmXtRRYjEzOw,14778
 datarobot_genai/drmcp/tools/predictive/model.py,sha256=BVxOMHh3--liwBU4VB1OWRrqkhJ4y_Rq053f7y94TF8,6276
 datarobot_genai/drmcp/tools/predictive/predict.py,sha256=Qoob2_t2crfWtyPzkXMRz2ITZumnczU6Dq4C7q9RBMI,9370
 datarobot_genai/drmcp/tools/predictive/predict_realtime.py,sha256=urq6rPyZFsAP-bPyclSNzrkvb6FTamdlFau8q0IWWJ0,13472
@@ -117,9 +117,9 @@ datarobot_genai/nat/datarobot_llm_clients.py,sha256=-_q_KlKOVQecIYJd8YRiYnS4ZNaz
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
 datarobot_genai/nat/datarobot_mcp_client.py,sha256=jL8sXb8g4gvt0VYgB2tfMGsMjpB1GV2XIbN0iv_LxVU,10701
 datarobot_genai/nat/helpers.py,sha256=Q7E3ADZdtFfS8E6OQPyw2wgA6laQ58N3bhLj5CBWwJs,3265
-datarobot_genai-0.2.37.dist-info/METADATA,sha256=JFxYZKfbfrbePIywe1WQ1ZfjQ5W9HAq_GNYAuXxTcB8,6301
-datarobot_genai-0.2.37.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.37.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.37.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.37.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.37.dist-info/RECORD,,
+datarobot_genai-0.2.39.dist-info/METADATA,sha256=V8b_XudaiugsXFgEhkrtEpEr0UaSHVHLL6xPV578UGc,6301
+datarobot_genai-0.2.39.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.39.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.39.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.39.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.39.dist-info/RECORD,,