langfun 0.0.2.dev20240429__py3-none-any.whl → 0.1.2.dev202501140804__py3-none-any.whl

langfun/core/modalities/mime.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The Langfun Authors
+# Copyright 2024 The Langfun Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,20 +13,30 @@
 # limitations under the License.
 """MIME type data."""
 
-import abc
-from typing import Annotated, Union
+import base64
+import functools
+from typing import Annotated, Any, Iterable, Type, Union
 import langfun.core as lf
 import pyglove as pg
 import requests
 
 
-class MimeType(lf.Modality):
-  """Base for MIME type data."""
+try:
+  import magic    # pylint: disable=g-import-not-at-top
+  from_buffer = magic.from_buffer
+except ImportError:
+  def from_buffer(*unused_args, **unused_kwargs):
+    raise RuntimeError(
+        'Please install "langfun[mime-auto]" to enable automatic MIME support.'
+    )
 
-  @property
-  @abc.abstractmethod
-  def mime_type(self) -> str:
-    """Returns the MIME type."""
+
+class Mime(lf.Modality):
+  """Base for MIME data."""
+
+  # The regular expression that describes the MIME type str.
+  # If None, the MIME type is dynamic. Subclass could override.
+  MIME_PREFIX = None
 
   uri: Annotated[str | None, 'The URI for locating the MIME data. '] = None
 
@@ -34,6 +44,86 @@ class MimeType(lf.Modality):
       Union[str, bytes, None], 'The raw content of the MIME type.'
   ] = None
 
+  @functools.cached_property
+  def mime_type(self) -> str:
+    """Returns the MIME type."""
+    mime = from_buffer((self.to_bytes()), mime=True)
+    if (
+        self.MIME_PREFIX
+        and not mime.lower().startswith(self.MIME_PREFIX)
+        # NOTE(daiyip): libmagic fails to detect the MIME type of some binary
+        # files.
+        and mime != 'application/octet-stream'
+    ):
+      raise ValueError(
+          f'Expected MIME type: {self.MIME_PREFIX}, Encountered: {mime}'
+      )
+    return mime
+
+  @functools.cached_property
+  def is_text(self) -> bool:
+    return self.mime_type.startswith((
+        'text/',
+        'application/javascript',
+        'application/json',
+        'application/ld+json',
+        'application/plain',
+        'application/rtf',
+        'application/xhtml+xml',
+        'application/xml',
+        'application/x-javascript',
+        'application/x-python-code',
+        'application/x-tex',
+        'application/x-typescript',
+        'application/x-yaml',
+    ))
+
+  @property
+  def is_binary(self) -> bool:
+    """Returns True if the MIME type is a binary type."""
+    return not self.is_text
+
+  def to_text(self) -> str:
+    """Returns the text content of the MIME type."""
+    if not self.is_text:
+      raise lf.ModalityError(
+          f'MIME type {self.mime_type!r} cannot be converted to text.'
+      )
+    return self.to_bytes().decode()
+
+  def is_compatible(
+      self, mime_types: str | Iterable[str]
+  ) -> bool:
+    """Returns True if this object is compatible to any of the MIME types."""
+    if isinstance(mime_types, str):
+      mime_types = {mime_types}
+    return self._is_compatible(mime_types)
+
+  def _is_compatible(self, mime_types: Iterable[str]):
+    return self.mime_type in mime_types
+
+  def make_compatible(
+      self,
+      mime_types: str | Iterable[str]
+      ) -> Union['Mime', list['Mime']]:
+    """Makes compatible MIME objects from this object."""
+    if isinstance(mime_types, str):
+      mime_types = {mime_types}
+    if not self._is_compatible(mime_types):
+      raise lf.ModalityError(
+          f'MIME type {self.mime_type!r} cannot be converted to supported '
+          f'types: {mime_types!r}.'
+      )
+    return self._make_compatible(mime_types)
+
+  def _make_compatible(
+      self,
+      mime_types: Iterable[str]
+  ) -> Union['Mime', list['Mime']]:
+    """Makes compatbile MIME objects from this object."""
+    del mime_types
+    return self
+
   def _on_bound(self):
     super()._on_bound()
     if self.uri is None and self.content is None:
@@ -43,39 +133,126 @@ class MimeType(lf.Modality):
     if self.content is not None:
       return self.content
 
-    assert self.uri is not None
-    if self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
-      content = requests.get(
-          self.uri,
-          headers={'User-Agent': 'Langfun'},
-      ).content
-    else:
-      content = pg.io.readfile(self.uri, mode='rb')
-    self.rebind(content=content, skip_notification=True)
+    self.rebind(content=self.download(self.uri), skip_notification=True)
     return self.content
 
+  @property
+  def content_uri(self) -> str:
+    """Returns the URI with encoded content."""
+    base64_content = base64.b64encode(self.to_bytes()).decode()
+    return f'data:{self.mime_type};base64,{base64_content}'
+
+  @property
+  def embeddable_uri(self) -> str:
+    """Returns the URI that can be embedded in HTML."""
+    if self.uri and self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
+      return self.uri
+    return self.content_uri
+
   @classmethod
-  def from_uri(cls, uri: str, **kwargs) -> 'MimeType':
+  def from_uri(cls, uri: str, **kwargs) -> 'Mime':
+    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 from_bytes(cls, content: bytes | str, **kwargs) -> 'MimeType':
+  def from_bytes(cls, content: bytes | str, **kwargs) -> 'Mime':
+    if cls is Mime:
+      mime = from_buffer(content, mime=True).lower()
+      return cls.class_from_mime_type(mime)(content=content, **kwargs)
     return cls(content=content, **kwargs)
 
+  @classmethod
+  def class_from_mime_type(cls, mime_type: str) -> Type['Mime']:
+    """Subclass from the given MIME type."""
+    for subcls in cls.__subclasses__():
+      if subcls.MIME_PREFIX is not None and mime_type.startswith(
+          subcls.MIME_PREFIX):
+        return subcls
+    return cls
+
+  @classmethod
+  def download(cls, uri: str) -> bytes | str:
+    """Downloads the content of the given URI."""
+    if uri.lower().startswith(('http:', 'https:', 'ftp:')):
+      return requests.get(
+          uri,
+          headers={'User-Agent': 'Mozilla/5.0'},
+      ).content
+    else:
+      content = pg.io.readfile(uri, mode='rb')
+      assert content is not None
+      return content
+
+  def _html_tree_view_content(
+      self,
+      **kwargs) -> str:
+    return self._raw_html()
 
-@pg.use_init_args(['type', 'content', 'uri'])
-class Custom(MimeType):
+  def _html_tree_view(
+      self,
+      view: pg.views.HtmlTreeView,
+      extra_flags: dict[str, Any] | None = None,
+      **kwargs
+  ):
+    extra_flags = extra_flags if extra_flags is not None else {}
+    raw_mime_content = extra_flags.get('raw_mime_content', False)
+    display_modality_when_hover = extra_flags.get(
+        'display_modality_when_hover', False
+    )
+    if raw_mime_content:
+      kwargs['enable_summary'] = False
+    elif display_modality_when_hover:
+      kwargs.update(
+          enable_summary=True,
+          enable_summary_tooltip=True,
+      )
+    return super()._html_tree_view(
+        view=view, extra_flags=extra_flags, **kwargs
+    )
+
+  def _html_tree_view_summary(
+      self,
+      *,
+      view: pg.views.HtmlTreeView,
+      extra_flags: dict[str, Any] | None = None,
+      **kwargs
+  ):
+    extra_flags = extra_flags or {}
+    if extra_flags.get('display_modality_when_hover', False):
+      def summary_tooltip(*args, content: str | None = None, **kwargs):
+        del content
+        return view.tooltip(*args, content=self._raw_html(), **kwargs)
+    else:
+      summary_tooltip = None
+    return super()._html_tree_view_summary(
+        view=view,
+        summary_tooltip_fn=summary_tooltip,
+        extra_flags=extra_flags,
+        **kwargs
+    )
+
+  def _raw_html(self) -> str:
+    if self.uri and self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
+      uri = self.uri
+    else:
+      uri = self.content_uri
+    return self._mime_control_for(uri)
+
+  def _mime_control_for(self, uri) -> str:
+    return f'<embed type="{self.mime_type}" src="{uri}"/>'
+
+
+@pg.use_init_args(['mime', 'content', 'uri'])
+class Custom(Mime):
   """Custom MIME data."""
 
-  type: Annotated[
+  mime: Annotated[
       str, 'The MIME type of the data. E.g. text/plain, or image/png. '
   ]
 
   @property
   def mime_type(self) -> str:
-    return self.type
-
-
-class PDF(Custom):
-  """PDF document."""
-  type = 'application/pdf'
+    return self.mime

langfun/core/modalities/mime_test.py

@@ -12,48 +12,129 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """MIME tests."""
+import inspect
 import unittest
 from unittest import mock
 
+import langfun.core as lf
 from langfun.core.modalities import mime
 import pyglove as pg
 
 
 def mock_request(*args, **kwargs):
   del args, kwargs
-  return pg.Dict(content='foo')
+  return pg.Dict(content=b'foo')
 
 
 def mock_readfile(*args, **kwargs):
   del args, kwargs
-  return 'bar'
+  return b'bar'
 
 
 class CustomMimeTest(unittest.TestCase):
 
-  def test_content(self):
-    content = mime.Custom('text/plain', 'foo')
-    self.assertEqual(content.to_bytes(), 'foo')
+  def test_is_text(self):
+    self.assertTrue(mime.Custom('text/plain', b'foo').is_text)
+    self.assertTrue(mime.Custom('text/xml', b'foo').is_text)
+    self.assertTrue(mime.Custom('application/json', b'foo').is_text)
+    self.assertTrue(mime.Custom('application/x-python-code', b'foo').is_text)
+    self.assertFalse(mime.Custom('application/pdf', b'foo').is_text)
+    self.assertFalse(mime.Custom('application/octet-stream', b'foo').is_text)
+
+  def test_from_byes(self):
+    content = mime.Mime.from_bytes(b'hello')
+    self.assertIs(content.__class__, mime.Mime)
+
+    content = mime.Custom('text/plain', b'foo')
+    self.assertEqual(content.to_bytes(), b'foo')
     self.assertEqual(content.mime_type, 'text/plain')
+    self.assertTrue(content.is_text)
+    self.assertFalse(content.is_binary)
+    self.assertEqual(content.to_text(), 'foo')
+    self.assertTrue(content.is_compatible('text/plain'))
+    self.assertFalse(content.is_compatible('text/xml'))
+    self.assertIs(content.make_compatible('text/plain'), content)
+
+    with self.assertRaisesRegex(
+        lf.ModalityError, '.* cannot be converted to supported types'
+    ):
+      content.make_compatible('application/pdf')
 
     with self.assertRaisesRegex(
         ValueError, 'Either uri or content must be provided.'
     ):
       mime.Custom('text/plain')
 
+  def test_uri(self):
+    content = mime.Custom.from_uri('http://mock/web/a.txt', mime='text/plain')
+    with mock.patch('requests.get') as mock_requests_stub:
+      mock_requests_stub.side_effect = mock_request
+      self.assertEqual(content.uri, 'http://mock/web/a.txt')
+      self.assertEqual(content.content_uri, 'data:text/plain;base64,Zm9v')
+      self.assertEqual(content.embeddable_uri, 'http://mock/web/a.txt')
+
+    content = mime.Custom.from_uri('a.txt', mime='text/plain')
+    with mock.patch('pyglove.io.readfile') as mock_readfile_stub:
+      mock_readfile_stub.side_effect = mock_readfile
+      self.assertEqual(content.uri, 'a.txt')
+      self.assertEqual(content.content_uri, 'data:text/plain;base64,YmFy')
+      self.assertEqual(content.embeddable_uri, 'data:text/plain;base64,YmFy')
+
   def test_from_uri(self):
-    content = mime.Custom.from_uri('http://mock/web/a.txt', type='text/plain')
+    content = mime.Custom.from_uri('http://mock/web/a.txt', mime='text/plain')
     with mock.patch('requests.get') as mock_requests_stub:
       mock_requests_stub.side_effect = mock_request
-      self.assertEqual(content.to_bytes(), 'foo')
+      self.assertEqual(content.to_bytes(), b'foo')
       self.assertEqual(content.mime_type, 'text/plain')
 
-    content = mime.Custom.from_uri('a.txt', type='text/plain')
+    content = mime.Custom.from_uri('a.txt', mime='text/plain')
     with mock.patch('pyglove.io.readfile') as mock_readfile_stub:
       mock_readfile_stub.side_effect = mock_readfile
-      self.assertEqual(content.to_bytes(), 'bar')
+      self.assertEqual(content.to_bytes(), b'bar')
       self.assertEqual(content.mime_type, 'text/plain')
 
+  def assert_html_content(self, html, expected):
+    expected = inspect.cleandoc(expected).strip()
+    actual = html.content.strip()
+    if actual != expected:
+      print(actual)
+    self.assertEqual(actual, expected)
+
+  def test_html(self):
+    self.assert_html_content(
+        mime.Custom('text/plain', b'foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+        ),
+        """
+        <details open class="pyglove custom"><summary><div class="summary-title">Custom(...)</div></summary><embed type="text/plain" src="data:text/plain;base64,Zm9v"/></details>
+        """
+    )
+    self.assert_html_content(
+        mime.Custom('text/plain', b'foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+            extra_flags=dict(
+                raw_mime_content=True,
+            )
+        ),
+        """
+        <embed type="text/plain" src="data:text/plain;base64,Zm9v"/>
+        """
+    )
+    self.assert_html_content(
+        mime.Custom('text/plain', b'foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+            extra_flags=dict(
+                display_modality_when_hover=True,
+            )
+        ),
+        """
+        <details open class="pyglove custom"><summary><div class="summary-title">Custom(...)</div><span class="tooltip"><embed type="text/plain" src="data:text/plain;base64,Zm9v"/></span></summary><embed type="text/plain" src="data:text/plain;base64,Zm9v"/></details>
+        """
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/modalities/ms_office.py

@@ -0,0 +1,117 @@
+# Copyright 2023 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Microsoft Office file types."""
+
+import base64
+import io
+import os
+from typing import Iterable
+from langfun.core.modalities import mime
+from langfun.core.modalities import pdf
+import requests
+
+
+class Xlsx(mime.Mime):
+  """Xlsx file type."""
+
+  MIME_PREFIX = (
+      'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
+  )
+
+  def _raw_html(self) -> str:
+    try:
+      import pandas as pd  # pylint: disable=g-import-not-at-top
+      import openpyxl      # pylint: disable=g-import-not-at-top, unused-import
+      df = pd.read_excel(io.BytesIO(self.to_bytes()))
+      return df.to_html()
+    except ImportError as e:
+      raise RuntimeError(
+          'Please install "langfun[mime-xlsx]" to enable XLSX support.'
+      ) from e
+
+  def _is_compatible(self, mime_types: Iterable[str]) -> bool:
+    return bool(set(mime_types).intersection([
+        'text/html',
+        'text/plain',
+    ]))
+
+  def _make_compatible(self, mime_types: Iterable[str]) -> mime.Mime:
+    """Returns the MimeType of the converted file."""
+    del mime_types
+    return mime.Mime(uri=self.uri, content=self._raw_html())
+
+
+class Docx(mime.Mime):
+  """Docx file type."""
+
+  MIME_PREFIX = (
+      'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
+  )
+
+  def to_xml(self) -> str:
+    try:
+      import docx  # pylint: disable=g-import-not-at-top
+      doc = docx.Document(io.BytesIO(self.to_bytes()))
+      return str(doc.element.xml)
+    except ImportError as e:
+      raise RuntimeError(
+          'Please install "langfun[mime-docx]" to enable Docx support.'
+      ) from e
+
+  def _repr_html_(self) -> str:
+    return self.to_xml()
+
+  def _is_compatible(self, mime_types: Iterable[str]) -> bool:
+    return bool(set(mime_types).intersection([
+        'application/xml',
+        'text/xml',
+        'text/plain',
+    ]))
+
+  def _make_compatible(self, mime_types: Iterable[str]) -> mime.Mime:
+    """Returns the MimeType of the converted file."""
+    del mime_types
+    return mime.Mime(uri=self.uri, content=self.to_xml())
+
+
+class Pptx(mime.Mime):
+  """Pptx file type."""
+
+  MIME_PREFIX = 'application/vnd.openxmlformats-officedocument.presentationml.presentation'
+  API_URL = 'https://v2.convertapi.com/convert/pptx/to/pdf'
+
+  def to_pdf(self, convert_api_key: str | None = None) -> pdf.PDF:
+    api_key = convert_api_key or os.environ.get('CONVERT_API_KEY')
+    url = f'{self.API_URL}?Secret={api_key}'
+
+    json = {
+        'Parameters': [{
+            'Name': 'File',
+            'FileValue': {
+                'Name': os.path.basename(self.uri) if self.uri else 'tmp.pptx',
+                'Data': base64.b64encode(self.to_bytes()).decode('utf-8'),
+            },
+        }]
+    }
+    response = requests.post(url, json=json).json()
+    base64_pdf = response['Files'][0]['FileData']
+    return pdf.PDF.from_bytes(base64.b64decode(base64_pdf))
+
+  def _is_compatible(self, mime_types: Iterable[str]) -> bool:
+    return 'application/pdf' in mime_types
+
+  def _make_compatible(self, mime_types: Iterable[str]) -> mime.Mime:
+    """Returns the MimeType of the converted file."""
+    del mime_types
+    return self.to_pdf()