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

datarobot_genai/drmcp/tools/clients/microsoft_graph.py

@@ -16,6 +16,7 @@
 
 import logging
 from typing import Any
+from urllib.parse import quote
 
 import httpx
 from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
@@ -415,6 +416,131 @@ class MicrosoftGraphClient:
 
         return base_resource
 
+    async def get_personal_drive_id(self) -> str:
+        """Get the current user's personal OneDrive drive ID.
+
+        Returns
+        -------
+            The drive ID string for the user's personal OneDrive.
+
+        Raises
+        ------
+            MicrosoftGraphError: If the drive cannot be retrieved.
+        """
+        try:
+            response = await self._client.get(f"{GRAPH_API_BASE}/me/drive")
+            response.raise_for_status()
+            data = response.json()
+            return data["id"]
+        except httpx.HTTPStatusError as e:
+            status_code = e.response.status_code
+            if status_code == 401:
+                raise MicrosoftGraphError(
+                    "Authentication failed. Access token may be expired or invalid."
+                ) from e
+            if status_code == 403:
+                raise MicrosoftGraphError(
+                    "Permission denied: cannot access personal OneDrive. "
+                    "Requires Files.Read or Files.ReadWrite permission."
+                ) from e
+            raise MicrosoftGraphError(f"Failed to get personal OneDrive: HTTP {status_code}") from e
+
+    async def create_file(
+        self,
+        drive_id: str,
+        file_name: str,
+        content: str,
+        parent_folder_id: str = "root",
+        conflict_behavior: str = "rename",
+    ) -> MicrosoftGraphItem:
+        """Create a text file in a drive (SharePoint document library or OneDrive).
+
+        Uses Microsoft Graph's simple upload endpoint for files < 4MB.
+        Files are created as text/plain content.
+
+        Args:
+            drive_id: The ID of the drive (document library) where the file will be created.
+            file_name: The name of the file to create (e.g., 'report.txt').
+            content: The text content to store in the file.
+            parent_folder_id: ID of the parent folder. Defaults to "root" (drive root folder).
+            conflict_behavior: How to handle name conflicts. Options:
+                - "rename" (default): Auto-renames to 'filename (1).txt', etc.
+                - "fail": Returns 409 Conflict error
+                - "replace": Overwrites existing file
+
+        Returns
+        -------
+            MicrosoftGraphItem representing the created file.
+
+        Raises
+        ------
+            MicrosoftGraphError: If file creation fails.
+        """
+        if not drive_id or not drive_id.strip():
+            raise MicrosoftGraphError("drive_id cannot be empty")
+        if not file_name or not file_name.strip():
+            raise MicrosoftGraphError("file_name cannot be empty")
+
+        # URL encode the filename for path-based addressing
+        encoded_name = quote(file_name, safe="")
+
+        # Simple upload endpoint for files < 4MB
+        # Reference: https://learn.microsoft.com/en-us/graph/api/driveitem-put-content
+        upload_url = (
+            f"{GRAPH_API_BASE}/drives/{drive_id}/items/{parent_folder_id}:/{encoded_name}:/content"
+        )
+
+        try:
+            response = await self._client.put(
+                upload_url,
+                content=content.encode("utf-8"),
+                headers={"Content-Type": "text/plain"},
+                params={"@microsoft.graph.conflictBehavior": conflict_behavior},
+            )
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            raise self._handle_create_file_error(e, drive_id, file_name, parent_folder_id) from e
+
+        return MicrosoftGraphItem.from_api_response(response.json())
+
+    def _handle_create_file_error(
+        self,
+        error: httpx.HTTPStatusError,
+        drive_id: str,
+        file_name: str,
+        parent_folder_id: str,
+    ) -> MicrosoftGraphError:
+        """Handle HTTP errors for file creation and return appropriate MicrosoftGraphError."""
+        status_code = error.response.status_code
+        error_msg = f"Failed to create file: HTTP {status_code}"
+
+        if status_code == 400:
+            try:
+                error_data = error.response.json()
+                api_message = error_data.get("error", {}).get("message", "Invalid request")
+                error_msg = f"Bad request creating file: {api_message}"
+            except Exception:
+                error_msg = "Bad request: invalid parameters for file creation."
+        elif status_code == 401:
+            error_msg = "Authentication failed. Access token may be expired or invalid."
+        elif status_code == 403:
+            error_msg = (
+                f"Permission denied: you don't have permission to create files in drive "
+                f"'{drive_id}'. Requires Files.ReadWrite.All permission."
+            )
+        elif status_code == 404:
+            error_msg = (
+                f"Parent folder '{parent_folder_id}' not found in drive '{drive_id}'."
+                if parent_folder_id != "root"
+                else f"Drive '{drive_id}' not found."
+            )
+        elif status_code == 409:
+            error_msg = f"File '{file_name}' already exists and conflict behavior is set to 'fail'."
+        elif status_code == 429:
+            error_msg = "Rate limit exceeded. Please try again later."
+
+        return MicrosoftGraphError(error_msg)
+
     async def __aenter__(self) -> "MicrosoftGraphClient":
         """Async context manager entry."""
         return self

datarobot_genai/drmcp/tools/microsoft_graph/tools.py

@@ -196,3 +196,82 @@ async def microsoft_graph_search_content(
             "count": n,
         },
     )
+
+
+@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/training.py

@@ -19,6 +19,7 @@ import logging
 from dataclasses import asdict
 from dataclasses import dataclass
 from typing import Annotated
+from typing import Any
 
 import pandas as pd
 from fastmcp.exceptions import ToolError
@@ -56,6 +57,36 @@ class DatasetInsight:
     missing_data_summary: dict[str, float]
 
 
+def _get_dataset_or_raise(client: Any, dataset_id: str) -> tuple[Any, pd.DataFrame]:
+    """Fetch dataset and return it with its dataframe, with proper error handling.
+
+    Args:
+        client: DataRobot SDK client instance
+        dataset_id: The ID of the dataset to fetch
+
+    Returns
+    -------
+        Tuple of (dataset object, dataframe)
+
+    Raises
+    ------
+        ToolError: If dataset is not found (404) or other error occurs
+    """
+    try:
+        dataset = client.Dataset.get(dataset_id)
+        return dataset, dataset.get_as_dataframe()
+    except Exception as e:
+        error_str = str(e)
+        # Check if it's a 404 error (dataset not found)
+        if "404" in error_str or "Not Found" in error_str:
+            raise ToolError(
+                f"Dataset '{dataset_id}' not found. Please verify the dataset ID exists "
+                "and you have access to it."
+            )
+        # For other errors, provide context
+        raise ToolError(f"Failed to retrieve dataset '{dataset_id}': {error_str}")
+
+
 @dr_mcp_tool(tags={"predictive", "training", "read", "analysis", "dataset"})
 async def analyze_dataset(
     *,
@@ -66,8 +97,7 @@ async def analyze_dataset(
         raise ToolError("Dataset ID must be provided")
 
     client = get_sdk_client()
-    dataset = client.Dataset.get(dataset_id)
-    df = dataset.get_as_dataframe()
+    dataset, df = _get_dataset_or_raise(client, dataset_id)
 
     # Analyze dataset structure
     numerical_cols = df.select_dtypes(include=["int64", "float64"]).columns.tolist()
@@ -119,12 +149,11 @@ async def suggest_use_cases(
         raise ToolError("Dataset ID must be provided")
 
     client = get_sdk_client()
-    dataset = client.Dataset.get(dataset_id)
-    df = dataset.get_as_dataframe()
+    dataset, df = _get_dataset_or_raise(client, dataset_id)
 
     # Get dataset insights first
-    insights_json = await analyze_dataset(dataset_id)
-    insights = json.loads(insights_json)
+    insights_result = await analyze_dataset(dataset_id=dataset_id)
+    insights = insights_result.structured_content
 
     suggestions = []
     for target_col in insights["potential_targets"]:
@@ -151,12 +180,11 @@ async def get_exploratory_insights(
         raise ToolError("Dataset ID must be provided")
 
     client = get_sdk_client()
-    dataset = client.Dataset.get(dataset_id)
-    df = dataset.get_as_dataframe()
+    dataset, df = _get_dataset_or_raise(client, dataset_id)
 
     # Get dataset insights first
-    insights_json = await analyze_dataset(dataset_id)
-    insights = json.loads(insights_json)
+    insights_result = await analyze_dataset(dataset_id=dataset_id)
+    insights = insights_result.structured_content
 
     eda_insights = {
         "dataset_summary": {

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

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

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

@@ -84,7 +84,7 @@ 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=PASGThDPE8zkBZqach8lurJL1y47DWUPLwvf9N6uLGM,19234
+datarobot_genai/drmcp/tools/clients/microsoft_graph.py,sha256=-g0EhaBVElKbujaO2cHdgc86hwFEkkyEyZVAw8pq7yM,24468
 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
@@ -93,16 +93,16 @@ datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=7bNrp7E3opKwsBDYfLIOsOGfPXW-A
 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/microsoft_graph/__init__.py,sha256=CuOaMt1AJo7cHx_GuhO3s_aqxZas_wlDsoBorBsvbeU,577
-datarobot_genai/drmcp/tools/microsoft_graph/tools.py,sha256=zJ-UA1TMhPOYcExvgWv0YBjDsSIDPA-U1SEbBrVfAc8,7744
+datarobot_genai/drmcp/tools/microsoft_graph/tools.py,sha256=cNctozv_4lRC5Kva3D2j4taZfeQHDE6LTAjcmeQXwWA,10446
 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=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
+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/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
 datarobot_genai/drmcp/tools/predictive/project.py,sha256=Mzf7rQogBV6h1-MWQYTwtDHOsMWfjOyyJpSYmmvNNuc,3253
-datarobot_genai/drmcp/tools/predictive/training.py,sha256=WWzzGibYMSvI8kqHnvav6qNIVjoe1EG4RyiYa3XhFYA,23984
+datarobot_genai/drmcp/tools/predictive/training.py,sha256=jeZGPWJ69PPOd2MhUbACgbllQ0CK7Kz-hNl596mJujQ,25021
 datarobot_genai/langgraph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datarobot_genai/langgraph/agent.py,sha256=DRnywmS9KDywyChtuIZZwNKbJs8BpC259EG_kxYbiQ8,15828
 datarobot_genai/langgraph/mcp.py,sha256=iA2_j46mZAaNaL7ntXT-LW6C-NMJkzr3VfKDDfe7mh8,2851
@@ -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.34.dist-info/METADATA,sha256=ypCJpwVf3XHRuudcLQZUwOOYDUwaJveVbk7133Wp1kQ,6301
-datarobot_genai-0.2.34.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.34.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.34.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.34.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.34.dist-info/RECORD,,
+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,,