From 86bb1b53debd0eb0e08441d1e40cfb4b7dcac2d3 Mon Sep 17 00:00:00 2001 From: HeshamHM28 Date: Wed, 2 Jul 2025 08:56:05 +0300 Subject: [PATCH 1/6] create new staging argument --- codeflash/api/cfapi.py | 54 +++++++- codeflash/cli_cmds/cli.py | 5 + codeflash/optimization/function_optimizer.py | 124 ++++++++++--------- 3 files changed, 125 insertions(+), 58 deletions(-) diff --git a/codeflash/api/cfapi.py b/codeflash/api/cfapi.py index 518aded68..aafa76e86 100644 --- a/codeflash/api/cfapi.py +++ b/codeflash/api/cfapi.py @@ -14,13 +14,14 @@ 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_repo_owner_and_name, git_root_dir, get_current_branch +from codeflash.result.explanation import Explanation from codeflash.version import __version__ +from codeflash.github.PrComment import FileDiffContent, PrComment if TYPE_CHECKING: from requests import Response - from codeflash.github.PrComment import FileDiffContent, PrComment from packaging import version if os.environ.get("CODEFLASH_CFAPI_SERVER", default="prod").lower() == "local": @@ -174,6 +175,55 @@ def create_pr( } return make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload) +def create_staging( + original_code:str, + new_code: 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 owner: The owner of the repository. + :param repo: The name of the repository. + :param base_branch: The base branch to target. + :param file_changes: A dictionary of file changes. + :param pr_comment: The pull request comment object, containing the optimization explanation, best runtime, etc. + :param generated_tests: The generated tests. + :return: The response object. + """ + # convert Path objects to strings + 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_behavioral_test_results=explanation.winning_behavioral_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) -> bool: """Check if the Codeflash GitHub App is installed on the specified repository. diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index c6aaebfe3..b5af3f386 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -47,6 +47,11 @@ 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", diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index acae4b9fb..cb6bff677 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -19,7 +19,7 @@ from rich.tree import Tree from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient -from codeflash.api.cfapi import add_code_context_hash, mark_optimization_success +from codeflash.api.cfapi import add_code_context_hash, mark_optimization_success, create_staging 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 @@ -997,64 +997,76 @@ 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() - ) - # Add runtime comments to generated tests before creating the PR - generated_tests = add_runtime_comments_to_generated_tests( - self.test_cfg, 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( - 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, - ) - 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.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) self.log_successful_optimization(explanation, generated_tests, exp_type) return best_optimization + def process_review( + self, 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 + ): + 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( + self.test_cfg, 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, + ) + + data = dict( + 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, + ) + + 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: + if 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, From f5f10407888802492a5f23978305e5c0f53efe6a Mon Sep 17 00:00:00 2001 From: HeshamHM28 Date: Wed, 2 Jul 2025 09:15:39 +0300 Subject: [PATCH 2/6] fix formatting issues --- codeflash/api/cfapi.py | 37 +++++----- codeflash/cli_cmds/cli.py | 6 +- codeflash/optimization/function_optimizer.py | 75 +++++++++++++------- 3 files changed, 70 insertions(+), 48 deletions(-) diff --git a/codeflash/api/cfapi.py b/codeflash/api/cfapi.py index aafa76e86..979d11a82 100644 --- a/codeflash/api/cfapi.py +++ b/codeflash/api/cfapi.py @@ -14,14 +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, git_root_dir, get_current_branch -from codeflash.result.explanation import Explanation -from codeflash.version import __version__ +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.result.explanation import Explanation + from packaging import version if os.environ.get("CODEFLASH_CFAPI_SERVER", default="prod").lower() == "local": @@ -175,8 +176,9 @@ def create_pr( } return make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload) + def create_staging( - original_code:str, + original_code: str, new_code: str, explanation: Explanation, existing_tests_source: str, @@ -204,20 +206,20 @@ def create_staging( for p in original_code } payload = { - "baseBranch": get_current_branch(), + "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_behavioral_test_results=explanation.winning_behavioral_test_results, - winning_benchmarking_test_results=explanation.winning_benchmarking_test_results, - benchmark_details=explanation.benchmark_details, - ).to_json(), + "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_behavioral_test_results=explanation.winning_behavioral_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, @@ -225,6 +227,7 @@ def create_staging( } return make_cfapi_request(endpoint="/create-staging", method="POST", payload=payload) + def is_github_app_installed_on_repo(owner: str, repo: str) -> bool: """Check if the Codeflash GitHub App is installed on the specified repository. diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index b5af3f386..8fdadafc4 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -47,11 +47,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("--staging-review", action="store_true", help="Upload optimizations to staging for review") parser.add_argument( "--verify-setup", action="store_true", diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index cb6bff677..30b8089a3 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -19,7 +19,7 @@ from rich.tree import Tree from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient -from codeflash.api.cfapi import add_code_context_hash, mark_optimization_success, create_staging +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 @@ -997,20 +997,40 @@ 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 - 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) + 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, + ) self.log_successful_optimization(explanation, generated_tests, exp_type) return best_optimization def process_review( - self, 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 - ): + 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], + ) -> None: coverage_message = ( original_code_baseline.coverage_results.build_message() - if original_code_baseline.coverage_results else "Coverage data not available" + if original_code_baseline.coverage_results + else "Coverage data not available" ) generated_tests = remove_functions_from_generated_tests( @@ -1018,7 +1038,9 @@ def process_review( ) 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() + optimized_runtime_by_test = ( + best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case() + ) generated_tests = add_runtime_comments_to_generated_tests( self.test_cfg, generated_tests, original_runtime_by_test, optimized_runtime_by_test @@ -1038,15 +1060,17 @@ def process_review( optimized_runtimes_all=optimized_runtime_by_test, ) - data = dict( - 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, - ) + data = { + "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, + } if not self.args.no_pr and not self.args.staging_review: data["git_remote"] = self.args.git_remote @@ -1059,13 +1083,12 @@ def process_review( 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: - if 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, - ) + 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, From 98ec87ea662e4f00e3cb3fd7d833071f20621bb7 Mon Sep 17 00:00:00 2001 From: HeshamHM28 Date: Tue, 22 Jul 2025 14:29:29 +0300 Subject: [PATCH 3/6] Resolve the conflict --- codeflash/api/cfapi.py | 25 ++-- codeflash/optimization/function_optimizer.py | 145 ++++++++++--------- 2 files changed, 93 insertions(+), 77 deletions(-) diff --git a/codeflash/api/cfapi.py b/codeflash/api/cfapi.py index 8eeb28145..f7b42b6a5 100644 --- a/codeflash/api/cfapi.py +++ b/codeflash/api/cfapi.py @@ -183,9 +183,10 @@ def create_pr( } return make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload) + def create_staging( - original_code: str, - new_code: str, + original_code: dict[Path, str], + new_code: dict[Path, str], explanation: Explanation, existing_tests_source: str, generated_original_test_source: str, @@ -194,15 +195,15 @@ def create_staging( ) -> Response: """Create a staging pull request, targeting the specified branch. (usually 'staging'). - :param owner: The owner of the repository. - :param repo: The name of the repository. - :param base_branch: The base branch to target. - :param file_changes: A dictionary of file changes. - :param pr_comment: The pull request comment object, containing the optimization explanation, best runtime, etc. - :param generated_tests: The generated tests. - :return: The response object. + :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. """ - # convert Path objects to strings relative_path = explanation.file_path.relative_to(git_root_dir()).as_posix() build_file_changes = { @@ -211,6 +212,7 @@ def create_staging( ) for p in original_code } + payload = { "baseBranch": get_current_branch(), "diffContents": build_file_changes, @@ -222,7 +224,7 @@ def create_staging( relative_file_path=relative_path, speedup_x=explanation.speedup_x, speedup_pct=explanation.speedup_pct, - winning_behavioral_test_results=explanation.winning_behavioral_test_results, + 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(), @@ -231,6 +233,7 @@ def create_staging( "traceId": function_trace_id, "coverage_message": coverage_message, } + return make_cfapi_request(endpoint="/create-staging", method="POST", payload=payload) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 94acc93b6..acf87227b 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -998,7 +998,6 @@ 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 - <<<<<<< feat-staging self.process_review( original_code_baseline, best_optimization, @@ -1012,71 +1011,85 @@ def find_and_process_best_optimization( exp_type, original_helper_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 - ) - >>>>>>> main - self.log_successful_optimization(explanation, generated_tests, exp_type) - return best_optimization + 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], + ) -> 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( + self.test_cfg, 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, + ) + + data = { + "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, + } + + 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, From d70525ca82bc5b55dffca8cca45ab847d187bf56 Mon Sep 17 00:00:00 2001 From: HeshamHM28 Date: Tue, 22 Jul 2025 21:07:59 +0300 Subject: [PATCH 4/6] fix typo --- codeflash/optimization/function_optimizer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index acf87227b..de2a8fd73 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -1234,7 +1234,7 @@ def establish_original_code_baseline( return Success( ( OriginalCodeBaseline( - behavioral_test_results=behavioral_results, + behavior_test_results=behavioral_results, benchmarking_test_results=benchmarking_results, replay_benchmarking_test_results=replay_benchmarking_test_results if self.args.benchmark @@ -1298,7 +1298,7 @@ def run_optimized_candidate( ) ) console.rule() - if compare_test_results(baseline_results.behavioral_test_results, candidate_behavior_results): + if compare_test_results(baseline_results.behavior_test_results, candidate_behavior_results): logger.info("Test results matched!") console.rule() else: From 7170b5c3c235c000d5c6bdf4fdce1c7a3a3f5477 Mon Sep 17 00:00:00 2001 From: HeshamHM28 Date: Tue, 22 Jul 2025 21:17:50 +0300 Subject: [PATCH 5/6] Fix the argument mismatch --- codeflash/optimization/function_optimizer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index de2a8fd73..86f075fa4 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -1043,9 +1043,7 @@ def process_review( best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case() ) - generated_tests = add_runtime_comments_to_generated_tests( - self.test_cfg, generated_tests, original_runtime_by_test, optimized_runtime_by_test - ) + 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] From 6c48296542c8b5d8cd2dd2e3a28b45f1f4bfa286 Mon Sep 17 00:00:00 2001 From: HeshamHM28 Date: Wed, 30 Jul 2025 16:52:57 +0300 Subject: [PATCH 6/6] fix generated tests comment --- codeflash/optimization/function_optimizer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 334017ae4..2634a61ae 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -17,6 +17,7 @@ from rich.panel import Panel from rich.syntax import Syntax from rich.tree import Tree + from codeflash.api.aiservice import AiServiceClient, AIServiceRefinerRequest, LocalAiServiceClient from codeflash.api.cfapi import add_code_context_hash, create_staging, mark_optimization_success from codeflash.benchmarking.utils import process_benchmark_data @@ -1132,7 +1133,9 @@ def process_review( best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case() ) - add_runtime_comments_to_generated_tests(generated_tests, original_runtime_by_test, optimized_runtime_by_test) + 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]