Merge branch 'main' into opt-impact-aseem

Author: aseembits93

codeflash/api/aiservice.py

@@ -11,6 +11,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.code_utils.time_utils import humanize_runtime
@@ -132,6 +133,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…")
@@ -193,6 +195,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/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

@@ -24,7 +24,7 @@
     get_functions_within_git_diff,
 )
 from codeflash.either import is_successful
-from codeflash.lsp.server import CodeflashLanguageServer, CodeflashLanguageServerProtocol
+from codeflash.lsp.server import CodeflashLanguageServer
 
 if TYPE_CHECKING:
     from argparse import Namespace
@@ -50,6 +50,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
@@ -60,7 +67,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")
@@ -160,17 +168,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)
@@ -183,7 +234,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(
@@ -328,7 +379,7 @@ def perform_function_optimization(  # noqa: PLR0911
 
         devnull_writer = open(os.devnull, "w")  # noqa
         with contextlib.redirect_stdout(devnull_writer):
-            function_to_tests, num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs)
+            function_to_tests, _num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs)
             function_optimizer.function_to_tests = function_to_tests
 
         test_setup_result = function_optimizer.generate_and_instrument_tests(

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