diff --git a/langfun/core/agentic/action.py b/langfun/core/agentic/action.py index 98bac89..0fa9607 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=pg.format(self.execution_summary(), verbose=False), 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'], ) @@ -263,6 +258,20 @@ def usage_summary(self) -> lf.UsageSummary: """Returns the usage summary of the action.""" return self._usage_summary + def execution_summary(self) -> dict[str, Any]: + """Execution summary string.""" + return pg.Dict( + num_queries=len(self.queries), + execution_breakdown=[ + dict( + action=action.action.__class__.__name__, + usage=action.usage_summary.total, + execution_time=action.execution.elapse, + ) + for action in self.actions + ] + ) + # # HTML views. # @@ -274,57 +283,47 @@ 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 = pg.format(self.execution_summary(), verbose=False) + 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): + def _html_tree_view_content( + self, + *, + extra_flags: dict[str, Any] | None = None, + **kwargs + ): del kwargs - self._tab_control = pg.views.html.controls.TabControl( - [self._execution_item_tab(item) for item in self.items], - tab_position='left' - ) - return pg.Html.element( - 'div', - [ - self._tab_control - ] - ) + extra_flags = extra_flags or {} + interactive = extra_flags.get('interactive', True) + if interactive or self.items: + self._tab_control = pg.views.html.controls.TabControl( + [self._execution_item_tab(item) for item in self.items], + tab_position='left' + ) + return self._tab_control.to_html() + return '(no tracked items)' def _execution_item_tab(self, item: TracedItem) -> pg.views.html.controls.Tab: if isinstance(item, ActionInvocation): @@ -381,7 +380,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 +420,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 +448,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 +464,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 +519,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,128 +576,88 @@ 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, + 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' + ) ) + if self.metadata: + tabs.append( + pg.views.html.controls.Tab( + 'metadata', + view.render( + self.metadata, + collapse_level=None, + enable_summary_tooltip=False + ), + name='metadata' + ) + ) - 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', - [ + tabs.append( + pg.views.html.controls.Tab( pg.Html.element( - 'div', + 'span', [ - view.render( - self.usage_summary, extra_flags=dict(as_badge=True) + 'execution', + self.execution._execution_badge(interactive), # pylint: disable=protected-access + ( + self.usage_summary.to_html( # pylint: disable=g-long-ternary + extra_flags=dict(as_badge=True) + ) + if (interactive + or self.usage_summary.total.num_requests > 0) + else None ), - result_badge, - result_metadata_badge, ], - css_classes=['invocation-badge-container'], + css_classes=['execution-tab-title'] ), - 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, + view.render(self.execution, extra_flags=extra_flags), + 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 + return tab_control @classmethod def _html_tree_view_css_styles(cls) -> list[str]: return super()._html_tree_view_css_styles() + [ """ - .invocation-badge-container { - display: flex; - padding-bottom: 5px; - } - .invocation-badge-container > .label-container { - margin-right: 3px; - } - .invocation-result.ready { - background-color: lightcyan; + .execution-tab-title { + text-align: left; } - .invocation-result-metadata.ready { - background-color: lightyellow; - } - details.pyglove.invocation-title { - background-color: aliceblue; - border: 0px solid white; + .execution-tab-title .usage-summary.label { + border-radius: 0px; + font-weight: normal; + color: #AAA; } """ ] @@ -722,7 +692,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 +715,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..5005992 100644 --- a/langfun/core/structured/prompting.py +++ b/langfun/core/structured/prompting.py @@ -15,6 +15,7 @@ import contextlib import functools +import time from typing import Annotated, Any, Callable, Iterator, Type, Union import langfun.core as lf @@ -221,6 +222,7 @@ class Flight(pg.Object): query_input = schema_lib.mark_missing(prompt) with lf.track_usages() as usage_summary: + start_time = time.time() if schema in (None, str): # Query with natural language output. output_message = lf.LangFunc.from_value(query_input, **kwargs)( @@ -250,6 +252,7 @@ class Flight(pg.Object): cache_seed=cache_seed, skip_lm=skip_lm, ) + end_time = time.time() def _result(message: lf.Message): return message.text if schema in (None, str) else message.result @@ -268,6 +271,8 @@ def _result(message: lf.Message): examples=pg.Ref(examples) if examples else [], lm_response=lf.AIMessage(output_message.text), usage_summary=usage_summary, + start_time=start_time, + end_time=end_time, ) for i, (tracker, include_child_scopes) in enumerate(trackers): if i == 0 or include_child_scopes: @@ -384,6 +389,14 @@ class QueryInvocation(pg.Object, pg.views.HtmlTreeView.Extension): lf.UsageSummary, 'Usage summary for `lf.query`.' ] + start_time: Annotated[ + float, + 'Start time of query.' + ] + end_time: Annotated[ + float, + 'End time of query.' + ] @functools.cached_property def lm_request(self) -> lf.Message: @@ -393,6 +406,11 @@ def lm_request(self) -> lf.Message: def output(self) -> Any: return query_output(self.lm_response, self.schema) + @property + def elapse(self) -> float: + """Returns query elapse in seconds.""" + return self.end_time - self.start_time + def _on_bound(self): super()._on_bound() self.__dict__.pop('lm_request', None) @@ -404,6 +422,8 @@ def _html_tree_view_summary( view: pg.views.HtmlTreeView, **kwargs: Any ) -> pg.Html | None: + kwargs.pop('title', None) + kwargs.pop('enable_summary_tooltip', None) return view.summary( value=self, title=pg.Html.element( @@ -423,11 +443,16 @@ def _html_tree_view_summary( ), css_classes=['query-invocation-lm'] ), + pg.views.html.controls.Badge( + f'{int(self.elapse)} seconds', + css_classes=['query-invocation-time'] + ), self.usage_summary.to_html(extra_flags=dict(as_badge=True)) ], css_classes=['query-invocation-title'] ), - enable_summary_tooltip=False + enable_summary_tooltip=False, + **kwargs ) def _html_tree_view_content( @@ -441,14 +466,14 @@ def _html_tree_view_content( 'input', pg.view(self.input, collapse_level=None), ), - pg.views.html.controls.Tab( - 'schema', - pg.view(self.schema), - ), pg.views.html.controls.Tab( 'output', pg.view(self.output, collapse_level=None), ), + pg.views.html.controls.Tab( + 'schema', + pg.view(self.schema), + ), pg.views.html.controls.Tab( 'lm_request', pg.view( @@ -463,24 +488,34 @@ def _html_tree_view_content( extra_flags=dict(include_message_metadata=False) ), ), - ], tab_position='top').to_html() + ], tab_position='top', selected=1).to_html() @classmethod def _html_tree_view_css_styles(cls) -> list[str]: return super()._html_tree_view_css_styles() + [ """ .query-invocation-title { - display: inline-block; - font-weight: normal; + display: inline-block; + font-weight: normal; } .query-invocation-type-name { - font-style: italic; - color: #888; + color: #888; } .query-invocation-lm.badge { - margin-left: 5px; - margin-right: 5px; - background-color: #fff0d6; + margin-left: 5px; + margin-right: 5px; + color: white; + background-color: mediumslateblue; + } + .query-invocation-time.badge { + margin-left: 5px; + border-radius: 0px; + font-weight: bold; + background-color: aliceblue; + } + .query-invocation-title .usage-summary.label { + border-radius: 0px; + color: #AAA; } """ ] diff --git a/langfun/core/structured/prompting_test.py b/langfun/core/structured/prompting_test.py index 5a93442..2e54b1d 100644 --- a/langfun/core/structured/prompting_test.py +++ b/langfun/core/structured/prompting_test.py @@ -996,6 +996,7 @@ def test_include_child_scopes(self): self.assertEqual(queries[1].schema.spec.cls, Activity) self.assertTrue(pg.eq(queries[1].output, Activity(description='hi'))) self.assertIs(queries[1].lm, lm) + self.assertGreater(queries[0].elapse, 0) self.assertGreater(queries[0].usage_summary.total.total_tokens, 0) self.assertGreater(queries[1].usage_summary.total.total_tokens, 0) diff --git a/langfun/core/template.py b/langfun/core/template.py index 412b8f9..e1b2f94 100644 --- a/langfun/core/template.py +++ b/langfun/core/template.py @@ -552,38 +552,33 @@ def render_template_str(): ) def render_fields(): - return pg.Html.element( - 'fieldset', - [ - pg.Html.element('legend', ['Template Variables']), - view.complex_value( - {k: v for k, v in self.sym_items()}, - name='fields', - root_path=root_path, - parent=self, - exclude_keys=['template_str', 'clean'], - collapse_level=max( - collapse_template_vars_level, collapse_level - ) if collapse_level is not None else None, - extra_flags=extra_flags, - debug=debug, - **view.get_passthrough_kwargs( - remove=['exclude_keys'], - **kwargs, - ) - ), - ], - css_classes=['template-fields'], + return view.complex_value( + {k: v for k, v in self.sym_items()}, + name='fields', + root_path=root_path, + parent=self, + exclude_keys=['template_str', 'clean'], + collapse_level=max( + collapse_template_vars_level, collapse_level + ) if collapse_level is not None else None, + extra_flags=extra_flags, + debug=debug, + **view.get_passthrough_kwargs( + remove=['exclude_keys'], + **kwargs, + ) ) - return pg.Html.element( - 'div', - [ + return pg.views.html.controls.TabControl([ + pg.views.html.controls.Tab( + 'template_str', render_template_str(), + ), + pg.views.html.controls.Tab( + 'variables', render_fields(), - ], - css_classes=['complex_value'], - ) + ), + ], selected=1) @classmethod def _html_tree_view_css_styles(cls) -> list[str]: @@ -601,15 +596,6 @@ def _html_tree_view_css_styles(cls) -> list[str]: background-color: #EEE; color: #cc2986; } - .template-fields { - margin: 0px 0px 5px 0px; - border: 1px solid #EEE; - padding: 5px; - } - .template-fields > legend { - font-size: 0.8em; - margin: 5px 0px 5px 0px; - } """ ] diff --git a/langfun/core/template_test.py b/langfun/core/template_test.py index c141df1..a3684ce 100644 --- a/langfun/core/template_test.py +++ b/langfun/core/template_test.py @@ -555,13 +555,6 @@ def on_event(self, event: TemplateRenderEvent): class HtmlTest(unittest.TestCase): - 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): class Foo(Template): @@ -594,36 +587,23 @@ class Bar(Template): background-color: #EEE; color: #cc2986; } - .template-fields { - margin: 0px 0px 5px 0px; - border: 1px solid #EEE; - padding: 5px; - } - .template-fields > legend { - font-size: 0.8em; - margin: 5px 0px 5px 0px; - } """ ), Foo(x=1, y=2).to_html().style_section, ) - self.assert_html_content( + self.assertIn( + 'template-str', Foo(x=Bar('{{y}} + {{z}}'), y=1).to_html( enable_summary_tooltip=False, - ), - """ -
Foo(...)
{{x}} + {{y}} = ?
Template Variables
xx
Bar(...)
{{y}} + {{z}}
Template Variables
yx.y
ContextualAttribute(...)
1
zx.z
ContextualAttribute(...)
(not available)
yy
int
1
- """ + ).content, ) - self.assert_html_content( + self.assertIn( + 'template-str', Foo(x=Bar('{{y}} + {{z}}'), y=1).to_html( enable_summary_tooltip=False, collapse_level=0, key_style='label', - ), - """ -
Foo(...)
{{x}} + {{y}} = ?
Template Variables
xx
Bar(...)
{{y}} + {{z}}
Template Variables
yx.y
ContextualAttribute(...)
1
zx.z
ContextualAttribute(...)
(not available)
yy1
- """ + ).content, )