langfun 0.1.2.dev202410140804__py3-none-any.whl → 0.1.2.dev202410180804__py3-none-any.whl

langfun/core/component.py

@@ -286,6 +286,54 @@ class ContextualAttribute(pg.symbolic.ValueFromParentChain):
     else:
       return pg.MISSING_VALUE
 
+  def _html_tree_view_content(
+      self,
+      *,
+      view: pg.views.HtmlTreeView,
+      parent: Any,
+      root_path: pg.KeyPath,
+      **kwargs,
+  ) -> pg.Html:
+    inferred_value = pg.MISSING_VALUE
+    if isinstance(parent, pg.Symbolic) and root_path:
+      inferred_value = parent.sym_inferred(root_path.key, pg.MISSING_VALUE)
+
+    if inferred_value is not pg.MISSING_VALUE:
+      kwargs.pop('name', None)
+      return view.render(
+          inferred_value, parent=self, root_path=root_path + '<inferred>',
+          **view.get_passthrough_kwargs(**kwargs)
+      )
+    return pg.Html.element(
+        'div',
+        [
+            '(not available)',
+        ],
+        css_classes=['unavailable-contextual'],
+    )
+
+  def _html_tree_view_config(self) -> dict[str, Any]:
+    return pg.views.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            collapse_level=1,
+        )
+    )
+
+  @classmethod
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        .contextual-attribute {
+          color: purple;
+        }
+        .unavailable-contextual {
+          color: gray;
+          font-style: italic;
+        }
+        """
+    ]
+
 
 # NOTE(daiyip): Returning Any instead of `lf.ContextualAttribute` to avoid
 # pytype check error as `contextual()` can be assigned to any type.

langfun/core/component_test.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 """Contextual component and app test."""
 
+import inspect
+from typing import Any
 import unittest
 import weakref
 
@@ -297,6 +299,52 @@ class ContextualAttributeTest(unittest.TestCase):
     self.assertEqual(c.z, 3)
     self.assertEqual(b.z, 3)
 
+  def test_to_html(self):
+    class A(lf.Component):
+      x: int = 1
+      y: int = lf.contextual()
+
+    def assert_content(html, expected):
+      expected = inspect.cleandoc(expected).strip()
+      actual = html.content.strip()
+      if actual != expected:
+        print(actual)
+      self.assertEqual(actual.strip(), expected)
+
+    self.assertIn(
+        inspect.cleandoc(
+            """
+            .contextual-attribute {
+              color: purple;
+            }
+            .unavailable-contextual {
+              color: gray;
+              font-style: italic;
+            }
+            """
+        ),
+        A().to_html().style_section,
+    )
+
+    assert_content(
+        A().to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove a"><summary><div class="summary-title">A(...)</div></summary><div class="complex-value a"><details open class="pyglove int"><summary><div class="summary-name">x<span class="tooltip">x</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details><details open class="pyglove contextual-attribute"><summary><div class="summary-name">y<span class="tooltip">y</span></div><div class="summary-title">ContextualAttribute(...)</div></summary><div class="unavailable-contextual">(not available)</div></details></div></details>
+        """
+    )
+
+    class B(lf.Component):
+      z: Any
+      y: int = 2
+
+    b = B(A())
+    assert_content(
+        b.z.to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove a"><summary><div class="summary-title">A(...)</div></summary><div class="complex-value a"><details open class="pyglove int"><summary><div class="summary-name">x<span class="tooltip">x</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details><details open class="pyglove contextual-attribute"><summary><div class="summary-name">y<span class="tooltip">y</span></div><div class="summary-title">ContextualAttribute(...)</div></summary><span class="simple-value int">2</span></details></div></details>
+        """
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/logging.py

@@ -15,8 +15,9 @@
 
 import contextlib
 import datetime
+import functools
 import typing
-from typing import Any, Iterator, Literal, Sequence
+from typing import Any, Iterator, Literal
 
 from langfun.core import component
 from langfun.core import console
@@ -57,11 +58,11 @@ class LogEntry(pg.Object):
       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
+      max_summary_len_for_str: int = 80,
       **kwargs
       ) -> str:
-    if len(self.message) > max_str_len_for_summary:
-      message = self.message[:max_str_len_for_summary] + '...'
+    if len(self.message) > max_summary_len_for_str:
+      message = self.message[:max_summary_len_for_str] + '...'
     else:
       message = self.message
 
@@ -69,18 +70,18 @@ class LogEntry(pg.Object):
         pg.Html.element(
             'span',
             [self.time.strftime('%H:%M:%S')],
-            css_class=['log-time']
+            css_classes=['log-time']
         ),
         pg.Html.element(
             'span',
             [pg.Html.escape(message)],
-            css_class=['log-summary'],
+            css_classes=['log-summary'],
         ),
     )
     return view.summary(
         self,
         title=title or s,
-        max_str_len_for_summary=max_str_len_for_summary,
+        max_summary_len_for_str=max_summary_len_for_str,
         **kwargs,
     )
 
@@ -89,43 +90,45 @@ class LogEntry(pg.Object):
       self,
       view: pg.views.HtmlTreeView,
       root_path: pg.KeyPath,
-      collapse_log_metadata_level: int | None = pg.View.PresetArgValue(0),
-      max_str_len_for_summary: int = pg.View.PresetArgValue(80),
-      collapse_level: int | None = pg.View.PresetArgValue(1),
+      max_summary_len_for_str: int = 80,
+      collapse_level: int | None = 1,
+      extra_flags: dict[str, Any] | None = None,
       **kwargs
   ) -> pg.Html:
     # pytype: enable=annotation-type-mismatch
+    extra_flags = extra_flags if extra_flags is not None else {}
+    collapse_log_metadata_level: int | None = extra_flags.get(
+        'collapse_log_metadata_level', None
+    )
     def render_message_text():
-      if len(self.message) < max_str_len_for_summary:
+      if len(self.message) < max_summary_len_for_str:
         return None
       return pg.Html.element(
           'span',
           [pg.Html.escape(self.message)],
-          css_class=['log-text'],
+          css_classes=['log-text'],
       )
 
     def render_metadata():
       if not self.metadata:
         return None
-      child_path = root_path + 'metadata'
       return pg.Html.element(
           'div',
           [
               view.render(
                   self.metadata,
                   name='metadata',
-                  root_path=child_path,
+                  root_path=root_path + 'metadata',
                   parent=self,
-                  collapse_level=(
-                      view.max_collapse_level(
-                          collapse_level,
-                          collapse_log_metadata_level,
-                          child_path
-                      )
-                  )
+                  collapse_level=view.get_collapse_level(
+                      (collapse_level, -1), collapse_log_metadata_level,
+                  ),
+                  max_summary_len_for_str=max_summary_len_for_str,
+                  extra_flags=extra_flags,
+                  **view.get_passthrough_kwargs(**kwargs),
               )
           ],
-          css_class=['log-metadata'],
+          css_classes=['log-metadata'],
       )
 
     return pg.Html.element(
@@ -134,12 +137,23 @@ class LogEntry(pg.Object):
             render_message_text(),
             render_metadata(),
         ],
-        css_class=['complex_value'],
+        css_classes=['complex_value'],
+    )
+
+  def _html_tree_view_config(self) -> dict[str, Any]:
+    return pg.views.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            css_classes=[f'log-{self.level}'],
+        )
     )
 
-  def _html_style(self) -> list[str]:
-    return super()._html_style() + [
+  @classmethod
+  @functools.cache
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
         """
+        /* Langfun LogEntry styles. */
         .log-time {
           color: #222;
           font-size: 12px;
@@ -203,9 +217,6 @@ class LogEntry(pg.Object):
         """
     ]
 
-  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

@@ -69,20 +69,36 @@ class LoggingTest(unittest.TestCase):
             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>
+        <details open class="pyglove log-entry log-info"><summary><div class="summary-title log-info"><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)
+            time=time, metadata=dict(x=dict(z=1), y=2)
         ).to_html(
-            max_str_len_for_summary=10,
+            extra_flags=dict(
+                collapse_log_metadata_level=1,
+            ),
+            max_summary_len_for_str=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>
+        <details open class="pyglove log-entry log-error"><summary><div class="summary-title log-error"><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<span class="tooltip">metadata</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details class="pyglove dict"><summary><div class="summary-name">x<span class="tooltip">metadata.x</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details open class="pyglove int"><summary><div class="summary-name">z<span class="tooltip">metadata.x.z</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details></div></details><details open class="pyglove int"><summary><div class="summary-name">y<span class="tooltip">metadata.y</span></div><div class="summary-title">int</div></summary><span class="simple-value int">2</span></details></div></details></div></div></details>
+        """
+    )
+    self.assert_html_content(
+        logging.LogEntry(
+            level='error', message='This is a longer message: 5 + 2 > 3',
+            time=time, metadata=dict(x=dict(z=1), y=2)
+        ).to_html(
+            extra_flags=dict(
+                max_summary_len_for_str=10,
+            ),
+            enable_summary_tooltip=False,
+        ),
+        """
+        <details open class="pyglove log-entry log-error"><summary><div class="summary-title log-error"><span class="log-time">12:30:45</span><span class="log-summary">This is a longer message: 5 + 2 &gt; 3</span></div></summary><div class="complex_value"><div class="log-metadata"><details open class="pyglove dict"><summary><div class="summary-name">metadata<span class="tooltip">metadata</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details open class="pyglove dict"><summary><div class="summary-name">x<span class="tooltip">metadata.x</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details open class="pyglove int"><summary><div class="summary-name">z<span class="tooltip">metadata.x.z</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details></div></details><details open class="pyglove int"><summary><div class="summary-name">y<span class="tooltip">metadata.y</span></div><div class="summary-title">int</div></summary><span class="simple-value int">2</span></details></div></details></div></div></details>
         """
     )
 

langfun/core/message.py

@@ -14,8 +14,9 @@
 """Messages that are exchanged between users and agents."""
 
 import contextlib
+import functools
 import io
-from typing import Annotated, Any, Optional, Sequence, Union
+from typing import Annotated, Any, Optional, Union
 
 from langfun.core import modality
 from langfun.core import natural_language
@@ -506,53 +507,70 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
     v = self.metadata[key]
     return v.value if isinstance(v, pg.Ref) else v
 
-  # pytype: disable=annotation-type-mismatch
   def _html_tree_view_content(
       self,
       *,
       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),
+      collapse_level: int | None = None,
+      extra_flags: dict[str, Any] | None = None,
       **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.
+      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_class=['message-tags'],
+          css_classes=['message-tags'],
       )
 
     def render_message_text():
@@ -573,12 +591,16 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
                           chunk,
                           name=chunk.referred_name,
                           root_path=child_path,
-                          collapse_level=child_path.depth + (
+                          collapse_level=(
                               0 if collapse_modalities_in_text else 1
-                          )
+                          ),
+                          extra_flags=dict(
+                              display_modality_when_hover=True,
+                          ),
+                          **passthrough_kwargs,
                       )
                   ],
-                  css_class=['modality-in-text'],
+                  css_classes=['modality-in-text'],
               )
           )
           referred_chunks[chunk.referred_name] = chunk
@@ -596,14 +618,15 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
                   self.result,
                   name='result',
                   root_path=child_path,
-                  collapse_level=view.max_collapse_level(
-                      collapse_level,
+                  collapse_level=view.get_collapse_level(
+                      (collapse_level, -1),
                       collapse_message_result_level,
-                      child_path,
-                  )
+                  ),
+                  extra_flags=extra_flags,
+                  **passthrough_kwargs,
               )
           ],
-          css_class=['message-result'],
+          css_classes=['message-result'],
       )
 
     def render_usage():
@@ -616,15 +639,19 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
               view.render(
                   self.usage,
                   name='llm usage',
+                  key_style='label',
                   root_path=child_path,
-                  collapse_level=view.max_collapse_level(
-                      collapse_level,
+                  collapse_level=view.get_collapse_level(
+                      (collapse_level, -1),
                       0 if collapse_llm_usage else 1,
-                      child_path,
-                  )
+                  ),
+                  extra_flags=extra_flags,
+                  **view.get_passthrough_kwargs(
+                      remove=['key_style'], **kwargs
+                  ),
               )
           ],
-          css_class=['message-usage'],
+          css_classes=['message-usage'],
       )
 
     def render_source_message():
@@ -635,21 +662,22 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
         source = source.source
       if source is not None:
         child_path = root_path + 'source'
+        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,
-            include_metadata=include_message_metadata,
-            collapse_level=view.max_collapse_level(
-                collapse_level,
+            collapse_level=view.get_collapse_level(
+                (collapse_level, -1),
                 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,
+            extra_flags=child_extra_flags,
+            **passthrough_kwargs,
         )
       return None
 
@@ -662,17 +690,20 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
           [
               view.render(
                   self.metadata,
-                  css_class=['message-metadata'],
+                  css_classes=['message-metadata'],
+                  exclude_keys=['usage', 'result'],
                   name='metadata',
                   root_path=child_path,
-                  collapse_level=view.max_collapse_level(
-                      collapse_level,
+                  collapse_level=view.get_collapse_level(
+                      (collapse_level, -1),
                       collapse_message_metadata_level,
-                      child_path,
-                  )
+                  ),
+                  **view.get_passthrough_kwargs(
+                      remove=['exclude_keys'], **kwargs
+                  ),
               )
           ],
-          css_class=['message-metadata'],
+          css_classes=['message-metadata'],
       )
 
     return pg.Html.element(
@@ -685,18 +716,30 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
             render_metadata(),
             render_source_message(),
         ],
-        css_class=['complex_value'],
+        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'],
+        )
     )
 
-  def _html_style(self) -> list[str]:
-    return super()._html_style() + [
+  @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;
         }
-        details.lf-message > summary > .summary_title::after {
+        .lf-message.summary-title::after {
             content: ' 💬';
         }
         details.pyglove.ai-message {
@@ -739,12 +782,12 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
             margin: 0px 5px 0px 5px;
         }
         .message-result {
-            color: purple;
+            color: dodgerblue;
         }
         .message-usage {
             color: orange;
         }
-        .message-usage .object_key.str {
+        .message-usage .object-key.str {
             border: 1px solid orange;
             background-color: orange;
             color: white;
@@ -752,9 +795,6 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
         """
     ]
 
-  def _html_element_class(self) -> list[str]:
-    return super()._html_element_class() + ['lf-message']
-
 
 #
 # Messages of different roles.