Group of Improvements in API Key CRUD

aixplain/factories/api_key_factory.py

@@ -4,12 +4,20 @@
 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:
     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"""
@@ -30,7 +38,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"],
@@ -46,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"""
@@ -84,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}"
@@ -112,12 +121,10 @@ 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:
-                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)
@@ -128,11 +135,16 @@ 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(
+                    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
+                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)}")

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

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,
@@ -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,
+        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
+
+            self.model = ModelFactory.get(model)
 
 
 class APIKey:
@@ -48,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,
@@ -59,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"],
@@ -68,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"],
@@ -110,7 +122,7 @@ def to_dict(self) -> Dict:
             "id": self.id,
             "name": self.name,
             "budget": self.budget,
-            "assetLimits": [],
+            "assetsLimits": [],
             "expiresAt": self.expires_at,
         }
 
@@ -126,7 +138,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,
@@ -157,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:
@@ -167,11 +177,48 @@ 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(
+                    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
+                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)}")
+
+    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")

aixplain/modules/model/utils.py

@@ -55,8 +55,9 @@ 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:
+        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:
@@ -66,7 +67,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}"

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", "")

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"

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,15 +16,15 @@ 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"],
                 request_per_day=api_key_data["asset_limits"][0]["request_per_day"],
                 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"),
     )
@@ -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
@@ -134,7 +139,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():
@@ -149,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,
@@ -166,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,