aixplain
diff --git a/‎.gitignore
Lines changed: 3 additions & 0 deletions b/‎.gitignore
Lines changed: 3 additions & 0 deletions
diff --git a/‎aixplain/enums/response_status.py
Lines changed: 3 additions & 0 deletions b/‎aixplain/enums/response_status.py
Lines changed: 3 additions & 0 deletions
diff --git a/‎aixplain/factories/__init__.py
Lines changed: 1 addition & 0 deletions b/‎aixplain/factories/__init__.py
Lines changed: 1 addition & 0 deletions
diff --git a/‎aixplain/factories/agent_factory/__init__.py
Lines changed: 15 additions & 3 deletions b/‎aixplain/factories/agent_factory/__init__.py
Lines changed: 15 additions & 3 deletions
diff --git a/‎aixplain/factories/agent_factory/utils.py
Lines changed: 14 additions & 7 deletions b/‎aixplain/factories/agent_factory/utils.py
Lines changed: 14 additions & 7 deletions
diff --git a/‎aixplain/factories/index_factory.py
Lines changed: 46 additions & 0 deletions b/‎aixplain/factories/index_factory.py
Lines changed: 46 additions & 0 deletions
diff --git a/‎aixplain/factories/model_factory/__init__.py
Lines changed: 5 additions & 3 deletions b/‎aixplain/factories/model_factory/__init__.py
Lines changed: 5 additions & 3 deletions
diff --git a/‎aixplain/factories/model_factory/utils.py
Lines changed: 7 additions & 0 deletions b/‎aixplain/factories/model_factory/utils.py
Lines changed: 7 additions & 0 deletions
diff --git a/‎aixplain/factories/pipeline_factory/utils.py
Lines changed: 47 additions & 14 deletions b/‎aixplain/factories/pipeline_factory/utils.py
Lines changed: 47 additions & 14 deletions
diff --git a/‎aixplain/modules/__init__.py
Lines changed: 1 addition & 0 deletions b/‎aixplain/modules/__init__.py
Lines changed: 1 addition & 0 deletions
@@ -2,9 +2,12 @@
 __pycache__/
 *.py[cod]
 *$py.class
+.aixplain_cache/
 
+setup_env_ahmet.sh
 # C extensions
 *.so
+*.ipynb
 
 # Distribution / packaging
 .Python
 
@@ -29,3 +29,6 @@ class ResponseStatus(Text, Enum):
     IN_PROGRESS = "IN_PROGRESS"
     SUCCESS = "SUCCESS"
     FAILED = "FAILED"
+
+    def __str__(self):
+        return self.value
@@ -33,3 +33,4 @@
 from .finetune_factory import FinetuneFactory
 from .wallet_factory import WalletFactory
 from .api_key_factory import APIKeyFactory
+from .index_factory import IndexFactory
@@ -29,10 +29,12 @@
 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.agent.tool.python_interpreter_tool import PythonInterpreterTool
+from aixplain.modules.agent.tool.custom_python_code_tool import CustomPythonCodeTool
 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
+from typing import Callable, Dict, List, Optional, Text, Union
 
 from aixplain.utils.file_utils import _request_with_retry
 from urllib.parse import urljoin
@@ -88,8 +90,8 @@ def create(
         agent.validate()
         response = "Unspecified error"
         try:
-            logging.debug(f"Start service for POST Create Agent  - {url} - {headers} - {json.dumps(payload)}")
-            r = _request_with_retry("post", url, headers=headers, json=payload)
+            logging.debug(f"Start service for POST Create Agent  - {url} - {headers} - {json.dumps(agent.to_dict())}")
+            r = _request_with_retry("post", url, headers=headers, json=agent.to_dict())
             response = r.json()
         except Exception:
             raise Exception("Agent Onboarding Error: Please contact the administrators.")
@@ -136,6 +138,16 @@ def create_pipeline_tool(cls, description: Text, pipeline: Union[Pipeline, Text]
         """Create a new pipeline tool."""
         return PipelineTool(description=description, pipeline=pipeline)
 
+    @classmethod
+    def create_python_interpreter_tool(cls) -> PythonInterpreterTool:
+        """Create a new python interpreter tool."""
+        return PythonInterpreterTool()
+
+    @classmethod
+    def create_custom_python_code_tool(cls, code: Union[Text, Callable], description: Text = "") -> CustomPythonCodeTool:
+        """Create a new custom python code tool."""
+        return CustomPythonCodeTool(description=description, code=code)
+
     @classmethod
     def list(cls) -> Dict:
         """List all agents available in the platform."""
 
@@ -3,7 +3,8 @@
 import aixplain.utils.config as config
 from aixplain.enums import Function, Supplier
 from aixplain.enums.asset_status import AssetStatus
-from aixplain.modules.agent import Agent, ModelTool, PipelineTool
+from aixplain.modules.agent import Agent, ModelTool, PipelineTool, PythonInterpreterTool
+from aixplain.modules.agent.tool.custom_python_code_tool import CustomPythonCodeTool
 from typing import Dict, Text
 from urllib.parse import urljoin
 
@@ -18,12 +19,13 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
         if tool["type"] == "model":
             supplier = "aixplain"
             for supplier_ in Supplier:
-                if tool["supplier"] is not None and tool["supplier"].lower() in [
-                    supplier_.value["code"].lower(),
-                    supplier_.value["name"].lower(),
-                ]:
-                    supplier = supplier_
-                    break
+                if isinstance(tool["supplier"], str):
+                    if tool["supplier"] is not None and tool["supplier"].lower() in [
+                        supplier_.value["code"].lower(),
+                        supplier_.value["name"].lower(),
+                    ]:
+                        supplier = supplier_
+                        break
 
             tool = ModelTool(
                 function=Function(tool.get("function", None)),
@@ -34,6 +36,11 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
             )
         elif tool["type"] == "pipeline":
             tool = PipelineTool(description=tool["description"], pipeline=tool["assetId"])
+        elif tool["type"] == "utility":
+            if tool.get("utilityCode", None) is not None:
+                tool = CustomPythonCodeTool(description=tool["description"], code=tool["utilityCode"])
+            else:
+                tool = PythonInterpreterTool()
         else:
             raise Exception("Agent Creation Error: Tool type not supported.")
         tools.append(tool)
 
@@ -0,0 +1,46 @@
+from aixplain.modules.model.index_model import IndexModel
+from aixplain.factories import ModelFactory
+from aixplain.enums import Function, ResponseStatus, SortBy, SortOrder, OwnershipType, Supplier
+from typing import Optional, Text, Union, List, Tuple
+
+
+class IndexFactory(ModelFactory):
+    @classmethod
+    def create(cls, name: Text, description: Text) -> IndexModel:
+        """Create a new index collection"""
+        model = cls.get("66eae6656eb56311f2595011")
+
+        data = {"data": name, "description": description}
+        response = model.run(data=data)
+        if response.status == ResponseStatus.SUCCESS:
+            model_id = response.data
+            model = cls.get(model_id)
+            return model
+
+        error_message = f"Index Factory Exception: {response.error_message}"
+        if error_message == "":
+            error_message = "Index Factory Exception:An error occurred while creating the index collection."
+        raise Exception(error_message)
+
+    @classmethod
+    def list(
+        cls,
+        query: Optional[Text] = "",
+        suppliers: Optional[Union[Supplier, List[Supplier]]] = None,
+        ownership: Optional[Tuple[OwnershipType, List[OwnershipType]]] = None,
+        sort_by: Optional[SortBy] = None,
+        sort_order: SortOrder = SortOrder.ASCENDING,
+        page_number: int = 0,
+        page_size: int = 20,
+    ) -> List[IndexModel]:
+        """List all indexes"""
+        return super().list(
+            function=Function.SEARCH,
+            query=query,
+            suppliers=suppliers,
+            ownership=ownership,
+            sort_by=sort_by,
+            sort_order=sort_order,
+            page_number=page_number,
+            page_size=page_size,
+        )
@@ -48,6 +48,7 @@ def create_utility_model(
         inputs: List[UtilityModelInput] = [],
         description: Optional[Text] = None,
         output_examples: Text = "",
+        api_key: Optional[Text] = None,
     ) -> UtilityModel:
         """Create a utility model
 
@@ -57,24 +58,25 @@ def create_utility_model(
             description (Text, optional): description of the model
             inputs (List[UtilityModelInput], optional): inputs of the model
             output_examples (Text, optional): output examples
-
+            api_key (Text, optional): Team API key. Defaults to None.
         Returns:
             UtilityModel: created utility model
         """
+        api_key = config.TEAM_API_KEY if api_key is None else api_key
         utility_model = UtilityModel(
             id="",
             name=name,
             description=description,
             inputs=inputs,
             code=code,
             function=Function.UTILITIES,
-            api_key=config.TEAM_API_KEY,
+            api_key=api_key,
             output_examples=output_examples,
         )
         utility_model.validate()
         payload = utility_model.to_dict()
         url = urljoin(cls.backend_url, "sdk/utilities")
-        headers = {"x-api-key": f"{config.TEAM_API_KEY}", "Content-Type": "application/json"}
+        headers = {"x-api-key": f"{api_key}", "Content-Type": "application/json"}
         try:
             logging.info(f"Start service for POST Utility Model - {url} - {headers} - {payload}")
             r = _request_with_retry("post", url, headers=headers, json=payload)
 
@@ -2,6 +2,7 @@
 import logging
 from aixplain.modules.model import Model
 from aixplain.modules.model.llm_model import LLM
+from aixplain.modules.model.index_model import IndexModel
 from aixplain.modules.model.utility_model import UtilityModel, UtilityModelInput
 from aixplain.enums import DataType, Function, Language, OwnershipType, Supplier, SortBy, SortOrder
 from aixplain.utils import config
@@ -29,6 +30,10 @@ def create_model_from_response(response: Dict) -> Model:
         for param in response["params"]:
             if "language" in param["name"]:
                 parameters[param["name"]] = [w["value"] for w in param["values"]]
+            else:
+                values = [w["value"] for w in param["defaultValues"]]
+                if len(values) > 0:
+                    parameters[param["name"]] = values
 
     function_id = response["function"]["id"]
     function = Function(function_id)
@@ -43,6 +48,8 @@ def create_model_from_response(response: Dict) -> Model:
         f = [p for p in response.get("params", []) if p["name"] == "temperature"]
         if len(f) > 0 and len(f[0].get("defaultValues", [])) > 0:
             temperature = float(f[0]["defaultValues"][0]["value"])
+    elif function == Function.SEARCH:
+        ModelClass = IndexModel
     elif function == Function.UTILITIES:
         ModelClass = UtilityModel
         inputs = [
 
@@ -14,6 +14,8 @@
     Route,
     Script,
     Link,
+    BareSegmentor,
+    BareReconstructor,
 )
 from typing import Dict
 
@@ -32,7 +34,9 @@ def build_from_response(response: Dict, load_architecture: bool = False) -> Pipe
         response["api_key"] = config.TEAM_API_KEY
 
     # instantiating pipeline generic info
-    pipeline = Pipeline(response["id"], response["name"], response["api_key"])
+    pipeline = Pipeline(
+        id=response["id"], name=response["name"], api_key=response["api_key"], status=response.get("status", "draft")
+    )
     if load_architecture is True:
         try:
             # instantiating nodes
@@ -45,39 +49,68 @@ def build_from_response(response: Dict, load_architecture: bool = False) -> Pipe
                 elif node_json["type"].lower() == "asset":
                     if node_json["functionType"] == "metric":
                         node = BareMetric(asset_id=node_json["assetId"])
+                    elif node_json["functionType"] == "reconstructor":
+                        node = BareReconstructor(asset_id=node_json["assetId"])
+                    elif node_json["functionType"] == "segmentor":
+                        node = BareSegmentor(asset_id=node_json["assetId"])
                     else:
                         node = BareAsset(asset_id=node_json["assetId"])
-                elif node_json["type"].lower() == "segmentor":
-                    raise NotImplementedError()
-                elif node_json["type"].lower() == "reconstructor":
-                    raise NotImplementedError()
                 elif node_json["type"].lower() == "decision":
-                    node = Decision(routes=[Route(**route) for route in node_json["routes"]])
+                    node = Decision(
+                        routes=[Route(**route) for route in node_json["routes"]]
+                    )
                 elif node_json["type"].lower() == "router":
-                    node = Router(routes=[Route(**route) for route in node_json["routes"]])
+                    node = Router(
+                        routes=[Route(**route) for route in node_json["routes"]]
+                    )
                 elif node_json["type"].lower() == "script":
-                    node = Script(fileId=node_json["fileId"], fileMetadata=node_json["fileMetadata"])
+                    node = Script(
+                        fileId=node_json["fileId"],
+                        fileMetadata=node_json["fileMetadata"],
+                    )
                 elif node_json["type"].lower() == "output":
                     node = Output()
 
                 if "inputValues" in node_json:
                     [
                         node.inputs.create_param(
-                            data_type=DataType(input_param["dataType"]) if "dataType" in input_param else None,
+                            data_type=(
+                                DataType(input_param["dataType"])
+                                if "dataType" in input_param
+                                else None
+                            ),
                             code=input_param["code"],
-                            value=input_param["value"] if "value" in input_param else None,
-                            is_required=input_param["isRequired"] if "isRequired" in input_param else False,
+                            value=(
+                                input_param["value"] if "value" in input_param else None
+                            ),
+                            is_required=(
+                                input_param["isRequired"]
+                                if "isRequired" in input_param
+                                else False
+                            ),
                         )
                         for input_param in node_json["inputValues"]
                         if input_param["code"] not in node.inputs
                     ]
                 if "outputValues" in node_json:
                     [
                         node.outputs.create_param(
-                            data_type=DataType(output_param["dataType"]) if "dataType" in output_param else None,
+                            data_type=(
+                                DataType(output_param["dataType"])
+                                if "dataType" in output_param
+                                else None
+                            ),
                             code=output_param["code"],
-                            value=output_param["value"] if "value" in output_param else None,
-                            is_required=output_param["isRequired"] if "isRequired" in output_param else False,
+                            value=(
+                                output_param["value"]
+                                if "value" in output_param
+                                else None
+                            ),
+                            is_required=(
+                                output_param["isRequired"]
+                                if "isRequired" in output_param
+                                else False
+                            ),
                         )
                         for output_param in node_json["outputValues"]
                         if output_param["code"] not in node.outputs
 
@@ -37,3 +37,4 @@
 from .agent.tool import Tool
 from .team_agent import TeamAgent
 from .api_key import APIKey, APIKeyLimits, APIKeyUsageLimit
+from .model.index_model import IndexModel