some async changes

Author: KRRT7

codeflash/api/aiservice.py

@@ -101,6 +101,7 @@ def optimize_python_code(  # noqa: D417
         trace_id: str,
         num_candidates: int = 10,
         experiment_metadata: ExperimentMetadata | None = None,
+        is_async: bool = False,
     ) -> list[OptimizedCandidate]:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -131,6 +132,7 @@ def optimize_python_code(  # noqa: D417
             "current_username": get_last_commit_author_if_pr_exists(None),
             "repo_owner": git_repo_owner,
             "repo_name": git_repo_name,
+            "is_async": is_async,
         }
 
         logger.info("Generating optimized candidates…")

codeflash/code_utils/edit_generated_tests.py

@@ -201,7 +201,7 @@ def remove_functions_from_generated_tests(
     for generated_test in generated_tests.generated_tests:
         for test_function in test_functions_to_remove:
             function_pattern = re.compile(
-                rf"(@pytest\.mark\.parametrize\(.*?\)\s*)?def\s+{re.escape(test_function)}\(.*?\):.*?(?=\ndef\s|$)",
+                rf"(@pytest\.mark\.parametrize\(.*?\)\s*)?(async\s+)?def\s+{re.escape(test_function)}\(.*?\):.*?(?=\n(async\s+)?def\s|$)",
                 re.DOTALL,
             )
 

codeflash/optimization/function_optimizer.py

@@ -983,6 +983,7 @@ def generate_tests_and_optimizations(
             self.function_trace_id[:-4] + "EXP0" if run_experiment else self.function_trace_id,
             N_CANDIDATES,
             ExperimentMetadata(id=self.experiment_id, group="control") if run_experiment else None,
+            self.function_to_optimize.is_async,
         )
         future_candidates_exp = None
 
@@ -998,6 +999,7 @@ def generate_tests_and_optimizations(
                 self.function_trace_id[:-4] + "EXP1",
                 N_CANDIDATES,
                 ExperimentMetadata(id=self.experiment_id, group="experiment"),
+                self.function_to_optimize.is_async,
             )
             futures.append(future_candidates_exp)