Merge branch 'main' into lsp/init-flow

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -577,6 +577,13 @@ def get_optimization_impact(
             ]
         )
         code_diff = f"```diff\n{diff_str}\n```"
+        # TODO get complexity metrics and fn call heuristics -> constructing a complete static call graph can be expensive for really large repos
+        # grep function name in codebase -> ast parser to get no of calls and no of calls in loop -> radon lib to get complexity metrics -> send as additional context to the AI service
+        # metric 1 -> call count - how many times the function is called in the codebase
+        # metric 2 -> loop call count - how many times the function is called in a loop in the codebase
+        # metric 3 -> presence of decorators like @profile, @cache -> this means the owner of the repo cares about the performance of this function
+        # metric 4 -> cyclomatic complexity (https://en.wikipedia.org/wiki/Cyclomatic_complexity)
+        # metric 5 (for future) -> halstead complexity (https://en.wikipedia.org/wiki/Halstead_complexity_measures)
         logger.info("!lsp|Computing Optimization Impact…")
         payload = {
             "code_diff": code_diff,

codeflash/api/cfapi.py

@@ -130,6 +130,7 @@ def suggest_changes(
     coverage_message: str,
     replay_tests: str = "",
     concolic_tests: str = "",
+    optimization_impact: str = "",
 ) -> Response:
     """Suggest changes to a pull request.
 
@@ -155,6 +156,7 @@ def suggest_changes(
         "coverage_message": coverage_message,
         "replayTests": replay_tests,
         "concolicTests": concolic_tests,
+        "optimizationImpact": optimization_impact,
     }
     return make_cfapi_request(endpoint="/suggest-pr-changes", method="POST", payload=payload)
 

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -2,9 +2,10 @@
 
 from typing import TYPE_CHECKING, Optional, Union
 
-import isort
 import libcst as cst
 
+from codeflash.code_utils.formatter import sort_imports
+
 if TYPE_CHECKING:
     from pathlib import Path
 
@@ -107,7 +108,7 @@ def instrument_codeflash_trace_decorator(file_to_funcs_to_optimize: dict[Path, l
         original_code = file_path.read_text(encoding="utf-8")
         new_code = add_codeflash_decorator_to_code(original_code, functions_to_optimize)
         # Modify the code
-        modified_code = isort.code(code=new_code, float_to_top=True)
+        modified_code = sort_imports(code=new_code, float_to_top=True)
 
         # Write the modified code back to the file
         file_path.write_text(modified_code, encoding="utf-8")

codeflash/benchmarking/replay_test.py

@@ -6,9 +6,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-import isort
-
 from codeflash.cli_cmds.console import logger
+from codeflash.code_utils.formatter import sort_imports
 from codeflash.discovery.functions_to_optimize import inspect_top_level_functions_or_methods
 from codeflash.verification.verification_utils import get_test_file_path
 
@@ -299,7 +298,7 @@ def generate_replay_test(
                 test_framework=test_framework,
                 max_run_count=max_run_count,
             )
-            test_code = isort.code(test_code)
+            test_code = sort_imports(code=test_code)
             output_file = get_test_file_path(
                 test_dir=Path(output_dir), function_name=benchmark_module_path, test_type="replay"
             )

codeflash/code_utils/code_replacer.py

@@ -5,13 +5,13 @@
 from functools import lru_cache
 from typing import TYPE_CHECKING, Optional, TypeVar
 
-import isort
 import libcst as cst
 from libcst.metadata import PositionProvider
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_extractor import add_global_assignments, add_needed_imports_from_module
 from codeflash.code_utils.config_parser import find_conftest_files
+from codeflash.code_utils.formatter import sort_imports
 from codeflash.code_utils.line_profile_utils import ImportAdder
 from codeflash.models.models import FunctionParent
 
@@ -226,7 +226,7 @@ def add_custom_marker_to_all_tests(test_paths: list[Path]) -> None:
         module = cst.parse_module(file_content)
         importadder = ImportAdder("import pytest")
         modified_module = module.visit(importadder)
-        modified_module = cst.parse_module(isort.code(modified_module.code, float_to_top=True))
+        modified_module = cst.parse_module(sort_imports(code=modified_module.code, float_to_top=True))
         pytest_mark_adder = PytestMarkAdder("codeflash_no_autouse")
         modified_module = modified_module.visit(pytest_mark_adder)
         test_path.write_text(modified_module.code, encoding="utf-8")

codeflash/code_utils/code_utils.py

@@ -304,7 +304,9 @@ def get_all_function_names(code: str) -> tuple[bool, list[str]]:
     return True, function_names
 
 
-def get_run_tmp_file(file_path: Path) -> Path:
+def get_run_tmp_file(file_path: Path | str) -> Path:
+    if isinstance(file_path, str):
+        file_path = Path(file_path)
     if not hasattr(get_run_tmp_file, "tmpdir"):
         get_run_tmp_file.tmpdir = TemporaryDirectory(prefix="codeflash_")
     return Path(get_run_tmp_file.tmpdir.name) / file_path

codeflash/code_utils/formatter.py

@@ -166,11 +166,11 @@ def format_code(
         return formatted_code
 
 
-def sort_imports(code: str) -> str:
+def sort_imports(code: str, *, float_to_top: bool = False) -> str:
     try:
         # Deduplicate and sort imports, modify the code in memory, not on disk
-        sorted_code = isort.code(code)
-    except Exception:
+        sorted_code = isort.code(code=code, float_to_top=float_to_top)
+    except Exception:  # this will also catch the FileSkipComment exception, use this fn everywhere
         logger.exception("Failed to sort imports with isort.")
         return code  # Fall back to original code if isort fails
 

codeflash/code_utils/instrument_existing_tests.py

@@ -10,6 +10,7 @@
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
+from codeflash.code_utils.formatter import sort_imports
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import FunctionParent, TestingMode, VerificationType
 
@@ -1129,7 +1130,7 @@ def add_async_decorator_to_function(
             import_transformer = AsyncDecoratorImportAdder(mode)
             module = module.visit(import_transformer)
 
-        return isort.code(module.code, float_to_top=True), decorator_transformer.added_decorator
+        return sort_imports(code=module.code, float_to_top=True), decorator_transformer.added_decorator
     except Exception as e:
         logger.exception(f"Error adding async decorator to function {function.qualified_name}: {e}")
         return source_code, False

codeflash/code_utils/line_profile_utils.py

@@ -6,10 +6,10 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Union
 
-import isort
 import libcst as cst
 
 from codeflash.code_utils.code_utils import get_run_tmp_file
+from codeflash.code_utils.formatter import sort_imports
 
 if TYPE_CHECKING:
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
@@ -213,7 +213,7 @@ def add_decorator_imports(function_to_optimize: FunctionToOptimize, code_context
         transformer = ImportAdder("from line_profiler import profile as codeflash_line_profile")
         # Apply the transformer to add the import
         module_node = module_node.visit(transformer)
-        modified_code = isort.code(module_node.code, float_to_top=True)
+        modified_code = sort_imports(code=module_node.code, float_to_top=True)
         # write to file
         with file_path.open("w", encoding="utf-8") as file:
             file.write(modified_code)