Merge branch 'main' into main

Author: Saga4

WARP.md

@@ -0,0 +1,191 @@
+# WARP.md
+
+This file provides guidance to WARP (warp.dev) when working with code in this repository.
+
+## Project Overview
+
+Codeflash is a general-purpose optimizer for Python that helps improve code performance while maintaining correctness. It uses advanced LLMs to generate optimization ideas, tests them for correctness, and benchmarks them for performance, then creates merge-ready pull requests.
+
+## Development Environment Setup
+
+### Prerequisites
+- Python 3.9+ (project uses uv for dependency management)
+- Git (for version control and PR creation)
+- Codeflash API key (for AI services)
+
+### Initial Setup
+```bash
+# Install dependencies using uv (preferred over pip)
+uv sync
+
+# Initialize codeflash configuration
+uv run codeflash init
+```
+
+## Core Development Commands
+
+### Code Quality & Linting
+```bash
+# Format code with ruff (includes check and format)
+uv run ruff check --fix codeflash/
+uv run ruff format codeflash/
+
+# Type checking with mypy
+uv run mypy codeflash/
+
+# Pre-commit hooks (ruff check + format)
+uv run pre-commit run --all-files
+```
+
+### Testing
+```bash
+# Run all tests
+uv run pytest
+
+# Run specific test file
+uv run pytest tests/test_specific_file.py
+
+# Run tests matching pattern
+uv run pytest -k "pattern"
+
+```
+
+### Running Codeflash
+```bash
+# Optimize entire codebase
+uv run codeflash --all
+
+# Optimize specific file
+uv run codeflash --file path/to/file.py
+
+# Optimize specific function
+uv run codeflash --function "module.function"
+
+# Optimize a script end-to-end
+uv run codeflash optimize script.py
+
+# Run with benchmarking
+uv run codeflash --benchmark
+
+# Verify setup
+uv run codeflash --verify-setup
+```
+
+## Architecture Overview
+
+### Main Components
+
+**Core Modules:**
+- `codeflash/main.py` - CLI entry point and command coordination
+- `codeflash/cli_cmds/` - Command-line interface implementations
+- `codeflash/optimization/` - Core optimization engine and algorithms
+- `codeflash/verification/` - Code correctness verification
+- `codeflash/benchmarking/` - Performance measurement and comparison
+- `codeflash/discovery/` - Code analysis and function discovery
+- `codeflash/tracing/` - Runtime tracing and profiling
+- `codeflash/context/` - Code context extraction and analysis
+- `codeflash/result/` - Result processing, PR creation, and explanations
+
+**Supporting Systems:**
+- `codeflash/api/` - Backend API communication
+- `codeflash/github/` - GitHub integration for PR creation
+- `codeflash/models/` - Data models and schemas
+- `codeflash/telemetry/` - Analytics and error reporting
+- `codeflash/code_utils/` - Code parsing, formatting, and manipulation utilities
+
+### Key Workflows
+
+1. **Code Discovery**: Analyzes codebase to identify optimization candidates
+2. **Context Extraction**: Extracts relevant code context and dependencies
+3. **Optimization Generation**: Uses LLMs to generate optimization candidates
+4. **Verification**: Tests optimizations for correctness using existing tests
+5. **Benchmarking**: Measures performance improvements
+6. **Result Processing**: Creates explanations and pull requests
+
+### Configuration
+
+Configuration is stored in `pyproject.toml` under `[tool.codeflash]`:
+- `module-root` - Source code location (default: "codeflash")
+- `tests-root` - Test location (default: "tests") 
+- `benchmarks-root` - Benchmark location (default: "tests/benchmarks")
+- `test-framework` - Testing framework ("pytest" or "unittest")
+- `formatter-cmds` - Commands for code formatting
+
+## Project Structure
+
+```
+codeflash/
+├── api/                 # Backend API communication
+├── benchmarking/        # Performance measurement
+├── cli_cmds/           # CLI command implementations
+├── code_utils/         # Code analysis and manipulation
+├── context/            # Code context extraction
+├── discovery/          # Function and test discovery  
+├── github/             # GitHub API integration
+├── lsp/                # Language server protocol support
+├── models/             # Data models and schemas
+├── optimization/       # Core optimization engine
+├── result/             # Result processing and PR creation
+├── telemetry/          # Analytics and monitoring
+├── tracing/            # Runtime tracing and profiling
+├── verification/       # Correctness verification
+└── main.py            # CLI entry point
+
+tests/                  # Test suite
+├── benchmarks/         # Performance benchmarks
+└── scripts/           # Test utilities
+
+docs/                   # Documentation
+code_to_optimize/       # Example code for optimization
+codeflash-benchmark/    # Benchmark workspace member
+```
+
+## Development Notes
+
+### Code Style
+- Uses ruff for linting and formatting (configured in pyproject.toml)
+- Strict mypy type checking enabled
+- Pre-commit hooks enforce code quality
+
+### Testing
+- pytest-based test suite with extensive coverage
+- Parameterized tests for multiple scenarios
+- Benchmarking tests for performance validation
+- Test discovery supports both pytest and unittest frameworks
+
+### Workspace Structure
+- Uses uv workspace with `codeflash-benchmark` as a member
+- Dependencies managed through uv.lock
+- Dynamic versioning from git tags using uv-dynamic-versioning
+
+### Build & Distribution
+- Uses hatchling as build backend
+- BSL-1.1 license
+- Excludes development files from distribution packages
+
+### CI/CD Integration
+- GitHub Actions workflow for automatic optimization of PR code
+- Pre-commit hooks for code quality enforcement
+- Automated testing and benchmarking
+
+## Important Patterns
+
+### Error Handling
+- Uses `either.py` for functional error handling patterns
+- Comprehensive error tracking through Sentry integration
+- Graceful degradation when AI services are unavailable
+
+### Instrumentation
+- Extensive tracing capabilities for performance analysis
+- Line profiler integration for detailed performance metrics
+- Custom tracer implementation for code execution analysis
+
+### AI Integration
+- Structured prompts and response handling for LLM interactions
+- Critic module for evaluating optimization quality
+- Context-aware code generation and explanation
+
+### Git Integration 
+- GitPython for repository operations
+- Automated PR creation with detailed explanations
+- Branch management for optimization experiments

codeflash-benchmark/codeflash_benchmark/plugin.py

@@ -8,21 +8,54 @@
 
 PYTEST_BENCHMARK_INSTALLED = importlib.util.find_spec("pytest_benchmark") is not None
 
+benchmark_options = [
+    ("--benchmark-columns", "store", None, "Benchmark columns"),
+    ("--benchmark-group-by", "store", None, "Benchmark group by"),
+    ("--benchmark-name", "store", None, "Benchmark name pattern"),
+    ("--benchmark-sort", "store", None, "Benchmark sort column"),
+    ("--benchmark-json", "store", None, "Benchmark JSON output file"),
+    ("--benchmark-save", "store", None, "Benchmark save name"),
+    ("--benchmark-warmup", "store", None, "Benchmark warmup"),
+    ("--benchmark-warmup-iterations", "store", None, "Benchmark warmup iterations"),
+    ("--benchmark-min-time", "store", None, "Benchmark minimum time"),
+    ("--benchmark-max-time", "store", None, "Benchmark maximum time"),
+    ("--benchmark-min-rounds", "store", None, "Benchmark minimum rounds"),
+    ("--benchmark-timer", "store", None, "Benchmark timer"),
+    ("--benchmark-calibration-precision", "store", None, "Benchmark calibration precision"),
+    ("--benchmark-disable", "store_true", False, "Disable benchmarks"),
+    ("--benchmark-skip", "store_true", False, "Skip benchmarks"),
+    ("--benchmark-only", "store_true", False, "Only run benchmarks"),
+    ("--benchmark-verbose", "store_true", False, "Verbose benchmark output"),
+    ("--benchmark-histogram", "store", None, "Benchmark histogram"),
+    ("--benchmark-compare", "store", None, "Benchmark compare"),
+    ("--benchmark-compare-fail", "store", None, "Benchmark compare fail threshold"),
+]
+
 
 def pytest_configure(config: pytest.Config) -> None:
     """Register the benchmark marker and disable conflicting plugins."""
     config.addinivalue_line("markers", "benchmark: mark test as a benchmark that should be run with codeflash tracing")
 
-    if config.getoption("--codeflash-trace") and PYTEST_BENCHMARK_INSTALLED:
-        config.option.benchmark_disable = True
-        config.pluginmanager.set_blocked("pytest_benchmark")
-        config.pluginmanager.set_blocked("pytest-benchmark")
+    if config.getoption("--codeflash-trace"):
+        # When --codeflash-trace is used, ignore all benchmark options by resetting them to defaults
+        for option, _, default, _ in benchmark_options:
+            option_name = option.replace("--", "").replace("-", "_")
+            if hasattr(config.option, option_name):
+                setattr(config.option, option_name, default)
+
+        if PYTEST_BENCHMARK_INSTALLED:
+            config.pluginmanager.set_blocked("pytest_benchmark")
+            config.pluginmanager.set_blocked("pytest-benchmark")
 
 
 def pytest_addoption(parser: pytest.Parser) -> None:
     parser.addoption(
         "--codeflash-trace", action="store_true", default=False, help="Enable CodeFlash tracing for benchmarks"
     )
+    # These options are ignored when --codeflash-trace is used
+    for option, action, default, help_text in benchmark_options:
+        help_suffix = " (ignored when --codeflash-trace is used)"
+        parser.addoption(option, action=action, default=default, help=help_text + help_suffix)
 
 
 @pytest.fixture
@@ -37,7 +70,7 @@ def benchmark(request: pytest.FixtureRequest) -> object:
     # If pytest-benchmark is installed and --codeflash-trace is not enabled,
     # return the normal pytest-benchmark fixture
     if PYTEST_BENCHMARK_INSTALLED:
-        from pytest_benchmark.fixture import BenchmarkFixture as BSF  # noqa: N814
+        from pytest_benchmark.fixture import BenchmarkFixture as BSF  # pyright: ignore[reportMissingImports]  # noqa: I001, N814
 
         bs = getattr(config, "_benchmarksession", None)
         if bs and bs.skip:

codeflash-benchmark/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "codeflash-benchmark"
-version = "0.1.0"
+version = "0.2.0"
 description = "Pytest benchmarking plugin for codeflash.ai - automatic code performance optimization"
 authors = [{ name = "CodeFlash Inc.", email = "[email protected]" }]
 requires-python = ">=3.9"
@@ -25,8 +25,8 @@ Repository = "https://github.com/codeflash-ai/codeflash-benchmark"
 codeflash-benchmark = "codeflash_benchmark.plugin"
 
 [build-system]
-requires = ["setuptools>=45", "wheel", "setuptools_scm"]
+requires = ["setuptools>=45", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools]
-packages = ["codeflash_benchmark"]
+packages = ["codeflash_benchmark"]

codeflash/api/aiservice.py

@@ -360,6 +360,7 @@ def log_results(  # noqa: D417
         is_correct: dict[str, bool] | None,
         optimized_line_profiler_results: dict[str, str] | None,
         metadata: dict[str, Any] | None,
+        optimizations_post: dict[str, str] | None = None,
     ) -> None:
         """Log features to the database.
 
@@ -372,6 +373,7 @@ def log_results(  # noqa: D417
         - is_correct (Optional[Dict[str, bool]]): Whether the optimized code is correct.
         - optimized_line_profiler_results: line_profiler results for every candidate mapped to their optimization_id
         - metadata: contains the best optimization id
+        - optimizations_post - dict mapping opt id to code str after postprocessing
 
         """
         payload = {
@@ -383,6 +385,7 @@ def log_results(  # noqa: D417
             "codeflash_version": codeflash_version,
             "optimized_line_profiler_results": optimized_line_profiler_results,
             "metadata": metadata,
+            "optimizations_post": optimizations_post,
         }
         try:
             self.make_ai_service_request("/log_features", payload=payload, timeout=5)

codeflash/cli_cmds/cmd_init.py

@@ -85,12 +85,16 @@ def init_codeflash() -> None:
 
         did_add_new_key = prompt_api_key()
 
-        if should_modify_pyproject_toml():
-            setup_info: SetupInfo = collect_setup_info()
+        should_modify, config = should_modify_pyproject_toml()
+
+        git_remote = config.get("git_remote", "origin") if config else "origin"
 
+        if should_modify:
+            setup_info: SetupInfo = collect_setup_info()
+            git_remote = setup_info.git_remote
             configure_pyproject_toml(setup_info)
 
-        install_github_app()
+        install_github_app(git_remote)
 
         install_github_actions(override_formatter_check=True)
 
@@ -151,7 +155,7 @@ def ask_run_end_to_end_test(args: Namespace) -> None:
         run_end_to_end_test(args, bubble_sort_path, bubble_sort_test_path)
 
 
-def should_modify_pyproject_toml() -> bool:
+def should_modify_pyproject_toml() -> tuple[bool, dict[str, Any] | None]:
     """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.
@@ -160,22 +164,22 @@ def should_modify_pyproject_toml() -> bool:
 
     pyproject_toml_path = Path.cwd() / "pyproject.toml"
     if not pyproject_toml_path.exists():
-        return True
+        return True, None
     try:
         config, config_file_path = parse_config_file(pyproject_toml_path)
     except Exception:
-        return True
+        return True, None
 
     if "module_root" not in config or config["module_root"] is None or not Path(config["module_root"]).is_dir():
-        return True
+        return True, None
     if "tests_root" not in config or config["tests_root"] is None or not Path(config["tests_root"]).is_dir():
-        return True
+        return True, None
 
     return Confirm.ask(
         "✅ A valid Codeflash config already exists in this project. Do you want to re-configure it?",
         default=False,
         show_default=True,
-    )
+    ), config
 
 
 # Custom theme for better UX
@@ -958,16 +962,18 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
     click.echo()
 
 
-def install_github_app() -> None:
+def install_github_app(git_remote: str) -> None:
     try:
         git_repo = git.Repo(search_parent_directories=True)
     except git.InvalidGitRepositoryError:
         click.echo("Skipping GitHub app installation because you're not in a git repository.")
         return
-    owner, repo = get_repo_owner_and_name(git_repo)
+    owner, repo = get_repo_owner_and_name(git_repo, git_remote)
 
     if is_github_app_installed_on_repo(owner, repo, suppress_errors=True):
-        click.echo("🐙 Looks like you've already installed the Codeflash GitHub app on this repository! Continuing…")
+        click.echo(
+            f"🐙 Looks like you've already installed the Codeflash GitHub app on this repository ({owner}/{repo})! Continuing…"
+        )
 
     else:
         click.prompt(

codeflash/discovery/pytest_new_process_discovery.py

@@ -41,7 +41,8 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s
 
     try:
         exitcode = pytest.main(
-            [tests_root, "-p no:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()]
+            [tests_root, "-p no:logging", "--collect-only", "-m", "not skip", "-p", "no:codeflash-benchmark"],
+            plugins=[PytestCollectionPlugin()],
         )
     except Exception as e:
         print(f"Failed to collect tests: {e!s}")