Merge to test

.gitignore

@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# Vscode
+.vscode

aixplain/enums/data_type.py

@@ -24,7 +24,7 @@
 from enum import Enum
 
 
-class DataType(Enum):
+class DataType(str, Enum):
     AUDIO = "audio"
     FLOAT = "float"
     IMAGE = "image"
@@ -33,6 +33,8 @@ class DataType(Enum):
     TENSOR = "tensor"
     TEXT = "text"
     VIDEO = "video"
+    EMBEDDING = "embedding"
+    NUMBER = "number"
 
     def __str__(self):
         return self._value_

aixplain/enums/function.py

@@ -55,6 +55,7 @@ def load_functions():
                 if input_data_object["required"] is True
             },
             "output": {output_data_object["dataType"] for output_data_object in function["output"]},
+            "spec": function
         }
         for function in resp["items"]
     }

aixplain/factories/agent_factory/__init__.py

@@ -24,10 +24,13 @@
 import json
 import logging
 
+from aixplain.enums.function import Function
 from aixplain.enums.supplier import Supplier
 from aixplain.modules.agent import Agent, Tool
 from aixplain.modules.agent.tool.model_tool import ModelTool
 from aixplain.modules.agent.tool.pipeline_tool import PipelineTool
+from aixplain.modules.model import Model
+from aixplain.modules.pipeline import Pipeline
 from aixplain.utils import config
 from typing import Dict, List, Optional, Text, Union
 
@@ -64,11 +67,12 @@ def create(
                 if isinstance(tool, ModelTool):
                     tool_payload.append(
                         {
-                            "function": tool.function.value,
+                            "function": tool.function.value if tool.function is not None else None,
                             "type": "model",
                             "description": tool.description,
                             "supplier": tool.supplier.value["code"] if tool.supplier else None,
                             "version": tool.version if tool.version else None,
+                            "assetId": tool.model,
                         }
                     )
                 elif isinstance(tool, PipelineTool):
@@ -113,6 +117,32 @@ def create(
             raise Exception(e)
         return agent
 
+    @classmethod
+    def create_model_tool(
+        cls,
+        model: Optional[Union[Model, Text]] = None,
+        function: Optional[Union[Function, Text]] = None,
+        supplier: Optional[Union[Supplier, Text]] = None,
+    ) -> ModelTool:
+        """Create a new model tool."""
+        if function is not None and isinstance(function, str):
+            function = Function(function)
+
+        if supplier is not None:
+            if isinstance(supplier, str):
+                for supplier_ in Supplier:
+                    if supplier.lower() in [supplier.value["code"].lower(), supplier.value["name"].lower()]:
+                        supplier = supplier_
+                        break
+                if isinstance(supplier, str):
+                    supplier = None
+        return ModelTool(function=function, supplier=supplier, model=model)
+
+    @classmethod
+    def create_pipeline_tool(cls, description: Text, pipeline: Union[Pipeline, Text]) -> PipelineTool:
+        """Create a new pipeline tool."""
+        return PipelineTool(description=description, pipeline=pipeline)
+
     @classmethod
     def list(cls) -> Dict:
         """List all agents available in the platform."""

aixplain/factories/agent_factory/utils.py

@@ -22,9 +22,10 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
                     break
 
             tool = ModelTool(
-                function=Function(tool["function"]),
+                function=Function(tool["function"]) if tool["function"] is not None else None,
                 supplier=tool["supplier"],
                 version=tool["version"],
+                model=tool["assetId"],
             )
         elif tool["type"] == "pipeline":
             tool = PipelineTool(description=tool["description"], pipeline=tool["assetId"])

aixplain/factories/file_factory.py

@@ -104,7 +104,7 @@ def check_storage_type(cls, input_link: Any) -> StorageType:
             return StorageType.TEXT
 
     @classmethod
-    def to_link(cls, data: Union[Text, Dict]) -> Union[Text, Dict]:
+    def to_link(cls, data: Union[Text, Dict], **kwargs) -> Union[Text, Dict]:
         """If user input data is a local file, upload to aiXplain platform
 
         Args:
@@ -117,10 +117,10 @@ def to_link(cls, data: Union[Text, Dict]) -> Union[Text, Dict]:
             for key in data:
                 if isinstance(data[key], str):
                     if cls.check_storage_type(data[key]) == StorageType.FILE:
-                        data[key] = cls.upload(local_path=data[key])
+                        data[key] = cls.upload(local_path=data[key], **kwargs)
         elif isinstance(data, str):
             if cls.check_storage_type(data) == StorageType.FILE:
-                data = cls.upload(local_path=data)
+                data = cls.upload(local_path=data, **kwargs)
         return data
 
     @classmethod

aixplain/factories/model_factory.py

@@ -361,7 +361,7 @@ def create_asset_repo(
             if function_dict["name"] == function:
                 function_id = function_dict["id"]
         if function_id is None:
-            raise Exception("Invalid function name")
+            raise Exception(f"Invalid function name {function}")
         create_url = urljoin(config.BACKEND_URL, f"sdk/models/onboard")
         logging.debug(f"URL: {create_url}")
         if api_key:

aixplain/factories/pipeline_factory.py → aixplain/factories/pipeline_factory/__init__.py

@@ -21,9 +21,10 @@
     Pipeline Factory Class
 """
 import json
-import logging
 import os
+import logging
 from typing import Dict, List, Optional, Text, Union
+from aixplain.factories.pipeline_factory.utils import build_from_response
 from aixplain.enums.data_type import DataType
 from aixplain.enums.function import Function
 from aixplain.enums.supplier import Supplier
@@ -45,27 +46,6 @@ class PipelineFactory:
     aixplain_key = config.AIXPLAIN_API_KEY
     backend_url = config.BACKEND_URL
 
-    @classmethod
-    def __get_typed_nodes(cls, response: Dict, type: str) -> List[Dict]:
-        # read "nodes" field from response and return the nodes that are marked by "type": type
-        return [node for node in response["nodes"] if node["type"].lower() == type.lower()]
-
-    @classmethod
-    def __from_response(cls, response: Dict) -> Pipeline:
-        """Converts response Json to 'Pipeline' object
-
-        Args:
-            response (Dict): Json from API
-
-        Returns:
-            Pipeline: Coverted 'Pipeline' object
-        """
-        if "api_key" not in response:
-            response["api_key"] = config.TEAM_API_KEY
-        input = cls.__get_typed_nodes(response, "input")
-        output = cls.__get_typed_nodes(response, "output")
-        return Pipeline(response["id"], response["name"], response["api_key"], input=input, output=output)
-
     @classmethod
     def get(cls, pipeline_id: Text, api_key: Optional[Text] = None) -> Pipeline:
         """Create a 'Pipeline' object from pipeline id
@@ -81,28 +61,38 @@ def get(cls, pipeline_id: Text, api_key: Optional[Text] = None) -> Pipeline:
         try:
             url = urljoin(cls.backend_url, f"sdk/pipelines/{pipeline_id}")
             if api_key is not None:
-                headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
+                headers = {
+                    "Authorization": f"Token {api_key}",
+                    "Content-Type": "application/json",
+                }
             elif cls.aixplain_key != "":
-                headers = {"x-aixplain-key": f"{cls.aixplain_key}", "Content-Type": "application/json"}
+                headers = {
+                    "x-aixplain-key": f"{cls.aixplain_key}",
+                    "Content-Type": "application/json",
+                }
             else:
-                headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+                headers = {
+                    "Authorization": f"Token {config.TEAM_API_KEY}",
+                    "Content-Type": "application/json",
+                }
             logging.info(f"Start service for GET Pipeline  - {url} - {headers}")
             r = _request_with_retry("get", url, headers=headers)
             resp = r.json()
             # set api key
             resp["api_key"] = config.TEAM_API_KEY
             if api_key is not None:
                 resp["api_key"] = api_key
-            pipeline = cls.__from_response(resp)
+            pipeline = build_from_response(resp, load_architecture=True)
             return pipeline
-        except Exception:
+        except Exception as e:
+            logging.exception(e)
             status_code = 400
             if resp is not None and "statusCode" in resp:
                 status_code = resp["statusCode"]
                 message = resp["message"]
                 message = f"Pipeline Creation: Status {status_code} - {message}"
             else:
-                message = "Pipeline Creation: Unspecified Error"
+                message = f"Pipeline Creation: Unspecified Error {e}"
             logging.error(message)
             raise Exception(f"Status {status_code}: {message}")
 
@@ -127,14 +117,20 @@ def get_assets_from_page(cls, page_number: int) -> List[Pipeline]:
         try:
             url = urljoin(cls.backend_url, f"sdk/pipelines/?pageNumber={page_number}")
             if cls.aixplain_key != "":
-                headers = {"x-aixplain-key": f"{cls.aixplain_key}", "Content-Type": "application/json"}
+                headers = {
+                    "x-aixplain-key": f"{cls.aixplain_key}",
+                    "Content-Type": "application/json",
+                }
             else:
-                headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+                headers = {
+                    "Authorization": f"Token {config.TEAM_API_KEY}",
+                    "Content-Type": "application/json",
+                }
             r = _request_with_retry("get", url, headers=headers)
             resp = r.json()
             logging.info(f"Listing Pipelines: Status of getting Pipelines on Page {page_number}: {resp}")
             all_pipelines = resp["items"]
-            pipeline_list = [cls.__from_response(pipeline_info_json) for pipeline_info_json in all_pipelines]
+            pipeline_list = [build_from_response(pipeline_info_json) for pipeline_info_json in all_pipelines]
             return pipeline_list
         except Exception as e:
             error_message = f"Listing Pipelines: Error in getting Pipelines on Page {page_number}: {e}"
@@ -177,9 +173,15 @@ def list(
 
         url = urljoin(cls.backend_url, "sdk/pipelines/paginate")
         if cls.aixplain_key != "":
-            headers = {"x-aixplain-key": f"{cls.aixplain_key}", "Content-Type": "application/json"}
+            headers = {
+                "x-aixplain-key": f"{cls.aixplain_key}",
+                "Content-Type": "application/json",
+            }
         else:
-            headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+            headers = {
+                "Authorization": f"Token {config.TEAM_API_KEY}",
+                "Content-Type": "application/json",
+            }
 
         assert 0 < page_size <= 100, "Pipeline List Error: Page size must be greater than 0 and not exceed 100."
         payload = {
@@ -228,11 +230,43 @@ def list(
             total = resp["total"]
             logging.info(f"Response for POST List Pipeline - Page Total: {page_total} / Total: {total}")
             for pipeline in results:
-                pipelines.append(cls.__from_response(pipeline))
-        return {"results": pipelines, "page_total": page_total, "page_number": page_number, "total": total}
+                pipelines.append(build_from_response(pipeline))
+        return {
+            "results": pipelines,
+            "page_total": page_total,
+            "page_number": page_number,
+            "total": total,
+        }
 
     @classmethod
-    def create(cls, name: Text, pipeline: Union[Text, Dict], api_key: Optional[Text] = None) -> Pipeline:
+    def init(cls, name: Text, api_key: Optional[Text] = None) -> Pipeline:
+        """Initialize a new Pipeline
+
+        Args:
+            name (Text): Pipeline Name
+            api_key (Optional[Text], optional): Team API Key to create the Pipeline. Defaults to None.
+
+        Returns:
+            Pipeline: instance of the new pipeline
+        """
+        if api_key is None:
+            api_key = config.TEAM_API_KEY
+        return Pipeline(
+            id="",
+            name=name,
+            api_key=api_key,
+            nodes=[],
+            links=[],
+            instance=None,
+        )
+
+    @classmethod
+    def create(
+        cls,
+        name: Text,
+        pipeline: Union[Text, Dict],
+        api_key: Optional[Text] = None,
+    ) -> Pipeline:
         """Draft Pipeline Creation
 
         Args:
@@ -259,10 +293,17 @@ def create(cls, name: Text, pipeline: Union[Text, Dict], api_key: Optional[Text]
                 if "functionType" in node and node["functionType"] == "AI":
                     pipeline["nodes"][i]["functionType"] = pipeline["nodes"][i]["functionType"].lower()
             # prepare payload
-            payload = {"name": name, "status": "draft", "architecture": pipeline}
+            payload = {
+                "name": name,
+                "status": "draft",
+                "architecture": pipeline,
+            }
             url = urljoin(cls.backend_url, "sdk/pipelines")
             api_key = api_key if api_key is not None else config.TEAM_API_KEY
-            headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
+            headers = {
+                "Authorization": f"Token {api_key}",
+                "Content-Type": "application/json",
+            }
             logging.info(f"Start service for POST Create Pipeline - {url} - {headers} - {json.dumps(payload)}")
             r = _request_with_retry("post", url, headers=headers, json=payload)
             response = r.json()