langfun 0.1.2.dev202410110804__py3-none-any.whl → 0.1.2.dev202410130803__py3-none-any.whl

langfun/core/llms/vertexai.py

@@ -40,24 +40,106 @@ except ImportError:
   Credentials = Any
 
 
+# https://cloud.google.com/vertex-ai/generative-ai/pricing
+# describes that the average number of characters per token is about 4.
+AVGERAGE_CHARS_PER_TOEKN = 4
+
+
+# Price in US dollars,
+# from https://cloud.google.com/vertex-ai/generative-ai/pricing
+# as of 2024-10-10.
 SUPPORTED_MODELS_AND_SETTINGS = {
-    'gemini-1.5-pro-001': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-pro-002': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-flash-002': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-flash-001': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-pro': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-flash': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-pro-latest': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-flash-latest': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-pro-preview-0514': pg.Dict(api='gemini', rpm=50),
-    'gemini-1.5-pro-preview-0409': pg.Dict(api='gemini', rpm=50),
-    'gemini-1.5-flash-preview-0514': pg.Dict(api='gemini', rpm=200),
-    'gemini-1.0-pro': pg.Dict(api='gemini', rpm=300),
-    'gemini-1.0-pro-vision': pg.Dict(api='gemini', rpm=100),
+    'gemini-1.5-pro-001': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-pro-002': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-flash-002': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.5-flash-001': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.5-pro': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-flash': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.5-pro-latest': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-flash-latest': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.5-pro-preview-0514': pg.Dict(
+        api='gemini',
+        rpm=50,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-pro-preview-0409': pg.Dict(
+        api='gemini',
+        rpm=50,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-flash-preview-0514': pg.Dict(
+        api='gemini',
+        rpm=200,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.0-pro': pg.Dict(
+        api='gemini',
+        rpm=300,
+        cost_per_1k_input_chars=0.000125,
+        cost_per_1k_output_chars=0.000375,
+    ),
+    'gemini-1.0-pro-vision': pg.Dict(
+        api='gemini',
+        rpm=100,
+        cost_per_1k_input_chars=0.000125,
+        cost_per_1k_output_chars=0.000375,
+    ),
     # PaLM APIs.
-    'text-bison': pg.Dict(api='palm', rpm=1600),
-    'text-bison-32k': pg.Dict(api='palm', rpm=300),
-    'text-unicorn': pg.Dict(api='palm', rpm=100),
+    'text-bison': pg.Dict(
+        api='palm',
+        rpm=1600
+    ),
+    'text-bison-32k': pg.Dict(
+        api='palm',
+        rpm=300
+    ),
+    'text-unicorn': pg.Dict(
+        api='palm',
+        rpm=100
+    ),
     # Endpoint
     # TODO(chengrun): Set a more appropriate rpm for endpoint.
     'custom': pg.Dict(api='endpoint', rpm=20),
@@ -161,6 +243,25 @@ class VertexAI(lf.LanguageModel):
         tokens_per_min=0,
     )
 
+  def estimate_cost(
+      self,
+      num_input_tokens: int,
+      num_output_tokens: int
+  ) -> float | None:
+    """Estimate the cost based on usage."""
+    cost_per_1k_input_chars = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
+        'cost_per_1k_input_chars', None
+    )
+    cost_per_1k_output_chars = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
+        'cost_per_1k_output_chars', None
+    )
+    if cost_per_1k_output_chars is None or cost_per_1k_input_chars is None:
+      return None
+    return (
+        cost_per_1k_input_chars * num_input_tokens
+        + cost_per_1k_output_chars * num_output_tokens
+    ) * AVGERAGE_CHARS_PER_TOEKN / 1000
+
   def _generation_config(
       self, prompt: lf.Message, options: lf.LMSamplingOptions
   ) -> Any:  # generative_models.GenerationConfig
@@ -285,6 +386,10 @@ class VertexAI(lf.LanguageModel):
         prompt_tokens=usage_metadata.prompt_token_count,
         completion_tokens=usage_metadata.candidates_token_count,
         total_tokens=usage_metadata.total_token_count,
+        estimated_cost=self.estimate_cost(
+            num_input_tokens=usage_metadata.prompt_token_count,
+            num_output_tokens=usage_metadata.candidates_token_count,
+        ),
     )
     return lf.LMSamplingResult(
         [

langfun/core/logging.py

@@ -89,8 +89,9 @@ class LogEntry(pg.Object):
       self,
       view: pg.views.HtmlTreeView,
       root_path: pg.KeyPath,
-      collapse_log_metadata_level: int = pg.View.PresetArgValue(0),
+      collapse_log_metadata_level: int | None = pg.View.PresetArgValue(0),
       max_str_len_for_summary: int = pg.View.PresetArgValue(80),
+      collapse_level: int | None = pg.View.PresetArgValue(1),
       **kwargs
   ) -> pg.Html:
     # pytype: enable=annotation-type-mismatch
@@ -106,16 +107,21 @@ class LogEntry(pg.Object):
     def render_metadata():
       if not self.metadata:
         return None
+      child_path = root_path + 'metadata'
       return pg.Html.element(
           'div',
           [
               view.render(
                   self.metadata,
                   name='metadata',
-                  root_path=root_path + 'metadata',
+                  root_path=child_path,
                   parent=self,
                   collapse_level=(
-                      root_path.depth + collapse_log_metadata_level + 1
+                      view.max_collapse_level(
+                          collapse_level,
+                          collapse_log_metadata_level,
+                          child_path
+                      )
                   )
               )
           ],

langfun/core/message.py

@@ -518,9 +518,10 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
       include_message_metadata: bool = pg.View.PresetArgValue(True),
       collapse_modalities_in_text: bool = pg.View.PresetArgValue(True),
       collapse_llm_usage: bool = pg.View.PresetArgValue(False),
-      collapse_message_result_level: int = pg.View.PresetArgValue(1),
-      collapse_message_metadata_level: int = pg.View.PresetArgValue(0),
-      collapse_source_message_level: int = pg.View.PresetArgValue(1),
+      collapse_message_result_level: int | None = pg.View.PresetArgValue(1),
+      collapse_message_metadata_level: int | None = pg.View.PresetArgValue(0),
+      collapse_source_message_level: int | None = pg.View.PresetArgValue(1),
+      collapse_level: int | None = pg.View.PresetArgValue(1),
       **kwargs,
   ) -> pg.Html:
   # pytype: enable=annotation-type-mismatch
@@ -541,6 +542,7 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
         message.
       collapse_source_message_level: The level to collapse the source in the
         message.
+      collapse_level: The global collapse level.
       **kwargs: Other keyword arguments.
 
     Returns:
@@ -594,8 +596,10 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
                   self.result,
                   name='result',
                   root_path=child_path,
-                  collapse_level=(
-                      child_path.depth + collapse_message_result_level
+                  collapse_level=view.max_collapse_level(
+                      collapse_level,
+                      collapse_message_result_level,
+                      child_path,
                   )
               )
           ],
@@ -613,8 +617,10 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
                   self.usage,
                   name='llm usage',
                   root_path=child_path,
-                  collapse_level=child_path.depth + (
-                      0 if collapse_llm_usage else 1
+                  collapse_level=view.max_collapse_level(
+                      collapse_level,
+                      0 if collapse_llm_usage else 1,
+                      child_path,
                   )
               )
           ],
@@ -628,13 +634,16 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
              and not source.has_tag(source_tag)):
         source = source.source
       if source is not None:
+        child_path = root_path + 'source'
         return view.render(
             self.source,
             name='source',
-            root_path=root_path + 'source',
+            root_path=child_path,
             include_metadata=include_message_metadata,
-            collapse_level=(
-                root_path.depth + 1 + collapse_source_message_level
+            collapse_level=view.max_collapse_level(
+                collapse_level,
+                collapse_source_message_level,
+                child_path
             ),
             collapse_source_level=max(0, collapse_source_message_level - 1),
             collapse_modalities=collapse_modalities_in_text,
@@ -656,8 +665,10 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
                   css_class=['message-metadata'],
                   name='metadata',
                   root_path=child_path,
-                  collapse_level=(
-                      child_path.depth + collapse_message_metadata_level
+                  collapse_level=view.max_collapse_level(
+                      collapse_level,
+                      collapse_message_metadata_level,
+                      child_path,
                   )
               )
           ],

langfun/core/message_test.py

@@ -385,7 +385,7 @@ class MessageTest(unittest.TestCase):
     self.assert_html_content(
         ai_message.to_html(enable_summary_tooltip=False),
         """
-        <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></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></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>
+        <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>
         """
     )
     self.assert_html_content(
@@ -399,7 +399,7 @@ class MessageTest(unittest.TestCase):
             source_tag=None,
         ),
         """
-        <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></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></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>
+        <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>
         """
     )
 

langfun/core/structured/completion_test.py

@@ -581,6 +581,7 @@ class CompleteStructureTest(unittest.TestCase):
             text='Activity(description="foo")',
             result=Activity(description='foo'),
             score=1.0,
+            is_cached=False,
             logprobs=None,
             usage=lf.LMSamplingUsage(553, 27, 580),
             tags=['lm-response', 'lm-output', 'transformed']

langfun/core/structured/mapping.py

@@ -218,7 +218,7 @@ class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
         **kwargs
     )
 
-  def _html_tree_view_collapse_level(self) -> int:
+  def _html_tree_view_uncollapse_level(self) -> int:
     return 2
 
 

langfun/core/structured/parsing_test.py

@@ -285,7 +285,7 @@ class ParseStructurePythonTest(unittest.TestCase):
     self.assertEqual(
         r,
         lf.AIMessage(
-            '1', score=1.0, result=1, logprobs=None,
+            '1', score=1.0, result=1, logprobs=None, is_cached=False,
             usage=lf.LMSamplingUsage(652, 1, 653),
             tags=['lm-response', 'lm-output', 'transformed']
         ),
@@ -645,6 +645,7 @@ class CallTest(unittest.TestCase):
             result=3,
             score=1.0,
             logprobs=None,
+            is_cached=False,
             usage=lf.LMSamplingUsage(315, 1, 316),
             tags=['lm-response', 'lm-output', 'transformed']
         ),

langfun/core/structured/prompting_test.py

@@ -82,6 +82,7 @@ class QueryTest(unittest.TestCase):
             result=1,
             score=1.0,
             logprobs=None,
+            is_cached=False,
             usage=lf.LMSamplingUsage(323, 1, 324),
             tags=['lm-response', 'lm-output', 'transformed'],
         ),

langfun/core/template.py

@@ -531,6 +531,8 @@ class Template(
       *,
       view: pg.views.HtmlTreeView,
       root_path: pg.KeyPath,
+      collapse_template_vars_level: int | None = pg.View.PresetArgValue(1),
+      collapse_level: int | None = pg.View.PresetArgValue(1),  # pytype: disable=annotation-type-mismatch
       **kwargs,
   ):
     def render_template_str():
@@ -574,7 +576,11 @@ class Template(
                   render_value_fn=render_value_fn,
                   exclude_keys=['template_str', 'clean'],
                   parent=self,
-                  collapse_level=root_path.depth + 1,
+                  collapse_level=view.max_collapse_level(
+                      collapse_level,
+                      collapse_template_vars_level,
+                      root_path
+                  ),
               ),
           ],
           css_class=['template-fields'],
@@ -628,10 +634,7 @@ class Template(
 
   # Additional CSS class to add to the root <details> element.
   def _html_element_class(self) -> Sequence[str] | None:
-    return [
-        pg.object_utils.camel_to_snake(self.__class__.__name__, '-'),
-        'lf-template'
-    ]
+    return super()._html_element_class() + ['lf-template']
 
 # Register converter from str to LangFunc, therefore we can always
 # pass strs to attributes that accept LangFunc.

langfun/core/templates/selfplay_test.py

@@ -59,7 +59,8 @@ class SelfPlayTest(unittest.TestCase):
       self.assertEqual(
           g(),
           lf.AIMessage(
-              '10', score=0.0, logprobs=None, usage=lf.UsageNotAvailable()
+              '10', score=0.0, logprobs=None, is_cached=False,
+              usage=lf.UsageNotAvailable()
           )
       )
 
@@ -72,7 +73,8 @@ class SelfPlayTest(unittest.TestCase):
       self.assertEqual(
           g(),
           lf.AIMessage(
-              '2', score=0.0, logprobs=None, usage=lf.UsageNotAvailable()
+              '2', score=0.0, logprobs=None, is_cached=False,
+              usage=lf.UsageNotAvailable()
           )
       )
 

{langfun-0.1.2.dev202410110804.dist-info → langfun-0.1.2.dev202410130803.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.2.dev202410110804
+Version: 0.1.2.dev202410130803
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors