langfun 0.1.2.dev202510230805__py3-none-any.whl → 0.1.2.dev202511270805__py3-none-any.whl

langfun/core/message_test.py

@@ -60,10 +60,19 @@ class MessageTest(unittest.TestCase):
     self.assertTrue(
         pg.eq(message.UserMessage.from_value('hi'), message.UserMessage('hi'))
     )
+    foo = CustomModality('foo')
     self.assertTrue(
         pg.eq(
-            message.UserMessage.from_value(CustomModality('foo')),
-            message.UserMessage('<<[[object]]>>', object=CustomModality('foo')),
+            message.UserMessage.from_value(foo),
+            message.UserMessage(f'<<[[{foo.id}]]>>', referred_modalities=[foo]),
+        )
+    )
+    self.assertTrue(
+        pg.eq(
+            message.UserMessage.from_value(foo),
+            message.UserMessage(
+                f'<<[[{foo.id}]]>>', referred_modalities={foo.id: foo}
+            ),
         )
     )
     m = message.UserMessage('hi')
@@ -224,76 +233,39 @@ class MessageTest(unittest.TestCase):
     self.assertEqual(str(m), m.text)
 
   def test_get_modality(self):
+    foo = CustomModality('foo')
+    bar = CustomModality('bar')
     m1 = message.UserMessage(
-        'hi, this is a {{img1}} and {{x.img2}}',
-        img1=CustomModality('foo'),
-        x=dict(img2=pg.Ref(CustomModality('bar'))),
+        'hi',
+        referred_modalities={
+            foo.id: foo,
+            bar.id: pg.Ref(bar),
+        },
     )
-    self.assertIs(m1.get_modality('img1'), m1.img1)
-    self.assertIs(m1.get_modality('x.img2'), m1.x.img2)
+    self.assertIs(m1.get_modality(foo.id), foo)
+    self.assertIs(m1.get_modality(bar.id), bar)
     self.assertIsNone(m1.get_modality('video'))
+    self.assertEqual(len(m1.modalities()), 2)
+    self.assertEqual(len(m1.modalities(CustomModality)), 2)
 
-    m2 = message.SystemMessage('class Question:\n  image={{img1}}', source=m1)
-    self.assertIs(m2.get_modality('img1'), m1.img1)
-    # We could get the modality object even it's not directly used by current
-    # message.
-    self.assertIs(m2.get_modality('x.img2'), m1.x.img2)
-    self.assertIsNone(m2.get_modality('video'))
-
-    m3 = message.AIMessage(
-        'This is the {{output_image}} based on {{x.img2}}',
-        output_image=CustomModality('bar'),
-        source=m2,
-    )
-    self.assertIs(m3.get_modality('x.img2'), m1.x.img2)
-    self.assertIs(m3.get_modality('output_image'), m3.output_image)
-    self.assertIsNone(m3.get_modality('video'))
-
-  def test_referred_modalities(self):
-    m1 = message.UserMessage(
-        '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}}'
-        ),
-        output_image=CustomModality('bar'),
-        source=m2,
-    )
-    self.assertEqual(
-        m3.referred_modalities(),
-        {
-            'output_image': m3.output_image,
-            'x.img2': m1.x.img2,
-        },
-    )
+    class MyModality(modality.Modality):
+      pass
 
-  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>'
-        )
-    )
+    self.assertEqual(len(m1.modalities(MyModality)), 0)
+    self.assertEqual(len(m1.modalities(lambda x: x.content == 'foo')), 1)
 
   def test_chunking(self):
+    foo = CustomModality('foo')
+    bar = CustomModality('bar')
     m = message.UserMessage(
-        inspect.cleandoc("""
-            Hi, this is <<[[a]]>> and this is {{b}}.
-            <<[[x.c]]>> {{something else
+        inspect.cleandoc(f"""
+            Hi, this is <<[[{foo.id}]]>> and this is {{b}}.
+            <<[[{bar.id}]]>> something else
             """),
-        a=CustomModality('foo'),
-        x=dict(c=CustomModality('bar')),
+        referred_modalities={
+            foo.id: pg.Ref(foo),
+            bar.id: pg.Ref(bar),
+        },
     )
     chunks = m.chunk()
     self.assertTrue(
@@ -301,10 +273,10 @@ class MessageTest(unittest.TestCase):
             chunks,
             [
                 'Hi, this is',
-                CustomModality('foo'),
-                'and this is {{b}}.\n',
-                CustomModality('bar'),
-                '{{something else',
+                foo,
+                'and this is {b}.\n',
+                bar,
+                'something else',
             ],
         )
     )
@@ -312,15 +284,17 @@ class MessageTest(unittest.TestCase):
         pg.eq(
             message.AIMessage.from_chunks(chunks),
             message.AIMessage(
-                inspect.cleandoc("""
-                    Hi, this is <<[[obj0]]>> and this is {{b}}.
-                    <<[[obj1]]>> {{something else
-                    """),
-                obj0=pg.Ref(m.a),
-                obj1=pg.Ref(m.x.c),
+                (
+                    f'Hi, this is <<[[{foo.id}]]>> and this '
+                    f'is {{b}}.\n<<[[{bar.id}]]>> '
+                    'something else'
+                ),
+                referred_modalities=[foo, bar],
             ),
         )
     )
+    with self.assertRaisesRegex(ValueError, 'Unknown modality reference'):
+      message.UserMessage('<<[[abc]]>>').chunk()
 
   def assert_html_content(self, html, expected):
     expected = inspect.cleandoc(expected).strip()
@@ -404,25 +378,26 @@ class MessageTest(unittest.TestCase):
         <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>
         """
     )
+    image = CustomModality('bird')
     self.assert_html_content(
         message.UserMessage(
-            'what is this <<[[image]]>>',
+            f'what is this <<[[{image.id}]]>>',
             tags=['lm-input'],
-            image=CustomModality('bird')
+            referred_modalities=[image],
         ).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>
+        <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">custom_modality:abaecf8c<span class="tooltip"></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">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),
+        f'What is in this image? <<[[{image.id}]]>> this is a test',
+        referred_modalities=[image],
         source=message.UserMessage('User input'),
         tags=['lm-input']
     )
@@ -438,7 +413,7 @@ class MessageTest(unittest.TestCase):
     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><tr><td><span class="object-key str">completion_tokens_details</span><span class="tooltip">metadata.usage.completion_tokens_details</span></td><td><span class="simple-value none-type">None</span></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>
+        <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><tr><td><span class="object-key str">completion_tokens_details</span><span class="tooltip">metadata.usage.completion_tokens_details</span></td><td><span class="simple-value none-type">None</span></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">custom_modality:acbd18db<span class="tooltip"></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">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"><span class="empty-container"></span></div></details></div></div></details></div></details>
         """
     )
     self.assert_html_content(
@@ -455,7 +430,7 @@ class MessageTest(unittest.TestCase):
             ),
         ),
         """
-        <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><tr><td><span class="object-key str">completion_tokens_details</span><span class="tooltip">metadata.usage.completion_tokens_details</span></td><td><span class="simple-value none-type">None</span></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>
+        <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><tr><td><span class="object-key str">completion_tokens_details</span><span class="tooltip">metadata.usage.completion_tokens_details</span></td><td><span class="simple-value none-type">None</span></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">custom_modality:acbd18db<span class="tooltip"></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">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"><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.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(
@@ -472,7 +447,7 @@ class MessageTest(unittest.TestCase):
             ),
         ),
         """
-        <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><tr><td><span class="object-key str">completion_tokens_details</span><span class="tooltip">metadata.usage.completion_tokens_details</span></td><td><span class="simple-value none-type">None</span></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>
+        <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><tr><td><span class="object-key str">completion_tokens_details</span><span class="tooltip">metadata.usage.completion_tokens_details</span></td><td><span class="simple-value none-type">None</span></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">custom_modality:acbd18db<span class="tooltip"></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">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"><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.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>
         """
     )
 
@@ -525,6 +500,12 @@ class MessageConverterTest(unittest.TestCase):
     self.assertIn('test_format2', message.Message.convertible_formats())
     self.assertIn('test_format3', message.Message.convertible_formats())
 
+    self.assertTrue(message.Message.is_convertible(int))
+    self.assertFalse(message.Message.is_convertible(dict))
+    self.assertTrue(message.Message.is_convertible('test_format1'))
+    self.assertTrue(message.Message.is_convertible('test_format2'))
+    self.assertTrue(message.Message.is_convertible('test_format3'))
+    self.assertFalse(message.Message.is_convertible('test_format4'))
     self.assertIn(int, message.Message.convertible_types())
     self.assertIn(tuple, message.Message.convertible_types())
     self.assertEqual(
@@ -565,6 +546,9 @@ class MessageConverterTest(unittest.TestCase):
         message.Message.from_value((1, 2, 3)),
         message.UserMessage('1,2,3')
     )
+    message.MessageConverter._REGISTRY.unregister(TestConverter)
+    message.MessageConverter._REGISTRY.unregister(TestConverter2)
+    message.MessageConverter._REGISTRY.unregister(TestConverter3)
 
   def test_get_role(self):
     self.assertEqual(

langfun/core/modalities/__init__.py

@@ -24,6 +24,14 @@ from langfun.core.modalities.image import Image
 from langfun.core.modalities.pdf import PDF
 from langfun.core.modalities.video import Video
 
+from langfun.core import message as _message_lib
+
+# Override the `images`, `videos` and `audios` properties of `Message` to
+# return the modalities of the corresponding types.
+_message_lib.Message.images = property(lambda self: self.modalities(Image))
+_message_lib.Message.videos = property(lambda self: self.modalities(Video))
+_message_lib.Message.audios = property(lambda self: self.modalities(Audio))
+
 # pylint: enable=g-import-not-at-top
 # pylint: enable=g-bad-import-order
 # pylint: enable=g-importing-member

langfun/core/modalities/audio.py

@@ -18,7 +18,27 @@ from langfun.core.modalities import mime
 
 
 class Audio(mime.Mime):
-  """Audio."""
+  """Represents audio for communicating with language models.
+
+  `lf.Audio` can be initialized from a URI (HTTP/HTTPS URL or local path)
+  using `lf.Audio.from_uri()` or from raw bytes using `lf.Audio.from_bytes()`.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  # Load audio from path
+  audio = lf.Audio.from_path('/path/to/audio.mp3')
+
+  # Use audio in a prompt
+  prompt = lf.Template(
+      'What is being said in this audio? {{audio}}', audio=audio
+  )
+  response = lf.query(prompt, lm=lf.llms.Gemini25Flash())
+  print(response)
+  ```
+  """
 
   MIME_PREFIX = 'audio'
 

langfun/core/modalities/image.py

@@ -33,7 +33,25 @@ except ImportError:
 
 
 class Image(mime.Mime):
-  """Image."""
+  """Represents an image for communicating with language models.
+
+  `lf.Image` can be initialized from a URI (HTTP/HTTPS URL or local path)
+  using `lf.Image.from_uri()` or from raw bytes using `lf.Image.from_bytes()`.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  # Load image from path
+  image = lf.Image.from_path('/path/to/image.png')
+
+  # Use image in a prompt
+  prompt = lf.Template('Describe this image: {{image}}', image=image)
+  response = lf.query(prompt, lm=lf.llms.Gemini25Flash())
+  print(response)
+  ```
+  """
 
   MIME_PREFIX = 'image'
 

langfun/core/modalities/mime.py

@@ -15,6 +15,7 @@
 
 import base64
 import functools
+import hashlib
 from typing import Annotated, Any, Iterable, Type, Union
 import langfun.core as lf
 # Placeholder for Google-internal internet access import.
@@ -36,7 +37,33 @@ def _detect_mime_type(content: bytes) -> str:
 
 
 class Mime(lf.Modality):
-  """Base for MIME data."""
+  """Base class for representing modality data based on MIME types.
+
+  `lf.Mime` is a subclass of `lf.Modality` that serves as a base for
+  handling various data types like images, audio, video, and PDFs,
+  identified by their MIME types. It provides unified methods for
+  loading data from URIs or bytes (`.from_uri()`, `.from_bytes()`) and
+  for accessing content (`.to_bytes()`).
+
+  Subclasses like `lf.Image`, `lf.Audio`, `lf.Video`, and `lf.PDF`
+  specialize in handling specific MIME type prefixes (e.g., 'image/', 'audio/').
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  # Load an image from a path
+  image = lf.Image.from_path('/path/to/image.png')
+  print(image.mime_type)
+  # Output: image/png
+
+  # Create a text document
+  text = lf.Custom.from_bytes(b'hello world', mime='text/plain')
+  print(text.mime_type)
+  # Output: text/plain
+  ```
+  """
 
   # The regular expression that describes the MIME type str.
   # If None, the MIME type is dynamic. Subclass could override.
@@ -48,6 +75,10 @@ class Mime(lf.Modality):
       Union[str, bytes, None], 'The raw content of the MIME type.'
   ] = None
 
+  metadata: Annotated[
+      dict[str, Any], 'Additional metadata attached to this object.'
+  ] = {}
+
   @functools.cached_property
   def mime_type(self) -> str:
     """Returns the MIME type."""
@@ -87,6 +118,17 @@ class Mime(lf.Modality):
     """Returns True if the MIME type is a binary type."""
     return not self.is_text
 
+  @property
+  def hash(self) -> str:
+    """Returns the hash of the MIME content."""
+    # Hash the URI to avoid downloading the content.
+    if self.uri is not None:
+      return hashlib.md5(self.uri.encode()).hexdigest()[:8]
+    if self.content is not None:
+      return super().hash
+    assert self.metadata
+    return hashlib.md5(str(self.metadata).encode()).hexdigest()[:8]
+
   def to_text(self) -> str:
     """Returns the text content of the MIME type."""
     if not self.is_text:
@@ -132,7 +174,7 @@ class Mime(lf.Modality):
 
   def _on_bound(self):
     super()._on_bound()
-    if self.uri is None and self.content is None:
+    if self.uri is None and self.content is None and not self.metadata:
       raise ValueError('Either uri or content must be provided.')
 
   def to_bytes(self) -> bytes:
@@ -162,6 +204,8 @@ class Mime(lf.Modality):
       return cls.class_from_mime_type(mime_type).from_bytes(content, **kwargs)
 
     if cls is Mime:
+      if 'youtube.com/watch' in uri:
+        return Custom(mime='text/html', uri=uri, **kwargs)
       content = cls.download(uri)
       mime = _detect_mime_type(content)
       return cls.class_from_mime_type(mime)(uri=uri, content=content, **kwargs)
@@ -272,7 +316,24 @@ class Mime(lf.Modality):
 
 @pg.use_init_args(['mime', 'content', 'uri'])
 class Custom(Mime):
-  """Custom MIME data."""
+  """Represents content of a custom MIME type.
+
+  `lf.modalities.Custom` is useful for representing data with MIME types
+  that do not have dedicated classes like `lf.Image` or `lf.Audio`.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  # Create a custom MIME object for plain text
+  text_data = lf.Custom.from_bytes(
+      b'This is a text document.', mime='text/plain'
+  )
+  print(text_data.mime_type)
+  # Output: text/plain
+  ```
+  """
 
   mime: Annotated[
       str, 'The MIME type of the data. E.g. text/plain, or image/png. '

langfun/core/modalities/mime_test.py

@@ -109,6 +109,17 @@ class CustomMimeTest(unittest.TestCase):
     with self.assertRaisesRegex(ValueError, 'Unsupported encoding'):
       mime.Mime.from_uri('data:text/plain;base16,abcd')
 
+    # Test YouTube URI
+    yt_uri = 'https://www.youtube.com/watch?v=dQw4w9WgXcQ'
+    with mock.patch(
+        'langfun.core.modalities.mime.Mime.download'
+    ) as mock_download:
+      content = mime.Mime.from_uri(yt_uri)
+      self.assertIsInstance(content, mime.Custom)
+      self.assertEqual(content.mime_type, 'text/html')
+      self.assertEqual(content.uri, yt_uri)
+      mock_download.assert_not_called()
+
   def assert_html_content(self, html, expected):
     expected = inspect.cleandoc(expected).strip()
     actual = html.content.strip()

langfun/core/modalities/pdf.py

@@ -17,6 +17,24 @@ from langfun.core.modalities import mime
 
 
 class PDF(mime.Mime):
-  """PDF document."""
+  """Represents a PDF document for communicating with language models.
+
+  `lf.PDF` can be initialized from a URI (HTTP/HTTPS URL or local path)
+  using `lf.PDF.from_uri()` or from raw bytes using `lf.PDF.from_bytes()`.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  # Load PDF from path
+  pdf = lf.PDF.from_path('/path/to/document.pdf')
+
+  # Use PDF in a prompt
+  prompt = lf.Template('Summarize this document: {{pdf}}', pdf=pdf)
+  response = lf.query(prompt, lm=lf.llms.Gemini25Flash())
+  print(response)
+  ```
+  """
 
   MIME_PREFIX = 'application/pdf'