quick todos, modify prompts and function arguments

Author: aseembits93

codeflash/api/aiservice.py

@@ -204,6 +204,65 @@ def optimize_python_code_line_profiler(
         console.rule()
         return []
 
+    def get_new_explanation(
+        self,
+        source_code: str,
+        dependency_code: str,
+        trace_id: str,
+        num_candidates: int = 10,
+        experiment_metadata: ExperimentMetadata | None = None,
+        existing_explanation: str = "",
+    ) -> str:
+        """Optimize the given python code for performance by making a request to the Django endpoint.
+
+        Parameters
+        ----------
+        - source_code (str): The python code to optimize.
+        - dependency_code (str): The dependency code used as read-only context for the optimization
+        - trace_id (str): Trace id of optimization run
+        - num_candidates (int): Number of optimization variants to generate. Default is 10.
+        - experiment_metadata (Optional[ExperimentalMetadata, None]): Any available experiment metadata for this optimization
+        - existing_explanation (str): Existing explanation from AIservice call
+
+        Returns
+        -------
+        - List[OptimizationCandidate]: A list of Optimization Candidates.
+
+        """
+        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,
+            "existing_explanation": existing_explanation,
+        }
+
+        logger.info("Generating optimized candidates…")
+        console.rule()
+        try:
+            response = self.make_ai_service_request("/explain", payload=payload, timeout=600)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error generating optimized candidates: {e}")
+            ph("cli-optimize-error-caught", {"error": str(e)})
+            return ""
+
+        if response.status_code == 200:
+            explanation = response.json()["explanation"]
+            logger.info(f"New Explanation: {explanation}")
+            console.rule()
+            return explanation
+        try:
+            error = response.json()["error"]
+        except Exception:
+            error = response.text
+        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 ""
+
 
     def log_results(
         self,

codeflash/optimization/function_optimizer.py

@@ -261,9 +261,14 @@ def optimize_function(self) -> Result[BestOptimization, str]:
                         best_optimization.candidate.explanation, title="Best Candidate Explanation", border_style="blue"
                     )
                 )
-                get_new_explanation = "abc"
+                #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)
                 explanation = Explanation(
-                    raw_explanation_message=best_optimization.candidate.explanation,
+                    raw_explanation_message=new_explanation if new_explanation!="" else 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,

tests/test_explain_api.py

@@ -0,0 +1,16 @@
+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"
+    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)
+    assert new_explanation.__len__()>0