langfun 0.1.1.dev20240805__tar.gz → 0.1.1.dev20240807__tar.gz

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.1.dev20240805
+Version: 0.1.1.dev20240807
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/langfun/core/__init__.py

@@ -123,6 +123,7 @@ from langfun.core import console
 
 # Helpers for implementing _repr_xxx_ methods.
 from langfun.core import repr_utils
+Html = repr_utils.Html
 
 # Utility for event logging.
 from langfun.core import logging

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/langfun/core/eval/base.py

@@ -18,7 +18,6 @@ import collections
 import dataclasses
 import functools
 import hashlib
-import html
 import inspect
 import io
 import os
@@ -531,53 +530,16 @@ class Evaluable(lf.Component):
   def _render_message(self, message: lf.Message, s: io.StringIO) -> None:
     for m in message.trace():
       if 'lm-input' in m.tags:
-        text_color = 'green'
+        color = 'green'
       elif 'lm-response' in m.tags:
-        text_color = 'blue'
+        color = 'blue'
       else:
-        text_color = 'black'
+        continue
 
       s.write(
-          f'<div style="color: {text_color}; white-space: pre-wrap;'
-          'padding: 10px; border: 1px solid; margin-top: 10px">'
+          f'<div style="color: {color}; border: 1px solid; margin-top: 10px">'
       )
-      s.write(html.escape(m.get('formatted_text', m.text)))
-
-      # Write output.
-      if m.result is not None:
-        s.write(
-            '<div style="color: magenta; white-space: pre-wrap;'
-            'padding: 10px; border: 1px solid; margin: 10px">'
-        )
-        s.write(html.escape(pg.format(m.result)))
-        s.write('</div>')
-
-      # Write modality information.
-      if 'lm-input' in m.tags or 'lm-response' in m.tags:
-        modalities = m.referred_modalities()
-        if modalities:
-          s.write(f'<div style="color: {text_color}; white-space: pre-wrap;'
-                  'padding: 10px; border: 1px solid; margin-top: 10px"><table>')
-          for name, modality in modalities.items():
-            s.write(f'<tr><td>{name}</td><td>')
-            if hasattr(modality, '_repr_html_'):
-              s.write(modality._repr_html_())   # pylint: disable=protected-access
-            else:
-              s.write(html.escape(pg.format(modality, max_bytes_len=32)))
-            s.write('</td></tr>')
-          s.write('</table></div>')
-
-      # Write usage information.
-      if m.metadata.get('usage', None):
-        s.write(
-            '<div style="background-color: #EEEEEE; color: black; '
-            'white-space: pre-wrap; padding: 10px; border: 0px solid; '
-            'margin: 10px">'
-            f'prompt: {m.usage.prompt_tokens} tokens, '
-            f'response: {m.usage.completion_tokens} tokens, '
-            f'total: {m.usage.total_tokens} tokens'
-            '</div>'
-        )
+      s.write(m._repr_html_())    # pylint: disable=protected-access
       s.write('</div>')
 
   @classmethod

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/langfun/core/logging.py

@@ -19,6 +19,7 @@ import typing
 from typing import Any, Literal, ContextManager
 
 from langfun.core import console
+from langfun.core import repr_utils
 import pyglove as pg
 
 
@@ -55,31 +56,19 @@ class LogEntry(pg.Object):
     s.write(f'<div style="padding-left: {padding_left}px;">')
     s.write(self._message_display)
     if self.metadata:
-      s.write('<div style="padding-left: 20px; margin-top: 10px">')
-      s.write('<table style="border-top: 1px solid #EEEEEE;">')
-      for k, v in self.metadata.items():
-        if hasattr(v, '_repr_html_'):
-          cs = v._repr_html_()  # pylint: disable=protected-access
-        else:
-          cs = f'<span style="white-space: pre-wrap">{str(v)}</span>'
-        key_span = self._round_text(k, color='#F1C40F', margin_bottom='0px')
-        s.write(
-            '<tr>'
-            '<td style="padding: 5px; vertical-align: top; '
-            f'border-bottom: 1px solid #EEEEEE">{key_span}</td>'
-            '<td style="padding: 5px; vertical-align: top; '
-            f'border-bottom: 1px solid #EEEEEE">{cs}</td></tr>'
-        )
-      s.write('</table></div>')
+      s.write(repr_utils.html_repr(self.metadata))
+    s.write('</div>')
     return s.getvalue()
 
   @property
-  def _message_text_color(self) -> str:
+  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':
@@ -99,25 +88,9 @@ class LogEntry(pg.Object):
 
   @property
   def _message_display(self) -> str:
-    return self._round_text(
+    return repr_utils.html_round_text(
         self._time_display + '&nbsp;' + self.message,
-        color=self._message_text_color,
-    )
-
-  def _round_text(
-      self,
-      text: str,
-      *,
-      color: str = '#EEEEEE',
-      display: str = 'inline-block',
-      margin_top: str = '5px',
-      margin_bottom: str = '5px',
-      whitespace: str = 'pre-wrap') -> str:
-    return (
-        f'<span style="background:{color}; display:{display};'
-        f'border-radius:10px; padding:5px; '
-        f'margin-top: {margin_top}; margin-bottom: {margin_bottom}; '
-        f'white-space: {whitespace}">{text}</span>'
+        background_color=self._message_text_bgcolor,
     )
 
 

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/langfun/core/logging_test.py

@@ -43,8 +43,14 @@ class LoggingTest(unittest.TestCase):
     self.assertEqual(logging.fatal('hi').level, 'fatal')
 
   def test_repr_html(self):
-    entry = logging.log('info', 'hi', indent=1, x=1, y=2)
-    self.assertIn('<div', entry._repr_html_())
+    def assert_color(entry, color):
+      self.assertIn(f'background-color: {color}', entry._repr_html_())
+
+    assert_color(logging.debug('hi', indent=0), '#EEEEEE')
+    assert_color(logging.info('hi', indent=1), '#A3E4D7')
+    assert_color(logging.warning('hi', x=1, y=2), '#F8C471')
+    assert_color(logging.error('hi', indent=2, x=1, y=2), '#F5C6CB')
+    assert_color(logging.fatal('hi', indent=2, x=1, y=2), '#F19CBB')
 
 
 if __name__ == '__main__':

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/langfun/core/message.py

@@ -14,11 +14,13 @@
 """Messages that are exchanged between users and agents."""
 
 import contextlib
+import html
 import io
 from typing import Annotated, Any, Optional, Union
 
 from langfun.core import modality
 from langfun.core import natural_language
+from langfun.core import repr_utils
 import pyglove as pg
 
 
@@ -304,15 +306,16 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
         m.referred_name: m for m in chunks if isinstance(m, modality.Modality)
     }
 
-  def chunk(self) -> list[str | modality.Modality]:
+  def chunk(self, text: str | None = None) -> list[str | modality.Modality]:
     """Chunk a message into a list of str or modality objects."""
     chunks = []
 
     def add_text_chunk(text_piece: str) -> None:
       if text_piece:
         chunks.append(text_piece)
+    if text is None:
+      text = self.text
 
-    text = self.text
     chunk_start = 0
     ref_end = 0
     while chunk_start < len(text):
@@ -490,6 +493,79 @@ 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(
+      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(),
+          )
+      )
+      s.write('<hr>')
+
+    # Body.
+    s.write(
+        f'<span style="color: {self._text_color()}; white-space: pre-wrap;">'
+    )
+
+    # 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'
 
 #
 # Messages of different roles.

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/langfun/core/message_test.py

@@ -26,6 +26,9 @@ 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):
 
@@ -315,6 +318,52 @@ 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')),
+    )
+    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.assertIn(
+        'background-color: blue',
+        message.AIMessage('hi').to_html().content,
+    )
+    self.assertIn(
+        'background-color: black',
+        message.SystemMessage('hi').to_html().content,
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun-0.1.1.dev20240807/langfun/core/repr_utils.py

@@ -0,0 +1,173 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Helpers for implementing _repr_xxx_ methods."""
+
+import collections
+import contextlib
+import io
+from typing import Any, Callable, Iterator
+
+from langfun.core import component
+import pyglove as pg
+
+
+class Html(pg.Object):
+  """A HTML adapter for rendering."""
+  content: str
+
+  def _repr_html_(self) -> str:
+    return self.content
+
+
+@contextlib.contextmanager
+def share_parts() -> Iterator[dict[str, int]]:
+  """Context manager for defining the context (scope) of shared content.
+
+  Under the context manager, call to `lf.write_shared` with the same content
+  will be written only once. This is useful for writing shared content such as
+  shared style and script sections in the HTML.
+
+  Example:
+    ```
+    class Foo(pg.Object):
+      def _repr_html_(self) -> str:
+        s = io.StringIO()
+        lf.repr_utils.write_shared_part(s, '<style>..</style>')
+        lf.repr_utils.write_shared_part(s, '<script>..</script>')
+        return s.getvalue()
+
+    with lf.repr_utils.share_parts() as share_parts:
+      # The <style> and <script> section will be written only once.
+      lf.console.display(Foo())
+      lf.console.display(Foo())
+
+    # Assert that the shared content is attempted to be written twice.
+    assert share_parts['<style>..</style>'] == 2
+    ```
+
+  Yields:
+    A dictionary mapping the shared content to the number of times it is
+    attempted to be written.
+  """
+  context = component.context_value(
+      '__shared_parts__', collections.defaultdict(int)
+  )
+  with component.context(__shared_parts__=context):
+    try:
+      yield context
+    finally:
+      pass
+
+
+def write_maybe_shared(s: io.StringIO, content: str) -> bool:
+  """Writes a maybe shared part to an string stream.
+
+  Args:
+    s: The string stream to write to.
+    content: A maybe shared content to write.
+
+  Returns:
+    True if the content is written to the string. False if the content is
+    already written under the same share context.
+  """
+  context = component.context_value('__shared_parts__', None)
+  if context is None:
+    s.write(content)
+    return True
+  written = content in context
+  if not written:
+    s.write(content)
+  context[content] += 1
+  return not written
+
+
+def html_repr(
+    value: Any,
+    item_color: Callable[
+        [str, str],
+        tuple[
+            str | None,   # Label text color
+            str | None,   # Label background color
+            str | None,   # Value text color
+            str | None,   # Value background color
+        ]
+    ] | None = None    # pylint: disable=bad-whitespace
+) -> str:
+  """Writes a list of key-value pairs to an string stream.
+
+  Args:
+    value: A value to be rendered in HTML.
+    item_color: A function that takes the key and value and returns a tuple
+      of four strings, the text color and background color of the label and
+      value respectively. If None, a default color scheme will be used.
+
+  Returns:
+    The HTML representation of the value.
+  """
+  s = io.StringIO()
+  s.write('<div style="padding-left: 20px; margin-top: 10px">')
+  s.write('<table style="border-top: 1px solid #EEEEEE;">')
+  item_color = item_color or (lambda k, v: (None, '#F1C40F', None, None))
+
+  for k, v in pg.object_utils.flatten(value).items():
+    if isinstance(v, pg.Ref):
+      v = v.value
+    if hasattr(v, '_repr_html_'):
+      cs = v._repr_html_()  # pylint: disable=protected-access
+    else:
+      cs = f'<span style="white-space: pre-wrap">{str(v)}</span>'
+
+    key_color, key_bg_color, value_color, value_bg_color = item_color(k, v)
+    key_span = html_round_text(
+        k,
+        text_color=key_color,
+        background_color=key_bg_color,
+        margin_bottom='0px'
+    )
+    value_color_style = f'color: {value_color};' if value_color else ''
+    value_bg_color_style = (
+        f'background-color: {value_bg_color};' if value_bg_color else ''
+    )
+    s.write(
+        '<tr>'
+        '<td style="padding: 5px; vertical-align: top; '
+        f'border-bottom: 1px solid #EEEEEE">{key_span}</td>'
+        '<td style="padding: 15px 5px 5px 5px; vertical-align: top; '
+        'border-bottom: 1px solid #EEEEEE;'
+        f'{value_color_style}{value_bg_color_style}">{cs}</td></tr>'
+    )
+  s.write('</table></div>')
+  return s.getvalue()
+
+
+def html_round_text(
+    text: str,
+    *,
+    text_color: str = 'black',
+    background_color: str = '#EEEEEE',
+    display: str = 'inline-block',
+    margin_top: str = '5px',
+    margin_bottom: str = '5px',
+    whitespace: str = 'pre-wrap') -> str:
+  """Renders a HTML span with rounded corners."""
+  color_style = f'color: {text_color};' if text_color else ''
+  bg_color_style = (
+      f'background-color: {background_color};' if background_color else ''
+  )
+  return (
+      f'<span style="{color_style}{bg_color_style}'
+      f'display:{display}; border-radius:10px; padding:5px; '
+      f'margin-top: {margin_top}; margin-bottom: {margin_bottom}; '
+      f'white-space: {whitespace}">{text}</span>'
+  )

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/langfun/core/repr_utils_test.py

@@ -17,6 +17,7 @@ import io
 import unittest
 
 from langfun.core import repr_utils
+import pyglove as pg
 
 
 class SharingContentTest(unittest.TestCase):
@@ -53,6 +54,19 @@ class SharingContentTest(unittest.TestCase):
     self.assertEqual(ctx1['<style>b</style>'], 2)
     self.assertEqual(ctx1['<style>a</style>'], 4)
 
+  def test_html(self):
+    html = repr_utils.Html('<div>foo</div>')
+    self.assertEqual(html.content, '<div>foo</div>')
+    self.assertEqual(html._repr_html_(), '<div>foo</div>')
+
+  def test_html_repr(self):
+    class Foo(pg.Object):
+      x: int
+
+    html = repr_utils.html_repr({'foo': pg.Ref(Foo(1))})
+    self.assertIn('foo</span>', html)
+    self.assertNotIn('Ref', html)
+
 
 if __name__ == '__main__':
   unittest.main()

{langfun-0.1.1.dev20240805 → langfun-0.1.1.dev20240807}/langfun.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.1.dev20240805
+Version: 0.1.1.dev20240807
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

langfun-0.1.1.dev20240805/langfun/core/repr_utils.py

@@ -1,83 +0,0 @@
-# Copyright 2024 The Langfun Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Helpers for implementing _repr_xxx_ methods."""
-
-import collections
-import contextlib
-import io
-from typing import Iterator
-
-from langfun.core import component
-
-
-@contextlib.contextmanager
-def share_parts() -> Iterator[dict[str, int]]:
-  """Context manager for defining the context (scope) of shared content.
-
-  Under the context manager, call to `lf.write_shared` with the same content
-  will be written only once. This is useful for writing shared content such as
-  shared style and script sections in the HTML.
-
-  Example:
-    ```
-    class Foo(pg.Object):
-      def _repr_html_(self) -> str:
-        s = io.StringIO()
-        lf.repr_utils.write_shared_part(s, '<style>..</style>')
-        lf.repr_utils.write_shared_part(s, '<script>..</script>')
-        return s.getvalue()
-
-    with lf.repr_utils.share_parts() as share_parts:
-      # The <style> and <script> section will be written only once.
-      lf.console.display(Foo())
-      lf.console.display(Foo())
-
-    # Assert that the shared content is attempted to be written twice.
-    assert share_parts['<style>..</style>'] == 2
-    ```
-
-  Yields:
-    A dictionary mapping the shared content to the number of times it is
-    attempted to be written.
-  """
-  context = component.context_value(
-      '__shared_parts__', collections.defaultdict(int)
-  )
-  with component.context(__shared_parts__=context):
-    try:
-      yield context
-    finally:
-      pass
-
-
-def write_maybe_shared(s: io.StringIO, content: str) -> bool:
-  """Writes a maybe shared part to an string stream.
-
-  Args:
-    s: The string stream to write to.
-    content: A maybe shared content to write.
-
-  Returns:
-    True if the content is written to the string. False if the content is
-    already written under the same share context.
-  """
-  context = component.context_value('__shared_parts__', None)
-  if context is None:
-    s.write(content)
-    return True
-  written = content in context
-  if not written:
-    s.write(content)
-  context[content] += 1
-  return not written