From 7bef094720df4635bbf2bb9fbf73350c88d09f46 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Tue, 15 Oct 2024 09:59:23 -0300 Subject: [PATCH 01/10] assetLimits -> assetsLimits --- aixplain/factories/api_key_factory.py | 2 +- aixplain/modules/api_key.py | 4 ++-- tests/unit/api_key_test.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/aixplain/factories/api_key_factory.py b/aixplain/factories/api_key_factory.py index 4ac8f00a..a0b0205b 100644 --- a/aixplain/factories/api_key_factory.py +++ b/aixplain/factories/api_key_factory.py @@ -30,7 +30,7 @@ def list(cls) -> List[APIKey]: name=key["name"], budget=key["budget"] if "budget" in key else None, global_limits=key["globalLimits"] if "globalLimits" in key else None, - asset_limits=key["assetLimits"] if "assetLimits" in key else [], + asset_limits=key["assetsLimits"] if "assetsLimits" in key else [], expires_at=key["expiresAt"] if "expiresAt" in key else None, access_key=key["accessKey"], is_admin=key["isAdmin"], diff --git a/aixplain/modules/api_key.py b/aixplain/modules/api_key.py index 886b0dab..abcc4571 100644 --- a/aixplain/modules/api_key.py +++ b/aixplain/modules/api_key.py @@ -110,7 +110,7 @@ def to_dict(self) -> Dict: "id": self.id, "name": self.name, "budget": self.budget, - "assetLimits": [], + "assetsLimits": [], "expiresAt": self.expires_at, } @@ -126,7 +126,7 @@ def to_dict(self) -> Dict: } for i, asset_limit in enumerate(self.asset_limits): - payload["assetLimits"].append( + payload["assetsLimits"].append( { "tpm": asset_limit.token_per_minute, "tpd": asset_limit.token_per_day, diff --git a/tests/unit/api_key_test.py b/tests/unit/api_key_test.py index 60d2371d..9c8dc1cd 100644 --- a/tests/unit/api_key_test.py +++ b/tests/unit/api_key_test.py @@ -25,7 +25,7 @@ def test_api_key_service(): "accessKey": "access-key", "budget": 1000, "globalLimits": {"tpm": 100, "tpd": 1000, "rpd": 1000, "rpm": 100}, - "assetLimits": [{"assetId": model_id, "tpm": 100, "tpd": 1000, "rpd": 1000, "rpm": 100}], + "assetsLimits": [{"assetId": model_id, "tpm": 100, "tpd": 1000, "rpd": 1000, "rpm": 100}], "expiresAt": "2024-10-07T00:00:00Z", "isAdmin": False, } From f824cf0c5aca3a4607b7d7bb77e1122120f0b128 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Thu, 17 Oct 2024 17:25:32 -0300 Subject: [PATCH 02/10] Update api key usage limit service --- aixplain/factories/api_key_factory.py | 20 ++++++++++------- aixplain/modules/api_key.py | 32 ++++++++++++++++++++------- tests/functional/apikey/test_api.py | 4 +++- 3 files changed, 39 insertions(+), 17 deletions(-) diff --git a/aixplain/factories/api_key_factory.py b/aixplain/factories/api_key_factory.py index a0b0205b..89759fcf 100644 --- a/aixplain/factories/api_key_factory.py +++ b/aixplain/factories/api_key_factory.py @@ -112,8 +112,8 @@ def update(cls, api_key: APIKey) -> APIKey: raise Exception(f"API Key Update Error: Failed to update API key with ID {api_key.id}. Error: {str(resp)}") @classmethod - def get_usage_limit(cls, api_key: Text = config.TEAM_API_KEY, asset_id: Optional[Text] = None) -> APIKeyUsageLimit: - """Get API key usage limit""" + def get_usage_limits(cls, api_key: Text = config.TEAM_API_KEY, asset_id: Optional[Text] = None) -> List[APIKeyUsageLimit]: + """Get API key usage limits""" try: url = f"{config.BACKEND_URL}/sdk/api-keys/usage-limits" if asset_id is not None: @@ -128,11 +128,15 @@ def get_usage_limit(cls, api_key: Text = config.TEAM_API_KEY, asset_id: Optional raise Exception(f"{message}") if 200 <= r.status_code < 300: - return APIKeyUsageLimit( - request_count=resp["requestCount"], - request_count_limit=resp["requestCountLimit"], - token_count=resp["tokenCount"], - token_count_limit=resp["tokenCountLimit"], - ) + return [ + APIKeyUsageLimit( + request_count=limit["requestCount"], + request_count_limit=limit["requestCountLimit"], + token_count=limit["tokenCount"], + token_count_limit=limit["tokenCountLimit"], + model=limit["assetId"] if "assetId" in limit else None, + ) + for limit in resp + ] else: raise Exception(f"API Key Usage Error: Failed to get usage. Error: {str(resp)}") diff --git a/aixplain/modules/api_key.py b/aixplain/modules/api_key.py index abcc4571..b6cc8d3b 100644 --- a/aixplain/modules/api_key.py +++ b/aixplain/modules/api_key.py @@ -27,19 +27,31 @@ def __init__( class APIKeyUsageLimit: - def __init__(self, request_count: int, request_count_limit: int, token_count: int, token_count_limit: int): - """Get the usage limits of an API key + def __init__( + self, + request_count: int, + request_count_limit: int, + token_count: int, + token_count_limit: int, + model: Optional[Union[Text, Model]] = None, + ): + """Get the usage limits of an API key globally (model equals to None) or for a specific model. Args: request_count (int): number of requests made request_count_limit (int): limit of requests token_count (int): number of tokens used token_count_limit (int): limit of tokens + model (Optional[Union[Text, Model]], optional): Model which the limits apply. Defaults to None. """ self.request_count = request_count self.request_count_limit = request_count_limit self.token_count = token_count self.token_count_limit = token_count_limit + if model is not None and isinstance(model, str): + from aixplain.factories import ModelFactory + + self.model = ModelFactory.get(model) class APIKey: @@ -167,11 +179,15 @@ def get_usage(self, asset_id: Optional[Text] = None) -> APIKeyUsageLimit: raise Exception(f"{message}") if 200 <= r.status_code < 300: - return APIKeyUsageLimit( - request_count=resp["requestCount"], - request_count_limit=resp["requestCountLimit"], - token_count=resp["tokenCount"], - token_count_limit=resp["tokenCountLimit"], - ) + return [ + APIKeyUsageLimit( + request_count=limit["requestCount"], + request_count_limit=limit["requestCountLimit"], + token_count=limit["tokenCount"], + token_count_limit=limit["tokenCountLimit"], + model=limit["assetId"] if "assetId" in limit else None, + ) + for limit in resp + ] else: raise Exception(f"API Key Usage Error: Failed to get usage. Error: {str(resp)}") diff --git a/tests/functional/apikey/test_api.py b/tests/functional/apikey/test_api.py index 80b75189..f2a50949 100644 --- a/tests/functional/apikey/test_api.py +++ b/tests/functional/apikey/test_api.py @@ -134,7 +134,9 @@ def test_list_api_keys(): if api_key.is_admin is False: usage = api_key.get_usage() - assert isinstance(usage, APIKeyUsageLimit) + assert isinstance(usage, list) + if len(usage) > 0: + assert isinstance(usage[0], APIKeyUsageLimit) def test_list_update_api_keys(): From a38b882ecab12e1560e7dc131bd14e97f70f177b Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Tue, 22 Oct 2024 15:25:37 -0300 Subject: [PATCH 03/10] Remove option for asset ID filtering in the URLs for getting usage --- aixplain/factories/api_key_factory.py | 3 +-- aixplain/modules/api_key.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/aixplain/factories/api_key_factory.py b/aixplain/factories/api_key_factory.py index 89759fcf..eac50ad2 100644 --- a/aixplain/factories/api_key_factory.py +++ b/aixplain/factories/api_key_factory.py @@ -116,8 +116,6 @@ def get_usage_limits(cls, api_key: Text = config.TEAM_API_KEY, asset_id: Optiona """Get API key usage limits""" try: url = f"{config.BACKEND_URL}/sdk/api-keys/usage-limits" - if asset_id is not None: - url += f"?assetId={asset_id}" headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"} logging.info(f"Start service for GET API Key Usage - {url} - {headers}") r = _request_with_retry("GET", url, headers=headers) @@ -137,6 +135,7 @@ def get_usage_limits(cls, api_key: Text = config.TEAM_API_KEY, asset_id: Optiona model=limit["assetId"] if "assetId" in limit else None, ) for limit in resp + if asset_id is None or ("assetId" in limit and limit["assetId"] == asset_id) ] else: raise Exception(f"API Key Usage Error: Failed to get usage. Error: {str(resp)}") diff --git a/aixplain/modules/api_key.py b/aixplain/modules/api_key.py index b6cc8d3b..98b32fe2 100644 --- a/aixplain/modules/api_key.py +++ b/aixplain/modules/api_key.py @@ -169,8 +169,6 @@ def get_usage(self, asset_id: Optional[Text] = None) -> APIKeyUsageLimit: url = f"{config.BACKEND_URL}/sdk/api-keys/{self.id}/usage-limits" headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"} logging.info(f"Start service for GET API Key Usage - {url} - {headers}") - if asset_id is not None: - url += f"?assetId={asset_id}" r = _request_with_retry("GET", url, headers=headers) resp = r.json() except Exception: @@ -188,6 +186,7 @@ def get_usage(self, asset_id: Optional[Text] = None) -> APIKeyUsageLimit: model=limit["assetId"] if "assetId" in limit else None, ) for limit in resp + if asset_id is None or ("assetId" in limit and limit["assetId"] == asset_id) ] else: raise Exception(f"API Key Usage Error: Failed to get usage. Error: {str(resp)}") From d8bd8f9f9dbbb575b336929a9effed21e1dafd72 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Tue, 22 Oct 2024 15:26:32 -0300 Subject: [PATCH 04/10] Update version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e0df02a2..1f034299 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ namespaces = true [project] name = "aiXplain" -version = "0.2.21rc0" +version = "0.2.21rc1" description = "aiXplain SDK adds AI functions to software." readme = "README.md" requires-python = ">=3.5, <4" From f4af442f03f839cdc2a9321f1578c41477407cee Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Tue, 22 Oct 2024 22:50:27 -0300 Subject: [PATCH 05/10] Full response on synchronous model execution --- aixplain/modules/model/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aixplain/modules/model/utils.py b/aixplain/modules/model/utils.py index a78455b7..3b812882 100644 --- a/aixplain/modules/model/utils.py +++ b/aixplain/modules/model/utils.py @@ -55,7 +55,7 @@ def call_run_endpoint(url: Text, api_key: Text, payload: Dict) -> Dict: "error_message": "Model Run: An error occurred while processing your request.", } else: - response = {"status": status, "data": data, "completed": True} + response = resp else: if r.status_code == 401: error = f"Unauthorized API key: Please verify the spelling of the API key and its current validity. Details: {resp}" From b52ce149b797cf0486c3fdd84539df73b096419a Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Thu, 24 Oct 2024 17:58:01 -0300 Subject: [PATCH 06/10] Get API Key service --- aixplain/factories/api_key_factory.py | 8 ++++++++ aixplain/modules/model/utils.py | 2 +- aixplain/utils/config.py | 2 +- tests/functional/apikey/test_api.py | 5 +++++ tests/unit/llm_test.py | 2 +- tests/unit/model_test.py | 2 +- 6 files changed, 17 insertions(+), 4 deletions(-) diff --git a/aixplain/factories/api_key_factory.py b/aixplain/factories/api_key_factory.py index eac50ad2..e83a4690 100644 --- a/aixplain/factories/api_key_factory.py +++ b/aixplain/factories/api_key_factory.py @@ -10,6 +10,14 @@ class APIKeyFactory: backend_url = config.BACKEND_URL + @classmethod + def get(cls, api_key: Text) -> APIKey: + """Get an API key""" + for api_key_obj in cls.list(): + if str(api_key_obj.access_key).startswith(api_key[:4]) and str(api_key_obj.access_key).endswith(api_key[-4:]): + return api_key_obj + raise Exception(f"API Key Error: API key {api_key} not found") + @classmethod def list(cls) -> List[APIKey]: """List all API keys""" diff --git a/aixplain/modules/model/utils.py b/aixplain/modules/model/utils.py index 3b812882..f3ae4851 100644 --- a/aixplain/modules/model/utils.py +++ b/aixplain/modules/model/utils.py @@ -66,7 +66,7 @@ def call_run_endpoint(url: Text, api_key: Text, payload: Dict) -> Dict: elif 480 <= r.status_code < 490: error = f"Supplier-related error: Please ensure that the selected supplier provides the model you are trying to access. Details: {resp}" elif 490 <= r.status_code < 500: - error = f"Validation-related error: Please ensure all required fields are provided and correctly formatted. Details: {resp}" + error = f"{resp}" else: status_code = str(r.status_code) error = f"Status {status_code} - Unspecified error: {resp}" diff --git a/aixplain/utils/config.py b/aixplain/utils/config.py index 59805c60..03bbdccf 100644 --- a/aixplain/utils/config.py +++ b/aixplain/utils/config.py @@ -19,7 +19,7 @@ logger = logging.getLogger(__name__) BACKEND_URL = os.getenv("BACKEND_URL", "https://platform-api.aixplain.com") -MODELS_RUN_URL = os.getenv("MODELS_RUN_URL", "https://models.aixplain.com") +MODELS_RUN_URL = os.getenv("MODELS_RUN_URL", "https://models.aixplain.com/api/v1/execute") # GET THE API KEY FROM CMD TEAM_API_KEY = os.getenv("TEAM_API_KEY", "") AIXPLAIN_API_KEY = os.getenv("AIXPLAIN_API_KEY", "") diff --git a/tests/functional/apikey/test_api.py b/tests/functional/apikey/test_api.py index f2a50949..32ce3cea 100644 --- a/tests/functional/apikey/test_api.py +++ b/tests/functional/apikey/test_api.py @@ -102,6 +102,11 @@ def test_create_update_api_key_from_dict(): assert api_key.id != "" assert api_key.name == api_key_name + api_key_ = APIKeyFactory.get(api_key=api_key.access_key) + assert isinstance(api_key_, APIKey) + assert api_key_.id != "" + assert api_key_.name == api_key_name + api_key.global_limits.token_per_day = 222 api_key.global_limits.token_per_minute = 222 api_key.global_limits.request_per_day = 222 diff --git a/tests/unit/llm_test.py b/tests/unit/llm_test.py index b0dbe19a..fce9d2a0 100644 --- a/tests/unit/llm_test.py +++ b/tests/unit/llm_test.py @@ -30,7 +30,7 @@ ), ( 495, - "Validation-related error: Please ensure all required fields are provided and correctly formatted. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "{'error': 'An unspecified error occurred while processing your request.'}", ), (501, "Status 501 - Unspecified error: {'error': 'An unspecified error occurred while processing your request.'}"), ], diff --git a/tests/unit/model_test.py b/tests/unit/model_test.py index 0907b8f1..d52baed6 100644 --- a/tests/unit/model_test.py +++ b/tests/unit/model_test.py @@ -130,7 +130,7 @@ def test_failed_poll(): ), ( 495, - "Validation-related error: Please ensure all required fields are provided and correctly formatted. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "{'error': 'An unspecified error occurred while processing your request.'}", ), (501, "Status 501 - Unspecified error: {'error': 'An unspecified error occurred while processing your request.'}"), ], From 880c032dd5cb04589ee70d6080d8aa976b5e4850 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Thu, 24 Oct 2024 19:46:11 -0300 Subject: [PATCH 07/10] Change the name to daily count limits --- aixplain/factories/api_key_factory.py | 8 +++---- aixplain/modules/api_key.py | 32 +++++++++++++-------------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/aixplain/factories/api_key_factory.py b/aixplain/factories/api_key_factory.py index e83a4690..757fedff 100644 --- a/aixplain/factories/api_key_factory.py +++ b/aixplain/factories/api_key_factory.py @@ -136,10 +136,10 @@ def get_usage_limits(cls, api_key: Text = config.TEAM_API_KEY, asset_id: Optiona if 200 <= r.status_code < 300: return [ APIKeyUsageLimit( - request_count=limit["requestCount"], - request_count_limit=limit["requestCountLimit"], - token_count=limit["tokenCount"], - token_count_limit=limit["tokenCountLimit"], + daily_request_count=limit["requestCount"], + daily_request_limit=limit["requestCountLimit"], + daily_token_count=limit["tokenCount"], + daily_token_limit=limit["tokenCountLimit"], model=limit["assetId"] if "assetId" in limit else None, ) for limit in resp diff --git a/aixplain/modules/api_key.py b/aixplain/modules/api_key.py index 98b32fe2..04aae08f 100644 --- a/aixplain/modules/api_key.py +++ b/aixplain/modules/api_key.py @@ -29,25 +29,25 @@ def __init__( class APIKeyUsageLimit: def __init__( self, - request_count: int, - request_count_limit: int, - token_count: int, - token_count_limit: int, + daily_request_count: int, + daily_request_limit: int, + daily_token_count: int, + daily_token_limit: int, model: Optional[Union[Text, Model]] = None, ): """Get the usage limits of an API key globally (model equals to None) or for a specific model. Args: - request_count (int): number of requests made - request_count_limit (int): limit of requests - token_count (int): number of tokens used - token_count_limit (int): limit of tokens + daily_request_count (int): number of requests made + daily_request_limit (int): limit of requests + daily_token_count (int): number of tokens used + daily_token_limit (int): limit of tokens model (Optional[Union[Text, Model]], optional): Model which the limits apply. Defaults to None. """ - self.request_count = request_count - self.request_count_limit = request_count_limit - self.token_count = token_count - self.token_count_limit = token_count_limit + self.daily_request_count = daily_request_count + self.daily_request_limit = daily_request_limit + self.daily_token_count = daily_token_count + self.daily_token_limit = daily_token_limit if model is not None and isinstance(model, str): from aixplain.factories import ModelFactory @@ -179,10 +179,10 @@ def get_usage(self, asset_id: Optional[Text] = None) -> APIKeyUsageLimit: if 200 <= r.status_code < 300: return [ APIKeyUsageLimit( - request_count=limit["requestCount"], - request_count_limit=limit["requestCountLimit"], - token_count=limit["tokenCount"], - token_count_limit=limit["tokenCountLimit"], + daily_request_count=limit["requestCount"], + daily_request_limit=limit["requestCountLimit"], + daily_token_count=limit["tokenCount"], + daily_token_limit=limit["tokenCountLimit"], model=limit["assetId"] if "assetId" in limit else None, ) for limit in resp From 2691eeb34e4dcb5e349d4e39597db61c751c3d11 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Fri, 25 Oct 2024 11:04:32 -0300 Subject: [PATCH 08/10] Explicit point to the error field when it exists --- aixplain/modules/model/utils.py | 1 + tests/unit/llm_test.py | 12 ++++++------ tests/unit/model_test.py | 12 ++++++------ 3 files changed, 13 insertions(+), 12 deletions(-) diff --git a/aixplain/modules/model/utils.py b/aixplain/modules/model/utils.py index f3ae4851..d29da68b 100644 --- a/aixplain/modules/model/utils.py +++ b/aixplain/modules/model/utils.py @@ -57,6 +57,7 @@ def call_run_endpoint(url: Text, api_key: Text, payload: Dict) -> Dict: else: response = resp else: + resp = resp["error"] if "error" in resp else resp if r.status_code == 401: error = f"Unauthorized API key: Please verify the spelling of the API key and its current validity. Details: {resp}" elif 460 <= r.status_code < 470: diff --git a/tests/unit/llm_test.py b/tests/unit/llm_test.py index fce9d2a0..54887950 100644 --- a/tests/unit/llm_test.py +++ b/tests/unit/llm_test.py @@ -14,25 +14,25 @@ [ ( 401, - "Unauthorized API key: Please verify the spelling of the API key and its current validity. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "Unauthorized API key: Please verify the spelling of the API key and its current validity. Details: An unspecified error occurred while processing your request.", ), ( 465, - "Subscription-related error: Please ensure that your subscription is active and has not expired. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "Subscription-related error: Please ensure that your subscription is active and has not expired. Details: An unspecified error occurred while processing your request.", ), ( 475, - "Billing-related error: Please ensure you have enough credits to run this model. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "Billing-related error: Please ensure you have enough credits to run this model. Details: An unspecified error occurred while processing your request.", ), ( 485, - "Supplier-related error: Please ensure that the selected supplier provides the model you are trying to access. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "Supplier-related error: Please ensure that the selected supplier provides the model you are trying to access. Details: An unspecified error occurred while processing your request.", ), ( 495, - "{'error': 'An unspecified error occurred while processing your request.'}", + "An unspecified error occurred while processing your request.", ), - (501, "Status 501 - Unspecified error: {'error': 'An unspecified error occurred while processing your request.'}"), + (501, "Status 501 - Unspecified error: An unspecified error occurred while processing your request."), ], ) def test_run_async_errors(status_code, error_message): diff --git a/tests/unit/model_test.py b/tests/unit/model_test.py index d52baed6..03dccdbe 100644 --- a/tests/unit/model_test.py +++ b/tests/unit/model_test.py @@ -114,25 +114,25 @@ def test_failed_poll(): [ ( 401, - "Unauthorized API key: Please verify the spelling of the API key and its current validity. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "Unauthorized API key: Please verify the spelling of the API key and its current validity. Details: An unspecified error occurred while processing your request.", ), ( 465, - "Subscription-related error: Please ensure that your subscription is active and has not expired. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "Subscription-related error: Please ensure that your subscription is active and has not expired. Details: An unspecified error occurred while processing your request.", ), ( 475, - "Billing-related error: Please ensure you have enough credits to run this model. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "Billing-related error: Please ensure you have enough credits to run this model. Details: An unspecified error occurred while processing your request.", ), ( 485, - "Supplier-related error: Please ensure that the selected supplier provides the model you are trying to access. Details: {'error': 'An unspecified error occurred while processing your request.'}", + "Supplier-related error: Please ensure that the selected supplier provides the model you are trying to access. Details: An unspecified error occurred while processing your request.", ), ( 495, - "{'error': 'An unspecified error occurred while processing your request.'}", + "An unspecified error occurred while processing your request.", ), - (501, "Status 501 - Unspecified error: {'error': 'An unspecified error occurred while processing your request.'}"), + (501, "Status 501 - Unspecified error: An unspecified error occurred while processing your request."), ], ) def test_run_async_errors(status_code, error_message): From 5c831a95a424d29d49176aa811b7c393240915e8 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Fri, 25 Oct 2024 11:15:16 -0300 Subject: [PATCH 09/10] Validate api key before updating it --- aixplain/factories/api_key_factory.py | 1 + 1 file changed, 1 insertion(+) diff --git a/aixplain/factories/api_key_factory.py b/aixplain/factories/api_key_factory.py index 757fedff..6a20a498 100644 --- a/aixplain/factories/api_key_factory.py +++ b/aixplain/factories/api_key_factory.py @@ -92,6 +92,7 @@ def create( @classmethod def update(cls, api_key: APIKey) -> APIKey: """Update an existing API key""" + api_key.validate() try: resp = "Unspecified error" url = f"{cls.backend_url}/sdk/api-keys/{api_key.id}" From c2deee4ad93dd6f127275aced5e07e170b7b7ce2 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Fri, 25 Oct 2024 15:57:29 -0300 Subject: [PATCH 10/10] Create Setters for token and request limits --- aixplain/factories/api_key_factory.py | 6 +-- aixplain/modules/__init__.py | 2 +- aixplain/modules/api_key.py | 42 ++++++++++++++--- tests/functional/apikey/test_api.py | 18 ++++---- tests/unit/api_key_test.py | 65 ++++++++++++++++++++++++--- 5 files changed, 109 insertions(+), 24 deletions(-) diff --git a/aixplain/factories/api_key_factory.py b/aixplain/factories/api_key_factory.py index 6a20a498..c719c26b 100644 --- a/aixplain/factories/api_key_factory.py +++ b/aixplain/factories/api_key_factory.py @@ -4,7 +4,7 @@ from datetime import datetime from typing import Text, List, Optional, Dict, Union from aixplain.utils.file_utils import _request_with_retry -from aixplain.modules.api_key import APIKey, APIKeyGlobalLimits, APIKeyUsageLimit +from aixplain.modules.api_key import APIKey, APIKeyLimits, APIKeyUsageLimit class APIKeyFactory: @@ -54,8 +54,8 @@ def create( cls, name: Text, budget: int, - global_limits: Union[Dict, APIKeyGlobalLimits], - asset_limits: List[Union[Dict, APIKeyGlobalLimits]], + global_limits: Union[Dict, APIKeyLimits], + asset_limits: List[Union[Dict, APIKeyLimits]], expires_at: datetime, ) -> APIKey: """Create a new API key""" diff --git a/aixplain/modules/__init__.py b/aixplain/modules/__init__.py index d49e29d4..4432e1ad 100644 --- a/aixplain/modules/__init__.py +++ b/aixplain/modules/__init__.py @@ -36,4 +36,4 @@ from .agent import Agent from .agent.tool import Tool from .team_agent import TeamAgent -from .api_key import APIKey, APIKeyGlobalLimits, APIKeyUsageLimit +from .api_key import APIKey, APIKeyLimits, APIKeyUsageLimit diff --git a/aixplain/modules/api_key.py b/aixplain/modules/api_key.py index 04aae08f..ae774c23 100644 --- a/aixplain/modules/api_key.py +++ b/aixplain/modules/api_key.py @@ -6,7 +6,7 @@ from typing import Dict, List, Optional, Text, Union -class APIKeyGlobalLimits: +class APIKeyLimits: def __init__( self, token_per_minute: int, @@ -60,8 +60,8 @@ def __init__( name: Text, expires_at: Optional[Union[datetime, Text]] = None, budget: Optional[float] = None, - asset_limits: List[APIKeyGlobalLimits] = [], - global_limits: Optional[Union[Dict, APIKeyGlobalLimits]] = None, + asset_limits: List[APIKeyLimits] = [], + global_limits: Optional[Union[Dict, APIKeyLimits]] = None, id: int = "", access_key: Optional[Text] = None, is_admin: bool = False, @@ -71,7 +71,7 @@ def __init__( self.budget = budget self.global_limits = global_limits if global_limits is not None and isinstance(global_limits, dict): - self.global_limits = APIKeyGlobalLimits( + self.global_limits = APIKeyLimits( token_per_minute=global_limits["tpm"], token_per_day=global_limits["tpd"], request_per_minute=global_limits["rpm"], @@ -80,7 +80,7 @@ def __init__( self.asset_limits = asset_limits for i, asset_limit in enumerate(self.asset_limits): if isinstance(asset_limit, dict): - self.asset_limits[i] = APIKeyGlobalLimits( + self.asset_limits[i] = APIKeyLimits( token_per_minute=asset_limit["tpm"], token_per_day=asset_limit["tpd"], request_per_minute=asset_limit["rpm"], @@ -190,3 +190,35 @@ def get_usage(self, asset_id: Optional[Text] = None) -> APIKeyUsageLimit: ] else: raise Exception(f"API Key Usage Error: Failed to get usage. Error: {str(resp)}") + + def __set_limit(self, limit: int, model: Optional[Union[Text, Model]], limit_type: Text) -> None: + """Set a limit for an API key""" + if model is None: + setattr(self.global_limits, limit_type, limit) + else: + if isinstance(model, Model): + model = model.id + is_found = False + for i, asset_limit in enumerate(self.asset_limits): + if asset_limit.model.id == model: + setattr(self.asset_limits[i], limit_type, limit) + is_found = True + break + if is_found is False: + raise Exception(f"Limit for Model {model} not found in the API key.") + + def set_token_per_day(self, token_per_day: int, model: Optional[Union[Text, Model]] = None) -> None: + """Set the token per day limit of an API key""" + self.__set_limit(token_per_day, model, "token_per_day") + + def set_token_per_minute(self, token_per_minute: int, model: Optional[Union[Text, Model]] = None) -> None: + """Set the token per minute limit of an API key""" + self.__set_limit(token_per_minute, model, "token_per_minute") + + def set_request_per_day(self, request_per_day: int, model: Optional[Union[Text, Model]] = None) -> None: + """Set the request per day limit of an API key""" + self.__set_limit(request_per_day, model, "request_per_day") + + def set_request_per_minute(self, request_per_minute: int, model: Optional[Union[Text, Model]] = None) -> None: + """Set the request per minute limit of an API key""" + self.__set_limit(request_per_minute, model, "request_per_minute") diff --git a/tests/functional/apikey/test_api.py b/tests/functional/apikey/test_api.py index 32ce3cea..2c228f6b 100644 --- a/tests/functional/apikey/test_api.py +++ b/tests/functional/apikey/test_api.py @@ -1,5 +1,5 @@ from aixplain.factories.api_key_factory import APIKeyFactory -from aixplain.modules import APIKey, APIKeyGlobalLimits, APIKeyUsageLimit +from aixplain.modules import APIKey, APIKeyLimits, APIKeyUsageLimit from datetime import datetime import json import pytest @@ -16,7 +16,7 @@ def test_create_api_key_from_json(): api_key = APIKeyFactory.create( name=api_key_data["name"], asset_limits=[ - APIKeyGlobalLimits( + APIKeyLimits( model=api_key_data["asset_limits"][0]["model"], token_per_minute=api_key_data["asset_limits"][0]["token_per_minute"], token_per_day=api_key_data["asset_limits"][0]["token_per_day"], @@ -24,7 +24,7 @@ def test_create_api_key_from_json(): request_per_minute=api_key_data["asset_limits"][0]["request_per_minute"], ) ], - global_limits=APIKeyGlobalLimits( + global_limits=APIKeyLimits( token_per_minute=api_key_data["global_limits"]["token_per_minute"], token_per_day=api_key_data["global_limits"]["token_per_day"], request_per_day=api_key_data["global_limits"]["request_per_day"], @@ -60,8 +60,8 @@ def test_create_api_key_from_dict(): api_key_name = "Test API Key" api_key = APIKeyFactory.create( name=api_key_name, - asset_limits=[APIKeyGlobalLimits(**limit) for limit in api_key_dict["asset_limits"]], - global_limits=APIKeyGlobalLimits(**api_key_dict["global_limits"]), + asset_limits=[APIKeyLimits(**limit) for limit in api_key_dict["asset_limits"]], + global_limits=APIKeyLimits(**api_key_dict["global_limits"]), budget=api_key_dict["budget"], expires_at=datetime.strptime(api_key_dict["expires_at"], "%Y-%m-%dT%H:%M:%SZ"), ) @@ -92,8 +92,8 @@ def test_create_update_api_key_from_dict(): api_key_name = "Test API Key" api_key = APIKeyFactory.create( name=api_key_name, - asset_limits=[APIKeyGlobalLimits(**limit) for limit in api_key_dict["asset_limits"]], - global_limits=APIKeyGlobalLimits(**api_key_dict["global_limits"]), + asset_limits=[APIKeyLimits(**limit) for limit in api_key_dict["asset_limits"]], + global_limits=APIKeyLimits(**api_key_dict["global_limits"]), budget=api_key_dict["budget"], expires_at=datetime.strptime(api_key_dict["expires_at"], "%Y-%m-%dT%H:%M:%SZ"), ) @@ -156,7 +156,7 @@ def test_list_update_api_keys(): number = randint(0, 10000) if api_key.global_limits is None: - api_key.global_limits = APIKeyGlobalLimits( + api_key.global_limits = APIKeyLimits( token_per_minute=number, token_per_day=number, request_per_day=number, @@ -173,7 +173,7 @@ def test_list_update_api_keys(): if len(api_key.asset_limits) == 0: api_key.asset_limits.append( - APIKeyGlobalLimits( + APIKeyLimits( model="640b517694bf816d35a59125", token_per_minute=number, token_per_day=number, diff --git a/tests/unit/api_key_test.py b/tests/unit/api_key_test.py index 9c8dc1cd..7da4e082 100644 --- a/tests/unit/api_key_test.py +++ b/tests/unit/api_key_test.py @@ -1,5 +1,5 @@ __author__ = "aixplain" -from aixplain.modules import APIKeyGlobalLimits +from aixplain.modules import APIKeyLimits from datetime import datetime import requests_mock import aixplain.utils.config as config @@ -13,7 +13,7 @@ def read_data(data_path): def test_api_key_service(): with requests_mock.Mocker() as mock: - model_id = "640b517694bf816d35a59125" + model_id = "test_asset_id" model_url = f"{config.BACKEND_URL}/sdk/models/{model_id}" model_map = read_data("tests/unit/mock_responses/model_response.json") mock.get(model_url, json=model_map) @@ -34,13 +34,11 @@ def test_api_key_service(): api_key = APIKeyFactory.create( name="Test API Key", asset_limits=[ - APIKeyGlobalLimits( + APIKeyLimits( model=model_id, token_per_minute=100, token_per_day=1000, request_per_day=1000, request_per_minute=100 ) ], - global_limits=APIKeyGlobalLimits( - token_per_minute=100, token_per_day=1000, request_per_day=1000, request_per_minute=100 - ), + global_limits=APIKeyLimits(token_per_minute=100, token_per_day=1000, request_per_day=1000, request_per_minute=100), budget=1000, expires_at=datetime(2024, 10, 7), ) @@ -65,3 +63,58 @@ def test_api_key_service(): mock.delete(delete_url, status_code=200) api_key.delete() + + +def test_setters(): + with requests_mock.Mocker() as mock: + model_id = "test_asset_id" + model_url = f"{config.BACKEND_URL}/sdk/models/{model_id}" + model_map = read_data("tests/unit/mock_responses/model_response.json") + mock.get(model_url, json=model_map) + + create_url = f"{config.BACKEND_URL}/sdk/api-keys" + api_key_response = { + "id": "key-id", + "name": "Name", + "accessKey": "access-key", + "budget": 1000, + "globalLimits": {"tpm": 100, "tpd": 1000, "rpd": 1000, "rpm": 100}, + "assetsLimits": [{"assetId": model_id, "tpm": 100, "tpd": 1000, "rpd": 1000, "rpm": 100}], + "expiresAt": "2024-10-07T00:00:00Z", + "isAdmin": False, + } + mock.post(create_url, json=api_key_response) + + api_key = APIKeyFactory.create( + name="Test API Key", + asset_limits=[ + APIKeyLimits( + model=model_id, token_per_minute=100, token_per_day=1000, request_per_day=1000, request_per_minute=100 + ) + ], + global_limits=APIKeyLimits(token_per_minute=100, token_per_day=1000, request_per_day=1000, request_per_minute=100), + budget=1000, + expires_at=datetime(2024, 10, 7), + ) + + api_key.set_token_per_day(1) + api_key.set_token_per_minute(1) + api_key.set_request_per_day(1) + api_key.set_request_per_minute(1) + api_key.set_token_per_day(1, model_id) + api_key.set_token_per_minute(1, model_id) + api_key.set_request_per_day(1, model_id) + api_key.set_request_per_minute(1, model_id) + + assert api_key.asset_limits[0].token_per_day == 1 + assert api_key.asset_limits[0].token_per_minute == 1 + assert api_key.asset_limits[0].request_per_day == 1 + assert api_key.asset_limits[0].request_per_minute == 1 + assert api_key.global_limits.token_per_day == 1 + assert api_key.global_limits.token_per_minute == 1 + assert api_key.global_limits.request_per_day == 1 + assert api_key.global_limits.request_per_minute == 1 + + +if __name__ == "__main__": + test_setters()