Enable api key parameter in data asset creation (#122)

Co-authored-by: Thiago Castro Ferreira <thiagocastroferreira@ip-192-168-0-3.ec2.internal>

aixplain/factories/corpus_factory.py

@@ -245,6 +245,7 @@ def create(
         functions: List[Function] = [],
         privacy: Privacy = Privacy.PRIVATE,
         error_handler: ErrorHandler = ErrorHandler.SKIP,
+        api_key: Optional[Text] = None
     ) -> Dict:
         """Asynchronous call to Upload a corpus to the user's dashboard.
 
@@ -259,6 +260,7 @@ def create(
             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.
+            api_key (Optional[Text]): team api key. Defaults to None.
 
         Returns:
             Dict: response dict
@@ -351,7 +353,7 @@ def create(
 
             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)
+            response = onboard_functions.create_data_asset(payload=corpus_payload, api_key=api_key)
             if response["success"] is True:
                 return_dict = {"status": response["status"], "asset_id": response["asset_id"]}
             else:

aixplain/factories/dataset_factory.py

@@ -280,8 +280,9 @@ def create(
         split_labels: Optional[List[Text]] = None,
         split_rate: Optional[List[float]] = None,
         error_handler: ErrorHandler = ErrorHandler.SKIP,
-        s3_link: Optional[str] = None,
-        aws_credentials: Optional[Dict[str, str]] = {"AWS_ACCESS_KEY_ID": None, "AWS_SECRET_ACCESS_KEY": None},
+        s3_link: Optional[Text] = None,
+        aws_credentials: Optional[Dict[Text, Text]] = {"AWS_ACCESS_KEY_ID": None, "AWS_SECRET_ACCESS_KEY": None},
+        api_key: Optional[Text] = None
     ) -> Dict:
         """Dataset Onboard
 
@@ -302,8 +303,9 @@ def create(
             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.
-            s3_link (Optional[str]): s3 url to files or directories
-            aws_credentials (Optional[Dict[str, str]]) : credentials for AWS and it should contains these two keys `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`
+            s3_link (Optional[Text]): s3 url to files or directories
+            aws_credentials (Optional[Dict[Text, Text]]) : credentials for AWS and it should contains these two keys `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`
+            api_key (Optional[Text]): team api key. Defaults to None.
         Returns:
             Dict: dataset onboard status
         """
@@ -485,7 +487,7 @@ def create(
             #     len(dataset_payload["output"]) > 0
             # ), "Data Asset Onboarding Error: Please specify the output data of your dataset."
 
-            response = onboard_functions.create_data_asset(dataset_payload, data_asset_type="dataset")
+            response = onboard_functions.create_data_asset(payload=dataset_payload, data_asset_type="dataset", api_key=api_key)
             if response["success"] is True:
                 return_dict = {"status": response["status"], "asset_id": response["asset_id"]}
             else:

aixplain/processes/data_onboarding/onboard_functions.py

@@ -288,17 +288,21 @@ def build_payload_dataset(
     return payload
 
 
-def create_data_asset(payload: Dict, data_asset_type: Text = "corpus") -> Dict:
+def create_data_asset(payload: Dict, data_asset_type: Text = "corpus", api_key: Optional[Text] = None) -> Dict:
     """Service to call onboard process in coreengine
 
     Args:
         payload (Dict): onboard payload
         data_asset_type (Text, optional): corpus or dataset. Defaults to "corpus".
+        api_key (Optional[Text]): team api key. Defaults to None.
 
     Returns:
         Dict: onboard status
     """
-    team_key = config.TEAM_API_KEY
+    if api_key is not None:
+        team_key = api_key
+    else:
+        team_key = config.TEAM_API_KEY
     headers = {"Authorization": "token " + team_key}
 
     url = urljoin(config.BACKEND_URL, f"sdk/{data_asset_type}/onboard")