Merge branch 'main' of github.com:codeflash-ai/codeflash into feat/markdown-read-writable-context

Author: mohammedahmed18

.github/workflows/deploy-docs-to-azure.yaml

@@ -278,6 +278,79 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
         console.rule()
         return []
 
+    def get_new_explanation(  # noqa: D417
+        self,
+        source_code: str,
+        optimized_code: str,
+        dependency_code: str,
+        trace_id: str,
+        original_line_profiler_results: str,
+        optimized_line_profiler_results: str,
+        original_code_runtime: str,
+        optimized_code_runtime: str,
+        speedup: str,
+        annotated_tests: str,
+        optimization_id: str,
+        original_explanation: str,
+    ) -> str:
+        """Optimize the given python code for performance by making a request to the Django endpoint.
+
+        Parameters
+        ----------
+        - source_code (str): The python code to optimize.
+        - optimized_code (str): The python code generated by the AI service.
+        - dependency_code (str): The dependency code used as read-only context for the optimization
+        - original_line_profiler_results: str - line profiler results for the baseline code
+        - optimized_line_profiler_results: str - line profiler results for the optimized code
+        - original_code_runtime: str - runtime for the baseline code
+        - optimized_code_runtime: str - runtime for the optimized code
+        - speedup: str - speedup of the optimized code
+        - annotated_tests: str - test functions annotated with runtime
+        - optimization_id: str - unique id of opt candidate
+        - original_explanation: str - original_explanation generated for the opt candidate
+
+        Returns
+        -------
+        - List[OptimizationCandidate]: A list of Optimization Candidates.
+
+        """
+        payload = {
+            "trace_id": trace_id,
+            "source_code": source_code,
+            "optimized_code": optimized_code,
+            "original_line_profiler_results": original_line_profiler_results,
+            "optimized_line_profiler_results": optimized_line_profiler_results,
+            "original_code_runtime": original_code_runtime,
+            "optimized_code_runtime": optimized_code_runtime,
+            "speedup": speedup,
+            "annotated_tests": annotated_tests,
+            "optimization_id": optimization_id,
+            "original_explanation": original_explanation,
+            "dependency_code": dependency_code,
+        }
+        logger.info("Generating explanation")
+        console.rule()
+        try:
+            response = self.make_ai_service_request("/explain", payload=payload, timeout=60)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error generating explanations: {e}")
+            ph("cli-optimize-error-caught", {"error": str(e)})
+            return ""
+
+        if response.status_code == 200:
+            explanation: str = response.json()["explanation"]
+            logger.debug(f"New Explanation: {explanation}")
+            console.rule()
+            return explanation
+        try:
+            error = response.json()["error"]
+        except Exception:
+            error = response.text
+        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
+        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        console.rule()
+        return ""
+
     def log_results(  # noqa: D417
         self,
         function_trace_id: str,

codeflash/api/aiservice.py

@@ -14,13 +14,15 @@
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
-from codeflash.code_utils.git_utils import get_repo_owner_and_name
+from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name, git_root_dir
+from codeflash.github.PrComment import FileDiffContent, PrComment
 from codeflash.version import __version__
 
 if TYPE_CHECKING:
     from requests import Response
 
-    from codeflash.github.PrComment import FileDiffContent, PrComment
+    from codeflash.result.explanation import Explanation
+
 from packaging import version
 
 if os.environ.get("CODEFLASH_CFAPI_SERVER", default="prod").lower() == "local":
@@ -182,6 +184,59 @@ def create_pr(
     return make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload)
 
 
+def create_staging(
+    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,
+) -> Response:
+    """Create a staging pull request, targeting the specified branch. (usually 'staging').
+
+    :param original_code: A mapping of file paths to original source code.
+    :param new_code: A mapping of file paths to optimized source code.
+    :param explanation: An Explanation object with optimization details.
+    :param existing_tests_source: Existing test code.
+    :param generated_original_test_source: Generated tests for the original function.
+    :param function_trace_id: Unique identifier for this optimization trace.
+    :param coverage_message: Coverage report or summary.
+    :return: The response object from the backend.
+    """
+    relative_path = explanation.file_path.relative_to(git_root_dir()).as_posix()
+
+    build_file_changes = {
+        Path(p).relative_to(git_root_dir()).as_posix(): FileDiffContent(
+            oldContent=original_code[p], newContent=new_code[p]
+        )
+        for p in original_code
+    }
+
+    payload = {
+        "baseBranch": get_current_branch(),
+        "diffContents": build_file_changes,
+        "prCommentFields": PrComment(
+            optimization_explanation=explanation.explanation_message(),
+            best_runtime=explanation.best_runtime_ns,
+            original_runtime=explanation.original_runtime_ns,
+            function_name=explanation.function_name,
+            relative_file_path=relative_path,
+            speedup_x=explanation.speedup_x,
+            speedup_pct=explanation.speedup_pct,
+            winning_behavior_test_results=explanation.winning_behavior_test_results,
+            winning_benchmarking_test_results=explanation.winning_benchmarking_test_results,
+            benchmark_details=explanation.benchmark_details,
+        ).to_json(),
+        "existingTests": existing_tests_source,
+        "generatedTests": generated_original_test_source,
+        "traceId": function_trace_id,
+        "coverage_message": coverage_message,
+    }
+
+    return make_cfapi_request(endpoint="/create-staging", method="POST", payload=payload)
+
+
 def is_github_app_installed_on_repo(owner: str, repo: str, *, suppress_errors: bool = False) -> bool:
     """Check if the Codeflash GitHub App is installed on the specified repository.
 

codeflash/api/cfapi.py

@@ -77,6 +77,7 @@ def parse_args() -> Namespace:
     parser.add_argument(
         "--no-pr", action="store_true", help="Do not create a PR for the optimization, only update the code locally."
     )
+    parser.add_argument("--staging-review", action="store_true", help="Upload optimizations to staging for review")
     parser.add_argument(
         "--verify-setup",
         action="store_true",

codeflash/cli_cmds/cli.py

@@ -19,7 +19,7 @@
 from rich.tree import Tree
 
 from codeflash.api.aiservice import AiServiceClient, AIServiceRefinerRequest, LocalAiServiceClient
-from codeflash.api.cfapi import add_code_context_hash, mark_optimization_success
+from codeflash.api.cfapi import add_code_context_hash, create_staging, mark_optimization_success
 from codeflash.benchmarking.utils import process_benchmark_data
 from codeflash.cli_cmds.console import code_print, console, logger, progress_bar
 from codeflash.code_utils import env_utils
@@ -1104,70 +1104,124 @@ def find_and_process_best_optimization(
                 original_code_combined[explanation.file_path] = self.function_to_optimize_source_code
                 new_code_combined = new_helper_code.copy()
                 new_code_combined[explanation.file_path] = new_code
-                if not self.args.no_pr:
-                    coverage_message = (
-                        original_code_baseline.coverage_results.build_message()
-                        if original_code_baseline.coverage_results
-                        else "Coverage data not available"
-                    )
-                    generated_tests = remove_functions_from_generated_tests(
-                        generated_tests=generated_tests, test_functions_to_remove=test_functions_to_remove
-                    )
-                    original_runtime_by_test = (
-                        original_code_baseline.benchmarking_test_results.usable_runtime_data_by_test_case()
-                    )
-                    optimized_runtime_by_test = (
-                        best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case()
-                    )
-                    qualified_name = self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)
-                    # Add runtime comments to generated tests before creating the PR
-                    generated_tests = add_runtime_comments_to_generated_tests(
-                        generated_tests, original_runtime_by_test, optimized_runtime_by_test
-                    )
-                    generated_tests_str = "\n\n".join(
-                        [test.generated_original_test_source for test in generated_tests.generated_tests]
-                    )
-                    existing_tests = existing_tests_source_for(
-                        qualified_name,
-                        function_to_all_tests,
-                        test_cfg=self.test_cfg,
-                        original_runtimes_all=original_runtime_by_test,
-                        optimized_runtimes_all=optimized_runtime_by_test,
-                    )
-                    if concolic_test_str:
-                        generated_tests_str += "\n\n" + concolic_test_str
-
-                    check_create_pr(
-                        original_code=original_code_combined,
-                        new_code=new_code_combined,
-                        explanation=explanation,
-                        existing_tests_source=existing_tests,
-                        generated_original_test_source=generated_tests_str,
-                        function_trace_id=self.function_trace_id[:-4] + exp_type
-                        if self.experiment_id
-                        else self.function_trace_id,
-                        coverage_message=coverage_message,
-                        git_remote=self.args.git_remote,
-                    )
-                    if (
-                        self.args.all
-                        or env_utils.get_pr_number()
-                        or self.args.replay_test
-                        or (self.args.file and not self.args.function)
-                    ):
-                        self.write_code_and_helpers(
-                            self.function_to_optimize_source_code,
-                            original_helper_code,
-                            self.function_to_optimize.file_path,
-                        )
-                else:
-                    # Mark optimization success since no PR will be created
-                    mark_optimization_success(
-                        trace_id=self.function_trace_id, is_optimization_found=best_optimization is not None
-                    )
+                self.process_review(
+                    original_code_baseline,
+                    best_optimization,
+                    generated_tests,
+                    test_functions_to_remove,
+                    concolic_test_str,
+                    original_code_combined,
+                    new_code_combined,
+                    explanation,
+                    function_to_all_tests,
+                    exp_type,
+                    original_helper_code,
+                    code_context,
+                )
                 self.log_successful_optimization(explanation, generated_tests, exp_type)
         return best_optimization
 
+    def process_review(
+        self,
+        original_code_baseline: OriginalCodeBaseline,
+        best_optimization: BestOptimization,
+        generated_tests: GeneratedTestsList,
+        test_functions_to_remove: list[str],
+        concolic_test_str: str | None,
+        original_code_combined: dict[Path, str],
+        new_code_combined: dict[Path, str],
+        explanation: Explanation,
+        function_to_all_tests: dict[str, set[FunctionCalledInTest]],
+        exp_type: str,
+        original_helper_code: dict[Path, str],
+        code_context: CodeOptimizationContext,
+    ) -> None:
+        coverage_message = (
+            original_code_baseline.coverage_results.build_message()
+            if original_code_baseline.coverage_results
+            else "Coverage data not available"
+        )
+
+        generated_tests = remove_functions_from_generated_tests(
+            generated_tests=generated_tests, test_functions_to_remove=test_functions_to_remove
+        )
+
+        original_runtime_by_test = original_code_baseline.benchmarking_test_results.usable_runtime_data_by_test_case()
+        optimized_runtime_by_test = (
+            best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case()
+        )
+
+        generated_tests = add_runtime_comments_to_generated_tests(
+            generated_tests, original_runtime_by_test, optimized_runtime_by_test
+        )
+
+        generated_tests_str = "\n#------------------------------------------------\n".join(
+            [test.generated_original_test_source for test in generated_tests.generated_tests]
+        )
+        if concolic_test_str:
+            generated_tests_str += "\n#------------------------------------------------\n" + concolic_test_str
+
+        existing_tests = existing_tests_source_for(
+            self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root),
+            function_to_all_tests,
+            test_cfg=self.test_cfg,
+            original_runtimes_all=original_runtime_by_test,
+            optimized_runtimes_all=optimized_runtime_by_test,
+        )
+        new_explanation_raw_str = self.aiservice_client.get_new_explanation(
+            source_code=code_context.read_writable_code,
+            dependency_code=code_context.read_only_context_code,
+            trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
+            optimized_code=best_optimization.candidate.source_code,
+            original_line_profiler_results=original_code_baseline.line_profile_results["str_out"],
+            optimized_line_profiler_results=best_optimization.line_profiler_test_results["str_out"],
+            original_code_runtime=humanize_runtime(original_code_baseline.runtime),
+            optimized_code_runtime=humanize_runtime(best_optimization.runtime),
+            speedup=f"{int(performance_gain(original_runtime_ns=original_code_baseline.runtime, optimized_runtime_ns=best_optimization.runtime) * 100)}%",
+            annotated_tests=generated_tests_str,
+            optimization_id=best_optimization.candidate.optimization_id,
+            original_explanation=best_optimization.candidate.explanation,
+        )
+        new_explanation = Explanation(
+            raw_explanation_message=new_explanation_raw_str or 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,
+        )
+        data = {
+            "original_code": original_code_combined,
+            "new_code": new_code_combined,
+            "explanation": new_explanation,
+            "existing_tests_source": existing_tests,
+            "generated_original_test_source": generated_tests_str,
+            "function_trace_id": self.function_trace_id[:-4] + exp_type
+            if self.experiment_id
+            else self.function_trace_id,
+            "coverage_message": coverage_message,
+        }
+
+        if not self.args.no_pr and not self.args.staging_review:
+            data["git_remote"] = self.args.git_remote
+            check_create_pr(**data)
+        elif self.args.staging_review:
+            create_staging(**data)
+        else:
+            # Mark optimization success since no PR will be created
+            mark_optimization_success(
+                trace_id=self.function_trace_id, is_optimization_found=best_optimization is not None
+            )
+
+        if ((not self.args.no_pr) or not self.args.staging_review) and (
+            self.args.all or env_utils.get_pr_number() or (self.args.file and not self.args.function)
+        ):
+            self.write_code_and_helpers(
+                self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
+            )
+
     def establish_original_code_baseline(
         self,
         code_context: CodeOptimizationContext,

codeflash/optimization/function_optimizer.py

@@ -1,2 +1,2 @@
 # These version placeholders will be replaced by uv-dynamic-versioning during build.
-__version__ = "0.16.0"
+__version__ = "0.16.1"