optimization source

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -18,7 +18,12 @@
 from codeflash.code_utils.time_utils import humanize_runtime
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
-from codeflash.models.models import AIServiceRefinerRequest, CodeStringsMarkdown, OptimizedCandidate
+from codeflash.models.models import (
+    AIServiceRefinerRequest,
+    CodeStringsMarkdown,
+    OptimizedCandidate,
+    OptimizedCandidateSource,
+)
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.version import __version__ as codeflash_version
 
@@ -86,15 +91,20 @@ def make_ai_service_request(
         # response.raise_for_status()  # Will raise an HTTPError if the HTTP request returned an unsuccessful status code
         return response
 
-    def _get_valid_candidates(self, optimizations_json: list[dict[str, Any]]) -> list[OptimizedCandidate]:
+    def _get_valid_candidates(
+        self, optimizations_json: list[dict[str, Any]], source: OptimizedCandidateSource
+    ) -> list[OptimizedCandidate]:
         candidates: list[OptimizedCandidate] = []
         for opt in optimizations_json:
             code = CodeStringsMarkdown.parse_markdown_code(opt["source_code"])
             if not code.code_strings:
                 continue
             candidates.append(
                 OptimizedCandidate(
-                    source_code=code, explanation=opt["explanation"], optimization_id=opt["optimization_id"]
+                    source_code=code,
+                    explanation=opt["explanation"],
+                    optimization_id=opt["optimization_id"],
+                    source=source,
                 )
             )
         return candidates
@@ -157,7 +167,7 @@ def optimize_python_code(  # noqa: D417
             console.rule()
             end_time = time.perf_counter()
             logger.debug(f"!lsp|Generating possible optimizations took {end_time - start_time:.2f} seconds.")
-            return self._get_valid_candidates(optimizations_json)
+            return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.OPTIMIZE)
         try:
             error = response.json()["error"]
         except Exception:
@@ -222,7 +232,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
                 f"!lsp|Generated {len(optimizations_json)} candidate optimizations using line profiler information."
             )
             console.rule()
-            return self._get_valid_candidates(optimizations_json)
+            return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.OPTIMIZE_LP)
         try:
             error = response.json()["error"]
         except Exception:
@@ -275,15 +285,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
             logger.debug(f"Generated {len(refined_optimizations)} candidate refinements.")
             console.rule()
 
-            refinements = self._get_valid_candidates(refined_optimizations)
-            return [
-                OptimizedCandidate(
-                    source_code=c.source_code,
-                    explanation=c.explanation,
-                    optimization_id=c.optimization_id[:-4] + "refi",
-                )
-                for c in refinements
-            ]
+            return self._get_valid_candidates(refined_optimizations, OptimizedCandidateSource.REFINE)
 
         try:
             error = response.json()["error"]
@@ -324,7 +326,7 @@ def optimize_python_code_repair(self, request: AIServiceCodeRepairRequest) -> Op
             fixed_optimization = response.json()
             console.rule()
 
-            valid_candidates = self._get_valid_candidates([fixed_optimization])
+            valid_candidates = self._get_valid_candidates([fixed_optimization], OptimizedCandidateSource.REPAIR)
             if not valid_candidates:
                 logger.error("Code repair failed to generate a valid candidate.")
                 return None

codeflash/models/models.py

@@ -383,11 +383,19 @@ class TestsInFile:
     test_type: TestType
 
 
+class OptimizedCandidateSource(enum.Enum, str):
+    OPTIMIZE = "OPTIMIZE"
+    OPTIMIZE_LP = "OPTIMIZE_LP"
+    REFINE = "REFINE"
+    REPAIR = "REPAIR"
+
+
 @dataclass(frozen=True)
 class OptimizedCandidate:
     source_code: CodeStringsMarkdown
     explanation: str
     optimization_id: str
+    source: OptimizedCandidateSource
 
 
 @dataclass(frozen=True)

codeflash/optimization/function_optimizer.py

@@ -77,6 +77,7 @@
     OptimizationSet,
     OptimizedCandidate,
     OptimizedCandidateResult,
+    OptimizedCandidateSource,
     OriginalCodeBaseline,
     TestFile,
     TestFiles,
@@ -739,7 +740,7 @@ def determine_best_candidate(
                         )
                         valid_optimizations.append(best_optimization)
                         # queue corresponding refined optimization for best optimization
-                        if not candidate.optimization_id.endswith("refi"):
+                        if candidate.source != OptimizedCandidateSource.REFINE:
                             self.future_all_refinements.append(
                                 self.refine_optimizations(
                                     valid_optimizations=[best_optimization],
@@ -1908,14 +1909,10 @@ def run_optimized_candidate(
                 console.rule()
             else:
                 result_unmatched_perc = len(diffs) / len(candidate_behavior_results)
-                if result_unmatched_perc > 0.5:
+                if candidate.source == OptimizedCandidateSource.REPAIR or result_unmatched_perc > 0.5:
                     # if the test unmatched percentage is greater than 50%, we can't fix it
                     return self.get_results_not_matched_error()
 
-                if candidate.optimization_id.endswith("cdrp"):
-                    # prevent looping for now
-                    return self.get_results_not_matched_error()
-
                 ai_service_client = self.aiservice_client if exp_type == "EXP0" else self.local_aiservice_client
                 logger.info("Adding this to the repair queue")
                 self.future_all_code_repair.append(