relevant context for new explanations

Author: aseembits93

codeflash/api/aiservice.py

@@ -207,11 +207,10 @@ def optimize_python_code_line_profiler(
     def get_new_explanation(
         self,
         source_code: str,
+        optimized_code: str,
         dependency_code: str,
         trace_id: str,
-        num_candidates: int = 10,
-        experiment_metadata: ExperimentMetadata | None = None,
-        existing_explanation: str = "",
+        existing_explanation: str,
     ) -> str:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -230,16 +229,12 @@ def get_new_explanation(
 
         """
         payload = {
-            "source_code": source_code,
-            "dependency_code": dependency_code,
-            "num_variants": num_candidates,
             "trace_id": trace_id,
-            "python_version": platform.python_version(),
-            "experiment_metadata": experiment_metadata,
-            "codeflash_version": codeflash_version,
+            "source_code": source_code,
+            "optimized_code":optimized_code,
             "existing_explanation": existing_explanation,
+            "dependency_code": dependency_code,
         }
-
         logger.info("Generating optimized candidates…")
         console.rule()
         try:

codeflash/optimization/function_optimizer.py

@@ -8,6 +8,7 @@
 import time
 import uuid
 from collections import defaultdict, deque
+from dataclasses import replace
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -254,29 +255,28 @@ def optimize_function(self) -> Result[BestOptimization, str]:
             )
 
             if best_optimization:
-                logger.info("Best candidate:")
-                code_print(best_optimization.candidate.source_code)
-                console.print(
-                    Panel(
-                        best_optimization.candidate.explanation, title="Best Candidate Explanation", border_style="blue"
-                    )
-                )
-                #could possibly have it in the best optimization dataclass
                 new_explanation = self.aiservice_client.get_new_explanation(source_code=code_context.read_writable_code,
                 dependency_code=code_context.read_only_context_code,
                 trace_id=self.function_trace_id,
                 num_candidates=1,
                 experiment_metadata=None, existing_explanation=best_optimization.candidate.explanation)
+                best_optimization.candidate = replace(best_optimization.candidate, explanation=new_explanation if new_explanation!="" else best_optimization.candidate.explanation)
                 explanation = Explanation(
-                    raw_explanation_message=new_explanation if new_explanation!="" else best_optimization.candidate.explanation,
+                    raw_explanation_message=best_optimization.candidate.explanation,
                     winning_behavioral_test_results=best_optimization.winning_behavioral_test_results,
                     winning_benchmarking_test_results=best_optimization.winning_benchmarking_test_results,
                     original_runtime_ns=original_code_baseline.runtime,
                     best_runtime_ns=best_optimization.runtime,
                     function_name=function_to_optimize_qualified_name,
                     file_path=self.function_to_optimize.file_path,
                 )
-
+                logger.info("Best candidate:")
+                code_print(best_optimization.candidate.source_code)
+                console.print(
+                    Panel(
+                        best_optimization.candidate.explanation, title="Best Candidate Explanation", border_style="blue"
+                    )
+                )
                 self.log_successful_optimization(explanation, generated_tests)
 
                 self.replace_function_and_helpers_with_optimized_code(

tests/test_explain_api.py

@@ -1,16 +1,22 @@
 from codeflash.api.aiservice import AiServiceClient
-from codeflash.models.ExperimentMetadata import ExperimentMetadata
 def test_explain_api():
     aiservice = AiServiceClient()
-    source_code: str = "a"
-    dependency_code: str = "b"
+    source_code: str = """def bubble_sort(arr):
+    n = len(arr)
+    for i in range(n):
+        for j in range(0, n-i-1):
+            if arr[j] > arr[j+1]:
+                arr[j], arr[j+1] = arr[j+1], arr[j]
+    return arr
+"""
+    dependency_code: str = "def helper(): return 1"
     trace_id: str = "d5822364-7617-4389-a4fc-64602a00b714"
-    num_candidates: int = 1
-    experiment_metadata: ExperimentMetadata | None = None
-    existing_explanation: str = "some explanation"
-    new_explanation = aiservice.get_new_explanation(source_code=source_code,
-                dependency_code=dependency_code,
-                trace_id=trace_id,
-                num_candidates=num_candidates,
-                experiment_metadata=experiment_metadata, existing_explanation=existing_explanation)
+    existing_explanation: str = "I used to numpy to optimize it"
+    optimized_code: str = """def bubble_sort(arr):
+    return arr.sort()
+"""
+    new_explanation = aiservice.get_new_explanation(source_code=source_code, optimized_code=optimized_code,
+                existing_explanation=existing_explanation, dependency_code=dependency_code,
+                trace_id=trace_id)
+    print("\nNew explanation: \n", new_explanation)
     assert new_explanation.__len__()>0