Merge pull request #1714 from codeflash-ai/testgen-review

feat: per-function test quality review and repair

Author: KRRT7

.github/workflows/codeflash-optimize.yaml

@@ -39,4 +39,4 @@ jobs:
       - name: ⚡️Codeflash Optimization
         id: optimize_code
         run: |
-          uv run codeflash --benchmark
+          uv run codeflash --benchmark --testgen-review

codeflash/api/aiservice.py

@@ -15,14 +15,15 @@
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
 from codeflash.code_utils.time_utils import humanize_runtime
 from codeflash.languages import Language, current_language
-from codeflash.languages.current import current_language_support
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import (
     AIServiceRefinerRequest,
     CodeStringsMarkdown,
+    FunctionRepairInfo,
     OptimizationReviewResult,
     OptimizedCandidate,
     OptimizedCandidateSource,
+    TestFileReview,
 )
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.version import __version__ as codeflash_version
@@ -57,13 +58,25 @@ def add_language_metadata(
         payload: dict[str, Any], language_version: str | None = None, module_system: str | None = None
     ) -> None:
         """Add language version and module system metadata to an API payload."""
+        from codeflash.languages.current import current_language_support
+
         payload["python_version"] = platform.python_version()
         default_lang_version = current_language_support().default_language_version
         if default_lang_version is not None:
             payload["language_version"] = language_version or default_lang_version
             if module_system:
                 payload["module_system"] = module_system
 
+    @staticmethod
+    def log_error_response(response: requests.Response, action: str, ph_event: str) -> None:
+        """Log and report an API error response."""
+        try:
+            error = response.json()["error"]
+        except Exception:
+            error = response.text
+        logger.error(f"Error {action}: {response.status_code} - {error}")
+        ph(ph_event, {"response_status_code": response.status_code, "error": error})
+
     def get_aiservice_base_url(self) -> str:
         if os.environ.get("CODEFLASH_AIS_SERVER", default="prod").lower() == "local":
             logger.info("Using local AI Service at http://localhost:8000")
@@ -95,14 +108,6 @@ def make_ai_service_request(
         ------
             requests.exceptions.RequestException: If the request fails
 
-        """
-        """Make an API request to the given endpoint on the AI service.
-
-        :param endpoint: The endpoint to call, e.g., "/optimize".
-        :param method: The HTTP method to use ('GET' or 'POST').
-        :param payload: Optional JSON payload to include in the POST request body.
-        :param timeout: The timeout for the request.
-        :return: The response object from the API.
         """
         url = f"{self.base_url}/ai{endpoint}"
         if method.upper() == "POST":
@@ -213,12 +218,7 @@ def optimize_code(
             logger.info(f"!lsp|Received {len(optimizations_json)} optimization candidates.")
             console.rule()
             return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.OPTIMIZE, language)
-        try:
-            error = response.json()["error"]
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
-        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating optimized candidates", "cli-optimize-error-response")
         console.rule()
         return []
 
@@ -285,12 +285,7 @@ def get_jit_rewritten_code(self, source_code: str, trace_id: str) -> list[Optimi
             end_time = time.perf_counter()
             logger.debug(f"!lsp|Generating jit rewritten code took {end_time - start_time:.2f} seconds.")
             return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.JIT_REWRITE)
-        try:
-            error = response.json()["error"]
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating jit rewritten candidate: {response.status_code} - {error}")
-        ph("cli-jit-rewrite-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating jit rewritten candidate", "cli-jit-rewrite-error-response")
         console.rule()
         return []
 
@@ -362,12 +357,7 @@ def optimize_python_code_line_profiler(
             logger.info(f"!lsp|Received {len(optimizations_json)} line profiler optimization candidates.")
             console.rule()
             return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.OPTIMIZE_LP)
-        try:
-            error = response.json()["error"]
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
-        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating optimized candidates", "cli-optimize-error-response")
         console.rule()
         return []
 
@@ -395,12 +385,7 @@ def adaptive_optimize(self, request: AIServiceAdaptiveOptimizeRequest) -> Optimi
 
             return valid_candidates[0]
 
-        try:
-            error = response.json()["error"]
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
-        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating optimized candidates", "cli-optimize-error-response")
         return None
 
     def optimize_code_refinement(self, request: list[AIServiceRefinerRequest]) -> list[OptimizedCandidate]:
@@ -456,12 +441,7 @@ def optimize_code_refinement(self, request: list[AIServiceRefinerRequest]) -> li
 
             return self._get_valid_candidates(refined_optimizations, OptimizedCandidateSource.REFINE)
 
-        try:
-            error = response.json()["error"]
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
-        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating optimized candidates", "cli-optimize-error-response")
         console.rule()
         return []
 
@@ -508,12 +488,7 @@ def code_repair(self, request: AIServiceCodeRepairRequest) -> OptimizedCandidate
 
             return valid_candidates[0]
 
-        try:
-            error = response.json()["error"]
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
-        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating optimized candidates", "cli-optimize-error-response")
         console.rule()
         return None
 
@@ -608,12 +583,7 @@ def get_new_explanation(
             explanation: str = response.json()["explanation"]
             console.rule()
             return explanation
-        try:
-            error = response.json()["error"]
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
-        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating optimized candidates", "cli-optimize-error-response")
         console.rule()
         return ""
 
@@ -660,12 +630,7 @@ def generate_ranking(
             ranking: list[int] = response.json()["ranking"]
             console.rule()
             return ranking
-        try:
-            error = response.json()["error"]
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating ranking: {response.status_code} - {error}")
-        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating ranking", "cli-optimize-error-response")
         console.rule()
         return None
 
@@ -726,7 +691,7 @@ def generate_regression_tests(
         language_version: str | None = None,
         module_system: str | None = None,
         is_numerical_code: bool | None = None,
-    ) -> tuple[str, str, str] | None:
+    ) -> tuple[str, str, str, str | None] | None:
         """Generate regression tests for the given function by making a request to the Django endpoint.
 
         Parameters
@@ -749,6 +714,8 @@ def generate_regression_tests(
 
         """
         # Validate test framework based on language
+        from codeflash.languages.current import current_language_support
+
         lang_support = current_language_support()
         valid_frameworks = lang_support.valid_test_frameworks
         assert test_framework in valid_frameworks, (
@@ -779,6 +746,8 @@ def generate_regression_tests(
         try:
             response = self.make_ai_service_request("/testgen", payload=payload, timeout=self.timeout)
         except requests.exceptions.RequestException as e:
+            from codeflash.telemetry.posthog_cf import ph
+
             logger.exception(f"Error generating tests: {e}")
             ph("cli-testgen-error-caught", {"error": str(e)})
             return None
@@ -792,17 +761,111 @@ def generate_regression_tests(
                 response_json["generated_tests"],
                 response_json["instrumented_behavior_tests"],
                 response_json["instrumented_perf_tests"],
+                response_json.get("raw_generated_tests"),
             )
+        self.log_error_response(response, "generating tests", "cli-testgen-error-response")
+        return None
+
+    def review_generated_tests(
+        self,
+        tests: list[dict[str, Any]],
+        function_source_code: str,
+        function_name: str,
+        trace_id: str,
+        coverage_summary: str = "",
+        coverage_details: dict[str, Any] | None = None,
+        language: str = "python",
+    ) -> list[TestFileReview]:
+        payload: dict[str, Any] = {
+            "tests": tests,
+            "function_source_code": function_source_code,
+            "function_name": function_name,
+            "trace_id": trace_id,
+            "language": language,
+            "codeflash_version": codeflash_version,
+            "call_sequence": self.get_next_sequence(),
+        }
+        if coverage_summary:
+            payload["coverage_summary"] = coverage_summary
+        if coverage_details:
+            payload["coverage_details"] = coverage_details
+        self.add_language_metadata(payload)
         try:
-            error = response.json()["error"]
-            logger.error(f"Error generating tests: {response.status_code} - {error}")
-            ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": error})
-            return None
-        except Exception:
-            logger.error(f"Error generating tests: {response.status_code} - {response.text}")
-            ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": response.text})
+            response = self.make_ai_service_request("/testgen_review", payload=payload, timeout=self.timeout)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error reviewing generated tests: {e}")
+            ph("cli-testgen-review-error-caught", {"error": str(e)})
+            return []
+
+        if response.status_code == 200:
+            data = response.json()
+            return [
+                TestFileReview(
+                    test_index=r["test_index"],
+                    functions_to_repair=[
+                        FunctionRepairInfo(function_name=f["function_name"], reason=f.get("reason", ""))
+                        for f in r.get("functions", [])
+                    ],
+                )
+                for r in data.get("reviews", [])
+            ]
+        self.log_error_response(response, "reviewing generated tests", "cli-testgen-review-error-response")
+        return []
+
+    def repair_generated_tests(
+        self,
+        test_source: str,
+        functions_to_repair: list[FunctionRepairInfo],
+        function_source_code: str,
+        function_to_optimize: FunctionToOptimize,
+        helper_function_names: list[str],
+        module_path: Path,
+        test_module_path: Path,
+        test_framework: str,
+        test_timeout: int,
+        trace_id: str,
+        language: str = "python",
+        coverage_details: dict[str, Any] | None = None,
+        previous_repair_errors: dict[str, str] | None = None,
+        module_source_code: str = "",
+    ) -> tuple[str, str, str] | None:
+        payload: dict[str, Any] = {
+            "test_source": test_source,
+            "functions_to_repair": [
+                {"function_name": f.function_name, "reason": f.reason} for f in functions_to_repair
+            ],
+            "function_source_code": function_source_code,
+            "function_to_optimize": function_to_optimize,
+            "helper_function_names": helper_function_names,
+            "module_path": module_path,
+            "test_module_path": test_module_path,
+            "test_framework": test_framework,
+            "test_timeout": test_timeout,
+            "trace_id": trace_id,
+            "language": language,
+            "codeflash_version": codeflash_version,
+            "call_sequence": self.get_next_sequence(),
+        }
+        if module_source_code:
+            payload["module_source_code"] = module_source_code
+        if coverage_details:
+            payload["coverage_details"] = coverage_details
+        if previous_repair_errors:
+            payload["previous_repair_errors"] = previous_repair_errors
+        self.add_language_metadata(payload)
+        try:
+            response = self.make_ai_service_request("/testgen_repair", payload=payload, timeout=self.timeout)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error repairing generated tests: {e}")
+            ph("cli-testgen-repair-error-caught", {"error": str(e)})
             return None
 
+        if response.status_code == 200:
+            data = response.json()
+            return (data["generated_tests"], data["instrumented_behavior_tests"], data["instrumented_perf_tests"])
+        self.log_error_response(response, "repairing generated tests", "cli-testgen-repair-error-response")
+        return None
+
     def get_optimization_review(
         self,
         original_code: dict[Path, str],
@@ -874,12 +937,7 @@ def get_optimization_review(
             return OptimizationReviewResult(
                 review=cast("str", data["review"]), explanation=cast("str", data.get("review_explanation", ""))
             )
-        try:
-            error = cast("str", response.json()["error"])
-        except Exception:
-            error = response.text
-        logger.error(f"Error generating optimization review: {response.status_code} - {error}")
-        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        self.log_error_response(response, "generating optimization review", "cli-optimize-error-response")
         console.rule()
         return OptimizationReviewResult(review="", explanation="")
 

codeflash/cli_cmds/cli.py

@@ -106,6 +106,12 @@ def parse_args() -> Namespace:
     )
     parser.add_argument("--no-draft", default=False, action="store_true", help="Skip optimization for draft PRs")
     parser.add_argument("--worktree", default=False, action="store_true", help="Use worktree for optimization")
+    parser.add_argument(
+        "--testgen-review", default=False, action="store_true", help="Enable AI review and repair of generated tests"
+    )
+    parser.add_argument(
+        "--testgen-review-turns", type=int, default=None, help="Number of review/repair cycles (default: 2)"
+    )
     parser.add_argument(
         "--async",
         default=False,
@@ -201,12 +207,6 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
         if env_utils.is_ci():
             args.no_pr = True
 
-    if getattr(args, "async", False):
-        logger.warning(
-            "The --async flag is deprecated and will be removed in a future version. "
-            "Async function optimization is now enabled by default."
-        )
-
     return args
 
 

codeflash/code_utils/config_consts.py

@@ -21,6 +21,7 @@
 COVERAGE_THRESHOLD = 60.0
 MIN_TESTCASE_PASSED_THRESHOLD = 6
 REPEAT_OPTIMIZATION_PROBABILITY = 0.1
+MAX_TEST_REPAIR_CYCLES = 2
 DEFAULT_IMPORTANCE_THRESHOLD = 0.001
 
 # pytest loop stability

codeflash/languages/python/context/code_context_extractor.py

@@ -439,6 +439,7 @@ def get_function_sources_from_jedi(
                                 fully_qualified_name=fqn,
                                 only_function_name=func_name,
                                 source_code=definition.get_line_code(),
+                                definition_type=definition.type,
                             )
                             file_path_to_function_source[definition_path].add(function_source)
                             function_source_list.append(function_source)