diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 2ea0f21dec..37e0d2c54b 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/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, ) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 26e4358fb1..ab68382a83 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/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) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 813667bb7a..93643da30f 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/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)