diff --git a/langfun/core/logging.py b/langfun/core/logging.py
index 6c6628f..941d87d 100644
--- a/langfun/core/logging.py
+++ b/langfun/core/logging.py
@@ -89,8 +89,9 @@ def _html_tree_view_content(
       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 @@ def render_message_text():
     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
+                      )
                   )
               )
           ],
diff --git a/langfun/core/message.py b/langfun/core/message.py
index fd7b083..a8b08de 100644
--- a/langfun/core/message.py
+++ b/langfun/core/message.py
@@ -518,9 +518,10 @@ def _html_tree_view_content(
       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 @@ def _html_tree_view_content(
         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 @@ def render_result():
                   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 @@ def render_usage():
                   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 @@ def render_source_message():
              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 @@ def render_metadata():
                   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,
                   )
               )
           ],
diff --git a/langfun/core/structured/mapping.py b/langfun/core/structured/mapping.py
index 176b021..5bcb2d3 100644
--- a/langfun/core/structured/mapping.py
+++ b/langfun/core/structured/mapping.py
@@ -218,7 +218,7 @@ def render_value(value, **kwargs):
         **kwargs
     )
 
-  def _html_tree_view_collapse_level(self) -> int:
+  def _html_tree_view_uncollapse_level(self) -> int:
     return 2
 
 
diff --git a/langfun/core/template.py b/langfun/core/template.py
index b377bcc..e640375 100644
--- a/langfun/core/template.py
+++ b/langfun/core/template.py
@@ -531,6 +531,8 @@ def _html_tree_view_content(
       *,
       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 @@ def render_value_fn(value, *, root_path, **kwargs):
                   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 @@ def _html_style(self) -> list[str]:
 
   # 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.