diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 23926658fb..d5f7d14d92 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -122,7 +122,11 @@ 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 - data_str = data_uri.split(";base64,", 1)[1] + if not data_uri.startswith("data:") or "," not in data_uri: + raise ValueError("Invalid data URI format.") + prefix, data_str = data_uri.split(",", 1) + if not prefix.endswith(";base64"): + raise ValueError("Data URI must use base64 encoding.") if data_str is not None: if data is not None: raise ValueError("Provide exactly one of data_bytes, data_str, or data_uri.") diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 19c77585e4..9b31dee238 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -208,6 +208,20 @@ def test_data_content_detect_image_format_from_base64(): detect_media_type_from_base64(data_str="data", data_uri="data:application/octet-stream;base64,AAA") +@mark.parametrize( + ("data_uri", "error_message"), + [ + ("data:text/plain,hello", "Data URI must use base64 encoding."), + ("data:image/png;base64", "Invalid data URI format."), + ("not-a-data-uri", "Invalid data URI format."), + ], +) +def test_detect_media_type_from_base64_rejects_malformed_data_uri(data_uri: str, error_message: str): + """Test malformed data URI inputs raise the documented ValueError.""" + with raises(ValueError, match=error_message): + detect_media_type_from_base64(data_uri=data_uri) + + def test_data_content_create_data_uri_from_base64(): """Test the create_data_uri_from_base64 class method.""" # Test with PNG data