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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
langfun/core/component.py CHANGED
@@ -286,6 +286,54 @@ class ContextualAttribute(pg.symbolic.ValueFromParentChain):
286
286
  else:
287
287
  return pg.MISSING_VALUE
288
288
 
289
+ def _html_tree_view_content(
290
+ self,
291
+ *,
292
+ view: pg.views.HtmlTreeView,
293
+ parent: Any,
294
+ root_path: pg.KeyPath,
295
+ **kwargs,
296
+ ) -> pg.Html:
297
+ inferred_value = pg.MISSING_VALUE
298
+ if isinstance(parent, pg.Symbolic) and root_path:
299
+ inferred_value = parent.sym_inferred(root_path.key, pg.MISSING_VALUE)
300
+
301
+ if inferred_value is not pg.MISSING_VALUE:
302
+ kwargs.pop('name', None)
303
+ return view.render(
304
+ inferred_value, parent=self, root_path=root_path + '<inferred>',
305
+ **view.get_passthrough_kwargs(**kwargs)
306
+ )
307
+ return pg.Html.element(
308
+ 'div',
309
+ [
310
+ '(not available)',
311
+ ],
312
+ css_classes=['unavailable-contextual'],
313
+ )
314
+
315
+ def _html_tree_view_config(self) -> dict[str, Any]:
316
+ return pg.views.HtmlTreeView.get_kwargs(
317
+ super()._html_tree_view_config(),
318
+ dict(
319
+ collapse_level=1,
320
+ )
321
+ )
322
+
323
+ @classmethod
324
+ def _html_tree_view_css_styles(cls) -> list[str]:
325
+ return super()._html_tree_view_css_styles() + [
326
+ """
327
+ .contextual-attribute {
328
+ color: purple;
329
+ }
330
+ .unavailable-contextual {
331
+ color: gray;
332
+ font-style: italic;
333
+ }
334
+ """
335
+ ]
336
+
289
337
 
290
338
  # NOTE(daiyip): Returning Any instead of `lf.ContextualAttribute` to avoid
291
339
  # pytype check error as `contextual()` can be assigned to any type.
@@ -13,6 +13,8 @@
13
13
  # limitations under the License.
14
14
  """Contextual component and app test."""
15
15
 
16
+ import inspect
17
+ from typing import Any
16
18
  import unittest
17
19
  import weakref
18
20
 
@@ -297,6 +299,52 @@ class ContextualAttributeTest(unittest.TestCase):
297
299
  self.assertEqual(c.z, 3)
298
300
  self.assertEqual(b.z, 3)
299
301
 
302
+ def test_to_html(self):
303
+ class A(lf.Component):
304
+ x: int = 1
305
+ y: int = lf.contextual()
306
+
307
+ def assert_content(html, expected):
308
+ expected = inspect.cleandoc(expected).strip()
309
+ actual = html.content.strip()
310
+ if actual != expected:
311
+ print(actual)
312
+ self.assertEqual(actual.strip(), expected)
313
+
314
+ self.assertIn(
315
+ inspect.cleandoc(
316
+ """
317
+ .contextual-attribute {
318
+ color: purple;
319
+ }
320
+ .unavailable-contextual {
321
+ color: gray;
322
+ font-style: italic;
323
+ }
324
+ """
325
+ ),
326
+ A().to_html().style_section,
327
+ )
328
+
329
+ assert_content(
330
+ A().to_html(enable_summary_tooltip=False),
331
+ """
332
+ <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>
333
+ """
334
+ )
335
+
336
+ class B(lf.Component):
337
+ z: Any
338
+ y: int = 2
339
+
340
+ b = B(A())
341
+ assert_content(
342
+ b.z.to_html(enable_summary_tooltip=False),
343
+ """
344
+ <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>
345
+ """
346
+ )
347
+
300
348
 
301
349
  if __name__ == '__main__':
302
350
  unittest.main()
langfun/core/logging.py CHANGED
@@ -15,8 +15,9 @@
15
15
 
16
16
  import contextlib
17
17
  import datetime
18
+ import functools
18
19
  import typing
19
- from typing import Any, Iterator, Literal, Sequence
20
+ from typing import Any, Iterator, Literal
20
21
 
21
22
  from langfun.core import component
22
23
  from langfun.core import console
@@ -57,11 +58,11 @@ class LogEntry(pg.Object):
57
58
  self,
58
59
  view: pg.views.HtmlTreeView,
59
60
  title: str | pg.Html | None = None,
60
- max_str_len_for_summary: int = pg.View.PresetArgValue(80), # pytype: disable=annotation-type-mismatch
61
+ max_summary_len_for_str: int = 80,
61
62
  **kwargs
62
63
  ) -> str:
63
- if len(self.message) > max_str_len_for_summary:
64
- message = self.message[:max_str_len_for_summary] + '...'
64
+ if len(self.message) > max_summary_len_for_str:
65
+ message = self.message[:max_summary_len_for_str] + '...'
65
66
  else:
66
67
  message = self.message
67
68
 
@@ -69,18 +70,18 @@ class LogEntry(pg.Object):
69
70
  pg.Html.element(
70
71
  'span',
71
72
  [self.time.strftime('%H:%M:%S')],
72
- css_class=['log-time']
73
+ css_classes=['log-time']
73
74
  ),
74
75
  pg.Html.element(
75
76
  'span',
76
77
  [pg.Html.escape(message)],
77
- css_class=['log-summary'],
78
+ css_classes=['log-summary'],
78
79
  ),
79
80
  )
80
81
  return view.summary(
81
82
  self,
82
83
  title=title or s,
83
- max_str_len_for_summary=max_str_len_for_summary,
84
+ max_summary_len_for_str=max_summary_len_for_str,
84
85
  **kwargs,
85
86
  )
86
87
 
@@ -89,43 +90,45 @@ class LogEntry(pg.Object):
89
90
  self,
90
91
  view: pg.views.HtmlTreeView,
91
92
  root_path: pg.KeyPath,
92
- collapse_log_metadata_level: int | None = pg.View.PresetArgValue(0),
93
- max_str_len_for_summary: int = pg.View.PresetArgValue(80),
94
- collapse_level: int | None = pg.View.PresetArgValue(1),
93
+ max_summary_len_for_str: int = 80,
94
+ collapse_level: int | None = 1,
95
+ extra_flags: dict[str, Any] | None = None,
95
96
  **kwargs
96
97
  ) -> pg.Html:
97
98
  # pytype: enable=annotation-type-mismatch
99
+ extra_flags = extra_flags if extra_flags is not None else {}
100
+ collapse_log_metadata_level: int | None = extra_flags.get(
101
+ 'collapse_log_metadata_level', None
102
+ )
98
103
  def render_message_text():
99
- if len(self.message) < max_str_len_for_summary:
104
+ if len(self.message) < max_summary_len_for_str:
100
105
  return None
101
106
  return pg.Html.element(
102
107
  'span',
103
108
  [pg.Html.escape(self.message)],
104
- css_class=['log-text'],
109
+ css_classes=['log-text'],
105
110
  )
106
111
 
107
112
  def render_metadata():
108
113
  if not self.metadata:
109
114
  return None
110
- child_path = root_path + 'metadata'
111
115
  return pg.Html.element(
112
116
  'div',
113
117
  [
114
118
  view.render(
115
119
  self.metadata,
116
120
  name='metadata',
117
- root_path=child_path,
121
+ root_path=root_path + 'metadata',
118
122
  parent=self,
119
- collapse_level=(
120
- view.max_collapse_level(
121
- collapse_level,
122
- collapse_log_metadata_level,
123
- child_path
124
- )
125
- )
123
+ collapse_level=view.get_collapse_level(
124
+ (collapse_level, -1), collapse_log_metadata_level,
125
+ ),
126
+ max_summary_len_for_str=max_summary_len_for_str,
127
+ extra_flags=extra_flags,
128
+ **view.get_passthrough_kwargs(**kwargs),
126
129
  )
127
130
  ],
128
- css_class=['log-metadata'],
131
+ css_classes=['log-metadata'],
129
132
  )
130
133
 
131
134
  return pg.Html.element(
@@ -134,12 +137,23 @@ class LogEntry(pg.Object):
134
137
  render_message_text(),
135
138
  render_metadata(),
136
139
  ],
137
- css_class=['complex_value'],
140
+ css_classes=['complex_value'],
141
+ )
142
+
143
+ def _html_tree_view_config(self) -> dict[str, Any]:
144
+ return pg.views.HtmlTreeView.get_kwargs(
145
+ super()._html_tree_view_config(),
146
+ dict(
147
+ css_classes=[f'log-{self.level}'],
148
+ )
138
149
  )
139
150
 
140
- def _html_style(self) -> list[str]:
141
- return super()._html_style() + [
151
+ @classmethod
152
+ @functools.cache
153
+ def _html_tree_view_css_styles(cls) -> list[str]:
154
+ return super()._html_tree_view_css_styles() + [
142
155
  """
156
+ /* Langfun LogEntry styles. */
143
157
  .log-time {
144
158
  color: #222;
145
159
  font-size: 12px;
@@ -203,9 +217,6 @@ class LogEntry(pg.Object):
203
217
  """
204
218
  ]
205
219
 
206
- def _html_element_class(self) -> Sequence[str] | None:
207
- return super()._html_element_class() + [f'log-{self.level}']
208
-
209
220
 
210
221
  def log(level: LogLevel,
211
222
  message: str,
@@ -69,20 +69,36 @@ class LoggingTest(unittest.TestCase):
69
69
  time=time, metadata={}
70
70
  ).to_html(enable_summary_tooltip=False),
71
71
  """
72
- <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>
72
+ <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>
73
73
  """
74
74
  )
75
75
  self.assert_html_content(
76
76
  logging.LogEntry(
77
77
  level='error', message='This is a longer message: 5 + 2 > 3',
78
- time=time, metadata=dict(x=1, y=2)
78
+ time=time, metadata=dict(x=dict(z=1), y=2)
79
79
  ).to_html(
80
- max_str_len_for_summary=10,
80
+ extra_flags=dict(
81
+ collapse_log_metadata_level=1,
82
+ ),
83
+ max_summary_len_for_str=10,
81
84
  enable_summary_tooltip=False,
82
- collapse_log_metadata_level=1
83
85
  ),
84
86
  """
85
- <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>
87
+ <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>
88
+ """
89
+ )
90
+ self.assert_html_content(
91
+ logging.LogEntry(
92
+ level='error', message='This is a longer message: 5 + 2 > 3',
93
+ time=time, metadata=dict(x=dict(z=1), y=2)
94
+ ).to_html(
95
+ extra_flags=dict(
96
+ max_summary_len_for_str=10,
97
+ ),
98
+ enable_summary_tooltip=False,
99
+ ),
100
+ """
101
+ <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>
86
102
  """
87
103
  )
88
104
 
langfun/core/message.py CHANGED
@@ -14,8 +14,9 @@
14
14
  """Messages that are exchanged between users and agents."""
15
15
 
16
16
  import contextlib
17
+ import functools
17
18
  import io
18
- from typing import Annotated, Any, Optional, Sequence, Union
19
+ from typing import Annotated, Any, Optional, Union
19
20
 
20
21
  from langfun.core import modality
21
22
  from langfun.core import natural_language
@@ -506,53 +507,70 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
506
507
  v = self.metadata[key]
507
508
  return v.value if isinstance(v, pg.Ref) else v
508
509
 
509
- # pytype: disable=annotation-type-mismatch
510
510
  def _html_tree_view_content(
511
511
  self,
512
512
  *,
513
513
  view: pg.views.HtmlTreeView,
514
514
  root_path: pg.KeyPath,
515
- source_tag: str | Sequence[str] | None = pg.View.PresetArgValue(
516
- ('lm-input', 'lm-output')
517
- ),
518
- include_message_metadata: bool = pg.View.PresetArgValue(True),
519
- collapse_modalities_in_text: bool = pg.View.PresetArgValue(True),
520
- collapse_llm_usage: bool = pg.View.PresetArgValue(False),
521
- collapse_message_result_level: int | None = pg.View.PresetArgValue(1),
522
- collapse_message_metadata_level: int | None = pg.View.PresetArgValue(0),
523
- collapse_source_message_level: int | None = pg.View.PresetArgValue(1),
524
- collapse_level: int | None = pg.View.PresetArgValue(1),
515
+ collapse_level: int | None = None,
516
+ extra_flags: dict[str, Any] | None = None,
525
517
  **kwargs,
526
518
  ) -> pg.Html:
527
- # pytype: enable=annotation-type-mismatch
528
519
  """Returns the HTML representation of the message.
529
520
 
530
521
  Args:
531
522
  view: The HTML tree view.
532
523
  root_path: The root path of the message.
533
- source_tag: tags to filter source messages. If None, the entire
534
- source chain will be included.
535
- include_message_metadata: Whether to include the metadata of the message.
536
- collapse_modalities_in_text: Whether to collapse the modalities in the
537
- message text.
538
- collapse_llm_usage: Whether to collapse the usage in the message.
539
- collapse_message_result_level: The level to collapse the result in the
540
- message.
541
- collapse_message_metadata_level: The level to collapse the metadata in the
542
- message.
543
- collapse_source_message_level: The level to collapse the source in the
544
- message.
545
524
  collapse_level: The global collapse level.
546
- **kwargs: Other keyword arguments.
525
+ extra_flags: Extra flags to control the rendering.
526
+ - source_tag: tags to filter source messages. If None, the entire
527
+ source chain will be included.
528
+ - include_message_metadata: Whether to include the metadata of the
529
+ message.
530
+ - collapse_modalities_in_text: Whether to collapse the modalities in the
531
+ message text.
532
+ - collapse_llm_usage: Whether to collapse the usage in the message.
533
+ - collapse_message_result_level: The level to collapse the result in the
534
+ message.
535
+ - collapse_message_metadata_level: The level to collapse the metadata in
536
+ the message.
537
+ - collapse_source_message_level: The level to collapse the source in the
538
+ message.
539
+ - collapse_level: The global collapse level.
540
+ **kwargs: Omitted keyword arguments.
547
541
 
548
542
  Returns:
549
543
  The HTML representation of the message content.
550
544
  """
545
+ extra_flags = extra_flags if extra_flags is not None else {}
546
+
547
+ include_message_metadata: bool = extra_flags.get(
548
+ 'include_message_metadata', True
549
+ )
550
+ source_tag: str | tuple[str, ...] | None = extra_flags.get(
551
+ 'source_tag', ('lm-input', 'lm-output')
552
+ )
553
+ collapse_modalities_in_text: bool = extra_flags.get(
554
+ 'collapse_modalities_in_text', True
555
+ )
556
+ collapse_llm_usage: bool = extra_flags.get(
557
+ 'collapse_llm_usage', False
558
+ )
559
+ collapse_message_result_level: int | None = extra_flags.get(
560
+ 'collapse_message_result_level', 1
561
+ )
562
+ collapse_message_metadata_level: int | None = extra_flags.get(
563
+ 'collapse_message_metadata_level', 1
564
+ )
565
+ collapse_source_message_level: int | None = extra_flags.get(
566
+ 'collapse_source_message_level', 1
567
+ )
568
+ passthrough_kwargs = view.get_passthrough_kwargs(**kwargs)
551
569
  def render_tags():
552
570
  return pg.Html.element(
553
571
  'div',
554
572
  [pg.Html.element('span', [tag]) for tag in self.tags],
555
- css_class=['message-tags'],
573
+ css_classes=['message-tags'],
556
574
  )
557
575
 
558
576
  def render_message_text():
@@ -573,12 +591,16 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
573
591
  chunk,
574
592
  name=chunk.referred_name,
575
593
  root_path=child_path,
576
- collapse_level=child_path.depth + (
594
+ collapse_level=(
577
595
  0 if collapse_modalities_in_text else 1
578
- )
596
+ ),
597
+ extra_flags=dict(
598
+ display_modality_when_hover=True,
599
+ ),
600
+ **passthrough_kwargs,
579
601
  )
580
602
  ],
581
- css_class=['modality-in-text'],
603
+ css_classes=['modality-in-text'],
582
604
  )
583
605
  )
584
606
  referred_chunks[chunk.referred_name] = chunk
@@ -596,14 +618,15 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
596
618
  self.result,
597
619
  name='result',
598
620
  root_path=child_path,
599
- collapse_level=view.max_collapse_level(
600
- collapse_level,
621
+ collapse_level=view.get_collapse_level(
622
+ (collapse_level, -1),
601
623
  collapse_message_result_level,
602
- child_path,
603
- )
624
+ ),
625
+ extra_flags=extra_flags,
626
+ **passthrough_kwargs,
604
627
  )
605
628
  ],
606
- css_class=['message-result'],
629
+ css_classes=['message-result'],
607
630
  )
608
631
 
609
632
  def render_usage():
@@ -616,15 +639,19 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
616
639
  view.render(
617
640
  self.usage,
618
641
  name='llm usage',
642
+ key_style='label',
619
643
  root_path=child_path,
620
- collapse_level=view.max_collapse_level(
621
- collapse_level,
644
+ collapse_level=view.get_collapse_level(
645
+ (collapse_level, -1),
622
646
  0 if collapse_llm_usage else 1,
623
- child_path,
624
- )
647
+ ),
648
+ extra_flags=extra_flags,
649
+ **view.get_passthrough_kwargs(
650
+ remove=['key_style'], **kwargs
651
+ ),
625
652
  )
626
653
  ],
627
- css_class=['message-usage'],
654
+ css_classes=['message-usage'],
628
655
  )
629
656
 
630
657
  def render_source_message():
@@ -635,21 +662,22 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
635
662
  source = source.source
636
663
  if source is not None:
637
664
  child_path = root_path + 'source'
665
+ child_extra_flags = extra_flags.copy()
666
+ child_extra_flags['collapse_source_message_level'] = (
667
+ view.get_collapse_level(
668
+ (collapse_source_message_level, -1), 0,
669
+ )
670
+ )
638
671
  return view.render(
639
672
  self.source,
640
673
  name='source',
641
674
  root_path=child_path,
642
- include_metadata=include_message_metadata,
643
- collapse_level=view.max_collapse_level(
644
- collapse_level,
675
+ collapse_level=view.get_collapse_level(
676
+ (collapse_level, -1),
645
677
  collapse_source_message_level,
646
- child_path
647
678
  ),
648
- collapse_source_level=max(0, collapse_source_message_level - 1),
649
- collapse_modalities=collapse_modalities_in_text,
650
- collapse_usage=collapse_llm_usage,
651
- collapse_metadata_level=collapse_message_metadata_level,
652
- collapse_result_level=collapse_message_result_level,
679
+ extra_flags=child_extra_flags,
680
+ **passthrough_kwargs,
653
681
  )
654
682
  return None
655
683
 
@@ -662,17 +690,20 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
662
690
  [
663
691
  view.render(
664
692
  self.metadata,
665
- css_class=['message-metadata'],
693
+ css_classes=['message-metadata'],
694
+ exclude_keys=['usage', 'result'],
666
695
  name='metadata',
667
696
  root_path=child_path,
668
- collapse_level=view.max_collapse_level(
669
- collapse_level,
697
+ collapse_level=view.get_collapse_level(
698
+ (collapse_level, -1),
670
699
  collapse_message_metadata_level,
671
- child_path,
672
- )
700
+ ),
701
+ **view.get_passthrough_kwargs(
702
+ remove=['exclude_keys'], **kwargs
703
+ ),
673
704
  )
674
705
  ],
675
- css_class=['message-metadata'],
706
+ css_classes=['message-metadata'],
676
707
  )
677
708
 
678
709
  return pg.Html.element(
@@ -685,18 +716,30 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
685
716
  render_metadata(),
686
717
  render_source_message(),
687
718
  ],
688
- css_class=['complex_value'],
719
+ css_classes=['complex_value'],
720
+ )
721
+
722
+ @classmethod
723
+ @functools.cache
724
+ def _html_tree_view_config(cls) -> dict[str, Any]:
725
+ return pg.views.HtmlTreeView.get_kwargs(
726
+ super()._html_tree_view_config(),
727
+ dict(
728
+ css_classes=['lf-message'],
729
+ )
689
730
  )
690
731
 
691
- def _html_style(self) -> list[str]:
692
- return super()._html_style() + [
732
+ @classmethod
733
+ @functools.cache
734
+ def _html_tree_view_css_styles(cls) -> list[str]:
735
+ return super()._html_tree_view_css_styles() + [
693
736
  """
694
737
  /* Langfun Message styles.*/
695
738
  [class^="message-"] > details {
696
739
  margin: 0px 0px 5px 0px;
697
740
  border: 1px solid #EEE;
698
741
  }
699
- details.lf-message > summary > .summary_title::after {
742
+ .lf-message.summary-title::after {
700
743
  content: ' 💬';
701
744
  }
702
745
  details.pyglove.ai-message {
@@ -739,12 +782,12 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
739
782
  margin: 0px 5px 0px 5px;
740
783
  }
741
784
  .message-result {
742
- color: purple;
785
+ color: dodgerblue;
743
786
  }
744
787
  .message-usage {
745
788
  color: orange;
746
789
  }
747
- .message-usage .object_key.str {
790
+ .message-usage .object-key.str {
748
791
  border: 1px solid orange;
749
792
  background-color: orange;
750
793
  color: white;
@@ -752,9 +795,6 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
752
795
  """
753
796
  ]
754
797
 
755
- def _html_element_class(self) -> list[str]:
756
- return super()._html_element_class() + ['lf-message']
757
-
758
798
 
759
799
  #
760
800
  # Messages of different roles.