Development 2 Test

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/factories/model_factory.py

@@ -80,6 +80,7 @@ def _create_model_from_response(cls, response: Dict) -> Model:
         return ModelClass(
             response["id"],
             response["name"],
+            description=response.get("description", ""),
             supplier=response["supplier"],
             api_key=response["api_key"],
             cost=response["pricing"],
@@ -221,8 +222,8 @@ def _get_assets_from_page(
     @classmethod
     def list(
         cls,
+        function: Function,
         query: Optional[Text] = "",
-        function: Optional[Function] = None,
         suppliers: Optional[Union[Supplier, List[Supplier]]] = None,
         source_languages: Optional[Union[Language, List[Language]]] = None,
         target_languages: Optional[Union[Language, List[Language]]] = None,
@@ -236,7 +237,7 @@ def list(
         """Gets the first k given models based on the provided task and language filters
 
         Args:
-            function (Optional[Function], optional): function filter. Defaults to None.
+            function (Function): function filter.
             source_languages (Optional[Union[Language, List[Language]]], optional): language filter of input data. Defaults to None.
             target_languages (Optional[Union[Language, List[Language]]], optional): language filter of output data. Defaults to None.
             is_finetunable (Optional[bool], optional): can be finetuned or not. Defaults to None.

aixplain/factories/pipeline_factory/utils.py

@@ -86,8 +86,23 @@ def build_from_response(response: Dict, load_architecture: bool = False) -> Pipe
                 node.label = node_json["label"]
                 pipeline.add_node(node)
 
+            # Decision nodes' output parameters are defined based on their
+            # input parameters linked. So here we have to make sure that
+            # decision nodes (having passthrough parameter) should be first
+            # linked
+            link_jsons = response["links"][:]
+            decision_links = []
+            for link_json in link_jsons:
+                for pm in link_json["paramMapping"]:
+                    if pm["to"] == "passthrough":
+                        decision_link_index = link_jsons.index(link_json)
+                        decision_link = link_jsons.pop(decision_link_index)
+                        decision_links.append(decision_link)
+
+            link_jsons = decision_links + link_jsons
+
             # instantiating links
-            for link_json in response["links"]:
+            for link_json in link_jsons:
                 for param_mapping in link_json["paramMapping"]:
                     link = Link(
                         from_node=pipeline.get_node(link_json["from"]),

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/modules/pipeline/designer/base.py

@@ -142,14 +142,34 @@ def __init__(
         pipeline: "DesignerPipeline" = None,
     ):
 
-        assert from_param in from_node.outputs, "Invalid from param"
-        assert to_param in to_node.inputs, "Invalid to param"
-
         if isinstance(from_param, Param):
             from_param = from_param.code
         if isinstance(to_param, Param):
             to_param = to_param.code
 
+        assert from_param in from_node.outputs, (
+            "Invalid from param. "
+            "Make sure all input params are already linked accordingly"
+        )
+
+        fp_instance = from_node.outputs[from_param]
+        from .nodes import Decision
+
+        if (
+            isinstance(to_node, Decision)
+            and to_param == to_node.inputs.passthrough.code
+        ):
+            if from_param not in to_node.outputs:
+                to_node.outputs.create_param(
+                    from_param,
+                    fp_instance.data_type,
+                    is_required=fp_instance.is_required,
+                )
+            else:
+                to_node.outputs[from_param].data_type = fp_instance.data_type
+
+        assert to_param in to_node.inputs, "Invalid to param"
+
         self.from_node = from_node
         self.to_node = to_node
         self.from_param = from_param
@@ -233,9 +253,7 @@ def __init__(self, node: "Node", *args, **kwargs):
     def add_param(self, param: Param) -> None:
         # check if param already registered
         if param in self:
-            raise ValueError(
-                f"Parameter with code '{param.code}' already exists."
-            )
+            raise ValueError(f"Parameter with code '{param.code}' already exists.")
         self._params.append(param)
         # also set attribute on the node dynamically if there's no
         # any attribute with the same name
@@ -353,9 +371,7 @@ def attach_to(self, pipeline: "DesignerPipeline"):
         :param pipeline: the pipeline
         """
         assert not self.pipeline, "Node already attached to a pipeline"
-        assert (
-            self not in pipeline.nodes
-        ), "Node already attached to a pipeline"
+        assert self not in pipeline.nodes, "Node already attached to a pipeline"
         assert self.type, "Node type not set"
 
         self.pipeline = pipeline