datarobot-genai 0.2.0__py3-none-any.whl

datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -0,0 +1,151 @@
+# Copyright 2025 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.
+
+import json
+from typing import Any
+
+from pydantic import BaseModel
+
+from .openai_llm_mcp_client import LLMResponse
+
+
+class ToolCallTestExpectations(BaseModel):
+    """Class to store tool call information."""
+
+    name: str
+    parameters: dict[str, Any]
+    result: str | dict[str, Any]
+
+
+class ETETestExpectations(BaseModel):
+    """Class to store test expectations for ETE tests."""
+
+    potential_no_tool_calls: bool = False
+    tool_calls_expected: list[ToolCallTestExpectations]
+    llm_response_content_contains_expectations: list[str]
+
+
+SHOULD_NOT_BE_EMPTY = "SHOULD_NOT_BE_EMPTY"
+
+
+def _check_dict_has_keys(
+    expected: dict[str, Any],
+    actual: dict[str, Any] | list[dict[str, Any]],
+    path: str = "",
+) -> bool:
+    """
+    Recursively check if all keys in expected dict exist in actual dict or in each item of
+    actual list.
+    Returns True if all expected keys exist, False otherwise.
+    """
+    # If actual is a list, check each item against the expected structure
+    if isinstance(actual, list):
+        if not actual:  # Empty list
+            return False
+        # Check first item against expected structure
+        return _check_dict_has_keys(expected, actual[0], path)
+
+    # Regular dict check
+    for key, value in expected.items():
+        current_path = f"{path}.{key}" if path else key
+        if key not in actual:
+            return False
+        if isinstance(value, dict):
+            if not isinstance(actual[key], dict):
+                return False
+            if not _check_dict_has_keys(value, actual[key], current_path):
+                return False
+    return True
+
+
+class ToolBaseE2E:
+    """Base class for end-to-end tests."""
+
+    async def _run_test_with_expectations(
+        self,
+        prompt: str,
+        test_expectations: ETETestExpectations,
+        openai_llm_client: Any,
+        mcp_session: Any,
+        test_name: str,
+    ) -> None:
+        """
+        Run a test with given expectations and validate the results.
+
+        Args:
+            prompt: The prompt to send to the LLM
+            test_expectations: ETETestExpectations object containing test expectations with keys:
+                - tool_calls_expected: List of expected tool calls with their parameters and results
+                - llm_response_content_contains_expectations: Expected content in the LLM response
+            openai_llm_client: The OpenAI LLM client
+            mcp_session: The test session
+            test_name: The name of the test (e.g. test_get_best_model_success)
+        """
+        # Get the test file name from the class name
+        file_name = self.__class__.__name__.lower().replace("e2e", "").replace("test", "")
+        output_file_name = f"{file_name}_{test_name}"
+
+        # Act
+        response: LLMResponse = await openai_llm_client.process_prompt_with_mcp_support(
+            prompt, mcp_session, output_file_name
+        )
+
+        # sometimes llm are too smart and doesn't call tools especially for the case when file
+        # doesn't exist
+        if test_expectations.potential_no_tool_calls and len(response.tool_calls) == 0:
+            pass
+        else:
+            # Verify LLM decided to use tools
+            assert len(response.tool_calls) == len(test_expectations.tool_calls_expected), (
+                "LLM should have decided to call tools"
+            )
+
+            for i, tool_call in enumerate(response.tool_calls):
+                assert tool_call.tool_name == test_expectations.tool_calls_expected[i].name, (
+                    f"Should have called {test_expectations.tool_calls_expected[i].name} tool, but "
+                    f"got: {tool_call.tool_name}"
+                )
+            assert tool_call.parameters == test_expectations.tool_calls_expected[i].parameters, (
+                f"Should have called {tool_call.tool_name} tool with the correct parameters, but "
+                f"got: {tool_call.parameters}"
+            )
+            if test_expectations.tool_calls_expected[i].result != SHOULD_NOT_BE_EMPTY:
+                expected_result = test_expectations.tool_calls_expected[i].result
+                if isinstance(expected_result, str):
+                    assert expected_result in response.tool_results[i], (
+                        f"Should have called {tool_call.tool_name} tool with the correct result, "
+                        f"but got: {response.tool_results[i]}"
+                    )
+                else:
+                    actual_result = json.loads(response.tool_results[i])
+                    assert _check_dict_has_keys(expected_result, actual_result), (
+                        f"Should have called {tool_call.tool_name} tool with the correct result "
+                        f"structure, but got: {response.tool_results[i]}"
+                    )
+            else:
+                assert len(response.tool_results[i]) > 0, (
+                    f"Should have called {tool_call.tool_name} tool with non-empty result, but "
+                    f"got: {response.tool_results[i]}"
+                )
+
+        # Verify LLM provided comprehensive response
+        assert len(response.content) > 100, "LLM should provide detailed response"
+        assert any(
+            expected_response.lower() in response.content
+            for expected_response in test_expectations.llm_response_content_contains_expectations
+        ), (
+            f"Response should mention "
+            f"{test_expectations.llm_response_content_contains_expectations}, "
+            f"but got: {response.content}"
+        )

datarobot_genai/drmcp/test_utils/utils.py

@@ -0,0 +1,91 @@
+# Copyright 2025 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.
+
+import datetime
+import json
+import os
+from typing import TYPE_CHECKING
+from typing import Any
+
+if TYPE_CHECKING:
+    from .openai_llm_mcp_client import LLMResponse
+
+from dotenv import load_dotenv
+
+
+def load_env() -> None:
+    load_dotenv(dotenv_path=".env", verbose=True, override=True)
+
+
+def format_tool_call(tool_call: dict[str, Any]) -> str:
+    """Format a single tool call in a readable way."""
+    return (
+        f"Tool: {tool_call['tool_name']}\n"
+        f"Parameters: {json.dumps(tool_call['parameters'], indent=2)}\n"
+        f"Reasoning: {tool_call['reasoning']}"
+    )
+
+
+def format_response(response: "LLMResponse") -> str:
+    """Format the LLM response in a readable way."""
+    formatted_parts = []
+
+    # Format the main content
+    formatted_parts.append("=== LLM Response ===\n")
+    formatted_parts.append(response.content)
+
+    # Format tool calls if any
+    if response.tool_calls:
+        formatted_parts.append("\n=== Tools Used ===")
+        for i, tool_call in enumerate(response.tool_calls, 1):
+            formatted_parts.append(f"\nTool Call #{i}:")
+            formatted_parts.append(
+                format_tool_call(
+                    {
+                        "tool_name": tool_call.tool_name,
+                        "parameters": tool_call.parameters,
+                        "reasoning": tool_call.reasoning,
+                    }
+                )
+            )
+
+    # Format tool results if any
+    if response.tool_results:
+        formatted_parts.append("\n=== Tool Results ===")
+        for i, result in enumerate(response.tool_results, 1):
+            formatted_parts.append(f"\nResult #{i}:")
+            formatted_parts.append(result)
+
+    return "\n".join(formatted_parts)
+
+
+def save_response_to_file(response: "LLMResponse", name: str | None = None) -> None:
+    """Save the response to a file in a readable format.
+
+    Args:
+        response: The LLM response to save
+        name: Optional name to use in the filename. If not provided,
+              will use a timestamp only.
+    """
+    # Create responses directory with timestamp
+    timestamp = datetime.datetime.now().strftime("%Y%m%d")
+    dir_path = "test_results/drmcp/.ete_responses/" + timestamp
+    os.makedirs(dir_path, exist_ok=True)
+
+    # Save both raw JSON and formatted text
+    base_name = f"{name}" if name else "response"
+
+    # Save formatted text
+    with open(f"{dir_path}/{base_name}.txt", "w") as f:
+        f.write(format_response(response))

datarobot_genai/drmcp/tools/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 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.
+

datarobot_genai/drmcp/tools/predictive/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2025 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.
+
+"""
+DataRobot MCP tools for interacting with DataRobot platform.
+
+This package contains pre-built tools for DataRobot operations:
+- data: Upload and manage datasets
+- project: Create and manage projects
+- model: List and interact with models
+- training: Train models with autopilot
+- deployment: Deploy and manage models
+- predict: Make batch predictions
+- predict_realtime: Make real-time predictions
+- deployment_info: Get deployment information
+"""

datarobot_genai/drmcp/tools/predictive/data.py

@@ -0,0 +1,97 @@
+# Copyright 2025 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.
+
+import logging
+import os
+
+from datarobot_genai.drmcp.core.clients import get_sdk_client
+from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"data", "management", "upload"})
+async def upload_dataset_to_ai_catalog(file_path: str) -> str:
+    """
+    Upload a dataset to the DataRobot AI Catalog.
+
+    Args:
+        file_path: Path to the file to upload.
+
+    Returns
+    -------
+        A string summary of the upload result.
+    """
+    client = get_sdk_client()
+    if not os.path.exists(file_path):
+        logger.error(f"File not found: {file_path}")
+        return f"File not found: {file_path}"
+    catalog_item = client.Dataset.create_from_file(file_path)
+    logger.info(f"Successfully uploaded dataset: {catalog_item.id}")
+    return f"AI Catalog ID: {catalog_item.id}"
+
+
+@dr_mcp_tool(tags={"data", "management", "list"})
+async def list_ai_catalog_items() -> str:
+    """
+    List all AI Catalog items (datasets) for the authenticated user.
+
+    Returns
+    -------
+        A string summary of the AI Catalog items with their IDs and names.
+    """
+    client = get_sdk_client()
+    datasets = client.Dataset.list()
+    if not datasets:
+        logger.info("No AI Catalog items found")
+        return "No AI Catalog items found."
+    result = "\n".join(f"{ds.id}: {ds.name}" for ds in datasets)
+    logger.info(f"Found {len(datasets)} AI Catalog items")
+    return result
+
+
+# from fastmcp import Context
+
+# from datarobot_genai.drmcp.core.memory_management import MemoryManager, get_memory_manager
+
+
+# @dr_mcp_tool()
+# async def list_ai_catalog_items(
+#     ctx: Context, agent_id: str = None, storage_id: str = None
+# ) -> str:
+#     """
+#     List all AI Catalog items (datasets) for the authenticated user.
+
+#     Returns:
+#         a resource id that can be used to retrieve the list of AI Catalog items using the
+#         get_resource tool
+#     """
+#     client = get_sdk_client()
+#     datasets = client.Dataset.list()
+#     if not datasets:
+#         logger.info("No AI Catalog items found")
+#         return "No AI Catalog items found."
+#     result = "\n".join(f"{ds.id}: {ds.name}" for ds in datasets)
+
+#     if MemoryManager.is_initialized():
+#         resource_id = await get_memory_manager().store_resource(
+#             data=result,
+#             memory_storage_id=storage_id,
+#             agent_identifier=agent_id,
+#         )
+#     else:
+#         raise ValueError("MemoryManager is not initialized")
+
+#     logger.info(f"Found {len(datasets)} AI Catalog items")
+#     return resource_id

datarobot_genai/drmcp/tools/predictive/deployment.py

@@ -0,0 +1,91 @@
+# Copyright 2025 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.
+
+import json
+import logging
+
+from datarobot_genai.drmcp.core.clients import get_sdk_client
+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.
+    """
+    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
+
+
+@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.
+
+    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)
+
+
+@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.
+
+    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"})
+        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})
+    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}"})