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

datarobot_genai/drmcp/tools/confluence/tools.py

@@ -252,3 +252,70 @@ async def confluence_search(
         content=f"Successfully executed CQL query and retrieved {n} result(s).",
         structured_content={"data": data, "count": n},
     )
+
+
+@dr_mcp_tool(tags={"confluence", "write", "update", "page"})
+async def confluence_update_page(
+    *,
+    page_id: Annotated[str, "The ID of the Confluence page to update."],
+    new_body_content: Annotated[
+        str,
+        "The full updated content of the page in Confluence Storage Format (XML) or raw text.",
+    ],
+    version_number: Annotated[
+        int,
+        "The current version number of the page, required to prevent update conflicts. "
+        "Get this from the confluence_get_page tool.",
+    ],
+) -> ToolResult:
+    """Update the content of an existing Confluence page.
+
+    Requires the current version number to ensure atomic updates.
+    Use this tool to update the body content of an existing Confluence page.
+    The version_number is required for optimistic locking - it prevents overwriting
+    changes made by others since you last fetched the page.
+
+    Usage:
+        page_id="856391684", new_body_content="<p>New content</p>", version_number=5
+
+    Important: Always fetch the page first using confluence_get_page to get the
+    current version number before updating.
+    """
+    if not page_id:
+        raise ToolError("Argument validation error: 'page_id' cannot be empty.")
+
+    if not new_body_content:
+        raise ToolError("Argument validation error: 'new_body_content' cannot be empty.")
+
+    if version_number < 1:
+        raise ToolError(
+            "Argument validation error: 'version_number' must be a positive integer (>= 1)."
+        )
+
+    access_token = await get_atlassian_access_token()
+    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)}"
+        )
+
+    return ToolResult(
+        content=f"Page ID {page_id} updated successfully to version {page_response.version}.",
+        structured_content={
+            "updated_page_id": page_response.page_id,
+            "new_version": page_response.version,
+        },
+    )

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -109,3 +109,69 @@ async def gdrive_find_contents(
             "nextPageToken": data.next_page_token,
         },
     )
+
+
+@dr_mcp_tool(tags={"google", "gdrive", "read", "content", "file", "download"})
+async def gdrive_read_content(
+    *,
+    file_id: Annotated[str, "The ID of the file to read."],
+    target_format: Annotated[
+        str | None,
+        "The preferred output format for Google Workspace files "
+        "(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:
+    """
+    Retrieve the content of a specific file by its ID. Google Workspace files are
+    automatically exported to LLM-readable formats (Push-Down).
+
+    Usage:
+        - Basic: gdrive_read_content(file_id="1ABC123def456")
+        - Custom format: gdrive_read_content(file_id="1ABC...", target_format="text/plain")
+        - First use gdrive_find_contents to discover file IDs
+
+    Supported conversions (defaults):
+        - Google Docs -> Markdown (text/markdown)
+        - Google Sheets -> CSV (text/csv)
+        - Google Slides -> Plain text (text/plain)
+        - PDF files -> Extracted text (text/plain)
+        - Other text files -> Downloaded as-is
+
+    Note: Binary files (images, videos, etc.) are not supported and will return an error.
+    Large Google Workspace files (>10MB) may fail to export due to API limits.
+
+    Refer to Google Drive export formats documentation:
+    https://developers.google.com/workspace/drive/api/guides/ref-export-formats
+    """
+    if not file_id or not file_id.strip():
+        raise ToolError("Argument validation error: 'file_id' 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:
+            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)}"
+        )
+
+    # Provide helpful context about the conversion
+    export_info = ""
+    if file_content.was_exported:
+        export_info = f" (exported from {file_content.original_mime_type})"
+
+    return ToolResult(
+        content=(
+            f"Successfully retrieved content of '{file_content.name}' "
+            f"({file_content.mime_type}){export_info}."
+        ),
+        structured_content=file_content.as_flat_dict(),
+    )

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

@@ -18,8 +18,11 @@ import json
 import logging
 from dataclasses import asdict
 from dataclasses import dataclass
+from typing import Annotated
 
 import pandas as pd
+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
@@ -53,22 +56,15 @@ class DatasetInsight:
     missing_data_summary: dict[str, float]
 
 
-@dr_mcp_tool(tags={"training", "analysis", "dataset"})
-async def analyze_dataset(dataset_id: str) -> str:
-    """
-    Analyze a dataset to understand its structure and potential use cases.
+@dr_mcp_tool(tags={"predictive", "training", "read", "analysis", "dataset"})
+async def analyze_dataset(
+    *,
+    dataset_id: Annotated[str, "The ID of the DataRobot dataset to analyze"] | None = None,
+) -> ToolError | ToolResult:
+    """Analyze a dataset to understand its structure and potential use cases."""
+    if not dataset_id:
+        return ToolError("Dataset ID must be provided")
 
-    Args:
-        dataset_id: The ID of the DataRobot dataset to analyze
-
-    Returns
-    -------
-        JSON string containing dataset insights including:
-        - Basic statistics (rows, columns)
-        - Column types (numerical, categorical, datetime, text)
-        - Potential target columns
-        - Missing data summary
-    """
     client = get_sdk_client()
     dataset = client.Dataset.get(dataset_id)
     df = dataset.get_as_dataframe()
@@ -105,27 +101,23 @@ async def analyze_dataset(dataset_id: str) -> str:
         potential_targets=potential_targets,
         missing_data_summary=missing_data,
     )
+    insights_dict = asdict(insights)
 
-    return json.dumps(asdict(insights), indent=2)
-
+    return ToolResult(
+        content=json.dumps(insights_dict, indent=2),
+        structured_content=insights_dict,
+    )
 
-@dr_mcp_tool(tags={"training", "analysis", "usecase"})
-async def suggest_use_cases(dataset_id: str) -> str:
-    """
-    Analyze a dataset and suggest potential machine learning use cases.
 
-    Args:
-        dataset_id: The ID of the DataRobot dataset to analyze
+@dr_mcp_tool(tags={"predictive", "training", "read", "analysis", "usecase"})
+async def suggest_use_cases(
+    *,
+    dataset_id: Annotated[str, "The ID of the DataRobot dataset to analyze"] | None = None,
+) -> ToolError | ToolResult:
+    """Analyze a dataset and suggest potential machine learning use cases."""
+    if not dataset_id:
+        return ToolError("Dataset ID must be provided")
 
-    Returns
-    -------
-        JSON string containing suggested use cases with:
-        - Use case name and description
-        - Suggested target column
-        - Problem type
-        - Confidence score
-        - Reasoning for the suggestion
-    """
     client = get_sdk_client()
     dataset = client.Dataset.get(dataset_id)
     df = dataset.get_as_dataframe()
@@ -141,27 +133,23 @@ async def suggest_use_cases(dataset_id: str) -> str:
 
     # Sort by confidence score
     suggestions.sort(key=lambda x: x["confidence"], reverse=True)
-    return json.dumps(suggestions, indent=2)
-
-
-@dr_mcp_tool(tags={"training", "analysis", "eda"})
-async def get_exploratory_insights(dataset_id: str, target_col: str | None = None) -> str:
-    """
-    Generate exploratory data insights for a dataset.
-
-    Args:
-        dataset_id: The ID of the DataRobot dataset to analyze
-        target_col: Optional target column to focus EDA insights on
-
-    Returns
-    -------
-        JSON string containing EDA insights including:
-        - Dataset summary statistics
-        - Target variable analysis (if specified)
-        - Feature correlations with target
-        - Missing data analysis
-        - Data type distribution
-    """
+
+    return ToolResult(
+        content=json.dumps(suggestions, indent=2),
+        structured_content={"use_case_suggestions": suggestions},
+    )
+
+
+@dr_mcp_tool(tags={"predictive", "training", "read", "analysis", "eda"})
+async def get_exploratory_insights(
+    *,
+    dataset_id: Annotated[str, "The ID of the DataRobot dataset to analyze"] | None = None,
+    target_col: Annotated[str, "Optional target column to focus EDA insights on"] | None = None,
+) -> ToolError | ToolResult:
+    """Generate exploratory data insights for a dataset."""
+    if not dataset_id:
+        return ToolError("Dataset ID must be provided")
+
     client = get_sdk_client()
     dataset = client.Dataset.get(dataset_id)
     df = dataset.get_as_dataframe()
@@ -238,8 +226,10 @@ async def get_exploratory_insights(dataset_id: str, target_col: str | None = Non
                     sorted(correlations.items(), key=lambda x: abs(x[1]), reverse=True)
                 )
 
-    eda_insights["ui_panel"] = ["eda"]
-    return json.dumps(eda_insights, indent=2)
+    return ToolResult(
+        content=json.dumps(eda_insights, indent=2),
+        structured_content=eda_insights,
+    )
 
 
 def _identify_potential_targets(
@@ -450,47 +440,50 @@ def _analyze_target_for_use_cases(df: pd.DataFrame, target_col: str) -> list[Use
     return suggestions
 
 
-@dr_mcp_tool(tags={"training", "autopilot", "model"})
+@dr_mcp_tool(tags={"predictive", "training", "write", "autopilot", "model"})
 async def start_autopilot(
-    target: str,
-    project_id: str | None = None,
-    mode: str | None = "quick",
-    dataset_url: str | None = None,
-    dataset_id: str | None = None,
-    project_name: str | None = "MCP Project",
-    use_case_id: str | None = None,
-) -> str:
-    """
-    Start automated model training (Autopilot) for a project.
-
-    Args:
-        target: Name of the target column for modeling.
-        project_id: Optional, the ID of the DataRobot project or a new project if no id is provided.
-        mode: Optional, Autopilot mode ('quick', 'comprehensive', or 'manual')
-        dataset_url: Optional, The URL to the dataset to upload (optional if dataset_id is provided)
-            for a new project.
-        dataset_id: Optional, The ID of an existing dataset in AI Catalog (optional if dataset_url
-            is provided) for a new project.
-        project_name: Optional, name for the project if no id is provided, creates a new project.
-        use_case_id: Optional, ID of the use case to associate this project with (required for
-            next-gen platform).
-
-    Returns
-    -------
-        JSON string containing:
-        - project_id: Project ID
-        - target: Target column name
-        - mode: Selected Autopilot mode
-        - status: Current project status
-        - ui_panel: List of recommended UI panels for visualization
-    """
+    *,
+    target: Annotated[str, "Name of the target column for modeling"] | None = None,
+    project_id: Annotated[
+        str, "Optional, the ID of the DataRobot project or a new project if no id is provided"
+    ]
+    | None = None,
+    mode: Annotated[str, "Optional, Autopilot mode ('quick', 'comprehensive', or 'manual')"]
+    | None = "quick",
+    dataset_url: Annotated[
+        str,
+        """
+        Optional, The URL to the dataset to upload
+        (optional if dataset_id is provided) for a new project.
+        """,
+    ]
+    | None = None,
+    dataset_id: Annotated[
+        str,
+        """
+        Optional, The ID of an existing dataset in AI Catalog
+        (optional if dataset_url is provided) for a new project.
+        """,
+    ]
+    | None = None,
+    project_name: Annotated[
+        str, "Optional, name for the project if no id is provided, creates a new project"
+    ]
+    | None = "MCP Project",
+    use_case_id: Annotated[
+        str,
+        "Optional, ID of the use case to associate this project (required for next-gen platform)",
+    ]
+    | None = None,
+) -> ToolError | ToolResult:
+    """Start automated model training (Autopilot) for a project."""
     client = get_sdk_client()
 
     if not project_id:
         if not dataset_url and not dataset_id:
-            return "Error: Either dataset_url or dataset_id must be provided"
+            return ToolError("Either dataset_url or dataset_id must be provided")
         if dataset_url and dataset_id:
-            return "Error: Please provide either dataset_url or dataset_id, not both"
+            return ToolError("Please provide either dataset_url or dataset_id, not both")
 
         if dataset_url:
             dataset = client.Dataset.create_from_url(dataset_url)
@@ -504,7 +497,7 @@ async def start_autopilot(
         project = client.Project.get(project_id)
 
     if not target:
-        return "Error: Target variable must be specified"
+        return ToolError("Target variable must be specified")
 
     try:
         # Start modeling
@@ -515,40 +508,48 @@ async def start_autopilot(
             "target": target,
             "mode": mode,
             "status": project.get_status(),
-            "ui_panel": ["eda", "model-training", "leaderboard"],
             "use_case_id": project.use_case_id,
         }
 
-        return json.dumps(result, indent=2)
-    except Exception as e:
-        return json.dumps(
-            {
-                "error": f"Failed to start Autopilot: {str(e)}",
-                "project_id": project.id,
-                "target": target,
-                "mode": mode,
-            },
-            indent=2,
+        return ToolResult(
+            content=json.dumps(result, indent=2),
+            structured_content=result,
         )
 
+    except Exception as e:
+        return ToolError(
+            content=json.dumps(
+                {
+                    "error": f"Failed to start Autopilot: {str(e)}",
+                    "project_id": project.id if project else None,
+                    "target": target,
+                    "mode": mode,
+                },
+                indent=2,
+            )
+        )
 
-@dr_mcp_tool(tags={"training", "model", "evaluation"})
-async def get_model_roc_curve(project_id: str, model_id: str, source: str = "validation") -> str:
-    """
-    Get detailed ROC curve for a specific model.
 
-    Args:
-        project_id: The ID of the DataRobot project
-        model_id: The ID of the model to analyze
-        source: The source of the data to use for the ROC curve ('validation' or 'holdout' or
-            'crossValidation')
+@dr_mcp_tool(tags={"prediction", "training", "read", "model", "evaluation"})
+async def get_model_roc_curve(
+    *,
+    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[
+        str,
+        """
+        The source of the data to use for the ROC curve
+        ('validation' or 'holdout' or 'crossValidation')
+        """,
+    ]
+    | str = "validation",
+) -> ToolError | ToolResult:
+    """Get detailed ROC curve for a specific model."""
+    if not project_id:
+        return ToolError("Project ID must be provided")
+    if not model_id:
+        return ToolError("Model ID must be provided")
 
-    Returns
-    -------
-        JSON string containing:
-        - roc_curve: ROC curve data
-        - ui_panel: List of recommended UI panels for visualization
-    """
     client = get_sdk_client()
     project = client.Project.get(project_id)
     model = client.Model.get(project=project, model_id=model_id)
@@ -581,26 +582,26 @@ async def get_model_roc_curve(project_id: str, model_id: str, source: str = "val
             "source": source,
         }
 
-        return json.dumps({"data": roc_data, "ui_panel": ["roc-curve"]}, indent=2)
+        return ToolResult(
+            content=json.dumps({"data": roc_data}, indent=2),
+            structured_content={"data": roc_data},
+        )
     except Exception as e:
-        return json.dumps({"error": f"Failed to get ROC curve: {str(e)}"}, indent=2)
+        return ToolError(f"Failed to get ROC curve: {str(e)}")
 
 
-@dr_mcp_tool(tags={"training", "model", "evaluation"})
-async def get_model_feature_impact(project_id: str, model_id: str) -> str:
-    """
-    Get detailed feature impact for a specific model.
-
-    Args:
-        project_id: The ID of the DataRobot project
-        model_id: The ID of the model to analyze
+@dr_mcp_tool(tags={"predictive", "training", "read", "model", "evaluation"})
+async def get_model_feature_impact(
+    *,
+    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,
+) -> ToolError | ToolResult:
+    """Get detailed feature impact for a specific model."""
+    if not project_id:
+        return ToolError("Project ID must be provided")
+    if not model_id:
+        return ToolError("Model ID must be provided")
 
-    Returns
-    -------
-        JSON string containing:
-        - feature_impact: Feature importance scores
-        - ui_panel: List of recommended UI panels for visualization
-    """
     client = get_sdk_client()
     project = client.Project.get(project_id)
     model = client.Model.get(project=project, model_id=model_id)
@@ -608,26 +609,31 @@ async def get_model_feature_impact(project_id: str, model_id: str) -> str:
     model.request_feature_impact()
     feature_impact = model.get_or_request_feature_impact()
 
-    return json.dumps({"data": feature_impact, "ui_panel": ["feature-impact"]}, indent=2)
-
+    return ToolResult(
+        content=json.dumps({"data": feature_impact}, indent=2),
+        structured_content={"data": feature_impact},
+    )
 
-@dr_mcp_tool(tags={"training", "model", "evaluation"})
-async def get_model_lift_chart(project_id: str, model_id: str, source: str = "validation") -> str:
-    """
-    Get detailed lift chart for a specific model.
 
-    Args:
-        project_id: The ID of the DataRobot project
-        model_id: The ID of the model to analyze
-        source: The source of the data to use for the lift chart ('validation' or 'holdout' or
-            'crossValidation')
+@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[
+        str,
+        """
+        The source of the data to use for the lift chart
+        ('validation' or 'holdout' or 'crossValidation')
+        """,
+    ]
+    | str = "validation",
+) -> ToolError | ToolResult:
+    """Get detailed lift chart for a specific model."""
+    if not project_id:
+        return ToolError("Project ID must be provided")
+    if not model_id:
+        return ToolError("Model ID must be provided")
 
-    Returns
-    -------
-        JSON string containing:
-        - lift_chart: Lift chart data
-        - ui_panel: List of recommended UI panels for visualization
-    """
     client = get_sdk_client()
     project = client.Project.get(project_id)
     model = client.Model.get(project=project, model_id=model_id)
@@ -648,4 +654,7 @@ async def get_model_lift_chart(project_id: str, model_id: str, source: str = "va
         "target_class": lift_chart.target_class,
     }
 
-    return json.dumps({"data": lift_chart_data, "ui_panel": ["lift-chart"]}, indent=2)
+    return ToolResult(
+        content=json.dumps({"data": lift_chart_data}, indent=2),
+        structured_content={"data": lift_chart_data},
+    )

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

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