Merge branch 'main' into chore/asyncio-optimization

Author: HeshamHM28

README.md

@@ -17,7 +17,7 @@ How to use Codeflash -
 - Automate optimizing all __future__ code you will write by installing Codeflash as a GitHub action.
 - Optimize a Python workflow `python myscript.py` end-to-end by running `codeflash optimize myscript.py`
 
-Codeflash is used by top engineering teams at [Pydantic](https://github.com/pydantic/pydantic/pulls?q=is%3Apr+author%3Amisrasaurabh1+is%3Amerged), [Langflow](https://github.com/langflow-ai/langflow/issues?q=state%3Aclosed%20is%3Apr%20author%3Amisrasaurabh1), [Roboflow](https://github.com/roboflow/inference/pulls?q=is%3Apr+is%3Amerged+codeflash+sort%3Acreated-asc), [Albumentations](https://github.com/albumentations-team/albumentations/issues?q=state%3Amerged%20is%3Apr%20author%3Akrrt7%20OR%20state%3Amerged%20is%3Apr%20author%3Aaseembits93%20) and many others to ship performant, expert level code.
+Codeflash is used by top engineering teams at **Pydantic** [(PRs Merged)](https://github.com/pydantic/pydantic/pulls?q=is%3Apr+author%3Amisrasaurabh1+is%3Amerged), **Roboflow** [(PRs Merged 1](https://github.com/roboflow/inference/issues?q=state%3Aclosed%20is%3Apr%20author%3Amisrasaurabh1%20is%3Amerged), [PRs Merged 2)](https://github.com/roboflow/inference/issues?q=state%3Amerged%20is%3Apr%20author%3Acodeflash-ai%5Bbot%5D), **Unstructured** [(PRs Merged 1](https://github.com/Unstructured-IO/unstructured/pulls?q=is%3Apr+Explanation+and+details+in%3Abody+is%3Amerged), [PRs Merged 2)](https://github.com/Unstructured-IO/unstructured-ingest/pulls?q=is%3Apr+Explanation+and+details+in%3Abody+is%3Amerged), **Langflow** [(PRs Merged)](https://github.com/langflow-ai/langflow/issues?q=state%3Aclosed%20is%3Apr%20author%3Amisrasaurabh1) and many others to ship performant, expert level code.
 
 Codeflash is great at optimizing AI Agents, Computer Vision algorithms, PyTorch code, numerical code, backend code or anything else you might write with Python.
 

codeflash/LICENSE

@@ -3,7 +3,7 @@ Business Source License 1.1
 Parameters
 
 Licensor:             CodeFlash Inc.
-Licensed Work:        Codeflash Client version 0.16.x
+Licensed Work:        Codeflash Client version 0.17.x
                       The Licensed Work is (c) 2024 CodeFlash Inc.
 
 Additional Use Grant: None. Production use of the Licensed Work is only permitted
@@ -13,7 +13,7 @@ Additional Use Grant: None. Production use of the Licensed Work is only permitte
                       Platform. Please visit codeflash.ai for further
                       information.
 
-Change Date:          2029-08-14
+Change Date:          2029-09-23
 
 Change License:       MIT
 

codeflash/api/aiservice.py

@@ -10,6 +10,7 @@
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
+from codeflash.code_utils.config_consts import N_CANDIDATES_EFFECTIVE, N_CANDIDATES_LP_EFFECTIVE
 from codeflash.code_utils.env_utils import get_codeflash_api_key
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
 from codeflash.lsp.helpers import is_LSP_enabled
@@ -131,6 +132,7 @@ def optimize_python_code(  # noqa: D417
             "current_username": get_last_commit_author_if_pr_exists(None),
             "repo_owner": git_repo_owner,
             "repo_name": git_repo_name,
+            "n_candidates": N_CANDIDATES_EFFECTIVE,
         }
 
         logger.info("!lsp|Generating optimized candidates…")
@@ -192,6 +194,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
             "experiment_metadata": experiment_metadata,
             "codeflash_version": codeflash_version,
             "lsp_mode": is_LSP_enabled(),
+            "n_candidates_lp": N_CANDIDATES_LP_EFFECTIVE,
         }
 
         console.rule()

codeflash/code_utils/config_consts.py

@@ -11,3 +11,25 @@
 MIN_TESTCASE_PASSED_THRESHOLD = 6
 REPEAT_OPTIMIZATION_PROBABILITY = 0.1
 DEFAULT_IMPORTANCE_THRESHOLD = 0.001
+N_CANDIDATES_LP = 6
+
+# LSP-specific
+N_CANDIDATES_LSP = 3
+N_TESTS_TO_GENERATE_LSP = 2
+TOTAL_LOOPING_TIME_LSP = 10.0  # Kept same timing for LSP mode to avoid in increase in performance reporting
+N_CANDIDATES_LP_LSP = 3
+
+MAX_N_CANDIDATES = 5
+MAX_N_CANDIDATES_LP = 6
+
+try:
+    from codeflash.lsp.helpers import is_LSP_enabled
+
+    _IS_LSP_ENABLED = is_LSP_enabled()
+except ImportError:
+    _IS_LSP_ENABLED = False
+
+N_CANDIDATES_EFFECTIVE = min(N_CANDIDATES_LSP if _IS_LSP_ENABLED else N_CANDIDATES, MAX_N_CANDIDATES)
+N_CANDIDATES_LP_EFFECTIVE = min(N_CANDIDATES_LP_LSP if _IS_LSP_ENABLED else N_CANDIDATES_LP, MAX_N_CANDIDATES_LP)
+N_TESTS_TO_GENERATE_EFFECTIVE = N_TESTS_TO_GENERATE_LSP if _IS_LSP_ENABLED else N_TESTS_TO_GENERATE
+TOTAL_LOOPING_TIME_EFFECTIVE = TOTAL_LOOPING_TIME_LSP if _IS_LSP_ENABLED else TOTAL_LOOPING_TIME

codeflash/code_utils/env_utils.py

@@ -42,6 +42,7 @@ def get_codeflash_api_key() -> str:
     if env_api_key and not shell_api_key:
         try:
             from codeflash.either import is_successful
+
             result = save_api_key_to_rc(env_api_key)
             if is_successful(result):
                 logger.debug(f"Automatically saved API key from environment to shell config: {result.unwrap()}")

codeflash/code_utils/git_utils.py

@@ -16,7 +16,7 @@
 from unidiff import PatchSet
 
 from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.config_consts import N_CANDIDATES
+from codeflash.code_utils.config_consts import N_CANDIDATES_EFFECTIVE
 
 if TYPE_CHECKING:
     from git import Repo
@@ -164,7 +164,7 @@ def create_git_worktrees(
 ) -> tuple[Path | None, list[Path]]:
     if git_root and worktree_root_dir:
         worktree_root = Path(tempfile.mkdtemp(dir=worktree_root_dir))
-        worktrees = [Path(tempfile.mkdtemp(dir=worktree_root)) for _ in range(N_CANDIDATES + 1)]
+        worktrees = [Path(tempfile.mkdtemp(dir=worktree_root)) for _ in range(N_CANDIDATES_EFFECTIVE + 1)]
         for worktree in worktrees:
             subprocess.run(["git", "worktree", "add", "-d", worktree], cwd=module_root, check=True)
     else:

codeflash/lsp/beta.py

@@ -21,6 +21,7 @@
 from codeflash.lsp.server import CodeflashLanguageServer, CodeflashLanguageServerProtocol
 from codeflash.lsp.service.perform_optimization import sync_perform_optimization
 
+
 if TYPE_CHECKING:
     from argparse import Namespace
 
@@ -45,6 +46,13 @@ class ProvideApiKeyParams:
     api_key: str
 
 
+@dataclass
+class ValidateProjectParams:
+    root_path_abs: str
+    config_file: Optional[str] = None
+    skip_validation: bool = False
+
+
 @dataclass
 class OnPatchAppliedParams:
     patch_id: str
@@ -55,7 +63,8 @@ class OptimizableFunctionsInCommitParams:
     commit_hash: str
 
 
-server = CodeflashLanguageServer("codeflash-language-server", "v1.0", protocol_cls=CodeflashLanguageServerProtocol)
+# server = CodeflashLanguageServer("codeflash-language-server", "v1.0", protocol_cls=CodeflashLanguageServerProtocol)
+server = CodeflashLanguageServer("codeflash-language-server", "v1.0")
 
 
 @server.feature("getOptimizableFunctionsInCurrentDiff")
@@ -155,17 +164,60 @@ def initialize_function_optimization(
     return {"functionName": params.functionName, "status": "success"}
 
 
-@server.feature("validateProject")
-def validate_project(server: CodeflashLanguageServer, _params: FunctionOptimizationParams) -> dict[str, str]:
+def _find_pyproject_toml(workspace_path: str) -> Path | None:
+    workspace_path_obj = Path(workspace_path)
+    max_depth = 2
+    base_depth = len(workspace_path_obj.parts)
+
+    for root, dirs, files in os.walk(workspace_path_obj):
+        depth = len(Path(root).parts) - base_depth
+        if depth > max_depth:
+            # stop going deeper into this branch
+            dirs.clear()
+            continue
+
+        if "pyproject.toml" in files:
+            file_path = Path(root) / "pyproject.toml"
+            with file_path.open("r", encoding="utf-8", errors="ignore") as f:
+                for line in f:
+                    if line.strip() == "[tool.codeflash]":
+                        return file_path.resolve()
+    return None
+
+
+# should be called the first thing to initialize and validate the project
+@server.feature("initProject")
+def init_project(server: CodeflashLanguageServer, params: ValidateProjectParams) -> dict[str, str]:
     from codeflash.cli_cmds.cmd_init import is_valid_pyproject_toml
 
+    pyproject_toml_path: Path | None = getattr(params, "config_file", None)
+
+    if server.args is None:
+        if pyproject_toml_path is not None:
+            # if there is a config file provided use it
+            server.prepare_optimizer_arguments(pyproject_toml_path)
+        else:
+            # otherwise look for it
+            pyproject_toml_path = _find_pyproject_toml(params.root_path_abs)
+            server.show_message_log(f"Found pyproject.toml at: {pyproject_toml_path}", "Info")
+            if pyproject_toml_path:
+                server.prepare_optimizer_arguments(pyproject_toml_path)
+            else:
+                return {
+                    "status": "error",
+                    "message": "No pyproject.toml found in workspace.",
+                }  # TODO: enhancec this message to say there is not tool.codeflash in pyproject.toml or smth
+
+    if getattr(params, "skip_validation", False):
+        return {"status": "success", "moduleRoot": server.args.module_root, "pyprojectPath": pyproject_toml_path}
+
     server.show_message_log("Validating project...", "Info")
-    config = is_valid_pyproject_toml(server.args.config_file)
+    config = is_valid_pyproject_toml(pyproject_toml_path)
     if config is None:
         server.show_message_log("pyproject.toml is not valid", "Error")
         return {
             "status": "error",
-            "message": "pyproject.toml is not valid",  # keep the error message the same, the extension is matching "pyproject.toml" in the error message to show the codeflash init instructions
+            "message": "pyproject.toml is not valid",  # keep the error message the same, the extension is matching "pyproject.toml" in the error message to show the codeflash init instructions,
         }
 
     args = process_args(server)
@@ -178,7 +230,7 @@ def validate_project(server: CodeflashLanguageServer, _params: FunctionOptimizat
     except Exception:
         return {"status": "error", "message": "Repository has no commits (unborn HEAD)"}
 
-    return {"status": "success", "moduleRoot": args.module_root}
+    return {"status": "success", "moduleRoot": args.module_root, "pyprojectPath": pyproject_toml_path}
 
 
 def _initialize_optimizer_if_api_key_is_valid(

codeflash/lsp/lsp_logger.py

@@ -124,7 +124,7 @@ def setup_logging() -> logging.Logger:
     logger = logging.getLogger()
     logger.handlers.clear()
 
-    # Set up stderr handler for VS Code output channel with [LSP-Server] prefix
+    # Set up stderr handler for VS Code output channel
     handler = logging.StreamHandler(sys.stderr)
     handler.setLevel(logging.DEBUG)
 

codeflash/lsp/server.py

@@ -1,46 +1,20 @@
 from __future__ import annotations
 
-from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-from lsprotocol.types import INITIALIZE, LogMessageParams, MessageType
-from pygls import uris
-from pygls.protocol import LanguageServerProtocol, lsp_method
+from lsprotocol.types import LogMessageParams, MessageType
+from pygls.protocol import LanguageServerProtocol
 from pygls.server import LanguageServer
 
 if TYPE_CHECKING:
-    from lsprotocol.types import InitializeParams, InitializeResult
+    from pathlib import Path
 
     from codeflash.optimization.optimizer import Optimizer
 
 
 class CodeflashLanguageServerProtocol(LanguageServerProtocol):
     _server: CodeflashLanguageServer
 
-    @lsp_method(INITIALIZE)
-    def lsp_initialize(self, params: InitializeParams) -> InitializeResult:
-        server = self._server
-        initialize_result: InitializeResult = super().lsp_initialize(params)
-
-        workspace_uri = params.root_uri
-        if workspace_uri:
-            workspace_path = uris.to_fs_path(workspace_uri)
-            pyproject_toml_path = self._find_pyproject_toml(workspace_path)
-            if pyproject_toml_path:
-                server.prepare_optimizer_arguments(pyproject_toml_path)
-            else:
-                server.show_message("No pyproject.toml found in workspace.")
-        else:
-            server.show_message("No workspace URI provided.")
-
-        return initialize_result
-
-    def _find_pyproject_toml(self, workspace_path: str) -> Path | None:
-        workspace_path_obj = Path(workspace_path)
-        for file_path in workspace_path_obj.rglob("pyproject.toml"):
-            return file_path.resolve()
-        return None
-
 
 class CodeflashLanguageServer(LanguageServer):
     def __init__(self, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401