fixes

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -296,15 +296,15 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
         console.rule()
         return []
 
-    def optimize_python_code_repair(self, request: AIServiceCodeRepairRequest) -> OptimizedCandidate | None:
-        """Optimize the given python code for performance by making a request to the Django endpoint.
+    def code_repair(self, request: AIServiceCodeRepairRequest) -> OptimizedCandidate | None:
+        """Repair the optimization candidate that is not matching the test result of the original code.
 
         Args:
-        request: optimization candidate details for refinement
+        request: candidate details for repair
 
         Returns:
         -------
-        - OptimizationCandidate: new fixed candidate.
+        - OptimizedCandidate: new fixed candidate.
 
         """
         console.rule()

codeflash/optimization/function_optimizer.py

@@ -872,6 +872,7 @@ def determine_best_candidate(
         eval_ctx = CandidateEvaluationContext()
         self.future_all_refinements.clear()
         self.future_all_code_repair.clear()
+        self.repair_counter = 0
 
         ai_service_client = self.aiservice_client if exp_type == "EXP0" else self.local_aiservice_client
         assert ai_service_client is not None, "AI service client must be set for optimization"
@@ -992,7 +993,7 @@ def repair_optimization(
             test_diffs=test_diffs,
             trace_id=trace_id,
         )
-        return executor.submit(ai_service_client.optimize_python_code_repair, request=request)
+        return executor.submit(ai_service_client.code_repair, request=request)
 
     def log_successful_optimization(
         self, explanation: Explanation, generated_tests: GeneratedTestsList, exp_type: str
@@ -1903,8 +1904,9 @@ def repair_if_possible(
         if self.repair_counter >= MAX_REPAIRS_PER_TRACE:
             logger.debug(f"Repair counter reached {MAX_REPAIRS_PER_TRACE}, skipping repair")
             return
-        if candidate.source == OptimizedCandidateSource.REPAIR:
-            logger.debug("Candidate is already a repair, skipping repair")
+        if candidate.source not in (OptimizedCandidateSource.OPTIMIZE, OptimizedCandidateSource.OPTIMIZE_LP):
+            # only repair the first pass of the candidates for now
+            logger.debug(f"Candidate is a result of {candidate.source.value}, skipping repair")
             return
         if not diffs:
             logger.debug("No diffs found, skipping repair")

codeflash/verification/equivalence.py

@@ -14,6 +14,9 @@
 
 INCREASED_RECURSION_LIMIT = 5000
 
+reprlib_repr = reprlib.Repr(maxstring=1500)
+test_diff_repr = reprlib_repr.repr
+
 
 def compare_test_results(original_results: TestResults, candidate_results: TestResults) -> tuple[bool, list[TestDiff]]:
     # This is meant to be only called with test results for the first loop index
@@ -68,27 +71,13 @@ def compare_test_results(original_results: TestResults, candidate_results: TestR
         )
         if original_pytest_error:
             original_pytest_error = shorten_pytest_error(original_pytest_error)
-        test_src_code = original_test_result.id.get_src_code(original_test_result.file_name)
-        test_diff = TestDiff(
-            scope=TestDiffScope.RETURN_VALUE,
-            original_value=reprlib.repr(original_test_result.return_value),
-            candidate_value=reprlib.repr(cdd_test_result.return_value),
-            test_src_code=test_src_code,
-            candidate_pytest_error=cdd_pytest_error,
-            original_pass=original_test_result.did_pass,
-            candidate_pass=cdd_test_result.did_pass,
-            original_pytest_error=original_pytest_error,
-        )
-        if not comparator(original_test_result.return_value, cdd_test_result.return_value, superset_obj=superset_obj):
-            test_diff.scope = TestDiffScope.RETURN_VALUE
-            test_diffs.append(test_diff)
 
         if not comparator(original_test_result.return_value, cdd_test_result.return_value, superset_obj=superset_obj):
             test_diffs.append(
                 TestDiff(
                     scope=TestDiffScope.RETURN_VALUE,
-                    original_value=repr(original_test_result.return_value),
-                    candidate_value=repr(cdd_test_result.return_value),
+                    original_value=test_diff_repr(repr(original_test_result.return_value)),
+                    candidate_value=test_diff_repr(repr(cdd_test_result.return_value)),
                     test_src_code=original_test_result.id.get_src_code(original_test_result.file_name),
                     candidate_pytest_error=cdd_pytest_error,
                     original_pass=original_test_result.did_pass,