langfun 0.0.2.dev20240429__py3-none-any.whl → 0.1.2.dev202501140804__py3-none-any.whl

langfun/core/message.py

@@ -14,6 +14,7 @@
 """Messages that are exchanged between users and agents."""
 
 import contextlib
+import functools
 import io
 from typing import Annotated, Any, Optional, Union
 
@@ -22,7 +23,11 @@ from langfun.core import natural_language
 import pyglove as pg
 
 
-class Message(natural_language.NaturalLanguageFormattable, pg.Object):
+class Message(
+    natural_language.NaturalLanguageFormattable,
+    pg.Object,
+    pg.views.HtmlTreeView.Extension
+):
   """Message.
 
   ``Message`` is the protocol for users and the system to interact with
@@ -144,7 +149,7 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
   def from_value(cls, value: Union[str, 'Message']) -> 'Message':
     """Creates a message from a value or return value itself if a Message."""
     if isinstance(value, modality.Modality):
-      return cls('{{object}}', object=value)
+      return cls('<<[[object]]>>', object=value)
     if isinstance(value, Message):
       return value
     return cls(value)
@@ -276,6 +281,16 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
   # API for supporting modalities.
   #
 
+  @property
+  def text_with_modality_hash(self) -> str:
+    """Returns text with modality object placeheld by their 8-byte MD5 hash."""
+    parts = [self.text]
+    for name, modality_obj in self.referred_modalities().items():
+      parts.append(
+          f'<{name}>{modality_obj.hash}</{name}>'
+      )
+    return ''.join(parts)
+
   def get_modality(
       self, var_name: str, default: Any = None, from_message_chain: bool = True
   ) -> modality.Modality | None:
@@ -304,21 +319,22 @@ 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):
       ref_start = text.find(modality.Modality.REF_START, ref_end)
       if ref_start == -1:
-        add_text_chunk(text[chunk_start:].strip())
+        add_text_chunk(text[chunk_start:].strip(' '))
         break
 
       var_start = ref_start + len(modality.Modality.REF_START)
@@ -330,29 +346,31 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
       var_name = text[var_start:ref_end].strip()
       var_value = self.get_modality(var_name)
       if var_value is not None:
-        add_text_chunk(text[chunk_start:ref_start].strip())
+        add_text_chunk(text[chunk_start:ref_start].strip(' '))
         chunks.append(var_value)
         chunk_start = ref_end + len(modality.Modality.REF_END)
     return chunks
 
   @classmethod
   def from_chunks(
-      cls, chunks: list[str | modality.Modality], separator: str = '\n'
+      cls, chunks: list[str | modality.Modality], separator: str = ' '
   ) -> 'Message':
     """Assembly a message from a list of string or modality objects."""
     fused_text = io.StringIO()
     ref_index = 0
     metadata = dict()
-
+    last_char = None
     for i, chunk in enumerate(chunks):
-      if i > 0:
+      if i > 0 and last_char not in ('\t', ' ', '\n'):
         fused_text.write(separator)
       if isinstance(chunk, str):
         fused_text.write(chunk)
+        last_char = chunk[-1]
       else:
         assert isinstance(chunk, modality.Modality), chunk
         var_name = f'obj{ref_index}'
         fused_text.write(modality.Modality.text_marker(var_name))
+        last_char = modality.Modality.REF_END[-1]
         # Make a reference if the chunk is already owned by another object
         # to avoid copy.
         metadata[var_name] = pg.maybe_ref(chunk)
@@ -391,6 +409,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.
   #
@@ -488,6 +511,293 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
     v = self.metadata[key]
     return v.value if isinstance(v, pg.Ref) else v
 
+  def _html_tree_view_content(
+      self,
+      *,
+      view: pg.views.HtmlTreeView,
+      root_path: pg.KeyPath | None = None,
+      collapse_level: int | None = None,
+      extra_flags: dict[str, Any] | None = None,
+      **kwargs,
+  ) -> pg.Html:
+    """Returns the HTML representation of the message.
+
+    Args:
+      view: The HTML tree view.
+      root_path: The root path of the message.
+      collapse_level: The global collapse level.
+      extra_flags: Extra flags to control the rendering.
+        - 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: Omitted keyword arguments.
+
+    Returns:
+      The HTML representation of the message content.
+    """
+    extra_flags = extra_flags if extra_flags is not None else {}
+
+    include_message_metadata: bool = extra_flags.get(
+        'include_message_metadata', True
+    )
+    source_tag: str | tuple[str, ...] | None = extra_flags.get(
+        'source_tag', ('lm-input', 'lm-output')
+    )
+    collapse_modalities_in_text: bool = extra_flags.get(
+        'collapse_modalities_in_text', True
+    )
+    collapse_llm_usage: bool = extra_flags.get(
+        'collapse_llm_usage', False
+    )
+    collapse_message_result_level: int | None = extra_flags.get(
+        'collapse_message_result_level', 1
+    )
+    collapse_message_metadata_level: int | None = extra_flags.get(
+        'collapse_message_metadata_level', 1
+    )
+    collapse_source_message_level: int | None = extra_flags.get(
+        'collapse_source_message_level', 1
+    )
+    passthrough_kwargs = view.get_passthrough_kwargs(**kwargs)
+    def render_tags():
+      return pg.Html.element(
+          'div',
+          [pg.Html.element('span', [tag]) for tag in self.tags],
+          css_classes=['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 = pg.KeyPath(['metadata', chunk.referred_name], root_path)
+          s.write(
+              pg.Html.element(
+                  'div',
+                  [
+                      view.render(
+                          chunk,
+                          name=chunk.referred_name,
+                          root_path=child_path,
+                          collapse_level=(
+                              0 if collapse_modalities_in_text else 1
+                          ),
+                          extra_flags=dict(
+                              display_modality_when_hover=True,
+                          ),
+                          **passthrough_kwargs,
+                      )
+                  ],
+                  css_classes=['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 = pg.KeyPath(['metadata', 'result'], root_path)
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.result,
+                  name='result',
+                  root_path=child_path,
+                  collapse_level=view.get_collapse_level(
+                      (collapse_level, -1),
+                      collapse_message_result_level,
+                  ),
+                  extra_flags=extra_flags,
+                  **passthrough_kwargs,
+              )
+          ],
+          css_classes=['message-result'],
+      )
+
+    def render_usage():
+      if 'usage' not in self.metadata:
+        return None
+      child_path = pg.KeyPath(['metadata', 'usage'], root_path)
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.usage,
+                  name='llm usage',
+                  key_style='label',
+                  root_path=child_path,
+                  collapse_level=view.get_collapse_level(
+                      (collapse_level, -1),
+                      0 if collapse_llm_usage else 1,
+                  ),
+                  extra_flags=extra_flags,
+                  **view.get_passthrough_kwargs(
+                      remove=['key_style'], **kwargs
+                  ),
+              )
+          ],
+          css_classes=['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 = pg.KeyPath('source', root_path)
+        child_extra_flags = extra_flags.copy()
+        child_extra_flags['collapse_source_message_level'] = (
+            view.get_collapse_level(
+                (collapse_source_message_level, -1), 0,
+            )
+        )
+        return view.render(
+            self.source,
+            name='source',
+            root_path=child_path,
+            collapse_level=view.get_collapse_level(
+                (collapse_level, -1),
+                collapse_source_message_level,
+            ),
+            extra_flags=child_extra_flags,
+            **passthrough_kwargs,
+        )
+      return None
+
+    def render_metadata():
+      if not include_message_metadata:
+        return None
+      child_path = pg.KeyPath('metadata', root_path)
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.metadata,
+                  css_classes=['message-metadata'],
+                  exclude_keys=['usage', 'result'],
+                  name='metadata',
+                  root_path=child_path,
+                  collapse_level=view.get_collapse_level(
+                      (collapse_level, -1),
+                      collapse_message_metadata_level,
+                  ),
+                  **view.get_passthrough_kwargs(
+                      remove=['exclude_keys'], **kwargs
+                  ),
+              )
+          ],
+          css_classes=['message-metadata'],
+      )
+
+    return pg.Html.element(
+        'div',
+        [
+            render_tags(),
+            render_message_text(),
+            render_result(),
+            render_usage(),
+            render_metadata(),
+            render_source_message(),
+        ],
+        css_classes=['complex_value'],
+    )
+
+  @classmethod
+  @functools.cache
+  def _html_tree_view_config(cls) -> dict[str, Any]:
+    return pg.views.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            css_classes=['lf-message'],
+        )
+    )
+
+  @classmethod
+  @functools.cache
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        /* Langfun Message styles.*/
+        [class^="message-"] > details {
+            margin: 0px 0px 5px 0px;
+            border: 1px solid #EEE;
+        }
+        .lf-message.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;
+            white-space: pre-wrap;
+            border: 1px solid #EEE;
+            border-radius: 5px;
+            background-color: #EEE;
+        }
+        .modality-in-text {
+            display: inline-block;
+        }
+        .modality-in-text > details.pyglove {
+            display: inline-block;
+            font-size: 0.8em;
+            border: 0;
+            background-color: #A6F1A6;
+            margin: 0px 5px 0px 5px;
+        }
+        .message-result {
+            color: dodgerblue;
+        }
+        .message-usage {
+            color: orange;
+        }
+        .message-usage .object-key.str {
+            border: 1px solid orange;
+            background-color: orange;
+            color: white;
+        }
+        """
+    ]
+
 
 #
 # Messages of different roles.