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

datarobot_genai/drmcp/tools/perplexity/tools.py

@@ -0,0 +1,117 @@
+# Copyright 2026 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Perplexity MCP tools."""
+
+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.perplexity import MAX_QUERIES
+from datarobot_genai.drmcp.tools.clients.perplexity import MAX_RESULTS
+from datarobot_genai.drmcp.tools.clients.perplexity import MAX_RESULTS_DEFAULT
+from datarobot_genai.drmcp.tools.clients.perplexity import MAX_SEARCH_DOMAIN_FILTER
+from datarobot_genai.drmcp.tools.clients.perplexity import MAX_TOKENS_PER_PAGE
+from datarobot_genai.drmcp.tools.clients.perplexity import MAX_TOKENS_PER_PAGE_DEFAULT
+from datarobot_genai.drmcp.tools.clients.perplexity import PerplexityClient
+from datarobot_genai.drmcp.tools.clients.perplexity import get_perplexity_access_token
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"perplexity", "web", "search", "websearch"})
+async def perplexity_search(
+    *,
+    query: Annotated[
+        str,
+        list[str],
+        f"The search query string OR "
+        f"a list of up to {MAX_QUERIES} sub-queries for multi-query research.",
+    ],
+    search_domain_filter: Annotated[
+        list[str] | None,
+        f"Up to {MAX_SEARCH_DOMAIN_FILTER} domains/URLs "
+        f"to allowlist or denylist (prefix with '-').",
+    ] = None,
+    recency: Annotated[
+        Literal["day", "week", "month", "year"] | None, "Filter results by time period."
+    ] = None,
+    max_results: Annotated[
+        int, f"Number of ranked results to return (1-{MAX_RESULTS})."
+    ] = MAX_RESULTS_DEFAULT,
+    max_tokens_per_page: Annotated[
+        int,
+        f"Content extraction cap per page (1-{MAX_TOKENS_PER_PAGE}) "
+        f"(default {MAX_TOKENS_PER_PAGE_DEFAULT}).",
+    ] = MAX_TOKENS_PER_PAGE_DEFAULT,
+) -> ToolResult:
+    """Perplexity web search tool combining multi-query research and content extraction control."""
+    if not query:
+        raise ToolError("Argument validation error: query cannot be empty.")
+    if query and isinstance(query, str) and not query.strip():
+        raise ToolError("Argument validation error: query cannot be empty.")
+    if query and isinstance(query, list) and len(query) > MAX_QUERIES:
+        raise ToolError(
+            f"Argument validation error: query list cannot be bigger than {MAX_QUERIES}."
+        )
+    if query and isinstance(query, list) and not all(q.strip() for q in query):
+        raise ToolError("Argument validation error: query cannot contain empty str.")
+    if search_domain_filter and len(search_domain_filter) > MAX_SEARCH_DOMAIN_FILTER:
+        raise ToolError(
+            f"Argument validation error: "
+            f"maximum number of search domain filters is {MAX_SEARCH_DOMAIN_FILTER}."
+        )
+    if max_results <= 0:
+        raise ToolError("Argument validation error: max_results must be greater than 0.")
+    if max_results > MAX_RESULTS:
+        raise ToolError(
+            f"Argument validation error: "
+            f"max_results must be smaller than or equal to {MAX_RESULTS}."
+        )
+    if max_tokens_per_page <= 0:
+        raise ToolError("Argument validation error: max_tokens_per_page must be greater than 0.")
+    if max_tokens_per_page > MAX_TOKENS_PER_PAGE:
+        raise ToolError(
+            f"Argument validation error: "
+            f"max_tokens_per_page must be smaller than or equal to {MAX_TOKENS_PER_PAGE}."
+        )
+
+    access_token = await get_perplexity_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    async with PerplexityClient(access_token=access_token) as perplexity_client:
+        results = await perplexity_client.search(
+            query=query,
+            search_domain_filter=search_domain_filter,
+            recency=recency,
+            max_results=max_results,
+            max_tokens_per_page=max_tokens_per_page,
+        )
+
+    return ToolResult(
+        structured_content={
+            "results": results,
+            "count": len(results),
+            "metadata": {
+                "queriesExecuted": len(query) if isinstance(query, list) else 1,
+                "filtersApplied": {"domains": search_domain_filter, "recency": recency},
+                "extractionLimit": max_tokens_per_page,
+            },
+        },
+    )

datarobot_genai/drmcp/tools/predictive/data.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import logging
 import os
 from typing import Annotated
@@ -60,7 +59,6 @@ async def upload_dataset_to_ai_catalog(
         raise ToolError("Failed to upload dataset.")
 
     return ToolResult(
-        content=f"Successfully uploaded dataset: {catalog_item.id}",
         structured_content={
             "dataset_id": catalog_item.id,
             "dataset_version_id": catalog_item.version_id,
@@ -78,21 +76,15 @@ async def list_ai_catalog_items() -> ToolResult:
     if not datasets:
         logger.info("No AI Catalog items found")
         return ToolResult(
-            content="No AI Catalog items found.",
             structured_content={"datasets": []},
         )
 
     datasets_dict = {ds.id: ds.name for ds in datasets}
-    datasets_count = len(datasets)
 
     return ToolResult(
-        content=(
-            f"Found {datasets_count} AI Catalog items, here are the details:\n"
-            f"{json.dumps(datasets_dict, indent=2)}"
-        ),
         structured_content={
             "datasets": datasets_dict,
-            "count": datasets_count,
+            "count": len(datasets),
         },
     )
 

datarobot_genai/drmcp/tools/predictive/deployment.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import logging
 from typing import Annotated
 
@@ -32,12 +31,10 @@ async def list_deployments() -> ToolResult:
     deployments = client.Deployment.list()
     if not deployments:
         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},
     )
 
@@ -54,10 +51,6 @@ async def get_model_info_from_deployment(
     client = get_sdk_client()
     deployment = client.Deployment.get(deployment_id)
     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,
     )
 
@@ -87,7 +80,6 @@ async def deploy_model(
             default_prediction_server_id=prediction_servers[0].id,
         )
         return ToolResult(
-            content=f"Created deployment {deployment.id} with label {label}",
             structured_content={
                 "deployment_id": deployment.id,
                 "label": label,

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,33 @@ 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(
+        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
@@ -203,64 +178,54 @@ async def generate_prediction_data_template(deployment_id: str, n_rows: int = 1)
     # Create DataFrame
     df = pd.DataFrame(template_data)
 
-    # Add metadata comments
-    result = f"# Prediction Data Template for Deployment: {deployment_id}\n"
-    result += f"# Model Type: {features_info['model_type']}\n"
-    result += f"# Target: {features_info['target']} (Type: {features_info['target_type']})\n"
+    # 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:
-        ts = features_info["time_series_config"]
-        result += f"# Time Series: datetime_column={ts['datetime_column']}, "
-        result += f"forecast_window=[{ts['forecast_window_start']}, {ts['forecast_window_end']}]\n"
-        if ts["series_id_columns"]:
-            result += f"# Multiseries ID Columns: {', '.join(ts['series_id_columns'])}\n"
-
-    result += f"# Total Features: {features_info['total_features']}\n"
-    result += df.to_csv(index=False)
+        structured_content["time_series_config"] = features_info["time_series_config"]
 
-    return str(result)
+    return ToolResult(
+        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 +324,28 @@ async def validate_prediction_data(
         "model_type": features_info["model_type"],
     }
 
-    return json.dumps(validation_report, indent=2)
+    return ToolResult(
+        structured_content=validation_report,
+    )
 
 
-@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)
+@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 +360,7 @@ 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(
+        structured_content=result,
+    )

datarobot_genai/drmcp/tools/predictive/model.py

@@ -93,33 +93,17 @@ async def get_best_model(
     best_model = leaderboard[0]
     logger.info(f"Found best model {best_model.id} for project {project_id}")
 
-    metric_info = ""
     metric_value = None
 
     if metric and best_model.metrics and metric in best_model.metrics:
         metric_value = best_model.metrics[metric].get("validation")
-        if metric_value is not None:
-            metric_info = f" with {metric}: {metric_value:.2f}"
 
     # Include full metrics in the response
     best_model_dict = model_to_dict(best_model)
     best_model_dict["metric"] = metric
     best_model_dict["metric_value"] = metric_value
 
-    # Format metrics for human-readable content
-    metrics_text = ""
-    if best_model.metrics:
-        metrics_list = []
-        for metric_name, metric_data in best_model.metrics.items():
-            if isinstance(metric_data, dict) and "validation" in metric_data:
-                val = metric_data["validation"]
-                if val is not None:
-                    metrics_list.append(f"{metric_name}: {val:.4f}")
-        if metrics_list:
-            metrics_text = "\nPerformance metrics:\n" + "\n".join(f"  - {m}" for m in metrics_list)
-
     return ToolResult(
-        content=f"Best model: {best_model.model_type}{metric_info}{metrics_text}",
         structured_content={
             "project_id": project_id,
             "best_model": best_model_dict,
@@ -148,7 +132,6 @@ async def score_dataset_with_model(
     job = model.score(dataset_url)
 
     return ToolResult(
-        content=f"Scoring job started: {job.id}",
         structured_content={
             "scoring_job_id": job.id,
             "project_id": project_id,
@@ -172,10 +155,6 @@ async def list_models(
     models = project.get_models()
 
     return ToolResult(
-        content=(
-            f"Found {len(models)} models in project {project_id}, here are the details:\n"
-            f"{json.dumps(models, indent=2, cls=ModelEncoder)}"
-        ),
         structured_content={
             "project_id": project_id,
             "models": [model_to_dict(model) for model in models],

datarobot_genai/drmcp/tools/predictive/predict_realtime.py

@@ -240,6 +240,9 @@ async def predict_realtime(
     else:
         raise ValueError("Either file_path or dataset must be provided.")
 
+    # Normalize column names: strip leading/trailing whitespace
+    df.columns = df.columns.str.strip()
+
     if series_id_column and series_id_column not in df.columns:
         raise ValueError(f"series_id_column '{series_id_column}' not found in input data.")
 

datarobot_genai/drmcp/tools/predictive/project.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import logging
 from typing import Annotated
 
@@ -33,11 +32,6 @@ async def list_projects() -> ToolResult:
     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,
     )
 
@@ -48,7 +42,7 @@ 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:
+) -> ToolResult:
     """Get a dataset ID by name for a given project.
 
     The dataset ID and the dataset type (source or prediction) as a string, or an error message.
@@ -70,21 +64,11 @@ async def get_project_dataset_by_name(
     for ds in all_datasets:
         if dataset_name.lower() in ds["dataset"].name.lower():
             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"],
                 },
             )
-    return ToolResult(
-        content=f"Dataset with name containing '{dataset_name}' not found in project {project_id}.",
-        structured_content={},
+    raise ToolError(
+        f"Dataset with name containing '{dataset_name}' not found in project {project_id}."
     )

datarobot_genai/drmcp/tools/predictive/training.py

@@ -14,7 +14,6 @@
 
 """Tools for analyzing datasets and suggesting ML use cases."""
 
-import json
 import logging
 from dataclasses import asdict
 from dataclasses import dataclass
@@ -134,7 +133,6 @@ async def analyze_dataset(
     insights_dict = asdict(insights)
 
     return ToolResult(
-        content=json.dumps(insights_dict, indent=2),
         structured_content=insights_dict,
     )
 
@@ -164,7 +162,6 @@ async def suggest_use_cases(
     suggestions.sort(key=lambda x: x["confidence"], reverse=True)
 
     return ToolResult(
-        content=json.dumps(suggestions, indent=2),
         structured_content={"use_case_suggestions": suggestions},
     )
 
@@ -255,7 +252,6 @@ async def get_exploratory_insights(
                 )
 
     return ToolResult(
-        content=json.dumps(eda_insights, indent=2),
         structured_content=eda_insights,
     )
 
@@ -540,22 +536,11 @@ async def start_autopilot(
         }
 
         return ToolResult(
-            content=json.dumps(result, indent=2),
             structured_content=result,
         )
 
     except Exception as e:
-        raise 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,
-            )
-        )
+        raise ToolError(f"Failed to start Autopilot: {str(e)}")
 
 
 @dr_mcp_tool(tags={"prediction", "training", "read", "model", "evaluation"})
@@ -611,7 +596,6 @@ async def get_model_roc_curve(
         }
 
         return ToolResult(
-            content=json.dumps({"data": roc_data}, indent=2),
             structured_content={"data": roc_data},
         )
     except Exception as e:
@@ -638,7 +622,6 @@ async def get_model_feature_impact(
     feature_impact = model.get_or_request_feature_impact()
 
     return ToolResult(
-        content=json.dumps({"data": feature_impact}, indent=2),
         structured_content={"data": feature_impact},
     )
 
@@ -684,6 +667,5 @@ async def get_model_lift_chart(
     }
 
     return ToolResult(
-        content=json.dumps({"data": lift_chart_data}, indent=2),
         structured_content={"data": lift_chart_data},
     )