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

langfun/core/data/conversion/gemini_test.py

@@ -0,0 +1,256 @@
+# Copyright 2025 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.
+import base64
+import unittest
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+from langfun.core.data.conversion import gemini  # pylint: disable=unused-import
+
+
+image_content = (
+    b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x18\x00\x00\x00\x18\x04'
+    b'\x03\x00\x00\x00\x12Y \xcb\x00\x00\x00\x18PLTE\x00\x00'
+    b'\x00fff_chaag_cg_ch^ci_ciC\xedb\x94\x00\x00\x00\x08tRNS'
+    b'\x00\n\x9f*\xd4\xff_\xf4\xe4\x8b\xf3a\x00\x00\x00>IDATx'
+    b'\x01c \x05\x08)"\xd8\xcc\xae!\x06pNz\x88k\x19\\Q\xa8"\x10'
+    b'\xc1\x14\x95\x01%\xc1\n\xa143Ta\xa8"D-\x84\x03QM\x98\xc3'
+    b'\x1a\x1a\x1a@5\x0e\x04\xa0q\x88\x05\x00\x07\xf8\x18\xf9'
+    b'\xdao\xd0|\x00\x00\x00\x00IEND\xaeB`\x82'
+)
+
+
+class GeminiConversionTest(unittest.TestCase):
+
+  def test_as_format_with_role(self):
+    self.assertEqual(
+        lf.UserMessage('hi').as_format('gemini'),
+        {
+            'role': 'user',
+            'parts': [{'text': 'hi'}],
+        },
+    )
+    self.assertEqual(
+        lf.AIMessage('hi').as_format('gemini'),
+        {
+            'role': 'model',
+            'parts': [{'text': 'hi'}],
+        },
+    )
+    self.assertEqual(
+        lf.SystemMessage('hi').as_format('gemini'),
+        {
+            'role': 'system',
+            'parts': [{'text': 'hi'}],
+        },
+    )
+
+  def test_as_format_with_image(self):
+    self.assertEqual(
+        lf.Template(
+            'What are the common words from {{image}}, {{pdf}} and {{video}}?',
+            image=lf_modalities.Image.from_bytes(image_content),
+            pdf=lf_modalities.Custom.from_uri(
+                'https://my.pdf', mime='application/pdf'
+            ),
+            video=lf_modalities.Custom.from_uri(
+                'https://www.youtube.com/watch?v=abcd', mime='text/html'
+            ),
+        ).render().as_gemini_format(),
+        {
+            'role': 'user',
+            'parts': [
+                {
+                    'text': 'What are the common words from'
+                },
+                {
+                    'inlineData': {
+                        'data': base64.b64encode(image_content).decode('utf-8'),
+                        'mimeType': 'image/png',
+                    }
+                },
+                {
+                    'text': ','
+                },
+                {
+                    'fileData': {
+                        'fileUri': 'https://my.pdf',
+                        'mimeType': 'application/pdf',
+                    }
+                },
+                {
+                    'text': 'and'
+                },
+                {
+                    'fileData': {
+                        'fileUri': 'https://www.youtube.com/watch?v=abcd',
+                        'mimeType': 'video/*',
+                    }
+                },
+                {
+                    'text': '?'
+                }
+            ],
+        },
+    )
+
+  def test_as_format_with_chunk_preprocessor(self):
+    self.assertEqual(
+        lf.Template(
+            'What is this {{image}}?',
+            image=lf_modalities.Image.from_bytes(image_content)
+        ).render().as_format(
+            'gemini',
+            chunk_preprocessor=lambda x: x if isinstance(x, str) else None
+        ),
+        {
+            'role': 'user',
+            'parts': [
+                {
+                    'text': 'What is this'
+                },
+                {
+                    'text': '?'
+                }
+            ],
+        },
+    )
+
+  def test_from_value_with_simple_text(self):
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'parts': [{'text': 'this is a text'}],
+            },
+            format='gemini',
+        ),
+        lf.AIMessage('this is a text'),
+    )
+
+  def test_from_value_with_role(self):
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'role': 'user',
+                'parts': [{'text': 'this is a text'}],
+            },
+            format='gemini',
+        ),
+        lf.UserMessage('this is a text'),
+    )
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'role': 'model',
+                'parts': [{'text': 'this is a text'}],
+            },
+            format='gemini',
+        ),
+        lf.AIMessage('this is a text'),
+    )
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'role': 'system',
+                'parts': [{'text': 'this is a text'}],
+            },
+            format='gemini',
+        ),
+        lf.SystemMessage('this is a text'),
+    )
+    with self.assertRaisesRegex(ValueError, 'Unsupported role: .*'):
+      lf.Message.from_value(
+          {
+              'role': 'function',
+              'parts': [{'text': 'this is a text'}],
+          },
+          format='gemini',
+      )
+
+  def test_from_value_with_thoughts(self):
+    message = lf.Message.from_value(
+        {
+            'role': 'user',
+            'parts': [
+                {
+                    'text': 'this is a red round object',
+                    'thought': True
+                },
+                {
+                    'text': 'this is a apple',
+                },
+            ],
+        },
+        format='gemini',
+    )
+    self.assertEqual(message.text, 'this is a apple')
+    self.assertEqual(message.thought, 'this is a red round object')
+
+  def test_from_value_with_modalities(self):
+    m = lf.Message.from_gemini_format(
+        {
+            'role': 'user',
+            'parts': [
+                {
+                    'text': 'What are the common words from'
+                },
+                {
+                    'inlineData': {
+                        'data': base64.b64encode(image_content).decode('utf-8'),
+                        'mimeType': 'image/png',
+                    }
+                },
+                {
+                    'text': ','
+                },
+                {
+                    'fileData': {
+                        'fileUri': 'https://my.pdf',
+                        'mimeType': 'application/pdf',
+                    }
+                },
+                {
+                    'text': 'and'
+                },
+                {
+                    'fileData': {
+                        'fileUri': 'https://www.youtube.com/watch?v=abcd',
+                        'mimeType': 'video/*',
+                    }
+                },
+                {
+                    'text': '?'
+                }
+            ],
+        },
+    )
+    self.assertEqual(
+        m.text,
+        (
+            'What are the common words from <<[[obj0]]>> , <<[[obj1]]>> '
+            'and <<[[obj2]]>> ?'
+        )
+    )
+    self.assertIsInstance(m.obj0, lf_modalities.Image)
+    self.assertEqual(m.obj0.mime_type, 'image/png')
+    self.assertEqual(m.obj0.to_bytes(), image_content)
+
+    self.assertIsInstance(m.obj1, lf_modalities.PDF)
+    self.assertEqual(m.obj1.uri, 'https://my.pdf')
+
+    self.assertIsInstance(m.obj2, lf_modalities.Video)
+    self.assertEqual(m.obj2.uri, 'https://www.youtube.com/watch?v=abcd')
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/data/conversion/openai.py

@@ -0,0 +1,131 @@
+# Copyright 2025 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.
+"""OpenAI API message conversion."""
+
+from typing import Annotated, Any, Callable
+
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+
+
+class OpenAIMessageConverter(lf.MessageConverter):
+  """Converter to OpenAI API."""
+
+  FORMAT_ID = 'openai'
+
+  chunk_preprocessor: Annotated[
+      Callable[[str | lf.Modality], Any] | None,
+      (
+          'Chunk preprocessor for Langfun chunk to OpenAI chunk conversion. '
+          'It will be applied before each Langfun chunk is converted. '
+          'If returns None, the chunk will be skipped.'
+      )
+  ] = None
+
+  def to_value(self, message: lf.Message) -> dict[str, Any]:
+    """Converts a Langfun message to OpenAI API."""
+    parts = []
+    for chunk in message.chunk():
+      if self.chunk_preprocessor is not None:
+        chunk = self.chunk_preprocessor(chunk)
+        if chunk is None:
+          continue
+
+      if isinstance(chunk, str):
+        item = dict(type='text', text=chunk)
+      elif isinstance(chunk, lf_modalities.Image):
+        item = dict(
+            type='image_url', image_url=dict(url=chunk.embeddable_uri)
+        )
+      # TODO(daiyip): Support audio_input.
+      else:
+        raise ValueError(f'Unsupported content type: {chunk!r}.')
+      parts.append(item)
+    return dict(
+        role=self.get_role(message),
+        content=parts,
+    )
+
+  def get_role(self, message: lf.Message) -> str:
+    """Returns the role of the message."""
+    if isinstance(message, lf.SystemMessage):
+      return 'system'
+    elif isinstance(message, lf.UserMessage):
+      return 'user'
+    elif isinstance(message, lf.AIMessage):
+      return 'assistant'
+    else:
+      raise ValueError(f'Unsupported message type: {message!r}.')
+
+  def get_message_cls(self, role: str) -> type[lf.Message]:
+    """Returns the message class of the message."""
+    match role:
+      case 'system':
+        return lf.SystemMessage
+      case 'user':
+        return lf.UserMessage
+      case 'assistant':
+        return lf.AIMessage
+      case _:
+        raise ValueError(f'Unsupported role: {role!r}.')
+
+  def from_value(self, value: dict[str, Any]) -> lf.Message:
+    """Returns a Langfun message from OpenAI message."""
+    message_cls = self.get_message_cls(
+        self._safe_read(value, 'role', default='assistant')
+    )
+    content = self._safe_read(value, 'content')
+    if isinstance(content, str):
+      return message_cls(content)
+
+    assert isinstance(content, list)
+    chunks = []
+    for item in content:
+      t = self._safe_read(item, 'type')
+      if t == 'text':
+        chunk = self._safe_read(item, 'text')
+      elif t == 'image_url':
+        chunk = lf_modalities.Image.from_uri(
+            self._safe_read(self._safe_read(item, 'image_url'), 'url')
+        )
+      else:
+        raise ValueError(f'Unsupported content type: {item!r}.')
+      chunks.append(chunk)
+    return message_cls.from_chunks(chunks)
+
+
+def _as_openai_format(
+    self,
+    chunk_preprocessor: Callable[[str | lf.Modality], Any] | None = None,
+    **kwargs
+) -> dict[str, Any]:
+  """Returns the OpenAI format of the chunk."""
+  return OpenAIMessageConverter(
+      chunk_preprocessor=chunk_preprocessor, **kwargs
+  ).to_value(self)
+
+
+@classmethod
+def _from_openai_format(
+    cls,
+    openai_message: dict[str, Any],
+    **kwargs
+) -> dict[str, Any]:
+  """Returns the OpenAI format of the chunk."""
+  del cls
+  return OpenAIMessageConverter(**kwargs).from_value(openai_message)
+
+# Set shortcut methods in lf.Message.
+lf.Message.as_openai_format = _as_openai_format
+lf.Message.from_openai_format = _from_openai_format

langfun/core/data/conversion/openai_test.py

@@ -0,0 +1,176 @@
+# Copyright 2025 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.
+import base64
+import unittest
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+from langfun.core.data.conversion import openai  # pylint: disable=unused-import
+
+
+image_content = (
+    b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x18\x00\x00\x00\x18\x04'
+    b'\x03\x00\x00\x00\x12Y \xcb\x00\x00\x00\x18PLTE\x00\x00'
+    b'\x00fff_chaag_cg_ch^ci_ciC\xedb\x94\x00\x00\x00\x08tRNS'
+    b'\x00\n\x9f*\xd4\xff_\xf4\xe4\x8b\xf3a\x00\x00\x00>IDATx'
+    b'\x01c \x05\x08)"\xd8\xcc\xae!\x06pNz\x88k\x19\\Q\xa8"\x10'
+    b'\xc1\x14\x95\x01%\xc1\n\xa143Ta\xa8"D-\x84\x03QM\x98\xc3'
+    b'\x1a\x1a\x1a@5\x0e\x04\xa0q\x88\x05\x00\x07\xf8\x18\xf9'
+    b'\xdao\xd0|\x00\x00\x00\x00IEND\xaeB`\x82'
+)
+
+
+class OpenAIConversionTest(unittest.TestCase):
+
+  def test_as_format_with_role(self):
+    self.assertEqual(
+        lf.UserMessage('hi').as_format('openai'),
+        {
+            'role': 'user',
+            'content': [{'type': 'text', 'text': 'hi'}],
+        },
+    )
+    self.assertEqual(
+        lf.AIMessage('hi').as_format('openai'),
+        {
+            'role': 'assistant',
+            'content': [{'type': 'text', 'text': 'hi'}],
+        },
+    )
+    self.assertEqual(
+        lf.SystemMessage('hi').as_format('openai'),
+        {
+            'role': 'system',
+            'content': [{'type': 'text', 'text': 'hi'}],
+        },
+    )
+
+  def test_as_format_with_image(self):
+    self.assertEqual(
+        lf.Template(
+            'What is this {{image}}?',
+            image=lf_modalities.Image.from_bytes(image_content)
+        ).render().as_format('openai'),
+        {
+            'role': 'user',
+            'content': [
+                {
+                    'type': 'text',
+                    'text': 'What is this'
+                },
+                {
+                    'type': 'image_url',
+                    'image_url': {
+                        'url': (
+                            'data:image/png;base64,'
+                            + base64.b64encode(image_content).decode('utf-8')
+                        )
+                    }
+                },
+                {
+                    'type': 'text',
+                    'text': '?'
+                }
+            ],
+        },
+    )
+
+  def test_as_format_with_chunk_preprocessor(self):
+    self.assertEqual(
+        lf.Template(
+            'What is this {{image}}?',
+            image=lf_modalities.Image.from_bytes(image_content)
+        ).render().as_openai_format(
+            chunk_preprocessor=lambda x: x if isinstance(x, str) else None
+        ),
+        {
+            'role': 'user',
+            'content': [
+                {
+                    'type': 'text',
+                    'text': 'What is this'
+                },
+                {
+                    'type': 'text',
+                    'text': '?'
+                }
+            ],
+        },
+    )
+
+  def test_from_value_with_simple_text(self):
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'content': 'this is a text',
+            },
+            format='openai',
+        ),
+        lf.AIMessage('this is a text'),
+    )
+
+  def test_from_value_with_role(self):
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'role': 'user',
+                'content': [{'type': 'text', 'text': 'hi'}],
+            },
+            format='openai',
+        ),
+        lf.UserMessage('hi'),
+    )
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'role': 'assistant',
+                'content': [{'type': 'text', 'text': 'hi'}],
+            },
+            format='openai',
+        ),
+        lf.AIMessage('hi'),
+    )
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'role': 'system',
+                'content': [{'type': 'text', 'text': 'hi'}],
+            },
+            format='openai',
+        ),
+        lf.SystemMessage('hi'),
+    )
+    with self.assertRaisesRegex(ValueError, 'Unsupported role: .*'):
+      lf.Message.from_value(
+          {
+              'role': 'function',
+              'content': [{'type': 'text', 'text': 'hi'}],
+          },
+          format='openai',
+      )
+
+  def test_from_value_with_image(self):
+    m = lf.Message.from_openai_format(
+        lf.Template(
+            'What is this {{image}}?',
+            image=lf_modalities.Image.from_bytes(image_content)
+        ).render().as_format('openai'),
+    )
+    self.assertEqual(m.text, 'What is this <<[[obj0]]>> ?')
+    self.assertIsInstance(m.obj0, lf_modalities.Image)
+    self.assertEqual(m.obj0.mime_type, 'image/png')
+    self.assertEqual(m.obj0.to_bytes(), image_content)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/llms/anthropic.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """Language models from Anthropic."""
 
-import base64
 import datetime
 import functools
 import os
@@ -21,6 +20,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 anthropic as anthropic_conversion  # pylint: disable=unused-import
 from langfun.core.llms import rest
 import pyglove as pg
 
@@ -502,10 +502,17 @@ class Anthropic(rest.REST):
     """Returns the JSON input for a message."""
     request = dict()
     request.update(self._request_args(sampling_options))
+
+    def modality_check(chunk: Any) -> Any:
+      if isinstance(chunk, lf_modalities.Mime):
+        if not self.supports_input(chunk.mime_type):
+          raise ValueError(f'Unsupported modality: {chunk!r}.')
+      return chunk
+
     request.update(
         dict(
             messages=[
-                dict(role='user', content=self._content_from_message(prompt))
+                prompt.as_format('anthropic', chunk_preprocessor=modality_check)
             ]
         )
     )
@@ -548,43 +555,8 @@ class Anthropic(rest.REST):
       args.pop('top_p', None)
     return args
 
-  def _content_from_message(self, prompt: lf.Message) -> list[dict[str, Any]]:
-    """Converts an message to Anthropic's content protocol (list of dicts)."""
-    # Refer: https://docs.anthropic.com/claude/reference/messages-examples
-    content = []
-    for chunk in prompt.chunk():
-      if isinstance(chunk, str):
-        content.append(dict(type='text', text=chunk))
-      elif isinstance(chunk, lf_modalities.Mime):
-        if not self.supports_input(chunk.mime_type):
-          raise ValueError(f'Unsupported modality: {chunk!r}.')
-        if isinstance(chunk, lf_modalities.Image):
-          item = dict(
-              type='image',
-              source=dict(
-                  type='base64',
-                  media_type=chunk.mime_type,
-                  data=base64.b64encode(chunk.to_bytes()).decode(),
-              ),
-          )
-        elif isinstance(chunk, lf_modalities.PDF):
-          item = dict(
-              type='document',
-              source=dict(
-                  type='base64',
-                  media_type=chunk.mime_type,
-                  data=base64.b64encode(chunk.to_bytes()).decode(),
-              ),
-          )
-        else:
-          raise NotImplementedError(
-              f'Modality conversion not implemented: {chunk!r}'
-          )
-        content.append(item)
-    return content
-
   def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
-    message = self._message_from_content(json['content'])
+    message = lf.Message.from_value(json, format='anthropic')
     input_tokens = json['usage']['input_tokens']
     output_tokens = json['usage']['output_tokens']
     return lf.LMSamplingResult(
@@ -596,20 +568,6 @@ class Anthropic(rest.REST):
         ),
     )
 
-  def _message_from_content(self, content: list[dict[str, Any]]) -> lf.Message:
-    """Converts Anthropic's content protocol to message."""
-    # Refer: https://docs.anthropic.com/claude/reference/messages-examples
-    # Thinking: https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#implementing-extended-thinking # pylint: disable=line-too-long
-    response = lf.AIMessage.from_chunks(
-        [x['text'] for x in content if x['type'] == 'text']
-    )
-    thinking = lf.AIMessage.from_chunks(
-        [x['thinking'] for x in content if x['type'] == 'thinking']
-    )
-    # thinking is added into the metadata.thinking field.
-    response.set('thinking', thinking)
-    return response
-
 
 class Claude37(Anthropic):
   """Base class for Claude 3.7 models."""