precommit mypy fix

Author: aseembits93

codeflash/api/aiservice.py

@@ -278,7 +278,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
         console.rule()
         return []
 
-    def get_new_explanation(
+    def get_new_explanation(  # noqa: D417
         self,
         source_code: str,
         optimized_code: str,
@@ -297,11 +297,15 @@ def get_new_explanation(
         Parameters
         ----------
         - source_code (str): The python code to optimize.
+        - optimized_code (str): The python code generated by the AI service.
         - 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
+        - original_line_profiler_results: str - line profiler results for the baseline code
+        - optimized_line_profiler_results: str - line profiler results for the optimized code
+        - original_code_runtime: str - runtime for the baseline code
+        - optimized_code_runtime: str - runtime for the optimized code
+        - speedup: str - speedup of the optimized code
+        - annotated_tests: str - test functions annotated with runtime
+        - optimization_id: str - unique id of opt candidate
 
         Returns
         -------
@@ -331,7 +335,7 @@ def get_new_explanation(
             return ""
 
         if response.status_code == 200:
-            explanation = response.json()["explanation"]
+            explanation: str = response.json()["explanation"]
             logger.info(f"New Explanation: {explanation}")
             console.rule()
             return explanation
@@ -344,8 +348,7 @@ def get_new_explanation(
         console.rule()
         return ""
 
-
-    def log_results(
+    def log_results(  # noqa: D417
         self,
         function_trace_id: str,
         speedup_ratio: dict[str, float | None] | None,

codeflash/optimization/function_optimizer.py

@@ -8,7 +8,6 @@
 import time
 import uuid
 from collections import defaultdict, deque
-from dataclasses import replace
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -1124,7 +1123,9 @@ def find_and_process_best_optimization(
                     new_explanation_raw_str = 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[:-4] + exp_type if self.experiment_id else self.function_trace_id,
+                        trace_id=self.function_trace_id[:-4] + exp_type
+                        if self.experiment_id
+                        else self.function_trace_id,
                         optimized_code=best_optimization.candidate.source_code,
                         original_line_profiler_results=original_code_baseline.line_profile_results["str_out"],
                         optimized_line_profiler_results=best_optimization.line_profiler_test_results["str_out"],
@@ -1134,13 +1135,16 @@ def find_and_process_best_optimization(
                         annotated_tests=generated_tests_str,
                         optimization_id=best_optimization.candidate.optimization_id,
                     )
-                    new_explanation = Explanation(raw_explanation_message=new_explanation_raw_str or explanation.raw_explanation_message, winning_behavior_test_results=explanation.winning_behavior_test_results,
-                                                    winning_benchmarking_test_results=explanation.winning_benchmarking_test_results,
-                                                    original_runtime_ns=explanation.original_runtime_ns,
-                                                    best_runtime_ns=explanation.best_runtime_ns,
-                                                    function_name=explanation.function_name,
-                                                    file_path=explanation.file_path,
-                                                    benchmark_details=explanation.benchmark_details)
+                    new_explanation = Explanation(
+                        raw_explanation_message=new_explanation_raw_str or explanation.raw_explanation_message,
+                        winning_behavior_test_results=explanation.winning_behavior_test_results,
+                        winning_benchmarking_test_results=explanation.winning_benchmarking_test_results,
+                        original_runtime_ns=explanation.original_runtime_ns,
+                        best_runtime_ns=explanation.best_runtime_ns,
+                        function_name=explanation.function_name,
+                        file_path=explanation.file_path,
+                        benchmark_details=explanation.benchmark_details,
+                    )
                     check_create_pr(
                         original_code=original_code_combined,
                         new_code=new_code_combined,