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

langfun/core/logging.py

@@ -13,16 +13,13 @@
 # limitations under the License.
 """Langfun event logging."""
 
-from collections.abc import Iterator
 import contextlib
 import datetime
-import io
 import typing
-from typing import Any, Literal
+from typing import Any, Iterator, Literal, Sequence
 
 from langfun.core import component
 from langfun.core import console
-from langfun.core import repr_utils
 import pyglove as pg
 
 
@@ -56,49 +53,153 @@ class LogEntry(pg.Object):
   def should_output(self, min_log_level: LogLevel) -> bool:
     return _LOG_LEVELS.index(self.level) >= _LOG_LEVELS.index(min_log_level)
 
-  def _repr_html_(self) -> str:
-    s = io.StringIO()
-    padding_left = 50 * self.indent
-    s.write(f'<div style="padding-left: {padding_left}px;">')
-    s.write(self._message_display)
-    if self.metadata:
-      s.write(repr_utils.html_repr(self.metadata))
-    s.write('</div>')
-    return s.getvalue()
-
-  @property
-  def _message_text_bgcolor(self) -> str:
-    match self.level:
-      case 'debug':
-        return '#EEEEEE'
-      case 'info':
-        return '#A3E4D7'
-      case 'warning':
-        return '#F8C471'
-      case 'error':
-        return '#F5C6CB'
-      case 'fatal':
-        return '#F19CBB'
-      case _:
-        raise ValueError(f'Unknown log level: {self.level}')
-
-  @property
-  def _time_display(self) -> str:
-    display_text = self.time.strftime('%H:%M:%S')
-    alt_text = self.time.strftime('%Y-%m-%d %H:%M:%S.%f')
-    return (
-        '<span style="background-color: #BBBBBB; color: white; '
-        'border-radius:5px; padding:0px 5px 0px 5px;" '
-        f'title="{alt_text}">{display_text}</span>'
+  def _html_tree_view_summary(
+      self,
+      view: pg.views.HtmlTreeView,
+      title: str | pg.Html | None = None,
+      max_str_len_for_summary: int = pg.View.PresetArgValue(80),  # pytype: disable=annotation-type-mismatch
+      **kwargs
+      ) -> str:
+    if len(self.message) > max_str_len_for_summary:
+      message = self.message[:max_str_len_for_summary] + '...'
+    else:
+      message = self.message
+
+    s = pg.Html(
+        pg.Html.element(
+            'span',
+            [self.time.strftime('%H:%M:%S')],
+            css_class=['log-time']
+        ),
+        pg.Html.element(
+            'span',
+            [pg.Html.escape(message)],
+            css_class=['log-summary'],
+        ),
+    )
+    return view.summary(
+        self,
+        title=title or s,
+        max_str_len_for_summary=max_str_len_for_summary,
+        **kwargs,
     )
 
-  @property
-  def _message_display(self) -> str:
-    return repr_utils.html_round_text(
-        self._time_display + '&nbsp;' + self.message,
-        background_color=self._message_text_bgcolor,
+  # pytype: disable=annotation-type-mismatch
+  def _html_tree_view_content(
+      self,
+      view: pg.views.HtmlTreeView,
+      root_path: pg.KeyPath,
+      collapse_log_metadata_level: int = pg.View.PresetArgValue(0),
+      max_str_len_for_summary: int = pg.View.PresetArgValue(80),
+      **kwargs
+  ) -> pg.Html:
+    # pytype: enable=annotation-type-mismatch
+    def render_message_text():
+      if len(self.message) < max_str_len_for_summary:
+        return None
+      return pg.Html.element(
+          'span',
+          [pg.Html.escape(self.message)],
+          css_class=['log-text'],
+      )
+
+    def render_metadata():
+      if not self.metadata:
+        return None
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.metadata,
+                  name='metadata',
+                  root_path=root_path + 'metadata',
+                  parent=self,
+                  collapse_level=(
+                      root_path.depth + collapse_log_metadata_level + 1
+                  )
+              )
+          ],
+          css_class=['log-metadata'],
+      )
+
+    return pg.Html.element(
+        'div',
+        [
+            render_message_text(),
+            render_metadata(),
+        ],
+        css_class=['complex_value'],
     )
 
+  def _html_style(self) -> list[str]:
+    return super()._html_style() + [
+        """
+        .log-time {
+          color: #222;
+          font-size: 12px;
+          padding-right: 10px;
+        }
+        .log-summary {
+          font-weight: normal;
+          font-style: italic;
+          padding: 4px;
+        }
+        .log-debug > summary > .summary_title::before {
+          content: '🛠️ '
+        }
+        .log-info > summary > .summary_title::before {
+          content: '💡 '
+        }
+        .log-warning > summary > .summary_title::before {
+          content: '❗ '
+        }
+        .log-error > summary > .summary_title::before {
+          content: '❌ '
+        }
+        .log-fatal > summary > .summary_title::before {
+          content: '💀 '
+        }
+        .log-text {
+          display: block;
+          color: black;
+          font-style: italic;
+          padding: 20px;
+          border-radius: 5px;
+          background: rgba(255, 255, 255, 0.5);
+          white-space: pre-wrap;
+        }
+        details.log-entry {
+          margin: 0px 0px 10px;
+          border: 0px;
+        }
+        div.log-metadata {
+          margin: 10px 0px 0px 0px; 
+        }
+        .log-metadata > details {
+          background-color: rgba(255, 255, 255, 0.5);
+          border: 1px solid transparent;
+        }
+        .log-debug {
+          background-color: #EEEEEE
+        }
+        .log-warning {
+          background-color: #F8C471
+        }
+        .log-info {
+          background-color: #A3E4D7
+        }
+        .log-error {
+          background-color: #F5C6CB
+        }
+        .log-fatal {
+          background-color: #F19CBB
+        }
+        """
+    ]
+
+  def _html_element_class(self) -> Sequence[str] | None:
+    return super()._html_element_class() + [f'log-{self.level}']
+
 
 def log(level: LogLevel,
         message: str,

langfun/core/logging_test.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 """Tests for langfun.core.logging."""
 
+import datetime
+import inspect
 import unittest
 
 from langfun.core import logging
@@ -52,6 +54,37 @@ class LoggingTest(unittest.TestCase):
     assert_color(logging.error('hi', indent=2, x=1, y=2), '#F5C6CB')
     assert_color(logging.fatal('hi', indent=2, x=1, y=2), '#F19CBB')
 
+  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):
+    time = datetime.datetime(2024, 10, 10, 12, 30, 45)
+    self.assert_html_content(
+        logging.LogEntry(
+            level='info', message='5 + 2 > 3',
+            time=time, metadata={}
+        ).to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove log-entry log-info"><summary><div class="summary_title"><span class="log-time">12:30:45</span><span class="log-summary">5 + 2 &gt; 3</span></div></summary><div class="complex_value"></div></details>
+        """
+    )
+    self.assert_html_content(
+        logging.LogEntry(
+            level='error', message='This is a longer message: 5 + 2 > 3',
+            time=time, metadata=dict(x=1, y=2)
+        ).to_html(
+            max_str_len_for_summary=10,
+            enable_summary_tooltip=False,
+            collapse_log_metadata_level=1
+        ),
+        """
+        <details open class="pyglove log-entry log-error"><summary><div class="summary_title"><span class="log-time">12:30:45</span><span class="log-summary">This is a ...</span></div></summary><div class="complex_value"><span class="log-text">This is a longer message: 5 + 2 &gt; 3</span><div class="log-metadata"><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">x</span><span class="tooltip key-path">metadata.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.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr></table></div></details></div></div></details>
+        """
+    )
 
 if __name__ == '__main__':
   unittest.main()

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,244 @@ 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 = pg.View.PresetArgValue(1),
+      collapse_message_metadata_level: int = pg.View.PresetArgValue(0),
+      collapse_source_message_level: int = 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.
+      **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=(
+                      child_path.depth + collapse_message_result_level
+                  )
+              )
+          ],
+          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=child_path.depth + (
+                      0 if collapse_llm_usage else 1
+                  )
+              )
+          ],
+          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:
+        return view.render(
+            self.source,
+            name='source',
+            root_path=root_path + 'source',
+            include_metadata=include_message_metadata,
+            collapse_level=(
+                root_path.depth + 1 + collapse_source_message_level
+            ),
+            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=(
+                      child_path.depth + collapse_message_metadata_level
+                  )
+              )
+          ],
+          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.