Merge branch 'main' into replay-test-save-dir

Author: aseembits93

codeflash/api/aiservice.py

@@ -11,6 +11,7 @@
 
 from codeflash.cli_cmds.console import console, logger
 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.models.models import OptimizedCandidate
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.version import __version__ as codeflash_version
@@ -97,6 +98,12 @@ def optimize_python_code(  # noqa: D417
 
         """
         start_time = time.perf_counter()
+        try:
+            git_repo_owner, git_repo_name = get_repo_owner_and_name()
+        except Exception as e:
+            logger.warning(f"Could not determine repo owner and name: {e}")
+            git_repo_owner, git_repo_name = None, None
+
         payload = {
             "source_code": source_code,
             "dependency_code": dependency_code,
@@ -105,6 +112,9 @@ def optimize_python_code(  # noqa: D417
             "python_version": platform.python_version(),
             "experiment_metadata": experiment_metadata,
             "codeflash_version": codeflash_version,
+            "current_username": get_last_commit_author_if_pr_exists(None),
+            "repo_owner": git_repo_owner,
+            "repo_name": git_repo_name,
         }
 
         logger.info("Generating optimized candidates…")

codeflash/api/cfapi.py

@@ -224,3 +224,14 @@ def add_code_context_hash(code_context_hash: str) -> None:
             "POST",
             {"owner": owner, "repo": repo, "pr_number": pr_number, "code_hash": code_context_hash},
         )
+
+
+def mark_optimization_success(trace_id: str, *, is_optimization_found: bool) -> Response:
+    """Mark an optimization event as success or not.
+
+    :param trace_id: The unique identifier for the optimization event.
+    :param is_optimization_found: Boolean indicating whether the optimization was found.
+    :return: The response object from the API.
+    """
+    payload = {"trace_id": trace_id, "is_optimization_found": is_optimization_found}
+    return make_cfapi_request(endpoint="/mark-as-success", method="POST", payload=payload)

codeflash/cli_cmds/cli.py

@@ -8,6 +8,7 @@
 from codeflash.cli_cmds.cmd_init import init_codeflash, install_github_actions
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils import env_utils
+from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.version import __version__ as version
 
@@ -42,7 +43,7 @@ def parse_args() -> Namespace:
     )
     parser.add_argument("--test-framework", choices=["pytest", "unittest"], default="pytest")
     parser.add_argument("--config-file", type=str, help="Path to the pyproject.toml with codeflash configs.")
-    parser.add_argument("--replay-test", type=str, help="Path to replay test to optimize functions from")
+    parser.add_argument("--replay-test", type=str, nargs="+", help="Paths to replay test to optimize functions from")
     parser.add_argument(
         "--no-pr", action="store_true", help="Do not create a PR for the optimization, only update the code locally."
     )
@@ -83,25 +84,22 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
         sys.exit()
     if not check_running_in_git_repo(module_root=args.module_root):
         if not confirm_proceeding_with_no_git_repo():
-            logger.critical("No git repository detected and user aborted run. Exiting...")
-            sys.exit(1)
+            exit_with_message("No git repository detected and user aborted run. Exiting...", error_on_exit=True)
         args.no_pr = True
     if args.function and not args.file:
-        logger.error("If you specify a --function, you must specify the --file it is in")
-        sys.exit(1)
+        exit_with_message("If you specify a --function, you must specify the --file it is in", error_on_exit=True)
     if args.file:
         if not Path(args.file).exists():
-            logger.error(f"File {args.file} does not exist")
-            sys.exit(1)
+            exit_with_message(f"File {args.file} does not exist", error_on_exit=True)
         args.file = Path(args.file).resolve()
         if not args.no_pr:
             owner, repo = get_repo_owner_and_name()
             require_github_app_or_exit(owner, repo)
     if args.replay_test:
-        if not Path(args.replay_test).is_file():
-            logger.error(f"Replay test file {args.replay_test} does not exist")
-            sys.exit(1)
-        args.replay_test = Path(args.replay_test).resolve()
+        for test_path in args.replay_test:
+            if not Path(test_path).is_file():
+                exit_with_message(f"Replay test file {test_path} does not exist", error_on_exit=True)
+        args.replay_test = [Path(replay_test).resolve() for replay_test in args.replay_test]
 
     return args
 
@@ -110,8 +108,7 @@ def process_pyproject_config(args: Namespace) -> Namespace:
     try:
         pyproject_config, pyproject_file_path = parse_config_file(args.config_file)
     except ValueError as e:
-        logger.error(e)
-        sys.exit(1)
+        exit_with_message(f"Error parsing config file: {e}", error_on_exit=True)
     supported_keys = [
         "module_root",
         "tests_root",
@@ -203,8 +200,7 @@ def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
             )
             apologize_and_exit()
         if not args.no_pr and not check_and_push_branch(git_repo):
-            logger.critical("❌ Branch is not pushed. Exiting...")
-            sys.exit(1)
+            exit_with_message("Branch is not pushed...", error_on_exit=True)
         owner, repo = get_repo_owner_and_name(git_repo)
         if not args.no_pr:
             require_github_app_or_exit(owner, repo)

codeflash/code_utils/code_utils.py

@@ -5,14 +5,15 @@
 import re
 import shutil
 import site
+import sys
 from contextlib import contextmanager
 from functools import lru_cache
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
 import tomlkit
 
-from codeflash.cli_cmds.console import logger
+from codeflash.cli_cmds.console import logger, paneled_text
 from codeflash.code_utils.config_parser import find_pyproject_toml
 
 ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
@@ -33,7 +34,7 @@ def custom_addopts() -> None:
             # Backup original addopts
             original_addopts = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "")
             # nothing to do if no addopts present
-            if original_addopts != "":
+            if original_addopts != "" and isinstance(original_addopts, list):
                 original_addopts = [x.strip() for x in original_addopts]
                 non_blacklist_plugin_args = re.sub(r"-n(?: +|=)\S+", "", " ".join(original_addopts)).split(" ")
                 non_blacklist_plugin_args = [x for x in non_blacklist_plugin_args if x != ""]
@@ -213,3 +214,9 @@ def cleanup_paths(paths: list[Path]) -> None:
 def restore_conftest(path_to_content_map: dict[Path, str]) -> None:
     for path, file_content in path_to_content_map.items():
         path.write_text(file_content, encoding="utf8")
+
+
+def exit_with_message(message: str, *, error_on_exit: bool = False) -> None:
+    paneled_text(message, panel_args={"style": "red"})
+
+    sys.exit(1 if error_on_exit else 0)

codeflash/code_utils/env_utils.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import os
-import sys
 import tempfile
 from functools import lru_cache
 from pathlib import Path
 from typing import Optional
 
 from codeflash.cli_cmds.console import logger
+from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.formatter import format_code
 from codeflash.code_utils.shell_utils import read_api_key_from_shell_config
 
@@ -24,11 +24,11 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool =
         try:
             format_code(formatter_cmds, tmp_file, print_status=False)
         except Exception:
-            print(
-                "⚠️ Codeflash requires a code formatter to be installed in your environment, but none was found. Please install a supported formatter, verify the formatter-cmds in your codeflash pyproject.toml config and try again."
+            exit_with_message(
+                "⚠️ Codeflash requires a code formatter to be installed in your environment, but none was found. Please install a supported formatter, verify the formatter-cmds in your codeflash pyproject.toml config and try again.",
+                error_on_exit=True,
             )
-            if exit_on_failure:
-                sys.exit(1)
+
     return return_code
 
 

codeflash/code_utils/git_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import shutil
 import subprocess
 import sys
@@ -176,3 +177,20 @@ def remove_git_worktrees(worktree_root: Path | None, worktrees: list[Path]) -> N
         logger.warning(f"Error removing worktrees: {e}")
     if worktree_root:
         shutil.rmtree(worktree_root)
+
+
+def get_last_commit_author_if_pr_exists(repo: Repo | None = None) -> str | None:
+    """Return the author's name of the last commit in the current branch if PR_NUMBER is set.
+
+    Otherwise, return None.
+    """
+    if "PR_NUMBER" not in os.environ:
+        return None
+    try:
+        repository: Repo = repo if repo else git.Repo(search_parent_directories=True)
+        last_commit = repository.head.commit
+    except Exception:
+        logger.exception("Failed to get last commit author.")
+        return None
+    else:
+        return last_commit.author.name