langfun 0.0.2.dev20240319__py3-none-any.whl → 0.0.2.dev20240429__py3-none-any.whl

langfun/core/structured/parsing.py

@@ -16,6 +16,7 @@ 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 schema as schema_lib
 import pyglove as pg
 
@@ -270,7 +271,7 @@ def call(
     return lm_output if returns_message else lm_output.text
 
   # Call `parsing_lm` for structured parsing.
-  return parse(
+  return prompting.query(
       lm_output,
       schema,
       examples=parsing_examples,

langfun/core/structured/parsing_test.py

@@ -17,11 +17,9 @@ import inspect
 import unittest
 
 import langfun.core as lf
-from langfun.core import coding
 from langfun.core.llms import fake
 from langfun.core.structured import mapping
 from langfun.core.structured import parsing
-from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
@@ -255,7 +253,7 @@ class ParseStructurePythonTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          coding.CodeError,
+          mapping.MappingError,
           'name .* is not defined',
       ):
         parsing.parse('three', int)
@@ -280,13 +278,15 @@ class ParseStructurePythonTest(unittest.TestCase):
         ),
         1,
     )
+    r = parsing.parse(
+        'the answer is 1', int, user_prompt='what is 0 + 1?', lm=lm,
+        returns_message=True
+    )
     self.assertEqual(
-        parsing.parse(
-            'the answer is 1', int, user_prompt='what is 0 + 1?', lm=lm,
-            returns_message=True
-        ),
+        r,
         lf.AIMessage(
             '1', score=1.0, result=1, logprobs=None,
+            usage=lf.LMSamplingUsage(652, 1, 653),
             tags=['lm-response', 'lm-output', 'transformed']
         ),
     )
@@ -544,7 +544,7 @@ class ParseStructureJsonTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          schema_lib.JsonError,
+          mapping.MappingError,
           'No JSON dict in the output',
       ):
         parsing.parse('three', int, protocol='json')
@@ -634,13 +634,18 @@ class CallTest(unittest.TestCase):
       )
 
   def test_call_with_returning_message(self):
+    r = parsing.call(
+        'Compute 1 + 2', int, lm=fake.StaticSequence(['three', '3']),
+        returns_message=True
+    )
     self.assertEqual(
-        parsing.call(
-            'Compute 1 + 2', int, lm=fake.StaticSequence(['three', '3']),
-            returns_message=True
-        ),
+        r,
         lf.AIMessage(
-            '3', result=3, score=1.0, logprobs=None,
+            '3',
+            result=3,
+            score=1.0,
+            logprobs=None,
+            usage=lf.LMSamplingUsage(315, 1, 316),
             tags=['lm-response', 'lm-output', 'transformed']
         ),
     )

langfun/core/structured/prompting.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Symbolic query."""
 
-from typing import Any, Type, Union
+from typing import Any, Callable, Type, Union
 
 import langfun.core as lf
 from langfun.core.structured import mapping
@@ -78,7 +78,9 @@ class QueryStructurePython(QueryStructure):
 
       {{ output_title }}:
         ```python
-        Answer(final_answer=2)
+        Answer(
+          final_answer=2
+        )
         ```
       """
   protocol = 'python'
@@ -107,6 +109,7 @@ def query(
     lm: lf.LanguageModel | None = None,
     examples: list[mapping.MappingExample] | None = None,
     cache_seed: int | None = 0,
+    response_postprocess: Callable[[str], str] | None = None,
     autofix: int = 0,
     autofix_lm: lf.LanguageModel | None = None,
     protocol: schema_lib.SchemaProtocol = 'python',
@@ -159,6 +162,9 @@ def query(
     cache_seed: Seed for computing cache key. The cache key is determined by a
       tuple of (lm, prompt, cache seed). If None, cache will be disabled for
       the query even cache is configured by the LM.
+    response_postprocess: An optional callable object to process the raw LM
+      response before parsing it into the final output object. If None, the
+      raw LM response will not be processed.
     autofix: Number of attempts to auto fix the generated code. If 0, autofix is
       disabled. Auto-fix is not supported for 'json' protocol.
     autofix_lm: The language model to use for autofix. If not specified, the
@@ -170,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.
@@ -188,13 +197,24 @@ def query(
     output = lf.LangFunc.from_value(prompt, **kwargs)(
         lm=lm, cache_seed=cache_seed, skip_lm=skip_lm
     )
+    if response_postprocess:
+      processed_text = response_postprocess(output.text)
+      if processed_text != output.text:
+        output = lf.AIMessage(processed_text, source=output)
     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)
@@ -206,6 +226,7 @@ def query(
       schema=schema,
       default=default,
       examples=examples,
+      response_postprocess=response_postprocess,
       autofix=autofix if protocol == 'python' else 0,
       **kwargs,
   )(

langfun/core/structured/prompting_test.py

@@ -17,12 +17,10 @@ import inspect
 import unittest
 
 import langfun.core as lf
-from langfun.core import coding
 from langfun.core import modalities
 from langfun.core.llms import fake
 from langfun.core.structured import mapping
 from langfun.core.structured import prompting
-from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
@@ -77,6 +75,7 @@ class QueryTest(unittest.TestCase):
             result=1,
             score=1.0,
             logprobs=None,
+            usage=lf.LMSamplingUsage(323, 1, 324),
             tags=['lm-response', 'lm-output', 'transformed'],
         ),
     )
@@ -116,12 +115,59 @@ class QueryTest(unittest.TestCase):
         y=2,
         lm=lm.clone(),
         expected_snippet=(
-            'Please respond to the last INPUT_OBJECT with OUTPUT_OBJECT'
-            ' according to OUTPUT_TYPE.\n\nINPUT_OBJECT:\n  1 + 1'
-            ' =\n\nOUTPUT_TYPE:\n  Answer\n\n  ```python\n  class Answer:\n   '
-            ' final_answer: int\n  ```\n\nOUTPUT_OBJECT:\n  ```python\n '
-            ' Answer(final_answer=2)\n  ```\n\nINPUT_OBJECT:\n  What is 1 +'
-            ' 2?\n\nOUTPUT_TYPE:\n  int\n\nOUTPUT_OBJECT:'
+            '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_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:!!'
         ),
     )
 
@@ -264,7 +310,9 @@ class QueryStructurePythonTest(unittest.TestCase):
 
             OUTPUT_OBJECT:
               ```python
-              Answer(final_answer=2)
+              Answer(
+                final_answer=2
+              )
               ```
 
             INPUT_OBJECT:
@@ -308,7 +356,9 @@ class QueryStructurePythonTest(unittest.TestCase):
 
             OUTPUT_OBJECT:
               ```python
-              Answer(final_answer=2)
+              Answer(
+                final_answer=2
+              )
               ```
 
             INPUT_OBJECT:
@@ -420,7 +470,7 @@ class QueryStructurePythonTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          coding.CodeError,
+          mapping.MappingError,
           'name .* is not defined',
       ):
         prompting.query('Compute 1 + 2', int)
@@ -436,6 +486,23 @@ class QueryStructurePythonTest(unittest.TestCase):
     ])
     self.assertEqual(prompting.query('what is 1 + 0', int, lm=lm, autofix=3), 1)
 
+  def test_response_postprocess(self):
+    with lf.context(
+        lm=fake.StaticResponse('<!-- some comment-->\n3'),
+        override_attrs=True,
+    ):
+      self.assertEqual(
+          prompting.query(
+              'Compute 1 + 2', response_postprocess=lambda x: x.split('\n')[1]),
+          '3'
+      )
+      self.assertEqual(
+          prompting.query(
+              'Compute 1 + 2', int,
+              response_postprocess=lambda x: x.split('\n')[1]),
+          3
+      )
+
 
 class QueryStructureJsonTest(unittest.TestCase):
 
@@ -641,7 +708,7 @@ class QueryStructureJsonTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          schema_lib.JsonError,
+          mapping.MappingError,
           'No JSON dict in the output',
       ):
         prompting.query('Compute 1 + 2', int, protocol='json')

langfun/core/structured/schema.py

@@ -55,10 +55,6 @@ def parse_value_spec(value) -> pg.typing.ValueSpec:
           ),
       ):
         raise ValueError(f'Unsupported schema specification: {v}')
-      if isinstance(spec, pg.typing.Object) and not issubclass(
-          spec.cls, pg.Symbolic
-      ):
-        raise ValueError(f'{v} must be a symbolic class to be parsable.')
       return spec
 
   return _parse_node(value)
@@ -208,7 +204,9 @@ def class_dependencies(
   if isinstance(value_or_spec, Schema):
     return value_or_spec.class_dependencies(include_subclasses)
 
-  if isinstance(value_or_spec, (pg.typing.ValueSpec, pg.symbolic.ObjectMeta)):
+  if inspect.isclass(value_or_spec) or isinstance(
+      value_or_spec, pg.typing.ValueSpec
+  ):
     value_or_spec = (value_or_spec,)
 
   if isinstance(value_or_spec, tuple):
@@ -216,7 +214,7 @@ def class_dependencies(
     for v in value_or_spec:
       if isinstance(v, pg.typing.ValueSpec):
         value_specs.append(v)
-      elif inspect.isclass(v) and issubclass(v, pg.Object):
+      elif inspect.isclass(v):
         value_specs.append(pg.typing.Object(v))
       else:
         raise TypeError(f'Unsupported spec type: {v!r}')
@@ -235,23 +233,20 @@ def class_dependencies(
 
   def _fill_dependencies(vs: pg.typing.ValueSpec, include_subclasses: bool):
     if isinstance(vs, pg.typing.Object):
-      if issubclass(vs.cls, pg.Object) and vs.cls not in seen:
+      if vs.cls not in seen:
         seen.add(vs.cls)
 
         # Add base classes as dependencies.
         for base_cls in vs.cls.__bases__:
           # We only keep track of user-defined symbolic classes.
-          if issubclass(
-              base_cls, pg.Object
-          ) and not base_cls.__module__.startswith('pyglove'):
+          if base_cls is not object and base_cls is not pg.Object:
             _fill_dependencies(
                 pg.typing.Object(base_cls), include_subclasses=False
             )
 
         # Add members as dependencies.
-        if hasattr(vs.cls, '__schema__'):
-          for field in vs.cls.__schema__.values():
-            _fill_dependencies(field.value, include_subclasses)
+        for field in _pg_schema(vs.cls).values():
+          _fill_dependencies(field.value, include_subclasses)
       _add_dependency(vs.cls)
 
       # Check subclasses if available.
@@ -364,17 +359,13 @@ def class_definition(
 ) -> str:
   """Returns the Python class definition."""
   out = io.StringIO()
-  if not issubclass(cls, pg.Object):
-    raise TypeError(
-        'Classes must be `pg.Object` subclasses to be used as schema. '
-        f'Encountered: {cls}.'
-    )
-  schema = cls.__schema__
+  schema = _pg_schema(cls)
   eligible_bases = []
   for base_cls in cls.__bases__:
-    if issubclass(base_cls, pg.Object):
+    if base_cls is not object:
       if include_pg_object_as_base or base_cls is not pg.Object:
         eligible_bases.append(base_cls.__name__)
+
   if eligible_bases:
     base_cls_str = ', '.join(eligible_bases)
     out.write(f'class {cls.__name__}({base_cls_str}):\n')
@@ -395,10 +386,12 @@ def class_definition(
   if schema.fields:
     for key, field in schema.items():
       if not isinstance(key, pg.typing.ConstStrKey):
-        raise TypeError(
+        pg.logging.warning(
             'Variable-length keyword arguments is not supported in '
-            f'structured parsing or query. Encountered: {field}'
+            f'structured parsing or query. Encountered: {cls}, Schema: {schema}'
         )
+        continue
+
       # Write field doc string as comments before the field definition.
       if field.description:
         for line in field.description.split('\n'):
@@ -839,3 +832,13 @@ class Unknown(pg.Object, pg.typing.CustomTyping):
 
 
 UNKNOWN = Unknown()
+
+
+def _pg_schema(cls: Type[Any]) -> pg.Schema:
+  """Returns PyGlove schema for the constructor of a class."""
+  schema = getattr(cls, '__schema__', None)
+  if schema is None:
+    schema = pg.symbolic.callable_schema(
+        cls.__init__, auto_typing=True, auto_doc=True, remove_self=True
+    )
+  return schema

langfun/core/structured/schema_generation.py

@@ -143,14 +143,14 @@ def generate_class(
 
 
 def classgen_example(
-    class_name: str, prompt: str | pg.Symbolic, cls: Type[Any]
+    prompt: str | pg.Symbolic, cls: Type[Any]
 ) -> mapping.MappingExample:
   """Creates a class generation example."""
   if isinstance(prompt, lf.Template):
     prompt = prompt.render()
   return mapping.MappingExample(
       input=prompt,
-      context=class_name,
+      context=cls.__name__,
       output=cls,
   )
 
@@ -168,7 +168,6 @@ def default_classgen_examples() -> list[mapping.MappingExample]:
 
   return [
       classgen_example(
-          'Solution',
           'How to evaluate an arithmetic expression?',
           Solution,
       )

langfun/core/structured/schema_generation_test.py

@@ -14,8 +14,8 @@
 import inspect
 import unittest
 
-import langfun.core.coding as lf_coding
 from langfun.core.llms import fake
+from langfun.core.structured import mapping
 from langfun.core.structured import schema_generation
 
 
@@ -92,7 +92,7 @@ class GenerateClassTest(unittest.TestCase):
     )
     self.assertIs(cls.__name__, 'B')
 
-    with self.assertRaises(lf_coding.CodeError):
+    with self.assertRaises(mapping.MappingError):
       schema_generation.generate_class(
           'Foo',
           'Generate a Foo class with a field pointing to another class A',

langfun/core/structured/schema_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Tests for structured parsing."""
 
+import dataclasses
 import inspect
 import typing
 import unittest
@@ -101,12 +102,7 @@ class SchemaTest(unittest.TestCase):
 
     self.assert_unsupported_annotation(typing.Type[int])
     self.assert_unsupported_annotation(typing.Union[int, str, bool])
-
-    class X:
-      pass
-
-    # X must be a symbolic type to be parsable.
-    self.assert_unsupported_annotation(X)
+    self.assert_unsupported_annotation(typing.Any)
 
   def test_schema_dict(self):
     schema = schema_lib.Schema([{'x': Itinerary}])
@@ -150,6 +146,25 @@ class SchemaTest(unittest.TestCase):
     schema = schema_lib.Schema([B])
     self.assertEqual(schema.class_dependencies(), [Foo, A, Bar, X, B])
 
+  def test_class_dependencies_non_pyglove(self):
+    class Baz:
+      def __init__(self, x: int):
+        pass
+
+    @dataclasses.dataclass(frozen=True)
+    class AA:
+      foo: tuple[Baz, int]
+
+    class XX(pg.Object):
+      pass
+
+    @dataclasses.dataclass(frozen=True)
+    class BB(AA):
+      foo2: Baz | XX
+
+    schema = schema_lib.Schema([AA])
+    self.assertEqual(schema.class_dependencies(), [Baz, AA, XX, BB])
+
   def test_schema_repr(self):
     schema = schema_lib.Schema([{'x': Itinerary}])
     self.assertEqual(
@@ -177,9 +192,9 @@ class SchemaTest(unittest.TestCase):
     self.assertEqual(schema.parse('{"result": 1}'), 1)
     schema = schema_lib.Schema(dict[str, int])
     self.assertEqual(
-        schema.parse(
-            '{"result": {"_type": "Unknown", "x": 1}}}', force_dict=True),
-        dict(x=1))
+        schema.parse('{"result": {"x": 1}}}'),
+        dict(x=1)
+    )
     with self.assertRaisesRegex(
         schema_lib.SchemaError, 'Expect .* but encountered .*'):
       schema.parse('{"result": "def"}')
@@ -440,28 +455,22 @@ class SchemaPythonReprTest(unittest.TestCase):
         'class A(Object):\n  pass\n',
     )
 
-    class B:
-      pass
-
-    with self.assertRaisesRegex(
-        TypeError, 'Classes must be `pg.Object` subclasses.*'):
-      schema_lib.class_definition(B)
-
     class C(pg.Object):
       x: str
       __kwargs__: typing.Any
 
-    with self.assertRaisesRegex(
-        TypeError, 'Variable-length keyword arguments is not supported'):
-      schema_lib.class_definition(C)
+    self.assertEqual(schema_lib.class_definition(C), 'class C:\n  x: str\n')
 
   def test_repr(self):
     class Foo(pg.Object):
       x: int
 
-    class Bar(pg.Object):
+    @dataclasses.dataclass(frozen=True)
+    class Bar:
+      """Class Bar."""
       y: str
 
+    @dataclasses.dataclass(frozen=True)
     class Baz(Bar):  # pylint: disable=unused-variable
       pass
 
@@ -475,7 +484,7 @@ class SchemaPythonReprTest(unittest.TestCase):
     schema = schema_lib.Schema([B])
     self.assertEqual(
         schema_lib.SchemaPythonRepr().class_definitions(schema),
-        inspect.cleandoc("""
+        inspect.cleandoc('''
             class Foo:
               x: int
 
@@ -483,16 +492,18 @@ class SchemaPythonReprTest(unittest.TestCase):
               foo: Foo
 
             class Bar:
+              """Class Bar."""
               y: str
 
             class Baz(Bar):
+              """Baz(y: str)"""
               y: str
 
             class B(A):
               foo: Foo
               bar: Bar
               foo2: Foo
-            """) + '\n',
+            ''') + '\n',
     )
 
     self.assertEqual(
@@ -501,7 +512,7 @@ class SchemaPythonReprTest(unittest.TestCase):
 
     self.assertEqual(
         schema_lib.SchemaPythonRepr().repr(schema),
-        inspect.cleandoc("""
+        inspect.cleandoc('''
             list[B]
 
             ```python
@@ -512,9 +523,11 @@ class SchemaPythonReprTest(unittest.TestCase):
               foo: Foo
 
             class Bar:
+              """Class Bar."""
               y: str
 
             class Baz(Bar):
+              """Baz(y: str)"""
               y: str
 
             class B(A):
@@ -522,7 +535,7 @@ class SchemaPythonReprTest(unittest.TestCase):
               bar: Bar
               foo2: Foo
             ```
-            """),
+            '''),
     )
     self.assertEqual(
         schema_lib.SchemaPythonRepr().repr(
@@ -531,24 +544,26 @@ class SchemaPythonReprTest(unittest.TestCase):
             include_pg_object_as_base=True,
             markdown=False,
         ),
-        inspect.cleandoc("""
+        inspect.cleandoc('''
             class Foo(Object):
               x: int
 
             class A(Object):
               foo: Foo
 
-            class Bar(Object):
+            class Bar:
+              """Class Bar."""
               y: str
 
             class Baz(Bar):
+              """Baz(y: str)"""
               y: str
 
             class B(A):
               foo: Foo
               bar: Bar
               foo2: Foo
-            """),
+            '''),
     )