diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 736474de60..879f54fcf7 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -121,6 +121,8 @@ def detect_media_type_from_base64( if data is not None: raise ValueError("Provide exactly one of data_bytes, data_str, or data_uri.") # Remove data URI prefix if present + if ";base64," not in data_uri: + raise ValueError("Data URI must use base64 encoding.") data_str = data_uri.split(";base64,", 1)[1] if data_str is not None: if data is not None: diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 23c7dd8cd6..1a1549ad9e 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -204,6 +204,11 @@ def test_data_content_detect_image_format_from_base64(): detect_media_type_from_base64(data_bytes=b"data", data_uri="data:application/octet-stream;base64,AAA") detect_media_type_from_base64(data_str="data", data_uri="data:application/octet-stream;base64,AAA") + with pytest.raises(ValueError, match="Data URI must use base64 encoding"): + detect_media_type_from_base64(data_uri="data:text/plain,Hello%20World") + with pytest.raises(ValueError, match="Data URI must use base64 encoding"): + detect_media_type_from_base64(data_uri="http://example.com/image.png") + def test_data_content_create_data_uri_from_base64(): """Test the create_data_uri_from_base64 class method."""