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

langfun/core/message.py

@@ -14,13 +14,11 @@
 """Messages that are exchanged between users and agents."""
 
 import contextlib
-import html
 import io
-from typing import Annotated, Any, Optional, Union
+from typing import Annotated, Any, Optional, Sequence, Union
 
 from langfun.core import modality
 from langfun.core import natural_language
-from langfun.core import repr_utils
 import pyglove as pg
 
 
@@ -406,6 +404,11 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
       with pg.notify_on_change(False):
         self.tags.append(tag)
 
+  def has_tag(self, tag: str | tuple[str, ...]) -> bool:
+    if isinstance(tag, str):
+      return tag in self.tags
+    return any(t in self.tags for t in tag)
+
   #
   # Message source chain.
   #
@@ -503,79 +506,255 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
     v = self.metadata[key]
     return v.value if isinstance(v, pg.Ref) else v
 
-  def _repr_html_(self):
-    return self.to_html().content
-
-  def to_html(
+  # pytype: disable=annotation-type-mismatch
+  def _html_tree_view_content(
       self,
-      include_message_type: bool = True
-  ) -> repr_utils.Html:
-    """Returns the HTML representation of the message."""
-    s = io.StringIO()
-    s.write('<div style="padding:0px 10px 0px 10px;">')
-    # Title bar.
-    if include_message_type:
-      s.write(
-          repr_utils.html_round_text(
-              self.__class__.__name__,
-              text_color='white',
-              background_color=self._text_color(),
+      *,
+      view: pg.views.HtmlTreeView,
+      root_path: pg.KeyPath,
+      source_tag: str | Sequence[str] | None = pg.View.PresetArgValue(
+          ('lm-input', 'lm-output')
+      ),
+      include_message_metadata: bool = pg.View.PresetArgValue(True),
+      collapse_modalities_in_text: bool = pg.View.PresetArgValue(True),
+      collapse_llm_usage: bool = pg.View.PresetArgValue(False),
+      collapse_message_result_level: int | None = pg.View.PresetArgValue(1),
+      collapse_message_metadata_level: int | None = pg.View.PresetArgValue(0),
+      collapse_source_message_level: int | None = pg.View.PresetArgValue(1),
+      collapse_level: int | None = pg.View.PresetArgValue(1),
+      **kwargs,
+  ) -> pg.Html:
+  # pytype: enable=annotation-type-mismatch
+    """Returns the HTML representation of the message.
+
+    Args:
+      view: The HTML tree view.
+      root_path: The root path of the message.
+      source_tag: tags to filter source messages. If None, the entire
+        source chain will be included.
+      include_message_metadata: Whether to include the metadata of the message.
+      collapse_modalities_in_text: Whether to collapse the modalities in the
+        message text.
+      collapse_llm_usage: Whether to collapse the usage in the message.
+      collapse_message_result_level: The level to collapse the result in the
+        message.
+      collapse_message_metadata_level: The level to collapse the metadata in the
+        message.
+      collapse_source_message_level: The level to collapse the source in the
+        message.
+      collapse_level: The global collapse level.
+      **kwargs: Other keyword arguments.
+
+    Returns:
+      The HTML representation of the message content.
+    """
+    def render_tags():
+      return pg.Html.element(
+          'div',
+          [pg.Html.element('span', [tag]) for tag in self.tags],
+          css_class=['message-tags'],
+      )
+
+    def render_message_text():
+      maybe_reformatted = self.get('formatted_text')
+      referred_chunks = {}
+      s = pg.Html('<div class="message-text">')
+      for chunk in self.chunk(maybe_reformatted):
+        if isinstance(chunk, str):
+          s.write(s.escape(chunk))
+        else:
+          assert isinstance(chunk, modality.Modality), chunk
+          child_path = root_path + 'metadata' + chunk.referred_name
+          s.write(
+              pg.Html.element(
+                  'div',
+                  [
+                      view.render(
+                          chunk,
+                          name=chunk.referred_name,
+                          root_path=child_path,
+                          collapse_level=child_path.depth + (
+                              0 if collapse_modalities_in_text else 1
+                          )
+                      )
+                  ],
+                  css_class=['modality-in-text'],
+              )
           )
+          referred_chunks[chunk.referred_name] = chunk
+      s.write('</div>')
+      return s
+
+    def render_result():
+      if 'result' not in self.metadata:
+        return None
+      child_path = root_path + 'metadata' + 'result'
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.result,
+                  name='result',
+                  root_path=child_path,
+                  collapse_level=view.max_collapse_level(
+                      collapse_level,
+                      collapse_message_result_level,
+                      child_path,
+                  )
+              )
+          ],
+          css_class=['message-result'],
+      )
+
+    def render_usage():
+      if 'usage' not in self.metadata:
+        return None
+      child_path = root_path + 'metadata' + 'usage'
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.usage,
+                  name='llm usage',
+                  root_path=child_path,
+                  collapse_level=view.max_collapse_level(
+                      collapse_level,
+                      0 if collapse_llm_usage else 1,
+                      child_path,
+                  )
+              )
+          ],
+          css_class=['message-usage'],
+      )
+
+    def render_source_message():
+      source = self.source
+      while (source is not None
+             and source_tag is not None
+             and not source.has_tag(source_tag)):
+        source = source.source
+      if source is not None:
+        child_path = root_path + 'source'
+        return view.render(
+            self.source,
+            name='source',
+            root_path=child_path,
+            include_metadata=include_message_metadata,
+            collapse_level=view.max_collapse_level(
+                collapse_level,
+                collapse_source_message_level,
+                child_path
+            ),
+            collapse_source_level=max(0, collapse_source_message_level - 1),
+            collapse_modalities=collapse_modalities_in_text,
+            collapse_usage=collapse_llm_usage,
+            collapse_metadata_level=collapse_message_metadata_level,
+            collapse_result_level=collapse_message_result_level,
+        )
+      return None
+
+    def render_metadata():
+      if not include_message_metadata:
+        return None
+      child_path = root_path + 'metadata'
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.metadata,
+                  css_class=['message-metadata'],
+                  name='metadata',
+                  root_path=child_path,
+                  collapse_level=view.max_collapse_level(
+                      collapse_level,
+                      collapse_message_metadata_level,
+                      child_path,
+                  )
+              )
+          ],
+          css_class=['message-metadata'],
       )
-      s.write('<hr>')
 
-    # Body.
-    s.write(
-        f'<span style="color: {self._text_color()}; white-space: pre-wrap;">'
+    return pg.Html.element(
+        'div',
+        [
+            render_tags(),
+            render_message_text(),
+            render_result(),
+            render_usage(),
+            render_metadata(),
+            render_source_message(),
+        ],
+        css_class=['complex_value'],
     )
 
-    # NOTE(daiyip): LLM may reformat the text from the input, therefore
-    # we proritize the formatted text if it's available.
-    maybe_reformatted = self.get('formatted_text')
-    referred_chunks = {}
-    for chunk in self.chunk(maybe_reformatted):
-      if isinstance(chunk, str):
-        s.write(html.escape(chunk))
-      else:
-        assert isinstance(chunk, modality.Modality), chunk
-        s.write('&nbsp;')
-        s.write(repr_utils.html_round_text(
-            chunk.referred_name,
-            text_color='black',
-            background_color='#f7dc6f'
-        ))
-        s.write('&nbsp;')
-        referred_chunks[chunk.referred_name] = chunk
-    s.write('</span>')
-
-    def item_color(k, v):
-      if isinstance(v, modality.Modality):
-        return ('black', '#f7dc6f', None, None)   # Light yellow
-      elif k == 'result':
-        return ('white', 'purple', 'purple', None)      # Blue.
-      elif k in ('usage',):
-        return ('white', '#e74c3c', None, None)   # Red.
-      else:
-        return ('white', '#17202a', None, None)   # Dark gray
-
-    # TODO(daiyip): Revisit the logic in deciding what metadata keys to
-    # expose to the user.
-    if referred_chunks:
-      s.write(repr_utils.html_repr(referred_chunks, item_color))
-
-    if 'lm-response' in self.tags:
-      s.write(repr_utils.html_repr(self.metadata, item_color))
-    s.write('</div>')
-    return repr_utils.Html(s.getvalue())
-
-  def _text_color(self) -> str:
-    match self.__class__.__name__:
-      case 'UserMessage':
-        return 'green'
-      case 'AIMessage':
-        return 'blue'
-      case _:
-        return 'black'
+  def _html_style(self) -> list[str]:
+    return super()._html_style() + [
+        """
+        /* Langfun Message styles.*/
+        [class^="message-"] > details {
+            margin: 0px 0px 5px 0px;
+            border: 1px solid #EEE;
+        }
+        details.lf-message > summary > .summary_title::after {
+            content: ' 💬';
+        }
+        details.pyglove.ai-message {
+            border: 1px solid blue;
+            color: blue;
+        }
+        details.pyglove.user-message {
+            border: 1px solid green;
+            color: green;
+        }
+        .message-tags {
+            margin: 5px 0px 5px 0px;
+            font-size: .8em;
+        }
+        .message-tags > span {
+            border-radius: 5px;
+            background-color: #CCC;
+            padding: 3px;
+            margin: 0px 2px 0px 2px;
+            color: white;
+        }
+        .message-text {
+            padding: 20px;
+            margin: 10px 5px 10px 5px;
+            font-style: italic;
+            font-size: 1.1em;
+            white-space: pre-wrap;
+            border: 1px solid #EEE;
+            border-radius: 5px;
+            background-color: #EEE;
+        }
+        .modality-in-text {
+            display: inline-block;
+        }
+        .modality-in-text > details {
+            display: inline-block;
+            font-size: 0.8em;
+            border: 0;
+            background-color: #A6F1A6;
+            margin: 0px 5px 0px 5px;
+        }
+        .message-result {
+            color: purple;
+        }
+        .message-usage {
+            color: orange;
+        }
+        .message-usage .object_key.str {
+            border: 1px solid orange;
+            background-color: orange;
+            color: white;
+        }
+        """
+    ]
+
+  def _html_element_class(self) -> list[str]:
+    return super()._html_element_class() + ['lf-message']
+
 
 #
 # Messages of different roles.

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><tr><td><span class="object_key str">estimated_cost</span><span class="tooltip key-path">metadata.usage.estimated_cost</span></td><td><div><span class="simple_value none-type">None</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><tr><td><span class="object_key str">estimated_cost</span><span class="tooltip key-path">metadata.usage.estimated_cost</span></td><td><div><span class="simple_value none-type">None</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><tr><td><span class="object_key str">estimated_cost</span><span class="tooltip key-path">metadata.usage.estimated_cost</span></td><td><div><span class="simple_value none-type">None</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><tr><td><span class="object_key str">estimated_cost</span><span class="tooltip key-path">metadata.usage.estimated_cost</span></td><td><div><span class="simple_value none-type">None</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):