cleaning up

Author: aseembits93

codeflash/api/aiservice.py

@@ -316,7 +316,6 @@ def log_results(  # noqa: D417
         original_runtime: float | None,
         optimized_runtime: dict[str, float | None] | None,
         is_correct: dict[str, bool] | None,
-        best_optimization_id: str | None,
         optimized_line_profiler_results: dict[str, str] | None,
     ) -> None:
         """Log features to the database.
@@ -328,7 +327,6 @@ 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.
-        - best_optimization_id (Optional[str]): The best optimization id.
         -optimized_line_profiler_results: line_profiler results for every candidate mapped to their optimization_id
 
         """
@@ -339,7 +337,6 @@ def log_results(  # noqa: D417
             "optimized_runtime": optimized_runtime,
             "is_correct": is_correct,
             "codeflash_version": codeflash_version,
-            "best_optimization_id": best_optimization_id,
             "optimized_line_profiler_results": optimized_line_profiler_results,
         }
         try:

codeflash/models/models.py

@@ -75,6 +75,7 @@ def __hash__(self) -> int:
 
 class BestOptimization(BaseModel):
     candidate: OptimizedCandidate
+    helper_functions: list[FunctionSource]
     code_context: CodeOptimizationContext
     runtime: int
     replay_performance_gain: Optional[dict[BenchmarkKey, float]] = None

codeflash/optimization/function_optimizer.py

@@ -361,8 +361,6 @@ def determine_best_candidate(
         file_path_to_helper_classes: dict[Path, set[str]],
         exp_type: str,
     ) -> BestOptimization | None:
-        # TODO remove
-
         best_optimization: BestOptimization | None = None
         _best_runtime_until_now = original_code_baseline.runtime
 
@@ -507,6 +505,7 @@ def determine_best_candidate(
 
                             best_optimization = BestOptimization(
                                 candidate=candidate,
+                                helper_functions=code_context.helper_functions,
                                 code_context=code_context,
                                 runtime=best_test_runtime,
                                 line_profiler_test_results=line_profile_test_results,