From 88af698664bb30ee08c5c79d6cc2cd419329a8e6 Mon Sep 17 00:00:00 2001 From: ahmetgunduz Date: Thu, 18 Sep 2025 18:32:16 +0300 Subject: [PATCH] fix: Centralize API key validation and fix misleading error messages Addresses confusing behavior where TeamAgentFactory showed misleading errors suggesting only TEAM_API_KEY was acceptable, when AIXPLAIN_API_KEY should also work. - Created centralized validation functions in config.py (single source of truth) - Updated api_key_checker.py to use centralized logic - Updated v2/core.py to accept either API key - Fixed error messages to indicate both keys are acceptable - Users can now use either AIXPLAIN_API_KEY or TEAM_API_KEY - Clear, consistent error messages across SDK - Backward compatible, no breaking changes Before: "TEAM_API_KEY has not been set properly" (misleading) After: "Neither 'AIXPLAIN_API_KEY' nor 'TEAM_API_KEY' has been set. Please set either environment variable." --- aixplain/decorators/api_key_checker.py | 15 +++---- aixplain/utils/config.py | 54 ++++++++++++++++++++------ aixplain/v2/core.py | 20 +++++++--- 3 files changed, 64 insertions(+), 25 deletions(-) diff --git a/aixplain/decorators/api_key_checker.py b/aixplain/decorators/api_key_checker.py index b4e7cf84..21fc72b7 100644 --- a/aixplain/decorators/api_key_checker.py +++ b/aixplain/decorators/api_key_checker.py @@ -1,11 +1,13 @@ -from aixplain.utils import config +"""API key validation decorator for aiXplain SDK.""" + +from aixplain.utils.config import check_api_keys_available def check_api_key(method): """Decorator to verify that an API key is set before executing the method. - This decorator checks if either TEAM_API_KEY or AIXPLAIN_API_KEY is set in the - configuration. If neither key is set, it raises an exception. + This decorator uses the centralized API key validation logic from config.py + to ensure consistent behavior across the entire SDK. Args: method (callable): The method to be decorated. @@ -22,11 +24,10 @@ def my_api_method(): # Method implementation pass """ + def wrapper(*args, **kwargs): - if config.TEAM_API_KEY == "" and config.AIXPLAIN_API_KEY == "": - raise Exception( - "A 'TEAM_API_KEY' is required to run an asset. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)" - ) + # Use centralized validation - single source of truth + check_api_keys_available() return method(*args, **kwargs) return wrapper diff --git a/aixplain/utils/config.py b/aixplain/utils/config.py index 6cb15775..fa3bc7aa 100644 --- a/aixplain/utils/config.py +++ b/aixplain/utils/config.py @@ -1,5 +1,4 @@ -""" -Copyright 2022 The aiXplain SDK authors +"""Copyright 2022 The aiXplain SDK authors. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -28,20 +27,51 @@ ENV = "dev" if "dev" in BACKEND_URL else "test" if "test" in BACKEND_URL else "prod" -if not TEAM_API_KEY and not AIXPLAIN_API_KEY: - raise Exception( - "'TEAM_API_KEY' has not been set properly and is empty. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)" - ) +def validate_api_keys(): + """Centralized API key validation function - single source of truth. -if AIXPLAIN_API_KEY and TEAM_API_KEY and AIXPLAIN_API_KEY != TEAM_API_KEY: - raise Exception( - "Conflicting API keys: 'AIXPLAIN_API_KEY' and 'TEAM_API_KEY' are both provided but do not match. Please provide only one API key." - ) + This function handles all API key validation logic: + 1. Ensures at least one API key is provided + 2. Prevents conflicting API keys + 3. Auto-normalizes AIXPLAIN_API_KEY to TEAM_API_KEY if needed + + Raises: + Exception: If no API keys are provided or if conflicting keys are detected + """ + global TEAM_API_KEY, AIXPLAIN_API_KEY + + if not TEAM_API_KEY and not AIXPLAIN_API_KEY: + raise Exception( + "Neither 'AIXPLAIN_API_KEY' nor 'TEAM_API_KEY' has been set. Please set either environment variable. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)" + ) + + if AIXPLAIN_API_KEY and TEAM_API_KEY and AIXPLAIN_API_KEY != TEAM_API_KEY: + raise Exception( + "Conflicting API keys: 'AIXPLAIN_API_KEY' and 'TEAM_API_KEY' are both provided but do not match. Please provide only one API key." + ) + + if AIXPLAIN_API_KEY and not TEAM_API_KEY: + TEAM_API_KEY = AIXPLAIN_API_KEY + + +def check_api_keys_available(): + """Runtime check to ensure API keys are available. + + This is used by decorators and other runtime validation. + Uses the same validation logic as the module-level check. + + Raises: + Exception: If no valid API keys are available + """ + if not TEAM_API_KEY and not AIXPLAIN_API_KEY: + raise Exception( + "An API key is required to run this operation. Please set either 'AIXPLAIN_API_KEY' or 'TEAM_API_KEY'. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)" + ) -if AIXPLAIN_API_KEY and not TEAM_API_KEY: - TEAM_API_KEY = AIXPLAIN_API_KEY +# Perform initial validation at module import time +validate_api_keys() PIPELINE_API_KEY = os.getenv("PIPELINE_API_KEY", "") MODEL_API_KEY = os.getenv("MODEL_API_KEY", "") diff --git a/aixplain/v2/core.py b/aixplain/v2/core.py index 209d87e7..e8190dd2 100644 --- a/aixplain/v2/core.py +++ b/aixplain/v2/core.py @@ -1,3 +1,5 @@ +"""Core module for aiXplain v2 API.""" + import os from typing import TypeVar from .client import AixplainClient @@ -98,8 +100,8 @@ class Aixplain: PIPELINES_RUN_URL = "https://platform-api.aixplain.com/assets/pipeline/execution/run" def __new__(cls, *args, **kwargs): - """ - Singleton pattern for the Aixplain class. + """Singleton pattern for the Aixplain class. + Otherwise, the environment variables will be overwritten in multiple instances. TODO: This should be removed once the factory classes are removed. @@ -124,10 +126,16 @@ def __init__( pipeline_url: str: The URL for the pipeline. model_url: str: The URL for the model. """ - self.api_key = api_key or os.getenv("TEAM_API_KEY") - assert ( - self.api_key - ), "API key is required. You should either pass it as an argument or set the TEAM_API_KEY environment variable." + # Use centralized API key validation logic + from aixplain.utils.config import ( + check_api_keys_available, + TEAM_API_KEY, + AIXPLAIN_API_KEY, + ) + + self.api_key = api_key or TEAM_API_KEY or AIXPLAIN_API_KEY + if not self.api_key: + check_api_keys_available() # This will raise the proper centralized error message self.base_url = backend_url or os.getenv("BACKEND_URL") or self.BACKEND_URL self.pipeline_url = pipeline_url or os.getenv("PIPELINES_RUN_URL") or self.PIPELINES_RUN_URL