langfun 0.0.2.dev20240425__py3-none-any.whl → 0.0.2.dev20240426__py3-none-any.whl

langfun/core/eval/base.py

@@ -540,7 +540,7 @@ class Evaluable(lf.Component):
           f'<div style="color: {text_color}; white-space: pre-wrap;'
           'padding: 10px; border: 1px solid; margin-top: 10px">'
       )
-      s.write(m.text)
+      s.write(m.get('formatted_text', m.text))
       if m.result is not None:
         s.write(
             '<div style="color: magenta; white-space: pre-wrap;'

langfun/core/langfunc.py

@@ -261,7 +261,6 @@ class LangFunc(
         if lm_input is None:
           lm_input = self.render(**kwargs)
 
-        lm_input.tag(message_lib.Message.TAG_LM_INPUT)
         if skip_lm:
           return lm_input
 
@@ -270,10 +269,6 @@ class LangFunc(
         # Send rendered text to LM.
         lm_output = self.lm(lm_input, cache_seed=cache_seed)
 
-        # Track the input as the source of the output.
-        lm_output.source = lm_input
-        lm_output.tag(message_lib.Message.TAG_LM_RESPONSE)
-
         # Transform the output message.
         lm_output = self.transform_output(lm_output)
         lm_output.tag(message_lib.Message.TAG_LM_OUTPUT)

langfun/core/language_model.py

@@ -346,9 +346,42 @@ class LanguageModel(component.Component):
 
     with component.context(override_attrs=True, **kwargs):
       if self.cache is None:
-        return self._sample(prompts)
+        results = self._sample(prompts)
       else:
-        return self._sample_with_cache_lookup(prompts, cache_seed)
+        results = self._sample_with_cache_lookup(prompts, cache_seed)
+
+      for prompt, result in zip(prompts, results):
+
+        # Tag LM input.
+        prompt.tag(message_lib.Message.TAG_LM_INPUT)
+
+        for sample in result.samples:
+          # Update metadata for response message.
+
+          response = sample.response
+          response.metadata.score = sample.score
+          response.metadata.logprobs = sample.logprobs
+
+          # NOTE(daiyip): Current usage is computed at per-result level,
+          # which is accurate when n=1. For n > 1, we average the usage across
+          # multiple samples.
+          usage = result.usage
+          if len(result.samples) == 1 or usage is None:
+            response.metadata.usage = usage
+          else:
+            n = len(result.samples)
+            response.metadata.usage = LMSamplingUsage(
+                prompt_tokens=usage.prompt_tokens // n,
+                completion_tokens=usage.completion_tokens // n,
+                total_tokens=usage.total_tokens // n,
+            )
+
+          # Track the prompt for corresponding response.
+          response.source = prompt
+
+          # Tag LM response.
+          response.tag(message_lib.Message.TAG_LM_RESPONSE)
+      return results
 
   def _sample_with_cache_lookup(
       self, prompts: list[str | message_lib.Message], cache_seed: int
@@ -436,13 +469,8 @@ class LanguageModel(component.Component):
       result = self.sample(
           [prompt], sampling_options=sampling_options, cache_seed=cache_seed
       )[0]
-      response = result.samples[0].response
-      logprobs = result.samples[0].logprobs
-      response.set('score', result.samples[0].score)
-      response.metadata.logprobs = logprobs
-      response.metadata.usage = result.usage
-
       elapse = time.time() - request_start
+      response = result.samples[0].response
       self._debug(prompt, response, call_counter, result.usage, elapse)
       return response
 
@@ -494,7 +522,9 @@ class LanguageModel(component.Component):
       title_suffix = console.colored(f' ({usage.prompt_tokens} tokens)', 'red')
 
     console.write(
-        prompt,
+        # We use metadata 'formatted_text' for scenarios where the prompt text
+        # is formatted by the LM.
+        prompt.get('formatted_text', prompt.text),
         title=f'\n[{call_counter}] PROMPT SENT TO LM{title_suffix}:',
         color='green',
     )

langfun/core/language_model_test.py

@@ -111,11 +111,35 @@ class LanguageModelTest(unittest.TestCase):
         lm.sample(prompts=['foo', 'bar']),
         [
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('foo', score=-1.0)],
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'foo',
+                            score=-1.0,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=-1.0,
+                        logprobs=None,
+                    )
+                ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('bar', score=-1.0)],
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'bar',
+                            score=-1.0,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=-1.0,
+                        logprobs=None,
+                    )
+                ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
         ],
@@ -128,41 +152,119 @@ class LanguageModelTest(unittest.TestCase):
         ),
         [
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('foo' * 2, score=0.5)],
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'foo' * 2,
+                            score=0.5,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=0.5,
+                        logprobs=None,
+                    ),
+                ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('bar' * 2, score=0.5)],
-                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'bar' * 2,
+                            score=0.5,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=0.5,
+                        logprobs=None,
+                    ),
+                ],
+                usage=lm_lib.LMSamplingUsage(
+                    prompt_tokens=100, completion_tokens=100, total_tokens=200
+                ),
             ),
-        ],
+        ]
     )
     # Test override individual flags within sampling_options.
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar'], temperature=1.0),
         [
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('foo', score=1.0)],
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'foo',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    ),
+                ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('bar', score=1.0)],
-                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'bar',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    ),
+                ],
+                usage=lm_lib.LMSamplingUsage(
+                    prompt_tokens=100, completion_tokens=100, total_tokens=200
+                ),
             ),
-        ],
+        ]
     )
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar'], top_k=2, temperature=0.7),
         [
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('foo' * 2, score=0.7)],
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'foo' * 2,
+                            score=0.7,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=0.7,
+                        logprobs=None,
+                    ),
+                ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('bar' * 2, score=0.7)],
-                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'bar' * 2,
+                            score=0.7,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=0.7,
+                        logprobs=None,
+                    ),
+                ],
+                usage=lm_lib.LMSamplingUsage(
+                    prompt_tokens=100, completion_tokens=100, total_tokens=200
+                ),
             ),
-        ],
+        ]
     )
 
   def test_call(self):
@@ -189,7 +291,16 @@ class LanguageModelTest(unittest.TestCase):
             lm_lib.LMSamplingResult(
                 [
                     lm_lib.LMSample(
-                        message_lib.AIMessage('foo', cache_seed=0), score=-1.0
+                        message_lib.AIMessage(
+                            'foo',
+                            cache_seed=0,
+                            score=-1.0,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=-1.0,
+                        logprobs=None,
                     )
                 ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),
@@ -197,7 +308,16 @@ class LanguageModelTest(unittest.TestCase):
             lm_lib.LMSamplingResult(
                 [
                     lm_lib.LMSample(
-                        message_lib.AIMessage('bar', cache_seed=0), score=-1.0
+                        message_lib.AIMessage(
+                            'bar',
+                            cache_seed=0,
+                            score=-1.0,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=-1.0,
+                        logprobs=None,
                     )
                 ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),
@@ -225,7 +345,16 @@ class LanguageModelTest(unittest.TestCase):
             lm_lib.LMSamplingResult(
                 [
                     lm_lib.LMSample(
-                        message_lib.AIMessage('foo', cache_seed=0), score=1.0
+                        message_lib.AIMessage(
+                            'foo',
+                            cache_seed=0,
+                            score=1.0,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
                     )
                 ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),
@@ -233,7 +362,16 @@ class LanguageModelTest(unittest.TestCase):
             lm_lib.LMSamplingResult(
                 [
                     lm_lib.LMSample(
-                        message_lib.AIMessage('baz', cache_seed=0), score=1.0
+                        message_lib.AIMessage(
+                            'baz',
+                            cache_seed=0,
+                            score=1.0,
+                            logprobs=None,
+                            usage=lm_lib.LMSamplingUsage(100, 100, 200),
+                            tags=[message_lib.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
                     )
                 ],
                 usage=lm_lib.LMSamplingUsage(100, 100, 200),

langfun/core/llms/fake_test.py

@@ -28,7 +28,19 @@ class EchoTest(unittest.TestCase):
         lm.sample(['hi']),
         [
             lf.LMSamplingResult(
-                [lf.LMSample('hi', 1.0)],
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'hi',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(2, 2, 4),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
                 lf.LMSamplingUsage(2, 2, 4))
         ]
     )
@@ -60,7 +72,19 @@ class StaticResponseTest(unittest.TestCase):
         lm.sample(['hi']),
         [
             lf.LMSamplingResult(
-                [lf.LMSample(canned_response, 1.0)],
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            canned_response,
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(2, 38, 40),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
                 usage=lf.LMSamplingUsage(2, 38, 40)
             )
         ],
@@ -69,7 +93,19 @@ class StaticResponseTest(unittest.TestCase):
         lm.sample(['Tell me a joke.']),
         [
             lf.LMSamplingResult(
-                [lf.LMSample(canned_response, 1.0)],
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            canned_response,
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(15, 38, 53),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
                 usage=lf.LMSamplingUsage(15, 38, 53)
             )
         ],
@@ -101,11 +137,35 @@ class StaticMappingTest(unittest.TestCase):
         lm.sample(['Hi', 'How are you?']),
         [
             lf.LMSamplingResult(
-                [lf.LMSample('Hello', 1.0)],
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'Hello',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(2, 5, 7),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
                 usage=lf.LMSamplingUsage(2, 5, 7)
             ),
             lf.LMSamplingResult(
-                [lf.LMSample('I am fine, how about you?', 1.0)],
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'I am fine, how about you?',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(12, 25, 37),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
                 usage=lf.LMSamplingUsage(12, 25, 37)
             )
         ]
@@ -126,11 +186,35 @@ class StaticSequenceTest(unittest.TestCase):
         lm.sample(['Hi', 'How are you?']),
         [
             lf.LMSamplingResult(
-                [lf.LMSample('Hello', 1.0)],
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'Hello',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(2, 5, 7),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
                 usage=lf.LMSamplingUsage(2, 5, 7)
             ),
             lf.LMSamplingResult(
-                [lf.LMSample('I am fine, how about you?', 1.0)],
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'I am fine, how about you?',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(12, 25, 37),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
                 usage=lf.LMSamplingUsage(12, 25, 37)
             )
         ]

langfun/core/llms/openai_test.py

@@ -184,23 +184,96 @@ class OpenAITest(unittest.TestCase):
         results[0],
         lf.LMSamplingResult(
             [
-                lf.LMSample('Sample 0 for prompt 0.', score=0.0),
-                lf.LMSample('Sample 1 for prompt 0.', score=0.1),
-                lf.LMSample('Sample 2 for prompt 0.', score=0.2),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 0 for prompt 0.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for prompt 0.',
+                        score=0.1,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.1,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 2 for prompt 0.',
+                        score=0.2,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.2,
+                    logprobs=None,
+                ),
             ],
             usage=lf.LMSamplingUsage(
                 prompt_tokens=100, completion_tokens=100, total_tokens=200
             ),
         ),
     )
-
     self.assertEqual(
         results[1],
-        lf.LMSamplingResult([
-            lf.LMSample('Sample 0 for prompt 1.', score=0.0),
-            lf.LMSample('Sample 1 for prompt 1.', score=0.1),
-            lf.LMSample('Sample 2 for prompt 1.', score=0.2),
-        ]),
+        lf.LMSamplingResult(
+            [
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 0 for prompt 1.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=None,
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for prompt 1.',
+                        score=0.1,
+                        logprobs=None,
+                        usage=None,
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.1,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 2 for prompt 1.',
+                        score=0.2,
+                        logprobs=None,
+                        usage=None,
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.2,
+                    logprobs=None,
+                ),
+            ],
+        ),
     )
 
   def test_sample_chat_completion(self):
@@ -216,9 +289,51 @@ class OpenAITest(unittest.TestCase):
         results[0],
         lf.LMSamplingResult(
             [
-                lf.LMSample('Sample 0 for message.', score=0.0),
-                lf.LMSample('Sample 1 for message.', score=0.0),
-                lf.LMSample('Sample 2 for message.', score=0.0),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 0 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 2 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
             ],
             usage=lf.LMSamplingUsage(
                 prompt_tokens=100, completion_tokens=100, total_tokens=200
@@ -229,9 +344,51 @@ class OpenAITest(unittest.TestCase):
         results[1],
         lf.LMSamplingResult(
             [
-                lf.LMSample('Sample 0 for message.', score=0.0),
-                lf.LMSample('Sample 1 for message.', score=0.0),
-                lf.LMSample('Sample 2 for message.', score=0.0),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 0 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 2 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
             ],
             usage=lf.LMSamplingUsage(
                 prompt_tokens=100, completion_tokens=100, total_tokens=200
@@ -251,8 +408,36 @@ class OpenAITest(unittest.TestCase):
         results[0],
         lf.LMSamplingResult(
             [
-                lf.LMSample('Sample 0 for prompt 0.', score=0.0),
-                lf.LMSample('Sample 1 for prompt 0.', score=0.1),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 0 for prompt 0.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=50,
+                            completion_tokens=50,
+                            total_tokens=100,
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for prompt 0.',
+                        score=0.1,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=50,
+                            completion_tokens=50,
+                            total_tokens=100,
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.1,
+                    logprobs=None,
+                ),
             ],
             usage=lf.LMSamplingUsage(
                 prompt_tokens=100, completion_tokens=100, total_tokens=200

langfun/core/structured/prompting.py

@@ -176,8 +176,11 @@ def query(
       returning the structured `message.result`.
     skip_lm: If True, returns the rendered prompt as a UserMessage object.
       otherwise return the LLM response based on the rendered prompt.
-    **kwargs: Keyword arguments passed to the
-      `lf.structured.NaturalLanguageToStructureed` transform.
+    **kwargs: Keyword arguments passed to render the prompt or configure the 
+      `lf.structured.Mapping` class. Notable kwargs are:
+      - template_str: Change the root template for query.
+      - preamble: Change the preamble for query.
+      - mapping_template: Change the template for each mapping examle.
 
   Returns:
     The result based on the schema.
@@ -201,10 +204,17 @@ def query(
     return output if returns_message else output.text
 
   # Query with structured output.
+  prompt_kwargs = kwargs.copy()
+
+  # NOTE(daiyip): when `template_str` is passed in, it's intended to modify the
+  # QueryStructure template string. Therefore, we pop out the argument for
+  # prompt rendering.
+  prompt_kwargs.pop('template_str', None)
+
   if isinstance(prompt, str):
-    prompt = lf.Template(prompt, **kwargs)
+    prompt = lf.Template(prompt, **prompt_kwargs)
   elif isinstance(prompt, lf.Template):
-    prompt = prompt.rebind(**kwargs)
+    prompt = prompt.rebind(**prompt_kwargs, raise_on_no_change=False)
 
   if isinstance(prompt, lf.Template):
     prompt = prompt.render(lm=lm)

langfun/core/structured/prompting_test.py

@@ -138,6 +138,39 @@ class QueryTest(unittest.TestCase):
         ),
     )
 
+  def test_str_to_structure_render_custom_template(self):
+    lm = fake.StaticResponse('1')
+    self.assert_render(
+        'What is {{x}} + {{y}}?',
+        int,
+        x=1,
+        y=2,
+        lm=lm.clone(),
+        template_str='!!{{ DEFAULT }}!!',
+        expected_snippet=(
+            '!!Please respond to the last INPUT_OBJECT with OUTPUT_OBJECT '
+            'according to OUTPUT_TYPE.\n\n'
+            'INPUT_OBJECT:\n  1 + 1 =\n\n'
+            'OUTPUT_TYPE:\n'
+            '  Answer\n\n'
+            '  ```python\n'
+            '  class Answer:\n'
+            '    final_answer: int\n'
+            '  ```\n\n'
+            'OUTPUT_OBJECT:\n'
+            '  ```python\n'
+            '  Answer(\n'
+            '    final_answer=2\n'
+            '  )\n'
+            '  ```\n\n'
+            'INPUT_OBJECT:\n'
+            '  What is 1 + 2?\n\n'
+            'OUTPUT_TYPE:\n'
+            '  int\n\n'
+            'OUTPUT_OBJECT:!!'
+        ),
+    )
+
   def test_str_to_str_render(self):
     lm = fake.StaticResponse('1')
     self.assert_render(

langfun/core/template.py

@@ -48,7 +48,12 @@ class Template(
     component.Component,
     pg.typing.CustomTyping,
 ):
-  """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,
@@ -101,6 +106,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):
@@ -153,7 +163,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)
@@ -398,6 +408,93 @@ 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
+
 
 # Register converter from str to LangFunc, therefore we can always
 # pass strs to attributes that accept LangFunc.

langfun/core/template_test.py

@@ -312,6 +312,72 @@ class RenderTest(unittest.TestCase):
         '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)

{langfun-0.0.2.dev20240425.dist-info → langfun-0.0.2.dev20240426.dist-info}/METADATA

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

{langfun-0.0.2.dev20240425.dist-info → langfun-0.0.2.dev20240426.dist-info}/RECORD

@@ -6,10 +6,10 @@ langfun/core/concurrent.py,sha256=TRc49pJ3HQro2kb5FtcWkHjhBm8UcgE8RJybU5cU3-0,24
 langfun/core/concurrent_test.py,sha256=mwFMZhDUdppnDr7vDSTwcbMHwrdsIoKJwRYNtl4ZWL4,15185
 langfun/core/console.py,sha256=bk5rNPNm9rMGW5YT2HixxU04p2umnoabn5SDz6Dqe88,2317
 langfun/core/console_test.py,sha256=5SYJdxpJGLgdSSQqqMPoA1X6jpsLD8rgcyk-EgI65oE,1077
-langfun/core/langfunc.py,sha256=WXdTc3QsmGD_n80KD9dFRr5MHpGZ9E_y_Rhtk4t9-3w,11852
+langfun/core/langfunc.py,sha256=bRujJfH4iTwKFtFxQf745uJkfltuFnPfOGLuP8ydcr4,11646
 langfun/core/langfunc_test.py,sha256=sQaKuZpGGmG80GRifhbxkj7nfzQLJKj4Vuw5y1s1K3U,8378
-langfun/core/language_model.py,sha256=mJfQ_Zqq9IyVyZUdYMQ1BPrpo4Gn8yxDJb_RghQFP_I,18911
-langfun/core/language_model_test.py,sha256=oWQjnyiJugSpHJKda-qLaSvmbm1sx_v-ZXrHvw_kNk4,14172
+langfun/core/language_model.py,sha256=dypSV3kr6BLC7hsvV1_QOiqrHUHWtOjQfyFqH79WZmU,20052
+langfun/core/language_model_test.py,sha256=T-itu7Li2smv2dkru0C0neCs2W4VJXlNTYahXU6jF54,19548
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
 langfun/core/message.py,sha256=QhvV9t5qaryPcruyxxcXi3gm9QDInkSldwTtK6sVJ3c,15734
 langfun/core/message_test.py,sha256=Z23pUM5vPnDrYkIIibe2KL73D5HKur_awI0ut_EQFQA,9501
@@ -21,8 +21,8 @@ langfun/core/sampling.py,sha256=vygWvgC8MFw0_AKNSmz-ywMXJYWf8cl0tI8QycvAmyI,5795
 langfun/core/sampling_test.py,sha256=U7PANpMsl9E_pa4_Y4FzesSjcwg-u-LKHGCWSgv-8FY,3663
 langfun/core/subscription.py,sha256=euawEuSZP-BHydaT-AQpfYFL0m5pWPGcW0upFhrojqc,10930
 langfun/core/subscription_test.py,sha256=Y4ZdbZEwm83YNZBxHff0QR4QUa4rdaNXA3_jfIcArBo,8717
-langfun/core/template.py,sha256=dr3tZCbXH2qWzigO_EFVHe0GDnnCu58Tru5Mvlzin4o,18447
-langfun/core/template_test.py,sha256=xty7PgdNhGpw7ZRZ6QGwhKZWG6dyRgI16Lg3p7IMLJg,13944
+langfun/core/template.py,sha256=FZByYq6mhVDjT4HJ3yY-_TUZ13BiURzTJSKLw6QyLY4,21462
+langfun/core/template_test.py,sha256=Mbv0dFjboGCVvbDkHD-HacZnlCi8Ku2Hpf2UjdwGSNo,15464
 langfun/core/text_formatting.py,sha256=ytjj7opnRJ6w-pkglL2CZUyfYDXLpNf65E42LBb31gc,5158
 langfun/core/text_formatting_test.py,sha256=nyKC6tn2L4hPJiqQHgxcbQsJJi4A4Nbj8FiO8iT6B80,1514
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
@@ -40,7 +40,7 @@ langfun/core/coding/python/parsing_test.py,sha256=9vAWF484kWIm6JZq8NFiMgKUDhXV-d
 langfun/core/coding/python/permissions.py,sha256=1QWGHvzL8MM0Ok_auQ9tURqZHtdOfJaDpBzZ29GUE-c,2544
 langfun/core/coding/python/permissions_test.py,sha256=w5EDb8QxpxgJyZkojyzVWQvDfg366zn99-g__6TbPQ0,2699
 langfun/core/eval/__init__.py,sha256=NSmPe2lxdxFoY4h8VkNyONPAFtOTUpK9WhmZRaqUgiI,1335
-langfun/core/eval/base.py,sha256=_XUc8KivNzgn39sJ0H5W-o_37yXrTBt_X_hKkckLCiA,60239
+langfun/core/eval/base.py,sha256=1svQoZ0C2DGCVLvr0Qt0TcrlJKtJptdoOBVAxkxnHoU,60264
 langfun/core/eval/base_test.py,sha256=g3lRp2dcq411cLYHpn8spI4feyv2nOccs5PlFBwav3g,22512
 langfun/core/eval/matching.py,sha256=Ks-L9vyMNDj4R8zFczzByT_4DK2wAFatyCZupdHzx_g,9932
 langfun/core/eval/matching_test.py,sha256=5Qs9ETaLoyNcJ43f-_bK2Bfe--2Y3U79DnSA55-l6pc,4932
@@ -50,7 +50,7 @@ langfun/core/llms/__init__.py,sha256=1bPg1QI8duOZCYINm-jWi094x0JtLmsk4KX60qIC_gs
 langfun/core/llms/anthropic.py,sha256=7W9YdPN3SlAFhAIQlihMkrpo7tTY_4NvD0KIlCrqcsk,8505
 langfun/core/llms/anthropic_test.py,sha256=TMM30myyEhwF99Le4RvJEXOn8RYl0q1FRkt9Q9nl1jk,5540
 langfun/core/llms/fake.py,sha256=b-Xk5IPTbUt-elsyzd_i3n1tqzc_kgETXrEvgJruSMk,2824
-langfun/core/llms/fake_test.py,sha256=ZlDQgL41EX3eYTfBQNp2nB2LciqCmtoHgCsGvW4XhwI,4184
+langfun/core/llms/fake_test.py,sha256=ipKfdOcuqVcJ8lDXVpnBVb9HHG0hAVkFkMoHpWjC2cI,7212
 langfun/core/llms/google_genai.py,sha256=n8zyJwh9UCTgb6-8LyvmjVNFGZQ4-zfzZ0ulkhHAnR8,8624
 langfun/core/llms/google_genai_test.py,sha256=_UcGTfl16-aDUlEWFC2W2F8y9jPUs53RBYA6MOCpGXw,7525
 langfun/core/llms/groq.py,sha256=NaGItVL_pkOpqPpI4bPGU27xLFRoaeizZ49v2s-4ERs,7844
@@ -58,7 +58,7 @@ langfun/core/llms/groq_test.py,sha256=M6GtlrsOvDun_j-sR8cPh4W_moHWZNSTiThu3kuwbb
 langfun/core/llms/llama_cpp.py,sha256=Y_KkMUf3Xfac49koMUtUslKl3h-HWp3-ntq7Jaa3bdo,2385
 langfun/core/llms/llama_cpp_test.py,sha256=ZxC6defGd_HX9SFRU9U4cJiQnBKundbOrchbXuC1Z2M,1683
 langfun/core/llms/openai.py,sha256=06nPhmw0zIA5Zqv3eqsrZtYLHnKwW7N8yt3LlFUFVpI,13247
-langfun/core/llms/openai_test.py,sha256=Yt_W6k8YXpT3bs0JroARofCGmn_Uq3u61LmZxqWS2DQ,8272
+langfun/core/llms/openai_test.py,sha256=MiLqBaYliAkWVEwOBmX3HTj_eAuWLv77q8-I3VyVEBU,14841
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=cFfYvOIUae842pncqCAsRvqXCk2AnAsRYVx0mcIoAeY,3338
 langfun/core/llms/cache/in_memory.py,sha256=YfFyJEhLs73cUiB0ZfhMxYpdE8Iuxxw-dvMFwGHTSHw,4742
@@ -84,8 +84,8 @@ langfun/core/structured/mapping.py,sha256=Vq3bQZWi4iYjcVn8D2kvPXTAm9jrQ-_1ueHLbX
 langfun/core/structured/mapping_test.py,sha256=PiXklMeIa8L6KtMi3ju7J9Y39gZy0hIGz-Oeq4A_7XE,3835
 langfun/core/structured/parsing.py,sha256=keoVqEfzAbdULh6GawWFsTQzU91MzJXYFZjXGXLaD8g,11492
 langfun/core/structured/parsing_test.py,sha256=34wDrXaQ-EYhJLfDL8mX9K53oQMSzh5pVYdKjnESmK8,20895
-langfun/core/structured/prompting.py,sha256=mOmCWNVMnBk4rI7KBlEm5kmusPXoAKiWcohhzaw-s2o,7427
-langfun/core/structured/prompting_test.py,sha256=1Ik-RWs2fixpUTUKoN7FfJX48hDaVScZ-E1CB7J7F3Y,19860
+langfun/core/structured/prompting.py,sha256=WICdX_figP2u98Hvg3BFSTF83nNKxM4x1STxkWe2_9Y,7925
+langfun/core/structured/prompting_test.py,sha256=vslaCAUikfwOvqsKzqs_oyEacrefFsr2SWSqu6OHi3w,20813
 langfun/core/structured/schema.py,sha256=mJXirgqx3N7SA9zBO_ISHrzcV-ZRshLhnMJyCcSjGjY,25057
 langfun/core/structured/schema_generation.py,sha256=U3nRQsqmMZg_qIVDh2fiY3K4JLfsAL1LcKzIFP1iXFg,5316
 langfun/core/structured/schema_generation_test.py,sha256=RM9s71kMNg2jTePwInkiW9fK1ACN37eyPeF8OII-0zw,2950
@@ -101,8 +101,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=DYVrkk7uNKCqJGEHH31HssU2BPuMItU1vJLzfcXIlYg,2156
-langfun-0.0.2.dev20240425.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240425.dist-info/METADATA,sha256=OZw2iKQkUx9uF-tFzARKjqCekvCCVYMXL00JvSpRxVE,3436
-langfun-0.0.2.dev20240425.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-langfun-0.0.2.dev20240425.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240425.dist-info/RECORD,,
+langfun-0.0.2.dev20240426.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240426.dist-info/METADATA,sha256=VIhHPTScfC_kHY3j2aCQ38CKW2BrPeSgDO37U2E2kYo,3436
+langfun-0.0.2.dev20240426.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+langfun-0.0.2.dev20240426.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240426.dist-info/RECORD,,