langfun 0.0.2.dev20240429__py3-none-any.whl → 0.1.2.dev202501150804__py3-none-any.whl

langfun/core/structured/function_generation_test.py

@@ -63,6 +63,42 @@ class FunctionGenerationTest(unittest.TestCase):
 
     lm = fake.StaticSequence([unittest_lm_response, function_gen_lm_response])
 
+    @function_generation.function_gen(lm=lm, unittest='auto')
+    def linear_search(items, target):  # pylint: disable=unused-argument
+      """Performs a linear search on a list to find a target value.
+
+      Args:
+          items (list): The list to search within.
+          target: The value to search for.
+
+      Returns:
+          int: The index of the target value if found, otherwise -1.
+      """
+
+    self.assertEqual(linear_search(['a', 'b', 'c'], 'c'), 2)
+    self.assertEqual(linear_search.source(), function_gen_lm_response)
+
+  def test_generate_function_without_unittest(self):
+    function_gen_lm_response = inspect.cleandoc("""
+        def linear_search(items, target):
+            \"\"\"
+            Performs a linear search on a list to find a target value.
+
+            Args:
+                items (list): The list to search within.
+                target: The value to search for.
+
+            Returns:
+                int: The index of the target value if found, otherwise -1.
+            \"\"\"
+            for i, item in enumerate(items):
+                if item == target:
+                    return i
+            return -1
+        """)
+
+    lm = fake.StaticSequence([function_gen_lm_response])
+
     @function_generation.function_gen(lm=lm)
     def linear_search(items, target):  # pylint: disable=unused-argument
       """Performs a linear search on a list to find a target value.
@@ -258,7 +294,9 @@ class FunctionGenerationTest(unittest.TestCase):
     cache_file = os.path.join(cache_file_dir, 'cache_file.json')
 
     @function_generation.function_gen(
-        lm=lm, unittest=_unittest_fn, cache_filename=cache_file
+        lm=lm,
+        unittest=_unittest_fn,
+        cache_filename=cache_file,
     )
     def linear_search(items, target):  # pylint: disable=unused-argument
       """Performs a linear search on a list to find a target value.
@@ -273,6 +311,36 @@ class FunctionGenerationTest(unittest.TestCase):
 
     self.assertEqual(linear_search(['a', 'b', 'c'], 'c'), 2)
 
+  def test_context_passthrough(self):
+
+    class Number(pg.Object):
+      value: int
+
+    function_gen_lm_response = inspect.cleandoc("""
+        ```python
+        def add(a: Number, b: Number) -> Number:
+            \"\"\"Adds two numbers together.\"\"\"
+            return Number(a.value + b.value)
+        ```
+        """)
+
+    lm = fake.StaticSequence(
+        [function_gen_lm_response]
+    )
+
+    def _unittest_fn(func):
+      assert func(Number(1), Number(2)) == Number(3)
+
+    custom_unittest = _unittest_fn
+
+    @function_generation.function_gen(
+        lm=lm, unittest=custom_unittest, num_retries=1
+    )
+    def add(a: Number, b: Number) -> Number:  # pylint: disable=unused-argument
+      """Adds two numbers together."""
+
+    self.assertEqual(add(Number(2), Number(3)), Number(5))
+
   def test_siganture_check(self):
     incorrect_signature_lm_response = inspect.cleandoc("""
         ```python
@@ -310,7 +378,9 @@ class FunctionGenerationTest(unittest.TestCase):
 
     custom_unittest = _unittest_fn
 
-    @function_generation.function_gen(lm=lm, unittest=custom_unittest)
+    @function_generation.function_gen(
+        lm=lm, unittest=custom_unittest, num_retries=2
+    )
     def linear_search(items, target):  # pylint: disable=unused-argument
       """Performs a linear search on a list to find a target value.
 

langfun/core/structured/mapping.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """The base of symbolic mapping methods."""
 
+import functools
 import io
 from typing import Annotated, Any, Callable
 import langfun.core as lf
@@ -45,20 +46,22 @@ class MappingError(Exception):  # pylint: disable=g-bad-exception-name
     r = io.StringIO()
     error_message = str(self.cause).rstrip()
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {error_message}', 'magenta'
         )
     )
     if include_lm_response:
       r.write('\n\n')
-      r.write(lf.colored('[LM Response]', 'blue', styles=['bold']))
+      r.write(pg.colored('[LM Response]', 'blue', styles=['bold']))
       r.write('\n')
-      r.write(lf.colored(self.lm_response.text, 'blue'))
+      r.write(pg.colored(self.lm_response.text, 'blue'))
     return r.getvalue()
 
 
 @pg.use_init_args(['input', 'output', 'schema', 'context'])
-class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
+class MappingExample(lf.NaturalLanguageFormattable,
+                     lf.Component,
+                     pg.views.HtmlTreeView.Extension):
   """Mapping example between text, schema and structured value."""
 
   input: pg.typing.Annotated[
@@ -92,6 +95,15 @@ class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
       'The natural language context for this mapping. ',
   ] = None
 
+  metadata: Annotated[
+      dict[str, Any],
+      (
+          'The metadata associated with the mapping example, '
+          'which chould carry structured data, such as tool function input. '
+          'It is a `pg.Dict` object whose keys can be accessed by attributes.'
+      ),
+  ] = pg.Dict()
+
   def schema_repr(
       self, protocol: schema_lib.SchemaProtocol = 'python', **kwargs
   ) -> str:
@@ -107,7 +119,11 @@ class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
 
   @classmethod
   def value_repr(
-      cls, value: Any, protocol: schema_lib.SchemaProtocol = 'python', **kwargs
+      cls,
+      value: Any,
+      protocol: schema_lib.SchemaProtocol = 'python',
+      use_modality_ref: bool = False,
+      **kwargs
   ) -> str:
     if isinstance(value, str):
       return value
@@ -116,7 +132,7 @@ class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
         return str(value)
 
     # Placehold modalities if they are present.
-    if pg.contains(value, type=lf.Modality):
+    if use_modality_ref and pg.contains(value, type=lf.Modality):
       value = lf.ModalityRef.placehold(value)
     return schema_lib.value_repr(protocol).repr(value, **kwargs)
 
@@ -147,24 +163,83 @@ class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
   def natural_language_format(self) -> str:
     result = io.StringIO()
     if self.context:
-      result.write(lf.colored('[CONTEXT]\n', styles=['bold']))
-      result.write(lf.colored(self.context, color='magenta'))
+      result.write(pg.colored('[CONTEXT]\n', styles=['bold']))
+      result.write(pg.colored(self.context, color='magenta'))
       result.write('\n\n')
 
-    result.write(lf.colored('[INPUT]\n', styles=['bold']))
-    result.write(lf.colored(self.input_repr(), color='green'))
-    result.write('\n\n')
+    result.write(pg.colored('[INPUT]\n', styles=['bold']))
+    result.write(pg.colored(self.input_repr(), color='green'))
 
     if self.schema is not None:
-      result.write(lf.colored('[SCHEMA]\n', styles=['bold']))
-      result.write(lf.colored(self.schema_repr(), color='red'))
       result.write('\n\n')
+      result.write(pg.colored('[SCHEMA]\n', styles=['bold']))
+      result.write(pg.colored(self.schema_repr(), color='red'))
 
     if schema_lib.MISSING != self.output:
-      result.write(lf.colored('[OUTPUT]\n', styles=['bold']))
-      result.write(lf.colored(self.output_repr(), color='blue'))
+      result.write('\n\n')
+      result.write(pg.colored('[OUTPUT]\n', styles=['bold']))
+      result.write(pg.colored(self.output_repr(), color='blue'))
+
+    if self.metadata:
+      result.write('\n\n')
+      result.write(pg.colored('[METADATA]\n', styles=['bold']))
+      result.write(pg.colored(str(self.metadata), color='cyan'))
     return result.getvalue().strip()
 
+  @classmethod
+  @functools.cache
+  def _html_tree_view_config(cls) -> dict[str, Any]:
+
+    def render_value(view, *, value, **kwargs):
+      if isinstance(value, lf.Template):
+        # Make a shallow copy to make sure modalities are rooted by
+        # the input.
+        value = value.clone().render()
+      if value is None:
+        return None
+      return view.render(value, **kwargs)
+
+    return pg.views.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            include_keys=['input', 'output', 'context', 'schema', 'metadata'],
+            extra_flags=dict(
+                render_value_fn=render_value,
+            ),
+            child_config=dict(
+                input=dict(
+                    collapse_level=1,
+                ),
+                output=dict(
+                    css_classes=['lf-example-output'],
+                    collapse_level=1,
+                ),
+                schema=dict(
+                    css_classes=['lf-example-schema'],
+                    collapse_level=1,
+                ),
+                metadata=dict(
+                    css_classes=['lf-example-metadata'],
+                    collapse_level=1,
+                ),
+            ),
+        )
+    )
+
+  @classmethod
+  @functools.cache
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        .lf-example-output {
+            color: dodgerblue;
+        }
+        .lf-example-schema {
+            color: blue;
+        }
+        """
+    ]
+
 
 class Mapping(lf.LangFunc):
   """Base class for mapping.
@@ -243,13 +318,7 @@ class Mapping(lf.LangFunc):
       {{ input_title }}:
       {{ example.input_repr(protocol, compact=False) | indent(2, True) }}
 
-      {% if has_modality_refs(example.input) -%}
-      {{ modality_refs_title }}:
-      {{ modality_refs_repr(example.input) | indent(2, True) }}
-
-      {% endif -%}
-
-      {%- if example.schema -%}
+      {% if example.schema -%}
       {{ schema_title }}:
       {{ example.schema_repr(protocol) | indent(2, True) }}
 
@@ -270,10 +339,6 @@ class Mapping(lf.LangFunc):
 
   schema_title: Annotated[str, 'The section title for schema.'] = 'SCHEMA'
 
-  modality_refs_title: Annotated[
-      str, 'The section title for modality refs.'
-  ] = 'MODALITY_REFERENCES'
-
   protocol: Annotated[
       schema_lib.SchemaProtocol,
       'The protocol for representing the schema and value.',
@@ -344,6 +409,12 @@ class Mapping(lf.LangFunc):
       lm_output = self.postprocess_response(lm_output)
       lm_output.result = self.postprocess_result(self.parse_result(lm_output))
     except Exception as e:  # pylint: disable=broad-exception-caught
+      if (self.lm.cache is not None
+          and lm_output.lm_input.cache_seed is not None):
+        success = self.lm.cache.delete(
+            self.lm, lm_output.lm_input, lm_output.lm_input.cache_seed
+        )
+        assert success
       if self.default == lf.RAISE_IF_HAS_ERROR:
         raise MappingError(lm_output, e) from e
       lm_output.result = self.default
@@ -378,24 +449,3 @@ class Mapping(lf.LangFunc):
     """Gets additional symbol definitions besides schema as globals."""
     return {'ModalityRef': lf.modality.ModalityRef}
 
-  #
-  # Helper methods for handling modalities.
-  #
-
-  def has_modality_refs(self, value: Any) -> bool:
-    """Returns true if the value has modalities."""
-    return not isinstance(value, lf.Modality) and pg.contains(
-        value, type=lf.Modality
-    )
-
-  def modalities(self, value: Any) -> dict[str, lf.Modality]:
-    return lf.Modality.from_value(value)
-
-  def modality_refs_repr(self, value: Any) -> str:
-    with lf.modality.format_modality_as_ref(True):
-      return pg.format(
-          self.modalities(value),
-          compact=False,
-          verbose=False,
-          python_format=True,
-      )

langfun/core/structured/mapping_test.py

@@ -14,6 +14,7 @@
 """Tests for structured mapping example."""
 
 import inspect
+from typing import Any
 import unittest
 
 import langfun.core as lf
@@ -28,11 +29,11 @@ class MappingErrorTest(unittest.TestCase):
         lf.AIMessage('hi'), ValueError('Cannot parse message.')
     )
     self.assertEqual(
-        lf.text_formatting.decolored(str(error)),
+        pg.decolor(str(error)),
         'ValueError: Cannot parse message.\n\n[LM Response]\nhi',
     )
     self.assertEqual(
-        lf.text_formatting.decolored(error.format(include_lm_response=False)),
+        pg.decolor(error.format(include_lm_response=False)),
         'ValueError: Cannot parse message.',
     )
 
@@ -129,6 +130,33 @@ class MappingExampleTest(unittest.TestCase):
             """),
     )
 
+  def test_str_with_metadata(self):
+    self.assertEqual(
+        str(
+            mapping.MappingExample(
+                '1 + 1 = 2',
+                schema=int,
+                context='Give the answer.',
+                metadata={'foo': 'bar'},
+            )
+        ),
+        inspect.cleandoc("""
+            \x1b[1m[CONTEXT]
+            \x1b[0m\x1b[35mGive the answer.\x1b[0m
+
+            \x1b[1m[INPUT]
+            \x1b[0m\x1b[32m1 + 1 = 2\x1b[0m
+
+            \x1b[1m[SCHEMA]
+            \x1b[0m\x1b[31mint\x1b[0m
+
+            \x1b[1m[METADATA]
+            \x1b[0m\x1b[36m{
+              foo = 'bar'
+            }\x1b[0m
+            """),
+    )
+
   def test_serialization(self):
     example = mapping.MappingExample(
         'the answer is 2', 2, int, context='compute 1 + 1'
@@ -137,6 +165,66 @@ class MappingExampleTest(unittest.TestCase):
         pg.eq(pg.from_json_str(example.to_json_str()), example)
     )
 
+  def assert_html_content(self, html, expected):
+    expected = inspect.cleandoc(expected).strip()
+    actual = html.content.strip()
+    if actual != expected:
+      print(actual)
+    self.assertEqual(actual, expected)
+
+  def test_html(self):
+
+    class Answer(pg.Object):
+      answer: int
+
+    class Addition(lf.Template):
+      """Template Addition.
+
+      {{x}} + {{y}} = ?
+      """
+      x: Any
+      y: Any
+
+    example = mapping.MappingExample(
+        input=Addition(x=1, y=2),
+        schema=Answer,
+        context='compute 1 + 1',
+        output=Answer(answer=3),
+        metadata={'foo': 'bar'},
+    )
+    self.assert_html_content(
+        example.to_html(
+            enable_summary_tooltip=False,
+            extra_flags=dict(
+                include_message_metadata=False
+            )
+        ),
+        """
+        <details open class="pyglove mapping-example"><summary><div class="summary-title">MappingExample(...)</div></summary><div class="complex-value mapping-example"><details open class="pyglove user-message lf-message"><summary><div class="summary-name lf-message">input<span class="tooltip lf-message">input</span></div><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>rendered</span></div><div class="message-text">1 + 2 = ?</div></div></details><details open class="pyglove answer lf-example-output"><summary><div class="summary-name lf-example-output">output<span class="tooltip lf-example-output">output</span></div><div class="summary-title lf-example-output">Answer(...)</div></summary><div class="complex-value answer"><details open class="pyglove int"><summary><div class="summary-name">answer<span class="tooltip">output.answer</span></div><div class="summary-title">int</div></summary><span class="simple-value int">3</span></details></div></details><details open class="pyglove str"><summary><div class="summary-name">context<span class="tooltip">context</span></div><div class="summary-title">str</div></summary><span class="simple-value str">&#x27;compute 1 + 1&#x27;</span></details><details open class="pyglove schema lf-example-schema"><summary><div class="summary-name lf-example-schema">schema<span class="tooltip lf-example-schema">schema</span></div><div class="summary-title lf-example-schema">Schema(...)</div></summary><div class="lf-schema-definition">Answer
+
+        ```python
+        class Answer:
+          answer: int
+        ```</div></details><details open class="pyglove dict lf-example-metadata"><summary><div class="summary-name lf-example-metadata">metadata<span class="tooltip lf-example-metadata">metadata</span></div><div class="summary-title lf-example-metadata">Dict(...)</div></summary><div class="complex-value dict"><details open class="pyglove str"><summary><div class="summary-name">foo<span class="tooltip">metadata.foo</span></div><div class="summary-title">str</div></summary><span class="simple-value str">&#x27;bar&#x27;</span></details></div></details></div></details>
+        """
+    )
+
+    example = mapping.MappingExample(
+        input=Addition(x=1, y=2),
+        output=Answer(answer=3),
+    )
+    self.assert_html_content(
+        example.to_html(
+            enable_summary_tooltip=False,
+            extra_flags=dict(
+                include_message_metadata=False
+            )
+        ),
+        """
+        <details open class="pyglove mapping-example"><summary><div class="summary-title">MappingExample(...)</div></summary><div class="complex-value mapping-example"><details open class="pyglove user-message lf-message"><summary><div class="summary-name lf-message">input<span class="tooltip lf-message">input</span></div><div class="summary-title lf-message">UserMessage(...)</div></summary><div class="complex_value"><div class="message-tags"><span>rendered</span></div><div class="message-text">1 + 2 = ?</div></div></details><details open class="pyglove answer lf-example-output"><summary><div class="summary-name lf-example-output">output<span class="tooltip lf-example-output">output</span></div><div class="summary-title lf-example-output">Answer(...)</div></summary><div class="complex-value answer"><details open class="pyglove int"><summary><div class="summary-name">answer<span class="tooltip">output.answer</span></div><div class="summary-title">int</div></summary><span class="simple-value int">3</span></details></div></details><details open class="pyglove contextual-attribute lf-example-schema"><summary><div class="summary-name lf-example-schema">schema<span class="tooltip lf-example-schema">schema</span></div><div class="summary-title lf-example-schema">ContextualAttribute(...)</div></summary><span class="simple-value none-type">None</span></details><details open class="pyglove dict lf-example-metadata"><summary><div class="summary-name lf-example-metadata">metadata<span class="tooltip lf-example-metadata">metadata</span></div><div class="summary-title lf-example-metadata">Dict(...)</div></summary><div class="complex-value dict"><span class="empty-container"></span></div></details></div></details>
+        """
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/structured/parsing.py

@@ -16,13 +16,13 @@ from typing import Any, Callable, Type, Union
 
 import langfun.core as lf
 from langfun.core.structured import mapping
-from langfun.core.structured import prompting
+from langfun.core.structured import querying
 from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
 @lf.use_init_args(['schema', 'default', 'examples'])
-class ParseStructure(mapping.Mapping):
+class _ParseStructure(mapping.Mapping):
   """Parse an object out from a natural language text."""
 
   context_title = 'USER_REQUEST'
@@ -37,7 +37,7 @@ class ParseStructure(mapping.Mapping):
   ]
 
 
-class ParseStructureJson(ParseStructure):
+class _ParseStructureJson(_ParseStructure):
   """Parse an object out from a NL text using JSON as the protocol."""
 
   preamble = """
@@ -53,7 +53,7 @@ class ParseStructureJson(ParseStructure):
   output_title = 'JSON'
 
 
-class ParseStructurePython(ParseStructure):
+class _ParseStructurePython(_ParseStructure):
   """Parse an object out from a NL text using Python as the protocol."""
 
   preamble = """
@@ -87,7 +87,7 @@ def parse(
     returns_message: bool = False,
     **kwargs,
 ) -> Any:
-  """Parse a natural langugage message based on schema.
+  """Parse a natural language message based on schema.
 
   Examples:
 
@@ -270,29 +270,41 @@ def call(
   if schema in (str, None):
     return lm_output if returns_message else lm_output.text
 
+  def _chain_nl_output_message(parsing_message: lf.Message):
+    """Chain the source of the parsed output to the LM output."""
+    parsing_message.root.source = lm_output
+    parsing_message.tag('parsing-lm-output')
+    parsing_message.lm_input.tag('parsing-lm-input')
+
   # Call `parsing_lm` for structured parsing.
-  return prompting.query(
-      lm_output,
-      schema,
-      examples=parsing_examples,
-      lm=parsing_lm or lm,
-      include_context=parsing_include_context,
-      cache_seed=cache_seed,
-      autofix=autofix,
-      autofix_lm=autofix_lm or lm,
-      protocol=protocol,
-      returns_message=returns_message,
-      **kwargs,
-  )
+  try:
+    parsing_message = querying.query(
+        lm_output.text,
+        schema,
+        examples=parsing_examples,
+        lm=parsing_lm or lm,
+        include_context=parsing_include_context,
+        cache_seed=cache_seed,
+        autofix=autofix,
+        autofix_lm=autofix_lm or lm,
+        protocol=protocol,
+        returns_message=True,
+        **kwargs,
+    )
+    _chain_nl_output_message(parsing_message)
+  except mapping.MappingError as e:
+    _chain_nl_output_message(e.lm_response)
+    raise e
+  return parsing_message if returns_message else parsing_message.result
 
 
 def _parse_structure_cls(
     protocol: schema_lib.SchemaProtocol,
-) -> Type[ParseStructure]:
+) -> Type[_ParseStructure]:
   if protocol == 'json':
-    return ParseStructureJson
+    return _ParseStructureJson
   elif protocol == 'python':
-    return ParseStructurePython
+    return _ParseStructurePython
   else:
     raise ValueError(f'Unknown protocol: {protocol!r}.')
 

langfun/core/structured/parsing_test.py

@@ -37,7 +37,7 @@ class Itinerary(pg.Object):
 class ParseStructurePythonTest(unittest.TestCase):
 
   def test_render_no_examples(self):
-    l = parsing.ParseStructurePython(int)
+    l = parsing._ParseStructurePython(int)
     m = lf.AIMessage('12 / 6 + 2 = 4')
     self.assertEqual(
         l.render(input=m, context='Compute 12 / 6 + 2.').text,
@@ -62,7 +62,7 @@ class ParseStructurePythonTest(unittest.TestCase):
     )
 
   def test_render_no_context(self):
-    l = parsing.ParseStructurePython(int)
+    l = parsing._ParseStructurePython(int)
     m = lf.AIMessage('12 / 6 + 2 = 4')
 
     self.assertEqual(
@@ -85,7 +85,7 @@ class ParseStructurePythonTest(unittest.TestCase):
     )
 
   def test_render(self):
-    l = parsing.ParseStructurePython(
+    l = parsing._ParseStructurePython(
         int,
         examples=[
             mapping.MappingExample(
@@ -212,7 +212,7 @@ class ParseStructurePythonTest(unittest.TestCase):
         ),
         override_attrs=True,
     ):
-      l = parsing.ParseStructurePython(
+      l = parsing._ParseStructurePython(
           [Itinerary],
           examples=[
               mapping.MappingExample(
@@ -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']
         ),
@@ -295,7 +295,7 @@ class ParseStructurePythonTest(unittest.TestCase):
 class ParseStructureJsonTest(unittest.TestCase):
 
   def test_render_no_examples(self):
-    l = parsing.ParseStructureJson(int)
+    l = parsing._ParseStructureJson(int)
     m = lf.AIMessage('12 / 6 + 2 = 4')
     self.assertEqual(
         l.render(input=m, context='Compute 12 / 6 + 2.').text,
@@ -320,7 +320,7 @@ class ParseStructureJsonTest(unittest.TestCase):
     )
 
   def test_render_no_context(self):
-    l = parsing.ParseStructureJson(int)
+    l = parsing._ParseStructureJson(int)
     m = lf.AIMessage('12 / 6 + 2 = 4')
 
     self.assertEqual(
@@ -343,7 +343,7 @@ class ParseStructureJsonTest(unittest.TestCase):
     )
 
   def test_render(self):
-    l = parsing.ParseStructureJson(
+    l = parsing._ParseStructureJson(
         int,
         examples=[
             mapping.MappingExample(
@@ -504,7 +504,7 @@ class ParseStructureJsonTest(unittest.TestCase):
         override_attrs=True,
     ):
       message = lf.LangFunc(lm_input)()
-      l = parsing.ParseStructureJson(
+      l = parsing._ParseStructureJson(
           [Itinerary],
           examples=[
               mapping.MappingExample(
@@ -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']
         ),
@@ -669,6 +670,49 @@ class CallTest(unittest.TestCase):
         3,
     )
 
+  def test_call_with_parsing_message_chaining(self):
+    output = parsing.call(
+        'Compute 1 + 2',
+        int,
+        lm=fake.StaticSequence(['three']),
+        parsing_lm=fake.StaticSequence(['3']),
+        parsing_examples=[
+            mapping.MappingExample(
+                context='Multiple four and five',
+                input='twenty',
+                schema=int,
+                output=20,
+            )
+        ],
+        returns_message=True,
+    )
+    self.assertIn('parsing-lm-output', output.tags)
+    self.assertIn('parsing-lm-input', output.source.tags)
+    self.assertEqual(output.root.text, 'Compute 1 + 2')
+
+  def test_call_with_parsing_message_chaining_on_parsing_error(self):
+    try:
+      output = parsing.call(
+          'Compute 1 + 2',
+          int,
+          lm=fake.StaticSequence(['three']),
+          parsing_lm=fake.StaticSequence(['abc']),
+          parsing_examples=[
+              mapping.MappingExample(
+                  context='Multiple four and five',
+                  input='twenty',
+                  schema=int,
+                  output=20,
+              )
+          ],
+          returns_message=True,
+      )
+    except mapping.MappingError as e:
+      output = e.lm_response
+    self.assertIn('parsing-lm-output', output.tags)
+    self.assertIn('parsing-lm-input', output.source.tags)
+    self.assertEqual(output.root.text, 'Compute 1 + 2')
+
   def test_call_with_autofix(self):
     lm = fake.StaticSequence(
         [