langfun 0.0.2.dev20240330__py3-none-any.whl → 0.1.2.dev202501140804__py3-none-any.whl

langfun/core/template.py

@@ -17,7 +17,7 @@ import contextlib
 import dataclasses
 import functools
 import inspect
-from typing import Annotated, Any, Callable, Iterator, Set, Tuple, Type
+from typing import Annotated, Any, Callable, Iterator, Set, Tuple, Type, Union
 
 import jinja2
 from jinja2 import meta as jinja2_meta
@@ -38,13 +38,23 @@ NO_TEMPLATE_DOCSTR_SIGN = 'THIS IS NOT A TEMPLATE'
 _TLS_RENDER_STACK = '_template_render_stack'
 _TLS_RENDER_RESULT_CACHE = '_template_render_result_cache'
 
+# The prefix for fields or contextual attributes to be treated as additional
+# metadata for rendered message.
+_ADDITIONAL_METADATA_PREFIX = 'metadata_'
+
 
 class Template(
     natural_language.NaturalLanguageFormattable,
     component.Component,
     pg.typing.CustomTyping,
+    pg.views.HtmlTreeView.Extension
 ):
-  """Langfun string template."""
+  """Langfun string template.
+
+  Langfun uses jinja2 as its template engine. Pleaes check out
+  https://jinja.palletsprojects.com/en/3.1.x/templates/ for detailed
+  explanation on the template language.
+  """
 
   template_str: Annotated[
       str,
@@ -97,6 +107,11 @@ class Template(
       # Declare template variables as symbolic attributes.
       template_vars = Template.resolve_vars(template_str)
       for var_name in template_vars:
+        if 'DEFAULT' == var_name:
+          raise ValueError(
+              '`{{ DEFAULT }}` cannot be used in pre-configured templates. '
+              f'Encountered: {template_str!r}'
+          )
         # NOTE(daiyip): This is to avoid warning from accessing
         # `pg.Object.schema`, which was replaced by `pg.Object.__schema__`.
         if var_name == 'schema' or not hasattr(cls, var_name):
@@ -149,7 +164,7 @@ class Template(
     # TODO(daiyip): Consider to delay template parsing upon usage.
     unassigned_vars = {}
     for k in self._variables:
-      if not hasattr(self, k):
+      if k not in ('DEFAULT',) and not hasattr(self, k):
         unassigned_vars[k] = component.contextual()
     if unassigned_vars:
       self.rebind(unassigned_vars, skip_notification=True)
@@ -217,6 +232,16 @@ class Template(
     """Returns the missing variable names."""
     return self.vars(closure=True, specified=False)
 
+  @classmethod
+  def raw_str(cls, text: str) -> str:
+    """Returns a template string that preserve the text as original."""
+    return '{% raw %}' + text + '{% endraw %}'
+
+  @classmethod
+  def from_raw_str(cls, text: str) -> 'Template':
+    """Returns a template that preserve the text as original."""
+    return cls(cls.raw_str(text), clean=False)
+
   def render(
       self,
       *,
@@ -303,19 +328,19 @@ class Template(
               with modality.format_modality_as_ref():
                 rendered_text = self._template.render(**inputs)
 
+            # Carry additional metadata.
+            metadata = self.additional_metadata()
+
         if self.clean:
           rendered_text = rendered_text.strip()
 
-        # Fill the variables for rendering the template as metadata.
-        message = message_cls(
-            text=rendered_text,
-            metadata={
-                k: pg.Ref(v)
-                for k, v in inputs.items()
-                if not inspect.ismethod(v)
-            },
+        metadata.update(
+            {k: pg.Ref(v) for k, v in inputs.items() if not inspect.ismethod(v)}
         )
 
+        # Fill the variables for rendering the template as metadata.
+        message = message_cls(text=rendered_text, metadata=metadata)
+
         # Tag input as rendered message.
         message.tag(message_lib.Message.TAG_RENDERED)
 
@@ -340,6 +365,20 @@ class Template(
       top = pg.object_utils.thread_local_pop(_TLS_RENDER_STACK)
       assert top is self, (top, self)
 
+  def additional_metadata(self) -> dict[str, Any]:
+    """Returns additional metadta to be carried in the rendered message."""
+    metadata = {}
+    # Carry metadata from `lf.context`.
+    for k, v in component.all_contextual_values().items():
+      if k.startswith(_ADDITIONAL_METADATA_PREFIX):
+        metadata[k.removeprefix(_ADDITIONAL_METADATA_PREFIX)] = v
+
+    # Carry metadata from fields.
+    for k, v in self.sym_init_args.items():
+      if k.startswith(_ADDITIONAL_METADATA_PREFIX):
+        metadata[k.removeprefix(_ADDITIONAL_METADATA_PREFIX)] = v
+    return metadata
+
   #
   # Implements `pg.typing.CustomTyping`.
   #
@@ -380,6 +419,196 @@ class Template(
     # Override __hash__ since __eq__ has changed.
     return object.__hash__(self)
 
+  #
+  # Special methods.
+  #
+
+  @property
+  def DEFAULT(self) -> 'Template':
+    """Referring to the default value used for this template.
+
+    This method is intended to be used in template for referring to the default
+    value of current template. There are two scenarios:
+
+    Scenario 1: Use instance-level template_str to override the class default.
+
+    ```
+    class Foo(lf.Template):
+       '''Foo template.
+
+       This is {{x}}.
+       '''
+
+    f = Foo(template_str='<h1>{{DEFAULT}}</h1>', x=1)
+    f.render()
+
+    >> <h1>This is 1.</h1>
+    ```
+
+    Scenario 2: Use an ad-hoc template to override a predefined field.
+
+    ```
+    class Bar(lf.Template):
+      '''Bar template.
+
+      {{preamble}}
+      {{prompt}}
+      '''
+      preamble: lf.Template = lf.Template('You are a chat bot.')
+      prompt: lf.Template = lf.Template('User: hi')
+
+    b = Bar(preamble=lf.Template('<h1>{{DEFAULT}}<h1>'),
+            prompt=lf.Template('<h2>{{DEFAULT}}</h2>')
+    b.render()
+
+    >> <h1>You are a chat bot.<h1>
+    >> <h2>User: hi</h2>
+    ```
+
+    Returns:
+      The default (pre-configured) value used for this template.
+    """
+    base_template = self.__class__.__schema__['template_str'].default_value
+    if base_template == pg.MISSING_VALUE:
+      if not self.sym_path:
+        raise ValueError(
+            f'No DEFAULT template found for {self!r}: '
+            'The template neither has a default `template_str` nor is '
+            'contained under another object.'
+        )
+      key = self.sym_path.key
+      assert self.sym_parent is not None
+      assigned_field = self.sym_parent.sym_attr_field(key)
+      container_cls = self.sym_parent.__class__
+
+      if (
+          assigned_field is None
+          or assigned_field.default_value == pg.MISSING_VALUE
+      ):
+        raise ValueError(
+            f'No DEFAULT template found for {self!r}: '
+            f'`{container_cls.__name__}.{key}` '
+            'does not have a default value. '
+        )
+      base_template = assigned_field.default_value
+      if isinstance(base_template, Template):
+        base_template = base_template.template_str
+      if not isinstance(base_template, str):
+        raise ValueError(
+            f'No DEFAULT template found for {self!r}: The default '
+            f'value {base_template!r} of '
+            f'`{container_cls.__name__}.{key}` is not a '
+            '`lf.Template` object or str.'
+        )
+    t = Template(base_template)
+    # NOTE(daiyip): Set the parent of the newly created template to self so
+    # it could access all the contextual variables.
+    t.sym_setparent(self)
+    return t
+
+  @classmethod
+  def from_value(
+      cls,
+      value: Union[str, message_lib.Message, 'Template'],
+      **kwargs
+  ) -> 'Template':
+    """Create a template object from a string or template."""
+    if isinstance(value, cls):
+      return value.clone(override=kwargs) if kwargs else value  # pylint: disable=no-value-for-parameter
+    if isinstance(value, str):
+      return cls(template_str=value, **kwargs)
+    if isinstance(value, message_lib.Message):
+      kwargs.update(value.metadata)
+      return cls(template_str=value.text, **kwargs)
+    if isinstance(value, Template):
+      lfun = cls(template_str=value.template_str, **kwargs)
+      # So lfun could acccess all attributes from value.
+      lfun.sym_setparent(value)
+      return lfun
+    return cls(template_str='{{input}}', input=value, **kwargs)
+
+  def _html_tree_view_content(
+      self,
+      *,
+      view: pg.views.HtmlTreeView,
+      root_path: pg.KeyPath | None = None,
+      collapse_level: int | None = None,
+      extra_flags: dict[str, Any] | None = None,
+      debug: bool = False,
+      **kwargs,
+  ):
+    extra_flags = extra_flags if extra_flags is not None else {}
+    collapse_template_vars_level: int | None = extra_flags.get(
+        'collapse_template_vars_level', 1
+    )
+
+    def render_template_str():
+      return pg.Html.element(
+          'div',
+          [
+              pg.Html.element('span', [self.template_str])
+          ],
+          css_classes=['template-str'],
+      )
+
+    def render_fields():
+      return view.complex_value(
+          {k: v for k, v in self.sym_items()},
+          name='fields',
+          root_path=root_path,
+          parent=self,
+          exclude_keys=['template_str', 'clean'],
+          collapse_level=max(
+              collapse_template_vars_level, collapse_level
+          ) if collapse_level is not None else None,
+          extra_flags=extra_flags,
+          debug=debug,
+          **view.get_passthrough_kwargs(
+              remove=['exclude_keys'],
+              **kwargs,
+          )
+      )
+
+    return pg.views.html.controls.TabControl([
+        pg.views.html.controls.Tab(
+            'template_str',
+            render_template_str(),
+        ),
+        pg.views.html.controls.Tab(
+            'variables',
+            render_fields(),
+        ),
+    ], selected=1)
+
+  @classmethod
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        /* Langfun Template styles. */
+        .template-str {
+            padding: 10px;
+            margin: 10px 5px 10px 5px;
+            font-style: italic;
+            font-size: 1.1em;
+            white-space: pre-wrap;
+            border: 1px solid #EEE;
+            border-radius: 5px;
+            background-color: #EEE;
+            color: #cc2986;
+        }
+        """
+    ]
+
+  @classmethod
+  @functools.cache
+  def _html_tree_view_config(cls) -> dict[str, Any]:
+    return pg.views.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            css_classes=['lf-template'],
+        )
+    )
+
 
 # Register converter from str to LangFunc, therefore we can always
 # pass strs to attributes that accept LangFunc.

langfun/core/template_test.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 """Template test."""
 import inspect
+from typing import Any
 import unittest
 
 from langfun.core import component
+from langfun.core import message as message_lib
 from langfun.core import modality
 from langfun.core import subscription
 from langfun.core.template import Template
@@ -97,6 +99,21 @@ class BasicTest(unittest.TestCase):
     self.assertEqual(d.z.render(), 'Bye, 1')
     self.assertEqual(d.p.render(), 'Again Hello, 1')
 
+  def test_raw_text(self):
+    self.assertEqual(
+        Template(
+            '{{a}}' + Template.raw_str('\n{{d}}, {%x%}\n') + '{{b}}',
+            a='hi', b=1
+        ).render().text,
+        'hi\n{{d}}, {%x%}\n1'
+    )
+
+  def test_from_raw_str(self):
+    self.assertEqual(
+        Template.from_raw_str('\n{{d}}, {%x%}\n').render().text,
+        '\n{{d}}, {%x%}\n'
+    )
+
 
 class DefinitionTest(unittest.TestCase):
 
@@ -308,9 +325,75 @@ class RenderTest(unittest.TestCase):
         Template(
             'This is {{ x }} and {{ a }}', x=1, a=CustomModality('foo')
         ).render(),
-        'This is 1 and {{a}}',
+        'This is 1 and <<[[a]]>>',
     )
 
+  def test_render_with_default(self):
+
+    class Foo(Template):
+      """Foo.
+
+      This is {{x}}
+      """
+
+    f = Foo(template_str='!{{DEFAULT}}!', x=1)
+    self.assertEqual(f.DEFAULT.x, 1)
+    self.assertEqual(
+        f.render(), '!This is 1!'
+    )
+
+    class Bar(Template):
+      """Bar.
+
+      {{preamble}}
+      {{prompt}}
+      """
+
+      preamble: Template = Template('You are a chat bot.')
+      prompt: Template = Template('User: hi! {{name}}')
+
+    b = Bar(
+        preamble=Template('<h1>{{DEFAULT}}</h1>'),
+        prompt=Template('<h2>{{DEFAULT}}</h2>'),
+        name='Tom',
+    )
+    # Test variable access.
+    self.assertEqual(
+        b.render(),
+        inspect.cleandoc("""
+            <h1>You are a chat bot.</h1>
+            <h2>User: hi! Tom</h2>
+            """),
+    )
+
+    with self.assertRaisesRegex(ValueError, '`{{ DEFAULT }}` cannot be used'):
+
+      class Baz(Template):  # pylint: disable=unused-variable
+        """Baz.
+
+        {{DEFAULT}}
+        """
+
+    with self.assertRaisesRegex(
+        ValueError, 'The template neither has a default `template_str` nor'
+    ):
+      Template('{{DEFAULT}}').render()
+
+    d = pg.Dict(x=Template('{{DEFAULT}}'))
+    with self.assertRaisesRegex(
+        ValueError, 'does not have a default value'
+    ):
+      _ = d.x.DEFAULT
+
+    class Tes(pg.Object):
+      x: str | None = None
+
+    t = Tes(x=Template('{{DEFAULT}}'))
+    with self.assertRaisesRegex(
+        ValueError, 'is not a `lf.Template` object or str'
+    ):
+      _ = t.x.DEFAULT
+
   def test_bad_render(self):
     with self.assertRaises(ValueError):
       Template('Hello {{x}}').render(allow_partial=False)
@@ -427,6 +510,14 @@ class RenderTest(unittest.TestCase):
     # Test len.
     self.assert_partial(Template('Hello {{len(x)}}'), 'Hello {{len(x)}}')
 
+  def test_additional_metadata(self):
+    t = Template('hi', metadata_weights=1.0, y=2)
+    self.assertEqual(t.render(), message_lib.UserMessage('hi', weights=1.0))
+
+    t = Template('hi')
+    with component.context(metadata_weights=1.0, y=2):
+      self.assertEqual(t.render(), message_lib.UserMessage('hi', weights=1.0))
+
 
 class TemplateRenderEventTest(unittest.TestCase):
 
@@ -462,5 +553,59 @@ class TemplateRenderEventTest(unittest.TestCase):
     self.assertEqual(render_stacks, [[l]])
 
 
+class HtmlTest(unittest.TestCase):
+
+  def test_html(self):
+
+    class Foo(Template):
+      """Template Foo.
+
+      {{x}} + {{y}} = ?
+      """
+      x: Any
+      y: Any
+
+    class Bar(Template):
+      """Template Bar.
+
+      {{y}} + {{z}}
+      """
+      y: Any
+
+    self.assertIn(
+        inspect.cleandoc(
+            """
+            /* Langfun Template styles. */
+            .template-str {
+                padding: 10px;
+                margin: 10px 5px 10px 5px;
+                font-style: italic;
+                font-size: 1.1em;
+                white-space: pre-wrap;
+                border: 1px solid #EEE;
+                border-radius: 5px;
+                background-color: #EEE;
+                color: #cc2986;
+            }
+            """
+        ),
+        Foo(x=1, y=2).to_html().style_section,
+    )
+    self.assertIn(
+        'template-str',
+        Foo(x=Bar('{{y}} + {{z}}'), y=1).to_html(
+            enable_summary_tooltip=False,
+        ).content,
+    )
+    self.assertIn(
+        'template-str',
+        Foo(x=Bar('{{y}} + {{z}}'), y=1).to_html(
+            enable_summary_tooltip=False,
+            collapse_level=0,
+            key_style='label',
+        ).content,
+    )
+
+
 if __name__ == '__main__':
   unittest.main()

langfun/core/templates/conversation.py

@@ -38,6 +38,11 @@ class Conversation(Completion):
       '(Optional) Preamble before beginning the conversation.',
   ] = None
 
+  role: Annotated[
+      str | None,
+      '(Optional) User defined role for the AI response in the conversation.',
+  ] = None
+
   conversation_context: Annotated[
       lf.LangFunc | None,
       (
@@ -71,6 +76,10 @@ class Conversation(Completion):
     with lf.context(**kwargs):
       # Call LM based on the prompt generated from `input_message`.
       lm_response = super().__call__()
+      if self.role is not None:
+        lm_response.rebind(
+            sender=self.role, skip_notification=True, raise_on_no_change=False
+        )
 
       # Add current turn to memory.
       self.add(self.input_message, lm_response)

langfun/core/templates/conversation_test.py

@@ -83,6 +83,7 @@ class ConversationTest(unittest.TestCase):
   def test_call(self):
     c = Conversation(
         lm=QuestionCounter(),
+        role='Agent',
         preamble="You are a helpful and joyful AI bot. Now let's chat.",
     )
     # First round.
@@ -102,7 +103,7 @@ class ConversationTest(unittest.TestCase):
         inspect.cleandoc("""
         You are a helpful and joyful AI bot. Now let's chat.
         User: Hello
-        AI: Response 1.
+        Agent: Response 1.
         User: How are you?
         """),
     )
@@ -114,9 +115,9 @@ class ConversationTest(unittest.TestCase):
         inspect.cleandoc("""
         You are a helpful and joyful AI bot. Now let's chat.
         User: Hello
-        AI: Response 1.
+        Agent: Response 1.
         User: How are you?
-        AI: Response 2.
+        Agent: Response 2.
         User: Okay, bye.
         """),
     )

langfun/core/templates/selfplay_test.py

@@ -56,7 +56,13 @@ class SelfPlayTest(unittest.TestCase):
     g = NumberGuess(target_num=10)
 
     with lf.context(lm=NumberGuesser(guesses=[50, 20, 5, 10])):
-      self.assertEqual(g(), lf.AIMessage('10', score=0.0, logprobs=None))
+      self.assertEqual(
+          g(),
+          lf.AIMessage(
+              '10', score=0.0, logprobs=None, is_cached=False,
+              usage=lf.UsageNotAvailable()
+          )
+      )
 
     self.assertEqual(g.num_turns, 4)
 
@@ -64,7 +70,13 @@ class SelfPlayTest(unittest.TestCase):
     g = NumberGuess(target_num=10, max_turns=10)
 
     with lf.context(lm=NumberGuesser(guesses=[50, 20, 5, 2, 5, 4])):
-      self.assertEqual(g(), lf.AIMessage('2', score=0.0, logprobs=None))
+      self.assertEqual(
+          g(),
+          lf.AIMessage(
+              '2', score=0.0, logprobs=None, is_cached=False,
+              usage=lf.UsageNotAvailable()
+          )
+      )
 
     self.assertEqual(g.num_turns, 10)