Python: Correction of MCP image type conversion in _mcp.py (#2901)

* Correction of MCP image type conversion in  _mcp.py

* Added a new overload to the init function of the DataContent() type of the Agent Framework, edited the test case to correctly test the usage of the data and uri fields while using DataContent()

* Fixed tests related to the changes of the DataContent type, added testing for both string and byte representations

Author: egeozanozyedek

python/packages/core/agent_framework/_mcp.py

@@ -152,7 +152,7 @@ def _mcp_type_to_ai_content(
             case types.ImageContent() | types.AudioContent():
                 return_types.append(
                     DataContent(
-                        uri=mcp_type.data,
+                        data=mcp_type.data,
                         media_type=mcp_type.mimeType,
                         raw_representation=mcp_type,
                     )

python/packages/core/agent_framework/_types.py

@@ -925,6 +925,10 @@ class DataContent(BaseContent):
             image_data = b"raw image bytes"
             data_content = DataContent(data=image_data, media_type="image/png")
 
+            # Create from base64-encoded string
+            base64_string = "iVBORw0KGgoAAAANS..."
+            data_content = DataContent(data=base64_string, media_type="image/png")
+
             # Create from data URI
             data_uri = "data:image/png;base64,iVBORw0KGgoAAAANS..."
             data_content = DataContent(uri=data_uri)
@@ -986,11 +990,38 @@ def __init__(
             **kwargs: Any additional keyword arguments.
         """
 
+    @overload
+    def __init__(
+        self,
+        *,
+        data: str,
+        media_type: str,
+        annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        raw_representation: Any | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes a DataContent instance with base64-encoded string data.
+
+        Important:
+            This is for binary data that is represented as a data URI, not for online resources.
+            Use ``UriContent`` for online resources.
+
+        Keyword Args:
+            data: The base64-encoded string data represented by this instance.
+                The data is used directly to construct a data URI.
+            media_type: The media type of the data.
+            annotations: Optional annotations associated with the content.
+            additional_properties: Optional additional properties associated with the content.
+            raw_representation: Optional raw representation of the content.
+            **kwargs: Any additional keyword arguments.
+        """
+
     def __init__(
         self,
         *,
         uri: str | None = None,
-        data: bytes | None = None,
+        data: bytes | str | None = None,
         media_type: str | None = None,
         annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
         additional_properties: dict[str, Any] | None = None,
@@ -1006,8 +1037,9 @@ def __init__(
         Keyword Args:
             uri: The URI of the data represented by this instance.
                 Should be in the form: "data:{media_type};base64,{base64_data}".
-            data: The binary data represented by this instance.
-                The data is transformed into a base64-encoded data URI.
+            data: The binary data or base64-encoded string represented by this instance.
+                If bytes, the data is transformed into a base64-encoded data URI.
+                If str, it is assumed to be already base64-encoded and used directly.
             media_type: The media type of the data.
             annotations: Optional annotations associated with the content.
             additional_properties: Optional additional properties associated with the content.
@@ -1017,7 +1049,9 @@ def __init__(
         if uri is None:
             if data is None or media_type is None:
                 raise ValueError("Either 'data' and 'media_type' or 'uri' must be provided.")
-            uri = f"data:{media_type};base64,{base64.b64encode(data).decode('utf-8')}"
+
+            base64_data: str = base64.b64encode(data).decode("utf-8") if isinstance(data, bytes) else data
+            uri = f"data:{media_type};base64,{base64_data}"
 
         # Validate URI format and extract media type if not provided
         validated_uri = self._validate_uri(uri)

python/packages/core/tests/core/test_mcp.py

@@ -75,17 +75,21 @@ def test_mcp_call_tool_result_to_ai_contents():
     mcp_result = types.CallToolResult(
         content=[
             types.TextContent(type="text", text="Result text"),
-            types.ImageContent(type="image", data="data:image/png;base64,xyz", mimeType="image/png"),
+            types.ImageContent(type="image", data="xyz", mimeType="image/png"),
+            types.ImageContent(type="image", data=b"abc", mimeType="image/webp"),
         ]
     )
     ai_contents = _mcp_call_tool_result_to_ai_contents(mcp_result)
 
-    assert len(ai_contents) == 2
+    assert len(ai_contents) == 3
     assert isinstance(ai_contents[0], TextContent)
     assert ai_contents[0].text == "Result text"
     assert isinstance(ai_contents[1], DataContent)
     assert ai_contents[1].uri == "data:image/png;base64,xyz"
     assert ai_contents[1].media_type == "image/png"
+    assert isinstance(ai_contents[2], DataContent)
+    assert ai_contents[2].uri == "data:image/webp;base64,abc"
+    assert ai_contents[2].media_type == "image/webp"
 
 
 def test_mcp_call_tool_result_with_meta_error():
@@ -183,7 +187,7 @@ def test_mcp_call_tool_result_regression_successful_workflow():
     mcp_result = types.CallToolResult(
         content=[
             types.TextContent(type="text", text="Success message"),
-            types.ImageContent(type="image", data="data:image/jpeg;base64,abc123", mimeType="image/jpeg"),
+            types.ImageContent(type="image", data="abc123", mimeType="image/jpeg"),
         ]
     )
 
@@ -218,7 +222,8 @@ def test_mcp_content_types_to_ai_content_text():
 
 def test_mcp_content_types_to_ai_content_image():
     """Test conversion of MCP image content to AI content."""
-    mcp_content = types.ImageContent(type="image", data="data:image/jpeg;base64,abc", mimeType="image/jpeg")
+    mcp_content = types.ImageContent(type="image", data="abc", mimeType="image/jpeg")
+    mcp_content = types.ImageContent(type="image", data=b"abc", mimeType="image/jpeg")
     ai_content = _mcp_type_to_ai_content(mcp_content)[0]
 
     assert isinstance(ai_content, DataContent)
@@ -229,7 +234,7 @@ def test_mcp_content_types_to_ai_content_image():
 
 def test_mcp_content_types_to_ai_content_audio():
     """Test conversion of MCP audio content to AI content."""
-    mcp_content = types.AudioContent(type="audio", data="data:audio/wav;base64,def", mimeType="audio/wav")
+    mcp_content = types.AudioContent(type="audio", data="def", mimeType="audio/wav")
     ai_content = _mcp_type_to_ai_content(mcp_content)[0]
 
     assert isinstance(ai_content, DataContent)