Merge branch 'main' of https://github.com/codeflash-ai/codeflash into uv-migrate

Author: KRRT7

.github/workflows/codeflash-optimize.yaml

@@ -12,7 +12,7 @@ concurrency:
 
 jobs:
   optimize:
-    name: Optimize new code in this PR
+    name: Optimize new Python code
     if: ${{ github.actor != 'codeflash-ai[bot]' }}
     runs-on: ubuntu-latest
     env:
@@ -22,16 +22,24 @@ jobs:
       CODEFLASH_PR_NUMBER: ${{ github.event.number }}
       COLUMNS: 110
     steps:
-      - uses: actions/checkout@v4
+      - name: 🛎️ Checkout
+        uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
-      - name: Set up Python 3.11 for CLI
+      - name: 🐍 Set up Python 3.11 for CLI
         uses: astral-sh/setup-uv@v5
         with:
           python-version: 3.11.6
 
-      - name: Run Codeflash to optimize code
+      - name: 📦 Install dependencies (CLI)
+        run: |
+          uv tool install poetry
+          uv venv
+          source .venv/bin/activate
+          poetry install --with dev
+
+      - name: ⚡️Codeflash Optimization
         id: optimize_code
         run: |
           uv run codeflash

code_to_optimize/bubble_sort_classmethod.py

@@ -0,0 +1,6 @@
+from code_to_optimize.bubble_sort_in_class import BubbleSortClass
+
+
+def sort_classmethod(x):
+    y = BubbleSortClass()
+    return y.sorter(x)

code_to_optimize/bubble_sort_nested_classmethod.py

@@ -0,0 +1,6 @@
+from code_to_optimize.bubble_sort_in_nested_class import WrapperClass
+
+
+def sort_classmethod(x):
+    y = WrapperClass.BubbleSortClass()
+    return y.sorter(x)

code_to_optimize/code_directories/simple_tracer_e2e/workload.py

@@ -1,4 +1,6 @@
 from concurrent.futures import ThreadPoolExecutor
+
+
 def funcA(number):
     k = 0
     for i in range(number * 100):
@@ -9,6 +11,7 @@ def funcA(number):
     # Use a generator expression directly in join for more efficiency
     return " ".join(str(i) for i in range(number))
 
+
 def test_threadpool() -> None:
     pool = ThreadPoolExecutor(max_workers=3)
     args = list(range(10, 31, 10))
@@ -19,4 +22,4 @@ def test_threadpool() -> None:
 
 
 if __name__ == "__main__":
-    test_threadpool()
+    test_threadpool()

codeflash/api/aiservice.py

@@ -135,6 +135,76 @@ def optimize_python_code(
         console.rule()
         return []
 
+    def optimize_python_code_line_profiler(
+        self,
+        source_code: str,
+        dependency_code: str,
+        trace_id: str,
+        line_profiler_results: str,
+        num_candidates: int = 10,
+        experiment_metadata: ExperimentMetadata | None = None,
+    ) -> list[OptimizedCandidate]:
+        """Optimize the given python code for performance by making a request to the Django endpoint.
+
+        Parameters
+        ----------
+        - source_code (str): The python code to optimize.
+        - dependency_code (str): The dependency code used as read-only context for the optimization
+        - trace_id (str): Trace id of optimization run
+        - num_candidates (int): Number of optimization variants to generate. Default is 10.
+        - experiment_metadata (Optional[ExperimentalMetadata, None]): Any available experiment metadata for this optimization
+
+        Returns
+        -------
+        - List[OptimizationCandidate]: A list of Optimization Candidates.
+
+        """
+        payload = {
+            "source_code": source_code,
+            "dependency_code": dependency_code,
+            "num_variants": num_candidates,
+            "line_profiler_results": line_profiler_results,
+            "trace_id": trace_id,
+            "python_version": platform.python_version(),
+            "experiment_metadata": experiment_metadata,
+            "codeflash_version": codeflash_version,
+        }
+
+        logger.info("Generating optimized candidates…")
+        console.rule()
+        if line_profiler_results=="":
+            logger.info("No LineProfiler results were provided, Skipping optimization.")
+            console.rule()
+            return []
+        try:
+            response = self.make_ai_service_request("/optimize-line-profiler", payload=payload, timeout=600)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error generating optimized candidates: {e}")
+            ph("cli-optimize-error-caught", {"error": str(e)})
+            return []
+
+        if response.status_code == 200:
+            optimizations_json = response.json()["optimizations"]
+            logger.info(f"Generated {len(optimizations_json)} candidates.")
+            console.rule()
+            return [
+                OptimizedCandidate(
+                    source_code=opt["source_code"],
+                    explanation=opt["explanation"],
+                    optimization_id=opt["optimization_id"],
+                )
+                for opt in optimizations_json
+            ]
+        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(
         self,
         function_trace_id: str,

codeflash/api/cfapi.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 import requests
+import sentry_sdk
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
@@ -194,7 +195,8 @@ def get_blocklisted_functions() -> dict[str, set[str]] | dict[str, Any]:
         req.raise_for_status()
         content: dict[str, list[str]] = req.json()
     except Exception as e:
-        logger.error(f"Error getting blocklisted functions: {e}", exc_info=True)
+        logger.error(f"Error getting blocklisted functions: {e}")
+        sentry_sdk.capture_exception(e)
         return {}
 
     return {Path(k).name: {v.replace("()", "") for v in values} for k, values in content.items()}

codeflash/cli_cmds/cli.py

@@ -128,14 +128,22 @@ def process_pyproject_config(args: Namespace) -> Namespace:
     assert args.tests_root is not None, "--tests-root must be specified"
     assert Path(args.tests_root).is_dir(), f"--tests-root {args.tests_root} must be a valid directory"
 
-    assert not (env_utils.get_pr_number() is not None and not env_utils.ensure_codeflash_api_key()), (
-        "Codeflash API key not found. When running in a Github Actions Context, provide the "
-        "'CODEFLASH_API_KEY' environment variable as a secret.\n"
-        "You can add a secret by going to your repository's settings page, then clicking 'Secrets' in the left sidebar.\n"
-        "Then, click 'New repository secret' and add your api key with the variable name CODEFLASH_API_KEY.\n"
-        f"Here's a direct link: {get_github_secrets_page_url()}\n"
-        "Exiting..."
-    )
+    if env_utils.get_pr_number() is not None:
+        assert env_utils.ensure_codeflash_api_key(), (
+            "Codeflash API key not found. When running in a Github Actions Context, provide the "
+            "'CODEFLASH_API_KEY' environment variable as a secret.\n"
+            "You can add a secret by going to your repository's settings page, then clicking 'Secrets' in the left sidebar.\n"
+            "Then, click 'New repository secret' and add your api key with the variable name CODEFLASH_API_KEY.\n"
+            f"Here's a direct link: {get_github_secrets_page_url()}\n"
+            "Exiting..."
+        )
+
+        repo = git.Repo(search_parent_directories=True)
+
+        owner, repo_name = get_repo_owner_and_name(repo)
+
+        require_github_app_or_exit(owner, repo_name)
+
     if hasattr(args, "ignore_paths") and args.ignore_paths is not None:
         normalized_ignore_paths = []
         for path in args.ignore_paths:

codeflash/cli_cmds/cmd_init.py

@@ -67,20 +67,27 @@ def init_codeflash() -> None:
 
         did_add_new_key = prompt_api_key()
 
-        setup_info: SetupInfo = collect_setup_info()
+        if should_modify_pyproject_toml():
 
-        configure_pyproject_toml(setup_info)
+            setup_info: SetupInfo = collect_setup_info()
+
+            configure_pyproject_toml(setup_info)
 
         install_github_app()
 
-        install_github_actions()
+        install_github_actions(override_formatter_check=True)
+
+        module_string = ""
+        if "setup_info" in locals():
+            module_string = f" you selected ({setup_info.module_root})"
+
 
         click.echo(
             f"{LF}"
             f"⚡️ Codeflash is now set up! You can now run:{LF}"
             f"    codeflash --file <path-to-file> --function <function-name> to optimize a function within a file{LF}"
             f"    codeflash --file <path-to-file> to optimize all functions in a file{LF}"
-            f"    codeflash --all to optimize all functions in all files in the module you selected ({setup_info.module_root}){LF}"
+            f"    codeflash --all to optimize all functions in all files in the module{module_string}{LF}"
             f"-or-{LF}"
             f"    codeflash --help to see all options{LF}"
         )
@@ -116,6 +123,30 @@ def ask_run_end_to_end_test(args: Namespace) -> None:
         bubble_sort_path, bubble_sort_test_path = create_bubble_sort_file_and_test(args)
         run_end_to_end_test(args, bubble_sort_path, bubble_sort_test_path)
 
+def should_modify_pyproject_toml() -> bool:
+    """
+    Check if the current directory contains a valid pyproject.toml file with codeflash config
+    If it does, ask the user if they want to re-configure it.
+    """
+    from rich.prompt import Confirm
+    pyproject_toml_path = Path.cwd() / "pyproject.toml"
+    if not pyproject_toml_path.exists():
+        return True
+    try:
+        config, config_file_path = parse_config_file(pyproject_toml_path)
+    except Exception as e:
+        return True
+
+    if "module_root" not in config or config["module_root"] is None or not Path(config["module_root"]).is_dir():
+        return True
+    if "tests_root" not in config or config["tests_root"] is None or not Path(config["tests_root"]).is_dir():
+        return True
+
+    create_toml = Confirm.ask(
+        f"✅ A valid Codeflash config already exists in this project. Do you want to re-configure it?", default=False, show_default=True
+    )
+    return create_toml
+
 
 def collect_setup_info() -> SetupInfo:
     curdir = Path.cwd()
@@ -362,9 +393,9 @@ def check_for_toml_or_setup_file() -> str | None:
     return cast(str, project_name)
 
 
-def install_github_actions() -> None:
+def install_github_actions(override_formatter_check: bool = False) -> None:
     try:
-        config, config_file_path = parse_config_file()
+        config, config_file_path = parse_config_file(override_formatter_check=override_formatter_check)
 
         ph("cli-github-actions-install-started")
         try:
@@ -504,11 +535,11 @@ def get_dependency_manager_installation_string(dep_manager: DependencyManager) -
     py_version = sys.version_info
     python_version_string = f"'{py_version.major}.{py_version.minor}'"
     if dep_manager == DependencyManager.UV:
-        return """name: Setup UV
+        return """name: 🐍 Setup UV
         uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true"""
-    return f"""name: Set up Python
+    return f"""name: 🐍 Set up Python
         uses: actions/setup-python@v5
         with:
           python-version: {python_version_string}"""

codeflash/cli_cmds/console.py

@@ -7,7 +7,16 @@
 
 from rich.console import Console
 from rich.logging import RichHandler
-from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
 
 from codeflash.cli_cmds.console_constants import SPINNER_TYPES
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
@@ -22,15 +31,26 @@
 console = Console()
 logging.basicConfig(
     level=logging.INFO,
-    handlers=[RichHandler(rich_tracebacks=True, markup=False, console=console, show_path=False, show_time=False)],
+    handlers=[
+        RichHandler(
+            rich_tracebacks=True,
+            markup=False,
+            console=console,
+            show_path=False,
+            show_time=False,
+        )
+    ],
     format=BARE_LOGGING_FORMAT,
 )
 
 logger = logging.getLogger("rich")
+logging.getLogger("parso").setLevel(logging.WARNING)
 
 
 def paneled_text(
-    text: str, panel_args: dict[str, str | bool] | None = None, text_args: dict[str, str] | None = None
+    text: str,
+    panel_args: dict[str, str | bool] | None = None,
+    text_args: dict[str, str] | None = None,
 ) -> None:
     """Print text in a panel."""
     from rich.panel import Panel
@@ -57,7 +77,9 @@ def code_print(code_str: str) -> None:
 
 
 @contextmanager
-def progress_bar(message: str, *, transient: bool = False) -> Generator[TaskID, None, None]:
+def progress_bar(
+    message: str, *, transient: bool = False
+) -> Generator[TaskID, None, None]:
     """Display a progress bar with a spinner and elapsed time."""
     progress = Progress(
         SpinnerColumn(next(spinners)),
@@ -69,3 +91,25 @@ def progress_bar(message: str, *, transient: bool = False) -> Generator[TaskID,
     task = progress.add_task(message, total=None)
     with progress:
         yield task
+
+
+@contextmanager
+def test_files_progress_bar(
+    total: int, description: str
+) -> Generator[tuple[Progress, TaskID], None, None]:
+    """Progress bar for test files."""
+    with Progress(
+        SpinnerColumn(next(spinners)),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(
+            complete_style="cyan",
+            finished_style="green",
+            pulse_style="yellow",
+        ),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TimeRemainingColumn(),
+        transient=True,
+    ) as progress:
+        task_id = progress.add_task(description, total=total)
+        yield progress, task_id

codeflash/cli_cmds/workflows/codeflash-optimize.yaml

@@ -15,7 +15,7 @@ concurrency:
 
 jobs:
   optimize:
-    name: Optimize new Python code in this PR
+    name: Optimize new Python code
     # Don't run codeflash on codeflash-ai[bot] commits, prevent duplicate optimizations
     if: ${{ github.actor != 'codeflash-ai[bot]' }}
     runs-on: ubuntu-latest
@@ -24,11 +24,12 @@ jobs:
       CODEFLASH_PR_NUMBER: ${{ github.event.number }}
     {{ working_directory }}
     steps:
-      - uses: actions/checkout@v4
+      - name: 🛎️ Checkout
+        uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - {{ setup_python_dependency_manager }}
-      - name: Install Project Dependencies
+      - name: 📦 Install Dependencies
         run: {{ install_dependencies_command }}
-      - name: Run Codeflash to optimize code
+      - name: ⚡️Codeflash Optimization
         run: {{ codeflash_command }}