Merge branch 'main' of github.com:codeflash-ai/codeflash into feat/markdown-read-writable-context

Author: mohammedahmed18

.pre-commit-config.yaml

@@ -1,7 +1,8 @@
 repos:
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.11.0"
-    hooks:
-      - id: ruff
-        args: [--fix, --exit-non-zero-on-fix, --config=pyproject.toml]
-      - id: ruff-format
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.12.7
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+    # Run the formatter.
+    - id: ruff-format

codeflash/api/cfapi.py

@@ -101,6 +101,9 @@ def get_user_id() -> Optional[str]:
             if min_version and version.parse(min_version) > version.parse(__version__):
                 msg = "Your Codeflash CLI version is outdated. Please update to the latest version using `pip install --upgrade codeflash`."
                 console.print(f"[bold red]{msg}[/bold red]")
+                if console.quiet:  # lsp
+                    logger.debug(msg)
+                    return f"Error: {msg}"
                 sys.exit(1)
             return userid
 

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -32,12 +32,14 @@ def leave_ClassDef(
     def visit_ClassDef(self, node: ClassDef) -> Optional[bool]:
         if self.class_name:  # Don't go into nested class
             return False
-        self.class_name = node.name.value  # noqa: RET503
+        self.class_name = node.name.value
+        return None
 
     def visit_FunctionDef(self, node: FunctionDef) -> Optional[bool]:
         if self.function_name:  # Don't go into nested function
             return False
-        self.function_name = node.name.value  # noqa: RET503
+        self.function_name = node.name.value
+        return None
 
     def leave_FunctionDef(self, original_node: FunctionDef, updated_node: FunctionDef) -> FunctionDef:
         if self.function_name == original_node.name.value:

codeflash/cli_cmds/cmd_init.py

@@ -38,7 +38,7 @@
     from argparse import Namespace
 
 CODEFLASH_LOGO: str = (
-    f"{LF}"  # noqa: ISC003
+    f"{LF}"
     r"                   _          ___  _               _     " + f"{LF}"
     r"                  | |        / __)| |             | |    " + f"{LF}"
     r"  ____   ___    _ | |  ____ | |__ | |  ____   ___ | | _  " + f"{LF}"

codeflash/code_utils/env_utils.py

@@ -34,7 +34,12 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool =
 
 @lru_cache(maxsize=1)
 def get_codeflash_api_key() -> str:
-    api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+    if console.quiet:  # lsp
+        # prefer shell config over env var in lsp mode
+        api_key = read_api_key_from_shell_config()
+    else:
+        api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+
     api_secret_docs_message = "For more information, refer to the documentation at [https://docs.codeflash.ai/getting-started/codeflash-github-actions#add-your-api-key-to-your-repository-secrets]."  # noqa
     if not api_key:
         msg = (

codeflash/code_utils/shell_utils.py

@@ -42,7 +42,7 @@ def get_shell_rc_path() -> Path:
 
 
 def get_api_key_export_line(api_key: str) -> str:
-    return f"{SHELL_RC_EXPORT_PREFIX}{api_key}"
+    return f'{SHELL_RC_EXPORT_PREFIX}"{api_key}"'
 
 
 def save_api_key_to_rc(api_key: str) -> Result[str, str]:

codeflash/lsp/beta.py

@@ -8,6 +8,8 @@
 
 from pygls import uris
 
+from codeflash.api.cfapi import get_codeflash_api_key, get_user_id
+from codeflash.code_utils.shell_utils import save_api_key_to_rc
 from codeflash.either import is_successful
 from codeflash.lsp.server import CodeflashLanguageServer, CodeflashLanguageServerProtocol
 
@@ -28,6 +30,11 @@ class FunctionOptimizationParams:
     functionName: str  # noqa: N815
 
 
+@dataclass
+class ProvideApiKeyParams:
+    api_key: str
+
+
 server = CodeflashLanguageServer("codeflash-language-server", "v1.0", protocol_cls=CodeflashLanguageServerProtocol)
 
 
@@ -118,6 +125,53 @@ def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOpt
     return {"functionName": params.functionName, "status": "success", "discovered_tests": num_discovered_tests}
 
 
+def _initialize_optimizer_if_valid(server: CodeflashLanguageServer) -> dict[str, str]:
+    user_id = get_user_id()
+    if user_id is None:
+        return {"status": "error", "message": "api key not found or invalid"}
+
+    if user_id.startswith("Error: "):
+        error_msg = user_id[7:]
+        return {"status": "error", "message": error_msg}
+
+    from codeflash.optimization.optimizer import Optimizer
+
+    server.optimizer = Optimizer(server.args)
+    return {"status": "success", "user_id": user_id}
+
+
+@server.feature("apiKeyExistsAndValid")
+def check_api_key(server: CodeflashLanguageServer, _params: any) -> dict[str, str]:
+    try:
+        return _initialize_optimizer_if_valid(server)
+    except Exception:
+        return {"status": "error", "message": "something went wrong while validating the api key"}
+
+
+@server.feature("provideApiKey")
+def provide_api_key(server: CodeflashLanguageServer, params: ProvideApiKeyParams) -> dict[str, str]:
+    try:
+        api_key = params.api_key
+        if not api_key.startswith("cf-"):
+            return {"status": "error", "message": "Api key is not valid"}
+
+        result = save_api_key_to_rc(api_key)
+        if not is_successful(result):
+            return {"status": "error", "message": result.failure()}
+
+        # clear cache to ensure the new api key is used
+        get_codeflash_api_key.cache_clear()
+        get_user_id.cache_clear()
+
+        init_result = _initialize_optimizer_if_valid(server)
+        if init_result["status"] == "error":
+            return {"status": "error", "message": "Api key is not valid"}
+
+        return {"status": "success", "message": "Api key saved successfully", "user_id": init_result["user_id"]}
+    except Exception:
+        return {"status": "error", "message": "something went wrong while saving the api key"}
+
+
 @server.feature("prepareOptimization")
 def prepare_optimization(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]:
     current_function = server.optimizer.current_function_being_optimized

codeflash/lsp/server.py

@@ -25,7 +25,7 @@ def lsp_initialize(self, params: InitializeParams) -> InitializeResult:
             workspace_path = uris.to_fs_path(workspace_uri)
             pyproject_toml_path = self._find_pyproject_toml(workspace_path)
             if pyproject_toml_path:
-                server.initialize_optimizer(pyproject_toml_path)
+                server.prepare_optimizer_arguments(pyproject_toml_path)
                 server.show_message(f"Found pyproject.toml at: {pyproject_toml_path}")
             else:
                 server.show_message("No pyproject.toml found in workspace.")
@@ -45,16 +45,17 @@ class CodeflashLanguageServer(LanguageServer):
     def __init__(self, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
         super().__init__(*args, **kwargs)
         self.optimizer = None
+        self.args = None
 
-    def initialize_optimizer(self, config_file: Path) -> None:
+    def prepare_optimizer_arguments(self, config_file: Path) -> None:
         from codeflash.cli_cmds.cli import parse_args, process_pyproject_config
-        from codeflash.optimization.optimizer import Optimizer
 
         args = parse_args()
         args.config_file = config_file
         args.no_pr = True  # LSP server should not create PRs
         args = process_pyproject_config(args)
-        self.optimizer = Optimizer(args)
+        self.args = args
+        # avoid initializing the optimizer during initialization, because it can cause an error if the api key is invalid
 
     def show_message_log(self, message: str, message_type: str) -> None:
         """Send a log message to the client's output channel.
@@ -70,6 +71,7 @@ def show_message_log(self, message: str, message_type: str) -> None:
             "Warning": MessageType.Warning,
             "Error": MessageType.Error,
             "Log": MessageType.Log,
+            "Debug": MessageType.Debug,
         }
 
         lsp_message_type = type_mapping.get(message_type, MessageType.Info)

codeflash/models/models.py

@@ -552,7 +552,7 @@ def unique_invocation_loop_id(self) -> str:
         return f"{self.loop_index}:{self.id.id()}"
 
 
-class TestResults(BaseModel):
+class TestResults(BaseModel):  # noqa: PLW1641
     # don't modify these directly, use the add method
     # also we don't support deletion of test results elements - caution is advised
     test_results: list[FunctionTestInvocation] = []

codeflash/result/critic.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils import env_utils
@@ -29,7 +29,8 @@ def speedup_critic(
     candidate_result: OptimizedCandidateResult,
     original_code_runtime: int,
     best_runtime_until_now: int | None,
-    disable_gh_action_noise: Optional[bool] = None,
+    *,
+    disable_gh_action_noise: bool = False,
 ) -> bool:
     """Take in a correct optimized Test Result and decide if the optimization should actually be surfaced to the user.
 
@@ -39,10 +40,8 @@ def speedup_critic(
     The noise floor is doubled when benchmarking on a (noisy) GitHub Action virtual instance, also we want to be more confident there.
     """
     noise_floor = 3 * MIN_IMPROVEMENT_THRESHOLD if original_code_runtime < 10000 else MIN_IMPROVEMENT_THRESHOLD
-    if not disable_gh_action_noise:
-        in_github_actions_mode = bool(env_utils.get_pr_number())
-        if in_github_actions_mode:
-            noise_floor = noise_floor * 2  # Increase the noise floor in GitHub Actions mode
+    if not disable_gh_action_noise and env_utils.is_ci():
+        noise_floor = noise_floor * 2  # Increase the noise floor in GitHub Actions mode
 
     perf_gain = performance_gain(
         original_runtime_ns=original_code_runtime, optimized_runtime_ns=candidate_result.best_test_runtime