precommit mypy fix

Author: aseembits93

codeflash/api/aiservice.py

@@ -54,7 +54,11 @@ def get_aiservice_base_url(self) -> str:
         return "https://app.codeflash.ai"
 
     def make_ai_service_request(
-        self, endpoint: str, method: str = "POST", payload: dict[str, Any] | None = None, timeout: float | None = None
+        self,
+        endpoint: str,
+        method: str = "POST",
+        payload: dict[str, Any] | list[dict[str, Any]] | None = None,
+        timeout: float | None = None,
     ) -> requests.Response:
         """Make an API request to the given endpoint on the AI service.
 
@@ -292,7 +296,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
                 OptimizedCandidate(
                     source_code=opt["source_code"],
                     explanation=opt["explanation"],
-                    optimization_id=opt["optimization_id"][:-4]+"refi",
+                    optimization_id=opt["optimization_id"][:-4] + "refi",
                 )
                 for opt in refined_optimizations
             ]
@@ -336,7 +340,7 @@ def log_results(  # noqa: D417
             "is_correct": is_correct,
             "codeflash_version": codeflash_version,
             "best_optimization_id": best_optimization_id,
-            "optimized_line_profiler_results": optimized_line_profiler_results
+            "optimized_line_profiler_results": optimized_line_profiler_results,
         }
         try:
             self.make_ai_service_request("/log_features", payload=payload, timeout=5)

codeflash/code_utils/code_utils.py

@@ -19,6 +19,7 @@
 
 ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
 
+
 def diff_length(a: str, b: str) -> int:
     """Compute the length (in characters) of the unified diff between two strings.
 
@@ -62,6 +63,7 @@ def create_rank_dictionary_compact(int_array: list[int]) -> dict[int, int]:
     # Create a dictionary mapping the original index to its rank (its position in the sorted list)
     return {original_index: rank for rank, original_index in enumerate(sorted_indices)}
 
+
 @contextmanager
 def custom_addopts() -> None:
     pyproject_file = find_pyproject_toml()

codeflash/optimization/function_optimizer.py

@@ -2,7 +2,6 @@
 
 import ast
 import concurrent.futures
-import difflib
 import os
 import random
 import subprocess
@@ -32,13 +31,13 @@
 from codeflash.code_utils.code_utils import (
     ImportErrorPattern,
     cleanup_paths,
+    create_rank_dictionary_compact,
+    diff_length,
     file_name_from_test_module_name,
     get_run_tmp_file,
     has_any_async_functions,
     module_name_from_file_path,
     restore_conftest,
-    diff_length,
-    create_rank_dictionary_compact,
 )
 from codeflash.code_utils.config_consts import (
     INDIVIDUAL_TESTCASE_TIMEOUT,
@@ -363,7 +362,6 @@ def determine_best_candidate(
         exp_type: str,
     ) -> BestOptimization | None:
         # TODO remove
-        from codeflash.models.models import OptimizedCandidate
 
         best_optimization: BestOptimization | None = None
         _best_runtime_until_now = original_code_baseline.runtime
@@ -482,7 +480,9 @@ def determine_best_candidate(
                                 original_helper_code=original_helper_code,
                                 candidate_index=candidate_index,
                             )
-                            optimized_line_profiler_results[candidate.optimization_id]=line_profile_test_results['str_out']
+                            optimized_line_profiler_results[candidate.optimization_id] = line_profile_test_results[
+                                "str_out"
+                            ]
                             replay_perf_gain = {}
                             if self.args.benchmark:
                                 test_results_by_benchmark = (
@@ -580,10 +580,12 @@ def determine_best_candidate(
         if not len(self.valid_optimizations):
             return None
         # need to figure out the best candidate here before we return best_optimization
-        diff_lens_list = [] # character level diff
+        diff_lens_list = []  # character level diff
         runtimes_list = []
         for valid_opt in self.valid_optimizations:
-            diff_lens_list.append(diff_length(valid_opt.candidate.source_code, code_context.read_writable_code)) #char level diff
+            diff_lens_list.append(
+                diff_length(valid_opt.candidate.source_code, code_context.read_writable_code)
+            )  # char level diff
             runtimes_list.append(valid_opt.runtime)
         diff_lens_ranking = create_rank_dictionary_compact(diff_lens_list)
         runtimes_ranking = create_rank_dictionary_compact(runtimes_list)
@@ -598,7 +600,7 @@ def determine_best_candidate(
             optimized_runtime=optimized_runtimes,
             is_correct=is_correct,
             best_optimization_id=best_optimization.candidate.optimization_id,
-            optimized_line_profiler_results= optimized_line_profiler_results
+            optimized_line_profiler_results=optimized_line_profiler_results,
         )
         return best_optimization
 
@@ -630,7 +632,7 @@ def refine_optimizations(
                 fto_name=fto_name,
             )
             for opt in valid_optimizations
-        ] # TODO: multiple workers for this?
+        ]  # TODO: multiple workers for this?
         future_refinement_results = executor.submit(ai_service_client.optimize_python_code_refinement, request=request)
         concurrent.futures.wait([future_refinement_results])
         return future_refinement_results.result()