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

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.39.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.34
+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.34.dist-info → datarobot_genai-0.2.39.dist-info}/RECORD

@@ -84,25 +84,25 @@ 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=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=zJ-UA1TMhPOYcExvgWv0YBjDsSIDPA-U1SEbBrVfAc8,7744
+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=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
-datarobot_genai/drmcp/tools/predictive/deployment_info.py,sha256=BGEF_dmbxOBJR0n1Tt9TO2-iNTQSBTr-oQUyaxLZ0ZI,15297
+datarobot_genai/drmcp/tools/predictive/deployment.py,sha256=Pc6lz9V2JOw3Ufw-SsGAhMKf6-YhvbjGoNLRFOIcSSY,3670
+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
 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.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,,