langfun 0.1.2.dev202510230805__py3-none-any.whl → 0.1.2.dev202510250803__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>
         """
     )
 

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/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.
@@ -87,6 +88,14 @@ 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]
+    return super().hash
+
   def to_text(self) -> str:
     """Returns the text content of the MIME type."""
     if not self.is_text:

langfun/core/modality.py

@@ -14,23 +14,15 @@
 """Interface for modality (e.g. Image, Video, etc.)."""
 
 import abc
+import contextlib
 import functools
 import hashlib
-from typing import Any, ContextManager
+import re
+from typing import Any, ContextManager, Iterator
 from langfun.core import component
 import pyglove as pg
 
 
-_TLS_MODALITY_AS_REF = '__format_modality_as_ref__'
-
-
-def format_modality_as_ref(enabled: bool = True) -> ContextManager[None]:
-  """A context manager that formats modality objects as references."""
-  return pg.object_utils.thread_local_value_scope(
-      _TLS_MODALITY_AS_REF, enabled, False
-  )
-
-
 class Modality(component.Component, pg.views.HtmlTreeView.Extension):
   """Base class for multimodal object."""
 
@@ -39,15 +31,18 @@ class Modality(component.Component, pg.views.HtmlTreeView.Extension):
 
   def _on_bound(self):
     super()._on_bound()
-    # Invalidate cached hash if modality member is changed.
+    # Invalidate cached hash and id if modality member is changed.
     self.__dict__.pop('hash', None)
+    self.__dict__.pop('id', None)
 
   def format(self, *args, **kwargs) -> str:
-    if self.referred_name is None or not pg.object_utils.thread_local_get(
-        _TLS_MODALITY_AS_REF, False
-    ):
+    if not pg.object_utils.thread_local_get(_TLS_MODALITY_AS_REF, False):
       return super().format(*args, **kwargs)
-    return Modality.text_marker(self.referred_name)
+
+    capture_scope = get_modality_capture_context()
+    if capture_scope is not None:
+      capture_scope.capture(self)
+    return Modality.text_marker(self.id)
 
   def __str_kwargs__(self) -> dict[str, Any]:
     # For modality objects, we don't want to use markdown format when they
@@ -70,14 +65,11 @@ class Modality(component.Component, pg.views.HtmlTreeView.Extension):
     """Returns a marker in the text for this object."""
     return Modality.REF_START + var_name + Modality.REF_END
 
-  @property
-  def referred_name(self) -> str | None:
+  @functools.cached_property
+  def id(self) -> str | None:
     """Returns the referred name of this object in its template."""
-    if not self.sym_path:
-      return None
-    # Strip the metadata prefix under message.
-    path = str(self.sym_path)
-    return path[9:] if path.startswith('metadata.') else path
+    modality_type = _camel_to_snake(self.__class__.__name__)
+    return f'{modality_type}:{self.hash}'
 
   @classmethod
   def from_value(cls, value: pg.Symbolic) -> dict[str, 'Modality']:
@@ -86,7 +78,7 @@ class Modality(component.Component, pg.views.HtmlTreeView.Extension):
     def _visit(k, v, p):
       del k, p
       if isinstance(v, Modality):
-        modalities[v.referred_name] = v
+        modalities[v.id] = v
         return pg.TraverseAction.CONTINUE
       return pg.TraverseAction.ENTER
 
@@ -102,7 +94,7 @@ class ModalityRef(pg.Object, pg.typing.CustomTyping):
   structure.
   """
 
-  name: str
+  id: str
 
   def custom_apply(
       self, path: pg.KeyPath, value_spec: pg.ValueSpec, *args, **kwargs
@@ -122,12 +114,97 @@ class ModalityRef(pg.Object, pg.typing.CustomTyping):
     """
 
     def _placehold(k, v, p):
-      del p
+      del k, p
       if isinstance(v, Modality):
-        return ModalityRef(name=value.sym_path + k)
+        return ModalityRef(id=v.id)
       return v
     return value.clone().rebind(_placehold, raise_on_no_change=False)
 
+  @classmethod
+  def restore(cls, value: pg.Symbolic, modalities: dict[str, Modality]) -> Any:
+    """Returns a copy of value by replacing refs with modality objects."""
+    def _restore(k, v, p):
+      del k, p
+      if isinstance(v, ModalityRef):
+        modality_object = modalities.get(v.id)
+        if modality_object is None:
+          raise ValueError(
+              f'Modality {v.id} not found in modalities {modalities.keys()}'
+          )
+        return modality_object
+      return v
+    return value.rebind(_restore, raise_on_no_change=False)
+
 
 class ModalityError(RuntimeError):  # pylint: disable=g-bad-exception-name
   """Exception raised when modality is not supported."""
+
+
+#
+# Context managers to deal with modality objects.
+#
+
+
+_TLS_MODALITY_CAPTURE_SCOPE = '__modality_capture_scope__'
+_TLS_MODALITY_AS_REF = '__format_modality_as_ref__'
+
+
+def format_modality_as_ref(enabled: bool = True) -> ContextManager[None]:
+  """A context manager that formats modality objects as references."""
+  return pg.object_utils.thread_local_value_scope(
+      _TLS_MODALITY_AS_REF, enabled, False
+  )
+
+
+class _ModalityCaptureContext:
+  """A context to capture modality objects when being rendered."""
+
+  def __init__(self):
+    self._references: dict[str, pg.Ref[Modality]] = {}
+
+  def capture(self, modality: Modality) -> None:
+    """Captures the modality object."""
+    self._references[modality.id] = pg.Ref(modality)
+
+  @property
+  def references(self) -> dict[str, pg.Ref[Modality]]:
+    """Returns the modality references captured in this context."""
+    return self._references
+
+
+@contextlib.contextmanager
+def capture_rendered_modalities() -> Iterator[dict[str, pg.Ref[Modality]]]:
+  """Capture modality objects whose references is being rendered.
+
+  Example:
+    ```
+    image = lf.Image.from_url(...)
+    with lf.modality.capture_rendered_modalities() as rendered_modalities:
+      with lf.modality.format_modality_as_ref():
+        print(f'Hello {image}')
+    self.assertEqual(rendered_modalities, {'image:<hash>': pg.Ref(image)})
+    ```
+  """
+  context = get_modality_capture_context()
+  top_level = context is None
+  if top_level:
+    context = _ModalityCaptureContext()
+    pg.object_utils.thread_local_set(_TLS_MODALITY_CAPTURE_SCOPE, context)
+
+  try:
+    yield context.references  # pylint: disable=attribute-error
+  finally:
+    if top_level:
+      pg.object_utils.thread_local_del(_TLS_MODALITY_CAPTURE_SCOPE)
+
+
+def get_modality_capture_context() -> _ModalityCaptureContext | None:
+  """Returns the current modality capture context."""
+  return pg.object_utils.thread_local_get(_TLS_MODALITY_CAPTURE_SCOPE, None)
+
+
+def _camel_to_snake(name: str) -> str:
+  """Converts a camelCase name to snake_case."""
+  return re.sub(
+      pattern=r'([A-Z]+)', repl=r'_\1', string=name
+  ).lower().lstrip('_')