Revert to state of branch main

Author: mohamedashrraf222

codeflash/api/aiservice.py

@@ -32,15 +32,49 @@
 
 class AiServiceClient:
     def __init__(self) -> None:
+        # API key validation is deferred to first use (in the headers property)
         self.base_url = self.get_aiservice_base_url()
-        self.headers = {"Authorization": f"Bearer {get_codeflash_api_key()}", "Connection": "close"}
-
+        
     def get_aiservice_base_url(self) -> str:
         if os.environ.get("CODEFLASH_AIS_SERVER", default="prod").lower() == "local":
             logger.info("Using local AI Service at http://localhost:8000")
             console.rule()
             return "http://localhost:8000"
         return "https://app.codeflash.ai"
+    
+    def validate_api_key(self) -> None:
+        """Validate API key. Raises OSError if invalid."""
+        if hasattr(self, '_api_key_validated') and self._api_key_validated:
+            return
+            
+        try:
+            from codeflash.api.cfapi import get_user_id  # noqa: PLC0415
+            
+            # Try to get user ID, if it fails the API key is invalid
+            user_id = get_user_id()
+            if user_id is None:
+                msg = (
+                    "Invalid Codeflash API key. The API key you provided is not valid.\n"
+                    "Please generate a new one at https://app.codeflash.ai/app/apikeys ,\n"
+                    "then set it as a CODEFLASH_API_KEY environment variable.\n"
+                    "For more information, refer to the documentation at \n"
+                    "https://docs.codeflash.ai/optimizing-with-codeflash/codeflash-github-actions#manual-setup\n"
+                    "or\n"
+                    "https://docs.codeflash.ai/optimizing-with-codeflash/codeflash-github-actions#automated-setup-recommended"
+                )
+                raise OSError(msg)
+        except ImportError:
+            # If cfapi is not available, skip validation
+            pass
+        
+        self._api_key_validated = True
+    
+    @property
+    def headers(self) -> dict[str, str]:
+        """Get headers with API key. Validates API key on first access."""
+        # Lazily validate API key on first use
+        self.validate_api_key()
+        return {"Authorization": f"Bearer {get_codeflash_api_key()}", "Connection": "close"}
 
     def make_ai_service_request(
         self,

codeflash/api/cfapi.py

@@ -89,6 +89,7 @@ def make_cfapi_request(
 def get_user_id(api_key: Optional[str] = None) -> Optional[str]:
     """Retrieve the user's userid by making a request to the /cfapi/cli-get-user endpoint.
 
+    :param api_key: The API key to use. If None, uses get_codeflash_api_key().
     :return: The userid or None if the request fails.
     """
     if not api_key and not ensure_codeflash_api_key():
@@ -119,7 +120,6 @@ def get_user_id(api_key: Optional[str] = None) -> Optional[str]:
     logger.error(f"Failed to look up your userid; is your CF API key valid? ({response.reason})")
     return None
 
-
 def suggest_changes(
     owner: str,
     repo: str,

codeflash/optimization/function_optimizer.py

@@ -213,7 +213,7 @@ def __init__(
     ) -> None:
         self.project_root = test_cfg.project_root_path
         self.test_cfg = test_cfg
-        self.aiservice_client = aiservice_client if aiservice_client else AiServiceClient()
+        self._aiservice_client = aiservice_client  # Can be None for lazy initialization
         self.function_to_optimize = function_to_optimize
         self.function_to_optimize_source_code = (
             function_to_optimize_source_code
@@ -247,6 +247,13 @@ def __init__(
             max_workers=n_tests + 2 if self.experiment_id is None else n_tests + 3
         )
 
+    @property
+    def aiservice_client(self) -> AiServiceClient:
+        """Lazy initialization of AiServiceClient to delay API key validation."""
+        if self._aiservice_client is None:
+            self._aiservice_client = AiServiceClient()
+        return self._aiservice_client
+
     def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]:
         should_run_experiment = self.experiment_id is not None
         logger.debug(f"Function Trace ID: {self.function_trace_id}")

codeflash/optimization/optimizer.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import contextlib
 import copy
 import os
 import tempfile
@@ -49,7 +50,7 @@ def __init__(self, args: Namespace) -> None:
             benchmark_tests_root=args.benchmarks_root if "benchmark" in args and "benchmarks_root" in args else None,
         )
 
-        self.aiservice_client = AiServiceClient()
+        self._aiservice_client: AiServiceClient | None = None
         self.experiment_id = os.getenv("CODEFLASH_EXPERIMENT_ID", None)
         self.local_aiservice_client = LocalAiServiceClient() if self.experiment_id else None
         self.replay_tests_dir = None
@@ -60,6 +61,13 @@ def __init__(self, args: Namespace) -> None:
         self.original_args_and_test_cfg: tuple[Namespace, TestConfig] | None = None
         self.patch_files: list[Path] = []
 
+    @property
+    def aiservice_client(self) -> AiServiceClient:
+        """Lazy initialization of AiServiceClient to delay API key validation."""
+        if self._aiservice_client is None:
+            self._aiservice_client = AiServiceClient()
+        return self._aiservice_client
+
     def run_benchmarks(
         self, file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]], num_optimizable_functions: int
     ) -> tuple[dict[str, dict[BenchmarkKey, float]], dict[BenchmarkKey, float]]:
@@ -260,6 +268,10 @@ def run(self) -> None:
         console.rule()
         if not env_utils.ensure_codeflash_api_key():
             return
+        
+        # Validate API key before starting optimization
+        self.aiservice_client.validate_api_key()
+        
         if self.args.no_draft and is_pr_draft():
             logger.warning("PR is in draft mode, skipping optimization")
             return
@@ -499,5 +511,13 @@ def run_with_args(args: Namespace) -> None:
         logger.warning("Keyboard interrupt received. Cleaning up and exiting, please wait…")
         if optimizer:
             optimizer.cleanup_temporary_paths()
-
-        raise SystemExit from None
+        raise SystemExit(0) from None
+    except (OSError, Exception) as e:
+        # Log the error if we have an optimizer instance with cleanup
+        if optimizer:
+            logger.error(f"Error during optimization: {e}")
+            with contextlib.suppress(Exception):
+                optimizer.cleanup_temporary_paths()
+        else:
+            logger.error(f"Error initializing optimizer: {e}")
+        raise