langfun 0.1.2.dev202410100804__py3-none-any.whl → 0.1.2.dev202410120803__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.
Files changed (42) hide show
  1. langfun/core/__init__.py +1 -0
  2. langfun/core/eval/base_test.py +1 -0
  3. langfun/core/langfunc_test.py +2 -2
  4. langfun/core/language_model.py +140 -24
  5. langfun/core/language_model_test.py +166 -36
  6. langfun/core/llms/__init__.py +8 -1
  7. langfun/core/llms/anthropic.py +72 -7
  8. langfun/core/llms/cache/in_memory_test.py +3 -2
  9. langfun/core/llms/fake_test.py +7 -0
  10. langfun/core/llms/groq.py +154 -6
  11. langfun/core/llms/openai.py +300 -42
  12. langfun/core/llms/openai_test.py +35 -8
  13. langfun/core/llms/vertexai.py +121 -16
  14. langfun/core/logging.py +150 -43
  15. langfun/core/logging_test.py +33 -0
  16. langfun/core/message.py +249 -70
  17. langfun/core/message_test.py +70 -45
  18. langfun/core/modalities/audio.py +1 -1
  19. langfun/core/modalities/audio_test.py +1 -1
  20. langfun/core/modalities/image.py +1 -1
  21. langfun/core/modalities/image_test.py +9 -3
  22. langfun/core/modalities/mime.py +39 -3
  23. langfun/core/modalities/mime_test.py +39 -0
  24. langfun/core/modalities/ms_office.py +2 -5
  25. langfun/core/modalities/ms_office_test.py +1 -1
  26. langfun/core/modalities/pdf_test.py +1 -1
  27. langfun/core/modalities/video.py +1 -1
  28. langfun/core/modalities/video_test.py +2 -2
  29. langfun/core/structured/completion_test.py +1 -0
  30. langfun/core/structured/mapping.py +38 -0
  31. langfun/core/structured/mapping_test.py +55 -0
  32. langfun/core/structured/parsing_test.py +2 -1
  33. langfun/core/structured/prompting_test.py +1 -0
  34. langfun/core/structured/schema.py +34 -0
  35. langfun/core/template.py +110 -1
  36. langfun/core/template_test.py +37 -0
  37. langfun/core/templates/selfplay_test.py +4 -2
  38. {langfun-0.1.2.dev202410100804.dist-info → langfun-0.1.2.dev202410120803.dist-info}/METADATA +1 -1
  39. {langfun-0.1.2.dev202410100804.dist-info → langfun-0.1.2.dev202410120803.dist-info}/RECORD +42 -42
  40. {langfun-0.1.2.dev202410100804.dist-info → langfun-0.1.2.dev202410120803.dist-info}/LICENSE +0 -0
  41. {langfun-0.1.2.dev202410100804.dist-info → langfun-0.1.2.dev202410120803.dist-info}/WHEEL +0 -0
  42. {langfun-0.1.2.dev202410100804.dist-info → langfun-0.1.2.dev202410120803.dist-info}/top_level.txt +0 -0
langfun/core/message.py CHANGED
@@ -14,13 +14,11 @@
14
14
  """Messages that are exchanged between users and agents."""
15
15
 
16
16
  import contextlib
17
- import html
18
17
  import io
19
- from typing import Annotated, Any, Optional, Union
18
+ from typing import Annotated, Any, Optional, Sequence, Union
20
19
 
21
20
  from langfun.core import modality
22
21
  from langfun.core import natural_language
23
- from langfun.core import repr_utils
24
22
  import pyglove as pg
25
23
 
26
24
 
@@ -406,6 +404,11 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
406
404
  with pg.notify_on_change(False):
407
405
  self.tags.append(tag)
408
406
 
407
+ def has_tag(self, tag: str | tuple[str, ...]) -> bool:
408
+ if isinstance(tag, str):
409
+ return tag in self.tags
410
+ return any(t in self.tags for t in tag)
411
+
409
412
  #
410
413
  # Message source chain.
411
414
  #
@@ -503,79 +506,255 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
503
506
  v = self.metadata[key]
504
507
  return v.value if isinstance(v, pg.Ref) else v
505
508
 
506
- def _repr_html_(self):
507
- return self.to_html().content
508
-
509
- def to_html(
509
+ # pytype: disable=annotation-type-mismatch
510
+ def _html_tree_view_content(
510
511
  self,
511
- include_message_type: bool = True
512
- ) -> repr_utils.Html:
513
- """Returns the HTML representation of the message."""
514
- s = io.StringIO()
515
- s.write('<div style="padding:0px 10px 0px 10px;">')
516
- # Title bar.
517
- if include_message_type:
518
- s.write(
519
- repr_utils.html_round_text(
520
- self.__class__.__name__,
521
- text_color='white',
522
- background_color=self._text_color(),
512
+ *,
513
+ view: pg.views.HtmlTreeView,
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),
525
+ **kwargs,
526
+ ) -> pg.Html:
527
+ # pytype: enable=annotation-type-mismatch
528
+ """Returns the HTML representation of the message.
529
+
530
+ Args:
531
+ view: The HTML tree view.
532
+ 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
+ collapse_level: The global collapse level.
546
+ **kwargs: Other keyword arguments.
547
+
548
+ Returns:
549
+ The HTML representation of the message content.
550
+ """
551
+ def render_tags():
552
+ return pg.Html.element(
553
+ 'div',
554
+ [pg.Html.element('span', [tag]) for tag in self.tags],
555
+ css_class=['message-tags'],
556
+ )
557
+
558
+ def render_message_text():
559
+ maybe_reformatted = self.get('formatted_text')
560
+ referred_chunks = {}
561
+ s = pg.Html('<div class="message-text">')
562
+ for chunk in self.chunk(maybe_reformatted):
563
+ if isinstance(chunk, str):
564
+ s.write(s.escape(chunk))
565
+ else:
566
+ assert isinstance(chunk, modality.Modality), chunk
567
+ child_path = root_path + 'metadata' + chunk.referred_name
568
+ s.write(
569
+ pg.Html.element(
570
+ 'div',
571
+ [
572
+ view.render(
573
+ chunk,
574
+ name=chunk.referred_name,
575
+ root_path=child_path,
576
+ collapse_level=child_path.depth + (
577
+ 0 if collapse_modalities_in_text else 1
578
+ )
579
+ )
580
+ ],
581
+ css_class=['modality-in-text'],
582
+ )
523
583
  )
584
+ referred_chunks[chunk.referred_name] = chunk
585
+ s.write('</div>')
586
+ return s
587
+
588
+ def render_result():
589
+ if 'result' not in self.metadata:
590
+ return None
591
+ child_path = root_path + 'metadata' + 'result'
592
+ return pg.Html.element(
593
+ 'div',
594
+ [
595
+ view.render(
596
+ self.result,
597
+ name='result',
598
+ root_path=child_path,
599
+ collapse_level=view.max_collapse_level(
600
+ collapse_level,
601
+ collapse_message_result_level,
602
+ child_path,
603
+ )
604
+ )
605
+ ],
606
+ css_class=['message-result'],
607
+ )
608
+
609
+ def render_usage():
610
+ if 'usage' not in self.metadata:
611
+ return None
612
+ child_path = root_path + 'metadata' + 'usage'
613
+ return pg.Html.element(
614
+ 'div',
615
+ [
616
+ view.render(
617
+ self.usage,
618
+ name='llm usage',
619
+ root_path=child_path,
620
+ collapse_level=view.max_collapse_level(
621
+ collapse_level,
622
+ 0 if collapse_llm_usage else 1,
623
+ child_path,
624
+ )
625
+ )
626
+ ],
627
+ css_class=['message-usage'],
628
+ )
629
+
630
+ def render_source_message():
631
+ source = self.source
632
+ while (source is not None
633
+ and source_tag is not None
634
+ and not source.has_tag(source_tag)):
635
+ source = source.source
636
+ if source is not None:
637
+ child_path = root_path + 'source'
638
+ return view.render(
639
+ self.source,
640
+ name='source',
641
+ root_path=child_path,
642
+ include_metadata=include_message_metadata,
643
+ collapse_level=view.max_collapse_level(
644
+ collapse_level,
645
+ collapse_source_message_level,
646
+ child_path
647
+ ),
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,
653
+ )
654
+ return None
655
+
656
+ def render_metadata():
657
+ if not include_message_metadata:
658
+ return None
659
+ child_path = root_path + 'metadata'
660
+ return pg.Html.element(
661
+ 'div',
662
+ [
663
+ view.render(
664
+ self.metadata,
665
+ css_class=['message-metadata'],
666
+ name='metadata',
667
+ root_path=child_path,
668
+ collapse_level=view.max_collapse_level(
669
+ collapse_level,
670
+ collapse_message_metadata_level,
671
+ child_path,
672
+ )
673
+ )
674
+ ],
675
+ css_class=['message-metadata'],
524
676
  )
525
- s.write('<hr>')
526
677
 
527
- # Body.
528
- s.write(
529
- f'<span style="color: {self._text_color()}; white-space: pre-wrap;">'
678
+ return pg.Html.element(
679
+ 'div',
680
+ [
681
+ render_tags(),
682
+ render_message_text(),
683
+ render_result(),
684
+ render_usage(),
685
+ render_metadata(),
686
+ render_source_message(),
687
+ ],
688
+ css_class=['complex_value'],
530
689
  )
531
690
 
532
- # NOTE(daiyip): LLM may reformat the text from the input, therefore
533
- # we proritize the formatted text if it's available.
534
- maybe_reformatted = self.get('formatted_text')
535
- referred_chunks = {}
536
- for chunk in self.chunk(maybe_reformatted):
537
- if isinstance(chunk, str):
538
- s.write(html.escape(chunk))
539
- else:
540
- assert isinstance(chunk, modality.Modality), chunk
541
- s.write('&nbsp;')
542
- s.write(repr_utils.html_round_text(
543
- chunk.referred_name,
544
- text_color='black',
545
- background_color='#f7dc6f'
546
- ))
547
- s.write('&nbsp;')
548
- referred_chunks[chunk.referred_name] = chunk
549
- s.write('</span>')
550
-
551
- def item_color(k, v):
552
- if isinstance(v, modality.Modality):
553
- return ('black', '#f7dc6f', None, None) # Light yellow
554
- elif k == 'result':
555
- return ('white', 'purple', 'purple', None) # Blue.
556
- elif k in ('usage',):
557
- return ('white', '#e74c3c', None, None) # Red.
558
- else:
559
- return ('white', '#17202a', None, None) # Dark gray
560
-
561
- # TODO(daiyip): Revisit the logic in deciding what metadata keys to
562
- # expose to the user.
563
- if referred_chunks:
564
- s.write(repr_utils.html_repr(referred_chunks, item_color))
565
-
566
- if 'lm-response' in self.tags:
567
- s.write(repr_utils.html_repr(self.metadata, item_color))
568
- s.write('</div>')
569
- return repr_utils.Html(s.getvalue())
570
-
571
- def _text_color(self) -> str:
572
- match self.__class__.__name__:
573
- case 'UserMessage':
574
- return 'green'
575
- case 'AIMessage':
576
- return 'blue'
577
- case _:
578
- return 'black'
691
+ def _html_style(self) -> list[str]:
692
+ return super()._html_style() + [
693
+ """
694
+ /* Langfun Message styles.*/
695
+ [class^="message-"] > details {
696
+ margin: 0px 0px 5px 0px;
697
+ border: 1px solid #EEE;
698
+ }
699
+ details.lf-message > summary > .summary_title::after {
700
+ content: ' 💬';
701
+ }
702
+ details.pyglove.ai-message {
703
+ border: 1px solid blue;
704
+ color: blue;
705
+ }
706
+ details.pyglove.user-message {
707
+ border: 1px solid green;
708
+ color: green;
709
+ }
710
+ .message-tags {
711
+ margin: 5px 0px 5px 0px;
712
+ font-size: .8em;
713
+ }
714
+ .message-tags > span {
715
+ border-radius: 5px;
716
+ background-color: #CCC;
717
+ padding: 3px;
718
+ margin: 0px 2px 0px 2px;
719
+ color: white;
720
+ }
721
+ .message-text {
722
+ padding: 20px;
723
+ margin: 10px 5px 10px 5px;
724
+ font-style: italic;
725
+ font-size: 1.1em;
726
+ white-space: pre-wrap;
727
+ border: 1px solid #EEE;
728
+ border-radius: 5px;
729
+ background-color: #EEE;
730
+ }
731
+ .modality-in-text {
732
+ display: inline-block;
733
+ }
734
+ .modality-in-text > details {
735
+ display: inline-block;
736
+ font-size: 0.8em;
737
+ border: 0;
738
+ background-color: #A6F1A6;
739
+ margin: 0px 5px 0px 5px;
740
+ }
741
+ .message-result {
742
+ color: purple;
743
+ }
744
+ .message-usage {
745
+ color: orange;
746
+ }
747
+ .message-usage .object_key.str {
748
+ border: 1px solid orange;
749
+ background-color: orange;
750
+ color: white;
751
+ }
752
+ """
753
+ ]
754
+
755
+ def _html_element_class(self) -> list[str]:
756
+ return super()._html_element_class() + ['lf-message']
757
+
579
758
 
580
759
  #
581
760
  # Messages of different roles.
@@ -15,6 +15,7 @@
15
15
 
16
16
  import inspect
17
17
  import unittest
18
+ from langfun.core import language_model
18
19
  from langfun.core import message
19
20
  from langfun.core import modality
20
21
  import pyglove as pg
@@ -26,9 +27,6 @@ class CustomModality(modality.Modality):
26
27
  def to_bytes(self):
27
28
  return self.content.encode()
28
29
 
29
- def _repr_html_(self):
30
- return f'<div>CustomModality: {self.content}</div>'
31
-
32
30
 
33
31
  class MessageTest(unittest.TestCase):
34
32
 
@@ -39,11 +37,19 @@ class MessageTest(unittest.TestCase):
39
37
 
40
38
  d = pg.Dict(x=A())
41
39
 
42
- m = message.UserMessage('hi', metadata=dict(x=1), x=pg.Ref(d.x), y=2)
40
+ m = message.UserMessage(
41
+ 'hi',
42
+ metadata=dict(x=1), x=pg.Ref(d.x),
43
+ y=2,
44
+ tags=['lm-input']
45
+ )
43
46
  self.assertEqual(m.metadata, {'x': pg.Ref(d.x), 'y': 2})
44
47
  self.assertEqual(m.sender, 'User')
45
48
  self.assertIs(m.x, d.x)
46
49
  self.assertEqual(m.y, 2)
50
+ self.assertTrue(m.has_tag('lm-input'))
51
+ self.assertTrue(m.has_tag(('lm-input', '')))
52
+ self.assertFalse(m.has_tag('lm-response'))
47
53
 
48
54
  with self.assertRaises(AttributeError):
49
55
  _ = m.z
@@ -332,50 +338,69 @@ class MessageTest(unittest.TestCase):
332
338
  )
333
339
  )
334
340
 
335
- def test_html(self):
336
- m = message.UserMessage(
337
- 'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>',
338
- img1=CustomModality('foo'),
339
- x=dict(img2=CustomModality('bar')),
341
+ def assert_html_content(self, html, expected):
342
+ expected = inspect.cleandoc(expected).strip()
343
+ actual = html.content.strip()
344
+ if actual != expected:
345
+ print(actual)
346
+ self.assertEqual(actual, expected)
347
+
348
+ def test_html_user_message(self):
349
+ self.assert_html_content(
350
+ message.UserMessage(
351
+ 'what is a <div>'
352
+ ).to_html(enable_summary_tooltip=False),
353
+ """
354
+ <details open class="pyglove user-message lf-message"><summary><div class="summary_title">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"></div><div class="message-text">what is a &lt;div&gt;</div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><span class="empty_container"></span></div></details></div></div></details>
355
+ """
340
356
  )
341
- self.assertEqual(
342
- m._repr_html_(),
343
- (
344
- '<div style="padding:0px 10px 0px 10px;"><span style="color: white;'
345
- 'background-color: green;display:inline-block; border-radius:10px; '
346
- 'padding:5px; margin-top: 5px; margin-bottom: 5px; white-space: '
347
- 'pre-wrap">UserMessage</span><hr><span style="color: green; '
348
- 'white-space: pre-wrap;">hi, this is a&nbsp;<span style="color: '
349
- 'black;background-color: #f7dc6f;display:inline-block; '
350
- 'border-radius:10px; padding:5px; margin-top: 5px; margin-bottom: '
351
- '5px; white-space: pre-wrap">img1</span>&nbsp;and&nbsp;<span style'
352
- '="color: black;background-color: #f7dc6f;display:inline-block; '
353
- 'border-radius:10px; padding:5px; margin-top: 5px; margin-bottom: '
354
- '5px; white-space: pre-wrap">x.img2</span>&nbsp;</span><div style='
355
- '"padding-left: 20px; margin-top: 10px"><table style="border-top: '
356
- '1px solid #EEEEEE;"><tr><td style="padding: 5px; vertical-align: '
357
- 'top; border-bottom: 1px solid #EEEEEE"><span style="color: black;'
358
- 'background-color: #f7dc6f;display:inline-block; border-radius:'
359
- '10px; padding:5px; margin-top: 5px; margin-bottom: 0px; '
360
- 'white-space: pre-wrap">img1</span></td><td style="padding: 15px '
361
- '5px 5px 5px; vertical-align: top; border-bottom: 1px solid '
362
- '#EEEEEE;"><div>CustomModality: foo</div></td></tr><tr><td style='
363
- '"padding: 5px; vertical-align: top; border-bottom: 1px solid '
364
- '#EEEEEE"><span style="color: black;background-color: #f7dc6f;'
365
- 'display:inline-block; border-radius:10px; padding:5px; margin-top:'
366
- ' 5px; margin-bottom: 0px; white-space: pre-wrap">x.img2</span>'
367
- '</td><td style="padding: 15px 5px 5px 5px; vertical-align: top; '
368
- 'border-bottom: 1px solid #EEEEEE;"><div>CustomModality: bar</div>'
369
- '</td></tr></table></div></div>'
370
- )
357
+ self.assert_html_content(
358
+ message.UserMessage(
359
+ 'what is this <<[[image]]>>',
360
+ tags=['lm-input'],
361
+ image=CustomModality('bird')
362
+ ).to_html(enable_summary_tooltip=False, include_message_metadata=False),
363
+ """
364
+ <details open class="pyglove user-message lf-message"><summary><div class="summary_title">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-input</span></div><div class="message-text">what is this<div class="modality-in-text"><details class="pyglove custom-modality"><summary><div class="summary_name">image</div><div class="summary_title">CustomModality(...)</div></summary><div class="complex_value custom-modality"><table><tr><td><span class="object_key str">content</span><span class="tooltip key-path">metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;bird&#x27;</span></div></td></tr></table></div></details></div></div></div></details>
365
+ """
371
366
  )
372
- self.assertIn(
373
- 'background-color: blue',
374
- message.AIMessage('hi').to_html().content,
367
+
368
+ def test_html_ai_message(self):
369
+ image = CustomModality('foo')
370
+ user_message = message.UserMessage(
371
+ 'What is in this image? <<[[image]]>> this is a test',
372
+ metadata=dict(image=image),
373
+ source=message.UserMessage('User input'),
374
+ tags=['lm-input']
375
375
  )
376
- self.assertIn(
377
- 'background-color: black',
378
- message.SystemMessage('hi').to_html().content,
376
+ ai_message = message.AIMessage(
377
+ 'My name is Gemini',
378
+ metadata=dict(
379
+ result=pg.Dict(x=1, y=2, z=pg.Dict(a=[12, 323])),
380
+ usage=language_model.LMSamplingUsage(10, 2, 12)
381
+ ),
382
+ tags=['lm-response', 'lm-output'],
383
+ source=user_message,
384
+ )
385
+ self.assert_html_content(
386
+ ai_message.to_html(enable_summary_tooltip=False),
387
+ """
388
+ <details open class="pyglove ai-message lf-message"><summary><div class="summary_title">AIMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-response</span><span>lm-output</span></div><div class="message-text">My name is Gemini</div><div class="message-result"><details open class="pyglove dict"><summary><div class="summary_name">result</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.result.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.result.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">z</span><span class="tooltip key-path">metadata.result.z</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict"><table><tr><td><span class="object_key str">a</span><span class="tooltip key-path">metadata.result.z.a</span></td><td><div><details class="pyglove list"><summary><div class="summary_title">List(...)</div></summary><div class="complex_value list"><table><tr><td><span class="object_key int">0</span><span class="tooltip key-path">metadata.result.z.a[0]</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key int">1</span><span class="tooltip key-path">metadata.result.z.a[1]</span></td><td><div><span class="simple_value int">323</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr></table></div></details></div><div class="message-usage"><details open class="pyglove lm-sampling-usage"><summary><div class="summary_name">llm usage</div><div class="summary_title">LMSamplingUsage(...)</div></summary><div class="complex_value lm-sampling-usage"><table><tr><td><span class="object_key str">prompt_tokens</span><span class="tooltip key-path">metadata.usage.prompt_tokens</span></td><td><div><span class="simple_value int">10</span></div></td></tr><tr><td><span class="object_key str">completion_tokens</span><span class="tooltip key-path">metadata.usage.completion_tokens</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">total_tokens</span><span class="tooltip key-path">metadata.usage.total_tokens</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key str">num_requests</span><span class="tooltip key-path">metadata.usage.num_requests</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">estimated_cost</span><span class="tooltip key-path">metadata.usage.estimated_cost</span></td><td><div><span class="simple_value none-type">None</span></div></td></tr></table></div></details></div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">result</span><span class="tooltip key-path">metadata.result</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">x</span><span class="tooltip key-path">metadata.result.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.result.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">z</span><span class="tooltip key-path">metadata.result.z</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">a</span><span class="tooltip key-path">metadata.result.z.a</span></td><td><div><details class="pyglove list"><summary><div class="summary_title">List(...)</div></summary><div class="complex_value list message-metadata"><table><tr><td><span class="object_key int">0</span><span class="tooltip key-path">metadata.result.z.a[0]</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key int">1</span><span class="tooltip key-path">metadata.result.z.a[1]</span></td><td><div><span class="simple_value int">323</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr><tr><td><span class="object_key str">usage</span><span class="tooltip key-path">metadata.usage</span></td><td><div><details class="pyglove lm-sampling-usage"><summary><div class="summary_title">LMSamplingUsage(...)</div></summary><div class="complex_value lm-sampling-usage message-metadata"><table><tr><td><span class="object_key str">prompt_tokens</span><span class="tooltip key-path">metadata.usage.prompt_tokens</span></td><td><div><span class="simple_value int">10</span></div></td></tr><tr><td><span class="object_key str">completion_tokens</span><span class="tooltip key-path">metadata.usage.completion_tokens</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">total_tokens</span><span class="tooltip key-path">metadata.usage.total_tokens</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key str">num_requests</span><span class="tooltip key-path">metadata.usage.num_requests</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">estimated_cost</span><span class="tooltip key-path">metadata.usage.estimated_cost</span></td><td><div><span class="simple_value none-type">None</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div><details open class="pyglove user-message lf-message"><summary><div class="summary_name">source</div><div class="summary_title">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-input</span></div><div class="message-text">What is in this image?<div class="modality-in-text"><details class="pyglove custom-modality"><summary><div class="summary_name">image</div><div class="summary_title">CustomModality(...)</div></summary><div class="complex_value custom-modality"><table><tr><td><span class="object_key str">content</span><span class="tooltip key-path">source.metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;foo&#x27;</span></div></td></tr></table></div></details></div>this is a test</div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">image</span><span class="tooltip key-path">source.metadata.image</span></td><td><div><details class="pyglove custom-modality"><summary><div class="summary_title">CustomModality(...)</div></summary><div class="complex_value custom-modality message-metadata"><table><tr><td><span class="object_key str">content</span><span class="tooltip key-path">source.metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;foo&#x27;</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></div></details></div></details>
389
+ """
390
+ )
391
+ self.assert_html_content(
392
+ ai_message.to_html(
393
+ enable_summary_tooltip=False,
394
+ collapse_modalities_in_text=False,
395
+ collapse_llm_usage=True,
396
+ collapse_message_result_level=0,
397
+ collapse_message_metadata_level=0,
398
+ collapse_source_message_level=0,
399
+ source_tag=None,
400
+ ),
401
+ """
402
+ <details open class="pyglove ai-message lf-message"><summary><div class="summary_title">AIMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-response</span><span>lm-output</span></div><div class="message-text">My name is Gemini</div><div class="message-result"><details class="pyglove dict"><summary><div class="summary_name">result</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.result.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.result.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">z</span><span class="tooltip key-path">metadata.result.z</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict"><table><tr><td><span class="object_key str">a</span><span class="tooltip key-path">metadata.result.z.a</span></td><td><div><details class="pyglove list"><summary><div class="summary_title">List(...)</div></summary><div class="complex_value list"><table><tr><td><span class="object_key int">0</span><span class="tooltip key-path">metadata.result.z.a[0]</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key int">1</span><span class="tooltip key-path">metadata.result.z.a[1]</span></td><td><div><span class="simple_value int">323</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr></table></div></details></div><div class="message-usage"><details class="pyglove lm-sampling-usage"><summary><div class="summary_name">llm usage</div><div class="summary_title">LMSamplingUsage(...)</div></summary><div class="complex_value lm-sampling-usage"><table><tr><td><span class="object_key str">prompt_tokens</span><span class="tooltip key-path">metadata.usage.prompt_tokens</span></td><td><div><span class="simple_value int">10</span></div></td></tr><tr><td><span class="object_key str">completion_tokens</span><span class="tooltip key-path">metadata.usage.completion_tokens</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">total_tokens</span><span class="tooltip key-path">metadata.usage.total_tokens</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key str">num_requests</span><span class="tooltip key-path">metadata.usage.num_requests</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">estimated_cost</span><span class="tooltip key-path">metadata.usage.estimated_cost</span></td><td><div><span class="simple_value none-type">None</span></div></td></tr></table></div></details></div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">result</span><span class="tooltip key-path">metadata.result</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">x</span><span class="tooltip key-path">metadata.result.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.result.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">z</span><span class="tooltip key-path">metadata.result.z</span></td><td><div><details class="pyglove dict"><summary><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">a</span><span class="tooltip key-path">metadata.result.z.a</span></td><td><div><details class="pyglove list"><summary><div class="summary_title">List(...)</div></summary><div class="complex_value list message-metadata"><table><tr><td><span class="object_key int">0</span><span class="tooltip key-path">metadata.result.z.a[0]</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key int">1</span><span class="tooltip key-path">metadata.result.z.a[1]</span></td><td><div><span class="simple_value int">323</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr></table></div></details></div></td></tr><tr><td><span class="object_key str">usage</span><span class="tooltip key-path">metadata.usage</span></td><td><div><details class="pyglove lm-sampling-usage"><summary><div class="summary_title">LMSamplingUsage(...)</div></summary><div class="complex_value lm-sampling-usage message-metadata"><table><tr><td><span class="object_key str">prompt_tokens</span><span class="tooltip key-path">metadata.usage.prompt_tokens</span></td><td><div><span class="simple_value int">10</span></div></td></tr><tr><td><span class="object_key str">completion_tokens</span><span class="tooltip key-path">metadata.usage.completion_tokens</span></td><td><div><span class="simple_value int">2</span></div></td></tr><tr><td><span class="object_key str">total_tokens</span><span class="tooltip key-path">metadata.usage.total_tokens</span></td><td><div><span class="simple_value int">12</span></div></td></tr><tr><td><span class="object_key str">num_requests</span><span class="tooltip key-path">metadata.usage.num_requests</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">estimated_cost</span><span class="tooltip key-path">metadata.usage.estimated_cost</span></td><td><div><span class="simple_value none-type">None</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div><details class="pyglove user-message lf-message"><summary><div class="summary_name">source</div><div class="summary_title">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>lm-input</span></div><div class="message-text">What is in this image?<div class="modality-in-text"><details open class="pyglove custom-modality"><summary><div class="summary_name">image</div><div class="summary_title">CustomModality(...)</div></summary><div class="complex_value custom-modality"><table><tr><td><span class="object_key str">content</span><span class="tooltip key-path">source.metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;foo&#x27;</span></div></td></tr></table></div></details></div>this is a test</div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><table><tr><td><span class="object_key str">image</span><span class="tooltip key-path">source.metadata.image</span></td><td><div><details class="pyglove custom-modality"><summary><div class="summary_title">CustomModality(...)</div></summary><div class="complex_value custom-modality message-metadata"><table><tr><td><span class="object_key str">content</span><span class="tooltip key-path">source.metadata.image.content</span></td><td><div><span class="simple_value str">&#x27;foo&#x27;</span></div></td></tr></table></div></details></div></td></tr></table></div></details></div><details class="pyglove user-message lf-message"><summary><div class="summary_name">source</div><div class="summary_title">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"></div><div class="message-text">User input</div><div class="message-metadata"><details class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict message-metadata"><span class="empty_container"></span></div></details></div></div></details></div></details></div></details>
403
+ """
379
404
  )
380
405
 
381
406
 
@@ -26,5 +26,5 @@ class Audio(mime.Mime):
26
26
  def audio_format(self) -> str:
27
27
  return self.mime_type.removeprefix(self.MIME_PREFIX + '/')
28
28
 
29
- def _html(self, uri: str) -> str:
29
+ def _mime_control_for(self, uri: str) -> str:
30
30
  return f'<audio controls> <source src="{uri}"> </audio>'
@@ -53,7 +53,7 @@ class AudioFileTest(unittest.TestCase):
53
53
  self.assertEqual(audio.audio_format, 'x-wav')
54
54
  self.assertEqual(audio.mime_type, 'audio/x-wav')
55
55
  self.assertEqual(
56
- audio._repr_html_(),
56
+ audio._raw_html(),
57
57
  '<audio controls> <source src="http://mock/web/a.wav"> </audio>',
58
58
  )
59
59
  self.assertEqual(audio.to_bytes(), content_bytes)
@@ -41,7 +41,7 @@ class Image(mime.Mime):
41
41
  def image_format(self) -> str:
42
42
  return self.mime_type.removeprefix(self.MIME_PREFIX + '/')
43
43
 
44
- def _html(self, uri: str) -> str:
44
+ def _mime_control_for(self, uri: str) -> str:
45
45
  return f'<img src="{uri}">'
46
46
 
47
47
  @functools.cached_property
@@ -45,7 +45,7 @@ class ImageTest(unittest.TestCase):
45
45
  def test_from_bytes(self):
46
46
  image = image_lib.Image.from_bytes(image_content)
47
47
  self.assertEqual(image.image_format, 'png')
48
- self.assertIn('data:image/png;base64,', image._repr_html_())
48
+ self.assertIn('data:image/png;base64,', image._raw_html())
49
49
  self.assertEqual(image.to_bytes(), image_content)
50
50
  with self.assertRaisesRegex(
51
51
  lf.ModalityError, '.* cannot be converted to text'
@@ -67,7 +67,10 @@ class ImageTest(unittest.TestCase):
67
67
  with mock.patch('requests.get') as mock_requests_get:
68
68
  mock_requests_get.side_effect = mock_request
69
69
  self.assertEqual(image.image_format, 'png')
70
- self.assertEqual(image._repr_html_(), '<img src="http://mock/web/a.png">')
70
+ self.assertEqual(
71
+ image._raw_html(),
72
+ '<img src="http://mock/web/a.png">'
73
+ )
71
74
  self.assertEqual(image.to_bytes(), image_content)
72
75
 
73
76
  def test_from_uri_base_cls(self):
@@ -76,7 +79,10 @@ class ImageTest(unittest.TestCase):
76
79
  image = mime_lib.Mime.from_uri('http://mock/web/a.png')
77
80
  self.assertIsInstance(image, image_lib.Image)
78
81
  self.assertEqual(image.image_format, 'png')
79
- self.assertEqual(image._repr_html_(), '<img src="http://mock/web/a.png">')
82
+ self.assertEqual(
83
+ image._raw_html(),
84
+ '<img src="http://mock/web/a.png">'
85
+ )
80
86
  self.assertEqual(image.to_bytes(), image_content)
81
87
 
82
88
  def test_image_size(self):