Merge pull request #739 from codeflash-ai/get-throughput-from-output

Get throughput from output  for async functions

Author: KRRT7

.github/workflows/e2e-async.yaml

@@ -0,0 +1,69 @@
+name: E2E - Async
+
+on:
+  pull_request:
+    paths:
+      - '**'  # Trigger for all paths
+
+  workflow_dispatch:
+
+jobs:
+  async-optimization:
+    # 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' || '' }}
+
+    runs-on: ubuntu-latest
+    env:
+      CODEFLASH_AIS_SERVER: prod
+      POSTHOG_API_KEY: ${{ secrets.POSTHOG_API_KEY }}
+      CODEFLASH_API_KEY: ${{ secrets.CODEFLASH_API_KEY }}
+      COLUMNS: 110
+      MAX_RETRIES: 3
+      RETRY_DELAY: 5
+      EXPECTED_IMPROVEMENT_PCT: 10
+      CODEFLASH_END_TO_END: 1
+    steps:
+      - 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
+          token: ${{ secrets.GITHUB_TOKEN }}
+
+      - 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 triggered by 'pull_request_target' and is open. Assuming protection rules are in place. Proceeding."
+            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
+        with:
+          python-version: 3.11.6
+
+      - name: Install dependencies (CLI)
+        run: |
+          uv sync
+          
+      - name: Run Codeflash to optimize async code
+        id: optimize_async_code
+        run: |
+          uv run python tests/scripts/end_to_end_test_async.py

.github/workflows/pre-commit.yaml

@@ -1,8 +1,9 @@
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.12.7
+  rev: v0.13.1
   hooks:
     # Run the linter.
     - id: ruff-check
+      args: [ --config=pyproject.toml ]
     # Run the formatter.
     - id: ruff-format

.pre-commit-config.yaml

@@ -1,4 +1,16 @@
 import time
-async def fake_api_call(delay, data):
-    time.sleep(0.0001)
-    return f"Processed: {data}"
+import asyncio
+
+
+async def retry_with_backoff(func, max_retries=3):
+    if max_retries < 1:
+        raise ValueError("max_retries must be at least 1")
+    last_exception = None
+    for attempt in range(max_retries):
+        try:
+            return await func()
+        except Exception as e:
+            last_exception = e
+            if attempt < max_retries - 1:
+                time.sleep(0.0001 * attempt)
+    raise last_exception

code_to_optimize/code_directories/async_e2e/main.py

@@ -70,7 +70,11 @@
                 "request": "launch",
                 "program": "${workspaceFolder:codeflash}/codeflash/main.py",
                 "args": [
-                    "--file", "src/async_examples/shocker.py", "--verbose"
+                    "--file",
+                    "src/async_examples/concurrency.py",
+                    "--function",
+                    "task",
+                    "--verbose"
                 ],
                 "cwd": "${input:chooseCwd}",
                 "console": "integratedTerminal",

codeflash.code-workspace

@@ -298,6 +298,9 @@ def get_new_explanation(  # noqa: D417
         annotated_tests: str,
         optimization_id: str,
         original_explanation: str,
+        original_throughput: str | None = None,
+        optimized_throughput: str | None = None,
+        throughput_improvement: str | None = None,
     ) -> str:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -314,6 +317,9 @@ def get_new_explanation(  # noqa: D417
         - 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
+        - original_throughput: str | None - throughput for the baseline code (operations per second)
+        - optimized_throughput: str | None - throughput for the optimized code (operations per second)
+        - throughput_improvement: str | None - throughput improvement percentage
 
         Returns
         -------
@@ -333,6 +339,9 @@ def get_new_explanation(  # noqa: D417
             "optimization_id": optimization_id,
             "original_explanation": original_explanation,
             "dependency_code": dependency_code,
+            "original_throughput": original_throughput,
+            "optimized_throughput": optimized_throughput,
+            "throughput_improvement": throughput_improvement,
         }
         logger.info("Generating explanation")
         console.rule()