refi optimization ids and original optimization ids

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -26,6 +26,7 @@
 
 @dataclass(frozen=True)
 class AIServiceRefinerRequest:
+    optimization_id: str
     original_source_code: str
     read_only_dependency_code: str
     original_code_runtime: str
@@ -232,10 +233,11 @@ def optimize_python_code_line_profiler(  # noqa: D417
         console.rule()
         return []
 
-    def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]) -> list[str]:
+    def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]) -> dict[str, str]:
         git_repo_owner, git_repo_name = safe_get_repo_owner_and_name()
         payload = [
             {
+                "optimization_id": opt.optimization_id,
                 "original_source_code": opt.original_source_code,
                 "read_only_dependency_code": opt.read_only_dependency_code,
                 "original_line_profiler_results": opt.original_line_profiler_results,
@@ -280,7 +282,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
         except requests.exceptions.RequestException as e:
             logger.exception(f"Error generating optimization refinements: {e}")
             ph("cli-optimize-error-caught", {"error": str(e)})
-            return []
+            return {}
 
         if response.status_code == 200:
             refined_optimizations = response.json()["result"]
@@ -294,7 +296,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
         logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
         ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
         console.rule()
-        return []
+        return {}
 
     def log_results(  # noqa: D417
         self,
@@ -303,6 +305,7 @@ def log_results(  # noqa: D417
         original_runtime: float | None,
         optimized_runtime: dict[str, float | None] | None,
         is_correct: dict[str, bool] | None,
+        metadata: dict[str, any] | None,
     ) -> None:
         """Log features to the database.
 
@@ -313,6 +316,7 @@ def log_results(  # noqa: D417
         - original_runtime (Optional[Dict[str, float]]): The original runtime.
         - optimized_runtime (Optional[Dict[str, float]]): The optimized runtime.
         - is_correct (Optional[Dict[str, bool]]): Whether the optimized code is correct.
+        - metadata (Optional[dict[str, any]]): metadata.
 
         """
         payload = {
@@ -322,6 +326,7 @@ def log_results(  # noqa: D417
             "optimized_runtime": optimized_runtime,
             "is_correct": is_correct,
             "codeflash_version": codeflash_version,
+            "metadata": metadata,
         }
         try:
             self.make_ai_service_request("/log_features", payload=payload, timeout=5)

codeflash/optimization/function_optimizer.py

@@ -547,7 +547,8 @@ def determine_best_candidate(
                         trace_id = self.function_trace_id
                         if trace_id.endswith(("EXP0", "EXP1")):
                             trace_id = trace_id[:-4] + exp_type
-                        refinement_diffs = self.refine_optimizations(
+                        # refinement_dict is a dictionary with optimization_id as a key and the refined code as a value
+                        refinement_dict = self.refine_optimizations(
                             valid_optimizations=self.valid_optimizations,
                             original_code_baseline=original_code_baseline,
                             code_context=code_context,
@@ -561,15 +562,18 @@ def determine_best_candidate(
                             executor=executor,
                             fto_name=self.function_to_optimize.qualified_name,
                         )
-                        # filter out empty strings of code
+
                         more_opt_candidates = [
                             OptimizedCandidate(
-                                source_code=refinement_diffs[i],
-                                explanation=self.valid_optimizations[i].candidate.explanation,
-                                optimization_id=self.valid_optimizations[i].candidate.optimization_id,
+                                source_code=code,
+                                explanation=self.valid_optimizations[
+                                    i
+                                ].candidate.explanation,  # TODO: handle the new explanation after the refinement
+                                optimization_id=opt_id,
                             )
-                            for i in range(len(refinement_diffs))
-                            if refinement_diffs[i] != ""
+                            for i, (opt_id, code) in enumerate(refinement_dict.items())
+                            # filter out empty strings of code
+                            if code != ""
                         ]
                         # we no longer need to apply diffs since we are generating the entire code again
                         candidates.extend(more_opt_candidates)
@@ -637,14 +641,16 @@ def create_rank_dictionary_compact(int_array: list[int]) -> dict[int, int]:
         runtimes_ranking = create_rank_dictionary_compact(runtimes_list)
         overall_ranking = {key: diff_lens_ranking[key] + runtimes_ranking[key] for key in diff_lens_ranking.keys()}  # noqa: SIM118
         min_key = min(overall_ranking, key=overall_ranking.get)
+        best_optimization = self.valid_optimizations[min_key]
         ai_service_client.log_results(
             function_trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
             speedup_ratio=speedup_ratios,
             original_runtime=original_code_baseline.runtime,
             optimized_runtime=optimized_runtimes,
             is_correct=is_correct,
+            metadata={"best_optimization_id": best_optimization.candidate.optimization_id},
         )
-        return self.valid_optimizations[min_key]
+        return best_optimization
 
     def refine_optimizations(
         self,
@@ -656,9 +662,10 @@ def refine_optimizations(
         ai_service_client: AiServiceClient,
         executor: concurrent.futures.ThreadPoolExecutor,
         fto_name: str,
-    ) -> list[str]:
+    ) -> dict[str, str]:
         request = [
             AIServiceRefinerRequest(
+                optimization_id=opt.candidate.optimization_id,
                 original_source_code=code_context.read_writable_code,
                 read_only_dependency_code=code_context.read_only_context_code,
                 original_code_runtime=humanize_runtime(original_code_baseline.runtime),