langfun 0.1.2.dev202410100804__py3-none-any.whl → 0.1.2.dev202410110804__py3-none-any.whl

langfun/core/message_test.py

@@ -15,6 +15,7 @@
 
 import inspect
 import unittest
+from langfun.core import language_model
 from langfun.core import message
 from langfun.core import modality
 import pyglove as pg
@@ -26,9 +27,6 @@ class CustomModality(modality.Modality):
   def to_bytes(self):
     return self.content.encode()
 
-  def _repr_html_(self):
-    return f'<div>CustomModality: {self.content}</div>'
-
 
 class MessageTest(unittest.TestCase):
 
@@ -39,11 +37,19 @@ class MessageTest(unittest.TestCase):
 
     d = pg.Dict(x=A())
 
-    m = message.UserMessage('hi', metadata=dict(x=1), x=pg.Ref(d.x), y=2)
+    m = message.UserMessage(
+        'hi',
+        metadata=dict(x=1), x=pg.Ref(d.x),
+        y=2,
+        tags=['lm-input']
+    )
     self.assertEqual(m.metadata, {'x': pg.Ref(d.x), 'y': 2})
     self.assertEqual(m.sender, 'User')
     self.assertIs(m.x, d.x)
     self.assertEqual(m.y, 2)
+    self.assertTrue(m.has_tag('lm-input'))
+    self.assertTrue(m.has_tag(('lm-input', '')))
+    self.assertFalse(m.has_tag('lm-response'))
 
     with self.assertRaises(AttributeError):
       _ = m.z
@@ -332,50 +338,69 @@ class MessageTest(unittest.TestCase):
         )
     )
 
-  def test_html(self):
-    m = message.UserMessage(
-        'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>',
-        img1=CustomModality('foo'),
-        x=dict(img2=CustomModality('bar')),
+  def assert_html_content(self, html, expected):
+    expected = inspect.cleandoc(expected).strip()
+    actual = html.content.strip()
+    if actual != expected:
+      print(actual)
+    self.assertEqual(actual, expected)
+
+  def test_html_user_message(self):
+    self.assert_html_content(
+        message.UserMessage(
+            'what is a <div>'
+        ).to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove user-message lf-message"><summary><div class="summary_title">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 class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><span class="empty_container"></span></div></details></div></div></details>
+        """
     )
-    self.assertEqual(
-        m._repr_html_(),
-        (
-            '<div style="padding:0px 10px 0px 10px;"><span style="color: white;'
-            'background-color: green;display:inline-block; border-radius:10px; '
-            'padding:5px; margin-top: 5px; margin-bottom: 5px; white-space: '
-            'pre-wrap">UserMessage</span><hr><span style="color: green; '
-            'white-space: pre-wrap;">hi, this is a&nbsp;<span style="color: '
-            'black;background-color: #f7dc6f;display:inline-block; '
-            'border-radius:10px; padding:5px; margin-top: 5px; margin-bottom: '
-            '5px; white-space: pre-wrap">img1</span>&nbsp;and&nbsp;<span style'
-            '="color: black;background-color: #f7dc6f;display:inline-block; '
-            'border-radius:10px; padding:5px; margin-top: 5px; margin-bottom: '
-            '5px; white-space: pre-wrap">x.img2</span>&nbsp;</span><div style='
-            '"padding-left: 20px; margin-top: 10px"><table style="border-top: '
-            '1px solid #EEEEEE;"><tr><td style="padding: 5px; vertical-align: '
-            'top; border-bottom: 1px solid #EEEEEE"><span style="color: black;'
-            'background-color: #f7dc6f;display:inline-block; border-radius:'
-            '10px; padding:5px; margin-top: 5px; margin-bottom: 0px; '
-            'white-space: pre-wrap">img1</span></td><td style="padding: 15px '
-            '5px 5px 5px; vertical-align: top; border-bottom: 1px solid '
-            '#EEEEEE;"><div>CustomModality: foo</div></td></tr><tr><td style='
-            '"padding: 5px; vertical-align: top; border-bottom: 1px solid '
-            '#EEEEEE"><span style="color: black;background-color: #f7dc6f;'
-            'display:inline-block; border-radius:10px; padding:5px; margin-top:'
-            ' 5px; margin-bottom: 0px; white-space: pre-wrap">x.img2</span>'
-            '</td><td style="padding: 15px 5px 5px 5px; vertical-align: top; '
-            'border-bottom: 1px solid #EEEEEE;"><div>CustomModality: bar</div>'
-            '</td></tr></table></div></div>'
-        )
+    self.assert_html_content(
+        message.UserMessage(
+            'what is this <<[[image]]>>',
+            tags=['lm-input'],
+            image=CustomModality('bird')
+        ).to_html(enable_summary_tooltip=False, include_message_metadata=False),
+        """
+        <details open class="pyglove user-message lf-message"><summary><div class="summary_title">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</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 key-path">metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;bird&#x27;</span></div></td></tr></table></div></details></div></div></div></details>
+        """
     )
-    self.assertIn(
-        'background-color: blue',
-        message.AIMessage('hi').to_html().content,
+
+  def test_html_ai_message(self):
+    image = CustomModality('foo')
+    user_message = message.UserMessage(
+        'What is in this image? <<[[image]]>> this is a test',
+        metadata=dict(image=image),
+        source=message.UserMessage('User input'),
+        tags=['lm-input']
     )
-    self.assertIn(
-        'background-color: black',
-        message.SystemMessage('hi').to_html().content,
+    ai_message = message.AIMessage(
+        'My name is Gemini',
+        metadata=dict(
+            result=pg.Dict(x=1, y=2, z=pg.Dict(a=[12, 323])),
+            usage=language_model.LMSamplingUsage(10, 2, 12)
+        ),
+        tags=['lm-response', 'lm-output'],
+        source=user_message,
+    )
+    self.assert_html_content(
+        ai_message.to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove ai-message lf-message"><summary><div class="summary_title">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</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 key-path">metadata.result.x</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">y</span><span class="tooltip key-path">metadata.result.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">z</span><span class="tooltip key-path">metadata.result.z</span></td><td><div><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 key-path">metadata.result.z.a</span></td><td><div><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 key-path">metadata.result.z.a[0]</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key int">1</span><span class="tooltip key-path">metadata.result.z.a[1]</span></td><td><div><span class="simple_value int">323</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr></table></div></details></div><div class="message-usage"><details open class="pyglove lm-sampling-usage"><summary><div class="summary_name">llm usage</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 key-path">metadata.usage.prompt_tokens</span></td><td><div><span class="simple_value int">10</span></div></td></tr><tr><td><span class="object_key str">completion_tokens</span><span class="tooltip key-path">metadata.usage.completion_tokens</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">total_tokens</span><span class="tooltip key-path">metadata.usage.total_tokens</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key str">num_requests</span><span class="tooltip key-path">metadata.usage.num_requests</span></td><td><div><span class="simple_value int">1</span></div></td></tr></table></div></details></div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">result</span><span class="tooltip key-path">metadata.result</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">x</span><span class="tooltip key-path">metadata.result.x</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">y</span><span class="tooltip key-path">metadata.result.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">z</span><span class="tooltip key-path">metadata.result.z</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">a</span><span class="tooltip key-path">metadata.result.z.a</span></td><td><div><details class="pyglove list"><summary><div class="summary_title">List(...)</div></summary><div class="complex_value list message-metadata"><table><tr><td><span class="object_key int">0</span><span class="tooltip key-path">metadata.result.z.a[0]</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key int">1</span><span class="tooltip key-path">metadata.result.z.a[1]</span></td><td><div><span class="simple_value int">323</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr><tr><td><span class="object_key str">usage</span><span class="tooltip key-path">metadata.usage</span></td><td><div><details class="pyglove lm-sampling-usage"><summary><div class="summary_title">LMSamplingUsage(...)</div></summary><div class="complex_value lm-sampling-usage message-metadata"><table><tr><td><span class="object_key str">prompt_tokens</span><span class="tooltip key-path">metadata.usage.prompt_tokens</span></td><td><div><span class="simple_value int">10</span></div></td></tr><tr><td><span class="object_key str">completion_tokens</span><span class="tooltip key-path">metadata.usage.completion_tokens</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">total_tokens</span><span class="tooltip key-path">metadata.usage.total_tokens</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key str">num_requests</span><span class="tooltip key-path">metadata.usage.num_requests</span></td><td><div><span class="simple_value int">1</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div><details open class="pyglove user-message lf-message"><summary><div class="summary_name">source</div><div class="summary_title">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</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 key-path">source.metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;foo&#x27;</span></div></td></tr></table></div></details></div>this is a test</div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">image</span><span class="tooltip key-path">source.metadata.image</span></td><td><div><details class="pyglove custom-modality"><summary><div class="summary_title">CustomModality(...)</div></summary><div class="complex_value custom-modality message-metadata"><table><tr><td><span class="object_key str">content</span><span class="tooltip key-path">source.metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;foo&#x27;</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></div></details></div></details>
+        """
+    )
+    self.assert_html_content(
+        ai_message.to_html(
+            enable_summary_tooltip=False,
+            collapse_modalities_in_text=False,
+            collapse_llm_usage=True,
+            collapse_message_result_level=0,
+            collapse_message_metadata_level=0,
+            collapse_source_message_level=0,
+            source_tag=None,
+        ),
+        """
+        <details open class="pyglove ai-message lf-message"><summary><div class="summary_title">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</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 key-path">metadata.result.x</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">y</span><span class="tooltip key-path">metadata.result.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">z</span><span class="tooltip key-path">metadata.result.z</span></td><td><div><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 key-path">metadata.result.z.a</span></td><td><div><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 key-path">metadata.result.z.a[0]</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key int">1</span><span class="tooltip key-path">metadata.result.z.a[1]</span></td><td><div><span class="simple_value int">323</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr></table></div></details></div><div class="message-usage"><details class="pyglove lm-sampling-usage"><summary><div class="summary_name">llm usage</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 key-path">metadata.usage.prompt_tokens</span></td><td><div><span class="simple_value int">10</span></div></td></tr><tr><td><span class="object_key str">completion_tokens</span><span class="tooltip key-path">metadata.usage.completion_tokens</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">total_tokens</span><span class="tooltip key-path">metadata.usage.total_tokens</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key str">num_requests</span><span class="tooltip key-path">metadata.usage.num_requests</span></td><td><div><span class="simple_value int">1</span></div></td></tr></table></div></details></div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">result</span><span class="tooltip key-path">metadata.result</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">x</span><span class="tooltip key-path">metadata.result.x</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">y</span><span class="tooltip key-path">metadata.result.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">z</span><span class="tooltip key-path">metadata.result.z</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">a</span><span class="tooltip key-path">metadata.result.z.a</span></td><td><div><details class="pyglove list"><summary><div class="summary_title">List(...)</div></summary><div class="complex_value list message-metadata"><table><tr><td><span class="object_key int">0</span><span class="tooltip key-path">metadata.result.z.a[0]</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key int">1</span><span class="tooltip key-path">metadata.result.z.a[1]</span></td><td><div><span class="simple_value int">323</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr><tr><td><span class="object_key str">usage</span><span class="tooltip key-path">metadata.usage</span></td><td><div><details class="pyglove lm-sampling-usage"><summary><div class="summary_title">LMSamplingUsage(...)</div></summary><div class="complex_value lm-sampling-usage message-metadata"><table><tr><td><span class="object_key str">prompt_tokens</span><span class="tooltip key-path">metadata.usage.prompt_tokens</span></td><td><div><span class="simple_value int">10</span></div></td></tr><tr><td><span class="object_key str">completion_tokens</span><span class="tooltip key-path">metadata.usage.completion_tokens</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">total_tokens</span><span class="tooltip key-path">metadata.usage.total_tokens</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key str">num_requests</span><span class="tooltip key-path">metadata.usage.num_requests</span></td><td><div><span class="simple_value int">1</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div><details class="pyglove user-message lf-message"><summary><div class="summary_name">source</div><div class="summary_title">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</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 key-path">source.metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;foo&#x27;</span></div></td></tr></table></div></details></div>this is a test</div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">image</span><span class="tooltip key-path">source.metadata.image</span></td><td><div><details class="pyglove custom-modality"><summary><div class="summary_title">CustomModality(...)</div></summary><div class="complex_value custom-modality message-metadata"><table><tr><td><span class="object_key str">content</span><span class="tooltip key-path">source.metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;foo&#x27;</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div><details class="pyglove user-message lf-message"><summary><div class="summary_name">source</div><div class="summary_title">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"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><span class="empty_container"></span></div></details></div></div></details></div></details></div></details>
+        """
     )
 
 

langfun/core/modalities/audio.py

@@ -26,5 +26,5 @@ class Audio(mime.Mime):
   def audio_format(self) -> str:
     return self.mime_type.removeprefix(self.MIME_PREFIX + '/')
 
-  def _html(self, uri: str) -> str:
+  def _mime_control_for(self, uri: str) -> str:
     return f'<audio controls> <source src="{uri}"> </audio>'

langfun/core/modalities/audio_test.py

@@ -53,7 +53,7 @@ class AudioFileTest(unittest.TestCase):
       self.assertEqual(audio.audio_format, 'x-wav')
       self.assertEqual(audio.mime_type, 'audio/x-wav')
       self.assertEqual(
-          audio._repr_html_(),
+          audio._raw_html(),
           '<audio controls> <source src="http://mock/web/a.wav"> </audio>',
       )
       self.assertEqual(audio.to_bytes(), content_bytes)

langfun/core/modalities/image.py

@@ -41,7 +41,7 @@ class Image(mime.Mime):
   def image_format(self) -> str:
     return self.mime_type.removeprefix(self.MIME_PREFIX + '/')
 
-  def _html(self, uri: str) -> str:
+  def _mime_control_for(self, uri: str) -> str:
     return f'<img src="{uri}">'
 
   @functools.cached_property

langfun/core/modalities/image_test.py

@@ -45,7 +45,7 @@ class ImageTest(unittest.TestCase):
   def test_from_bytes(self):
     image = image_lib.Image.from_bytes(image_content)
     self.assertEqual(image.image_format, 'png')
-    self.assertIn('data:image/png;base64,', image._repr_html_())
+    self.assertIn('data:image/png;base64,', image._raw_html())
     self.assertEqual(image.to_bytes(), image_content)
     with self.assertRaisesRegex(
         lf.ModalityError, '.* cannot be converted to text'
@@ -67,7 +67,10 @@ class ImageTest(unittest.TestCase):
     with mock.patch('requests.get') as mock_requests_get:
       mock_requests_get.side_effect = mock_request
       self.assertEqual(image.image_format, 'png')
-      self.assertEqual(image._repr_html_(), '<img src="http://mock/web/a.png">')
+      self.assertEqual(
+          image._raw_html(),
+          '<img src="http://mock/web/a.png">'
+      )
       self.assertEqual(image.to_bytes(), image_content)
 
   def test_from_uri_base_cls(self):
@@ -76,7 +79,10 @@ class ImageTest(unittest.TestCase):
       image = mime_lib.Mime.from_uri('http://mock/web/a.png')
       self.assertIsInstance(image, image_lib.Image)
       self.assertEqual(image.image_format, 'png')
-      self.assertEqual(image._repr_html_(), '<img src="http://mock/web/a.png">')
+      self.assertEqual(
+          image._raw_html(),
+          '<img src="http://mock/web/a.png">'
+      )
       self.assertEqual(image.to_bytes(), image_content)
 
   def test_image_size(self):

langfun/core/modalities/mime.py

@@ -178,14 +178,50 @@ class Mime(lf.Modality):
       assert content is not None
       return content
 
-  def _repr_html_(self) -> str:
+  def _html_tree_view_content(
+      self,
+      **kwargs) -> str:
+    return self._raw_html()
+
+  def _html_tree_view_render(
+      self,
+      view: pg.views.HtmlTreeView,
+      raw_mime_content: bool = pg.View.PresetArgValue(False),                 # pytype: disable=annotation-type-mismatch
+      display_modality_when_hover: bool = pg.View.PresetArgValue(False),   # pytype: disable=annotation-type-mismatch
+      **kwargs
+  ):
+    if raw_mime_content:
+      return pg.Html(self._raw_html())
+    else:
+      if display_modality_when_hover:
+        kwargs.update(
+            display_modality_when_hover=True,
+            enable_summary_tooltip=True,
+        )
+      return super()._html_tree_view_render(view=view, **kwargs)
+
+  def _html_tree_view_tooltip(
+      self,
+      *,
+      view: pg.views.HtmlTreeView,
+      content: pg.Html | str | None = None,
+      display_modality_when_hover: bool = pg.View.PresetArgValue(False),    # pytype: disable=annotation-type-mismatch
+      **kwargs
+  ):
+    if content is None and display_modality_when_hover:
+      content = self._raw_html()
+    return super()._html_tree_view_tooltip(
+        view=view, content=content, **kwargs
+    )
+
+  def _raw_html(self) -> str:
     if self.uri and self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
       uri = self.uri
     else:
       uri = self.content_uri
-    return self._html(uri)
+    return self._mime_control_for(uri)
 
-  def _html(self, uri) -> str:
+  def _mime_control_for(self, uri) -> str:
     return f'<embed type="{self.mime_type}" src="{uri}"/>'
 
 

langfun/core/modalities/mime_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """MIME tests."""
+import inspect
 import unittest
 from unittest import mock
 
@@ -77,6 +78,44 @@ class CustomMimeTest(unittest.TestCase):
       self.assertEqual(content.to_bytes(), 'bar')
       self.assertEqual(content.mime_type, 'text/plain')
 
+  def assert_html_content(self, html, expected):
+    expected = inspect.cleandoc(expected).strip()
+    actual = html.content.strip()
+    if actual != expected:
+      print(actual)
+    self.assertEqual(actual, expected)
+
+  def test_html(self):
+    self.assert_html_content(
+        mime.Custom('text/plain', b'foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+        ),
+        """
+        <details open class="pyglove custom"><summary><div class="summary_title">Custom(...)</div></summary><embed type="text/plain" src="data:text/plain;base64,Zm9v"/></details>
+        """
+    )
+    self.assert_html_content(
+        mime.Custom('text/plain', b'foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+            raw_mime_content=True,
+        ),
+        """
+        <embed type="text/plain" src="data:text/plain;base64,Zm9v"/>
+        """
+    )
+    self.assert_html_content(
+        mime.Custom('text/plain', b'foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+            display_modality_when_hover=True,
+        ),
+        """
+        <details open class="pyglove custom"><summary><div class="summary_title">Custom(...)</div><span class="tooltip custom"><embed type="text/plain" src="data:text/plain;base64,Zm9v"/></span></summary><embed type="text/plain" src="data:text/plain;base64,Zm9v"/></details>
+        """
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/modalities/ms_office.py

@@ -29,7 +29,7 @@ class Xlsx(mime.Mime):
       'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
   )
 
-  def to_html(self) -> str:
+  def _raw_html(self) -> str:
     try:
       import pandas as pd  # pylint: disable=g-import-not-at-top
       import openpyxl      # pylint: disable=g-import-not-at-top, unused-import
@@ -40,9 +40,6 @@ class Xlsx(mime.Mime):
           'Please install "langfun[mime-xlsx]" to enable XLSX support.'
       ) from e
 
-  def _repr_html_(self) -> str:
-    return self.to_html()
-
   def _is_compatible(self, mime_types: Iterable[str]) -> bool:
     return bool(set(mime_types).intersection([
         'text/html',
@@ -52,7 +49,7 @@ class Xlsx(mime.Mime):
   def _make_compatible(self, mime_types: Iterable[str]) -> mime.Mime:
     """Returns the MimeType of the converted file."""
     del mime_types
-    return mime.Mime(uri=self.uri, content=self.to_html())
+    return mime.Mime(uri=self.uri, content=self._raw_html())
 
 
 class Docx(mime.Mime):

langfun/core/modalities/ms_office_test.py

@@ -347,7 +347,7 @@ class XlsxTest(unittest.TestCase):
           ),
       )
       self.assertEqual(content.to_bytes(), xlsx_bytes)
-      self.assertEqual(content.to_html(), expected_xlsx_html)
+      self.assertEqual(content._raw_html(), expected_xlsx_html)
 
 
 class PptxTest(unittest.TestCase):

langfun/core/modalities/pdf_test.py

@@ -49,7 +49,7 @@ class PdfTest(unittest.TestCase):
     pdf = pdf_lib.PDF.from_bytes(pdf_bytes)
     self.assertIn(
         '<embed type="application/pdf" src="data:application/pdf;base64,',
-        pdf._repr_html_()
+        pdf._raw_html()
     )
 
 

langfun/core/modalities/video.py

@@ -26,5 +26,5 @@ class Video(mime.Mime):
   def video_format(self) -> str:
     return self.mime_type.removeprefix(self.MIME_PREFIX + '/')
 
-  def _html(self, uri: str) -> str:
+  def _mime_control_for(self, uri: str) -> str:
     return f'<video controls> <source src="{uri}"> </video>'

langfun/core/modalities/video_test.py

@@ -38,7 +38,7 @@ class VideoContentTest(unittest.TestCase):
     video = video_lib.Video.from_bytes(mp4_bytes)
     self.assertEqual(video.mime_type, 'video/mp4')
     self.assertEqual(video.video_format, 'mp4')
-    self.assertIn('data:video/mp4;base64,', video._repr_html_())
+    self.assertIn('data:video/mp4;base64,', video._raw_html())
     self.assertEqual(video.to_bytes(), mp4_bytes)
 
   def test_bad_video(self):
@@ -56,7 +56,7 @@ class VideoFileTest(unittest.TestCase):
       self.assertEqual(video.video_format, 'mp4')
       self.assertEqual(video.mime_type, 'video/mp4')
       self.assertEqual(
-          video._repr_html_(),
+          video._raw_html(),
           '<video controls> <source src="http://mock/web/a.mp4"> </video>',
       )
       self.assertEqual(video.to_bytes(), mp4_bytes)

langfun/core/structured/mapping.py

@@ -183,6 +183,44 @@ class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
       result.write(lf.colored(str(self.metadata), color='cyan'))
     return result.getvalue().strip()
 
+  def _html_tree_view_content(
+      self,
+      *,
+      parent: Any,
+      view: pg.views.HtmlTreeView,
+      root_path: pg.KeyPath,
+      **kwargs,
+  ):
+    def render_value(value, **kwargs):
+      if isinstance(value, lf.Template):
+        # Make a shallow copy to make sure modalities are rooted by
+        # the input.
+        value = value.clone().render()
+      return view.render(value, **kwargs)
+
+    exclude_keys = []
+    if not self.context:
+      exclude_keys.append('context')
+    if not self.schema:
+      exclude_keys.append('schema')
+    if not self.metadata:
+      exclude_keys.append('metadata')
+
+    kwargs.pop('special_keys', None)
+    kwargs.pop('exclude_keys', None)
+    return view.complex_value(
+        self.sym_init_args,
+        parent=self,
+        root_path=root_path,
+        render_value_fn=render_value,
+        special_keys=['input', 'output', 'context', 'schema', 'metadata'],
+        exclude_keys=exclude_keys,
+        **kwargs
+    )
+
+  def _html_tree_view_collapse_level(self) -> int:
+    return 2
+
 
 class Mapping(lf.LangFunc):
   """Base class for mapping.

langfun/core/structured/mapping_test.py

@@ -14,6 +14,7 @@
 """Tests for structured mapping example."""
 
 import inspect
+from typing import Any
 import unittest
 
 import langfun.core as lf
@@ -164,6 +165,60 @@ class MappingExampleTest(unittest.TestCase):
         pg.eq(pg.from_json_str(example.to_json_str()), example)
     )
 
+  def assert_html_content(self, html, expected):
+    expected = inspect.cleandoc(expected).strip()
+    actual = html.content.strip()
+    if actual != expected:
+      print(actual)
+    self.assertEqual(actual, expected)
+
+  def test_html(self):
+
+    class Answer(pg.Object):
+      answer: int
+
+    class Addition(lf.Template):
+      """Template Addition.
+
+      {{x}} + {{y}} = ?
+      """
+      x: Any
+      y: Any
+
+    example = mapping.MappingExample(
+        input=Addition(x=1, y=2),
+        schema=Answer,
+        context='compute 1 + 1',
+        output=Answer(answer=3),
+        metadata={'foo': 'bar'},
+    )
+    self.assert_html_content(
+        example.to_html(
+            enable_summary_tooltip=False, include_message_metadata=False
+        ),
+        """
+        <details open class="pyglove mapping-example"><summary><div class="summary_title">MappingExample(...)</div></summary><div class="complex_value mapping-example"><div class="special_value"><details open class="pyglove user-message lf-message"><summary><div class="summary_name">input</div><div class="summary_title">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>rendered</span></div><div class="message-text">1 + 2 = ?</div></div></details></div><div class="special_value"><details open class="pyglove answer"><summary><div class="summary_name">output</div><div class="summary_title">Answer(...)</div></summary><div class="complex_value answer"><table><tr><td><span class="object_key str">answer</span><span class="tooltip key-path">output.answer</span></td><td><div><span class="simple_value int">3</span></div></td></tr></table></div></details></div><div class="special_value"><details open class="pyglove str"><summary><div class="summary_name">context</div><div class="summary_title">&#x27;compute 1 + 1&#x27;</div></summary><span class="simple_value str">&#x27;compute 1 + 1&#x27;</span></details></div><div class="special_value"><details open class="pyglove schema"><summary><div class="summary_name">schema</div><div class="summary_title">Schema(...)</div></summary><div class="lf-schema-definition">Answer
+
+        ```python
+        class Answer:
+          answer: int
+        ```</div></details></div><div class="special_value"><details open class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict"><table><tr><td><span class="object_key str">foo</span><span class="tooltip key-path">metadata.foo</span></td><td><div><span class="simple_value str">&#x27;bar&#x27;</span></div></td></tr></table></div></details></div></div></details>
+        """
+    )
+
+    example = mapping.MappingExample(
+        input=Addition(x=1, y=2),
+        output=Answer(answer=3),
+    )
+    self.assert_html_content(
+        example.to_html(
+            enable_summary_tooltip=False, include_message_metadata=False
+        ),
+        """
+        <details open class="pyglove mapping-example"><summary><div class="summary_title">MappingExample(...)</div></summary><div class="complex_value mapping-example"><div class="special_value"><details open class="pyglove user-message lf-message"><summary><div class="summary_name">input</div><div class="summary_title">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>rendered</span></div><div class="message-text">1 + 2 = ?</div></div></details></div><div class="special_value"><details open class="pyglove answer"><summary><div class="summary_name">output</div><div class="summary_title">Answer(...)</div></summary><div class="complex_value answer"><table><tr><td><span class="object_key str">answer</span><span class="tooltip key-path">output.answer</span></td><td><div><span class="simple_value int">3</span></div></td></tr></table></div></details></div></div></details>
+        """
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/structured/schema.py

@@ -189,6 +189,40 @@ class Schema(lf.NaturalLanguageFormattable, pg.Object):
       return value
     return cls(parse_value_spec(value))
 
+  def _html_tree_view_content(
+      self,
+      *,
+      view: pg.views.HtmlTreeView,
+      root_path: pg.KeyPath,
+      **kwargs,
+  ):
+    return pg.Html.element(
+        'div',
+        [self.schema_str(protocol='python')],
+        css_class=['lf-schema-definition']
+    ).add_style(
+        """
+        .lf-schema-definition {
+            color: blue;
+            margin: 5px;
+            white-space: pre-wrap;
+        }
+        """
+    )
+
+  def _html_tree_view_tooltip(
+      self,
+      *,
+      view: pg.views.HtmlTreeView,
+      content: pg.Html | str | None = None,
+      **kwargs,
+  ):
+    return view.tooltip(
+        self,
+        content=content or pg.Html.escape(self.schema_str(protocol='python')),
+        **kwargs
+    )
+
 
 def _top_level_object_specs_from_value(value: pg.Symbolic) -> list[Type[Any]]:
   """Returns a list of top level value specs from a symbolic value."""