experiment

Author: aseembits93

codeflash/api/aiservice.py

@@ -544,7 +544,8 @@ def get_optimization_impact(
         replay_tests: str,
         root_dir: Path,
         concolic_tests: str,  # noqa: ARG002
-    ) -> str:
+        calling_fn_details: str,
+    ) -> tuple[str, str]:
         """Compute the optimization impact of current Pull Request.
 
         Args:
@@ -558,6 +559,7 @@ def get_optimization_impact(
         replay_tests: str -> replay test table
         root_dir: Path -> path of git directory
         concolic_tests: str -> concolic_tests (not used)
+        calling_fn_details: str -> filenames and definitions of functions which call the function_to_optimize
 
         Returns:
         -------
@@ -577,13 +579,6 @@ 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,
@@ -598,25 +593,26 @@ def get_optimization_impact(
             "benchmark_details": explanation.benchmark_details if explanation.benchmark_details else None,
             "optimized_runtime": humanize_runtime(explanation.best_runtime_ns),
             "original_runtime": humanize_runtime(explanation.original_runtime_ns),
+            "calling_fn_details": calling_fn_details,
         }
         console.rule()
         try:
             response = self.make_ai_service_request("/optimization_impact", payload=payload, timeout=600)
         except requests.exceptions.RequestException as e:
             logger.exception(f"Error generating optimization refinements: {e}")
             ph("cli-optimize-error-caught", {"error": str(e)})
-            return ""
+            return ("", str(e))
 
         if response.status_code == 200:
-            return cast("str", response.json()["impact"])
+            return (cast("str", response.json()["impact"]), cast("str", response.json()["impact_explanation"]))
         try:
             error = cast("str", response.json()["error"])
         except Exception:
             error = response.text
         logger.error(f"Error generating impact candidates: {response.status_code} - {error}")
         ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
         console.rule()
-        return ""
+        return ("", error)
 
 
 class LocalAiServiceClient(AiServiceClient):

codeflash/api/cfapi.py

@@ -173,6 +173,7 @@ def create_pr(
     coverage_message: str,
     replay_tests: str = "",
     concolic_tests: str = "",
+    optimization_impact: str = "",
 ) -> Response:
     """Create a pull request, targeting the specified branch. (usually 'main').
 
@@ -197,6 +198,7 @@ def create_pr(
         "coverage_message": coverage_message,
         "replayTests": replay_tests,
         "concolicTests": concolic_tests,
+        "optimizationImpact": optimization_impact,
     }
     return make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload)
 
@@ -212,6 +214,7 @@ def create_staging(
     replay_tests: str,
     concolic_tests: str,
     root_dir: Path,
+    optimization_impact: str = "",
 ) -> Response:
     """Create a staging pull request, targeting the specified branch. (usually 'staging').
 
@@ -252,6 +255,7 @@ def create_staging(
         "coverage_message": coverage_message,
         "replayTests": replay_tests,
         "concolicTests": concolic_tests,
+        "optimizationImpact": optimization_impact,
     }
 
     return make_cfapi_request(endpoint="/create-staging", method="POST", payload=payload)

codeflash/code_utils/code_extractor.py

@@ -1,4 +1,3 @@
-# ruff: noqa: ARG002
 from __future__ import annotations
 
 import ast
@@ -10,15 +9,13 @@
 
 import jedi
 import libcst as cst
-import radon.visitors
 from libcst.codemod import CodemodContext
 from libcst.codemod.visitors import AddImportsVisitor, GatherImportsVisitor, RemoveImportsVisitor
 from libcst.helpers import calculate_module_and_package
-from radon.complexity import cc_visit
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.config_consts import MAX_CONTEXT_LEN_IMPACT, TIME_LIMIT_FOR_OPT_IMPACT
-from codeflash.models.models import CodePosition, FunctionParent, ImpactMetrics
+from codeflash.models.models import CodePosition, FunctionParent
 
 if TYPE_CHECKING:
     from libcst.helpers import ModuleNameAndPackage
@@ -38,28 +35,28 @@ def __init__(self) -> None:
         self.scope_depth = 0
         self.if_else_depth = 0
 
-    def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:  # noqa: ARG002
         self.scope_depth += 1
         return True
 
-    def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:
+    def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:  # noqa: ARG002
         self.scope_depth -= 1
 
     def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:
         self.scope_depth += 1
         return True
 
-    def leave_ClassDef(self, original_node: cst.ClassDef) -> None:
+    def leave_ClassDef(self, original_node: cst.ClassDef) -> None:  # noqa: ARG002
         self.scope_depth -= 1
 
-    def visit_If(self, node: cst.If) -> Optional[bool]:
+    def visit_If(self, node: cst.If) -> Optional[bool]:  # noqa: ARG002
         self.if_else_depth += 1
         return True
 
-    def leave_If(self, original_node: cst.If) -> None:
+    def leave_If(self, original_node: cst.If) -> None:  # noqa: ARG002
         self.if_else_depth -= 1
 
-    def visit_Else(self, node: cst.Else) -> Optional[bool]:
+    def visit_Else(self, node: cst.Else) -> Optional[bool]:  # noqa: ARG002
         # Else blocks are already counted as part of the if statement
         return True
 
@@ -86,21 +83,21 @@ def __init__(self, new_assignments: dict[str, cst.Assign], new_assignment_order:
         self.scope_depth = 0
         self.if_else_depth = 0
 
-    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:  # noqa: ARG002
         self.scope_depth += 1
 
     def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
         self.scope_depth -= 1
         return updated_node
 
-    def visit_ClassDef(self, node: cst.ClassDef) -> None:
+    def visit_ClassDef(self, node: cst.ClassDef) -> None:  # noqa: ARG002
         self.scope_depth += 1
 
-    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:
+    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:  # noqa: ARG002
         self.scope_depth -= 1
         return updated_node
 
-    def visit_If(self, node: cst.If) -> None:
+    def visit_If(self, node: cst.If) -> None:  # noqa: ARG002
         self.if_else_depth += 1
 
     def leave_If(self, original_node: cst.If, updated_node: cst.If) -> cst.If:
@@ -1156,8 +1153,7 @@ def get_fn_references_jedi(
 
 def get_opt_impact_metrics(
     source_code: str, file_path: Path, qualified_name: str, project_root: Path, tests_root: Path
-) -> ImpactMetrics:
-    metrics = ImpactMetrics()
+) -> str:
     try:
         qualified_name_split = qualified_name.rsplit(".", maxsplit=1)
         if len(qualified_name_split) == 1:
@@ -1167,26 +1163,8 @@ def get_opt_impact_metrics(
         matches = get_fn_references_jedi(
             source_code, file_path, project_root, target_function, target_class
         )  # jedi is not perfect, it doesn't capture aliased references
-        cyclomatic_complexity_results = cc_visit(source_code)
-        match_found = False
-        for result in cyclomatic_complexity_results:
-            if match_found:
-                break
-            if isinstance(result, radon.visitors.Function) and not target_class:
-                if result.name == target_function:
-                    metrics.cyclomatic_complexity = result.complexity
-                    metrics.cyclomatic_complexity_rating = result.letter
-                    match_found = True
-            elif isinstance(result, radon.visitors.Class) and target_class:  # noqa: SIM102
-                if result.name == target_class:
-                    for method in result.methods:
-                        if match_found:
-                            break
-                        if method.name == target_function:
-                            metrics.cyclomatic_complexity = method.complexity
-                            metrics.cyclomatic_complexity_rating = method.letter
-                            match_found = True
-        metrics.calling_fns = find_occurances(qualified_name, str(file_path), matches, project_root, tests_root)
+        calling_fns_details = find_occurances(qualified_name, str(file_path), matches, project_root, tests_root)
     except Exception as e:
+        calling_fns_details = ""
         logger.debug(f"Investigate {e}")
-    return metrics
+    return calling_fns_details

codeflash/optimization/function_optimizer.py

@@ -24,6 +24,7 @@
 from codeflash.benchmarking.utils import process_benchmark_data
 from codeflash.cli_cmds.console import code_print, console, logger, lsp_log, progress_bar
 from codeflash.code_utils import env_utils
+from codeflash.code_utils.code_extractor import get_opt_impact_metrics
 from codeflash.code_utils.code_replacer import (
     add_custom_marker_to_all_tests,
     modify_autouse_fixture,
@@ -1461,12 +1462,36 @@ def process_review(
 
         if raise_pr or staging_review:
             data["root_dir"] = git_root_dir()
+            calling_fn_details = get_opt_impact_metrics(
+                self.function_to_optimize_source_code,
+                self.function_to_optimize.file_path,
+                self.function_to_optimize.qualified_name,
+                self.project_root,
+                self.test_cfg.tests_root,
+            )
             opt_impact_response = ""
             try:
-                opt_impact_response = self.aiservice_client.get_optimization_impact(**data)
+                opt_impact_response = self.aiservice_client.get_optimization_impact(
+                    **data, calling_fn_details=calling_fn_details
+                )
             except Exception as e:
                 logger.debug(f"optimization impact response failed, investigate {e}")
-            data["optimization_impact"] = opt_impact_response
+            data["optimization_impact"] = opt_impact_response[0]
+            new_explanation_with_opt_explanation = Explanation(
+                raw_explanation_message=f"Impact: {opt_impact_response[0]}\n Impact_explanation: {opt_impact_response[1]} END OF IMPACT EXPLANATION\n"
+                + new_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,
+                original_async_throughput=explanation.original_async_throughput,
+                best_async_throughput=explanation.best_async_throughput,
+            )
+            best_optimization.explanation_v2 = new_explanation_with_opt_explanation.explanation_message()
+            data["explanation"] = new_explanation_with_opt_explanation
         if raise_pr and not staging_review:
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)

codeflash/result/create_pr.py

@@ -277,6 +277,7 @@ def check_create_pr(
             coverage_message=coverage_message,
             replay_tests=replay_tests,
             concolic_tests=concolic_tests,
+            optimization_impact=optimization_impact,
         )
         if response.ok:
             pr_id = response.text