langfun 0.1.2.dev202412060804__py3-none-any.whl → 0.1.2.dev202412080804__py3-none-any.whl

langfun/core/structured/function_generation.py

@@ -16,7 +16,7 @@
 import functools
 import inspect
 import re
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, Literal, Optional, Tuple
 
 from langfun.core import language_model
 from langfun.core import template
@@ -25,7 +25,7 @@ from langfun.core.structured import prompting
 import pyglove as pg
 
 
-def unittest_gen(signature, lm, num_retries=10):
+def unittest_gen(signature, lm, num_retries=1):
   """Generates unit tests for a python function signature."""
 
   class UnitTest(pg.Object):
@@ -76,12 +76,16 @@ def unittest_with_test_cases(f, unittests):
 
 def _function_gen(
     func: Callable[..., Any],
+    context: dict[str, Any],
     signature: str,
     lm: language_model.LanguageModel,
-    num_retries: int = 10,
+    num_retries: int = 1,
     unittest: Optional[
-        Callable[[Callable[..., Any]], None] | list[Tuple[Any, Any]]
+        Callable[[Callable[..., Any]], None]
+        | list[Tuple[Any, Any]]
+        | Literal["auto"]
     ] = None,
+    unittest_num_retries: int = 1,
 ):
   """Generates a python function with LLM and verify its quality with unit testing."""
 
@@ -131,32 +135,43 @@ def _function_gen(
     """
 
   unittest_examples = None
-  if unittest is None:
-    unittest_examples = unittest_gen(signature, lm=lm)
-  elif not callable(unittest):
+  if unittest == "auto":
+    unittest_examples = unittest_gen(
+        signature, lm=lm, num_retries=unittest_num_retries
+    )
+  elif isinstance(unittest, list):
     unittest_examples = unittest
 
+  last_error = None
   for _ in range(num_retries):
     try:
       source_code = prompting.query(
           PythonFunctionPrompt(signature=signature), lm=lm
       )
-      f = python.evaluate(source_code)
+      f = python.evaluate(source_code, global_vars=context)
 
       # Check whether the sigantures are the same.
       if inspect.signature(f) != inspect.signature(func):
-        continue
+        raise python.CodeError(
+            code=source_code,
+            cause=TypeError(
+                f"Signature mismatch: Expected: {inspect.signature(func)}, "
+                f"Actual: {inspect.signature(f)}.",
+            ),
+        )
 
       if callable(unittest):
         unittest(f)
-      else:
+      elif unittest_examples:
         unittest_with_test_cases(f, unittest_examples)
 
       return f, source_code
-    except Exception:  # pylint: disable=broad-exception-caught
-      pass
-
-  return None, None
+    except python.CodeError as e:
+      last_error = e
+      pg.logging.warning(
+          f"Bad code generated: {e}",
+      )
+  raise last_error
 
 
 def _process_signature(signature):
@@ -172,10 +187,13 @@ def _process_signature(signature):
 def function_gen(
     lm: language_model.LanguageModel,
     cache_filename: str | None = None,
-    num_retries: int = 10,
+    num_retries: int = 1,
     unittest: Optional[
-        Callable[[Callable[..., Any]], None] | list[Tuple[Any, Any]]
+        Callable[[Callable[..., Any]], None]
+        | list[Tuple[Any, Any]]
+        | Literal["auto"]
     ] = None,
+    unittest_num_retries: int = 1,
 ):
   """A decorator for automating function generation using a language model.
 
@@ -192,9 +210,12 @@ def function_gen(
         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.
+        and outputs, or a function that throws an error if a test fails, or let
+        LLM automatically create the unit test cases. If a generated function is
+        and returned, it should pass all the unittests.
+      unittest_num_retries: If unittest is set to "auto", this parameter
+        specifies the number of times the LLM's attempts to generate unit test
+        cases.
 
   Returns:
       The implemented function object.
@@ -204,6 +225,13 @@ def function_gen(
     setattr(func, "__function__", None)
     setattr(func, "__source_code__", None)
 
+    # Prepare the globals/locals for the generated code to be evaluated against.
+    callstack = inspect.stack()
+    assert len(callstack) > 1
+    context = dict(callstack[1][0].f_globals)
+    context.update(callstack[1][0].f_locals)
+    context.pop(func.__name__, None)
+
     @functools.wraps(func)
     def lm_generated_func(*args, **kwargs):
       if func.__function__ is not None:
@@ -222,15 +250,20 @@ def function_gen(
 
         if signature in cache:
           func.__source_code__ = cache[signature]
-          func.__function__ = python.evaluate(func.__source_code__)
+          func.__function__ = python.evaluate(
+              func.__source_code__, global_vars=context
+          )
           return func.__function__(*args, **kwargs)
 
       func.__function__, func.__source_code__ = _function_gen(
-          func, signature, lm, num_retries=num_retries, unittest=unittest
+          func,
+          context,
+          signature,
+          lm,
+          num_retries=num_retries,
+          unittest=unittest,
+          unittest_num_retries=unittest_num_retries,
       )
-      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)

langfun/core/structured/function_generation_test.py

@@ -63,6 +63,42 @@ class FunctionGenerationTest(unittest.TestCase):
 
     lm = fake.StaticSequence([unittest_lm_response, function_gen_lm_response])
 
+    @function_generation.function_gen(lm=lm, unittest='auto')
+    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_generate_function_without_unittest(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
+        """)
+
+    lm = fake.StaticSequence([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.
@@ -258,7 +294,9 @@ class FunctionGenerationTest(unittest.TestCase):
     cache_file = os.path.join(cache_file_dir, 'cache_file.json')
 
     @function_generation.function_gen(
-        lm=lm, unittest=_unittest_fn, cache_filename=cache_file
+        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.
@@ -273,6 +311,36 @@ class FunctionGenerationTest(unittest.TestCase):
 
     self.assertEqual(linear_search(['a', 'b', 'c'], 'c'), 2)
 
+  def test_context_passthrough(self):
+
+    class Number(pg.Object):
+      value: int
+
+    function_gen_lm_response = inspect.cleandoc("""
+        ```python
+        def add(a: Number, b: Number) -> Number:
+            \"\"\"Adds two numbers together.\"\"\"
+            return Number(a.value + b.value)
+        ```
+        """)
+
+    lm = fake.StaticSequence(
+        [function_gen_lm_response]
+    )
+
+    def _unittest_fn(func):
+      assert func(Number(1), Number(2)) == Number(3)
+
+    custom_unittest = _unittest_fn
+
+    @function_generation.function_gen(
+        lm=lm, unittest=custom_unittest, num_retries=1
+    )
+    def add(a: Number, b: Number) -> Number:  # pylint: disable=unused-argument
+      """Adds two numbers together."""
+
+    self.assertEqual(add(Number(2), Number(3)), Number(5))
+
   def test_siganture_check(self):
     incorrect_signature_lm_response = inspect.cleandoc("""
         ```python
@@ -310,7 +378,9 @@ class FunctionGenerationTest(unittest.TestCase):
 
     custom_unittest = _unittest_fn
 
-    @function_generation.function_gen(lm=lm, unittest=custom_unittest)
+    @function_generation.function_gen(
+        lm=lm, unittest=custom_unittest, num_retries=2
+    )
     def linear_search(items, target):  # pylint: disable=unused-argument
       """Performs a linear search on a list to find a target value.
 

{langfun-0.1.2.dev202412060804.dist-info → langfun-0.1.2.dev202412080804.dist-info}/METADATA

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

{langfun-0.1.2.dev202412060804.dist-info → langfun-0.1.2.dev202412080804.dist-info}/RECORD

@@ -123,8 +123,8 @@ langfun/core/structured/completion.py,sha256=cS2PjG7sqzDu5x0xoTk8RmNcoeX55iVwH38
 langfun/core/structured/completion_test.py,sha256=lendf6nPsNfAmd5A7k3v_HS2At9F_jjbKBcV7OEt94o,19310
 langfun/core/structured/description.py,sha256=SXW4MJvshFjbR-0gw6rE21o6WXq12UlRXawvDBXMZFA,5211
 langfun/core/structured/description_test.py,sha256=UtZGjSFUaQ6130t1E5tcL7ODu0xIefkapb53TbnqsK8,7362
-langfun/core/structured/function_generation.py,sha256=pFgS3vcRAWiuFBol2x5Eeip3XqoudONsOpeJpWyjT3s,7479
-langfun/core/structured/function_generation_test.py,sha256=ZJI-aaGgWWszn92u7h5IZ9Pl70N2DgAGGJrIxPzsvwg,10065
+langfun/core/structured/function_generation.py,sha256=c2KB3M86GE1H-8vSZlilko0OJRem8uqvX7CgqTzLdVw,8558
+langfun/core/structured/function_generation_test.py,sha256=LaXYDXf9GlqUrR6v_gtmK_H4kxzonmU7SYbn7XXMgjU,12128
 langfun/core/structured/mapping.py,sha256=vLKH79UT-j0qkQdvqlQBO7SkXXuM-yr2Idm8_HH8qwM,13649
 langfun/core/structured/mapping_test.py,sha256=bHm2ZCXBITq_G8Lvw_olFHeUUc4s_lGXZm9v9JhoPB4,9630
 langfun/core/structured/parsing.py,sha256=D58wBWOC6r6DCJNychCDkiHPrsy1XJfBDCDDZtug00k,11765
@@ -148,8 +148,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202412060804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202412060804.dist-info/METADATA,sha256=7jXlkJbHKaJuKLg8E5_BZ0jS99pZozItRnVWolNqAtg,8281
-langfun-0.1.2.dev202412060804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-langfun-0.1.2.dev202412060804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202412060804.dist-info/RECORD,,
+langfun-0.1.2.dev202412080804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202412080804.dist-info/METADATA,sha256=71KlK3WFqJkT60wgcB53JeY2zHOqHqX1ucAyX6r05CI,8281
+langfun-0.1.2.dev202412080804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+langfun-0.1.2.dev202412080804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202412080804.dist-info/RECORD,,