langfun 0.0.2.dev20240330__py3-none-any.whl → 0.0.2.dev20240511__py3-none-any.whl

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,9 +13,10 @@
 # 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.llms import fake
 from langfun.core.structured import mapping
 from langfun.core.structured import schema as schema_lib
 import pyglove as pg
@@ -78,7 +79,9 @@ class QueryStructurePython(QueryStructure):
 
       {{ output_title }}:
         ```python
-        Answer(final_answer=2)
+        Answer(
+          final_answer=2
+        )
         ```
       """
   protocol = 'python'
@@ -107,9 +110,11 @@ 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',
+    include_methods: bool = False,
     returns_message: bool = False,
     skip_lm: bool = False,
     **kwargs,
@@ -157,8 +162,11 @@ def query(
     examples: An optional list of fewshot examples for helping parsing. If None,
       the default one-shot example will be added.
     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.
+      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
@@ -166,12 +174,17 @@ def query(
       will use `lm`.
     protocol: The protocol for schema/value representation. Applicable values
       are 'json' and 'python'. By default `python` will be used.
+    include_methods: If True, include method definitions in the output type
+      during prompting.
     returns_message: If True, returns `lf.Message` as the output, instead of
       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,16 +201,22 @@ 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.
-  if isinstance(prompt, str):
-    prompt = lf.Template(prompt, **kwargs)
-  elif isinstance(prompt, lf.Template):
-    prompt = prompt.rebind(**kwargs)
+  prompt_kwargs = kwargs.copy()
 
-  if isinstance(prompt, lf.Template):
-    prompt = prompt.render(lm=lm)
+  # 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, lf.Message, lf.Template)):
+    prompt = lf.Template.from_value(prompt, **prompt_kwargs).render(lm=lm)
   else:
     prompt = schema_lib.mark_missing(prompt)
 
@@ -206,6 +225,8 @@ def query(
       schema=schema,
       default=default,
       examples=examples,
+      include_methods=include_methods,
+      response_postprocess=response_postprocess,
       autofix=autofix if protocol == 'python' else 0,
       **kwargs,
   )(
@@ -215,3 +236,31 @@ def query(
       skip_lm=skip_lm,
   )
   return output if returns_message else output.result
+
+
+def query_prompt(
+    prompt: Union[str, pg.Symbolic],
+    schema: Union[
+        schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
+    ] = None,
+    **kwargs,
+) -> lf.Message:
+  """Returns the final prompt sent to LLM for `lf.query`."""
+  kwargs.pop('returns_message', None)
+  kwargs.pop('skip_lm', None)
+  return query(prompt, schema, skip_lm=True, returns_message=True, **kwargs)
+
+
+def query_output(
+    response: Union[str, lf.Message],
+    schema: Union[
+        schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
+    ],
+    **kwargs,
+) -> Any:
+  """Returns the final output of `lf.query` from a provided LLM response."""
+  kwargs.pop('prompt', None)
+  kwargs.pop('lm', None)
+  return query(
+      'Unused prompt', schema, lm=fake.StaticResponse(response), **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:!!'
         ),
     )
 
@@ -239,6 +285,49 @@ class QueryTest(unittest.TestCase):
     with self.assertRaisesRegex(ValueError, 'Unknown protocol'):
       prompting.query('what is 1 + 1', int, protocol='text')
 
+  def test_query_prompt(self):
+    self.assertEqual(
+        prompting.query_prompt('what is this?', int),
+        inspect.cleandoc("""
+            Please respond to the last INPUT_OBJECT with OUTPUT_OBJECT according to OUTPUT_TYPE.
+
+            INPUT_OBJECT:
+              1 + 1 =
+
+            OUTPUT_TYPE:
+              Answer
+
+              ```python
+              class Answer:
+                final_answer: int
+              ```
+
+            OUTPUT_OBJECT:
+              ```python
+              Answer(
+                final_answer=2
+              )
+              ```
+
+            INPUT_OBJECT:
+              what is this?
+
+            OUTPUT_TYPE:
+              int
+
+            OUTPUT_OBJECT:
+            """),
+    )
+
+  def test_query_output(self):
+    self.assertEqual(
+        prompting.query_output(
+            lf.AIMessage('1'),
+            int,
+        ),
+        1,
+    )
+
 
 class QueryStructurePythonTest(unittest.TestCase):
 
@@ -264,7 +353,9 @@ class QueryStructurePythonTest(unittest.TestCase):
 
             OUTPUT_OBJECT:
               ```python
-              Answer(final_answer=2)
+              Answer(
+                final_answer=2
+              )
               ```
 
             INPUT_OBJECT:
@@ -308,7 +399,9 @@ class QueryStructurePythonTest(unittest.TestCase):
 
             OUTPUT_OBJECT:
               ```python
-              Answer(final_answer=2)
+              Answer(
+                final_answer=2
+              )
               ```
 
             INPUT_OBJECT:
@@ -420,7 +513,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 +529,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 +751,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

@@ -301,6 +301,7 @@ class SchemaPythonRepr(SchemaRepr):
       schema: Schema,
       *,
       include_result_definition: bool = True,
+      include_methods: bool = False,
       markdown: bool = True,
       **kwargs,
   ) -> str:
@@ -308,7 +309,7 @@ class SchemaPythonRepr(SchemaRepr):
     if include_result_definition:
       ret += self.result_definition(schema)
     class_definition_str = self.class_definitions(
-        schema, markdown=markdown, **kwargs
+        schema, markdown=markdown, include_methods=include_methods, **kwargs
     )
     if class_definition_str:
       ret += f'\n\n{class_definition_str}'
@@ -331,6 +332,7 @@ def class_definitions(
     classes: Sequence[Type[Any]],
     *,
     include_pg_object_as_base: bool = False,
+    include_methods: bool = False,
     strict: bool = False,
     markdown: bool = False,
 ) -> str | None:
@@ -346,6 +348,7 @@ def class_definitions(
             cls,
             strict=strict,
             include_pg_object_as_base=include_pg_object_as_base,
+            include_methods=include_methods,
         )
     )
   ret = def_str.getvalue()
@@ -355,7 +358,10 @@ def class_definitions(
 
 
 def class_definition(
-    cls, strict: bool = False, include_pg_object_as_base: bool = False
+    cls,
+    strict: bool = False,
+    include_pg_object_as_base: bool = False,
+    include_methods: bool = False,
 ) -> str:
   """Returns the Python class definition."""
   out = io.StringIO()
@@ -383,13 +389,16 @@ def class_definition(
         out.write('\n')
       out.write('  """\n')
 
+  empty_class = True
   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'):
@@ -399,11 +408,33 @@ def class_definition(
             out.write('\n')
       out.write(f'  {field.key}: {annotation(field.value, strict=strict)}')
       out.write('\n')
-  else:
+      empty_class = False
+
+  if include_methods:
+    for method in _iter_newly_defined_methods(cls):
+      out.write('\n')
+      out.write(
+          textwrap.indent(
+              inspect.cleandoc('\n' + inspect.getsource(method)), ' ' * 2)
+      )
+      out.write('\n')
+      empty_class = False
+
+  if empty_class:
     out.write('  pass\n')
   return out.getvalue()
 
 
+def _iter_newly_defined_methods(cls):
+  names = set(dir(cls))
+  for base in cls.__bases__:
+    names -= set(dir(base))
+  for name in names:
+    attr = getattr(cls, name)
+    if callable(attr):
+      yield attr
+
+
 def annotation(
     vs: pg.typing.ValueSpec,
     annotate_optional: bool = True,
@@ -491,7 +522,8 @@ def annotation(
 class SchemaJsonRepr(SchemaRepr):
   """JSON-representation for a schema."""
 
-  def repr(self, schema: Schema) -> str:
+  def repr(self, schema: Schema, **kwargs) -> str:
+    del kwargs
     out = io.StringIO()
     def _visit(node: Any) -> None:
       if isinstance(node, str):

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

@@ -192,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"}')
@@ -459,9 +459,24 @@ class SchemaPythonReprTest(unittest.TestCase):
       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')
+
+    class D(pg.Object):
+      x: str
+      def __call__(self, y: int) -> int:
+        return len(self.x) + y
+
+    self.assertEqual(
+        schema_lib.class_definition(D, include_methods=True),
+        inspect.cleandoc(
+            """
+            class D:
+              x: str
+
+              def __call__(self, y: int) -> int:
+                return len(self.x) + y
+            """) + '\n'
+    )
 
   def test_repr(self):
     class Foo(pg.Object):
@@ -479,13 +494,21 @@ class SchemaPythonReprTest(unittest.TestCase):
     class A(pg.Object):
       foo: Foo
 
+      def foo_value(self) -> int:
+        return self.foo.x
+
     class B(A):
       bar: Bar
       foo2: Foo
 
+      def bar_value(self) -> str:
+        return self.bar.y
+
     schema = schema_lib.Schema([B])
     self.assertEqual(
-        schema_lib.SchemaPythonRepr().class_definitions(schema),
+        schema_lib.SchemaPythonRepr().class_definitions(
+            schema, include_methods=True
+        ),
         inspect.cleandoc('''
             class Foo:
               x: int
@@ -493,6 +516,9 @@ class SchemaPythonReprTest(unittest.TestCase):
             class A:
               foo: Foo
 
+              def foo_value(self) -> int:
+                return self.foo.x
+
             class Bar:
               """Class Bar."""
               y: str
@@ -505,6 +531,9 @@ class SchemaPythonReprTest(unittest.TestCase):
               foo: Foo
               bar: Bar
               foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
             ''') + '\n',
     )
 

langfun/core/structured/scoring.py

@@ -32,8 +32,9 @@ def score(
     lm: lf.LanguageModel | None = None,
     examples: list[mapping.MappingExample] | None = None,
     protocol: schema_lib.SchemaProtocol = 'python',
+    return_scoring_results: bool = False,
     **kwargs,
-) -> list[float]:
+) -> list[float] | list[lf.LMScoringResult]:
   """Scores the outputs based on the prompt."""
   if not completions:
     raise ValueError('`completions` must not be empty.')
@@ -72,4 +73,6 @@ def score(
           for c in completions
       ],
   )
+  if return_scoring_results:
+    return results
   return [r.score for r in results]

langfun/core/structured/scoring_test.py

@@ -35,6 +35,12 @@ class ScoringTest(unittest.TestCase):
   def test_score(self):
     self.assertEqual(scoring.score('hi', [1, 2], lm=fake.Echo()), [0.0, -1.0])
 
+  def test_score_returning_scoring_results(self):
+    self.assertEqual(scoring.score(
+        'hi', [1, 2], lm=fake.Echo(), return_scoring_results=True),
+                     [lf.LMScoringResult(score=0.0, gradients=None),
+                      lf.LMScoringResult(score=-1.0, gradients=None)])
+
   def test_scope_with_lm_from_the_context(self):
     with lf.context(lm=fake.Echo()):
       self.assertEqual(scoring.score('hi', [1, 2]), [0.0, -1.0])