Error handling and Team Key validation

README.md

@@ -34,6 +34,11 @@ export TEAM_API_KEY=YOUR_API_KEY
 ```bash
 set TEAM_API_KEY=YOUR_API_KEY
 ```
+#### Jupyter Notebook
+```
+%env TEAM_API_KEY=YOUR_API_KEY
+```
+
 ## Usage
 
 Let’s see how we can use aixplain to run a machine translation model.

aixplain/__init__.py

@@ -25,3 +25,16 @@
 from logging import NullHandler
 
 logging.getLogger(__name__).addHandler(NullHandler())
+
+# Load Environment Variables
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Validate API keys
+from aixplain.utils import config
+
+if config.TEAM_API_KEY == "" and config.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)"
+    )

aixplain/__version__.py

@@ -1,7 +1,7 @@
 __title__ = "aixplain"
 __description__ = "aiXplain SDK adds AI functions to software."
 __url__ = "https://github.com/aixplain/pipelines/tree/main/docs"
-__version__ = "0.2"
+__version__ = "0.2.1"
 __author__ = "aiXplain"
 __author_email__ = "thiago.ferreira@aixplain.com, krishna.durai@aixplain.com, lucas.pavanelli@aixplain.com"
 __license__ = "http://www.apache.org/licenses/LICENSE-2.0"

aixplain/enums/__init__.py

@@ -1,6 +1,7 @@
 from .data_split import DataSplit
 from .data_subtype import DataSubtype
 from .data_type import DataType
+from .error_handler import ErrorHandler
 from .file_type import FileType
 from .function import Function
 from .language import Language

aixplain/enums/error_handler.py

@@ -0,0 +1,37 @@
+__author__ = "aiXplain"
+
+"""
+Copyright 2023 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.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Author: aiXplain team
+Date: May 26th 2023
+Description:
+    Error Handler Enum
+"""
+
+from enum import Enum
+
+
+class ErrorHandler(Enum):
+    """
+    Enumeration class defining different error handler strategies.
+
+    Attributes:
+        SKIP (str): skip failed rows.
+        FAIL (str): raise an exception.
+    """
+
+    SKIP = "skip"
+    FAIL = "fail"

aixplain/enums/function.py

@@ -36,10 +36,15 @@ def load_functions():
 
     url = urljoin(backend_url, "sdk/functions")
     if aixplain_key != "":
+        api_key = aixplain_key
         headers = {"x-aixplain-key": aixplain_key, "Content-Type": "application/json"}
     else:
         headers = {"x-api-key": api_key, "Content-Type": "application/json"}
     r = _request_with_retry("get", url, headers=headers)
+    if not 200 <= r.status_code < 300:
+        raise Exception(
+            f'Functions could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
+        )
     resp = r.json()
     return Enum("Function", {w["id"].upper().replace("-", "_"): w["id"] for w in resp["items"]}, type=str)
 

aixplain/enums/language.py

@@ -36,10 +36,15 @@ def load_languages():
 
     url = urljoin(backend_url, "sdk/languages")
     if aixplain_key != "":
+        api_key = aixplain_key
         headers = {"x-aixplain-key": aixplain_key, "Content-Type": "application/json"}
     else:
         headers = {"x-api-key": api_key, "Content-Type": "application/json"}
     r = _request_with_retry("get", url, headers=headers)
+    if not 200 <= r.status_code < 300:
+        raise Exception(
+            f'Languages could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
+        )
     resp = r.json()
     languages = {}
     for w in resp:

aixplain/enums/license.py

@@ -37,10 +37,15 @@ def load_licenses():
 
         url = urljoin(backend_url, "sdk/licenses")
         if aixplain_key != "":
+            api_key = aixplain_key
             headers = {"x-aixplain-key": aixplain_key, "Content-Type": "application/json"}
         else:
             headers = {"x-api-key": api_key, "Content-Type": "application/json"}
         r = _request_with_retry("get", url, headers=headers)
+        if not 200 <= r.status_code < 300:
+            raise Exception(
+                f'Licenses could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
+            )
         resp = r.json()
         return Enum("License", {"_".join(w["name"].split()): w["id"] for w in resp}, type=str)
     except Exception as e:

aixplain/factories/corpus_factory.py

@@ -34,6 +34,7 @@
 from aixplain.modules.metadata import MetaData
 from aixplain.enums.data_subtype import DataSubtype
 from aixplain.enums.data_type import DataType
+from aixplain.enums.error_handler import ErrorHandler
 from aixplain.enums.function import Function
 from aixplain.enums.language import Language
 from aixplain.enums.license import License
@@ -239,6 +240,7 @@ def create(
         tags: List[Text] = [],
         functions: List[Function] = [],
         privacy: Privacy = Privacy.PRIVATE,
+        error_handler: ErrorHandler = ErrorHandler.SKIP,
     ) -> Dict:
         """Asynchronous call to Upload a corpus to the user's dashboard.
 
@@ -252,6 +254,7 @@ def create(
             tags (Optional[List[Text]], optional): tags that explain the corpus. Defaults to [].
             functions (Optional[List[Function]], optional): AI functions for which the corpus may be used. Defaults to [].
             privacy (Optional[Privacy], optional): visibility of the corpus. Defaults to Privacy.PRIVATE.
+            error_handler (ErrorHandler, optional): how to handle failed rows in the data asset. Defaults to ErrorHandler.SKIP.
 
         Returns:
             Dict: response dict
@@ -347,7 +350,7 @@ def create(
                 onboard_status="onboarding",
             )
 
-            corpus_payload = onboard_functions.build_payload_corpus(corpus, [ref.id for ref in ref_data])
+            corpus_payload = onboard_functions.build_payload_corpus(corpus, [ref.id for ref in ref_data], error_handler)
 
             response = onboard_functions.create_data_asset(corpus_payload)
             if response["success"] is True:

aixplain/factories/dataset_factory.py

@@ -34,6 +34,7 @@
 from aixplain.modules.metadata import MetaData
 from aixplain.enums.data_subtype import DataSubtype
 from aixplain.enums.data_type import DataType
+from aixplain.enums.error_handler import ErrorHandler
 from aixplain.enums.function import Function
 from aixplain.enums.language import Language
 from aixplain.enums.license import License
@@ -312,6 +313,7 @@ def create(
         privacy: Privacy = Privacy.PRIVATE,
         split_labels: Optional[List[Text]] = None,
         split_rate: Optional[List[float]] = None,
+        error_handler: ErrorHandler = ErrorHandler.SKIP,
     ) -> Dict:
         """Dataset Onboard
 
@@ -331,6 +333,7 @@ def create(
             meta_ref_data (Dict[Text, Any], optional): metadata which is already in the platform. Defaults to {}.
             tags (List[Text], optional): datasets description tags. Defaults to [].
             privacy (Privacy, optional): dataset privacy. Defaults to Privacy.PRIVATE.
+            error_handler (ErrorHandler, optional): how to handle failed rows in the data asset. Defaults to ErrorHandler.SKIP.
 
         Returns:
             Dict: dataset onboard status
@@ -509,7 +512,7 @@ def create(
             ), f"Data Asset Onboarding Error: All data must have the same number of rows. Lengths: {str(set(sizes))}"
 
             dataset_payload = onboard_functions.build_payload_dataset(
-                dataset, input_ref_data, output_ref_data, hypotheses_ref_data, meta_ref_data, tags
+                dataset, input_ref_data, output_ref_data, hypotheses_ref_data, meta_ref_data, tags, error_handler
             )
             assert (
                 len(dataset_payload["input"]) > 0

aixplain/modules/model.py

@@ -162,7 +162,14 @@ def poll(self, poll_url: Text, name: Text = "model_process") -> Dict:
             logging.error(f"Single Poll for Model: Error of polling for {name}: {e}")
         return resp
 
-    def run(self, data: Union[Text, Dict], name: Text = "model_process", timeout: float = 300, parameters: Dict = {}, wait_time: float = 0.5) -> Dict:
+    def run(
+        self,
+        data: Union[Text, Dict],
+        name: Text = "model_process",
+        timeout: float = 300,
+        parameters: Dict = {},
+        wait_time: float = 0.5,
+    ) -> Dict:
         """Runs a model call.
 
         Args:
@@ -203,6 +210,10 @@ def run_async(self, data: Union[Text, Dict], name: Text = "model_process", param
         Returns:
             dict: polling URL in response
         """
+        if self.api_key == "":
+            raise Exception(
+                "A 'TEAM_API_KEY' is required to run a model. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)"
+            )
         headers = {"x-api-key": self.api_key, "Content-Type": "application/json"}
 
         data = FileFactory.to_link(data)

aixplain/modules/pipeline.py

@@ -133,7 +133,9 @@ def poll(self, poll_url: Text, name: Text = "pipeline_process") -> Dict:
             resp = {"status": "FAILED"}
         return resp
 
-    def run(self, data: Union[Text, Dict], name: Text = "pipeline_process", timeout: float = 20000.0, wait_time: float = 1.0) -> Dict:
+    def run(
+        self, data: Union[Text, Dict], name: Text = "pipeline_process", timeout: float = 20000.0, wait_time: float = 1.0
+    ) -> Dict:
         """Runs a pipeline call.
 
         Args:
@@ -173,7 +175,10 @@ def run_async(self, data: Union[Text, Dict], name: str = "pipeline_process") ->
         Returns:
             Dict: polling URL in response
         """
-
+        if self.api_key == "":
+            raise Exception(
+                "A 'TEAM_API_KEY' is required to run a pipeline. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)"
+            )
         headers = {"x-api-key": self.api_key, "Content-Type": "application/json"}
 
         data = FileFactory.to_link(data)

aixplain/processes/data_onboarding/onboard_functions.py

@@ -10,6 +10,7 @@
 
 from aixplain.enums.data_subtype import DataSubtype
 from aixplain.enums.data_type import DataType
+from aixplain.enums.error_handler import ErrorHandler
 from aixplain.enums.file_type import FileType
 from aixplain.enums.storage_type import StorageType
 from aixplain.modules.corpus import Corpus
@@ -128,12 +129,13 @@ def build_payload_data(data: Data) -> Dict:
     return data_json
 
 
-def build_payload_corpus(corpus: Corpus, ref_data: List[Text]) -> Dict:
+def build_payload_corpus(corpus: Corpus, ref_data: List[Text], error_handler: ErrorHandler) -> Dict:
     """Create corpus payload to call coreengine on the onboard process
 
     Args:
         corpus (Corpus): corpus object
         ref_data (List[Text]): list of referred data
+        error_handler (ErrorHandler): how to handle failed rows
 
     Returns:
         Dict: payload
@@ -142,6 +144,7 @@ def build_payload_corpus(corpus: Corpus, ref_data: List[Text]) -> Dict:
         "name": corpus.name,
         "description": corpus.description,
         "suggestedFunctions": [f.value for f in corpus.functions],
+        "onboardingErrorsPolicy": error_handler.value,
         "tags": corpus.tags,
         "pricing": {"type": "FREE", "cost": 0},
         "privacy": corpus.privacy.value,
@@ -163,6 +166,7 @@ def build_payload_dataset(
     hypotheses_ref_data: Dict[Text, Any],
     meta_ref_data: Dict[Text, Any],
     tags: List[Text],
+    error_handler: ErrorHandler,
 ) -> Dict:
     """Generate onboard payload to coreengine
 
@@ -173,6 +177,7 @@ def build_payload_dataset(
         hypotheses_ref_data (Dict[Text, Any]): reference to existent hypotheses to the target data
         meta_ref_data (Dict[Text, Any]): reference to existent metadata
         tags (List[Text]): description tags
+        error_handler (ErrorHandler): how to handle failed rows
 
     Returns:
         Dict: onboard payload
@@ -188,6 +193,7 @@ def build_payload_dataset(
         "name": dataset.name,
         "description": dataset.description,
         "function": dataset.function.value,
+        "onboardingErrorsPolicy": error_handler.value,
         "tags": dataset.tags,
         "privacy": dataset.privacy.value,
         "license": {"typeId": dataset.license.value},

aixplain/utils/config.py

@@ -13,24 +13,17 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-import warnings
 import os
 import logging
-from dotenv import load_dotenv
 
 logger = logging.getLogger(__name__)
 
-load_dotenv()
 BACKEND_URL = os.getenv("BACKEND_URL", "https://platform-api.aixplain.com/")
 PIPELINES_RUN_URL = os.getenv("PIPELINES_RUN_URL", "https://platform-api.aixplain.com/assets/pipeline/execution/run")
 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", "")
-if TEAM_API_KEY == "" and AIXPLAIN_API_KEY == "":
-    logger.warning(
-        "'TEAM_API_KEY' has not been set properly. For help, please refer to the documentation(https://github.com/aixplain/aixplain#api-key-setup)"
-    )
 PIPELINE_API_KEY = os.getenv("PIPELINE_API_KEY", "")
 MODEL_API_KEY = os.getenv("MODEL_API_KEY", "")
 LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")