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

langfun/core/structured/function_generation.py

@@ -0,0 +1,245 @@
+# Copyright 2023 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""LLM-based function generation."""
+
+import functools
+import inspect
+import re
+from typing import Any, Callable, Optional, Tuple
+
+from langfun.core import language_model
+from langfun.core import template
+from langfun.core.coding import python
+from langfun.core.structured import prompting
+import pyglove as pg
+
+
+def unittest_gen(signature, lm, num_retries=10):
+  """Generates unit tests for a python function signature."""
+
+  class UnitTest(pg.Object):
+    """A valid unit test for a python function."""
+
+    input: dict[str, Any]
+    expected_output: Any
+
+  class PythonFunctionSignature(pg.Object):
+    signature: str
+
+  unittest_examples = None
+  for _ in range(num_retries):
+    r = prompting.query(
+        PythonFunctionSignature(signature=signature),
+        list[UnitTest],
+        lm=lm,
+        default=None,
+    )
+    if isinstance(r, list) and r:
+      unittest_examples = []
+      for unit_test in r:
+        unittest_examples.append((unit_test.input, unit_test.expected_output))
+      break
+
+  return unittest_examples
+
+
+def unittest_with_test_cases(f, unittests):
+  """Applies unit tests to a python function to be tested."""
+  if not unittests:
+    raise ValueError(f"No unit tests provided: {unittests}")
+
+  for unit_test in unittests:
+    inputs = unit_test[0]
+    if isinstance(inputs, dict):
+      actual = f(**inputs)
+    elif isinstance(inputs, tuple):
+      actual = f(*inputs)
+    else:
+      actual = f(inputs)
+
+    expected = unit_test[1]
+    assert (
+        actual == expected
+    ), f"Test FAILED: Inputs: {inputs}, Expected: {expected}, Actual: {actual}"
+
+
+def _function_gen(
+    func: Callable[..., Any],
+    signature: str,
+    lm: language_model.LanguageModel,
+    num_retries: int = 10,
+    unittest: Optional[
+        Callable[[Callable[..., Any]], None] | list[Tuple[Any, Any]]
+    ] = None,
+):
+  """Generates a python function with LLM and verify its quality with unit testing."""
+
+  class PythonFunctionPrompt(template.Template):
+    r"""A template for a python function generation.
+
+    Please reply to the last PYTHON_FUNCTION_SIGNATURE with a self-sufficient,
+    error-free, and efficiently coded PYTHON_FUNCTION, crafted to the standards
+    of a world-class programmer.
+
+      PYTHON_FUNCTION_SIGNATURE:
+        ```python
+        def calculate_area_circle(radius: float) -> float:
+        \"\"\"Calculates the area of a circle given its radius.
+
+        Args:
+            radius: The radius of the circle.
+
+        Returns:
+            The area of the circle.
+        \"\"\"
+        ```
+
+      PYTHON_FUNCTION:
+        ```python
+        def calculate_area_circle(radius: float) -> float:
+        \"\"\"Calculates the area of a circle given its radius.
+
+        Args:
+            radius: The radius of the circle.
+
+        Returns:
+            The area of the circle.
+        \"\"\"
+        import math
+
+        area = math.pi * radius**2
+        return area
+        ```
+
+      PYTHON_FUNCTION_SIGNATURE:
+        ```python
+        {{signature}}
+        ```
+
+      PYTHON_FUNCTION:
+    """
+
+  unittest_examples = None
+  if unittest is None:
+    unittest_examples = unittest_gen(signature, lm=lm)
+  elif not callable(unittest):
+    unittest_examples = unittest
+
+  for _ in range(num_retries):
+    try:
+      source_code = prompting.query(
+          PythonFunctionPrompt(signature=signature), lm=lm
+      )
+      f = python.evaluate(source_code)
+
+      # Check whether the sigantures are the same.
+      if inspect.signature(f) != inspect.signature(func):
+        continue
+
+      if callable(unittest):
+        unittest(f)
+      else:
+        unittest_with_test_cases(f, unittest_examples)
+
+      return f, source_code
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass
+
+  return None, None
+
+
+def _process_signature(signature):
+  # Remove the decorator.
+  pattern = r"^\@.*function_gen.*$"
+  signature = re.sub(pattern, "", signature, flags=re.MULTILINE)
+  # Remove the possible 'pass' in an empty function.
+  pattern = r"^\s*pass\s*$"
+  signature = re.sub(pattern, "", signature, flags=re.MULTILINE)
+  return signature.strip()
+
+
+def function_gen(
+    lm: language_model.LanguageModel,
+    cache_filename: str | None = None,
+    num_retries: int = 10,
+    unittest: Optional[
+        Callable[[Callable[..., Any]], None] | list[Tuple[Any, Any]]
+    ] = None,
+):
+  """A decorator for automating function generation using a language model.
+
+  This decorator should be applied to functions that are not yet implemented. It
+  facilitates the implementation via the specified LLM, ensuring
+  quality through unit tests.
+
+  Args:
+      lm (lf.LanguageModel): The language model used for generating function
+        implementations.
+      cache_filename (str | None): Optional. The path of the file where
+        generated function implementations are loaded from or saved to.
+      num_retries (int): Maximum number of attempts the language model should
+        make to generate a suitable function implementation.
+      unittest: This optional parameter enables the definition of custom unit
+        tests. You can either provide a list of test cases as tuples of inputs
+        and outputs, or a function that throws an error if a test fails. If left
+        as None (the default setting), the LLM will automatically create the
+        unit test cases.
+
+  Returns:
+      The implemented function object.
+  """
+
+  def _decorate(func):
+    setattr(func, "__function__", None)
+    setattr(func, "__source_code__", None)
+
+    @functools.wraps(func)
+    def lm_generated_func(*args, **kwargs):
+      if func.__function__ is not None:
+        return func.__function__(*args, **kwargs)
+
+      signature = _process_signature(inspect.getsource(func))
+      cache = pg.Dict()
+      if cache_filename is not None:
+        try:
+          cache = pg.load(cache_filename)
+        except FileNotFoundError:
+          pg.logging.warning(
+              "Creating a new cache as cache file '%s' does not exist.",
+              cache_filename,
+          )
+
+        if signature in cache:
+          func.__source_code__ = cache[signature]
+          func.__function__ = python.evaluate(func.__source_code__)
+          return func.__function__(*args, **kwargs)
+
+      func.__function__, func.__source_code__ = _function_gen(
+          func, signature, lm, num_retries=num_retries, unittest=unittest
+      )
+      if func.__function__ is None:
+        raise ValueError(f"Function generation failed. Signature:\n{signature}")
+
+      if cache_filename is not None:
+        cache[signature] = func.__source_code__
+        cache.save(cache_filename)
+      return func.__function__(*args, **kwargs)
+
+    lm_generated_func.__name__ = func.__name__
+    lm_generated_func.__qualname__ = func.__qualname__
+    lm_generated_func.__module__ = func.__module__
+    lm_generated_func.source = lambda: func.__source_code__
+    return lm_generated_func
+
+  return _decorate

langfun/core/structured/function_generation_test.py

@@ -0,0 +1,329 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import inspect
+import os
+import tempfile
+import unittest
+from langfun.core.llms import fake
+from langfun.core.structured import function_generation
+import pyglove as pg
+
+
+class FunctionGenerationTest(unittest.TestCase):
+
+  def test_generate_function(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
+        """)
+    unittest_lm_response = inspect.cleandoc("""
+        ```python
+          [
+            UnitTest(
+              input={
+                'items': [1, 2, 3, 4, 5],
+                'target': 3
+              },
+              expected_output=2
+            ),
+            UnitTest(
+              input={
+                'items': [1, 2, 3, 4, 5],
+                'target': 6
+              },
+              expected_output=-1
+            )
+          ]
+        ```
+        """)
+
+    lm = fake.StaticSequence([unittest_lm_response, 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.
+
+      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_custom_unittest_examples(self):
+    function_gen_lm_response = inspect.cleandoc("""
+        ```python
+        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])
+
+    custom_unittest = [(([1, 2, 3, 4, 5], 3), 2)]
+
+    @function_generation.function_gen(lm=lm, unittest=custom_unittest)
+    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)
+
+  def test_custom_unittest_fn(self):
+    function_gen_lm_response = inspect.cleandoc("""
+        ```python
+        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])
+
+    def _unittest_fn(func):
+      assert func([1, 2, 3, 4, 5], 3) == 2
+      assert func([1, 2, 3, 4, 5], 6) == -1
+
+    custom_unittest = _unittest_fn
+
+    @function_generation.function_gen(lm=lm, unittest=custom_unittest)
+    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)
+
+  def test_load_function_from_cache_file(self):
+    lm = fake.StaticSequence([])
+
+    def _unittest_fn(func):
+      assert func([1, 2, 3, 4, 5], 3) == 2
+      assert func([1, 2, 3, 4, 5], 6) == -1
+
+    cache_file_dir = tempfile.gettempdir()
+    cache_file = os.path.join(cache_file_dir, 'cache_file.json')
+
+    cache_key = """@function_generation.function_gen(
+        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.
+
+      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.
+      \"\"\""""
+    cache_value = """
+        ```python
+        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
+        ```
+    """
+    cache = pg.Dict()
+    cache[cache_key] = cache_value
+    cache.save(cache_file)
+
+    @function_generation.function_gen(
+        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.
+
+      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(['a', 'b', 'c'], 'd'), -1)
+
+  def test_empty_cache_file(self):
+    function_gen_lm_response = inspect.cleandoc("""
+        ```python
+        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])
+
+    def _unittest_fn(func):
+      assert func([1, 2, 3, 4, 5], 3) == 2
+      assert func([1, 2, 3, 4, 5], 6) == -1
+
+    cache_file_dir = tempfile.gettempdir()
+    cache_file = os.path.join(cache_file_dir, 'cache_file.json')
+
+    @function_generation.function_gen(
+        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.
+
+      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)
+
+  def test_siganture_check(self):
+    incorrect_signature_lm_response = inspect.cleandoc("""
+        ```python
+        def dummy():
+          pass
+        ```
+        """)
+    function_gen_lm_response = inspect.cleandoc("""
+        ```python
+        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(
+        [incorrect_signature_lm_response, function_gen_lm_response]
+    )
+
+    def _unittest_fn(func):
+      assert func([1, 2, 3, 4, 5], 3) == 2
+      assert func([1, 2, 3, 4, 5], 6) == -1
+
+    custom_unittest = _unittest_fn
+
+    @function_generation.function_gen(lm=lm, unittest=custom_unittest)
+    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)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/structured/mapping.py

@@ -14,12 +14,49 @@
 """The base of symbolic mapping methods."""
 
 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(
+        lf.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('\n')
+      r.write(lf.colored(self.lm_response.text, 'blue'))
+    return r.getvalue()
+
+
 @pg.use_init_args(['input', 'output', 'schema', 'context'])
 class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
   """Mapping example between text, schema and structured value."""
@@ -214,7 +251,7 @@ class Mapping(lf.LangFunc):
 
       {%- if example.schema -%}
       {{ schema_title }}:
-      {{ example.schema_repr(protocol) | indent(2, True) }}
+      {{ example.schema_repr(protocol, include_methods=include_methods) | indent(2, True) }}
 
       {% endif -%}
 
@@ -242,6 +279,10 @@ class Mapping(lf.LangFunc):
       'The protocol for representing the schema and value.',
   ] = 'python'
 
+  include_methods: Annotated[
+      bool, 'If True, include method definitions in the schema.'
+  ] = False
+
   #
   # Other user-provided flags.
   #
@@ -278,6 +319,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 +345,11 @@ 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.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 +366,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

langfun/core/structured/mapping_test.py

@@ -16,10 +16,27 @@
 import inspect
 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(
+        lf.text_formatting.decolored(str(error)),
+        'ValueError: Cannot parse message.\n\n[LM Response]\nhi',
+    )
+    self.assertEqual(
+        lf.text_formatting.decolored(error.format(include_lm_response=False)),
+        'ValueError: Cannot parse message.',
+    )
+
+
 class MappingExampleTest(unittest.TestCase):
 
   def test_basics(self):