datarobot-genai 0.2.24__py3-none-any.whl → 0.2.29__py3-none-any.whl

datarobot_genai/drmcp/core/config.py

@@ -31,6 +31,124 @@ from .constants import DEFAULT_DATAROBOT_ENDPOINT
 from .constants import RUNTIME_PARAM_ENV_VAR_NAME_PREFIX
 
 
+class MCPToolConfig(BaseSettings):
+    """Tool configuration for MCP server."""
+
+    enable_predictive_tools: bool = Field(
+        default=True,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
+            "ENABLE_PREDICTIVE_TOOLS",
+        ),
+        description="Enable/disable predictive tools",
+    )
+
+    enable_jira_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
+            "ENABLE_JIRA_TOOLS",
+        ),
+        description="Enable/disable Jira tools",
+    )
+
+    enable_confluence_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
+            "ENABLE_CONFLUENCE_TOOLS",
+        ),
+        description="Enable/disable Confluence tools",
+    )
+
+    enable_gdrive_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
+            "ENABLE_GDRIVE_TOOLS",
+        ),
+        description="Enable/disable GDrive tools",
+    )
+
+    enable_microsoft_graph_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_MICROSOFT_GRAPH_TOOLS",
+            "ENABLE_MICROSOFT_GRAPH_TOOLS",
+        ),
+        description="Enable/disable Sharepoint tools",
+    )
+
+    is_atlassian_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
+            "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Atlassian OAuth provider is configured for Atlassian integration",
+    )
+
+    @property
+    def is_atlassian_oauth_configured(self) -> bool:
+        """Check if Atlassian OAuth is configured via provider flag or environment variables."""
+        return self.is_atlassian_oauth_provider_configured or bool(
+            os.getenv("ATLASSIAN_CLIENT_ID") and os.getenv("ATLASSIAN_CLIENT_SECRET")
+        )
+
+    is_google_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Google OAuth provider is configured for Google integration",
+    )
+
+    @property
+    def is_google_oauth_configured(self) -> bool:
+        return self.is_google_oauth_provider_configured or bool(
+            os.getenv("GOOGLE_CLIENT_ID") and os.getenv("GOOGLE_CLIENT_SECRET")
+        )
+
+    is_microsoft_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
+            "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Microsoft OAuth provider is configured for Microsoft integration",
+    )
+
+    @property
+    def is_microsoft_oauth_configured(self) -> bool:
+        return self.is_microsoft_oauth_provider_configured or bool(
+            os.getenv("MICROSOFT_CLIENT_ID") and os.getenv("MICROSOFT_CLIENT_SECRET")
+        )
+
+    @field_validator(
+        "enable_predictive_tools",
+        "enable_jira_tools",
+        "enable_confluence_tools",
+        "enable_gdrive_tools",
+        "enable_microsoft_graph_tools",
+        "is_atlassian_oauth_provider_configured",
+        "is_google_oauth_provider_configured",
+        "is_microsoft_oauth_provider_configured",
+        mode="before",
+    )
+    @classmethod
+    def validate_runtime_params(cls, v: Any) -> Any:
+        """Validate runtime parameters."""
+        return extract_datarobot_runtime_param_payload(v)
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        case_sensitive=False,
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+
 class MCPServerConfig(BaseSettings):
     """MCP Server configuration using pydantic settings."""
 
@@ -188,86 +306,11 @@ class MCPServerConfig(BaseSettings):
         ),
         description="Enable/disable memory management",
     )
-    enable_predictive_tools: bool = Field(
-        default=True,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
-            "ENABLE_PREDICTIVE_TOOLS",
-        ),
-        description="Enable/disable predictive tools",
-    )
 
-    # Jira tools
-    enable_jira_tools: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
-            "ENABLE_JIRA_TOOLS",
-        ),
-        description="Enable/disable Jira tools",
+    tool_config: MCPToolConfig = Field(
+        default_factory=MCPToolConfig,
+        description="Tool configuration",
     )
-    is_jira_oauth_provider_configured: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
-            "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
-        ),
-        description="Whether Jira OAuth provider is configured for Jira integration",
-    )
-
-    @property
-    def is_jira_oauth_configured(self) -> bool:
-        return self.is_jira_oauth_provider_configured or bool(
-            os.getenv("JIRA_CLIENT_ID") and os.getenv("JIRA_CLIENT_SECRET")
-        )
-
-    # Confluence tools
-    enable_confluence_tools: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
-            "ENABLE_CONFLUENCE_TOOLS",
-        ),
-        description="Enable/disable Confluence tools",
-    )
-    is_confluence_oauth_provider_configured: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
-            "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
-        ),
-        description="Whether Confluence OAuth provider is configured for Confluence integration",
-    )
-
-    @property
-    def is_confluence_oauth_configured(self) -> bool:
-        return self.is_confluence_oauth_provider_configured or bool(
-            os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
-        )
-
-    # Gdrive tools
-    enable_gdrive_tools: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
-            "ENABLE_GDRIVE_TOOLS",
-        ),
-        description="Enable/disable GDrive tools",
-    )
-    is_gdrive_oauth_provider_configured: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
-            "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
-        ),
-        description="Whether GDrive OAuth provider is configured for GDrive integration",
-    )
-
-    @property
-    def is_gdrive_oauth_configured(self) -> bool:
-        return self.is_gdrive_oauth_provider_configured or bool(
-            os.getenv("GDRIVE_CLIENT_ID") and os.getenv("GDRIVE_CLIENT_SECRET")
-        )
 
     @field_validator(
         "otel_attributes",
@@ -291,11 +334,6 @@ class MCPServerConfig(BaseSettings):
         "mcp_server_register_dynamic_tools_on_startup",
         "tool_registration_duplicate_behavior",
         "mcp_server_register_dynamic_prompts_on_startup",
-        "enable_predictive_tools",
-        "enable_jira_tools",
-        "is_jira_oauth_provider_configured",
-        "enable_confluence_tools",
-        "is_confluence_oauth_provider_configured",
         mode="before",
     )
     @classmethod

datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -31,9 +31,6 @@ from .dynamic_prompts.register import register_prompts_from_datarobot_prompt_man
 from .dynamic_tools.deployment.register import register_tools_of_datarobot_deployments
 from .logging import MCPLogging
 from .mcp_instance import mcp
-from .mcp_server_tools import get_all_available_tags  # noqa # pylint: disable=unused-import
-from .mcp_server_tools import get_tool_info_by_name  # noqa # pylint: disable=unused-import
-from .mcp_server_tools import list_tools_by_tags  # noqa # pylint: disable=unused-import
 from .memory_management.manager import MemoryManager
 from .routes import register_routes
 from .routes_utils import prefix_mount_path

datarobot_genai/drmcp/core/mcp_instance.py

@@ -16,20 +16,18 @@ import logging
 from collections.abc import Callable
 from functools import wraps
 from typing import Any
-from typing import overload
+from typing import TypedDict
 
 from fastmcp import Context
 from fastmcp import FastMCP
 from fastmcp.exceptions import NotFoundError
 from fastmcp.prompts.prompt import Prompt
 from fastmcp.server.dependencies import get_context
-from fastmcp.tools import FunctionTool
 from fastmcp.tools import Tool
-from fastmcp.utilities.types import NotSet
-from fastmcp.utilities.types import NotSetT
 from mcp.types import AnyFunction
 from mcp.types import Tool as MCPTool
 from mcp.types import ToolAnnotations
+from typing_extensions import Unpack
 
 from .config import MCPServerConfig
 from .config import get_config
@@ -120,86 +118,6 @@ class TaggedFastMCP(FastMCP):
                 "In stateless mode, clients will see changes on next request."
             )
 
-    @overload
-    def tool(
-        self,
-        name_or_fn: AnyFunction,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-    ) -> FunctionTool: ...
-
-    @overload
-    def tool(
-        self,
-        name_or_fn: str | None = None,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-    ) -> Callable[[AnyFunction], FunctionTool]: ...
-
-    def tool(
-        self,
-        name_or_fn: str | Callable[..., Any] | None = None,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-        **kwargs: Any,
-    ) -> Callable[[AnyFunction], FunctionTool] | FunctionTool:
-        """
-        Extend tool decorator that supports tags and other annotations, while remaining
-        signature-compatible with FastMCP.tool to avoid recursion issues with partials.
-        """
-        if isinstance(annotations, dict):
-            annotations = ToolAnnotations(**annotations)
-
-        # Ensure tags are available both via native fastmcp `tags` and inside annotations
-        if tags is not None:
-            tags_ = sorted(tags)
-            if annotations is None:
-                annotations = ToolAnnotations()  # type: ignore[call-arg]
-                annotations.tags = tags_  # type: ignore[attr-defined, union-attr]
-            else:
-                # At this point, annotations is ToolAnnotations (not dict)
-                assert isinstance(annotations, ToolAnnotations)
-                annotations.tags = tags_  # type: ignore[attr-defined]
-
-        return super().tool(
-            name_or_fn,
-            name=name,
-            title=title,
-            description=description,
-            tags=tags,
-            output_schema=output_schema
-            if output_schema is not None
-            else kwargs.get("output_schema"),
-            annotations=annotations,
-            exclude_args=exclude_args,
-            meta=meta,
-            enabled=enabled,
-        )
-
     async def list_tools(
         self, tags: list[str] | None = None, match_all: bool = False
     ) -> list[MCPTool]:
@@ -371,16 +289,37 @@ mcp = TaggedFastMCP(
 )
 
 
+class ToolKwargs(TypedDict, total=False):
+    """Keyword arguments passed through to FastMCP's mcp.tool() decorator.
+
+    All parameters are optional and forwarded directly to FastMCP tool registration.
+    See FastMCP documentation for full details on each parameter.
+    """
+
+    name: str | None
+    title: str | None
+    description: str | None
+    icons: list[Any] | None
+    tags: set[str] | None
+    output_schema: dict[str, Any] | None
+    annotations: Any | None
+    exclude_args: list[str] | None
+    meta: dict[str, Any] | None
+    enabled: bool | None
+
+
 def dr_core_mcp_tool(
-    name: str | None = None,
-    description: str | None = None,
-    tags: set[str] | None = None,
+    **kwargs: Unpack[ToolKwargs],
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
-    """Combine decorator that includes mcp.tool() and dr_mcp_extras()."""
+    """Combine decorator that includes mcp.tool() and dr_mcp_extras().
+
+    All keyword arguments are passed through to FastMCP's mcp.tool() decorator.
+    See ToolKwargs for available parameters.
+    """
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         instrumented = dr_mcp_extras()(func)
-        mcp.tool(name=name, description=description, tags=tags)(instrumented)
+        mcp.tool(**kwargs)(instrumented)
         return instrumented
 
     return decorator
@@ -413,27 +352,23 @@ async def memory_aware_wrapper(func: Callable[..., Any], *args: Any, **kwargs: A
 
 
 def dr_mcp_tool(
-    name: str | None = None,
-    description: str | None = None,
-    tags: set[str] | None = None,
+    **kwargs: Unpack[ToolKwargs],
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     """Combine decorator that includes mcp.tool(), dr_mcp_extras(), and capture memory ids from
     the request headers if they exist.
 
-    Args:
-        name: Tool name
-        description: Tool description
-        tags: Optional set of tags to apply to the tool
+    All keyword arguments are passed through to FastMCP's mcp.tool() decorator.
+    See ToolKwargs for available parameters.
     """
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         @wraps(func)
-        async def wrapper(*args: Any, **kwargs: Any) -> Any:
-            return await memory_aware_wrapper(func, *args, **kwargs)
+        async def wrapper(*args: Any, **inner_kwargs: Any) -> Any:
+            return await memory_aware_wrapper(func, *args, **inner_kwargs)
 
         # Apply the MCP decorators
         instrumented = dr_mcp_extras()(wrapper)
-        mcp.tool(name=name, description=description, tags=tags)(instrumented)
+        mcp.tool(**kwargs)(instrumented)
         return instrumented
 
     return decorator
@@ -488,11 +423,10 @@ async def register_tools(
     # Apply dr_mcp_extras to the memory-aware function
     wrapped_fn = dr_mcp_extras()(memory_aware_fn)
 
-    # Create annotations with tags, deployment_id if provided
-    annotations = ToolAnnotations()  # type: ignore[call-arg]
-    if tags is not None:
-        annotations.tags = tags  # type: ignore[attr-defined]
+    # Create annotations only when additional metadata is required
+    annotations: ToolAnnotations | None = None  # type: ignore[assignment]
     if deployment_id is not None:
+        annotations = ToolAnnotations()  # type: ignore[call-arg]
         annotations.deployment_id = deployment_id  # type: ignore[attr-defined]
 
     tool = Tool.from_function(

datarobot_genai/drmcp/core/tool_config.py

@@ -30,6 +30,7 @@ class ToolType(str, Enum):
     JIRA = "jira"
     CONFLUENCE = "confluence"
     GDRIVE = "gdrive"
+    MICROSOFT_GRAPH = "microsoft_graph"
 
 
 class ToolConfig(TypedDict):
@@ -39,7 +40,7 @@ class ToolConfig(TypedDict):
     oauth_check: Callable[["MCPServerConfig"], bool] | None
     directory: str
     package_prefix: str
-    config_field_name: str  # Name of the config field (e.g., "enable_predictive_tools")
+    config_field_name: str
 
 
 # Tool configuration registry
@@ -53,25 +54,32 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
     ),
     ToolType.JIRA: ToolConfig(
         name="jira",
-        oauth_check=lambda config: config.is_jira_oauth_configured,
+        oauth_check=lambda config: config.tool_config.is_atlassian_oauth_configured,
         directory="jira",
         package_prefix="datarobot_genai.drmcp.tools.jira",
         config_field_name="enable_jira_tools",
     ),
     ToolType.CONFLUENCE: ToolConfig(
         name="confluence",
-        oauth_check=lambda config: config.is_confluence_oauth_configured,
+        oauth_check=lambda config: config.tool_config.is_atlassian_oauth_configured,
         directory="confluence",
         package_prefix="datarobot_genai.drmcp.tools.confluence",
         config_field_name="enable_confluence_tools",
     ),
     ToolType.GDRIVE: ToolConfig(
         name="gdrive",
-        oauth_check=lambda config: config.is_gdrive_oauth_configured,
+        oauth_check=lambda config: config.tool_config.is_google_oauth_configured,
         directory="gdrive",
         package_prefix="datarobot_genai.drmcp.tools.gdrive",
         config_field_name="enable_gdrive_tools",
     ),
+    ToolType.MICROSOFT_GRAPH: ToolConfig(
+        name="microsoft_graph",
+        oauth_check=lambda config: config.tool_config.is_microsoft_oauth_configured,
+        directory="microsoft_graph",
+        package_prefix="datarobot_genai.drmcp.tools.microsoft_graph",
+        config_field_name="enable_microsoft_graph_tools",
+    ),
 }
 
 
@@ -92,12 +100,12 @@ def is_tool_enabled(tool_type: ToolType, config: "MCPServerConfig") -> bool:
     -------
         True if the tool is enabled, False otherwise
     """
-    tool_config = TOOL_CONFIGS[tool_type]
-    enable_config_name = tool_config["config_field_name"]
-    is_enabled = getattr(config, enable_config_name)
+    tool_config_registry = TOOL_CONFIGS[tool_type]
+    enable_config_name = tool_config_registry["config_field_name"]
+    is_enabled = getattr(config.tool_config, enable_config_name)
 
     # If tool is enabled, check OAuth requirements if needed
-    if is_enabled and tool_config["oauth_check"] is not None:
-        return tool_config["oauth_check"](config)
+    if is_enabled and tool_config_registry["oauth_check"] is not None:
+        return tool_config_registry["oauth_check"](config)
 
     return is_enabled

datarobot_genai/drmcp/core/tool_filter.py

@@ -41,7 +41,7 @@ def filter_tools_by_tags(
     filtered_tools = []
 
     for tool in tools:
-        tool_tags = getattr(tool.annotations, "tags", []) if tool.annotations else []
+        tool_tags = get_tool_tags(tool)
 
         if not tool_tags:
             continue
@@ -68,9 +68,18 @@ def get_tool_tags(tool: Tool | MCPTool) -> list[str]:
     -------
         List of tags for the tool
     """
+    # Primary: native FastMCP meta location
+    if hasattr(tool, "meta") and getattr(tool, "meta"):
+        fastmcp_meta = tool.meta.get("_fastmcp", {})
+        meta_tags = fastmcp_meta.get("tags", [])
+        if isinstance(meta_tags, list):
+            return meta_tags
+
+    # Fallback: annotations.tags (for compatibility during transition)
     if tool.annotations and hasattr(tool.annotations, "tags"):
         tags = getattr(tool.annotations, "tags", [])
         return tags if isinstance(tags, list) else []
+
     return []
 
 

datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -39,6 +39,54 @@ class ETETestExpectations(BaseModel):
 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]],
@@ -130,7 +178,26 @@ class ToolBaseE2E:
                             f"result, but got: {response.tool_results[i]}"
                         )
                     else:
-                        actual_result = json.loads(response.tool_results[i])
+                        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]}"