lf.Mime.from_uri to support data URLs.

E.g.:

```
lf.Mime.from_uri('data:text/plain;base64,...')
```

PiperOrigin-RevId: 739602781

Author: daiyip

langfun/core/modalities/mime.py

@@ -152,12 +152,33 @@ def embeddable_uri(self) -> str:
 
   @classmethod
   def from_uri(cls, uri: str, **kwargs) -> 'Mime':
+    if uri.startswith('data:'):
+      mime_type, content = cls._parse_data_uri(uri)
+      return cls.class_from_mime_type(mime_type).from_bytes(content, **kwargs)
+
     if cls is Mime:
       content = cls.download(uri)
       mime = from_buffer(content, mime=True).lower()
       return cls.class_from_mime_type(mime)(uri=uri, content=content, **kwargs)
     return cls(uri=uri, content=None, **kwargs)
 
+  @classmethod
+  def _parse_data_uri(cls, uri: str) -> tuple[str, bytes]:
+    """Returns the MIME type and content from the given data URI."""
+    assert uri.startswith('data:'), uri
+    mime_end_pos = uri.find(';', 0)
+    if mime_end_pos == -1:
+      raise ValueError(f'Invalid data URI: {uri!r}.')
+    mime_type = uri[5: mime_end_pos].strip().lower()
+    encoding_end_pos = uri.find(',', mime_end_pos + 1)
+    if encoding_end_pos == -1:
+      raise ValueError(f'Invalid data URI: {uri!r}.')
+    encoding = uri[mime_end_pos + 1: encoding_end_pos].strip().lower()
+    if encoding != 'base64':
+      raise ValueError(f'Unsupported encoding: {encoding!r}.')
+    base64_content = uri[encoding_end_pos + 1:].strip().encode()
+    return mime_type, base64.b64decode(base64_content)
+
   @classmethod
   def from_bytes(cls, content: bytes | str, **kwargs) -> 'Mime':
     if cls is Mime:

langfun/core/modalities/mime_test.py

@@ -93,6 +93,19 @@ def test_from_uri(self):
       self.assertEqual(content.to_bytes(), b'bar')
       self.assertEqual(content.mime_type, 'text/plain')
 
+    content = mime.Mime.from_uri('data:text/plain;base64,Zm9v')
+    self.assertIsNone(content.uri)
+    self.assertEqual(content.mime_type, 'text/plain')
+    self.assertEqual(content.content, b'foo')
+    self.assertEqual(content.content_uri, 'data:text/plain;base64,Zm9v')
+    self.assertEqual(content.embeddable_uri, 'data:text/plain;base64,Zm9v')
+
+    with self.assertRaisesRegex(ValueError, 'Invalid data URI'):
+      mime.Mime.from_uri('data:text/plain')
+
+    with self.assertRaisesRegex(ValueError, 'Invalid data URI'):
+      mime.Mime.from_uri('data:text/plain;base64,')
+
   def assert_html_content(self, html, expected):
     expected = inspect.cleandoc(expected).strip()
     actual = html.content.strip()