Merge branch 'main' into misrasaurabh1-patch-1

Author: misrasaurabh1

.github/workflows/codeflash-optimize.yaml

@@ -1,7 +1,7 @@
 name: CodeFlash
 
 on:
-  pull_request_target:
+  pull_request:
     paths:
       - '**'  # Trigger for all paths
 
@@ -14,8 +14,6 @@ concurrency:
 jobs:
   optimize:
     name: Optimize new Python code
-    # Dynamically determine if environment is needed only when workflow files change and contributor is external
-    environment: ${{ (github.event_name == 'workflow_dispatch' || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
     if: ${{ github.actor != 'codeflash-ai[bot]' }}
     runs-on: ubuntu-latest
     env:
@@ -28,31 +26,7 @@ jobs:
       - name: 🛎️ Checkout
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event.pull_request.head.ref }}
-          repository: ${{ github.event.pull_request.head.repo.full_name }}
           fetch-depth: 0
-      - name: Validate PR
-        run: |
-          # Check for any workflow changes
-          if git diff --name-only "${{ github.event.pull_request.base.sha }}" "${{ github.event.pull_request.head.sha }}" | grep -q "^.github/workflows/"; then
-            echo "⚠️ Workflow changes detected."
-
-            # Get the PR author
-            AUTHOR="${{ github.event.pull_request.user.login }}"
-            echo "PR Author: $AUTHOR"
-
-            # Allowlist check
-            if [[ "$AUTHOR" == "misrasaurabh1" || "$AUTHOR" == "KRRT7" ]]; then
-              echo "✅ Authorized user ($AUTHOR). Proceeding."
-            elif [[ "${{ github.event.pull_request.state }}" == "open" ]]; then
-              echo "✅ PR is open. Proceeding with appropriate protections."
-            else
-              echo "⛔ Unauthorized user ($AUTHOR) attempting to modify workflows. Exiting."
-              exit 1
-            fi
-          else
-            echo "✅ No workflow file changes detected. Proceeding."
-          fi
 
       - name: 🐍 Set up Python 3.11 for CLI
         uses: astral-sh/setup-uv@v5

.github/workflows/end-to-end-test-init-optim.yaml

@@ -20,7 +20,7 @@ jobs:
       COLUMNS: 110
       MAX_RETRIES: 3
       RETRY_DELAY: 5
-      EXPECTED_IMPROVEMENT_PCT: 300
+      EXPECTED_IMPROVEMENT_PCT: 30
       CODEFLASH_END_TO_END: 1
     steps:
       - name: 🛎️ Checkout

code_to_optimize/code_directories/simple_tracer_e2e/workload.py

@@ -2,6 +2,7 @@
 
 
 def funcA(number):
+    number = number if number < 1000 else 1000
     k = 0
     for i in range(number * 100):
         k += i

codeflash/api/aiservice.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import time
+
 import json
 import os
 import platform
@@ -95,6 +97,7 @@ def optimize_python_code(
         - List[OptimizationCandidate]: A list of Optimization Candidates.
 
         """
+        start_time = time.perf_counter()
         payload = {
             "source_code": source_code,
             "dependency_code": dependency_code,
@@ -118,6 +121,8 @@ def optimize_python_code(
             optimizations_json = response.json()["optimizations"]
             logger.info(f"Generated {len(optimizations_json)} candidates.")
             console.rule()
+            end_time = time.perf_counter()
+            logger.debug(f"Optimization took {end_time - start_time:.2f} seconds.")
             return [
                 OptimizedCandidate(
                     source_code=opt["source_code"],

codeflash/cli_cmds/cmd_init.py

@@ -676,13 +676,16 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
         )
         apologize_and_exit()
 
+    enable_telemetry = ask_for_telemetry()
+
     codeflash_section = tomlkit.table()
     codeflash_section.add(tomlkit.comment("All paths are relative to this pyproject.toml's directory."))
     codeflash_section["module-root"] = setup_info.module_root
     codeflash_section["tests-root"] = setup_info.tests_root
     codeflash_section["test-framework"] = setup_info.test_framework
     codeflash_section["benchmarks-root"] = setup_info.benchmarks_root if setup_info.benchmarks_root else ""
     codeflash_section["ignore-paths"] = setup_info.ignore_paths
+    codeflash_section["disable-telemetry"] = not enable_telemetry
     if setup_info.git_remote not in ["", "origin"]:
         codeflash_section["git-remote"] = setup_info.git_remote
     formatter = setup_info.formatter
@@ -934,3 +937,16 @@ def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test
             console.rule()
             Path(path).unlink(missing_ok=True)
             logger.info(f"🗑️  Deleted {path}")
+
+
+def ask_for_telemetry() -> bool:
+    """Prompt the user to enable or disable telemetry."""
+    from rich.prompt import Confirm
+
+    enable_telemetry = Confirm.ask(
+        "⚡️ Would you like to enable telemetry to help us improve the Codeflash experience?",
+        default=True,
+        show_default=True,
+    )
+
+    return enable_telemetry

codeflash/code_utils/code_utils.py

@@ -2,6 +2,7 @@
 
 import ast
 import os
+import shutil
 import site
 from functools import lru_cache
 from pathlib import Path
@@ -118,4 +119,8 @@ def has_any_async_functions(code: str) -> bool:
 
 def cleanup_paths(paths: list[Path]) -> None:
     for path in paths:
-        path.unlink(missing_ok=True)
+        if path and path.exists():
+            if path.is_dir():
+                shutil.rmtree(path, ignore_errors=True)
+            else:
+                path.unlink(missing_ok=True)

codeflash/code_utils/compat.py

@@ -2,6 +2,8 @@
 import sys
 from pathlib import Path
 
+from platformdirs import user_config_dir
+
 # os-independent newline
 # important for any user-facing output or files we write
 # make sure to use this in f-strings e.g. f"some string{LF}"
@@ -12,3 +14,8 @@
 SAFE_SYS_EXECUTABLE: str = Path(sys.executable).as_posix()
 
 IS_POSIX = os.name != "nt"
+
+
+codeflash_cache_dir = Path(user_config_dir(appname="codeflash", appauthor="codeflash-ai", ensure_exists=True))
+
+codeflash_cache_db = codeflash_cache_dir / "codeflash_cache.db"

codeflash/code_utils/github_utils.py

@@ -26,3 +26,6 @@ def require_github_app_or_exit(owner: str, repo: str) -> None:
             f"Note: if you want to find optimizations without opening PRs, you can run Codeflash with the --no-pr flag.{LF}"
         )
         apologize_and_exit()
+
+def github_pr_url(owner: str, repo: str, pr_number: str) -> str:
+    return f"https://github.com/{owner}/{repo}/pull/{pr_number}"

codeflash/context/code_context_extractor.py

@@ -217,7 +217,9 @@ def extract_code_string_context_from_files(
             continue
         try:
             qualified_helper_function_names = {func.qualified_name for func in helper_function_sources}
-            code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, qualified_helper_function_names)
+            code_without_unused_defs = remove_unused_definitions_by_function_names(
+                original_code, qualified_helper_function_names
+            )
             code_context = parse_code_and_prune_cst(
                 code_without_unused_defs, code_context_type, set(), qualified_helper_function_names, remove_docstrings
             )
@@ -325,7 +327,9 @@ def extract_code_markdown_context_from_files(
             continue
         try:
             qualified_helper_function_names = {func.qualified_name for func in helper_function_sources}
-            code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, qualified_helper_function_names)
+            code_without_unused_defs = remove_unused_definitions_by_function_names(
+                original_code, qualified_helper_function_names
+            )
             code_context = parse_code_and_prune_cst(
                 code_without_unused_defs, code_context_type, set(), qualified_helper_function_names, remove_docstrings
             )
@@ -403,12 +407,8 @@ def get_function_sources_from_jedi(
             for name in names:
                 try:
                     definitions: list[Name] = name.goto(follow_imports=True, follow_builtin_imports=False)
-                except Exception as e:
-                    try:
-                        logger.exception(f"Error while getting definition for {name.full_name}: {e}")
-                    except Exception as e:
-                        # name.full_name can also throw exceptions sometimes
-                        logger.exception(f"Error while getting definition: {e}")
+                except Exception:  # noqa: BLE001
+                    logger.debug(f"Error while getting definitions for {qualified_function_name}")
                     definitions = []
                 if definitions:
                     # TODO: there can be multiple definitions, see how to handle such cases
@@ -424,7 +424,12 @@ def get_function_sources_from_jedi(
                         and not belongs_to_function_qualified(definition, qualified_function_name)
                         and definition.full_name.startswith(definition.module_name)
                         # Avoid nested functions or classes. Only class.function is allowed
-                        and len((qualified_name := get_qualified_name(definition.module_name, definition.full_name)).split(".")) <= 2
+                        and len(
+                            (qualified_name := get_qualified_name(definition.module_name, definition.full_name)).split(
+                                "."
+                            )
+                        )
+                        <= 2
                     ):
                         function_source = FunctionSource(
                             file_path=definition_path,