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

langfun/__init__.py

@@ -48,6 +48,7 @@ query_output = structured.query_output
 source_form = structured.source_form
 function_gen = structured.function_gen
 
+from langfun.core import data
 from langfun.core import eval  # pylint: disable=redefined-builtin
 from langfun.core import templates
 from langfun.core import coding

langfun/core/__init__.py

@@ -87,6 +87,8 @@ from langfun.core.message import AIMessage
 from langfun.core.message import SystemMessage
 from langfun.core.message import MemoryRecord
 
+from langfun.core.message import MessageConverter
+
 # Interface for modality.
 from langfun.core.modality import Modality
 from langfun.core.modality import ModalityRef

langfun/core/data/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+"""langfun data sub-module."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=g-bad-import-order
+
+from langfun.core.data import conversion

langfun/core/data/conversion/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+"""langfun data conversions."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=g-bad-import-order
+
+from langfun.core.data.conversion import anthropic
+from langfun.core.data.conversion import gemini
+from langfun.core.data.conversion import openai

langfun/core/data/conversion/anthropic.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.
+"""Anthropic API message conversion."""
+
+import base64
+from typing import Annotated, Any, Callable
+
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+
+
+class AnthropicMessageConverter(lf.MessageConverter):
+  """Converter to Anthropic public API."""
+
+  FORMAT_ID = 'anthropic'
+
+  chunk_preprocessor: Annotated[
+      Callable[[str | lf.Modality], Any] | None,
+      (
+          'Chunk preprocessor for Langfun chunk to Anthropic 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 Gemini API."""
+    content = []
+    for chunk in message.chunk():
+      if self.chunk_preprocessor:
+        chunk = self.chunk_preprocessor(chunk)
+        if chunk is None:
+          continue
+
+      if isinstance(chunk, str):
+        content.append({'type': 'text', 'text': chunk})
+      elif isinstance(chunk, lf_modalities.Mime):
+        if isinstance(chunk, lf_modalities.Image):
+          content.append(
+              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):
+          content.append(
+              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}'
+          )
+    return dict(role=self.get_role(message), content=content)
+
+  def from_value(self, value: dict[str, Any]) -> lf.Message:
+    """Returns a Langfun message from Anthropic message."""
+    message_cls = self.get_message_cls(
+        self._safe_read(value, 'role', default='assistant')
+    )
+    content = self._safe_read(value, 'content', default=[])
+    assert isinstance(content, list)
+
+    chunks = []
+    thought_chunks = []
+    for part in content:
+      t = self._safe_read(part, 'type')
+      if t == 'text':
+        chunks.append(self._safe_read(part, 'text'))
+      elif t == 'thinking':
+        thought_chunks.append(self._safe_read(part, 'thinking'))
+      elif t in ('image', 'document'):
+        source = self._safe_read(part, 'source')
+        chunks.append(
+            lf_modalities.Mime.class_from_mime_type(
+                self._safe_read(source, 'media_type')
+            ).from_bytes(base64.b64decode(self._safe_read(source, 'data')))
+        )
+      else:
+        raise ValueError(f'Unsupported content part: {part!r}.')
+    message = message_cls.from_chunks(chunks)
+    if thought_chunks:
+      message.set('thought', message_cls.from_chunks(thought_chunks))
+    return message
+
+
+def _as_anthropic_format(
+    self,
+    chunk_preprocessor: Callable[[str | lf.Modality], Any] | None = None,
+    **kwargs
+) -> dict[str, Any]:
+  """Returns the Anthropic format of the chunk."""
+  return AnthropicMessageConverter(
+      chunk_preprocessor=chunk_preprocessor, **kwargs
+  ).to_value(self)
+
+
+@classmethod
+def _from_anthropic_format(
+    cls,
+    anthropic_message: dict[str, Any],
+    **kwargs
+) -> dict[str, Any]:
+  """Returns the Anthropic format of the chunk."""
+  del cls
+  return AnthropicMessageConverter(**kwargs).from_value(anthropic_message)
+
+# Set shortcut methods in lf.Message.
+lf.Message.as_anthropic_format = _as_anthropic_format
+lf.Message.from_anthropic_format = _from_anthropic_format

langfun/core/data/conversion/anthropic_test.py

@@ -0,0 +1,267 @@
+# 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 anthropic  # 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'
+)
+
+pdf_content = (
+    b'%PDF-1.1\n%\xc2\xa5\xc2\xb1\xc3\xab\n\n1 0 obj\n'
+    b'<< /Type /Catalog\n     /Pages 2 0 R\n  >>\nendobj\n\n2 0 obj\n '
+    b'<< /Type /Pages\n     /Kids [3 0 R]\n     '
+    b'/Count 1\n     /MediaBox [0 0 300 144]\n  '
+    b'>>\nendobj\n\n3 0 obj\n  '
+    b'<<  /Type /Page\n      /Parent 2 0 R\n      /Resources\n       '
+    b'<< /Font\n'
+    b'<< /F1\n'
+    b'<< /Type /Font\n'
+    b'/Subtype /Type1\n'
+    b'/BaseFont /Times-Roman\n'
+    b'>>\n>>\n>>\n      '
+    b'/Contents 4 0 R\n  >>\nendobj\n\n4 0 obj\n  '
+    b'<< /Length 55 >>\nstream\n  BT\n    /F1 18 Tf\n    0 0 Td\n    '
+    b'(Hello World) Tj\n  ET\nendstream\nendobj\n\nxref\n0 5\n0000000000 '
+    b'65535 f \n0000000018 00000 n \n0000000077 00000 n \n0000000178 00000 n '
+    b'\n0000000457 00000 n \ntrailer\n  <<  /Root 1 0 R\n      /Size 5\n  '
+    b'>>\nstartxref\n565\n%%EOF\n'
+)
+
+
+class AnthropicConversionTest(unittest.TestCase):
+
+  def test_as_format_with_role(self):
+    self.assertEqual(
+        lf.UserMessage('hi').as_format('anthropic'),
+        {
+            'role': 'user',
+            'content': [{'type': 'text', 'text': 'hi'}],
+        },
+    )
+    self.assertEqual(
+        lf.AIMessage('hi').as_format('anthropic'),
+        {
+            'role': 'assistant',
+            'content': [{'type': 'text', 'text': 'hi'}],
+        },
+    )
+    self.assertEqual(
+        lf.SystemMessage('hi').as_format('anthropic'),
+        {
+            'role': 'system',
+            'content': [{'type': 'text', 'text': 'hi'}],
+        },
+    )
+
+  def test_as_format_with_image(self):
+    self.assertEqual(
+        lf.Template(
+            'What are the common words from {{image}} and {{pdf}}?',
+            image=lf_modalities.Image.from_bytes(image_content),
+            pdf=lf_modalities.PDF.from_bytes(pdf_content),
+        ).render().as_anthropic_format(),
+        {
+            'role': 'user',
+            'content': [
+                {
+                    'type': 'text',
+                    'text': 'What are the common words from'
+                },
+                {
+                    'type': 'image',
+                    'source': {
+                        'type': 'base64',
+                        'media_type': 'image/png',
+                        'data': base64.b64encode(
+                            image_content
+                        ).decode('utf-8'),
+                    }
+                },
+                {
+                    'type': 'text',
+                    'text': 'and'
+                },
+                {
+                    'type': 'document',
+                    'source': {
+                        'type': 'base64',
+                        'media_type': 'application/pdf',
+                        'data': base64.b64encode(
+                            pdf_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_format(
+            'anthropic',
+            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': [{'type': 'text', 'text': 'this is a text'}],
+            },
+            format='anthropic',
+        ),
+        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': 'this is a text'}],
+            },
+            format='anthropic',
+        ),
+        lf.UserMessage('this is a text'),
+    )
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'role': 'assistant',
+                'content': [{'type': 'text', 'text': 'this is a text'}],
+            },
+            format='anthropic',
+        ),
+        lf.AIMessage('this is a text'),
+    )
+    self.assertEqual(
+        lf.Message.from_value(
+            {
+                'role': 'system',
+                'content': [{'type': 'text', 'text': 'this is a text'}],
+            },
+            format='anthropic',
+        ),
+        lf.SystemMessage('this is a text'),
+    )
+    with self.assertRaisesRegex(ValueError, 'Unsupported role: .*'):
+      lf.Message.from_value(
+          {
+              'role': 'function',
+              'content': [{'type': 'text', 'text': 'this is a text'}],
+          },
+          format='anthropic',
+      )
+
+  def test_from_value_with_thoughts(self):
+    message = lf.Message.from_anthropic_format(
+        {
+            'role': 'user',
+            'content': [
+                {
+                    'type': 'thinking',
+                    'thinking': 'this is a red round object',
+                },
+                {
+                    'type': 'text',
+                    'text': 'this is a apple',
+                },
+            ],
+        },
+    )
+    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_value(
+        {
+            'role': 'user',
+            'content': [
+                {
+                    'type': 'text',
+                    'text': 'What are the common words from'
+                },
+                {
+                    'type': 'image',
+                    'source': {
+                        'type': 'base64',
+                        'media_type': 'image/png',
+                        'data': base64.b64encode(image_content).decode('utf-8'),
+                    }
+                },
+                {
+                    'type': 'text',
+                    'text': 'and'
+                },
+                {
+                    'type': 'document',
+                    'source': {
+                        'type': 'base64',
+                        'media_type': 'application/pdf',
+                        'data': base64.b64encode(pdf_content).decode('utf-8'),
+                    }
+                },
+                {
+                    'type': 'text',
+                    'text': '?'
+                },
+            ],
+        },
+        format='anthropic',
+    )
+    self.assertEqual(
+        m.text,
+        'What are the common words from <<[[obj0]]>> and <<[[obj1]]>> ?'
+    )
+    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.to_bytes(), pdf_content)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/data/conversion/gemini.py

@@ -0,0 +1,168 @@
+# 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.
+"""Gemini API message conversion."""
+
+import base64
+from typing import Annotated, Any, Callable
+
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+
+
+class GeminiMessageConverter(lf.MessageConverter):
+  """Converter to Gemini public API."""
+
+  FORMAT_ID = 'gemini'
+
+  chunk_preprocessor: Annotated[
+      Callable[[str | lf.Modality], Any] | None,
+      (
+          'Chunk preprocessor for Langfun chunk to Gemini 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 Gemini API."""
+    parts = []
+    for chunk in message.chunk():
+      if self.chunk_preprocessor:
+        chunk = self.chunk_preprocessor(chunk)
+        if chunk is None:
+          continue
+
+      if isinstance(chunk, str):
+        parts.append(self._convert_chunk(chunk))
+      else:
+        if isinstance(chunk, lf_modalities.Mime):
+          modalities = [chunk]
+        # NOTE(daiyip): preprocessing may convert a single chunk into
+        # a list of chunks
+        elif isinstance(chunk, list):
+          modalities = chunk
+        else:
+          raise ValueError(f'Unsupported content type: {chunk!r}.')
+        parts.extend(self._convert_chunk(c) for c in modalities)
+    return dict(role=self.get_role(message), parts=parts)
+
+  def _convert_chunk(self, chunk: str | lf.Modality) -> Any:
+    """Converts a Langfun chunk to Gemini chunk."""
+    if isinstance(chunk, str):
+      return {'text': chunk}
+    if not isinstance(chunk, lf_modalities.Mime):
+      raise ValueError(f'Unsupported content chunk: {chunk!r}.')
+    # NOTE(daiyip): special handling for YouTube video.
+    if chunk.uri and chunk.uri.startswith('https://www.youtube.com/watch?v='):
+      return {
+          'fileData': {
+              'mimeType': 'video/*',
+              'fileUri': chunk.uri
+          }
+      }
+    if chunk.is_text:
+      return {'text': chunk.to_text()}
+    if chunk.uri and chunk.uri.lower().startswith(
+        ('http:', 'https:', 'ftp:')
+    ):
+      return {
+          'fileData': {
+              'mimeType': chunk.mime_type,
+              'fileUri': chunk.uri,
+          }
+      }
+    return {
+        'inlineData': {
+            'data': base64.b64encode(chunk.to_bytes()).decode(),
+            'mimeType': chunk.mime_type,
+        }
+    }
+
+  def from_value(self, value: dict[str, Any]) -> lf.Message:
+    """Returns a Langfun message from Gemini message."""
+    message_cls = self.get_message_cls(
+        self._safe_read(value, 'role', default='model')
+    )
+    parts = self._safe_read(value, 'parts', default=[])
+    assert isinstance(parts, list)
+
+    chunks = []
+    thought_chunks = []
+    for part in parts:
+      if 'text' in part:
+        text = self._safe_read(part, 'text')
+        if 'thought' in part:
+          thought_chunks.append(text)
+        else:
+          chunks.append(text)
+      elif 'inlineData' in part:
+        data = self._safe_read(part, 'inlineData')
+        chunks.append(
+            lf_modalities.Mime.class_from_mime_type(
+                self._safe_read(data, 'mimeType')
+            ).from_bytes(base64.b64decode(self._safe_read(data, 'data')))
+        )
+      elif 'fileData' in part:
+        data = self._safe_read(part, 'fileData')
+        chunks.append(
+            lf_modalities.Mime.class_from_mime_type(
+                self._safe_read(data, 'mimeType')
+            ).from_uri(self._safe_read(data, 'fileUri'))
+        )
+      else:
+        raise ValueError(f'Unsupported content part: {part!r}.')
+    message = message_cls.from_chunks(chunks)
+    if thought_chunks:
+      message.set('thought', message_cls.from_chunks(thought_chunks))
+    return message
+
+  @classmethod
+  def get_role(cls, message: lf.Message) -> str:
+    """Returns the role of the message."""
+    if isinstance(message, lf.AIMessage):
+      return 'model'
+    return super().get_role(message)
+
+  @classmethod
+  def get_message_cls(cls, role: str) -> type[lf.Message]:
+    """Returns the message class of the message."""
+    if role == 'model':
+      return lf.AIMessage
+    return super().get_message_cls(role)
+
+
+def _as_gemini_format(
+    self,
+    chunk_preprocessor: Callable[[str | lf.Modality], Any] | None = None,
+    **kwargs
+) -> dict[str, Any]:
+  """Returns the Gemini format of the chunk."""
+  return GeminiMessageConverter(
+      chunk_preprocessor=chunk_preprocessor, **kwargs
+  ).to_value(self)
+
+
+@classmethod
+def _from_gemini_format(
+    cls,
+    gemini_message: dict[str, Any],
+    **kwargs
+) -> dict[str, Any]:
+  """Returns the Gemini format of the chunk."""
+  del cls
+  return GeminiMessageConverter(**kwargs).from_value(gemini_message)
+
+# Set shortcut methods in lf.Message.
+lf.Message.as_gemini_format = _as_gemini_format
+lf.Message.from_gemini_format = _from_gemini_format