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

langfun/core/structured/function_generation.py

@@ -76,6 +76,7 @@ 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 = 1,
@@ -141,21 +142,23 @@ def _function_gen(
   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):
-        pg.logging.warning(
-            "Signature mismatch. Expected: %s, Actual: %s",
-            inspect.signature(func),
-            inspect.signature(f),
+        raise python.CodeError(
+            code=source_code,
+            cause=TypeError(
+                f"Signature mismatch: Expected: {inspect.signature(func)}, "
+                f"Actual: {inspect.signature(f)}.",
+            ),
         )
-        continue
 
       if callable(unittest):
         unittest(f)
@@ -163,10 +166,12 @@ def _function_gen(
         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):
@@ -220,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:
@@ -238,20 +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,
+          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

@@ -311,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

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

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

{langfun-0.1.2.dev202412070804.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=gOV5B4KXzN6ng1P1QtZ8aOAEQB8eAbgwWGj57tnzWJY,8159
-langfun/core/structured/function_generation_test.py,sha256=1OtstouOYyYOd_gmZtL8RRbh-FcYGEvBNju6lNrJrOA,11331
+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.dev202412070804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202412070804.dist-info/METADATA,sha256=Ao9kzm8AFw7749nP7p_m3k41rVcBBzG8_QuMnLLN9_U,8281
-langfun-0.1.2.dev202412070804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-langfun-0.1.2.dev202412070804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202412070804.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,,