ENG-1332: Started implementing new aixplain programmatic api

aixplain/__init__.py

@@ -20,24 +20,23 @@
 limitations under the License.
 """
 
-# Set default logging handler to avoid "No handler found" warnings.
 import os
 import logging
-from logging import NullHandler
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from .v2.core import Aixplain  # noqa
 
 LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO").upper()
 logging.basicConfig(level=LOG_LEVEL)
-logging.getLogger(__name__).addHandler(NullHandler())
 
-# Load Environment Variables
-from dotenv import load_dotenv
 
-load_dotenv()
+aixplain_v2 = None
+try:
+    aixplain_v2 = Aixplain()
+except Exception:
+    pass
 
-# 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)"
-    )
+__all__ = ["Aixplain", "aixplain_v2"]

aixplain/enums/supplier.py

@@ -54,6 +54,7 @@ def load_suppliers():
     suppliers = Enum(
         "Supplier", {clean_name(w["name"]): {"id": w["id"], "name": w["name"], "code": w["code"]} for w in resp}, type=dict
     )
+    suppliers.__str__ = lambda self: self.value["name"]
 
     return suppliers
 

aixplain/modules/agent/__init__.py

@@ -389,3 +389,6 @@ def deploy(self) -> None:
         assert self.status != AssetStatus.ONBOARDED, "Agent is already deployed."
         self.status = AssetStatus.ONBOARDED
         self.update()
+
+    def __repr__(self):
+        return f"Agent(id={self.id}, name={self.name}, function={self.function})"

aixplain/modules/pipeline/asset.py

@@ -112,9 +112,7 @@ def __polling(
         while not completed and (end - start) < timeout:
             try:
                 response_body = self.poll(poll_url, name=name)
-                logging.debug(
-                    f"Polling for Pipeline: Status of polling for {name} : {response_body}"
-                )
+                logging.debug(f"Polling for Pipeline: Status of polling for {name} : {response_body}")
                 completed = response_body["completed"]
 
                 end = time.time()
@@ -126,13 +124,9 @@ def __polling(
                 logging.error(f"Polling for Pipeline: polling for {name} : Continue")
         if response_body and response_body["status"] == "SUCCESS":
             try:
-                logging.debug(
-                    f"Polling for Pipeline: Final status of polling for {name} : SUCCESS - {response_body}"
-                )
+                logging.debug(f"Polling for Pipeline: Final status of polling for {name} : SUCCESS - {response_body}")
             except Exception:
-                logging.error(
-                    f"Polling for Pipeline: Final status of polling for {name} : ERROR - {response_body}"
-                )
+                logging.error(f"Polling for Pipeline: Final status of polling for {name} : ERROR - {response_body}")
         else:
             logging.error(
                 f"Polling for Pipeline: Final status of polling for {name} : No response in {timeout} seconds - {response_body}"
@@ -162,9 +156,7 @@ def poll(self, poll_url: Text, name: Text = "pipeline_process") -> Dict:
                     resp["data"] = json.loads(resp["data"])["response"]
                 except Exception:
                     resp = r.json()
-            logging.info(
-                f"Single Poll for Pipeline: Status of polling for {name} : {resp}"
-            )
+            logging.info(f"Single Poll for Pipeline: Status of polling for {name} : {resp}")
         except Exception:
             resp = {"status": "FAILED"}
         return resp
@@ -186,9 +178,7 @@ def _should_fallback_to_v2(self, response: Dict, version: str) -> bool:
         should_fallback = False
         if "status" not in response or response["status"] == "FAILED":
             should_fallback = True
-        elif response["status"] == "SUCCESS" and (
-            "data" not in response or not response["data"]
-        ):
+        elif response["status"] == "SUCCESS" and ("data" not in response or not response["data"]):
             should_fallback = True
         # Check for conditions that require a fallback
 
@@ -294,10 +284,7 @@ def __prepare_payload(
                 try:
                     payload = json.loads(data)
                     if isinstance(payload, dict) is False:
-                        if (
-                            isinstance(payload, int) is True
-                            or isinstance(payload, float) is True
-                        ):
+                        if isinstance(payload, int) is True or isinstance(payload, float) is True:
                             payload = str(payload)
                         payload = {"data": payload}
                 except Exception:
@@ -335,9 +322,7 @@ def __prepare_payload(
                         asset_payload["dataAsset"]["dataset_id"] = dasset.id
 
                         source_data_list = [
-                            dfield
-                            for dfield in dasset.source_data
-                            if dasset.source_data[dfield].id == data[node_label]
+                            dfield for dfield in dasset.source_data if dasset.source_data[dfield].id == data[node_label]
                         ]
 
                         if len(source_data_list) > 0:
@@ -410,9 +395,7 @@ def run_async(
         try:
             if 200 <= r.status_code < 300:
                 resp = r.json()
-                logging.info(
-                    f"Result of request for {name}  - {r.status_code} - {resp}"
-                )
+                logging.info(f"Result of request for {name}  - {r.status_code} - {resp}")
                 poll_url = resp["url"]
                 response = {"status": "IN_PROGRESS", "url": poll_url}
             else:
@@ -428,7 +411,9 @@ def run_async(
                     error = "Validation-related error: Please ensure all required fields are provided and correctly formatted."
                 else:
                     status_code = str(r.status_code)
-                    error = f"Status {status_code}: Unspecified error: An unspecified error occurred while processing your request."
+                    error = (
+                        f"Status {status_code}: Unspecified error: An unspecified error occurred while processing your request."
+                    )
                 response = {"status": "FAILED", "error_message": error}
                 logging.error(f"Error in request for {name} - {r.status_code}: {error}")
         except Exception:
@@ -477,9 +462,7 @@ def update(
 
             for i, node in enumerate(pipeline["nodes"]):
                 if "functionType" in node:
-                    pipeline["nodes"][i]["functionType"] = pipeline["nodes"][i][
-                        "functionType"
-                    ].lower()
+                    pipeline["nodes"][i]["functionType"] = pipeline["nodes"][i]["functionType"].lower()
             # prepare payload
             status = "draft"
             if save_as_asset is True:
@@ -497,9 +480,7 @@ def update(
                 "Authorization": f"Token {api_key}",
                 "Content-Type": "application/json",
             }
-            logging.info(
-                f"Start service for PUT Update Pipeline - {url} - {headers} - {json.dumps(payload)}"
-            )
+            logging.info(f"Start service for PUT Update Pipeline - {url} - {headers} - {json.dumps(payload)}")
             r = _request_with_retry("put", url, headers=headers, json=payload)
             response = r.json()
             logging.info(f"Pipeline {response['id']} Updated.")
@@ -554,9 +535,7 @@ def save(
 
             for i, node in enumerate(pipeline["nodes"]):
                 if "functionType" in node:
-                    pipeline["nodes"][i]["functionType"] = pipeline["nodes"][i][
-                        "functionType"
-                    ].lower()
+                    pipeline["nodes"][i]["functionType"] = pipeline["nodes"][i]["functionType"].lower()
             # prepare payload
             status = "draft"
             if save_as_asset is True:
@@ -589,3 +568,6 @@ def deploy(self, api_key: Optional[Text] = None) -> None:
         pipeline = self.to_dict()
         self.update(pipeline=pipeline, save_as_asset=True, api_key=api_key, name=self.name)
         self.status = AssetStatus.ONBOARDED
+
+    def __repr__(self):
+        return f"Pipeline(id={self.id}, name={self.name})"