diff --git a/aixplain/factories/api_key_factory.py b/aixplain/factories/api_key_factory.py index 750ce0b2..f5d0a6b2 100644 --- a/aixplain/factories/api_key_factory.py +++ b/aixplain/factories/api_key_factory.py @@ -2,9 +2,9 @@ import logging import aixplain.utils.config as config from datetime import datetime -from typing import Text, List, Dict, Union +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 +from aixplain.modules.api_key import APIKey, APIKeyGlobalLimits, APIKeyUsageLimit class APIKeyFactory: @@ -72,7 +72,7 @@ def create( name=resp["name"], budget=resp["budget"] if "budget" in resp else None, global_limits=resp["globalLimits"] if "globalLimits" in resp else None, - asset_limits=resp["assetLimits"] if "assetLimits" in resp else [], + asset_limits=resp["assetsLimits"] if "assetsLimits" in resp else [], expires_at=resp["expiresAt"] if "expiresAt" in resp else None, access_key=resp["accessKey"], is_admin=resp["isAdmin"], @@ -110,3 +110,29 @@ def update(cls, api_key: APIKey) -> APIKey: return api_key else: 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""" + 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) + resp = r.json() + except Exception: + message = "API Key Usage Error: Make sure the API Key exists and you are the owner." + logging.error(message) + 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"], + ) + else: + raise Exception(f"API Key Usage Error: Failed to get usage. Error: {str(resp)}") diff --git a/aixplain/modules/__init__.py b/aixplain/modules/__init__.py index 6ac5ae9e..d49e29d4 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 +from .api_key import APIKey, APIKeyGlobalLimits, APIKeyUsageLimit diff --git a/aixplain/modules/api_key.py b/aixplain/modules/api_key.py index 15b1bb68..884735b6 100644 --- a/aixplain/modules/api_key.py +++ b/aixplain/modules/api_key.py @@ -1,3 +1,6 @@ +import logging +from aixplain.utils import config +from aixplain.utils.file_utils import _request_with_retry from aixplain.modules import Model from datetime import datetime from typing import Dict, List, Optional, Text, Union @@ -23,6 +26,22 @@ def __init__( self.model = ModelFactory.get(model) +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 + + 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 + """ + self.request_count = request_count + self.request_count_limit = request_count_limit + self.token_count = token_count + self.token_count_limit = token_count_limit + + class APIKey: def __init__( self, @@ -54,7 +73,7 @@ def __init__( token_per_day=asset_limit["tpd"], request_per_minute=asset_limit["rpm"], request_per_day=asset_limit["rpd"], - model=asset_limit["model"], + model=asset_limit["assetId"], ) self.expires_at = expires_at self.access_key = access_key @@ -110,17 +129,13 @@ def to_dict(self) -> Dict: "tpd": asset_limit.token_per_day, "rpm": asset_limit.request_per_minute, "rpd": asset_limit.request_per_day, - "model": asset_limit.model.id, + "assetId": asset_limit.model.id, } ) return payload def delete(self) -> None: """Delete an API key by its ID""" - import logging - from aixplain.utils import config - from aixplain.utils.file_utils import _request_with_retry - try: url = f"{config.BACKEND_URL}/sdk/api-keys/{self.id}" headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"} @@ -132,3 +147,28 @@ def delete(self) -> None: message = "API Key Deletion Error: Make sure the API Key exists and you are the owner." logging.error(message) raise Exception(f"{message}") + + def get_usage(self, asset_id: Optional[Text] = None) -> APIKeyUsageLimit: + """Get the usage limits of an API key""" + try: + 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: + message = "API Key Usage Error: Make sure the API Key exists and you are the owner." + logging.error(message) + 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"], + ) + else: + raise Exception(f"API Key Usage Error: Failed to get usage. Error: {str(resp)}") diff --git a/pyproject.toml b/pyproject.toml index be397bdd..44f4d45e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ namespaces = true [project] name = "aiXplain" -version = "0.2.18" +version = "0.2.20" description = "aiXplain SDK adds AI functions to software." readme = "README.md" requires-python = ">=3.5, <4" diff --git a/tests/functional/agent/agent_functional_test.py b/tests/functional/agent/agent_functional_test.py index 0b0327e1..648f4f28 100644 --- a/tests/functional/agent/agent_functional_test.py +++ b/tests/functional/agent/agent_functional_test.py @@ -55,7 +55,7 @@ def test_end2end(run_input_map): if "pipeline_tools" in run_input_map: for tool in run_input_map["pipeline_tools"]: tools.append(AgentFactory.create_pipeline_tool(pipeline=tool["pipeline_id"], description=tool["description"])) - + agent = AgentFactory.create( name=run_input_map["agent_name"], description=run_input_map["agent_name"], llm_id=run_input_map["llm_id"], tools=tools ) diff --git a/tests/functional/apikey/test_api.py b/tests/functional/apikey/test_api.py index 31094a59..9359eddc 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 +from aixplain.modules import APIKey, APIKeyGlobalLimits, APIKeyUsageLimit from datetime import datetime import json import pytest @@ -81,6 +81,9 @@ def test_list_api_keys(): assert isinstance(api_key, APIKey) assert api_key.id != "" + if api_key.is_admin is False: + usage = api_key.get_usage() + assert isinstance(usage, APIKeyUsageLimit) def test_create_api_key_wrong_input(): api_key_name = "Test API Key" diff --git a/tests/unit/api_key_test.py b/tests/unit/api_key_test.py index fa610cae..60d2371d 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": [{"model": model_id, "tpm": 100, "tpd": 1000, "rpd": 1000, "rpm": 100}], + "assetLimits": [{"assetId": model_id, "tpm": 100, "tpd": 1000, "rpd": 1000, "rpm": 100}], "expiresAt": "2024-10-07T00:00:00Z", "isAdmin": False, }