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

langfun/core/message_test.py

@@ -15,6 +15,7 @@
 
 import inspect
 import unittest
+from langfun.core import language_model
 from langfun.core import message
 from langfun.core import modality
 import pyglove as pg
@@ -36,11 +37,19 @@ class MessageTest(unittest.TestCase):
 
     d = pg.Dict(x=A())
 
-    m = message.UserMessage('hi', metadata=dict(x=1), x=pg.Ref(d.x), y=2)
+    m = message.UserMessage(
+        'hi',
+        metadata=dict(x=1), x=pg.Ref(d.x),
+        y=2,
+        tags=['lm-input']
+    )
     self.assertEqual(m.metadata, {'x': pg.Ref(d.x), 'y': 2})
     self.assertEqual(m.sender, 'User')
     self.assertIs(m.x, d.x)
     self.assertEqual(m.y, 2)
+    self.assertTrue(m.has_tag('lm-input'))
+    self.assertTrue(m.has_tag(('lm-input', '')))
+    self.assertFalse(m.has_tag('lm-response'))
 
     with self.assertRaises(AttributeError):
       _ = m.z
@@ -54,7 +63,7 @@ class MessageTest(unittest.TestCase):
     self.assertTrue(
         pg.eq(
             message.UserMessage.from_value(CustomModality('foo')),
-            message.UserMessage('{{object}}', object=CustomModality('foo')),
+            message.UserMessage('<<[[object]]>>', object=CustomModality('foo')),
         )
     )
     m = message.UserMessage('hi')
@@ -258,13 +267,16 @@ class MessageTest(unittest.TestCase):
 
   def test_referred_modalities(self):
     m1 = message.UserMessage(
-        'hi, this is a {{img1}} and {{x.img2}}',
+        'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>',
         img1=CustomModality('foo'),
         x=dict(img2=CustomModality('bar')),
     )
     m2 = message.SystemMessage('class Question:\n  image={{img1}}', source=m1)
     m3 = message.AIMessage(
-        'This is the {{output_image}} based on {{x.img2}}, {{unknown_var}}',
+        (
+            'This is the <<[[output_image]]>> based on <<[[x.img2]]>>, '
+            '{{unknown_var}}'
+        ),
         output_image=CustomModality('bar'),
         source=m2,
     )
@@ -276,11 +288,25 @@ class MessageTest(unittest.TestCase):
         },
     )
 
+  def test_text_with_modality_hash(self):
+    m = message.UserMessage(
+        'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>',
+        img1=CustomModality('foo'),
+        x=dict(img2=CustomModality('bar')),
+    )
+    self.assertEqual(
+        m.text_with_modality_hash,
+        (
+            'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>'
+            '<img1>acbd18db</img1><x.img2>37b51d19</x.img2>'
+        )
+    )
+
   def test_chunking(self):
     m = message.UserMessage(
         inspect.cleandoc("""
-            Hi, this is {{a}} and this is {{b}}.
-            {{x.c}} {{something else
+            Hi, this is <<[[a]]>> and this is {{b}}.
+            <<[[x.c]]>> {{something else
             """),
         a=CustomModality('foo'),
         x=dict(c=CustomModality('bar')),
@@ -292,9 +318,9 @@ class MessageTest(unittest.TestCase):
             [
                 'Hi, this is',
                 CustomModality('foo'),
-                'and this is {{b}}.',
+                'and this is {{b}}.\n',
                 CustomModality('bar'),
-                ' {{something else',
+                '{{something else',
             ],
         )
     )
@@ -303,11 +329,8 @@ class MessageTest(unittest.TestCase):
             message.AIMessage.from_chunks(chunks),
             message.AIMessage(
                 inspect.cleandoc("""
-                    Hi, this is
-                    {{obj0}}
-                    and this is {{b}}.
-                    {{obj1}}
-                     {{something else
+                    Hi, this is <<[[obj0]]>> and this is {{b}}.
+                    <<[[obj1]]>> {{something else
                     """),
                 obj0=pg.Ref(m.a),
                 obj1=pg.Ref(m.x.c),
@@ -315,6 +338,160 @@ class MessageTest(unittest.TestCase):
         )
     )
 
+  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_style(self):
+    self.assertIn(
+        inspect.cleandoc(
+            """
+            /* Langfun Message styles.*/
+            [class^="message-"] > details {
+                margin: 0px 0px 5px 0px;
+                border: 1px solid #EEE;
+            }
+            .lf-message.summary-title::after {
+                content: ' 💬';
+            }
+            details.pyglove.ai-message {
+                border: 1px solid blue;
+                color: blue;
+            }
+            details.pyglove.user-message {
+                border: 1px solid green;
+                color: green;
+            }
+            .message-tags {
+                margin: 5px 0px 5px 0px;
+                font-size: .8em;
+            }
+            .message-tags > span {
+                border-radius: 5px;
+                background-color: #CCC;
+                padding: 3px;
+                margin: 0px 2px 0px 2px;
+                color: white;
+            }
+            .message-text {
+                padding: 20px;
+                margin: 10px 5px 10px 5px;
+                font-style: italic;
+                white-space: pre-wrap;
+                border: 1px solid #EEE;
+                border-radius: 5px;
+                background-color: #EEE;
+            }
+            .modality-in-text {
+                display: inline-block;
+            }
+            .modality-in-text > details.pyglove {
+                display: inline-block;
+                font-size: 0.8em;
+                border: 0;
+                background-color: #A6F1A6;
+                margin: 0px 5px 0px 5px;
+            }
+            .message-result {
+                color: dodgerblue;
+            }
+            .message-usage {
+                color: orange;
+            }
+            .message-usage .object-key.str {
+                border: 1px solid orange;
+                background-color: orange;
+                color: white;
+            }
+            """
+        ),
+        message.UserMessage('hi').to_html().style_section,
+    )
+
+  def test_html_user_message(self):
+    self.assert_html_content(
+        message.UserMessage(
+            'what is a <div>'
+        ).to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove user-message lf-message"><summary><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"></div><div class="message-text">what is a &lt;div&gt;</div><div class="message-metadata"><details open class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></div></div></details>
+        """
+    )
+    self.assert_html_content(
+        message.UserMessage(
+            'what is this <<[[image]]>>',
+            tags=['lm-input'],
+            image=CustomModality('bird')
+        ).to_html(
+            enable_summary_tooltip=False,
+            extra_flags=dict(include_message_metadata=False)
+        ),
+        """
+        <details open class="pyglove user-message lf-message"><summary><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-input</span></div><div class="message-text">what is this<div class="modality-in-text"><details class="pyglove custom-modality"><summary><div class="summary-name">image<span class="tooltip">metadata.image</span></div><div class="summary-title">CustomModality(...)</div></summary><div class="complex-value custom-modality"><details open class="pyglove str"><summary><div class="summary-name">content<span class="tooltip">metadata.image.content</span></div><div class="summary-title">str</div></summary><span class="simple-value str">&#x27;bird&#x27;</span></details></div></details></div></div></div></details>
+        """
+    )
+
+  def test_html_ai_message(self):
+    image = CustomModality('foo')
+    user_message = message.UserMessage(
+        'What is in this image? <<[[image]]>> this is a test',
+        metadata=dict(image=image),
+        source=message.UserMessage('User input'),
+        tags=['lm-input']
+    )
+    ai_message = message.AIMessage(
+        'My name is Gemini',
+        metadata=dict(
+            result=pg.Dict(x=1, y=2, z=pg.Dict(a=[12, 323])),
+            usage=language_model.LMSamplingUsage(10, 2, 12)
+        ),
+        tags=['lm-response', 'lm-output'],
+        source=user_message,
+    )
+    self.assert_html_content(
+        ai_message.to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove ai-message lf-message"><summary><div class="summary-title lf-message">AIMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-response</span><span>lm-output</span></div><div class="message-text">My name is Gemini</div><div class="message-result"><details open class="pyglove dict"><summary><div class="summary-name">result<span class="tooltip">metadata.result</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details open class="pyglove int"><summary><div class="summary-name">x<span class="tooltip">metadata.result.x</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details><details open class="pyglove int"><summary><div class="summary-name">y<span class="tooltip">metadata.result.y</span></div><div class="summary-title">int</div></summary><span class="simple-value int">2</span></details><details class="pyglove dict"><summary><div class="summary-name">z<span class="tooltip">metadata.result.z</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details class="pyglove list"><summary><div class="summary-name">a<span class="tooltip">metadata.result.z.a</span></div><div class="summary-title">List(...)</div></summary><div class="complex-value list"><table><tr><td><span class="object-key int">0</span><span class="tooltip">metadata.result.z.a[0]</span></td><td><span class="simple-value int">12</span></td></tr><tr><td><span class="object-key int">1</span><span class="tooltip">metadata.result.z.a[1]</span></td><td><span class="simple-value int">323</span></td></tr></table></div></details></div></details></div></details></div><div class="message-usage"><details open class="pyglove lm-sampling-usage"><summary><div class="summary-name">llm usage<span class="tooltip">metadata.usage</span></div><div class="summary-title">LMSamplingUsage(...)</div></summary><div class="complex-value lm-sampling-usage"><table><tr><td><span class="object-key str">prompt_tokens</span><span class="tooltip">metadata.usage.prompt_tokens</span></td><td><span class="simple-value int">10</span></td></tr><tr><td><span class="object-key str">completion_tokens</span><span class="tooltip">metadata.usage.completion_tokens</span></td><td><span class="simple-value int">2</span></td></tr><tr><td><span class="object-key str">total_tokens</span><span class="tooltip">metadata.usage.total_tokens</span></td><td><span class="simple-value int">12</span></td></tr><tr><td><span class="object-key str">num_requests</span><span class="tooltip">metadata.usage.num_requests</span></td><td><span class="simple-value int">1</span></td></tr><tr><td><span class="object-key str">estimated_cost</span><span class="tooltip">metadata.usage.estimated_cost</span></td><td><span class="simple-value none-type">None</span></td></tr><tr><td><span class="object-key str">retry_stats</span><span class="tooltip">metadata.usage.retry_stats</span></td><td><details class="pyglove retry-stats"><summary><div class="summary-title">RetryStats(...)</div></summary><div class="complex-value retry-stats"><table><tr><td><span class="object-key str">num_occurences</span><span class="tooltip">metadata.usage.retry_stats.num_occurences</span></td><td><span class="simple-value int">0</span></td></tr><tr><td><span class="object-key str">total_wait_interval</span><span class="tooltip">metadata.usage.retry_stats.total_wait_interval</span></td><td><span class="simple-value float">0.0</span></td></tr><tr><td><span class="object-key str">total_call_interval</span><span class="tooltip">metadata.usage.retry_stats.total_call_interval</span></td><td><span class="simple-value float">0.0</span></td></tr><tr><td><span class="object-key str">errors</span><span class="tooltip">metadata.usage.retry_stats.errors</span></td><td><details class="pyglove dict"><summary><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></td></tr></table></div></details></td></tr></table></div></details></div><div class="message-metadata"><details open class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></div><details open class="pyglove user-message lf-message"><summary><div class="summary-name lf-message">source<span class="tooltip lf-message">source</span></div><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-input</span></div><div class="message-text">What is in this image?<div class="modality-in-text"><details class="pyglove custom-modality"><summary><div class="summary-name">image<span class="tooltip">source.metadata.image</span></div><div class="summary-title">CustomModality(...)</div></summary><div class="complex-value custom-modality"><details open class="pyglove str"><summary><div class="summary-name">content<span class="tooltip">source.metadata.image.content</span></div><div class="summary-title">str</div></summary><span class="simple-value str">&#x27;foo&#x27;</span></details></div></details></div>this is a test</div><div class="message-metadata"><details open class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">source.metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><details class="pyglove custom-modality"><summary><div class="summary-name">image<span class="tooltip">source.metadata.image</span></div><div class="summary-title">CustomModality(...)</div></summary><div class="complex-value custom-modality"><details open class="pyglove str"><summary><div class="summary-name">content<span class="tooltip">source.metadata.image.content</span></div><div class="summary-title">str</div></summary><span class="simple-value str">&#x27;foo&#x27;</span></details></div></details></div></details></div></div></details></div></details>
+        """
+    )
+    self.assert_html_content(
+        ai_message.to_html(
+            key_style='label',
+            enable_summary_tooltip=False,
+            extra_flags=dict(
+                collapse_modalities_in_text=False,
+                collapse_llm_usage=True,
+                collapse_message_result_level=0,
+                collapse_message_metadata_level=0,
+                collapse_source_message_level=0,
+                source_tag=None,
+            ),
+        ),
+        """
+        <details open class="pyglove ai-message lf-message"><summary><div class="summary-title lf-message">AIMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-response</span><span>lm-output</span></div><div class="message-text">My name is Gemini</div><div class="message-result"><details class="pyglove dict"><summary><div class="summary-name">result<span class="tooltip">metadata.result</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><table><tr><td><span class="object-key str">x</span><span class="tooltip">metadata.result.x</span></td><td><span class="simple-value int">1</span></td></tr><tr><td><span class="object-key str">y</span><span class="tooltip">metadata.result.y</span></td><td><span class="simple-value int">2</span></td></tr><tr><td><span class="object-key str">z</span><span class="tooltip">metadata.result.z</span></td><td><details class="pyglove dict"><summary><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><table><tr><td><span class="object-key str">a</span><span class="tooltip">metadata.result.z.a</span></td><td><details class="pyglove list"><summary><div class="summary-title">List(...)</div></summary><div class="complex-value list"><table><tr><td><span class="object-key int">0</span><span class="tooltip">metadata.result.z.a[0]</span></td><td><span class="simple-value int">12</span></td></tr><tr><td><span class="object-key int">1</span><span class="tooltip">metadata.result.z.a[1]</span></td><td><span class="simple-value int">323</span></td></tr></table></div></details></td></tr></table></div></details></td></tr></table></div></details></div><div class="message-usage"><details class="pyglove lm-sampling-usage"><summary><div class="summary-name">llm usage<span class="tooltip">metadata.usage</span></div><div class="summary-title">LMSamplingUsage(...)</div></summary><div class="complex-value lm-sampling-usage"><table><tr><td><span class="object-key str">prompt_tokens</span><span class="tooltip">metadata.usage.prompt_tokens</span></td><td><span class="simple-value int">10</span></td></tr><tr><td><span class="object-key str">completion_tokens</span><span class="tooltip">metadata.usage.completion_tokens</span></td><td><span class="simple-value int">2</span></td></tr><tr><td><span class="object-key str">total_tokens</span><span class="tooltip">metadata.usage.total_tokens</span></td><td><span class="simple-value int">12</span></td></tr><tr><td><span class="object-key str">num_requests</span><span class="tooltip">metadata.usage.num_requests</span></td><td><span class="simple-value int">1</span></td></tr><tr><td><span class="object-key str">estimated_cost</span><span class="tooltip">metadata.usage.estimated_cost</span></td><td><span class="simple-value none-type">None</span></td></tr><tr><td><span class="object-key str">retry_stats</span><span class="tooltip">metadata.usage.retry_stats</span></td><td><details class="pyglove retry-stats"><summary><div class="summary-title">RetryStats(...)</div></summary><div class="complex-value retry-stats"><table><tr><td><span class="object-key str">num_occurences</span><span class="tooltip">metadata.usage.retry_stats.num_occurences</span></td><td><span class="simple-value int">0</span></td></tr><tr><td><span class="object-key str">total_wait_interval</span><span class="tooltip">metadata.usage.retry_stats.total_wait_interval</span></td><td><span class="simple-value float">0.0</span></td></tr><tr><td><span class="object-key str">total_call_interval</span><span class="tooltip">metadata.usage.retry_stats.total_call_interval</span></td><td><span class="simple-value float">0.0</span></td></tr><tr><td><span class="object-key str">errors</span><span class="tooltip">metadata.usage.retry_stats.errors</span></td><td><details class="pyglove dict"><summary><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></td></tr></table></div></details></td></tr></table></div></details></div><div class="message-metadata"><details class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></div><details class="pyglove user-message lf-message"><summary><div class="summary-name lf-message">source<span class="tooltip lf-message">source</span></div><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-input</span></div><div class="message-text">What is in this image?<div class="modality-in-text"><details open class="pyglove custom-modality"><summary><div class="summary-name">image<span class="tooltip">source.metadata.image</span></div><div class="summary-title">CustomModality(...)</div></summary><div class="complex-value custom-modality"><table><tr><td><span class="object-key str">content</span><span class="tooltip">source.metadata.image.content</span></td><td><span class="simple-value str">&#x27;foo&#x27;</span></td></tr></table></div></details></div>this is a test</div><div class="message-metadata"><details class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">source.metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><table><tr><td><span class="object-key str">image</span><span class="tooltip">source.metadata.image</span></td><td><details class="pyglove custom-modality"><summary><div class="summary-title">CustomModality(...)</div></summary><div class="complex-value custom-modality"><table><tr><td><span class="object-key str">content</span><span class="tooltip">source.metadata.image.content</span></td><td><span class="simple-value str">&#x27;foo&#x27;</span></td></tr></table></div></details></td></tr></table></div></details></div><details class="pyglove user-message lf-message"><summary><div class="summary-name lf-message">source<span class="tooltip lf-message">source.source</span></div><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"></div><div class="message-text">User input</div><div class="message-metadata"><details class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">source.source.metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></div></div></details></div></details></div></details>
+        """
+    )
+    self.assert_html_content(
+        ai_message.to_html(
+            key_style='label',
+            enable_summary_tooltip=False,
+            extra_flags=dict(
+                collapse_modalities_in_text=True,
+                collapse_llm_usage=False,
+                collapse_message_result_level=1,
+                collapse_message_metadata_level=1,
+                collapse_source_message_level=2,
+                source_tag=None,
+            ),
+        ),
+        """
+        <details open class="pyglove ai-message lf-message"><summary><div class="summary-title lf-message">AIMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-response</span><span>lm-output</span></div><div class="message-text">My name is Gemini</div><div class="message-result"><details open class="pyglove dict"><summary><div class="summary-name">result<span class="tooltip">metadata.result</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><table><tr><td><span class="object-key str">x</span><span class="tooltip">metadata.result.x</span></td><td><span class="simple-value int">1</span></td></tr><tr><td><span class="object-key str">y</span><span class="tooltip">metadata.result.y</span></td><td><span class="simple-value int">2</span></td></tr><tr><td><span class="object-key str">z</span><span class="tooltip">metadata.result.z</span></td><td><details class="pyglove dict"><summary><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><table><tr><td><span class="object-key str">a</span><span class="tooltip">metadata.result.z.a</span></td><td><details class="pyglove list"><summary><div class="summary-title">List(...)</div></summary><div class="complex-value list"><table><tr><td><span class="object-key int">0</span><span class="tooltip">metadata.result.z.a[0]</span></td><td><span class="simple-value int">12</span></td></tr><tr><td><span class="object-key int">1</span><span class="tooltip">metadata.result.z.a[1]</span></td><td><span class="simple-value int">323</span></td></tr></table></div></details></td></tr></table></div></details></td></tr></table></div></details></div><div class="message-usage"><details open class="pyglove lm-sampling-usage"><summary><div class="summary-name">llm usage<span class="tooltip">metadata.usage</span></div><div class="summary-title">LMSamplingUsage(...)</div></summary><div class="complex-value lm-sampling-usage"><table><tr><td><span class="object-key str">prompt_tokens</span><span class="tooltip">metadata.usage.prompt_tokens</span></td><td><span class="simple-value int">10</span></td></tr><tr><td><span class="object-key str">completion_tokens</span><span class="tooltip">metadata.usage.completion_tokens</span></td><td><span class="simple-value int">2</span></td></tr><tr><td><span class="object-key str">total_tokens</span><span class="tooltip">metadata.usage.total_tokens</span></td><td><span class="simple-value int">12</span></td></tr><tr><td><span class="object-key str">num_requests</span><span class="tooltip">metadata.usage.num_requests</span></td><td><span class="simple-value int">1</span></td></tr><tr><td><span class="object-key str">estimated_cost</span><span class="tooltip">metadata.usage.estimated_cost</span></td><td><span class="simple-value none-type">None</span></td></tr><tr><td><span class="object-key str">retry_stats</span><span class="tooltip">metadata.usage.retry_stats</span></td><td><details class="pyglove retry-stats"><summary><div class="summary-title">RetryStats(...)</div></summary><div class="complex-value retry-stats"><table><tr><td><span class="object-key str">num_occurences</span><span class="tooltip">metadata.usage.retry_stats.num_occurences</span></td><td><span class="simple-value int">0</span></td></tr><tr><td><span class="object-key str">total_wait_interval</span><span class="tooltip">metadata.usage.retry_stats.total_wait_interval</span></td><td><span class="simple-value float">0.0</span></td></tr><tr><td><span class="object-key str">total_call_interval</span><span class="tooltip">metadata.usage.retry_stats.total_call_interval</span></td><td><span class="simple-value float">0.0</span></td></tr><tr><td><span class="object-key str">errors</span><span class="tooltip">metadata.usage.retry_stats.errors</span></td><td><details class="pyglove dict"><summary><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></td></tr></table></div></details></td></tr></table></div></details></div><div class="message-metadata"><details open class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></div><details open class="pyglove user-message lf-message"><summary><div class="summary-name lf-message">source<span class="tooltip lf-message">source</span></div><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-input</span></div><div class="message-text">What is in this image?<div class="modality-in-text"><details class="pyglove custom-modality"><summary><div class="summary-name">image<span class="tooltip">source.metadata.image</span></div><div class="summary-title">CustomModality(...)</div></summary><div class="complex-value custom-modality"><table><tr><td><span class="object-key str">content</span><span class="tooltip">source.metadata.image.content</span></td><td><span class="simple-value str">&#x27;foo&#x27;</span></td></tr></table></div></details></div>this is a test</div><div class="message-metadata"><details open class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">source.metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><table><tr><td><span class="object-key str">image</span><span class="tooltip">source.metadata.image</span></td><td><details class="pyglove custom-modality"><summary><div class="summary-title">CustomModality(...)</div></summary><div class="complex-value custom-modality"><table><tr><td><span class="object-key str">content</span><span class="tooltip">source.metadata.image.content</span></td><td><span class="simple-value str">&#x27;foo&#x27;</span></td></tr></table></div></details></td></tr></table></div></details></div><details open class="pyglove user-message lf-message"><summary><div class="summary-name lf-message">source<span class="tooltip lf-message">source.source</span></div><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"></div><div class="message-text">User input</div><div class="message-metadata"><details open class="pyglove dict message-metadata"><summary><div class="summary-name message-metadata">metadata<span class="tooltip message-metadata">source.source.metadata</span></div><div class="summary-title message-metadata">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></div></div></details></div></details></div></details>
+        """
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/modalities/__init__.py

@@ -17,10 +17,14 @@
 # pylint: disable=g-bad-import-order
 # pylint: disable=g-import-not-at-top
 
-from langfun.core.modalities.mime import MimeType
+from langfun.core.modalities.audio import Audio
+from langfun.core.modalities.mime import Mime
 from langfun.core.modalities.mime import Custom
-from langfun.core.modalities.mime import PDF
+from langfun.core.modalities.ms_office import Docx
+from langfun.core.modalities.ms_office import Pptx
+from langfun.core.modalities.ms_office import Xlsx
 from langfun.core.modalities.image import Image
+from langfun.core.modalities.pdf import PDF
 from langfun.core.modalities.video import Video
 
 # pylint: enable=g-import-not-at-top

langfun/core/modalities/audio.py

@@ -0,0 +1,30 @@
+# 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.
+# 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.
+"""Audio types."""
+
+import functools
+from langfun.core.modalities import mime
+
+
+class Audio(mime.Mime):
+  """Audio."""
+
+  MIME_PREFIX = 'audio'
+
+  @functools.cached_property
+  def audio_format(self) -> str:
+    return self.mime_type.removeprefix(self.MIME_PREFIX + '/')
+
+  def _mime_control_for(self, uri: str) -> str:
+    return f'<audio controls> <source src="{uri}"> </audio>'

langfun/core/modalities/audio_test.py

@@ -0,0 +1,63 @@
+# 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.
+# 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.
+"""Audio tests."""
+import unittest
+from unittest import mock
+
+from langfun.core.modalities import audio as audio_lib
+import pyglove as pg
+
+
+content_bytes = (
+    b'RIFF$\x00\x00\x00WAVEfmt'
+    b' \x10\x00\x00\x00\x01\x00\x01\x00\x00\x04\x00\x00\x00\x04\x00\x00\x01\x00\x08\x00data\x00\x00\x00\x00'
+)
+
+
+def mock_request(*args, **kwargs):
+  del args, kwargs
+  return pg.Dict(content=content_bytes)
+
+
+class AudioTest(unittest.TestCase):
+
+  def test_audio_content(self):
+    audio = audio_lib.Audio.from_bytes(content_bytes)
+    self.assertEqual(audio.mime_type, 'audio/x-wav')
+    self.assertEqual(audio.audio_format, 'x-wav')
+    self.assertEqual(audio.to_bytes(), content_bytes)
+
+  def test_bad_audio(self):
+    audio = audio_lib.Audio.from_bytes(b'bad')
+    with self.assertRaisesRegex(ValueError, 'Expected MIME type'):
+      _ = audio.audio_format
+
+
+class AudioFileTest(unittest.TestCase):
+
+  def test_audio_file(self):
+    audio = audio_lib.Audio.from_uri('http://mock/web/a.wav')
+    with mock.patch('requests.get') as mock_requests_get:
+      mock_requests_get.side_effect = mock_request
+      self.assertEqual(audio.audio_format, 'x-wav')
+      self.assertEqual(audio.mime_type, 'audio/x-wav')
+      self.assertEqual(
+          audio._raw_html(),
+          '<audio controls> <source src="http://mock/web/a.wav"> </audio>',
+      )
+      self.assertEqual(audio.to_bytes(), content_bytes)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/modalities/image.py

@@ -13,28 +13,47 @@
 # limitations under the License.
 """Image modality."""
 
-import base64
-import imghdr
-from typing import cast
+import functools
+import io
+from typing import Any
+
 from langfun.core.modalities import mime
 
+try:
+  from PIL import Image as pil_image    # pylint: disable=g-import-not-at-top
+  PILImage = pil_image.Image
+  pil_open = pil_image.open
+except ImportError:
+  PILImage = Any
+
+  def pil_open(*unused_args, **unused_kwargs):
+    raise RuntimeError(
+        'Please install "langfun[mime-pil]" to enable PIL image support.'
+    )
+
+
+class Image(mime.Mime):
+  """Image."""
 
-class Image(mime.MimeType):
-  """Base class for image."""
+  MIME_PREFIX = 'image'
 
-  @property
+  @functools.cached_property
   def image_format(self) -> str:
-    iformat = imghdr.what(None, self.to_bytes())
-    if iformat not in ['png', 'jpeg']:
-      raise ValueError(f'Unsupported image format: {iformat!r}.')
-    return cast(str, iformat)
-
-  @property
-  def mime_type(self) -> str:
-    return f'image/{self.image_format}'
-
-  def _repr_html_(self) -> str:
-    if self.uri and self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
-      return f'<img src="{self.uri}">'
-    image_raw = base64.b64encode(self.to_bytes()).decode()
-    return f'<img src="data:image/{self.image_format};base64,{image_raw}">'
+    return self.mime_type.removeprefix(self.MIME_PREFIX + '/')
+
+  def _mime_control_for(self, uri: str) -> str:
+    return f'<img src="{uri}">'
+
+  @functools.cached_property
+  def size(self) -> tuple[int, int]:
+    img = pil_open(io.BytesIO(self.to_bytes()))
+    return img.size
+
+  def to_pil_image(self) -> PILImage:   # pytype: disable=invalid-annotation
+    return pil_open(io.BytesIO(self.to_bytes()))
+
+  @classmethod
+  def from_pil_image(cls, img: PILImage) -> 'Image':  # pytype: disable=invalid-annotation
+    buf = io.BytesIO()
+    img.save(buf, format='PNG')
+    return cls.from_bytes(buf.getvalue())

langfun/core/modalities/image_test.py

@@ -12,10 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Image tests."""
+import io
 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 PIL.Image as pil_image
 import pyglove as pg
 
 
@@ -36,30 +40,69 @@ 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.assertIn('data:image/png;base64,', image._raw_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, 'Unsupported image format'):
+    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
       self.assertEqual(image.image_format, 'png')
-      self.assertEqual(image._repr_html_(), '<img src="http://mock/web/a.png">')
+      self.assertEqual(
+          image._raw_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._raw_html(),
+          '<img src="http://mock/web/a.png">'
+      )
       self.assertEqual(image.to_bytes(), image_content)
 
+  def test_image_size(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
+      self.assertEqual(image.size, (24, 24))
+
+  def test_to_pil_image(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
+      self.assertIsInstance(image.to_pil_image(), pil_image.Image)
+
+  def test_from_pil_image(self):
+    image = pil_image.open(io.BytesIO(image_content))
+    self.assertIsInstance(
+        image_lib.Image.from_pil_image(image), image_lib.Image
+    )
+
 
 if __name__ == '__main__':
   unittest.main()