datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/drmcp/test_utils/test_interactive.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python3
+
+# 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.
+
+"""Interactive MCP Client Test Script.
+
+This script allows you to test arbitrary commands with the MCP server
+using an LLM agent that can decide which tools to call.
+
+Supports elicitation - when tools require user input (like authentication tokens),
+the script will prompt you interactively.
+"""
+
+import asyncio
+import json
+import os
+import sys
+import traceback
+from pathlib import Path
+from typing import Any
+
+from dotenv import load_dotenv
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.context import RequestContext
+from mcp.types import ElicitRequestParams
+from mcp.types import ElicitResult
+
+from datarobot_genai.drmcp import get_dr_mcp_server_url
+from datarobot_genai.drmcp import get_headers
+from datarobot_genai.drmcp.test_utils.clients.base import LLMResponse
+from datarobot_genai.drmcp.test_utils.clients.base import ToolCall
+from datarobot_genai.drmcp.test_utils.clients.dr_gateway import DRLLMGatewayMCPClient
+
+# Re-export for backwards compatibility
+__all__ = ["DRLLMGatewayMCPClient", "LLMResponse", "ToolCall", "test_mcp_interactive"]
+
+
+async def test_mcp_interactive() -> None:
+    """Test the MCP server interactively with LLM agent."""
+    # Check for required environment variables
+    datarobot_api_token = os.environ.get("DATAROBOT_API_TOKEN")
+    if not datarobot_api_token:
+        print("❌ Error: DATAROBOT_API_TOKEN environment variable is required")
+        print("Please set it in your .env file or export it")
+        return
+
+    # Optional DataRobot settings
+    datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT")
+    model = os.environ.get("MODEL")
+
+    print("🤖 Initializing LLM MCP Client...")
+
+    # Initialize the LLM client with elicitation handler
+    config = {
+        "datarobot_api_token": datarobot_api_token,
+        "save_llm_responses": False,
+    }
+    if datarobot_endpoint:
+        config["datarobot_endpoint"] = datarobot_endpoint
+    if model:
+        config["model"] = model
+
+    llm_client = DRLLMGatewayMCPClient(str(config))
+
+    # Get MCP server URL
+    mcp_server_url = get_dr_mcp_server_url()
+    if not mcp_server_url:
+        print("❌ Error: MCP server URL is not configured")
+        print("Please set DR_MCP_SERVER_URL environment variable or run: task test-interactive")
+        return
+
+    print(f"🔗 Connecting to MCP server at: {mcp_server_url}")
+
+    # Elicitation handler: prompt user for required values
+    async def elicitation_handler(
+        context: RequestContext[ClientSession, Any], params: ElicitRequestParams
+    ) -> ElicitResult:
+        print(f"\n📋 Elicitation Request: {params.message}")
+        if params.requestedSchema:
+            print(f"   Schema: {params.requestedSchema}")
+
+        while True:
+            try:
+                response = input("   Enter value (or 'decline'/'cancel'): ").strip()
+            except (EOFError, KeyboardInterrupt):
+                return ElicitResult(action="cancel")
+
+            if response.lower() == "decline":
+                return ElicitResult(action="decline")
+            if response.lower() == "cancel":
+                return ElicitResult(action="cancel")
+            if response:
+                return ElicitResult(action="accept", content={"value": response})
+            print("   Please enter a value or 'decline'/'cancel'")
+
+    try:
+        async with streamablehttp_client(
+            url=mcp_server_url,
+            headers=get_headers(),
+        ) as (read_stream, write_stream, _):
+            async with ClientSession(
+                read_stream,
+                write_stream,
+                elicitation_callback=elicitation_handler,
+            ) as session:
+                await session.initialize()
+
+                print("✅ Connected to MCP server!")
+                print("📋 Available tools:")
+
+                tools_result = await session.list_tools()
+                for i, tool in enumerate(tools_result.tools, 1):
+                    print(f"  {i}. {tool.name}: {tool.description}")
+
+                print("\n" + "=" * 60)
+                print("🎯 Interactive Testing Mode")
+                print("=" * 60)
+                print("Type your questions/commands. The AI will decide which tools to use.")
+                print("If a tool requires additional information, you will be prompted.")
+                print("Type 'quit' or 'exit' to stop.")
+                print()
+
+                while True:
+                    try:
+                        user_input = input("🤔 You: ").strip()
+
+                        if user_input.lower() in ["quit", "exit", "q"]:
+                            print("👋 Goodbye!")
+                            break
+
+                        if not user_input:
+                            continue
+                    except (EOFError, KeyboardInterrupt):
+                        print("\n👋 Goodbye!")
+                        break
+
+                    print("🤖 AI is thinking...")
+
+                    response = await llm_client.process_prompt_with_mcp_support(
+                        prompt=user_input,
+                        mcp_session=session,
+                    )
+
+                    print("\n🤖 AI Response:")
+                    print("-" * 40)
+                    print(response.content)
+
+                    if response.tool_calls:
+                        print("\n🔧 Tools Used:")
+                        for i, tool_call in enumerate(response.tool_calls, 1):
+                            print(f"  {i}. {tool_call.tool_name}")
+                            print(f"     Parameters: {tool_call.parameters}")
+                            print(f"     Reasoning: {tool_call.reasoning}")
+
+                            if i <= len(response.tool_results):
+                                result = response.tool_results[i - 1]
+                                try:
+                                    result_data = json.loads(result)
+                                    if result_data.get("status") == "error":
+                                        error_msg = result_data.get("error", "Unknown error")
+                                        print(f"     ❌ Error: {error_msg}")
+                                    elif result_data.get("status") == "success":
+                                        print("     ✅ Success")
+                                except json.JSONDecodeError:
+                                    if len(result) > 100:
+                                        print(f"     Result: {result[:100]}...")
+                                    else:
+                                        print(f"     Result: {result}")
+
+                    print("\n" + "=" * 60)
+    except Exception as e:
+        print(f"❌ Connection Error: {e}")
+        print(f"   Server URL: {mcp_server_url}")
+        traceback.print_exc()
+        return
+
+
+if __name__ == "__main__":
+    # Ensure we're in the right directory
+    if not Path("src").exists():
+        print("❌ Error: Please run this script from the project root")
+        sys.exit(1)
+
+    # Load environment variables from .env file
+    print("📄 Loading environment variables...")
+    load_dotenv()
+
+    print("🚀 Starting Interactive MCP Client Test")
+    print("Make sure the MCP server is running with: task drmcp-dev")
+    print()
+
+    asyncio.run(test_mcp_interactive())

datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -0,0 +1,220 @@
+# 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 .clients.base 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 _extract_structured_content(tool_result: str) -> Any:
+    r"""
+    Extract and parse structured content from tool result string.
+
+    Tool results are formatted as:
+    "Content: {content}\nStructured content: {structured_content}"
+
+    Structured content can be:
+    1. A JSON object with a "result" key: {"result": "..."} or {"result": "{...}"}
+    2. A direct JSON object: {"key": "value", ...}
+    3. Empty or missing
+
+    Args:
+        tool_result: The tool result string
+
+    Returns
+    -------
+        Parsed structured content, or None if not available
+    """
+    # Early returns for invalid inputs
+    if not tool_result or "Structured content: " not in tool_result:
+        return None
+
+    structured_part = tool_result.split("Structured content: ", 1)[1].strip()
+    # Parse JSON, return None on failure or empty structured_part
+    if not structured_part:
+        return None
+    try:
+        structured_data = json.loads(structured_part)
+    except json.JSONDecodeError:
+        return None
+
+    # If structured data has a "result" key, extract and parse that
+    if isinstance(structured_data, dict) and "result" in structured_data:
+        result_value = structured_data["result"]
+        # If result is a JSON string (starts with { or [), try to parse it
+        if isinstance(result_value, str) and result_value.strip().startswith(("{", "[")):
+            try:
+                parsed_result = json.loads(result_value)
+            except json.JSONDecodeError:
+                parsed_result = result_value  # Return string as-is if parsing fails
+            return parsed_result
+        return result_value  # Return result value directly
+
+    # If it's a direct JSON object (not wrapped in {"result": ...}), return it as-is
+    return structured_data
+
+
+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, "
+                    f"but 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 "
+                            f"result, but got: {response.tool_results[i]}"
+                        )
+                    else:
+                        actual_result = _extract_structured_content(response.tool_results[i])
+                        if actual_result is None:
+                            # Fallback: try to parse the entire tool result as JSON
+                            try:
+                                actual_result = json.loads(response.tool_results[i])
+                            except json.JSONDecodeError:
+                                # If that fails, try to extract content part
+                                if "Content: " in response.tool_results[i]:
+                                    content_part = response.tool_results[i].split("Content: ", 1)[1]
+                                    if "\nStructured content: " in content_part:
+                                        content_part = content_part.split(
+                                            "\nStructured content: ", 1
+                                        )[0]
+                                    try:
+                                        actual_result = json.loads(content_part.strip())
+                                    except json.JSONDecodeError:
+                                        raise AssertionError(
+                                            f"Could not parse tool result for "
+                                            f"{tool_call.tool_name}: {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 "
+                            f"result 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 .clients.base 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/clients/__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.
+