fix: include module-level variables and second-degree helpers in testgen helper FQNs

codeflash/languages/python/context/code_context_extractor.py

@@ -107,8 +107,7 @@ def get_code_optimization_context(
     for qualified_names in helpers_of_fto_qualified_names_dict.values():
         qualified_names.update({f"{qn.rsplit('.', 1)[0]}.__init__" for qn in qualified_names if "." in qn})
 
-    # Get FunctionSource representation of helpers of helpers of FTO
-    helpers_of_helpers_dict, _helpers_of_helpers_list = get_function_sources_from_jedi(
+    helpers_of_helpers_dict, helpers_of_helpers_list = get_function_sources_from_jedi(
         helpers_of_fto_qualified_names_dict, project_root_path
     )
 
@@ -186,13 +185,16 @@ def get_code_optimization_context(
     code_hash_context = hashing_code_context.markdown
     code_hash = hashlib.sha256(code_hash_context.encode("utf-8")).hexdigest()
 
+    all_helper_fqns = list({fs.fully_qualified_name for fs in helpers_of_fto_list + helpers_of_helpers_list})
+
     return CodeOptimizationContext(
         testgen_context=testgen_context,
         read_writable_code=final_read_writable_code,
         read_only_context_code=read_only_context_code,
         hashing_code_context=code_hash_context,
         hashing_code_context_hash=code_hash,
         helper_functions=helpers_of_fto_list,
+        testgen_helper_fqns=all_helper_fqns,
         preexisting_objects=preexisting_objects,
     )
 
@@ -317,13 +319,12 @@ def get_code_optimization_context_for_language(
     return CodeOptimizationContext(
         testgen_context=testgen_context,
         read_writable_code=read_writable_code,
-        # Pass type definitions and globals as read-only context for the AI
-        # This way the AI sees them as context but doesn't include them in optimized output
         read_only_context_code=code_context.read_only_context,
         hashing_code_context=read_writable_code.flat,
         hashing_code_context_hash=code_hash,
         helper_functions=helper_function_sources,
-        preexisting_objects=set(),  # Not implemented for non-Python yet
+        testgen_helper_fqns=[fs.fully_qualified_name for fs in helper_function_sources],
+        preexisting_objects=set(),
     )
 
 
@@ -519,15 +520,16 @@ def get_function_sources_from_jedi(
                         and not belongs_to_function_qualified(definition, qualified_function_name)
                         and definition.full_name.startswith(definition.module_name)
                     )
-                    if is_valid_definition and definition.type in ("function", "class"):
+                    if is_valid_definition and definition.type in ("function", "class", "statement"):
                         if definition.type == "function":
                             fqn = definition.full_name
                             func_name = definition.name
-                        else:
-                            # When a class is instantiated (e.g., MyClass()), track its __init__ as a helper
-                            # This ensures the class definition with constructor is included in testgen context
+                        elif definition.type == "class":
                             fqn = f"{definition.full_name}.__init__"
                             func_name = "__init__"
+                        else:
+                            fqn = definition.full_name
+                            func_name = definition.name
                         qualified_name = get_qualified_name(definition.module_name, fqn)
                         # Avoid nested functions or classes. Only class.function is allowed
                         if len(qualified_name.split(".")) <= 2:

codeflash/models/models.py

@@ -380,6 +380,7 @@ class CodeOptimizationContext(BaseModel):
     hashing_code_context: str = ""
     hashing_code_context_hash: str = ""
     helper_functions: list[FunctionSource]
+    testgen_helper_fqns: list[str] = []
     preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]]
 
 

codeflash/optimization/function_optimizer.py

@@ -579,6 +579,7 @@ def generate_and_instrument_tests(
         test_results = self.generate_tests(
             testgen_context=code_context.testgen_context,
             helper_functions=code_context.helper_functions,
+            testgen_helper_fqns=code_context.testgen_helper_fqns,
             generated_test_paths=generated_test_paths,
             generated_perf_test_paths=generated_perf_test_paths,
         )
@@ -1521,7 +1522,8 @@ def get_code_optimization_context(self) -> Result[CodeOptimizationContext, str]:
                 read_only_context_code=new_code_ctx.read_only_context_code,
                 hashing_code_context=new_code_ctx.hashing_code_context,
                 hashing_code_context_hash=new_code_ctx.hashing_code_context_hash,
-                helper_functions=new_code_ctx.helper_functions,  # only functions that are read writable
+                helper_functions=new_code_ctx.helper_functions,
+                testgen_helper_fqns=new_code_ctx.testgen_helper_fqns,
                 preexisting_objects=new_code_ctx.preexisting_objects,
             )
         )
@@ -1727,6 +1729,7 @@ def generate_tests(
         self,
         testgen_context: CodeStringsMarkdown,
         helper_functions: list[FunctionSource],
+        testgen_helper_fqns: list[str],
         generated_test_paths: list[Path],
         generated_perf_test_paths: list[Path],
     ) -> Result[tuple[int, GeneratedTestsList, dict[str, set[FunctionCalledInTest]], str], str]:
@@ -1735,13 +1738,9 @@ def generate_tests(
         assert len(generated_test_paths) == n_tests
 
         if not self.args.no_gen_tests:
-            # Submit test generation tasks
+            helper_fqns = testgen_helper_fqns or [definition.fully_qualified_name for definition in helper_functions]
             future_tests = self.submit_test_generation_tasks(
-                self.executor,
-                testgen_context.markdown,
-                [definition.fully_qualified_name for definition in helper_functions],
-                generated_test_paths,
-                generated_perf_test_paths,
+                self.executor, testgen_context.markdown, helper_fqns, generated_test_paths, generated_perf_test_paths
             )
 
         future_concolic_tests = self.executor.submit(

pyproject.toml

@@ -310,6 +310,9 @@ split-on-trailing-comma = false
 docstring-code-format = true
 skip-magic-trailing-comma = true
 
+[tool.ty.src]
+exclude = ["tests", "code_to_optimize", "pie_test_set", "experiments"]
+
 [tool.hatch.version]
 source = "uv-dynamic-versioning"
 

tests/test_code_context_extractor.py

@@ -978,7 +978,12 @@ def test_repo_helper() -> None:
     code_ctx = get_code_optimization_context(function_to_optimize, project_root)
     read_write_context, read_only_context = code_ctx.read_writable_code, code_ctx.read_only_context_code
     hashing_context = code_ctx.hashing_code_context
+    path_to_globals = project_root / "globals.py"
     expected_read_write_context = f"""
+```python:{path_to_globals.relative_to(project_root)}
+# Define a global variable
+API_URL = "https://api.example.com/data"
+```
 ```python:{path_to_utils.relative_to(project_root)}
 import math
 
@@ -1071,7 +1076,12 @@ def test_repo_helper_of_helper() -> None:
     code_ctx = get_code_optimization_context(function_to_optimize, project_root)
     read_write_context, read_only_context = code_ctx.read_writable_code, code_ctx.read_only_context_code
     hashing_context = code_ctx.hashing_code_context
+    path_to_globals = project_root / "globals.py"
     expected_read_write_context = f"""
+```python:{path_to_globals.relative_to(project_root)}
+# Define a global variable
+API_URL = "https://api.example.com/data"
+```
 ```python:{path_to_utils.relative_to(project_root)}
 import math
 from transform_utils import DataTransformer
@@ -1798,6 +1808,8 @@ def calculate(self, operation, x, y):
 """
         expected_read_only_context = """
 ```python:utility_module.py
+import sys
+
 DEFAULT_PRECISION = "medium"
 
 # Try-except block with variable definitions
@@ -1808,6 +1820,17 @@ def calculate(self, operation, x, y):
     # Used variable in except block
     CALCULATION_BACKEND = "python"
 
+# Nested if-else with variable definitions
+if sys.platform.startswith('win'):
+    # Used variable in outer if
+    SYSTEM_TYPE = "windows"
+elif sys.platform.startswith('linux'):
+    # Used variable in outer elif
+    SYSTEM_TYPE = "linux"
+else:
+    # Used variable in outer else
+    SYSTEM_TYPE = "other"
+
 # Function that will be used in the main code
 def select_precision(precision, fallback_precision):
     if precision is None:
@@ -2014,6 +2037,8 @@ def get_system_details():
         relative_path = file_path.relative_to(project_root)
         expected_read_write_context = f"""
 ```python:utility_module.py
+import sys
+
 DEFAULT_PRECISION = "medium"
 
 # Try-except block with variable definitions
@@ -2024,6 +2049,17 @@ def get_system_details():
     # Used variable in except block
     CALCULATION_BACKEND = "python"
 
+# Nested if-else with variable definitions
+if sys.platform.startswith('win'):
+    # Used variable in outer if
+    SYSTEM_TYPE = "windows"
+elif sys.platform.startswith('linux'):
+    # Used variable in outer elif
+    SYSTEM_TYPE = "linux"
+else:
+    # Used variable in outer else
+    SYSTEM_TYPE = "other"
+
 # Function that will be used in the main code
 def select_precision(precision, fallback_precision):
     if precision is None:
@@ -2064,6 +2100,8 @@ def __init__(self, precision="high", fallback_precision=None, mode="standard"):
 """
         expected_read_only_context = """
 ```python:utility_module.py
+import sys
+
 DEFAULT_PRECISION = "medium"
 
 # Try-except block with variable definitions
@@ -2073,6 +2111,17 @@ def __init__(self, precision="high", fallback_precision=None, mode="standard"):
 except ImportError:
     # Used variable in except block
     CALCULATION_BACKEND = "python"
+
+# Nested if-else with variable definitions
+if sys.platform.startswith('win'):
+    # Used variable in outer if
+    SYSTEM_TYPE = "windows"
+elif sys.platform.startswith('linux'):
+    # Used variable in outer elif
+    SYSTEM_TYPE = "linux"
+else:
+    # Used variable in outer else
+    SYSTEM_TYPE = "other"
 ```
 """
         assert read_write_context.markdown.strip() == expected_read_write_context.strip()