Merge branch 'main' into local-setup

Author: misrasaurabh1

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/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)

codeflash/optimization/function_optimizer.py

@@ -12,7 +12,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-import isort
 import libcst as cst
 from rich.console import Group
 from rich.panel import Panel
@@ -900,7 +899,7 @@ def reformat_code_and_helpers(
         optimized_context: CodeStringsMarkdown,
     ) -> tuple[str, dict[Path, str]]:
         should_sort_imports = not self.args.disable_imports_sorting
-        if should_sort_imports and isort.code(original_code) != original_code:
+        if should_sort_imports and sort_imports(code=original_code) != original_code:
             should_sort_imports = False
 
         optimized_code = ""
@@ -1461,14 +1460,12 @@ def process_review(
 
         if raise_pr or staging_review:
             data["root_dir"] = git_root_dir()
-            # try:
-            #     # modify argument of staging vs pr based on the impact
-            #     opt_impact_response = self.aiservice_client.get_optimization_impact(**data)
-            #     if opt_impact_response == "low":
-            #         raise_pr = False
-            #         staging_review = True
-            # except Exception as e:
-            #     logger.debug(f"optimization impact response failed, investigate {e}")
+            opt_impact_response = ""
+            try:
+                opt_impact_response = self.aiservice_client.get_optimization_impact(**data)
+            except Exception as e:
+                logger.debug(f"optimization impact response failed, investigate {e}")
+            data["optimization_impact"] = opt_impact_response
         if raise_pr and not staging_review:
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)

codeflash/result/create_pr.py

@@ -185,6 +185,7 @@ def check_create_pr(
     concolic_tests: str,
     root_dir: Path,
     git_remote: Optional[str] = None,
+    optimization_impact: str = "",
 ) -> None:
     pr_number: Optional[int] = env_utils.get_pr_number()
     git_repo = git.Repo(search_parent_directories=True)
@@ -226,6 +227,7 @@ def check_create_pr(
             coverage_message=coverage_message,
             replay_tests=replay_tests,
             concolic_tests=concolic_tests,
+            optimization_impact=optimization_impact,
         )
         if response.ok:
             logger.info(f"Suggestions were successfully made to PR #{pr_number}")