From d5ff5ba51325de4300c0b763458234e7b3e4484b Mon Sep 17 00:00:00 2001
From: Daiyi Peng <daiyip@google.com>
Date: Tue, 10 Dec 2024 23:33:00 -0800
Subject: [PATCH] Improve the HTML view for agentic actions.

1) Use tab control for rendering action/result/metadata/execution, with current action displayed in the tab header.
2) Select `output` as the default tab for QueryInvocation.

PiperOrigin-RevId: 704987034
---
 langfun/core/agentic/action.py       | 337 +++++++++++----------------
 langfun/core/agentic/action_test.py  |   6 +-
 langfun/core/structured/prompting.py |   2 +-
 3 files changed, 146 insertions(+), 199 deletions(-)

diff --git a/langfun/core/agentic/action.py b/langfun/core/agentic/action.py
index 98bac89..fd25943 100644
--- a/langfun/core/agentic/action.py
+++ b/langfun/core/agentic/action.py
@@ -32,7 +32,7 @@ def _on_bound(self):
     super()._on_bound()
     self._session = None
     self._result = None
-    self._result_metadata = {}
+    self._metadata = {}
 
   @property
   def session(self) -> Optional['Session']:
@@ -45,9 +45,9 @@ def result(self) -> Any:
     return self._result
 
   @property
-  def result_metadata(self) -> dict[str, Any] | None:
+  def metadata(self) -> dict[str, Any] | None:
     """Returns the metadata associated with the result from previous call."""
-    return self._result_metadata
+    return self._metadata
 
   def __call__(
       self,
@@ -69,7 +69,7 @@ def __call__(
       if new_session:
         self._session = session
       self._result = result
-      self._result_metadata = session.current_action.result_metadata
+      self._metadata = session.current_action.metadata
       return self._result
 
   @abc.abstractmethod
@@ -139,7 +139,7 @@ def start(self) -> None:
     if self._time_badge is not None:
       self._time_badge.update(
           'Starting',
-          add_class=['running'],
+          add_class=['starting'],
           remove_class=['not-started'],
       )
 
@@ -149,10 +149,14 @@ def stop(self) -> None:
     if self._time_badge is not None:
       self._time_badge.update(
           f'{int(self.elapse)} seconds',
+          tooltip=f'Execution finished in {int(self.elapse)} seconds.',
           add_class=['finished'],
           remove_class=['running'],
       )
 
+  def __len__(self) -> int:
+    return len(self.items)
+
   @property
   def has_started(self) -> bool:
     return self.start_time is not None
@@ -238,17 +242,8 @@ def append(self, item: TracedItem) -> None:
         and not isinstance(item, lf.logging.LogEntry)):
       sub_task_label = self._execution_item_label(item)
       self._time_badge.update(
-          pg.Html.element(
-              'span',
-              [
-                  'Running',
-                  pg.views.html.controls.Badge(
-                      sub_task_label.text,
-                      tooltip=sub_task_label.tooltip,
-                      css_classes=['task-in-progress']
-                  )
-              ]
-          ),
+          text=sub_task_label.text,
+          tooltip=sub_task_label.tooltip.content,
           add_class=['running'],
           remove_class=['not-started'],
       )
@@ -274,44 +269,30 @@ def _html_tree_view_summary(
       extra_flags: dict[str, Any] | None = None,
       view: pg.views.html.HtmlTreeView, **kwargs
   ):
-    extra_flags = extra_flags or {}
-    interactive = extra_flags.get('interactive', True)
-    def time_badge():
-      if not self.has_started:
-        label = '(Not started)'
-        css_class = 'not-started'
-      elif not self.has_stopped:
-        label = 'Starting'
-        css_class = 'running'
-      else:
-        label = f'{int(self.elapse)} seconds'
-        css_class = 'finished'
-      return pg.views.html.controls.Badge(
-          label,
-          css_classes=['execution-time', css_class],
-          interactive=interactive,
-      )
-    time_badge = time_badge()
+    return None
+
+  def _execution_badge(self, interactive: bool = True):
+    if not self.has_started:
+      label = '(Not started)'
+      tooltip = 'Execution not started.'
+      css_class = 'not-started'
+    elif not self.has_stopped:
+      label = 'Starting'
+      tooltip = 'Execution starting.'
+      css_class = 'running'
+    else:
+      label = f'{int(self.elapse)} seconds'
+      tooltip = f'Execution finished in {int(self.elapse)} seconds.'
+      css_class = 'finished'
+    time_badge = pg.views.html.controls.Badge(
+        label,
+        tooltip=tooltip,
+        css_classes=['execution-time', css_class],
+        interactive=interactive,
+    )
     if interactive:
       self._time_badge = time_badge
-    title = pg.Html.element(
-        'div',
-        [
-            'ExecutionTrace',
-            time_badge,
-        ],
-        css_classes=['execution-trace-title'],
-    )
-    kwargs.pop('title', None)
-    kwargs['enable_summary_tooltip'] = False
-    kwargs['enable_key_tooltip'] = False
-    return view.summary(
-        self,
-        name=name,
-        title=title,
-        extra_flags=extra_flags,
-        **kwargs
-    )
+    return time_badge
 
   def _html_tree_view_content(self, **kwargs):
     del kwargs
@@ -319,12 +300,7 @@ def _html_tree_view_content(self, **kwargs):
         [self._execution_item_tab(item) for item in self.items],
         tab_position='left'
     )
-    return pg.Html.element(
-        'div',
-        [
-            self._tab_control
-        ]
-    )
+    return self._tab_control.to_html()
 
   def _execution_item_tab(self, item: TracedItem) -> pg.views.html.controls.Tab:
     if isinstance(item, ActionInvocation):
@@ -381,7 +357,8 @@ def _execution_item_label(
       )
     elif isinstance(item, ExecutionTrace):
       return pg.views.html.controls.Label(
-          item.name or 'Phase'
+          item.name or 'Phase',
+          tooltip=f'Execution phase {item.name!r}.'
       )
     else:
       raise ValueError(f'Unsupported item type: {type(item)}')
@@ -420,21 +397,21 @@ def _html_tree_view_css_styles(self) -> list[str]:
           display: inline-block;
         }
         .badge.execution-time {
-          margin-left: 5px;
+          margin-left: 4px;
+          border-radius: 0px;
+        }
+        .execution-time.starting {
+          background-color: ghostwhite;
+          font-weight: normal;
         }
         .execution-time.running {
-          background-color: lavender;
+          background-color: ghostwhite;
           font-weight: normal;
         }
         .execution-time.finished {
           background-color: aliceblue;
           font-weight: bold;
         }
-        .badge.task-in-progress {
-          margin-left: 5px;
-          background-color: azure;
-          font-weight: bold;
-        }
        """
     ]
 
@@ -448,7 +425,7 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
       'The result of the action.'
   ] = None
 
-  result_metadata: Annotated[
+  metadata: Annotated[
       dict[str, Any],
       'The metadata returned by the action.'
   ] = {}
@@ -464,8 +441,7 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
   def _on_bound(self):
     super()._on_bound()
     self._current_phase = self.execution
-    self._result_badge = None
-    self._result_metadata_badge = None
+    self._tab_control = None
 
   @property
   def current_phase(self) -> ExecutionTrace:
@@ -520,31 +496,42 @@ def start(self) -> None:
     """Starts the execution of the action."""
     self.execution.start()
 
-  def end(self, result: Any, result_metadata: dict[str, Any]) -> None:
+  def end(self, result: Any, metadata: dict[str, Any]) -> None:
     """Ends the execution of the action with result and metadata."""
     self.execution.stop()
     self.rebind(
         result=result,
-        result_metadata=result_metadata,
+        metadata=metadata,
         skip_notification=True,
         raise_on_no_change=False
     )
-    if self._result_badge is not None:
-      self._result_badge.update(
-          self._result_badge_label(result),
-          tooltip=self._result_badge_tooltip(result),
-          add_class=['ready'],
-          remove_class=['not-ready'],
-      )
-    if self._result_metadata_badge is not None:
-      result_metadata = dict(result_metadata)
-      result_metadata.pop('session', None)
-      self._result_metadata_badge.update(
-          '{...}',
-          tooltip=self._result_metadata_badge_tooltip(result_metadata),
-          add_class=['ready'],
-          remove_class=['not-ready'],
+    if self._tab_control is not None:
+      if self.metadata:
+        self._tab_control.insert(
+            1,
+            pg.views.html.controls.Tab(
+                'metadata',
+                pg.view(
+                    self.metadata,
+                    collapse_level=None,
+                    enable_summary_tooltip=False
+                ),
+                name='metadata',
+            )
+        )
+      self._tab_control.insert(
+          1,
+          pg.views.html.controls.Tab(
+              'result',
+              pg.view(
+                  self.result,
+                  collapse_level=None,
+                  enable_summary_tooltip=False
+              ),
+              name='result',
+          ),
       )
+      self._tab_control.select(['metadata', 'result'])
 
   #
   # HTML views.
@@ -566,107 +553,77 @@ def _html_tree_view_content(
     interactive = extra_flags.get('interactive', True)
     if (isinstance(self.action, RootAction)
         and self.execution.has_stopped
-        and len(self.execution.items) == 1):
+        and len(self.execution) == 1):
       return view.content(self.execution.items[0], extra_flags=extra_flags)
 
-    def _result_badge():
-      if not self.execution.has_stopped:
-        label = '(n/a)'
-        tooltip = 'Result is not available yet.'
-        css_class = 'not-ready'
-      else:
-        label = self._result_badge_label(self.result)
-        tooltip = self._result_badge_tooltip(self.result)
-        css_class = 'ready'
-      return pg.views.html.controls.Badge(
-          label,
-          tooltip=tooltip,
-          css_classes=['invocation-result', css_class],
-          interactive=interactive,
+    tabs = []
+    if not isinstance(self.action, RootAction):
+      tabs.append(
+          pg.views.html.controls.Tab(
+              'action',
+              view.render(  # pylint: disable=g-long-ternary
+                  self.action,
+                  collapse_level=None,
+                  root_path=self.action.sym_path,
+                  css_classes='invocation-title',
+                  enable_summary_tooltip=False,
+              ),
+              name='action',
+          )
       )
-
-    def _result_metadata_badge():
-      if not self.execution.has_stopped:
-        label = '(n/a)'
-        tooltip = 'Result metadata is not available yet.'
-        css_class = 'not-ready'
-      else:
-        label = '{...}' if self.result_metadata else '(empty)'
-        tooltip = self._result_metadata_badge_tooltip(self.result_metadata)
-        css_class = 'ready'
-      return pg.views.html.controls.Badge(
-          label,
-          tooltip=tooltip,
-          css_classes=['invocation-result-metadata', css_class],
-          interactive=interactive,
+    if self.execution.has_stopped:
+      tabs.append(
+          pg.views.html.controls.Tab(
+              'result',
+              view.render(
+                  self.result,
+                  collapse_level=None,
+                  enable_summary_tooltip=False
+              ),
+              name='result'
+          )
       )
-
-    result_badge = _result_badge()
-    result_metadata_badge = _result_metadata_badge()
-    if interactive:
-      self._result_badge = result_badge
-      self._result_metadata_badge = result_metadata_badge
-
-    return pg.Html.element(
-        'div',
-        [
-            pg.Html.element(
-                'div',
-                [
-                    view.render(
-                        self.usage_summary, extra_flags=dict(as_badge=True)
-                    ),
-                    result_badge,
-                    result_metadata_badge,
-                ],
-                css_classes=['invocation-badge-container'],
-            ),
-            view.render(  # pylint: disable=g-long-ternary
-                self.action,
-                name='action',
-                collapse_level=None,
-                root_path=self.action.sym_path,
-                css_classes='invocation-title',
-                enable_summary_tooltip=False,
-            ) if not isinstance(self.action, RootAction) else None,
-            view.render(self.execution, name='execution'),
-        ]
-    )
-
-  def _result_badge_label(self, result: Any) -> str:
-    label = pg.format(
-        result, python_format=True, verbose=False
-    )
-    if len(label) > 40:
-      if isinstance(result, str):
-        label = label[:40] + '...'
-      else:
-        label = f'{result.__class__.__name__}(...)'
-    return label
-
-  def _result_badge_tooltip(self, result: Any) -> pg.Html:
-    return typing.cast(
-        pg.Html,
-        pg.view(
-            result, name='result',
-            collapse_level=None,
-            enable_summary_tooltip=False,
-            enable_key_tooltip=False,
-        )
-    )
-
-  def _result_metadata_badge_tooltip(
-      self, result_metadata: dict[str, Any]
-  ) -> pg.Html:
-    return typing.cast(
-        pg.Html,
-        pg.view(
-            result_metadata,
-            name='result_metadata',
-            collapse_level=None,
-            enable_summary_tooltip=False,
+      if self.metadata:
+        tabs.append(
+            pg.views.html.controls.Tab(
+                'metadata',
+                view.render(
+                    self.metadata,
+                    collapse_level=None,
+                    enable_summary_tooltip=False
+                ),
+                name='metadata'
+            )
         )
-    )
+    if not self.execution.has_stopped or self.execution.items:
+      execution_title = pg.Html.element(
+          'span',
+          [
+              'execution',
+              self.execution._execution_badge(interactive)  # pylint: disable=protected-access
+          ]
+      )
+      tabs.append(
+          pg.views.html.controls.Tab(
+              execution_title,
+              view.render(self.execution),
+              name='execution',
+          )
+      )
+    tab_control = pg.views.html.controls.TabControl(tabs)
+    # Select the tab following a priority: metadata, result, action, execution.
+    tab_control.select(['metadata', 'result', 'action', 'execution'])
+    if interactive:
+      self._tab_control = tab_control
+    if interactive or self.usage_summary.total.num_requests > 0:
+      return pg.Html.element(
+          'div',
+          [
+              self.usage_summary.to_html(extra_flags=dict(as_badge=True)),
+              tab_control,
+          ],
+      )
+    return tab_control
 
   @classmethod
   def _html_tree_view_css_styles(cls) -> list[str]:
@@ -679,16 +636,6 @@ def _html_tree_view_css_styles(cls) -> list[str]:
         .invocation-badge-container > .label-container {
           margin-right: 3px;
         }
-        .invocation-result.ready {
-          background-color: lightcyan;
-        }
-        .invocation-result-metadata.ready {
-          background-color: lightyellow;
-        }
-        details.pyglove.invocation-title {
-          background-color: aliceblue;
-          border: 0px solid white;
-        }
         """
     ]
 
@@ -722,7 +669,7 @@ def current_action(self) -> ActionInvocation:
   def add_metadata(self, **kwargs: Any) -> None:
     """Adds metadata to the current invocation."""
     with pg.notify_on_change(False):
-      self._current_action.result_metadata.update(kwargs)
+      self._current_action.metadata.update(kwargs)
 
   def phase(self, name: str) -> ContextManager[ExecutionTrace]:
     """Context manager for starting a new execution phase."""
@@ -745,11 +692,11 @@ def track_action(self, action: Action) -> Iterator[ActionInvocation]:
       yield invocation
     finally:
       # Stop the execution of the current action.
-      self._current_action.end(action.result, action.result_metadata)
+      self._current_action.end(action.result, action.metadata)
       self._current_action = parent_action
       if parent_action is self.root:
         parent_action.end(
-            result=action.result, result_metadata=action.result_metadata,
+            result=action.result, metadata=action.metadata,
         )
 
   @contextlib.contextmanager
diff --git a/langfun/core/agentic/action_test.py b/langfun/core/agentic/action_test.py
index f2b544e..97e3eed 100644
--- a/langfun/core/agentic/action_test.py
+++ b/langfun/core/agentic/action_test.py
@@ -73,7 +73,7 @@ def make_additional_query(self, lm):
     self.assertTrue(root.execution.has_stopped)
     self.assertGreater(root.execution.elapse, 0)
     self.assertEqual(root.result, 3)
-    self.assertEqual(root.result_metadata, dict(note='foo'))
+    self.assertEqual(root.metadata, dict(note='foo'))
 
     # The root space should have one action (foo), no queries, and no logs.
     self.assertEqual(len(list(root.actions)), 1)
@@ -109,11 +109,11 @@ def make_additional_query(self, lm):
     self.assertIsInstance(bar_invocation, action_lib.ActionInvocation)
     self.assertIsInstance(bar_invocation.action, Bar)
     self.assertEqual(bar_invocation.result, 2)
-    self.assertEqual(bar_invocation.result_metadata, dict(note='bar'))
+    self.assertEqual(bar_invocation.metadata, dict(note='bar'))
     self.assertEqual(len(bar_invocation.execution.items), 2)
 
     # Save to HTML
-    self.assertIn('invocation-result', session.to_html().content)
+    self.assertIn('result', session.to_html().content)
 
     # Save session to JSON
     json_str = session.to_json_str(save_ref_value=True)
diff --git a/langfun/core/structured/prompting.py b/langfun/core/structured/prompting.py
index 68d7fc0..2cc3896 100644
--- a/langfun/core/structured/prompting.py
+++ b/langfun/core/structured/prompting.py
@@ -463,7 +463,7 @@ def _html_tree_view_content(
                 extra_flags=dict(include_message_metadata=False)
             ),
         ),
-    ], tab_position='top').to_html()
+    ], tab_position='top', selected=2).to_html()
 
   @classmethod
   def _html_tree_view_css_styles(cls) -> list[str]: