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

langfun/core/structured/mapping.py

@@ -13,15 +13,55 @@
 # limitations under the License.
 """The base of symbolic mapping methods."""
 
+import functools
 import io
-from typing import Annotated, Any
+from typing import Annotated, Any, Callable
 import langfun.core as lf
 from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
+class MappingError(Exception):  # pylint: disable=g-bad-exception-name
+  """Mapping error."""
+
+  def __init__(self, lm_response: lf.Message, cause: Exception):
+    self._lm_response = lm_response
+    self._cause = cause
+
+  @property
+  def lm_response(self) -> lf.Message:
+    """Returns the LM response that failed to be mapped."""
+    return self._lm_response
+
+  @property
+  def cause(self) -> Exception:
+    """Returns the cause of the error."""
+    return self._cause
+
+  def __str__(self) -> str:
+    return self.format(include_lm_response=True)
+
+  def format(self, include_lm_response: bool = True) -> str:
+    """Formats the mapping error."""
+    r = io.StringIO()
+    error_message = str(self.cause).rstrip()
+    r.write(
+        pg.colored(
+            f'{self.cause.__class__.__name__}: {error_message}', 'magenta'
+        )
+    )
+    if include_lm_response:
+      r.write('\n\n')
+      r.write(pg.colored('[LM Response]', 'blue', styles=['bold']))
+      r.write('\n')
+      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[
@@ -55,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:
@@ -70,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
@@ -79,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)
 
@@ -110,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.
@@ -206,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) }}
 
@@ -233,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.',
@@ -278,6 +380,14 @@ class Mapping(lf.LangFunc):
       ),
   ] = lf.RAISE_IF_HAS_ERROR
 
+  response_postprocess: Annotated[
+      Callable[[str], str] | None,
+      (
+          'A callable object that post process the raw LLM response before '
+          'parsing it into the output Python object.'
+      )
+  ] = None
+
   #
   # Key methods for implementing specific mappings.
   #
@@ -296,10 +406,17 @@ class Mapping(lf.LangFunc):
   def transform_output(self, lm_output: lf.Message) -> lf.Message:
     """Transforms LM response into structure if schema is present."""
     try:
+      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 e
+        raise MappingError(lm_output, e) from e
       lm_output.result = self.default
     return lm_output
 
@@ -316,6 +433,14 @@ class Mapping(lf.LangFunc):
         autofix_lm=self.autofix_lm or self.lm,
     )
 
+  def postprocess_response(self, response: lf.Message) -> lf.Message:
+    """Post process LLM response."""
+    if self.response_postprocess is not None:
+      postprocessed_text = self.response_postprocess(response.text)
+      if postprocessed_text != response.text:
+        return lf.AIMessage(postprocessed_text, source=response)
+    return response
+
   def postprocess_result(self, result: Any) -> Any:
     """Post process structured output."""
     return result
@@ -324,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,12 +14,30 @@
 """Tests for structured mapping example."""
 
 import inspect
+from typing import Any
 import unittest
 
+import langfun.core as lf
 from langfun.core.structured import mapping
 import pyglove as pg
 
 
+class MappingErrorTest(unittest.TestCase):
+
+  def test_format(self):
+    error = mapping.MappingError(
+        lf.AIMessage('hi'), ValueError('Cannot parse message.')
+    )
+    self.assertEqual(
+        pg.decolor(str(error)),
+        'ValueError: Cannot parse message.\n\n[LM Response]\nhi',
+    )
+    self.assertEqual(
+        pg.decolor(error.format(include_lm_response=False)),
+        'ValueError: Cannot parse message.',
+    )
+
+
 class MappingExampleTest(unittest.TestCase):
 
   def test_basics(self):
@@ -112,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'
@@ -120,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}.')