langfun 0.0.2.dev20240414__py3-none-any.whl → 0.0.2.dev20240418__py3-none-any.whl

langfun/core/structured/__init__.py

@@ -48,6 +48,8 @@ from langfun.core.structured.schema_generation import generate_class
 from langfun.core.structured.schema_generation import classgen_example
 from langfun.core.structured.schema_generation import default_classgen_examples
 
+from langfun.core.structured.function_generation import function_gen
+
 from langfun.core.structured.mapping import Mapping
 from langfun.core.structured.mapping import MappingExample
 

langfun/core/structured/completion_test.py

@@ -583,6 +583,7 @@ class CompleteStructureTest(unittest.TestCase):
             result=Activity(description='foo'),
             score=1.0,
             logprobs=None,
+            usage=lf.LMSamplingUsage(553, 27, 580),
             tags=['lm-response', 'lm-output', 'transformed']
         )
     )

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/parsing_test.py

@@ -280,13 +280,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']
         ),
     )
@@ -634,13 +636,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_test.py

@@ -77,6 +77,7 @@ class QueryTest(unittest.TestCase):
             result=1,
             score=1.0,
             logprobs=None,
+            usage=lf.LMSamplingUsage(323, 1, 324),
             tags=['lm-response', 'lm-output', 'transformed'],
         ),
     )

langfun/core/templates/selfplay_test.py

@@ -56,7 +56,9 @@ class SelfPlayTest(unittest.TestCase):
     g = NumberGuess(target_num=10)
 
     with lf.context(lm=NumberGuesser(guesses=[50, 20, 5, 10])):
-      self.assertEqual(g(), lf.AIMessage('10', score=0.0, logprobs=None))
+      self.assertEqual(
+          g(), lf.AIMessage('10', score=0.0, logprobs=None, usage=None)
+      )
 
     self.assertEqual(g.num_turns, 4)
 
@@ -64,7 +66,9 @@ class SelfPlayTest(unittest.TestCase):
     g = NumberGuess(target_num=10, max_turns=10)
 
     with lf.context(lm=NumberGuesser(guesses=[50, 20, 5, 2, 5, 4])):
-      self.assertEqual(g(), lf.AIMessage('2', score=0.0, logprobs=None))
+      self.assertEqual(
+          g(), lf.AIMessage('2', score=0.0, logprobs=None, usage=None)
+      )
 
     self.assertEqual(g.num_turns, 10)
 

{langfun-0.0.2.dev20240414.dist-info → langfun-0.0.2.dev20240418.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.0.2.dev20240414
+Version: 0.0.2.dev20240418
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors