langfun 0.1.2.dev202503240804__py3-none-any.whl → 0.1.2.dev202503260804__py3-none-any.whl

langfun/core/llms/gemini.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 """Gemini REST API (Shared by Google GenAI and Vertex AI)."""
 
-import base64
 import datetime
 import functools
 from typing import Annotated, Any
 
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
+from langfun.core.data.conversion import gemini as gemini_conversion  # pylint: disable=unused-import
 from langfun.core.llms import rest
 import pyglove as pg
 
@@ -559,7 +559,19 @@ class Gemini(rest.REST):
     request = dict(
         generationConfig=self._generation_config(prompt, sampling_options)
     )
-    request['contents'] = [self._content_from_message(prompt)]
+    def modality_conversion(chunk: str | lf.Modality) -> Any:
+      if isinstance(chunk, lf_modalities.Mime):
+        try:
+          return chunk.make_compatible(
+              self.model_info.input_modalities + ['text/plain']
+          )
+        except lf.ModalityError as e:
+          raise lf.ModalityError(f'Unsupported modality: {chunk!r}') from e
+      return chunk
+
+    request['contents'] = [
+        prompt.as_format('gemini', chunk_preprocessor=modality_conversion)
+    ]
     return request
 
   def _generation_config(
@@ -593,49 +605,9 @@ class Gemini(rest.REST):
       )
     return config
 
-  def _content_from_message(self, prompt: lf.Message) -> dict[str, Any]:
-    """Gets generation content from langfun message."""
-    parts = []
-    for lf_chunk in prompt.chunk():
-      if isinstance(lf_chunk, str):
-        parts.append({'text': lf_chunk})
-      elif isinstance(lf_chunk, lf_modalities.Mime):
-        try:
-          modalities = lf_chunk.make_compatible(
-              self.model_info.input_modalities + ['text/plain']
-          )
-          if isinstance(modalities, lf_modalities.Mime):
-            modalities = [modalities]
-          for modality in modalities:
-            if modality.is_text:
-              # Add YouTube video into the context window.
-              # https://ai.google.dev/gemini-api/docs/vision?lang=python#youtube
-              if modality.mime_type == 'text/html' and modality.uri.startswith(
-                  'https://www.youtube.com/watch?v='
-              ):
-                parts.append({
-                    'fileData': {'mimeType': 'video/*', 'fileUri': modality.uri}
-                })
-              else:
-                parts.append({'text': modality.to_text()})
-            else:
-              parts.append({
-                  'inlineData': {
-                      'data': base64.b64encode(modality.to_bytes()).decode(),
-                      'mimeType': modality.mime_type,
-                  }
-              })
-        except lf.ModalityError as e:
-          raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}') from e
-      else:
-        raise NotImplementedError(
-            f'Input conversion not implemented: {lf_chunk!r}'
-        )
-    return dict(role='user', parts=parts)
-
   def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
     messages = [
-        self._message_from_content_parts(candidate['content'].get('parts', []))
+        lf.Message.from_value(candidate['content'], format='gemini')
         for candidate in json['candidates']
     ]
     usage = json['usageMetadata']
@@ -652,22 +624,3 @@ class Gemini(rest.REST):
             total_tokens=input_tokens + output_tokens,
         ),
     )
-
-  def _message_from_content_parts(
-      self, parts: list[dict[str, Any]]
-  ) -> lf.Message:
-    """Converts Vertex AI's content parts protocol to message."""
-    chunks = []
-    thought_chunks = []
-    for part in parts:
-      if text_part := part.get('text'):
-        if part.get('thought'):
-          thought_chunks.append(text_part)
-        else:
-          chunks.append(text_part)
-      else:
-        raise ValueError(f'Unsupported part: {part}')
-    message = lf.AIMessage.from_chunks(chunks)
-    if thought_chunks:
-      message.set('thought', lf.AIMessage.from_chunks(thought_chunks))
-    return message

langfun/core/llms/gemini_test.py

@@ -13,13 +13,11 @@
 # limitations under the License.
 """Tests for Gemini API."""
 
-import base64
 from typing import Any
 import unittest
 from unittest import mock
 
 import langfun.core as lf
-from langfun.core import modalities as lf_modalities
 from langfun.core.llms import gemini
 import pyglove as pg
 import requests
@@ -105,36 +103,6 @@ class GeminiTest(unittest.TestCase):
         0.51
     )
 
-  def test_content_from_message_text_only(self):
-    text = 'This is a beautiful day'
-    model = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
-    chunks = model._content_from_message(lf.UserMessage(text))
-    self.assertEqual(chunks, {'role': 'user', 'parts': [{'text': text}]})
-
-  def test_content_from_message_mm(self):
-    image = lf_modalities.Image.from_bytes(example_image)
-    message = lf.UserMessage(
-        'This is an <<[[image]]>>, what is it?', image=image
-    )
-    model = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
-    content = model._content_from_message(message)
-    self.assertEqual(
-        content,
-        {
-            'role': 'user',
-            'parts': [
-                {'text': 'This is an'},
-                {
-                    'inlineData': {
-                        'data': base64.b64encode(example_image).decode(),
-                        'mimeType': 'image/png',
-                    }
-                },
-                {'text': ', what is it?'},
-            ],
-        },
-    )
-
   def test_generation_config(self):
     model = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
     json_schema = {

langfun/core/llms/openai_compatible.py

@@ -17,6 +17,7 @@ from typing import Annotated, Any
 
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
+from langfun.core.data.conversion import openai as openai_conversion  # pylint: disable=unused-import
 from langfun.core.llms import rest
 import pyglove as pg
 
@@ -61,20 +62,6 @@ class OpenAICompatible(rest.REST):
       args['seed'] = options.random_seed
     return args
 
-  def _content_from_message(self, message: lf.Message) -> list[dict[str, Any]]:
-    """Returns a OpenAI content object from a Langfun message."""
-    content = []
-    for chunk in message.chunk():
-      if isinstance(chunk, str):
-        item = dict(type='text', text=chunk)
-      elif (isinstance(chunk, lf_modalities.Image)
-            and self.supports_input(chunk.mime_type)):
-        item = dict(type='image_url', image_url=dict(url=chunk.embeddable_uri))
-      else:
-        raise ValueError(f'Unsupported modality: {chunk!r}.')
-      content.append(item)
-    return content
-
   def request(
       self,
       prompt: lf.Message,
@@ -114,16 +101,25 @@ class OpenAICompatible(rest.REST):
 
     # Prepare messages.
     messages = []
+
+    def modality_check(chunk: str | lf.Modality) -> Any:
+      if (isinstance(chunk, lf_modalities.Mime)
+          and not self.supports_input(chunk.mime_type)):
+        raise ValueError(
+            f'Unsupported modality: {chunk!r}.'
+        )
+      return chunk
+
     # Users could use `metadata_system_message` to pass system message.
     system_message = prompt.metadata.get('system_message')
     if system_message:
-      system_message = lf.SystemMessage.from_value(system_message)
       messages.append(
-          dict(role='system',
-               content=self._content_from_message(system_message))
+          lf.SystemMessage.from_value(system_message).as_format(
+              'openai', chunk_preprocessor=modality_check
+          )
       )
     messages.append(
-        dict(role='user', content=self._content_from_message(prompt))
+        prompt.as_format('openai', chunk_preprocessor=modality_check)
     )
     request = dict()
     request.update(request_args)
@@ -145,7 +141,7 @@ class OpenAICompatible(rest.REST):
           for t in choice_logprobs['content']
       ]
     return lf.LMSample(
-        choice['message']['content'],
+        lf.Message.from_value(choice['message'], format='openai'),
         score=0.0,
         logprobs=logprobs,
     )

langfun/core/message.py

@@ -13,10 +13,14 @@
 # limitations under the License.
 """Messages that are exchanged between users and agents."""
 
+import abc
+import collections
+from collections.abc import Mapping
 import contextlib
 import functools
+import inspect
 import io
-from typing import Annotated, Any, Optional, Union
+from typing import Annotated, Any, ClassVar, Optional, Type, Union
 
 from langfun.core import modality
 from langfun.core import natural_language
@@ -146,13 +150,91 @@ class Message(
     self._source = source
 
   @classmethod
-  def from_value(cls, value: Union[str, 'Message']) -> 'Message':
-    """Creates a message from a value or return value itself if a Message."""
+  def from_value(
+      cls,
+      value: Union[str, 'Message', Any],
+      *,
+      format: str | None = None,    # pylint: disable=redefined-builtin
+      **kwargs
+  ) -> 'Message':
+    """Creates a message from a str, message or object of registered format.
+
+    Example:
+      # Create a user message from a str.
+      lf.UserMessage.from_value('hi')
+
+      # Create a user message from a multi-modal object.
+      lf.UserMessage.from_value(lf.Image.from_uri('...'))
+
+      # Create a user message from OpenAI API format.
+      lf.Message.from_value(
+          {
+              'role': 'user',
+              'content': [{'type': 'text', 'text': 'hi'}],
+          },
+          format='openai.api',
+      )
+
+    Args:
+      value: The value to create a message from.
+      format: The format ID to convert to. If None, the conversion will be
+        performed according to the type of `value`. Otherwise, the converter
+        registered for the format will be used.
+      **kwargs: The keyword arguments passed to the __init__ of the converter.
+
+    Returns:
+      A message created from the value.
+    """
     if isinstance(value, modality.Modality):
       return cls('<<[[object]]>>', object=value)
     if isinstance(value, Message):
       return value
-    return cls(value)
+    if isinstance(value, str):
+      return cls(value)
+    value_type = type(value)
+    if format is None:
+      converter = MessageConverter.get_by_type(value_type, **kwargs)
+    else:
+      converter = MessageConverter.get_by_format(format, **kwargs)
+      if (converter.OUTPUT_TYPE is not None
+          and not isinstance(value, converter.OUTPUT_TYPE)):
+        raise ValueError(f'{format!r} is not applicable to {value!r}.')
+    return converter.from_value(value)
+
+  #
+  # Conversion to other formats or types.
+  #
+
+  def as_format(self, format_or_type: str | type[Any], **kwargs) -> Any:
+    """Converts the message to a registered format or type.
+
+    Example:
+
+      m = lf.Template('What is this {{image}}?').render()
+      m.as_format('openai')  # Convert to OpenAI message format.
+      m.as_format('gemini')  # Convert to Gemini message format.
+      m.as_format('anthropic')  # Convert to Anthropic message format.
+
+    Args:
+      format_or_type: The format ID or type to convert to.
+      **kwargs: The conversion arguments.
+
+    Returns:
+      The converted object according to the format or type.
+    """
+    return MessageConverter.get(format_or_type, **kwargs).to_value(self)
+
+  @classmethod
+  @property
+  def convertible_formats(cls) -> list[str]:
+    """Returns supported format for message conversion."""
+    return MessageConverter.convertible_formats()
+
+  @classmethod
+  @property
+  def convertible_types(cls) -> list[str]:
+    """Returns supported types for message conversion."""
+    return MessageConverter.convertible_types()
 
   #
   # Unified interface for accessing text, result and metadata.
@@ -377,29 +459,6 @@ class Message(
         ref_index += 1
     return cls(fused_text.getvalue().strip(), metadata=metadata)
 
-  #
-  # API for testing the message types.
-  #
-
-  @property
-  def from_user(self) -> bool:
-    """Returns True if it's user message."""
-    return isinstance(self, UserMessage)
-
-  @property
-  def from_agent(self) -> bool:
-    """Returns True if it's agent message."""
-    return isinstance(self, AIMessage)
-
-  @property
-  def from_system(self) -> bool:
-    """Returns True if it's agent message."""
-    return isinstance(self, SystemMessage)
-
-  @property
-  def from_memory(self) -> bool:
-    return isinstance(self, MemoryRecord)
-
   #
   # Tagging
   #
@@ -799,6 +858,152 @@ class Message(
     ]
 
 
+class _MessageConverterRegistry:
+  """Message converter registry."""
+
+  def __init__(self):
+    self._name_to_converter: dict[str, Type[MessageConverter]] = {}
+    self._type_to_converters: dict[Type[Any], list[Type[MessageConverter]]] = (
+        collections.defaultdict(list)
+    )
+
+  def register(self, converter: Type['MessageConverter']) -> None:
+    """Registers a message converter."""
+    self._name_to_converter[converter.FORMAT_ID] = converter
+    if converter.OUTPUT_TYPE is not None:
+      self._type_to_converters[converter.OUTPUT_TYPE].append(converter)
+
+  def get_by_type(self, t: Type[Any], **kwargs) -> 'MessageConverter':
+    """Returns a message converter for the given type."""
+    t = self._type_to_converters[t]
+    if not t:
+      raise TypeError(
+          f'Cannot convert Message to {t!r}.'
+      )
+    if len(t) > 1:
+      raise TypeError(
+          f'More than one converters found for output type {t!r}. '
+          f'Please specify one for this conversion: {[x.FORMAT_ID for x in t]}.'
+      )
+    return t[0](**kwargs)
+
+  def get_by_format(self, format: str, **kwargs) -> 'MessageConverter':   # pylint: disable=redefined-builtin
+    """Returns a message converter for the given format."""
+    if format not in self._name_to_converter:
+      raise ValueError(f'Unsupported format: {format!r}.')
+    return self._name_to_converter[format](**kwargs)
+
+  def get(
+      self,
+      format_or_type: str | Type[Any], **kwargs
+  ) -> 'MessageConverter':
+    """Returns a message converter for the given format or type."""
+    if isinstance(format_or_type, str):
+      return self.get_by_format(format_or_type, **kwargs)
+    assert isinstance(format_or_type, type), format_or_type
+    return self.get_by_type(format_or_type, **kwargs)
+
+  def convertible_formats(self) -> list[str]:
+    """Returns a list of converter names."""
+    return sorted(list(self._name_to_converter.keys()))
+
+  def convertible_types(self) -> list[Type[Any]]:
+    """Returns a list of converter types."""
+    return list(self._type_to_converters.keys())
+
+
+class MessageConverter(pg.Object):
+  """Interface for converting a Langfun message to other formats."""
+
+  # A global unique identifier for the converter.
+  FORMAT_ID: ClassVar[str]
+
+  # The output type of the converter.
+  # If None, the converter will not be registered to handle
+  # `lf.Message.to_value(output_type)`.
+  OUTPUT_TYPE: ClassVar[Type[Any] | None] = None
+
+  _REGISTRY = _MessageConverterRegistry()
+
+  def __init_subclass__(cls, *args, **kwargs):
+    super().__init_subclass__(*args, **kwargs)
+    if not inspect.isabstract(cls):
+      cls._REGISTRY.register(cls)
+
+  @abc.abstractmethod
+  def to_value(self, message: Message) -> Any:
+    """Converts a Langfun message to other formats."""
+
+  @abc.abstractmethod
+  def from_value(self, value: Message) -> Message:
+    """Returns a MessageConverter from a Langfun message."""
+
+  @classmethod
+  def _safe_read(
+      cls,
+      data: Mapping[str, Any],
+      key: str,
+      default: Any = pg.MISSING_VALUE
+  ) -> Any:
+    """Safe reads a key from a mapping."""
+    if not isinstance(data, Mapping):
+      raise ValueError(f'Invalid data type: {data!r}.')
+    if key not in data:
+      if pg.MISSING_VALUE == default:
+        raise ValueError(f'Missing key {key!r} in {data!r}')
+      return default
+    return data[key]
+
+  @classmethod
+  def get_role(cls, message: Message) -> str:
+    """Returns the role of the message."""
+    if isinstance(message, SystemMessage):
+      return 'system'
+    elif isinstance(message, UserMessage):
+      return 'user'
+    elif isinstance(message, AIMessage):
+      return 'assistant'
+    else:
+      raise ValueError(f'Unsupported message type: {message!r}.')
+
+  @classmethod
+  def get_message_cls(cls, role: str) -> type[Message]:
+    """Returns the message class of the message."""
+    match role:
+      case 'system':
+        return SystemMessage
+      case 'user':
+        return UserMessage
+      case 'assistant':
+        return AIMessage
+      case _:
+        raise ValueError(f'Unsupported role: {role!r}.')
+
+  @classmethod
+  def get(cls, format_or_type: str | Type[Any], **kwargs) -> 'MessageConverter':
+    """Returns a message converter."""
+    return cls._REGISTRY.get(format_or_type, **kwargs)
+
+  @classmethod
+  def get_by_format(cls, format: str, **kwargs) -> 'MessageConverter':  # pylint: disable=redefined-builtin
+    """Returns a message converter for the given format."""
+    return cls._REGISTRY.get_by_format(format, **kwargs)
+
+  @classmethod
+  def get_by_type(cls, t: Type[Any], **kwargs) -> 'MessageConverter':
+    """Returns a message converter for the given type."""
+    return cls._REGISTRY.get_by_type(t, **kwargs)
+
+  @classmethod
+  def convertible_formats(cls) -> list[str]:
+    """Returns a list of converter names."""
+    return cls._REGISTRY.convertible_formats()
+
+  @classmethod
+  def convertible_types(cls) -> list[Type[Any]]:
+    """Returns a list of converter types."""
+    return cls._REGISTRY.convertible_types()
+
 #
 # Messages of different roles.
 #

langfun/core/message_test.py

@@ -187,10 +187,6 @@ class MessageTest(unittest.TestCase):
     m = message.UserMessage('hi')
     self.assertEqual(m.text, 'hi')
     self.assertEqual(m.sender, 'User')
-    self.assertTrue(m.from_user)
-    self.assertFalse(m.from_agent)
-    self.assertFalse(m.from_system)
-    self.assertFalse(m.from_memory)
     self.assertEqual(str(m), m.text)
 
     m = message.UserMessage('hi', sender='Tom')
@@ -201,10 +197,6 @@ class MessageTest(unittest.TestCase):
     m = message.AIMessage('hi')
     self.assertEqual(m.text, 'hi')
     self.assertEqual(m.sender, 'AI')
-    self.assertFalse(m.from_user)
-    self.assertTrue(m.from_agent)
-    self.assertFalse(m.from_system)
-    self.assertFalse(m.from_memory)
     self.assertEqual(str(m), m.text)
 
     m = message.AIMessage('hi', sender='Model')
@@ -215,10 +207,6 @@ class MessageTest(unittest.TestCase):
     m = message.SystemMessage('hi')
     self.assertEqual(m.text, 'hi')
     self.assertEqual(m.sender, 'System')
-    self.assertFalse(m.from_user)
-    self.assertFalse(m.from_agent)
-    self.assertTrue(m.from_system)
-    self.assertFalse(m.from_memory)
     self.assertEqual(str(m), m.text)
 
     m = message.SystemMessage('hi', sender='Environment1')
@@ -229,10 +217,6 @@ class MessageTest(unittest.TestCase):
     m = message.MemoryRecord('hi')
     self.assertEqual(m.text, 'hi')
     self.assertEqual(m.sender, 'Memory')
-    self.assertFalse(m.from_user)
-    self.assertFalse(m.from_agent)
-    self.assertFalse(m.from_system)
-    self.assertTrue(m.from_memory)
     self.assertEqual(str(m), m.text)
 
     m = message.MemoryRecord('hi', sender="Someone's Memory")
@@ -493,5 +477,135 @@ class MessageTest(unittest.TestCase):
     )
 
 
+class MessageConverterTest(unittest.TestCase):
+
+  def test_basics(self):
+
+    class IntConverter(message.MessageConverter):
+      OUTPUT_TYPE = int
+
+    class TestConverter(IntConverter):  # pylint: disable=unused-variable
+      FORMAT_ID = 'test_format1'
+
+      def to_value(self, m: message.Message) -> int:
+        return int(m.text)
+
+      def from_value(self, value: int) -> message.Message:
+        return message.UserMessage(str(value))
+
+    class TestConverter2(IntConverter):  # pylint: disable=unused-variable
+      FORMAT_ID = 'test_format2'
+
+      def to_value(self, m: message.Message) -> int:
+        return int(m.text) + 1
+
+      def from_value(self, value: int) -> message.Message:
+        return message.UserMessage(str(value - 1))
+
+    class TestConverter3(message.MessageConverter):  # pylint: disable=unused-variable
+      FORMAT_ID = 'test_format3'
+      OUTPUT_TYPE = tuple
+
+      def to_value(self, m: message.Message) -> tuple[int, ...]:
+        return tuple(int(x) for x in m.text.split(','))
+
+      def from_value(self, value: tuple[int, ...]) -> message.Message:
+        return message.UserMessage(','.join(str(x) for x in value))
+
+    self.assertIn('test_format1', message.Message.convertible_formats)
+    self.assertIn('test_format2', message.Message.convertible_formats)
+    self.assertIn('test_format3', message.Message.convertible_formats)
+
+    self.assertIn(int, message.Message.convertible_types)
+    self.assertIn(tuple, message.Message.convertible_types)
+    self.assertEqual(
+        message.Message.from_value(1, format='test_format1'),
+        message.UserMessage('1')
+    )
+    self.assertEqual(
+        message.UserMessage('1').as_format('test_format1'),
+        1
+    )
+    self.assertEqual(
+        message.Message.from_value(1, format='test_format2'),
+        message.UserMessage('0')
+    )
+    self.assertEqual(
+        message.UserMessage('1').as_format('test_format2'),
+        2
+    )
+    with self.assertRaisesRegex(ValueError, 'Unsupported format: .*'):
+      message.UserMessage('1').as_format('test4')
+
+    with self.assertRaisesRegex(TypeError, 'Cannot convert Message to .*'):
+      message.UserMessage('1').as_format(float)
+
+    with self.assertRaisesRegex(
+        TypeError, 'More than one converters found for output type .*'
+    ):
+      message.UserMessage('1').as_format(int)
+    self.assertEqual(
+        message.UserMessage('1,2,3').as_format('test_format3'),
+        (1, 2, 3)
+    )
+    self.assertEqual(
+        message.UserMessage('1,2,3').as_format(tuple),
+        (1, 2, 3)
+    )
+    self.assertEqual(
+        message.Message.from_value((1, 2, 3)),
+        message.UserMessage('1,2,3')
+    )
+
+  def test_get_role(self):
+    self.assertEqual(
+        message.MessageConverter.get_role(message.SystemMessage('hi')),
+        'system',
+    )
+    self.assertEqual(
+        message.MessageConverter.get_role(message.UserMessage('hi')),
+        'user',
+    )
+    self.assertEqual(
+        message.MessageConverter.get_role(message.AIMessage('hi')),
+        'assistant',
+    )
+    with self.assertRaisesRegex(ValueError, 'Unsupported message type: .*'):
+      message.MessageConverter.get_role(message.MemoryRecord('hi'))
+
+  def test_get_message_cls(self):
+    self.assertEqual(
+        message.MessageConverter.get_message_cls('system'),
+        message.SystemMessage,
+    )
+    self.assertEqual(
+        message.MessageConverter.get_message_cls('user'),
+        message.UserMessage,
+    )
+    self.assertEqual(
+        message.MessageConverter.get_message_cls('assistant'),
+        message.AIMessage,
+    )
+    with self.assertRaisesRegex(ValueError, 'Unsupported role: .*'):
+      message.MessageConverter.get_message_cls('foo')
+
+  def test_safe_read(self):
+    self.assertEqual(
+        message.MessageConverter._safe_read({'a': 1}, 'a'),
+        1,
+    )
+    self.assertEqual(
+        message.MessageConverter._safe_read({'a': 1}, 'a', default=2),
+        1,
+    )
+    self.assertEqual(
+        message.MessageConverter._safe_read({'a': 1}, 'b', default=2),
+        2,
+    )
+    with self.assertRaisesRegex(ValueError, 'Invalid data type: .*'):
+      message.MessageConverter._safe_read(1, 'a')
+    with self.assertRaisesRegex(ValueError, 'Missing key .*'):
+      message.MessageConverter._safe_read({'a': 1}, 'b')
+
 if __name__ == '__main__':
   unittest.main()