Skip to content

datarobot_genai.drmcp.test_utils.tool_base_ete

tool_base_ete

ToolCallTestExpectations

Bases: BaseModel

Class to store tool call information.

Source code in datarobot_genai/drmcp/test_utils/tool_base_ete.py
class ToolCallTestExpectations(BaseModel):
    """Class to store tool call information."""

    name: str
    acceptable_tool_names: list[str] = Field(
        default_factory=list,
        description=(
            "Logical tool names treated as equivalent for this step (same parameters/result "
            "checks). Example: deployment_get_features vs deployment_get_info for deployment "
            "scoring metadata."
        ),
    )
    parameters: dict[str, Any]
    result: str | dict[str, Any]

    def allowed_tool_names(self) -> set[str]:
        return {self.name, *self.acceptable_tool_names}

ETETestExpectations

Bases: BaseModel

Class to store test expectations for ETE tests.

By default allow_unexpected_tool_calls is True so models may call extra tools (e.g. modeling_list_projects after an error, catalog_get_preview after resolving an id). Set it to False when a test must assert an exact tool-call count and order with no additional calls.

Source code in datarobot_genai/drmcp/test_utils/tool_base_ete.py
class ETETestExpectations(BaseModel):
    """Class to store test expectations for ETE tests.

    By default ``allow_unexpected_tool_calls`` is True so models may call extra tools
    (e.g. modeling_list_projects after an error, catalog_get_preview after resolving an id).
    Set it to False when a test must assert an exact tool-call count and order with
    no additional calls.
    """

    potential_no_tool_calls: bool = False
    allow_unexpected_tool_calls: bool = True
    tool_calls_expected: list[ToolCallTestExpectations]
    llm_response_content_contains_expectations: list[str]

ToolBaseE2E

Base class for end-to-end tests.

Source code in datarobot_genai/drmcp/test_utils/tool_base_ete.py
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
                - allow_unexpected_tool_calls: Default True — expected calls must appear in order,
                  with optional extra calls allowed. Set False for strict exact count.
                - 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_models_get_bestmodel_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:
            diagnostics = _build_failure_diagnostics(
                test_expectations.tool_calls_expected,
                response,
            )

            expected_calls = test_expectations.tool_calls_expected
            expected_actual_indices: list[tuple[int, int]] = []

            # Verify LLM decided to use tools
            if test_expectations.allow_unexpected_tool_calls:
                assert len(response.tool_calls) >= len(expected_calls), (
                    f"LLM should have decided to call tools\n{diagnostics}"
                )
                search_start = 0
                for expected_idx, expected_call in enumerate(expected_calls):
                    matched_actual_idx = None
                    for actual_idx in range(search_start, len(response.tool_calls)):
                        actual_call = response.tool_calls[actual_idx]
                        canonical = _canonical_tool_name_for_expectation(
                            actual_call.tool_name, expected_call
                        )
                        if canonical is None:
                            continue
                        if not _check_dict_params_match(
                            expected_call.parameters, actual_call.parameters
                        ):
                            continue
                        matched_actual_idx = actual_idx
                        break

                    allowed = ", ".join(sorted(expected_call.allowed_tool_names()))
                    assert matched_actual_idx is not None, (
                        f"Should have called one of [{allowed}] with the correct "
                        f"parameters. Expected (subset): {expected_call.parameters}\n"
                        f"{diagnostics}"
                    )
                    expected_actual_indices.append((expected_idx, matched_actual_idx))
                    search_start = matched_actual_idx + 1
            else:
                assert len(response.tool_calls) == len(expected_calls), (
                    f"LLM should have decided to call tools\n{diagnostics}"
                )
                expected_actual_indices = [(i, i) for i in range(len(expected_calls))]

            for expected_idx, actual_idx in expected_actual_indices:
                tool_call = response.tool_calls[actual_idx]
                expected_call = expected_calls[expected_idx]
                canonical = _canonical_tool_name_for_expectation(tool_call.tool_name, expected_call)
                assert canonical is not None, (
                    f"Should have called one of {sorted(expected_call.allowed_tool_names())}, "
                    f"but got: {tool_call.tool_name}\n"
                    f"{diagnostics}"
                )

                assert _check_dict_params_match(expected_call.parameters, tool_call.parameters), (
                    f"Should have called {canonical} tool with the correct parameters. "
                    f"Expected (subset): {expected_call.parameters}, "
                    f"but got: {tool_call.parameters}\n"
                    f"{diagnostics}"
                )
                if expected_call.result != SHOULD_NOT_BE_EMPTY:
                    expected_result = expected_call.result
                    if isinstance(expected_result, str):
                        assert expected_result in response.tool_results[actual_idx], (
                            f"Should have called {canonical} tool with the correct "
                            f"result, but got: {response.tool_results[actual_idx]}\n"
                            f"{diagnostics}"
                        )
                    else:
                        actual_result = _extract_structured_content(
                            response.tool_results[actual_idx]
                        )
                        if actual_result is None:
                            # Fallback: try to parse the entire tool result as JSON
                            try:
                                actual_result = json.loads(response.tool_results[actual_idx])
                            except json.JSONDecodeError:
                                # If that fails, try to extract content part
                                if "Content: " in response.tool_results[actual_idx]:
                                    content_part = response.tool_results[actual_idx].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"{canonical}: "
                                            f"{response.tool_results[actual_idx]}"
                                        )
                        assert _check_dict_has_keys(expected_result, actual_result), (
                            f"Should have called {canonical} tool with the correct "
                            f"result structure, but got: {response.tool_results[actual_idx]}\n"
                            f"{diagnostics}"
                        )
                else:
                    assert len(response.tool_results[actual_idx]) > 0, (
                        f"Should have called {canonical} tool with non-empty result, but "
                        f"got: {response.tool_results[actual_idx]}\n"
                        f"{diagnostics}"
                    )

        # Verify LLM provided comprehensive response
        assert len(response.content) > 60, "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}"
        )