Optimization Impact

codeflash/api/aiservice.py

@@ -4,7 +4,8 @@
 import os
 import platform
 import time
-from typing import TYPE_CHECKING, Any
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
 
 import requests
 from pydantic.json import pydantic_encoder
@@ -13,18 +14,18 @@
 from codeflash.code_utils.config_consts import N_CANDIDATES_EFFECTIVE, N_CANDIDATES_LP_EFFECTIVE
 from codeflash.code_utils.env_utils import get_codeflash_api_key
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
+from codeflash.code_utils.time_utils import humanize_runtime
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import AIServiceRefinerRequest, CodeStringsMarkdown, OptimizedCandidate
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.version import __version__ as codeflash_version
 
 if TYPE_CHECKING:
-    from pathlib import Path
-
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.models.ExperimentMetadata import ExperimentMetadata
     from codeflash.models.models import AIServiceRefinerRequest
+    from codeflash.result.explanation import Explanation
 
 
 class AiServiceClient:
@@ -516,6 +517,90 @@ def generate_regression_tests(  # noqa: D417
             ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": response.text})
             return None
 
+    def get_optimization_impact(
+        self,
+        original_code: dict[Path, str],
+        new_code: dict[Path, str],
+        explanation: Explanation,
+        existing_tests_source: str,
+        generated_original_test_source: str,
+        function_trace_id: str,
+        coverage_message: str,
+        replay_tests: str,
+        concolic_tests: str,
+        root_dir: Path,
+        original_line_profiler_results: str,
+        optimized_line_profiler_results: str,
+    ) -> str:
+        """Optimize the given python code for performance by making a request to the Django endpoint.
+
+        Args:
+        original_code: dict,
+        new_code: dict,
+        explanation: Explanation,
+        existing_tests_source: str,
+        generated_original_test_source: str,
+        function_trace_id: str,
+        coverage_message: str,
+        replay_tests: str,
+        concolic_tests: str,
+        root_dir: Path,
+        original_line_profiler_results: str,
+        optimized_line_profiler_results: str,
+
+        Returns:
+        -------
+        - 'high','medium' or 'low' optimization impact
+
+        """
+        logger.info("!lsp|Computing Optimization Impact…")
+        original_code_str = ""
+        new_code_str = ""
+        for p, code in original_code.items():
+            original_code_str += f"```python:{Path(p).relative_to(root_dir).as_posix()}"
+            original_code_str += "\n"
+            original_code_str += code
+        for p, code in new_code.items():
+            new_code_str += f"```python:{Path(p).relative_to(root_dir).as_posix()}"
+            new_code_str += "\n"
+            new_code_str += code
+
+        payload = {
+            "original_code": original_code_str,
+            "optimized_code": new_code_str,
+            "existing_tests": existing_tests_source,
+            "generated_tests": generated_original_test_source,
+            "trace_id": function_trace_id,
+            "coverage_message": coverage_message,
+            "replay_tests": replay_tests,
+            "concolic_tests": concolic_tests,
+            "speedup": f"{1 + float(explanation.speedup):.2f}x",
+            "loop_count": explanation.winning_benchmarking_test_results.number_of_loops(),
+            "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),
+            "original_line_profiler_results": original_line_profiler_results,
+            "optimized_line_profiler_results": optimized_line_profiler_results,
+        }
+        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 ""
+
+        if response.status_code == 200:
+            return cast("str", response.json()["impact"])
+        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 ""
+
 
 class LocalAiServiceClient(AiServiceClient):
     """Client for interacting with the local AI service."""

codeflash/optimization/function_optimizer.py

@@ -1410,6 +1410,17 @@ def process_review(
         if raise_pr or self.args.staging_review:
             data["root_dir"] = git_root_dir()
 
+        if raise_pr:
+            # modify argument of staging vs pr based on the impact
+            opt_impact_response = self.aiservice_client.get_optimization_impact(
+                **data,
+                original_line_profiler_results=original_code_baseline.line_profile_results["str_out"],
+                optimized_line_profiler_results=best_optimization.line_profiler_test_results["str_out"],
+            )
+            if opt_impact_response in ["low", "medium"]:
+                raise_pr = False
+                self.args.staging_review = True
+
         if raise_pr and not self.args.staging_review:
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)