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

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 = []
@@ -196,3 +187,144 @@ async def microsoft_graph_search_content(
             "count": n,
         },
     )
+
+
+@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",
+        "graph api",
+        "sharepoint",
+        "onedrive",
+        "document library",
+        "create",
+        "file",
+        "write",
+    }
+)
+async def microsoft_create_file(
+    *,
+    file_name: Annotated[str, "The name of the file to create (e.g., 'report.txt')."],
+    content_text: Annotated[str, "The raw text content to be stored in the file."],
+    document_library_id: Annotated[
+        str | None,
+        "The ID of the document library (Drive). If not provided, saves to personal OneDrive.",
+    ] = None,
+    parent_folder_id: Annotated[
+        str | None,
+        "ID of the parent folder. Defaults to 'root' (root of the drive).",
+    ] = "root",
+) -> ToolResult | ToolError:
+    """
+    Create a new text file in SharePoint or OneDrive.
+
+    **Personal OneDrive:** Just provide file_name and content_text.
+    The file saves to your personal OneDrive root folder.
+
+    **SharePoint:** Provide document_library_id to save to a specific
+    SharePoint site. Get the ID from microsoft_graph_search_content
+    results ('documentLibraryId' field).
+
+    **Conflict Resolution:** If a file with the same name exists,
+    it will be automatically renamed (e.g., 'report (1).txt').
+    """
+    if not file_name or not file_name.strip():
+        raise ToolError("Error: file_name is required.")
+    if not content_text:
+        raise ToolError("Error: content_text is required.")
+
+    access_token = await get_microsoft_graph_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    folder_id = parent_folder_id if parent_folder_id else "root"
+
+    async with MicrosoftGraphClient(access_token=access_token) as client:
+        # Auto-fetch personal OneDrive if no library specified
+        if document_library_id is None:
+            drive_id = await client.get_personal_drive_id()
+            is_personal_onedrive = True
+        else:
+            drive_id = document_library_id
+            is_personal_onedrive = False
+
+        created_file = await client.create_file(
+            drive_id=drive_id,
+            file_name=file_name.strip(),
+            content=content_text,
+            parent_folder_id=folder_id,
+            conflict_behavior="rename",
+        )
+
+    return ToolResult(
+        content=f"File '{created_file.name}' created successfully.",
+        structured_content={
+            "file_name": created_file.name,
+            "destination": "onedrive" if is_personal_onedrive else "sharepoint",
+            "driveId": drive_id,
+            "id": created_file.id,
+            "webUrl": created_file.web_url,
+            "parentFolderId": created_file.parent_folder_id,
+        },
+    )

datarobot_genai/drmcp/tools/predictive/deployment.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,71 +25,73 @@ from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
 logger = logging.getLogger(__name__)
 
 
-@dr_mcp_tool(tags={"deployment", "management", "list"})
-async def list_deployments() -> str:
-    """
-    List all DataRobot deployments for the authenticated user.
-
-    Returns
-    -------
-        A string summary of the user's DataRobot deployments.
-    """
+@dr_mcp_tool(tags={"predictive", "deployment", "read", "management", "list"})
+async def list_deployments() -> ToolResult:
+    """List all DataRobot deployments for the authenticated user."""
     client = get_sdk_client()
     deployments = client.Deployment.list()
     if not deployments:
-        logger.info("No deployments found")
-        return "No deployments found."
-    result = "\n".join(f"{d.id}: {d.label}" for d in deployments)
-    logger.info(f"Found {len(deployments)} deployments")
-    return result
-
+        return ToolResult(
+            content="No deployments found.",
+            structured_content={"deployments": []},
+        )
+    deployments_dict = {d.id: d.label for d in deployments}
+    return ToolResult(
+        content="\n".join(f"{d.id}: {d.label}" for d in deployments),
+        structured_content={"deployments": deployments_dict},
+    )
 
-@dr_mcp_tool(tags={"deployment", "model", "info"})
-async def get_model_info_from_deployment(deployment_id: str) -> str:
-    """
-    Get model info associated with a given deployment ID.
 
-    Args:
-        deployment_id: The ID of the DataRobot deployment.
+@dr_mcp_tool(tags={"predictive", "deployment", "read", "model", "info"})
+async def get_model_info_from_deployment(
+    *,
+    deployment_id: Annotated[str, "The ID of the DataRobot deployment"] | None = None,
+) -> ToolError | ToolResult:
+    """Retrieve model info associated with a given deployment ID."""
+    if not deployment_id:
+        raise ToolError("Deployment ID must be provided")
 
-    Returns
-    -------
-        The model info associated with the deployment as a JSON string.
-    """
     client = get_sdk_client()
     deployment = client.Deployment.get(deployment_id)
-    logger.info(f"Retrieved model info for deployment {deployment_id}")
-    return json.dumps(deployment.model, indent=2)
-
+    return ToolResult(
+        content=(
+            f"Retrieved model info for deployment {deployment_id}, here are the details:\n"
+            f"{json.dumps(deployment.model, indent=2)}"
+        ),
+        structured_content=deployment.model,
+    )
 
-@dr_mcp_tool(tags={"deployment", "model", "create"})
-async def deploy_model(model_id: str, label: str, description: str = "") -> str:
-    """
-    Deploy a model by creating a new DataRobot deployment.
 
-    Args:
-        model_id: The ID of the DataRobot model to deploy.
-        label: The label/name for the deployment.
-        description: Optional description for the deployment.
+@dr_mcp_tool(tags={"predictive", "deployment", "write", "model", "create"})
+async def deploy_model(
+    *,
+    model_id: Annotated[str, "The ID of the DataRobot model to deploy"] | None = None,
+    label: Annotated[str, "The label/name for the deployment"] | None = None,
+    description: Annotated[str, "Optional description for the deployment"] | None = None,
+) -> ToolError | ToolResult:
+    """Deploy a model by creating a new DataRobot deployment."""
+    if not model_id:
+        raise ToolError("Model ID must be provided")
+    if not label:
+        raise ToolError("Model label must be provided")
 
-    Returns
-    -------
-        JSON string with deployment ID and label, or error message.
-    """
     client = get_sdk_client()
     try:
         prediction_servers = client.PredictionServer.list()
         if not prediction_servers:
-            logger.error("No prediction servers available")
-            return json.dumps({"error": "No prediction servers available"})
+            raise ToolError("No prediction servers available for deployment.")
         deployment = client.Deployment.create_from_learning_model(
             model_id=model_id,
             label=label,
             description=description,
             default_prediction_server_id=prediction_servers[0].id,
         )
-        logger.info(f"Created deployment {deployment.id} with label {label}")
-        return json.dumps({"deployment_id": deployment.id, "label": label})
+        return ToolResult(
+            content=f"Created deployment {deployment.id} with label {label}",
+            structured_content={
+                "deployment_id": deployment.id,
+                "label": label,
+            },
+        )
     except Exception as e:
-        logger.error(f"Error deploying model {model_id}: {type(e).__name__}: {e}")
-        return json.dumps({"error": f"Error deploying model {model_id}: {type(e).__name__}: {e}"})
+        raise ToolError(f"Error deploying model {model_id}: {type(e).__name__}: {e}")

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,
+    )