refactor: restructure codebase for locality and faster CLI startup

Move files closer to their consumers:
- function_context.py merged into code_context_extractor.py
- FunctionOptimizer base class to languages/function_optimizer.py
- test_runner, instrument_codeflash_capture, parse_line_profile to languages/python/
- oauth_handler.py to cli_cmds/

Split cmd_init.py (1993 lines) into focused modules:
- init_config.py: config types, validation, writing, shared UI
- init_auth.py: API key management + GitHub app installation
- github_workflow.py: GitHub Actions workflow generation
- cmd_init.py: init orchestrator + Python setup (639 lines)

Defer heavy imports (cmd_init, posthog, sentry) from module-level to
usage sites, reducing CLI startup from ~600ms to ~250ms. Replace
set_defaults(func=) with direct args.command dispatch in main().

.claude/rules/architecture.md

@@ -6,10 +6,15 @@ When adding, moving, or deleting source files, update this doc to match.
 codeflash/
 ├── main.py                 # CLI entry point
 ├── cli_cmds/               # Command handling, console output (Rich)
+│   ├── cmd_init.py         # Init orchestrator + Python-specific setup
+│   ├── init_config.py      # Config types, validation, writing, shared UI helpers
+│   ├── init_auth.py        # API key management + GitHub app installation
+│   ├── github_workflow.py  # GitHub Actions workflow generation
+│   ├── init_javascript.py  # JavaScript/TypeScript project initialization
+│   └── oauth_handler.py    # OAuth PKCE flow for CodeFlash authentication
 ├── discovery/              # Find optimizable functions
 ├── optimization/           # Generate optimized code via AI
-│   ├── optimizer.py        # Main optimization orchestration
+│   └── optimizer.py        # Main optimization orchestration
-│   └── function_optimizer.py  # Per-function optimization logic
 ├── verification/           # Run deterministic tests (pytest plugin)
 ├── benchmarking/           # Performance measurement
 ├── github/                 # PR creation
@@ -20,12 +25,16 @@ codeflash/
 │   ├── base.py                    # LanguageSupport protocol and shared data types
 │   ├── registry.py                # Language registration and lookup by extension/enum
 │   ├── current.py                 # Current language singleton (set_current_language / current_language_support)
+│   ├── function_optimizer.py       # FunctionOptimizer base class for per-function optimization
 │   ├── code_replacer.py           # Language-agnostic code replacement
 │   ├── python/
 │   │   ├── support.py             # PythonSupport (LanguageSupport implementation)
 │   │   ├── function_optimizer.py  # PythonFunctionOptimizer subclass
 │   │   ├── optimizer.py           # Python module preparation & AST resolution
-│   │   └── normalizer.py          # Python code normalization for deduplication
+│   │   ├── normalizer.py          # Python code normalization for deduplication
+│   │   ├── test_runner.py         # Test subprocess execution for Python
+│   │   ├── instrument_codeflash_capture.py  # Instrument __init__ with capture decorators
+│   │   └── parse_line_profile_test_output.py  # Parse line profiler output
 │   └── javascript/
 │       ├── support.py             # JavaScriptSupport (LanguageSupport implementation)
 │       ├── function_optimizer.py  # JavaScriptFunctionOptimizer subclass
@@ -46,9 +55,9 @@ codeflash/
 
 | Task | Start here |
 |------|------------|
-| CLI arguments & commands | `cli_cmds/cli.py` |
+| CLI arguments & commands | `cli_cmds/cli.py` (parsing), `main.py` (subcommand dispatch) |
 | Optimization orchestration | `optimization/optimizer.py` → `run()` |
-| Per-function optimization | `optimization/function_optimizer.py` (base), `languages/python/function_optimizer.py`, `languages/javascript/function_optimizer.py` |
+| Per-function optimization | `languages/function_optimizer.py` (base), `languages/python/function_optimizer.py`, `languages/javascript/function_optimizer.py` |
 | Function discovery | `discovery/functions_to_optimize.py` |
 | Context extraction | `languages/<lang>/context/code_context_extractor.py` |
 | Test execution | `languages/<lang>/support.py` (`run_behavioral_tests`, etc.), `verification/pytest_plugin.py` |

codeflash/cli_cmds/cli.py

@@ -5,10 +5,7 @@ from argparse import SUPPRESS, ArgumentParser, Namespace
 from pathlib import Path
 
 from codeflash.cli_cmds import logging_config
-from codeflash.cli_cmds.cli_common import apologize_and_exit
-from codeflash.cli_cmds.cmd_init import init_codeflash, install_github_actions
 from codeflash.cli_cmds.console import logger
-from codeflash.cli_cmds.extension import install_vscode_extension
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_utils import exit_with_message, normalize_ignore_paths
 from codeflash.code_utils.config_parser import parse_config_file
@@ -21,20 +18,12 @@ def parse_args() -> Namespace:
     parser = ArgumentParser()
     subparsers = parser.add_subparsers(dest="command", help="Sub-commands")
 
-    init_parser = subparsers.add_parser("init", help="Initialize Codeflash for your project.")
+    subparsers.add_parser("init", help="Initialize Codeflash for your project.")
-    init_parser.set_defaults(func=init_codeflash)
-
     subparsers.add_parser("vscode-install", help="Install the Codeflash VSCode extension")
-
+    subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow")
-    init_actions_parser = subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow")
-    init_actions_parser.set_defaults(func=install_github_actions)
 
     trace_optimize = subparsers.add_parser("optimize", help="Trace and optimize your project.")
 
-    from codeflash.tracer import main as tracer_main
-
-    trace_optimize.set_defaults(func=tracer_main)
-
     trace_optimize.add_argument(
         "--max-function-count",
         type=int,
@@ -182,10 +171,6 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
         _handle_reset_config(confirm=not getattr(args, "yes", False))
         sys.exit()
 
-    if args.command == "vscode-install":
-        install_vscode_extension()
-        sys.exit()
-
     if not check_running_in_git_repo(module_root=args.module_root):
         if not confirm_proceeding_with_no_git_repo():
             exit_with_message("No git repository detected and user aborted run. Exiting...", error_on_exit=True)
@@ -337,6 +322,8 @@ def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
                     f"I couldn't find a git repository in the current directory. "
                     f"I need a git repository to run {mode} and open PRs for optimizations. Exiting..."
                 )
+                from codeflash.cli_cmds.cli_common import apologize_and_exit
+
                 apologize_and_exit()
             git_remote = getattr(args, "git_remote", None)
             if not check_and_push_branch(git_repo, git_remote=git_remote):

codeflash/cli_cmds/console_constants.py

@@ -1,3 +1,18 @@
+from codeflash.code_utils.compat import LF
+from codeflash.version import __version__ as version
+
+CODEFLASH_LOGO: str = (
+    f"{LF}"
+    r"                   _          ___  _               _     " + f"{LF}"
+    r"                  | |        / __)| |             | |    " + f"{LF}"
+    r"  ____   ___    _ | |  ____ | |__ | |  ____   ___ | | _  " + f"{LF}"
+    r" / ___) / _ \  / || | / _  )|  __)| | / _  | /___)| || \ " + f"{LF}"
+    r"( (___ | |_| |( (_| |( (/ / | |   | |( ( | ||___ || | | |" + f"{LF}"
+    r" \____) \___/  \____| \____)|_|   |_| \_||_|(___/ |_| |_|" + f"{LF}"
+    f"{('v' + version).rjust(66)}{LF}"
+    f"{LF}"
+)
+
 SPINNER_TYPES = {
     "point",
     "simpleDots",

codeflash/cli_cmds/github_workflow.py

@@ -0,0 +1,898 @@
+from __future__ import annotations
+
+import sys
+from enum import Enum, auto
+from pathlib import Path
+from typing import Any
+
+import click
+import git
+import inquirer
+import tomlkit
+from git import Repo
+from rich.panel import Panel
+from rich.text import Text
+
+from codeflash.api.aiservice import AiServiceClient
+from codeflash.api.cfapi import setup_github_actions
+from codeflash.cli_cmds.cli_common import apologize_and_exit
+from codeflash.cli_cmds.console import console, logger
+from codeflash.cli_cmds.init_config import CodeflashTheme
+from codeflash.code_utils.compat import LF
+from codeflash.code_utils.config_parser import parse_config_file
+from codeflash.code_utils.env_utils import get_codeflash_api_key
+from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name
+from codeflash.code_utils.github_utils import get_github_secrets_page_url
+from codeflash.telemetry.posthog_cf import ph
+
+
+class DependencyManager(Enum):
+    """Python dependency managers."""
+
+    PIP = auto()
+    POETRY = auto()
+    UV = auto()
+    UNKNOWN = auto()
+
+
+def install_github_actions(override_formatter_check: bool = False) -> None:
+    try:
+        config, _config_file_path = parse_config_file(override_formatter_check=override_formatter_check)
+
+        ph("cli-github-actions-install-started")
+        try:
+            repo = Repo(config["module_root"], search_parent_directories=True)
+        except git.InvalidGitRepositoryError:
+            click.echo(
+                "Skipping GitHub action installation for continuous optimization because you're not in a git repository."
+            )
+            return
+
+        git_root = Path(repo.git.rev_parse("--show-toplevel"))
+        workflows_path = git_root / ".github" / "workflows"
+        optimize_yaml_path = workflows_path / "codeflash.yaml"
+
+        # Check if workflow file already exists locally BEFORE showing prompt
+        if optimize_yaml_path.exists():
+            # Workflow file already exists locally - skip prompt and setup
+            already_exists_message = "✅ GitHub Actions workflow file already exists.\n\n"
+            already_exists_message += "No changes needed - your repository is already configured!"
+
+            already_exists_panel = Panel(
+                Text(already_exists_message, style="green", justify="center"),
+                title="✅ Already Configured",
+                border_style="bright_green",
+            )
+            console.print(already_exists_panel)
+            console.print()
+
+            logger.info("[github_workflow.py:install_github_actions] Workflow file already exists locally, skipping setup")
+            return
+
+        # Get repository information for API call
+        git_remote = config.get("git_remote", "origin")
+        # get_current_branch handles detached HEAD and other edge cases internally
+        try:
+            base_branch = get_current_branch(repo)
+        except Exception as e:
+            logger.warning(
+                f"[github_workflow.py:install_github_actions] Could not determine current branch: {e}. Falling back to 'main'."
+            )
+            base_branch = "main"
+
+        # Generate workflow content
+        from importlib.resources import files
+
+        benchmark_mode = False
+        benchmarks_root = config.get("benchmarks_root", "").strip()
+        if benchmarks_root and benchmarks_root != "":
+            benchmark_panel = Panel(
+                Text(
+                    "📊 Benchmark Mode Available\n\n"
+                    "I noticed you've configured a benchmarks_root in your config. "
+                    "Benchmark mode will show the performance impact of Codeflash's optimizations on your benchmarks.",
+                    style="cyan",
+                ),
+                title="📊 Benchmark Mode",
+                border_style="bright_cyan",
+            )
+            console.print(benchmark_panel)
+            console.print()
+
+            benchmark_questions = [
+                inquirer.Confirm("benchmark_mode", message="Run GitHub Actions in benchmark mode?", default=True)
+            ]
+
+            benchmark_answers = inquirer.prompt(benchmark_questions, theme=CodeflashTheme())
+            benchmark_mode = benchmark_answers["benchmark_mode"] if benchmark_answers else False
+
+        # Show prompt only if workflow doesn't exist locally
+        actions_panel = Panel(
+            Text(
+                "🤖 GitHub Actions Setup\n\n"
+                "GitHub Actions will automatically optimize your code in every pull request. "
+                "This is the recommended way to use Codeflash for continuous optimization.",
+                style="blue",
+            ),
+            title="🤖 Continuous Optimization",
+            border_style="bright_blue",
+        )
+        console.print(actions_panel)
+        console.print()
+
+        creation_questions = [
+            inquirer.Confirm(
+                "confirm_creation",
+                message="Set up GitHub Actions for continuous optimization? We'll open a pull request with the workflow file.",
+                default=True,
+            )
+        ]
+
+        creation_answers = inquirer.prompt(creation_questions, theme=CodeflashTheme())
+        if not creation_answers or not creation_answers["confirm_creation"]:
+            skip_panel = Panel(
+                Text("⏩️ Skipping GitHub Actions setup.", style="yellow"), title="⏩️ Skipped", border_style="yellow"
+            )
+            console.print(skip_panel)
+            ph("cli-github-workflow-skipped")
+            return
+        ph(
+            "cli-github-optimization-confirm-workflow-creation",
+            {"confirm_creation": creation_answers["confirm_creation"]},
+        )
+
+        # Generate workflow content AFTER user confirmation
+        logger.info("[github_workflow.py:install_github_actions] User confirmed, generating workflow content...")
+
+        # Select the appropriate workflow template based on project language
+        project_language = detect_project_language_for_workflow(Path.cwd())
+        if project_language in ("javascript", "typescript"):
+            workflow_template = "codeflash-optimize-js.yaml"
+        else:
+            workflow_template = "codeflash-optimize.yaml"
+
+        optimize_yml_content = (
+            files("codeflash").joinpath("cli_cmds", "workflows", workflow_template).read_text(encoding="utf-8")
+        )
+        materialized_optimize_yml_content = generate_dynamic_workflow_content(
+            optimize_yml_content, config, git_root, benchmark_mode
+        )
+
+        workflows_path.mkdir(parents=True, exist_ok=True)
+
+        pr_created_via_api = False
+        pr_url = None
+
+        try:
+            owner, repo_name = get_repo_owner_and_name(repo, git_remote)
+        except Exception as e:
+            logger.error(f"[github_workflow.py:install_github_actions] Failed to get repository owner and name: {e}")
+            # Fall back to local file creation
+            workflows_path.mkdir(parents=True, exist_ok=True)
+            with optimize_yaml_path.open("w", encoding="utf8") as optimize_yml_file:
+                optimize_yml_file.write(materialized_optimize_yml_content)
+            workflow_success_panel = Panel(
+                Text(
+                    f"✅ Created GitHub action workflow at {optimize_yaml_path}\n\n"
+                    "Your repository is now configured for continuous optimization!",
+                    style="green",
+                    justify="center",
+                ),
+                title="🎉 Workflow Created!",
+                border_style="bright_green",
+            )
+            console.print(workflow_success_panel)
+            console.print()
+        else:
+            # Try to create PR via API
+            try:
+                # Workflow file doesn't exist on remote or content differs - proceed with PR creation
+                console.print("Creating PR with GitHub Actions workflow...")
+                logger.info(
+                    f"[github_workflow.py:install_github_actions] Calling setup_github_actions API for {owner}/{repo_name} on branch {base_branch}"
+                )
+
+                response = setup_github_actions(
+                    owner=owner,
+                    repo=repo_name,
+                    base_branch=base_branch,
+                    workflow_content=materialized_optimize_yml_content,
+                )
+
+                if response.status_code == 200:
+                    response_data = response.json()
+                    if response_data.get("success"):
+                        pr_url = response_data.get("pr_url")
+
+                        if pr_url:
+                            pr_created_via_api = True
+                            success_message = f"✅ PR created: {pr_url}\n\n"
+                            success_message += "Your repository is now configured for continuous optimization!"
+
+                            workflow_success_panel = Panel(
+                                Text(success_message, style="green", justify="center"),
+                                title="🎉 Workflow PR Created!",
+                                border_style="bright_green",
+                            )
+                            console.print(workflow_success_panel)
+                            console.print()
+
+                            logger.info(
+                                f"[github_workflow.py:install_github_actions] Successfully created PR #{response_data.get('pr_number')} for {owner}/{repo_name}"
+                            )
+                        else:
+                            # File already exists with same content
+                            pr_created_via_api = True  # Mark as handled (no PR needed)
+                            already_exists_message = "✅ Workflow file already exists with the same content.\n\n"
+                            already_exists_message += "No changes needed - your repository is already configured!"
+
+                            already_exists_panel = Panel(
+                                Text(already_exists_message, style="green", justify="center"),
+                                title="✅ Already Configured",
+                                border_style="bright_green",
+                            )
+                            console.print(already_exists_panel)
+                            console.print()
+                    else:
+                        # API returned success=false, extract error details
+                        error_data = response_data
+                        error_msg = error_data.get("error", "Unknown error")
+                        error_message = error_data.get("message", error_msg)
+                        error_help = error_data.get("help", "")
+                        installation_url = error_data.get("installation_url")
+
+                        # For permission errors, don't fall back - show a focused message and abort early
+                        if response.status_code == 403:
+                            logger.error(
+                                f"[github_workflow.py:install_github_actions] Permission denied for {owner}/{repo_name}"
+                            )
+                            # Extract installation_url if available, otherwise use default
+                            installation_url_403 = error_data.get(
+                                "installation_url", "https://github.com/apps/codeflash-ai/installations/select_target"
+                            )
+
+                            permission_error_panel = Panel(
+                                Text(
+                                    "❌ Access Denied\n\n"
+                                    f"The GitHub App may not be installed on {owner}/{repo_name}, or it doesn't have the required permissions.\n\n"
+                                    "💡 To fix this:\n"
+                                    "1. Install the CodeFlash GitHub App on your repository\n"
+                                    "2. Ensure the app has 'Contents: write', 'Workflows: write', and 'Pull requests: write' permissions\n"
+                                    "3. Make sure you have write access to the repository\n\n"
+                                    f"🔗 Install GitHub App: {installation_url_403}",
+                                    style="red",
+                                ),
+                                title="❌ Setup Failed",
+                                border_style="red",
+                            )
+                            console.print(permission_error_panel)
+                            console.print()
+                            click.echo(
+                                f"Please install the CodeFlash GitHub App and ensure it has the required permissions.{LF}"
+                                f"Visit: {installation_url_403}{LF}"
+                            )
+                            apologize_and_exit()
+
+                        # Show detailed error panel for all other errors
+                        error_panel_text = f"❌ {error_msg}\n\n{error_message}\n"
+                        if error_help:
+                            error_panel_text += f"\n💡 {error_help}\n"
+                        if installation_url:
+                            error_panel_text += f"\n🔗 Install GitHub App: {installation_url}"
+
+                        error_panel = Panel(
+                            Text(error_panel_text, style="red"), title="❌ Setup Failed", border_style="red"
+                        )
+                        console.print(error_panel)
+                        console.print()
+
+                        # For GitHub App not installed, don't fall back - show clear instructions
+                        if response.status_code == 404 and installation_url:
+                            logger.error(
+                                f"[github_workflow.py:install_github_actions] GitHub App not installed on {owner}/{repo_name}"
+                            )
+                            click.echo(
+                                f"Please install the CodeFlash GitHub App on your repository to continue.{LF}"
+                                f"Visit: {installation_url}{LF}"
+                            )
+                            return
+
+                        # For other errors, fall back to local file creation
+                        raise Exception(error_message)  # noqa: TRY002, TRY301
+                else:
+                    # API call returned non-200 status, try to parse error response
+                    try:
+                        error_data = response.json()
+                        error_msg = error_data.get("error", "API request failed")
+                        error_message = error_data.get("message", f"API returned status {response.status_code}")
+                        error_help = error_data.get("help", "")
+                        installation_url = error_data.get("installation_url")
+
+                        # For permission errors, don't fall back - show a focused message and abort early
+                        if response.status_code == 403:
+                            logger.error(
+                                f"[github_workflow.py:install_github_actions] Permission denied for {owner}/{repo_name}"
+                            )
+                            # Extract installation_url if available, otherwise use default
+                            installation_url_403 = error_data.get(
+                                "installation_url", "https://github.com/apps/codeflash-ai/installations/select_target"
+                            )
+
+                            permission_error_panel = Panel(
+                                Text(
+                                    "❌ Access Denied\n\n"
+                                    f"The GitHub App may not be installed on {owner}/{repo_name}, or it doesn't have the required permissions.\n\n"
+                                    "💡 To fix this:\n"
+                                    "1. Install the CodeFlash GitHub App on your repository\n"
+                                    "2. Ensure the app has 'Contents: write', 'Workflows: write', and 'Pull requests: write' permissions\n"
+                                    "3. Make sure you have write access to the repository\n\n"
+                                    f"🔗 Install GitHub App: {installation_url_403}",
+                                    style="red",
+                                ),
+                                title="❌ Setup Failed",
+                                border_style="red",
+                            )
+                            console.print(permission_error_panel)
+                            console.print()
+                            click.echo(
+                                f"Please install the CodeFlash GitHub App and ensure it has the required permissions.{LF}"
+                                f"Visit: {installation_url_403}{LF}"
+                            )
+                            apologize_and_exit()
+
+                        # Show detailed error panel for all other errors
+                        error_panel_text = f"❌ {error_msg}\n\n{error_message}\n"
+                        if error_help:
+                            error_panel_text += f"\n💡 {error_help}\n"
+                        if installation_url:
+                            error_panel_text += f"\n🔗 Install GitHub App: {installation_url}"
+
+                        error_panel = Panel(
+                            Text(error_panel_text, style="red"), title="❌ Setup Failed", border_style="red"
+                        )
+                        console.print(error_panel)
+                        console.print()
+
+                        # For GitHub App not installed, don't fall back - show clear instructions
+                        if response.status_code == 404 and installation_url:
+                            logger.error(
+                                f"[github_workflow.py:install_github_actions] GitHub App not installed on {owner}/{repo_name}"
+                            )
+                            click.echo(
+                                f"Please install the CodeFlash GitHub App on your repository to continue.{LF}"
+                                f"Visit: {installation_url}{LF}"
+                            )
+                            return
+
+                        # For authentication errors, don't fall back
+                        if response.status_code == 401:
+                            logger.error(
+                                f"[github_workflow.py:install_github_actions] Authentication failed for {owner}/{repo_name}"
+                            )
+                            click.echo(f"Authentication failed. Please check your API key and try again.{LF}")
+                            return
+
+                        # For other errors, fall back to local file creation
+                        raise Exception(error_message)  # noqa: TRY002
+                    except (ValueError, KeyError) as parse_error:
+                        # Couldn't parse error response, use generic message
+                        status_msg = f"API returned status {response.status_code}"
+                        raise Exception(status_msg) from parse_error  # noqa: TRY002
+
+            except Exception as api_error:
+                # Fall back to local file creation if API call fails (for non-critical errors)
+                logger.warning(
+                    f"[github_workflow.py:install_github_actions] API call failed, falling back to local file creation: {api_error}"
+                )
+                workflows_path.mkdir(parents=True, exist_ok=True)
+                with optimize_yaml_path.open("w", encoding="utf8") as optimize_yml_file:
+                    optimize_yml_file.write(materialized_optimize_yml_content)
+                workflow_success_panel = Panel(
+                    Text(
+                        f"✅ Created GitHub action workflow at {optimize_yaml_path}\n\n"
+                        "Your repository is now configured for continuous optimization!",
+                        style="green",
+                        justify="center",
+                    ),
+                    title="🎉 Workflow Created!",
+                    border_style="bright_green",
+                )
+                console.print(workflow_success_panel)
+                console.print()
+
+        # Show appropriate message based on whether PR was created via API
+        if pr_created_via_api:
+            if pr_url:
+                click.echo(
+                    f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}"
+                    f"Once you merge the PR, the workflow will be active.{LF}"
+                )
+            else:
+                # File already exists
+                click.echo(
+                    f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}"
+                    f"The workflow is ready to use.{LF}"
+                )
+        else:
+            # Fell back to local file creation
+            click.echo(
+                f"Please edit, commit and push this GitHub actions file to your repo, and you're all set!{LF}"
+                f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}"
+            )
+
+        # Show GitHub secrets setup panel (needed in both cases - PR created via API or local file)
+        try:
+            existing_api_key = get_codeflash_api_key()
+        except OSError:
+            existing_api_key = None
+
+        # GitHub secrets setup panel - always shown since secrets are required for the workflow to work
+        secrets_message = (
+            "🔐 Next Step: Add API Key as GitHub Secret\n\n"
+            "You'll need to add your CODEFLASH_API_KEY as a secret to your GitHub repository.\n\n"
+            "📋 Steps:\n"
+            "1. Press Enter to open your repo's secrets page\n"
+            "2. Click 'New repository secret'\n"
+            "3. Add your API key with the variable name CODEFLASH_API_KEY"
+        )
+
+        if existing_api_key:
+            secrets_message += f"\n\n🔑 Your API Key: {existing_api_key}"
+
+        secrets_panel = Panel(
+            Text(secrets_message, style="blue"), title="🔐 GitHub Secrets Setup", border_style="bright_blue"
+        )
+        console.print(secrets_panel)
+
+        console.print(f"\n📍 Press Enter to open: {get_github_secrets_page_url(repo)}")
+        console.input()
+
+        click.launch(get_github_secrets_page_url(repo))
+
+        # Post-launch message panel
+        launch_panel = Panel(
+            Text(
+                "🐙 I opened your GitHub secrets page!\n\n"
+                "Note: If you see a 404, you probably don't have access to this repo's secrets. "
+                "Ask a repo admin to add it for you, or (not recommended) you can temporarily "
+                "hard-code your API key into the workflow file.",
+                style="cyan",
+            ),
+            title="🌐 Browser Opened",
+            border_style="bright_cyan",
+        )
+        console.print(launch_panel)
+        click.pause()
+        console.print()
+        ph("cli-github-workflow-created")
+    except KeyboardInterrupt:
+        apologize_and_exit()
+
+
+def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager:
+    """Determine which dependency manager is being used based on pyproject.toml contents."""
+    cwd = Path.cwd()
+    if (cwd / "poetry.lock").exists():
+        return DependencyManager.POETRY
+    if (cwd / "uv.lock").exists():
+        return DependencyManager.UV
+    if "tool" not in pyproject_data:
+        return DependencyManager.PIP
+
+    tool_section = pyproject_data["tool"]
+
+    # Check for poetry
+    if "poetry" in tool_section:
+        return DependencyManager.POETRY
+
+    # Check for uv
+    if any(key.startswith("uv") for key in tool_section):
+        return DependencyManager.UV
+
+    # Look for pip-specific markers
+    if "pip" in tool_section or "setuptools" in tool_section:
+        return DependencyManager.PIP
+
+    return DependencyManager.UNKNOWN
+
+
+def get_codeflash_github_action_command(dep_manager: DependencyManager) -> str:
+    """Generate the appropriate codeflash command based on the dependency manager."""
+    if dep_manager == DependencyManager.POETRY:
+        return """|
+          poetry env use python
+          poetry run codeflash"""
+    if dep_manager == DependencyManager.UV:
+        return "uv run codeflash"
+    # PIP or UNKNOWN
+    return "codeflash"
+
+
+def get_dependency_installation_commands(dep_manager: DependencyManager) -> tuple[str, str]:
+    """Generate commands to install the dependency manager and project dependencies."""
+    if dep_manager == DependencyManager.POETRY:
+        return """|
+          python -m pip install --upgrade pip
+          pip install poetry
+          poetry install --all-extras"""
+    if dep_manager == DependencyManager.UV:
+        return """|
+          uv sync --all-extras
+          uv pip install --upgrade codeflash"""
+    # PIP or UNKNOWN
+    return """|
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+          pip install codeflash"""
+
+
+def get_dependency_manager_installation_string(dep_manager: DependencyManager) -> str:
+    py_version = sys.version_info
+    python_version_string = f"'{py_version.major}.{py_version.minor}'"
+    if dep_manager == DependencyManager.UV:
+        return """name: 🐍 Setup UV
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true"""
+    return f"""name: 🐍 Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: {python_version_string}"""
+
+
+def get_github_action_working_directory(toml_path: Path, git_root: Path) -> str:
+    if toml_path.parent == git_root:
+        return ""
+    working_dir = str(toml_path.parent.relative_to(git_root))
+    return f"""defaults:
+      run:
+        working-directory: ./{working_dir}"""
+
+
+def detect_project_language_for_workflow(project_root: Path) -> str:
+    """Detect the primary language of the project for workflow generation.
+
+    Returns: 'python', 'javascript', or 'typescript'
+    """
+    # Check for TypeScript config
+    if (project_root / "tsconfig.json").exists():
+        return "typescript"
+
+    # Check for JavaScript/TypeScript indicators
+    has_package_json = (project_root / "package.json").exists()
+    has_pyproject = (project_root / "pyproject.toml").exists()
+
+    if has_package_json and not has_pyproject:
+        # Pure JS/TS project
+        return "javascript"
+    if has_pyproject and not has_package_json:
+        # Pure Python project
+        return "python"
+
+    # Both exist - count files to determine primary language
+    js_count = 0
+    py_count = 0
+    for file in project_root.rglob("*"):
+        if file.is_file():
+            suffix = file.suffix.lower()
+            if suffix in {".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs"}:
+                js_count += 1
+            elif suffix == ".py":
+                py_count += 1
+
+    if js_count > py_count:
+        return "javascript"
+    return "python"
+
+
+def collect_repo_files_for_workflow(git_root: Path) -> dict[str, Any]:
+    """Collect important repository files and directory structure for workflow generation.
+
+    :param git_root: Root directory of the git repository
+    :return: Dictionary with 'files' (path -> content) and 'directory_structure' (nested dict)
+    """
+    # Important files to collect with contents
+    important_files = [
+        "pyproject.toml",
+        "requirements.txt",
+        "requirements-dev.txt",
+        "requirements/requirements.txt",
+        "requirements/dev.txt",
+        "Pipfile",
+        "Pipfile.lock",
+        "poetry.lock",
+        "uv.lock",
+        "setup.py",
+        "setup.cfg",
+        "Dockerfile",
+        "docker-compose.yml",
+        "docker-compose.yaml",
+        "Makefile",
+        "README.md",
+        "README.rst",
+    ]
+
+    # Also collect GitHub workflows
+    workflows_path = git_root / ".github" / "workflows"
+    if workflows_path.exists():
+        important_files.extend(
+            str(workflow_file.relative_to(git_root)) for workflow_file in workflows_path.glob("*.yml")
+        )
+        important_files.extend(
+            str(workflow_file.relative_to(git_root)) for workflow_file in workflows_path.glob("*.yaml")
+        )
+
+    files_dict: dict[str, str] = {}
+    max_file_size = 8 * 1024  # 8KB limit per file
+
+    for file_path_str in important_files:
+        file_path = git_root / file_path_str
+        if file_path.exists() and file_path.is_file():
+            try:
+                content = file_path.read_text(encoding="utf-8", errors="ignore")
+                # Limit file size
+                if len(content) > max_file_size:
+                    content = content[:max_file_size] + "\n... (truncated)"
+                files_dict[file_path_str] = content
+            except Exception as e:
+                logger.warning(f"[github_workflow.py:collect_repo_files_for_workflow] Failed to read {file_path_str}: {e}")
+
+    # Collect 2-level directory structure
+    directory_structure: dict[str, Any] = {}
+    try:
+        for item in sorted(git_root.iterdir()):
+            if item.name.startswith(".") and item.name not in [".github", ".git"]:
+                continue  # Skip hidden files/folders except .github
+
+            if item.is_dir():
+                # Level 1: directory
+                dir_dict: dict[str, Any] = {"type": "directory", "contents": {}}
+                try:
+                    # Level 2: contents of directory
+                    for subitem in sorted(item.iterdir()):
+                        if subitem.name.startswith("."):
+                            continue
+                        if subitem.is_dir():
+                            dir_dict["contents"][subitem.name] = {"type": "directory"}
+                        else:
+                            dir_dict["contents"][subitem.name] = {"type": "file"}
+                except PermissionError:
+                    pass  # Skip directories we can't read
+                directory_structure[item.name] = dir_dict
+            elif item.is_file():
+                directory_structure[item.name] = {"type": "file"}
+    except Exception as e:
+        logger.warning(f"[github_workflow.py:collect_repo_files_for_workflow] Error collecting directory structure: {e}")
+
+    return {"files": files_dict, "directory_structure": directory_structure}
+
+
+def generate_dynamic_workflow_content(
+    optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False
+) -> str:
+    """Generate workflow content with dynamic steps from AI service, falling back to static template."""
+    # First, do the basic replacements that are always needed
+    module_path = str(Path(config["module_root"]).relative_to(git_root) / "**")
+    optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path)
+
+    # Detect project language
+    project_language = detect_project_language_for_workflow(Path.cwd())
+
+    # For JavaScript/TypeScript projects, use static template customization
+    # (AI-generated steps are currently Python-only)
+    if project_language in ("javascript", "typescript"):
+        return customize_codeflash_yaml_content(optimize_yml_content, config, git_root, benchmark_mode)
+
+    # Python project - try AI-generated steps
+    toml_path = Path.cwd() / "pyproject.toml"
+    try:
+        with toml_path.open(encoding="utf8") as pyproject_file:
+            pyproject_data = tomlkit.parse(pyproject_file.read())
+    except FileNotFoundError:
+        click.echo(
+            f"I couldn't find a pyproject.toml in the current directory.{LF}"
+            f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file."
+        )
+        apologize_and_exit()
+
+    working_dir = get_github_action_working_directory(toml_path, git_root)
+    optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir)
+
+    # Try to generate dynamic steps using AI service
+    try:
+        repo_data = collect_repo_files_for_workflow(git_root)
+
+        # Prepare codeflash config for AI
+        codeflash_config = {
+            "module_root": config["module_root"],
+            "tests_root": config.get("tests_root", ""),
+            "benchmark_mode": benchmark_mode,
+        }
+
+        aiservice_client = AiServiceClient()
+        dynamic_steps = aiservice_client.generate_workflow_steps(
+            repo_files=repo_data["files"],
+            directory_structure=repo_data["directory_structure"],
+            codeflash_config=codeflash_config,
+        )
+
+        if dynamic_steps:
+            # Replace the entire steps section with AI-generated steps
+            # Find the steps section in the template
+            steps_start = optimize_yml_content.find("    steps:")
+            if steps_start != -1:
+                # Find the end of the steps section (next line at same or less indentation)
+                lines = optimize_yml_content.split("\n")
+                steps_start_line = optimize_yml_content[:steps_start].count("\n")
+                steps_end_line = len(lines)
+
+                # Find where steps section ends (next job or end of file)
+                for i in range(steps_start_line + 1, len(lines)):
+                    line = lines[i]
+                    # Stop if we hit a line that's not indented (new job or end of jobs)
+                    if line and not line.startswith(" ") and not line.startswith("\t"):
+                        steps_end_line = i
+                        break
+
+                # Extract steps content from AI response (remove "steps:" prefix if present)
+                steps_content = dynamic_steps
+                if steps_content.startswith("steps:"):
+                    # Remove "steps:" and leading newline
+                    steps_content = steps_content[6:].lstrip("\n")
+
+                # Ensure proper indentation (8 spaces for steps section in YAML)
+                indented_steps = []
+                for line in steps_content.split("\n"):
+                    if line.strip():
+                        # If line doesn't start with enough spaces, add them
+                        if not line.startswith(" "):
+                            indented_steps.append("        " + line)
+                        else:
+                            # Preserve existing indentation but ensure minimum 8 spaces
+                            current_indent = len(line) - len(line.lstrip())
+                            if current_indent < 8:
+                                indented_steps.append(" " * 8 + line.lstrip())
+                            else:
+                                indented_steps.append(line)
+                    else:
+                        indented_steps.append("")
+
+                # Add codeflash command step at the end
+                dep_manager = determine_dependency_manager(pyproject_data)
+                codeflash_cmd = get_codeflash_github_action_command(dep_manager)
+                if benchmark_mode:
+                    codeflash_cmd += " --benchmark"
+
+                # Format codeflash command properly
+                if "|" in codeflash_cmd:
+                    # Multi-line command
+                    cmd_lines = codeflash_cmd.split("\n")
+                    codeflash_step = f"      - name: ⚡️Codeflash Optimization\n        run: {cmd_lines[0].strip()}"
+                    for cmd_line in cmd_lines[1:]:
+                        codeflash_step += f"\n          {cmd_line.strip()}"
+                else:
+                    codeflash_step = f"      - name: ⚡️Codeflash Optimization\n        run: {codeflash_cmd}"
+
+                indented_steps.append(codeflash_step)
+
+                # Reconstruct the workflow
+                return "\n".join([*lines[:steps_start_line], "    steps:", *indented_steps, *lines[steps_end_line:]])
+            logger.warning("[github_workflow.py:generate_dynamic_workflow_content] Could not find steps section in template")
+        else:
+            logger.debug(
+                "[github_workflow.py:generate_dynamic_workflow_content] AI service returned no steps, falling back to static"
+            )
+
+    except Exception as e:
+        logger.warning(
+            f"[github_workflow.py:generate_dynamic_workflow_content] Error generating dynamic workflow, falling back to static: {e}"
+        )
+
+    # Fallback to static template
+    return customize_codeflash_yaml_content(optimize_yml_content, config, git_root, benchmark_mode)
+
+
+def customize_codeflash_yaml_content(
+    optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False
+) -> str:
+    module_path = str(Path(config["module_root"]).relative_to(git_root) / "**")
+    optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path)
+
+    # Detect project language
+    project_language = detect_project_language_for_workflow(Path.cwd())
+
+    if project_language in ("javascript", "typescript"):
+        # JavaScript/TypeScript project
+        return _customize_js_workflow_content(optimize_yml_content, git_root, benchmark_mode)
+
+    # Python project (default)
+    return _customize_python_workflow_content(optimize_yml_content, git_root, benchmark_mode)
+
+
+def _customize_python_workflow_content(optimize_yml_content: str, git_root: Path, benchmark_mode: bool = False) -> str:
+    """Customize workflow content for Python projects."""
+    # Get dependency installation commands
+    toml_path = Path.cwd() / "pyproject.toml"
+    try:
+        with toml_path.open(encoding="utf8") as pyproject_file:
+            pyproject_data = tomlkit.parse(pyproject_file.read())
+    except FileNotFoundError:
+        click.echo(
+            f"I couldn't find a pyproject.toml in the current directory.{LF}"
+            f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file."
+        )
+        apologize_and_exit()
+
+    working_dir = get_github_action_working_directory(toml_path, git_root)
+    optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir)
+    dep_manager = determine_dependency_manager(pyproject_data)
+
+    python_depmanager_installation = get_dependency_manager_installation_string(dep_manager)
+    optimize_yml_content = optimize_yml_content.replace(
+        "{{ setup_runtime_environment }}", python_depmanager_installation
+    )
+    install_deps_cmd = get_dependency_installation_commands(dep_manager)
+
+    optimize_yml_content = optimize_yml_content.replace("{{ install_dependencies_command }}", install_deps_cmd)
+
+    # Add codeflash command
+    codeflash_cmd = get_codeflash_github_action_command(dep_manager)
+
+    if benchmark_mode:
+        codeflash_cmd += " --benchmark"
+    return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd)
+
+
+def _customize_js_workflow_content(optimize_yml_content: str, git_root: Path, benchmark_mode: bool = False) -> str:
+    """Customize workflow content for JavaScript/TypeScript projects."""
+    from codeflash.cli_cmds.init_javascript import (
+        determine_js_package_manager,
+        get_js_codeflash_install_step,
+        get_js_codeflash_run_command,
+        get_js_dependency_installation_commands,
+        get_js_runtime_setup_steps,
+        is_codeflash_dependency,
+    )
+
+    project_root = Path.cwd()
+    package_json_path = project_root / "package.json"
+
+    if not package_json_path.exists():
+        click.echo(
+            f"I couldn't find a package.json in the current directory.{LF}"
+            f"Please run `npm init` or create a package.json file first."
+        )
+        apologize_and_exit()
+
+    # Determine working directory relative to git root
+    if project_root == git_root:
+        working_dir = ""
+    else:
+        rel_path = str(project_root.relative_to(git_root))
+        working_dir = f"""defaults:
+      run:
+        working-directory: ./{rel_path}"""
+
+    optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir)
+
+    # Determine package manager and codeflash dependency status
+    pkg_manager = determine_js_package_manager(project_root)
+    codeflash_is_dep = is_codeflash_dependency(project_root)
+
+    # Setup runtime environment (Node.js/Bun)
+    runtime_setup = get_js_runtime_setup_steps(pkg_manager)
+    optimize_yml_content = optimize_yml_content.replace("{{ setup_runtime_steps }}", runtime_setup)
+
+    # Install dependencies
+    install_deps_cmd = get_js_dependency_installation_commands(pkg_manager)
+    optimize_yml_content = optimize_yml_content.replace("{{ install_dependencies_command }}", install_deps_cmd)
+
+    # Install codeflash step (only if not a dependency)
+    install_codeflash = get_js_codeflash_install_step(pkg_manager, is_dependency=codeflash_is_dep)
+    optimize_yml_content = optimize_yml_content.replace("{{ install_codeflash_step }}", install_codeflash)
+
+    # Codeflash run command
+    codeflash_cmd = get_js_codeflash_run_command(pkg_manager, is_dependency=codeflash_is_dep)
+    if benchmark_mode:
+        codeflash_cmd += " --benchmark"
+    return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd)

codeflash/cli_cmds/init_auth.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import os
+
+import click
+import git
+import inquirer
+
+from codeflash.api.cfapi import get_user_id, is_github_app_installed_on_repo
+from codeflash.cli_cmds.cli_common import apologize_and_exit
+from codeflash.cli_cmds.console import console
+from codeflash.cli_cmds.init_config import CodeflashTheme
+from codeflash.cli_cmds.oauth_handler import perform_oauth_signin
+from codeflash.code_utils.compat import LF
+from codeflash.code_utils.env_utils import get_codeflash_api_key
+from codeflash.code_utils.git_utils import get_git_remotes, get_repo_owner_and_name
+from codeflash.code_utils.shell_utils import get_shell_rc_path, save_api_key_to_rc
+from codeflash.either import is_successful
+from codeflash.telemetry.posthog_cf import ph
+
+
+class CFAPIKeyType(click.ParamType):
+    name = "cfapi-key"
+
+    def convert(self, value: str, param: click.Parameter | None, ctx: click.Context | None) -> str | None:
+        value = value.strip()
+        if not value.startswith("cf-") and value != "":
+            self.fail(
+                f"That key [{value}] seems to be invalid. It should start with a 'cf-' prefix. Please try again.",
+                param,
+                ctx,
+            )
+        return value
+
+
+# Returns True if the user entered a new API key, False if they used an existing one
+def prompt_api_key() -> bool:
+    """Prompt user for API key via OAuth or manual entry."""
+    from rich.panel import Panel
+    from rich.text import Text
+
+    # Check for existing API key
+    try:
+        existing_api_key = get_codeflash_api_key()
+    except OSError:
+        existing_api_key = None
+
+    if existing_api_key:
+        display_key = f"{existing_api_key[:3]}****{existing_api_key[-4:]}"
+        api_key_panel = Panel(
+            Text(
+                f"🔑 I found a CODEFLASH_API_KEY in your environment [{display_key}]!\n\n"
+                "✅ You're all set with API authentication!",
+                style="green",
+                justify="center",
+            ),
+            title="🔑 API Key Found",
+            border_style="bright_green",
+        )
+        console.print(api_key_panel)
+        console.print()
+        return False
+
+    # Prompt for authentication method
+    auth_choices = ["🔐 Login in with Codeflash", "🔑 Use Codeflash API key"]
+
+    questions = [
+        inquirer.List(
+            "auth_method",
+            message="How would you like to authenticate?",
+            choices=auth_choices,
+            default=auth_choices[0],
+            carousel=True,
+        )
+    ]
+
+    answers = inquirer.prompt(questions, theme=CodeflashTheme())
+    if not answers:
+        apologize_and_exit()
+
+    method = answers["auth_method"]
+
+    if method == auth_choices[1]:
+        enter_api_key_and_save_to_rc()
+        ph("cli-new-api-key-entered")
+        return True
+
+    # Perform OAuth sign-in
+    api_key = perform_oauth_signin()
+
+    if not api_key:
+        apologize_and_exit()
+
+    # Save API key
+    shell_rc_path = get_shell_rc_path()
+    if not shell_rc_path.exists() and os.name == "nt":
+        shell_rc_path.touch()
+        click.echo(f"✅ Created {shell_rc_path}")
+
+    result = save_api_key_to_rc(api_key)
+    if is_successful(result):
+        click.echo(result.unwrap())
+        click.echo("✅ Signed in successfully and API key saved!")
+    else:
+        click.echo(result.failure())
+        click.pause()
+
+    os.environ["CODEFLASH_API_KEY"] = api_key
+    ph("cli-oauth-signin-completed")
+    return True
+
+
+def enter_api_key_and_save_to_rc() -> None:
+    browser_launched = False
+    api_key = ""
+    while api_key == "":
+        api_key = click.prompt(
+            f"Enter your Codeflash API key{' [or press Enter to open your API key page]' if not browser_launched else ''}",
+            hide_input=False,
+            default="",
+            type=CFAPIKeyType(),
+            show_default=False,
+        ).strip()
+        if api_key:
+            break
+        if not browser_launched:
+            click.echo(
+                f"Opening your Codeflash API key page. Grab a key from there!{LF}"
+                "You can also open this link manually: https://app.codeflash.ai/app/apikeys"
+            )
+            click.launch("https://app.codeflash.ai/app/apikeys")
+            browser_launched = True  # This does not work on remote consoles
+    shell_rc_path = get_shell_rc_path()
+    if not shell_rc_path.exists() and os.name == "nt":
+        # On Windows, create the appropriate file (PowerShell .ps1 or CMD .bat) in the user's home directory
+        shell_rc_path.parent.mkdir(parents=True, exist_ok=True)
+        shell_rc_path.touch()
+        click.echo(f"✅ Created {shell_rc_path}")
+    get_user_id(api_key=api_key)  # Used to verify whether the API key is valid.
+    result = save_api_key_to_rc(api_key)
+    if is_successful(result):
+        click.echo(result.unwrap())
+    else:
+        click.echo(result.failure())
+        click.pause()
+
+    os.environ["CODEFLASH_API_KEY"] = api_key
+
+
+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
+
+    if git_remote not in get_git_remotes(git_repo):
+        click.echo(f"Skipping GitHub app installation, remote ({git_remote}) does not exist in this repository.")
+        return
+
+    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(
+            f"🐙 Looks like you've already installed the Codeflash GitHub app on this repository ({owner}/{repo})! Continuing…"
+        )
+
+    else:
+        try:
+            click.prompt(
+                f"Finally, you'll need to install the Codeflash GitHub app by choosing the repository you want to install Codeflash on.{LF}"
+                f"I will attempt to open the github app page - https://github.com/apps/codeflash-ai/installations/select_target {LF}"
+                f"Please, press ENTER to open the app installation page{LF}",
+                default="",
+                type=click.STRING,
+                prompt_suffix=">>> ",
+                show_default=False,
+            )
+            click.launch("https://github.com/apps/codeflash-ai/installations/select_target")
+            click.prompt(
+                f"Please, press ENTER once you've finished installing the github app from https://github.com/apps/codeflash-ai/installations/select_target{LF}",
+                default="",
+                type=click.STRING,
+                prompt_suffix=">>> ",
+                show_default=False,
+            )
+
+            count = 2
+            while not is_github_app_installed_on_repo(owner, repo, suppress_errors=True):
+                if count == 0:
+                    click.echo(
+                        f"❌ It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo}.{LF}"
+                        f"You won't be able to create PRs with Codeflash until you install the app.{LF}"
+                        f"In the meantime you can make local only optimizations by using the '--no-pr' flag with codeflash.{LF}"
+                    )
+                    break
+                click.prompt(
+                    f"❌ It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo}.{LF}"
+                    f"Please install it from https://github.com/apps/codeflash-ai/installations/select_target {LF}"
+                    f"Please, press ENTER to continue once you've finished installing the github app…{LF}",
+                    default="",
+                    type=click.STRING,
+                    prompt_suffix=">>> ",
+                    show_default=False,
+                )
+                count -= 1
+        except (KeyboardInterrupt, EOFError, click.exceptions.Abort):
+            # leave empty line for the next prompt to be properly rendered
+            click.echo()

codeflash/cli_cmds/init_config.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+import os
+from enum import Enum
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import click
+import inquirer
+import inquirer.themes
+import tomlkit
+from pydantic.dataclasses import dataclass
+
+from codeflash.cli_cmds.cli_common import apologize_and_exit
+from codeflash.cli_cmds.console import console
+from codeflash.code_utils.compat import LF
+from codeflash.code_utils.config_parser import parse_config_file
+from codeflash.code_utils.env_utils import check_formatter_installed
+from codeflash.lsp.helpers import is_LSP_enabled
+from codeflash.telemetry.posthog_cf import ph
+
+
+@dataclass(frozen=True)
+class CLISetupInfo:
+    """Setup info for Python projects."""
+
+    module_root: str
+    tests_root: str
+    benchmarks_root: Union[str, None]
+    ignore_paths: list[str]
+    formatter: Union[str, list[str]]
+    git_remote: str
+    enable_telemetry: bool
+
+
+@dataclass(frozen=True)
+class VsCodeSetupInfo:
+    """Setup info for VSCode extension initialization."""
+
+    module_root: str
+    tests_root: str
+    formatter: Union[str, list[str]]
+
+
+# Custom theme for better UX
+class CodeflashTheme(inquirer.themes.Default):
+    def __init__(self) -> None:
+        super().__init__()
+        self.Question.mark_color = inquirer.themes.term.yellow
+        self.Question.brackets_color = inquirer.themes.term.bright_blue
+        self.Question.default_color = inquirer.themes.term.bright_cyan
+        self.List.selection_color = inquirer.themes.term.bright_blue
+        self.Checkbox.selection_color = inquirer.themes.term.bright_blue
+        self.Checkbox.selected_icon = "✅"
+        self.Checkbox.unselected_icon = "⬜"
+
+
+# common sections between normal mode and lsp mode
+class CommonSections(Enum):
+    module_root = "module_root"
+    tests_root = "tests_root"
+    formatter_cmds = "formatter_cmds"
+
+    def get_toml_key(self) -> str:
+        return self.value.replace("_", "-")
+
+
+@lru_cache(maxsize=1)
+def get_valid_subdirs(current_dir: Optional[Path] = None) -> list[str]:
+    ignore_subdirs = [
+        "venv",
+        "node_modules",
+        "dist",
+        "build",
+        "build_temp",
+        "build_scripts",
+        "env",
+        "logs",
+        "tmp",
+        "__pycache__",
+    ]
+    path_str = str(current_dir) if current_dir else "."
+    return [
+        d
+        for d in next(os.walk(path_str))[1]
+        if not d.startswith(".") and not d.startswith("__") and d not in ignore_subdirs
+    ]
+
+
+def get_suggestions(section: str) -> tuple[list[str], Optional[str]]:
+    valid_subdirs = get_valid_subdirs()
+    if section == CommonSections.module_root:
+        return [d for d in valid_subdirs if d != "tests"], None
+    if section == CommonSections.tests_root:
+        default = "tests" if "tests" in valid_subdirs else None
+        return valid_subdirs, default
+    if section == CommonSections.formatter_cmds:
+        return ["disabled", "ruff", "black"], "disabled"
+    msg = f"Unknown section: {section}"
+    raise ValueError(msg)
+
+
+def config_found(pyproject_toml_path: Union[str, Path]) -> tuple[bool, str]:
+    pyproject_toml_path = Path(pyproject_toml_path)
+
+    if not pyproject_toml_path.exists():
+        return False, f"Configuration file not found: {pyproject_toml_path}"
+
+    if not pyproject_toml_path.is_file():
+        return False, f"Configuration file is not a file: {pyproject_toml_path}"
+
+    if pyproject_toml_path.suffix != ".toml":
+        return False, f"Configuration file is not a .toml file: {pyproject_toml_path}"
+
+    return True, ""
+
+
+def is_valid_pyproject_toml(pyproject_toml_path: Union[str, Path]) -> tuple[bool, dict[str, Any] | None, str]:
+    pyproject_toml_path = Path(pyproject_toml_path)
+    try:
+        config, _ = parse_config_file(pyproject_toml_path)
+    except Exception as e:
+        return False, None, f"Failed to parse configuration: {e}"
+
+    module_root = config.get("module_root")
+    if not module_root:
+        return False, config, "Missing required field: 'module_root'"
+
+    if not Path(module_root).is_dir():
+        return False, config, f"Invalid 'module_root': directory does not exist at {module_root}"
+
+    tests_root = config.get("tests_root")
+    if not tests_root:
+        return False, config, "Missing required field: 'tests_root'"
+
+    if not Path(tests_root).is_dir():
+        return False, config, f"Invalid 'tests_root': directory does not exist at {tests_root}"
+
+    return True, config, ""
+
+
+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.
+    """
+    from rich.prompt import Confirm
+
+    pyproject_toml_path = Path.cwd() / "pyproject.toml"
+
+    found, _ = config_found(pyproject_toml_path)
+    if not found:
+        return True, None
+
+    valid, config, _message = is_valid_pyproject_toml(pyproject_toml_path)
+    if not valid:
+        # needs to be re-configured
+        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
+
+
+def get_formatter_cmds(formatter: str) -> list[str]:
+    if formatter == "black":
+        return ["black $file"]
+    if formatter == "ruff":
+        return ["ruff check --exit-zero --fix $file", "ruff format $file"]
+    if formatter == "other":
+        click.echo(
+            "🔧 In pyproject.toml, please replace 'your-formatter' with the command you use to format your code."
+        )
+        return ["your-formatter $file"]
+    if formatter in {"don't use a formatter", "disabled"}:
+        return ["disabled"]
+    if " && " in formatter:
+        return formatter.split(" && ")
+    return [formatter]
+
+
+# Create or update the pyproject.toml file with the Codeflash dependency & configuration
+def configure_pyproject_toml(
+    setup_info: Union[VsCodeSetupInfo, CLISetupInfo], config_file: Optional[Path] = None
+) -> bool:
+    for_vscode = isinstance(setup_info, VsCodeSetupInfo)
+    toml_path = config_file or Path.cwd() / "pyproject.toml"
+    try:
+        with toml_path.open(encoding="utf8") as pyproject_file:
+            pyproject_data = tomlkit.parse(pyproject_file.read())
+    except FileNotFoundError:
+        click.echo(
+            f"I couldn't find a pyproject.toml in the current directory.{LF}"
+            f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file."
+        )
+        return False
+
+    codeflash_section = tomlkit.table()
+    codeflash_section.add(tomlkit.comment("All paths are relative to this pyproject.toml's directory."))
+
+    if for_vscode:
+        for section in CommonSections:
+            if hasattr(setup_info, section.value):
+                codeflash_section[section.get_toml_key()] = getattr(setup_info, section.value)
+    else:
+        codeflash_section["module-root"] = setup_info.module_root
+        codeflash_section["tests-root"] = setup_info.tests_root
+        codeflash_section["ignore-paths"] = setup_info.ignore_paths
+        if not setup_info.enable_telemetry:
+            codeflash_section["disable-telemetry"] = not setup_info.enable_telemetry
+        if setup_info.git_remote not in ["", "origin"]:
+            codeflash_section["git-remote"] = setup_info.git_remote
+
+    formatter = setup_info.formatter
+
+    formatter_cmds = formatter if isinstance(formatter, list) else get_formatter_cmds(formatter)
+
+    check_formatter_installed(formatter_cmds, exit_on_failure=False)
+    codeflash_section["formatter-cmds"] = formatter_cmds
+    # Add the 'codeflash' section, ensuring 'tool' section exists
+    tool_section = pyproject_data.get("tool", tomlkit.table())
+
+    if for_vscode:
+        # merge the existing codeflash section, instead of overwriting it
+        existing_codeflash = tool_section.get("codeflash", tomlkit.table())
+
+        for key, value in codeflash_section.items():
+            existing_codeflash[key] = value
+        tool_section["codeflash"] = existing_codeflash
+    else:
+        tool_section["codeflash"] = codeflash_section
+
+    pyproject_data["tool"] = tool_section
+
+    with toml_path.open("w", encoding="utf8") as pyproject_file:
+        pyproject_file.write(tomlkit.dumps(pyproject_data))
+    click.echo(f"Added Codeflash configuration to {toml_path}")
+    click.echo()
+    return True
+
+
+def create_empty_pyproject_toml(pyproject_toml_path: Path) -> None:
+    ph("cli-create-pyproject-toml")
+    lsp_mode = is_LSP_enabled()
+    # Define a minimal pyproject.toml content
+    new_pyproject_toml = tomlkit.document()
+    new_pyproject_toml["tool"] = {"codeflash": {}}
+    try:
+        pyproject_toml_path.write_text(tomlkit.dumps(new_pyproject_toml), encoding="utf8")
+
+        # Check if the pyproject.toml file was created
+        if pyproject_toml_path.exists() and not lsp_mode:
+            from rich.panel import Panel
+            from rich.text import Text
+
+            success_panel = Panel(
+                Text(
+                    f"✅ Created a pyproject.toml file at {pyproject_toml_path}\n\n"
+                    "Your project is now ready for Codeflash configuration!",
+                    style="green",
+                    justify="center",
+                ),
+                title="🎉 Success!",
+                border_style="bright_green",
+            )
+            console.print(success_panel)
+            console.print("\n📍 Press any key to continue...")
+            console.input()
+        ph("cli-created-pyproject-toml")
+    except OSError:
+        click.echo("❌ Failed to create pyproject.toml. Please check your disk permissions and available space.")
+        apologize_and_exit()
+
+
+def ask_for_telemetry() -> bool:
+    """Prompt the user to enable or disable telemetry."""
+    from rich.prompt import Confirm
+
+    return Confirm.ask(
+        "⚡️ Help us improve Codeflash by sharing anonymous usage data (e.g. errors encountered)?",
+        default=True,
+        show_default=True,
+    )

codeflash/cli_cmds/init_javascript.py

@@ -69,7 +69,7 @@ class JSSetupInfo:
 # Import theme from cmd_init to avoid duplication
 def _get_theme():
     """Get the CodeflashTheme - imported lazily to avoid circular imports."""
-    from codeflash.cli_cmds.cmd_init import CodeflashTheme
+    from codeflash.cli_cmds.init_config import CodeflashTheme
 
     return CodeflashTheme()
 
@@ -210,7 +210,8 @@ def get_package_install_command(project_root: Path, package: str, dev: bool = Tr
 
 def init_js_project(language: ProjectLanguage, *, skip_confirm: bool = False, skip_api_key: bool = False) -> None:
     """Initialize Codeflash for a JavaScript/TypeScript project."""
-    from codeflash.cli_cmds.cmd_init import install_github_actions, install_github_app, prompt_api_key
+    from codeflash.cli_cmds.github_workflow import install_github_actions
+    from codeflash.cli_cmds.init_auth import install_github_app, prompt_api_key
 
     lang_name = "TypeScript" if language == ProjectLanguage.TYPESCRIPT else "JavaScript"
 
@@ -325,7 +326,7 @@ def collect_js_setup_info(language: ProjectLanguage, *, skip_confirm: bool = Fal
     Uses auto-detection for most settings and only asks for overrides if needed.
     When skip_confirm is True, uses all auto-detected defaults without prompting.
     """
-    from codeflash.cli_cmds.cmd_init import ask_for_telemetry, get_valid_subdirs
+    from codeflash.cli_cmds.init_config import ask_for_telemetry, get_valid_subdirs
     from codeflash.code_utils.config_js import (
         detect_formatter,
         detect_module_root,

codeflash/languages/base.py

@@ -678,7 +678,7 @@ class LanguageSupport(Protocol):
     @property
     def function_optimizer_class(self) -> type:
         """Return the FunctionOptimizer subclass for this language."""
-        from codeflash.optimization.function_optimizer import FunctionOptimizer
+        from codeflash.languages.function_optimizer import FunctionOptimizer
 
         return FunctionOptimizer
 

codeflash/languages/javascript/function_optimizer.py

@@ -15,6 +15,7 @@ from codeflash.code_utils.config_consts import (
     TOTAL_LOOPING_TIME_EFFECTIVE,
 )
 from codeflash.either import Failure, Success
+from codeflash.languages.function_optimizer import FunctionOptimizer
 from codeflash.models.models import (
     CodeOptimizationContext,
     CodeString,
@@ -23,7 +24,6 @@ from codeflash.models.models import (
     TestingMode,
     TestResults,
 )
-from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.equivalence import compare_test_results
 
 if TYPE_CHECKING:

codeflash/languages/python/context/code_context_extractor.py

@@ -37,7 +37,6 @@ from codeflash.models.models import (
     CodeStringsMarkdown,
     FunctionSource,
 )
-from codeflash.optimization.function_context import belongs_to_function_qualified
 
 if TYPE_CHECKING:
     from jedi.api.classes import Name
@@ -1470,3 +1469,41 @@ def prune_cst(
             include_init_dunder=include_init_dunder,
         ),
     )
+
+
+def belongs_to_method(name: Name, class_name: str, method_name: str) -> bool:
+    """Check if the given name belongs to the specified method."""
+    return belongs_to_function(name, method_name) and belongs_to_class(name, class_name)
+
+
+def belongs_to_function(name: Name, function_name: str) -> bool:
+    """Check if the given jedi Name is a direct child of the specified function."""
+    if name.name == function_name:  # Handles function definition and recursive function calls
+        return False
+    if (name := name.parent()) and name.type == "function":
+        return name.name == function_name
+    return False
+
+
+def belongs_to_class(name: Name, class_name: str) -> bool:
+    """Check if given jedi Name is a direct child of the specified class."""
+    while name := name.parent():
+        if name.type == "class":
+            return name.name == class_name
+    return False
+
+
+def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> bool:
+    """Check if the given jedi Name is a direct child of the specified function, matched by qualified function name."""
+    try:
+        if (
+            name.full_name.startswith(name.module_name)
+            and get_qualified_name(name.module_name, name.full_name) == qualified_function_name
+        ):
+            # Handles function definition and recursive function calls
+            return False
+        if (name := name.parent()) and name.type == "function":
+            return get_qualified_name(name.module_name, name.full_name) == qualified_function_name
+        return False
+    except ValueError:
+        return False

codeflash/languages/python/function_optimizer.py

@@ -10,6 +10,7 @@ from codeflash.cli_cmds.console import code_print, console, logger
 from codeflash.code_utils.code_utils import unified_diff_strings
 from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE
 from codeflash.either import Failure, Success
+from codeflash.languages.function_optimizer import FunctionOptimizer
 from codeflash.languages.python.context.unused_definition_remover import (
     detect_unused_helper_functions,
     revert_unused_helper_functions,
@@ -22,7 +23,6 @@ from codeflash.languages.python.static_analysis.code_replacer import (
 )
 from codeflash.languages.python.static_analysis.line_profile_utils import add_decorator_imports, contains_jit_decorator
 from codeflash.models.models import TestingMode, TestResults
-from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.parse_test_output import calculate_function_throughput_from_test_results
 
 if TYPE_CHECKING:
@@ -87,7 +87,7 @@ class PythonFunctionOptimizer(FunctionOptimizer):
         return original_conftest_content
 
     def instrument_capture(self, file_path_to_helper_classes: dict[Path, set[str]]) -> None:
-        from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
+        from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture
 
         instrument_codeflash_capture(self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root)
 
@@ -181,7 +181,7 @@ class PythonFunctionOptimizer(FunctionOptimizer):
     def parse_line_profile_test_results(
         self, line_profiler_output_file: Path | None
     ) -> tuple[TestResults | dict, CoverageData | None]:
-        from codeflash.verification.parse_line_profile_test_output import parse_line_profile_results
+        from codeflash.languages.python.parse_line_profile_test_output import parse_line_profile_results
 
         return parse_line_profile_results(line_profiler_output_file=line_profiler_output_file)
 

codeflash/languages/python/reference_graph.py

@@ -79,7 +79,7 @@ def _is_valid_definition(definition: Name, caller_qualified_name: str, project_r
         return False
 
     try:
-        from codeflash.optimization.function_context import belongs_to_function_qualified
+        from codeflash.languages.python.context.code_context_extractor import belongs_to_function_qualified
 
         if belongs_to_function_qualified(definition, caller_qualified_name):
             return False

codeflash/languages/python/support.py

@@ -1051,8 +1051,8 @@ class PythonSupport:
         from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE
         from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE
         from codeflash.languages.python.static_analysis.coverage_utils import prepare_coverage_files
+        from codeflash.languages.python.test_runner import execute_test_subprocess
         from codeflash.models.models import TestType
-        from codeflash.verification.test_runner import execute_test_subprocess
 
         blocklisted_plugins = ["benchmark", "codspeed", "xdist", "sugar"]
 
@@ -1156,7 +1156,7 @@ class PythonSupport:
 
         from codeflash.code_utils.code_utils import get_run_tmp_file
         from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE
-        from codeflash.verification.test_runner import execute_test_subprocess
+        from codeflash.languages.python.test_runner import execute_test_subprocess
 
         blocklisted_plugins = ["codspeed", "cov", "benchmark", "profiling", "xdist", "sugar"]
 
@@ -1200,7 +1200,7 @@ class PythonSupport:
         from codeflash.code_utils.code_utils import get_run_tmp_file
         from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE
         from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE
-        from codeflash.verification.test_runner import execute_test_subprocess
+        from codeflash.languages.python.test_runner import execute_test_subprocess
 
         blocklisted_plugins = ["codspeed", "cov", "benchmark", "profiling", "xdist", "sugar"]
 

codeflash/lsp/beta.py

@@ -11,13 +11,13 @@ from typing import TYPE_CHECKING, Optional, Union
 
 from codeflash.api.cfapi import get_codeflash_api_key, get_user_id
 from codeflash.cli_cmds.cli import process_pyproject_config
-from codeflash.cli_cmds.cmd_init import (
+from codeflash.cli_cmds.cmd_init import create_find_common_tags_file
+from codeflash.cli_cmds.init_config import (
     CommonSections,
     VsCodeSetupInfo,
     config_found,
     configure_pyproject_toml,
     create_empty_pyproject_toml,
-    create_find_common_tags_file,
     get_formatter_cmds,
     get_suggestions,
     get_valid_subdirs,

codeflash/main.py

@@ -18,14 +18,11 @@ if "--subagent" in sys.argv:
     warnings.filterwarnings("ignore")
 
 from codeflash.cli_cmds.cli import parse_args, process_pyproject_config
-from codeflash.cli_cmds.cmd_init import CODEFLASH_LOGO, ask_run_end_to_end_test
 from codeflash.cli_cmds.console import paneled_text
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.checkpoint import ask_should_use_checkpoint_get_functions
 from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.code_utils.version_check import check_for_newer_minor_version
-from codeflash.telemetry import posthog_cf
-from codeflash.telemetry.sentry import init_sentry
 
 if TYPE_CHECKING:
     from argparse import Namespace
@@ -33,6 +30,9 @@ if TYPE_CHECKING:
 
 def main() -> None:
     """Entry point for the codeflash command-line interface."""
+    from codeflash.telemetry import posthog_cf
+    from codeflash.telemetry.sentry import init_sentry
+
     args = parse_args()
     print_codeflash_banner()
 
@@ -46,11 +46,30 @@ def main() -> None:
             disable_telemetry = pyproject_config.get("disable_telemetry", False)
         init_sentry(enabled=not disable_telemetry, exclude_errors=True)
         posthog_cf.initialize_posthog(enabled=not disable_telemetry)
-        args.func()
+
+        if args.command == "init":
+            from codeflash.cli_cmds.cmd_init import init_codeflash
+
+            init_codeflash()
+        elif args.command == "init-actions":
+            from codeflash.cli_cmds.github_workflow import install_github_actions
+
+            install_github_actions()
+        elif args.command == "vscode-install":
+            from codeflash.cli_cmds.extension import install_vscode_extension
+
+            install_vscode_extension()
+        elif args.command == "optimize":
+            from codeflash.tracer import main as tracer_main
+
+            tracer_main(args)
     elif args.verify_setup:
         args = process_pyproject_config(args)
         init_sentry(enabled=not args.disable_telemetry, exclude_errors=True)
         posthog_cf.initialize_posthog(enabled=not args.disable_telemetry)
+
+        from codeflash.cli_cmds.cmd_init import ask_run_end_to_end_test
+
         ask_run_end_to_end_test(args)
     else:
         # Check for first-run experience (no config exists)
@@ -117,6 +136,8 @@ def _handle_config_loading(args: Namespace) -> Namespace | None:
 
 
 def print_codeflash_banner() -> None:
+    from codeflash.cli_cmds.console_constants import CODEFLASH_LOGO
+
     paneled_text(
         CODEFLASH_LOGO, panel_args={"title": "https://codeflash.ai", "expand": False}, text_args={"style": "bold gold3"}
     )

codeflash/optimization/function_context.py

@@ -1,46 +0,0 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-from codeflash.code_utils.code_utils import get_qualified_name
-
-if TYPE_CHECKING:
-    from jedi.api.classes import Name
-
-
-def belongs_to_method(name: Name, class_name: str, method_name: str) -> bool:
-    """Check if the given name belongs to the specified method."""
-    return belongs_to_function(name, method_name) and belongs_to_class(name, class_name)
-
-
-def belongs_to_function(name: Name, function_name: str) -> bool:
-    """Check if the given jedi Name is a direct child of the specified function."""
-    if name.name == function_name:  # Handles function definition and recursive function calls
-        return False
-    if (name := name.parent()) and name.type == "function":
-        return name.name == function_name
-    return False
-
-
-def belongs_to_class(name: Name, class_name: str) -> bool:
-    """Check if given jedi Name is a direct child of the specified class."""
-    while name := name.parent():
-        if name.type == "class":
-            return name.name == class_name
-    return False
-
-
-def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> bool:
-    """Check if the given jedi Name is a direct child of the specified function, matched by qualified function name."""
-    try:
-        if (
-            name.full_name.startswith(name.module_name)
-            and get_qualified_name(name.module_name, name.full_name) == qualified_function_name
-        ):
-            # Handles function definition and recursive function calls
-            return False
-        if (name := name.parent()) and name.type == "function":
-            return get_qualified_name(name.module_name, name.full_name) == qualified_function_name
-        return False
-    except ValueError:
-        return False

codeflash/optimization/optimizer.py

@@ -42,8 +42,8 @@ if TYPE_CHECKING:
     from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.languages.base import DependencyResolver
+    from codeflash.languages.function_optimizer import FunctionOptimizer
     from codeflash.models.models import BenchmarkKey, FunctionCalledInTest, ValidCode
-    from codeflash.optimization.function_optimizer import FunctionOptimizer
 
 
 class Optimizer:

codeflash/tracer.py

@@ -199,8 +199,8 @@ def main(args: Namespace | None = None) -> ArgumentParser:
                     result_pickle_file_path.unlink(missing_ok=True)
             if not parsed_args.trace_only and replay_test_paths:
                 from codeflash.cli_cmds.cli import parse_args, process_pyproject_config
-                from codeflash.cli_cmds.cmd_init import CODEFLASH_LOGO
                 from codeflash.cli_cmds.console import paneled_text
+                from codeflash.cli_cmds.console_constants import CODEFLASH_LOGO
                 from codeflash.languages import set_current_language
                 from codeflash.languages.base import Language
                 from codeflash.telemetry import posthog_cf

tests/benchmarks/test_benchmark_code_extract_code_context.py

@@ -22,7 +22,7 @@ def test_benchmark_extract(benchmark) -> None:
     )
     function_to_optimize = FunctionToOptimize(
         function_name="replace_function_and_helpers_with_optimized_code",
-        file_path=file_path / "optimization" / "function_optimizer.py",
+        file_path=file_path / "languages" / "function_optimizer.py",
         parents=[FunctionParent(name="FunctionOptimizer", type="ClassDef")],
         starting_line=None,
         ending_line=None,

tests/test_async_run_and_parse_tests.py

@@ -16,7 +16,7 @@ from codeflash.code_utils.instrument_existing_tests import (
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestType
 from codeflash.optimization.optimizer import Optimizer
-from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
+from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows")

tests/test_cmd_init.py

@@ -4,7 +4,7 @@ from pathlib import Path
 
 import pytest
 
-from codeflash.cli_cmds.cmd_init import (
+from codeflash.cli_cmds.init_config import (
     CLISetupInfo,
     VsCodeSetupInfo,
     configure_pyproject_toml,

tests/test_codeflash_capture.py

@@ -10,8 +10,8 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import FunctionParent, TestFile, TestFiles, TestingMode, TestType, VerificationType
 from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer
 from codeflash.verification.equivalence import compare_test_results
-from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
+from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture
-from codeflash.verification.test_runner import execute_test_subprocess
+from codeflash.languages.python.test_runner import execute_test_subprocess
 from codeflash.verification.verification_utils import TestConfig
 
 

tests/test_formatter.py

@@ -9,7 +9,7 @@ from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.code_utils.formatter import format_code, format_generated_code, sort_imports
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import CodeString, CodeStringsMarkdown
-from codeflash.optimization.function_optimizer import FunctionOptimizer
+from codeflash.languages.function_optimizer import FunctionOptimizer
 from codeflash.verification.verification_utils import TestConfig
 
 

tests/test_instrument_all_and_run.py

@@ -12,7 +12,7 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestType
 from codeflash.optimization.optimizer import Optimizer
 from codeflash.verification.equivalence import compare_test_results
-from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
+from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture
 
 # Used by cli instrumentation
 codeflash_wrap_string = """def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, codeflash_cur, codeflash_con, *args, **kwargs):

tests/test_instrument_codeflash_capture.py

@@ -3,7 +3,7 @@ from pathlib import Path
 from codeflash.code_utils.code_utils import get_run_tmp_file
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import FunctionParent
-from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
+from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture
 
 
 def test_add_codeflash_capture():

tests/test_instrumentation_run_results_aiservice.py

@@ -14,7 +14,7 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import FunctionParent, TestFile, TestFiles, TestingMode, TestType, VerificationType
 from codeflash.optimization.optimizer import Optimizer
 from codeflash.verification.equivalence import compare_test_results
-from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
+from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture
 
 # Used by aiservice instrumentation
 behavior_logging_code = """

tests/test_languages/test_javascript_optimization_flow.py

@@ -305,7 +305,7 @@ describe('fibonacci', () => {
         """Test FunctionOptimizer can be instantiated for JavaScript."""
         skip_if_js_not_supported()
         from codeflash.discovery.functions_to_optimize import find_all_functions_in_file
-        from codeflash.optimization.function_optimizer import FunctionOptimizer
+        from codeflash.languages.function_optimizer import FunctionOptimizer
 
         src_file = js_project / "utils.js"
         functions = find_all_functions_in_file(src_file)
@@ -340,7 +340,7 @@ describe('fibonacci', () => {
         """Test FunctionOptimizer can be instantiated for TypeScript."""
         skip_if_js_not_supported()
         from codeflash.discovery.functions_to_optimize import find_all_functions_in_file
-        from codeflash.optimization.function_optimizer import FunctionOptimizer
+        from codeflash.languages.function_optimizer import FunctionOptimizer
 
         src_file = ts_project / "utils.ts"
         functions = find_all_functions_in_file(src_file)