langfun 0.0.2.dev20240531__tar.gz → 0.0.2.dev20240601__tar.gz

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.0.2.dev20240531
+Version: 0.0.2.dev20240601
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -25,13 +25,14 @@ Requires-Dist: google-cloud-aiplatform>=1.5.0
 Requires-Dist: google-generativeai>=0.3.2
 Requires-Dist: jinja2>=3.1.2
 Requires-Dist: openai==0.27.2
+Requires-Dist: openpyxl>=3.1.0
+Requires-Dist: pandas>=2.1.4
 Requires-Dist: pyglove>=0.4.5.dev20240423
+Requires-Dist: python-docx>=0.8.11
 Requires-Dist: python-magic>=0.4.27
 Requires-Dist: requests>=2.31.0
 Requires-Dist: termcolor==1.1.0
 Requires-Dist: tqdm>=4.64.1
-Requires-Dist: python-docx>=0.8.11
-Requires-Dist: pandas>=2.1.4
 
 <div align="center">
 <img src="https://raw.githubusercontent.com/google/langfun/main/docs/_static/logo.svg" width="520px" alt="logo"></img>

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/__init__.py

@@ -57,6 +57,8 @@ from langfun.core import memories
 
 from langfun.core import modalities
 
+Mime = modalities.Mime
+MimeType = Mime  # For backwards compatibility.
 Image = modalities.Image
 Video = modalities.Video
 PDF = modalities.PDF

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/__init__.py

@@ -94,6 +94,7 @@ from langfun.core.message import MemoryRecord
 # Interface for modality.
 from langfun.core.modality import Modality
 from langfun.core.modality import ModalityRef
+from langfun.core.modality import ModalityError
 
 # Interfaces for languge models.
 from langfun.core.language_model import LanguageModel

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/llms/google_genai.py

@@ -49,9 +49,10 @@ class GenAI(lf.LanguageModel):
       ),
   ] = None
 
-  multimodal: Annotated[bool, 'Whether this model has multimodal support.'] = (
-      False
-  )
+  supported_modalities: Annotated[
+      list[str],
+      'A list of MIME types for supported modalities'
+  ] = []
 
   # Set the default max concurrency to 8 workers.
   max_concurrency = 8
@@ -118,14 +119,27 @@ class GenAI(lf.LanguageModel):
     chunks = []
     for lf_chunk in formatted.chunk():
       if isinstance(lf_chunk, str):
-        chunk = lf_chunk
-      elif self.multimodal and isinstance(lf_chunk, lf_modalities.MimeType):
-        chunk = genai.types.BlobDict(
-            data=lf_chunk.to_bytes(), mime_type=lf_chunk.mime_type
-        )
+        chunks.append(lf_chunk)
+      elif isinstance(lf_chunk, lf_modalities.Mime):
+        try:
+          modalities = lf_chunk.make_compatible(
+              self.supported_modalities + ['text/plain']
+          )
+          if isinstance(modalities, lf_modalities.Mime):
+            modalities = [modalities]
+          for modality in modalities:
+            if modality.is_text:
+              chunk = modality.to_text()
+            else:
+              chunk = genai.types.BlobDict(
+                  data=modality.to_bytes(),
+                  mime_type=modality.mime_type
+              )
+            chunks.append(chunk)
+        except lf.ModalityError as e:
+          raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}') from e
       else:
-        raise ValueError(f'Unsupported modality: {lf_chunk!r}')
-      chunks.append(chunk)
+        raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}')
     return chunks
 
   def _response_to_result(
@@ -264,18 +278,57 @@ _GOOGLE_GENAI_MODEL_HUB = _ModelHub()
 #
 
 
+_IMAGE_TYPES = [
+    'image/png',
+    'image/jpeg',
+    'image/webp',
+    'image/heic',
+    'image/heif',
+]
+
+_AUDIO_TYPES = [
+    'audio/aac',
+    'audio/flac',
+    'audio/mp3',
+    'audio/m4a',
+    'audio/mpeg',
+    'audio/mpga',
+    'audio/mp4',
+    'audio/opus',
+    'audio/pcm',
+    'audio/wav',
+    'audio/webm'
+]
+
+_VIDEO_TYPES = [
+    'video/mov',
+    'video/mpeg',
+    'video/mpegps',
+    'video/mpg',
+    'video/mp4',
+    'video/webm',
+    'video/wmv',
+    'video/x-flv',
+    'video/3gpp',
+]
+
+_PDF = [
+    'application/pdf',
+]
+
+
 class GeminiPro1_5(GenAI):  # pylint: disable=invalid-name
   """Gemini Pro latest model."""
 
   model = 'gemini-1.5-pro-latest'
-  multimodal = True
+  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
 
 
 class GeminiFlash1_5(GenAI):  # pylint: disable=invalid-name
   """Gemini Flash latest model."""
 
   model = 'gemini-1.5-flash-latest'
-  multimodal = True
+  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
 
 
 class GeminiPro(GenAI):
@@ -288,7 +341,7 @@ class GeminiProVision(GenAI):
   """Gemini Pro vision model."""
 
   model = 'gemini-pro-vision'
-  multimodal = True
+  supported_modalities = _IMAGE_TYPES + _VIDEO_TYPES
 
 
 class Palm2(GenAI):

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/llms/google_genai_test.py

@@ -107,7 +107,7 @@ class GenAITest(unittest.TestCase):
     )
 
     # Non-multimodal model.
-    with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
+    with self.assertRaisesRegex(lf.ModalityError, 'Unsupported modality'):
       google_genai.GeminiPro()._content_from_message(message)
 
     model = google_genai.GeminiProVision()

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/llms/vertexai.py

@@ -75,9 +75,10 @@ class VertexAI(lf.LanguageModel):
       ),
   ] = None
 
-  multimodal: Annotated[bool, 'Whether this model has multimodal support.'] = (
-      False
-  )
+  supported_modalities: Annotated[
+      list[str],
+      'A list of MIME types for supported modalities'
+  ] = []
 
   def _on_bound(self):
     super()._on_bound()
@@ -142,16 +143,29 @@ class VertexAI(lf.LanguageModel):
     """Gets generation input from langfun message."""
     from vertexai import generative_models
     chunks = []
+
     for lf_chunk in prompt.chunk():
       if isinstance(lf_chunk, str):
-        chunk = lf_chunk
-      elif self.multimodal and isinstance(lf_chunk, lf_modalities.MimeType):
-        chunk = generative_models.Part.from_data(
-            lf_chunk.to_bytes(), lf_chunk.mime_type
-        )
+        chunks.append(lf_chunk)
+      elif isinstance(lf_chunk, lf_modalities.Mime):
+        try:
+          modalities = lf_chunk.make_compatible(
+              self.supported_modalities + ['text/plain']
+          )
+          if isinstance(modalities, lf_modalities.Mime):
+            modalities = [modalities]
+          for modality in modalities:
+            if modality.is_text:
+              chunk = modality.to_text()
+            else:
+              chunk = generative_models.Part.from_data(
+                  modality.to_bytes(), modality.mime_type
+              )
+            chunks.append(chunk)
+        except lf.ModalityError as e:
+          raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}') from e
       else:
-        raise ValueError(f'Unsupported modality: {lf_chunk!r}')
-      chunks.append(chunk)
+        raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}')
     return chunks
 
   def _generation_response_to_message(
@@ -265,25 +279,64 @@ class _ModelHub:
 _VERTEXAI_MODEL_HUB = _ModelHub()
 
 
+_IMAGE_TYPES = [
+    'image/png',
+    'image/jpeg',
+    'image/webp',
+    'image/heic',
+    'image/heif',
+]
+
+_AUDIO_TYPES = [
+    'audio/aac',
+    'audio/flac',
+    'audio/mp3',
+    'audio/m4a',
+    'audio/mpeg',
+    'audio/mpga',
+    'audio/mp4',
+    'audio/opus',
+    'audio/pcm',
+    'audio/wav',
+    'audio/webm'
+]
+
+_VIDEO_TYPES = [
+    'video/mov',
+    'video/mpeg',
+    'video/mpegps',
+    'video/mpg',
+    'video/mp4',
+    'video/webm',
+    'video/wmv',
+    'video/x-flv',
+    'video/3gpp',
+]
+
+_PDF = [
+    'application/pdf',
+]
+
+
 class VertexAIGeminiPro1_5(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-preview-0514'
-  multimodal = True
+  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
 
 
 class VertexAIGeminiPro1_5_0409(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-preview-0409'
-  multimodal = True
+  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
 
 
 class VertexAIGeminiFlash1_5(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-preview-0514'
-  multimodal = True
+  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
 
 
 class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
@@ -296,7 +349,7 @@ class VertexAIGeminiPro1Vision(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.0 Pro model."""
 
   model = 'gemini-1.0-pro-vision'
-  multimodal = True
+  supported_modalities = _IMAGE_TYPES + _VIDEO_TYPES
 
 
 class VertexAIPalm2(VertexAI):  # pylint: disable=invalid-name

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/llms/vertexai_test.py

@@ -79,7 +79,7 @@ class VertexAITest(unittest.TestCase):
     )
 
     # Non-multimodal model.
-    with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
+    with self.assertRaisesRegex(lf.ModalityError, 'Unsupported modality'):
       vertexai.VertexAIGeminiPro1()._content_from_message(message)
 
     model = vertexai.VertexAIGeminiPro1Vision()

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/modalities/__init__.py

@@ -18,7 +18,7 @@
 # pylint: disable=g-import-not-at-top
 
 from langfun.core.modalities.audio import Audio
-from langfun.core.modalities.mime import MimeType
+from langfun.core.modalities.mime import Mime
 from langfun.core.modalities.mime import Custom
 from langfun.core.modalities.ms_office import Docx
 from langfun.core.modalities.ms_office import Pptx

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/modalities/audio.py

@@ -17,7 +17,7 @@ import functools
 from langfun.core.modalities import mime
 
 
-class Audio(mime.MimeType):
+class Audio(mime.Mime):
   """Audio."""
 
   MIME_PREFIX = 'audio'

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/modalities/image.py

@@ -17,7 +17,7 @@ import functools
 from langfun.core.modalities import mime
 
 
-class Image(mime.MimeType):
+class Image(mime.Mime):
   """Image."""
 
   MIME_PREFIX = 'image'

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/modalities/image_test.py

@@ -15,7 +15,9 @@
 import unittest
 from unittest import mock
 
+import langfun.core as lf
 from langfun.core.modalities import image as image_lib
+from langfun.core.modalities import mime as mime_lib
 import pyglove as pg
 
 
@@ -36,23 +38,29 @@ def mock_request(*args, **kwargs):
   return pg.Dict(content=image_content)
 
 
-class ImageContentTest(unittest.TestCase):
+class ImageTest(unittest.TestCase):
 
-  def test_image_content(self):
+  def test_from_bytes(self):
     image = image_lib.Image.from_bytes(image_content)
     self.assertEqual(image.image_format, 'png')
     self.assertIn('data:image/png;base64,', image._repr_html_())
     self.assertEqual(image.to_bytes(), image_content)
+    with self.assertRaisesRegex(
+        lf.ModalityError, '.* cannot be converted to text'
+    ):
+      image.to_text()
 
-  def test_bad_image(self):
+  def test_from_bytes_invalid(self):
     image = image_lib.Image.from_bytes(b'bad')
     with self.assertRaisesRegex(ValueError, 'Expected MIME type'):
       _ = image.image_format
 
+  def test_from_bytes_base_cls(self):
+    self.assertIsInstance(
+        mime_lib.Mime.from_bytes(image_content), image_lib.Image
+    )
 
-class ImageFileTest(unittest.TestCase):
-
-  def test_image_file(self):
+  def test_from_uri(self):
     image = image_lib.Image.from_uri('http://mock/web/a.png')
     with mock.patch('requests.get') as mock_requests_get:
       mock_requests_get.side_effect = mock_request
@@ -60,6 +68,15 @@ class ImageFileTest(unittest.TestCase):
       self.assertEqual(image._repr_html_(), '<img src="http://mock/web/a.png">')
       self.assertEqual(image.to_bytes(), image_content)
 
+  def test_from_uri_base_cls(self):
+    with mock.patch('requests.get') as mock_requests_get:
+      mock_requests_get.side_effect = mock_request
+      image = mime_lib.Mime.from_uri('http://mock/web/a.png')
+      self.assertIsInstance(image, image_lib.Image)
+      self.assertEqual(image.image_format, 'png')
+      self.assertEqual(image._repr_html_(), '<img src="http://mock/web/a.png">')
+      self.assertEqual(image.to_bytes(), image_content)
+
 
 if __name__ == '__main__':
   unittest.main()

langfun-0.0.2.dev20240601/langfun/core/modalities/mime.py

@@ -0,0 +1,191 @@
+# 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.
+"""MIME type data."""
+
+import base64
+import functools
+from typing import Annotated, Iterable, Type, Union
+import langfun.core as lf
+import magic
+import pyglove as pg
+import requests
+
+
+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
+
+  content: Annotated[
+      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 = magic.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/xhtml+xml',
+            'application/xml',
+            'application/x-tex',
+            '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:
+      raise ValueError('Either uri or content must be provided.')
+
+  def to_bytes(self) -> bytes:
+    if self.content is not None:
+      return self.content
+
+    self.rebind(content=self.download(self.uri), skip_notification=True)
+    return self.content
+
+  @property
+  def content_uri(self) -> str:
+    base64_content = base64.b64encode(self.to_bytes()).decode()
+    return f'data:{self.mime_type};base64,{base64_content}'
+
+  @classmethod
+  def from_uri(cls, uri: str, **kwargs) -> 'Mime':
+    if cls is Mime:
+      content = cls.download(uri)
+      mime = magic.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) -> 'Mime':
+    if cls is Mime:
+      mime = magic.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 _repr_html_(self) -> str:
+    if self.uri and self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
+      uri = self.uri
+    else:
+      uri = self.content_uri
+    return self._html(uri)
+
+  def _html(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."""
+
+  mime: Annotated[
+      str, 'The MIME type of the data. E.g. text/plain, or image/png. '
+  ]
+
+  @property
+  def mime_type(self) -> str:
+    return self.mime

{langfun-0.0.2.dev20240531 → langfun-0.0.2.dev20240601}/langfun/core/modalities/mime_test.py

@@ -15,6 +15,7 @@
 import unittest
 from unittest import mock
 
+import langfun.core as lf
 from langfun.core.modalities import mime
 import pyglove as pg
 
@@ -31,10 +32,24 @@ def mock_readfile(*args, **kwargs):
 
 class CustomMimeTest(unittest.TestCase):
 
-  def test_content(self):
-    content = mime.Custom('text/plain', 'foo')
-    self.assertEqual(content.to_bytes(), 'foo')
+  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.'