diff --git a/codeflash/api/cfapi.py b/codeflash/api/cfapi.py index 6bf0448c3..a27e85bff 100644 --- a/codeflash/api/cfapi.py +++ b/codeflash/api/cfapi.py @@ -91,7 +91,7 @@ def get_user_id(api_key: Optional[str] = None) -> Optional[str]: :return: The userid or None if the request fails. """ - if not ensure_codeflash_api_key(): + if not api_key and not ensure_codeflash_api_key(): return None response = make_cfapi_request( diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index 67b4bacd2..c54ceaafd 100644 --- a/codeflash/lsp/beta.py +++ b/codeflash/lsp/beta.py @@ -263,6 +263,15 @@ def init_project(server: CodeflashLanguageServer, params: ValidateProjectParams) def _initialize_optimizer_if_api_key_is_valid( server: CodeflashLanguageServer, api_key: Optional[str] = None ) -> dict[str, str]: + key_check_result = _check_api_key_validity(api_key) + if key_check_result.get("status") != "success": + return key_check_result + + _initialize_optimizer(server) + return key_check_result + + +def _check_api_key_validity(api_key: Optional[str]) -> dict[str, str]: user_id = get_user_id(api_key=api_key) if user_id is None: return {"status": "error", "message": "api key not found or invalid"} @@ -271,11 +280,15 @@ def _initialize_optimizer_if_api_key_is_valid( error_msg = user_id[7:] return {"status": "error", "message": error_msg} + return {"status": "success", "user_id": user_id} + + +def _initialize_optimizer(server: CodeflashLanguageServer) -> None: from codeflash.optimization.optimizer import Optimizer new_args = process_args(server) - server.optimizer = Optimizer(new_args) - return {"status": "success", "user_id": user_id} + if not server.optimizer: + server.optimizer = Optimizer(new_args) def process_args(server: CodeflashLanguageServer) -> Namespace: @@ -302,16 +315,16 @@ def provide_api_key(server: CodeflashLanguageServer, params: ProvideApiKeyParams if not api_key.startswith("cf-"): return {"status": "error", "message": "Api key is not valid"} - # clear cache to ensure the new api key is used + # # 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_api_key_is_valid(server, api_key) - if init_result["status"] == "error": - return {"status": "error", "message": "Api key is not valid"} - - user_id = init_result["user_id"] + key_check_result = _check_api_key_validity(api_key) + if key_check_result.get("status") != "success": + return key_check_result + user_id = key_check_result["user_id"] result = save_api_key_to_rc(api_key) + # initialize optimizer with the new api key + _initialize_optimizer(server) if not is_successful(result): return {"status": "error", "message": result.failure()} return {"status": "success", "message": "Api key saved successfully", "user_id": user_id} # noqa: TRY300