Development 2 Test

aixplain/factories/agent_factory/utils.py

@@ -19,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)),

aixplain/factories/pipeline_factory/utils.py

@@ -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

aixplain/modules/agent/__init__.py

@@ -92,6 +92,8 @@ def __init__(
         super().__init__(id, name, description, api_key, supplier, version, cost=cost)
         self.additional_info = additional_info
         self.tools = tools
+        for i, _ in enumerate(tools):
+            self.tools[i].api_key = api_key
         self.llm_id = llm_id
         if isinstance(status, str):
             try:
@@ -110,7 +112,7 @@ def validate(self) -> None:
         ), "Agent Creation Error: Agent name must not contain special characters."
 
         try:
-            llm = ModelFactory.get(self.llm_id)
+            llm = ModelFactory.get(self.llm_id, api_key=self.api_key)
             assert llm.function == Function.TEXT_GENERATION, "Large Language Model must be a text generation model."
         except Exception:
             raise Exception(f"Large Language Model with ID '{self.llm_id}' not found.")
@@ -307,19 +309,19 @@ def delete(self) -> None:
                 message = f"Agent Deletion Error (HTTP {r.status_code}): There was an error in deleting the agent."
             logging.error(message)
             raise Exception(f"{message}")
-        
+
     def update(self) -> None:
         """Update agent."""
         import warnings
         import inspect
+
         # Get the current call stack
         stack = inspect.stack()
-        if len(stack) > 2 and stack[1].function != 'save':
+        if len(stack) > 2 and stack[1].function != "save":
             warnings.warn(
-                "update() is deprecated and will be removed in a future version. "
-                "Please use save() instead.",
+                "update() is deprecated and will be removed in a future version. " "Please use save() instead.",
                 DeprecationWarning,
-                stacklevel=2
+                stacklevel=2,
             )
         from aixplain.factories.agent_factory.utils import build_agent
 
@@ -343,10 +345,9 @@ def update(self) -> None:
             error_msg = f"Agent Update Error (HTTP {r.status_code}): {resp}"
             raise Exception(error_msg)
 
-
     def save(self) -> None:
         """Save the Agent."""
-        self.update() 
+        self.update()
 
     def deploy(self) -> None:
         assert self.status == AssetStatus.DRAFT, "Agent must be in draft status to be deployed."

aixplain/modules/agent/tool/__init__.py

@@ -22,6 +22,7 @@
 """
 from abc import ABC
 from typing import Optional, Text
+from aixplain.utils import config
 
 
 class Tool(ABC):
@@ -38,6 +39,7 @@ def __init__(
         name: Text,
         description: Text,
         version: Optional[Text] = None,
+        api_key: Optional[Text] = config.TEAM_API_KEY,
         **additional_info,
     ) -> None:
         """Specialized software or resource designed to assist the AI in executing specific tasks or functions based on user commands.
@@ -46,10 +48,12 @@ def __init__(
             name (Text): name of the tool
             description (Text): descriptiion of the tool
             version (Text): version of the tool
+            api_key (Text): api key of the tool. Defaults to config.TEAM_API_KEY.
         """
         self.name = name
         self.description = description
         self.version = version
+        self.api_key = api_key
         self.additional_info = additional_info
 
     def to_dict(self):

aixplain/modules/agent/tool/model_tool.py

@@ -108,7 +108,7 @@ def validate(self) -> Model:
         try:
             model = None
             if self.model is not None:
-                model = ModelFactory.get(self.model)
+                model = ModelFactory.get(self.model, api_key=self.api_key)
             return model
         except Exception:
             raise Exception(f"Model Tool Unavailable. Make sure Model '{self.model}' exists or you have access to it.")

aixplain/modules/agent/tool/pipeline_tool.py

@@ -62,6 +62,6 @@ def validate(self):
         from aixplain.factories.pipeline_factory import PipelineFactory
 
         try:
-            PipelineFactory.get(self.pipeline)
+            PipelineFactory.get(self.pipeline, api_key=self.api_key)
         except Exception:
             raise Exception(f"Pipeline Tool Unavailable. Make sure Pipeline '{self.pipeline}' exists or you have access to it.")

aixplain/modules/model/index_model.py

@@ -3,7 +3,7 @@
 from aixplain.utils import config
 from aixplain.modules.model.response import ModelResponse
 from typing import Text, Optional, Union, Dict
-from aixplain.modules.model.document_index import DocumentIndex
+from aixplain.modules.model.record import Record
 from typing import List
 
 
@@ -51,30 +51,18 @@ def __init__(
         self.url = config.MODELS_RUN_URL
         self.backend_url = config.BACKEND_URL
 
-    def search(self, query: str, top_k: int = 10) -> ModelResponse:
-        data = {"action": "search", "data": query, "payload": {"filters": {}, "top_k": top_k}}
+    def search(self, query: str, top_k: int = 10, filters: Dict = {}) -> ModelResponse:
+        data = {"action": "search", "data": query, "payload": {"filters": filters, "top_k": top_k}}
         return self.run(data=data)
 
-    def add(self, documents: List[DocumentIndex]) -> ModelResponse:
+    def upsert(self, documents: List[Record]) -> ModelResponse:
         payloads = [doc.to_dict() for doc in documents]
         data = {"action": "ingest", "data": "", "payload": {"payloads": payloads}}
         response = self.run(data=data)
         if response.status == ResponseStatus.SUCCESS:
             response.data = payloads
             return response
-        raise Exception(f"Failed to add documents: {response.error_message}")
-
-    def update(self, documents: List[DocumentIndex]) -> ModelResponse:
-        payloads = [
-            {"value": doc.value, "value_type": doc.value_type, "id": str(doc.id), "uri": doc.uri, "attributes": doc.attributes}
-            for doc in documents
-        ]
-        data = {"action": "update", "data": "", "payload": {"payloads": payloads}}
-        response = self.run(data=data)
-        if response.status == ResponseStatus.SUCCESS:
-            response.data = payloads
-            return response
-        raise Exception(f"Failed to update documents: {response.error_message}")
+        raise Exception(f"Failed to upsert documents: {response.error_message}")
 
     def count(self) -> float:
         data = {"action": "count", "data": ""}

aixplain/modules/model/document_index.py → aixplain/modules/model/record.py

@@ -2,7 +2,7 @@
 from uuid import uuid4
 
 
-class DocumentIndex:
+class Record:
     def __init__(self, value: str, value_type: str = "text", id: Optional[str] = None, uri: str = "", attributes: dict = {}):
         self.value = value
         self.value_type = value_type

aixplain/modules/model/response.py

@@ -22,6 +22,10 @@ def __init__(
         self.data = data
         self.details = details
         self.completed = completed
+        if error_message == "":
+            error_message = kwargs.get("error", "")
+            if "supplierError" in kwargs:
+                error_message = f"{error_message} - {kwargs.get('supplierError', '')}"
         self.error_message = error_message
         self.used_credits = used_credits
         self.run_time = run_time

aixplain/modules/model/utility_model.py

@@ -115,18 +115,39 @@ def __init__(
         self.output_examples = output_examples
 
     def validate(self):
-        self.code, inputs, description = parse_code(self.code)
+        """Validate the Utility Model."""
+        description = None
+        inputs = []
+        # check if the model exists and if the code is strring with s3:// 
+        # if not, parse the code and update the description and inputs and do the validation
+        # if yes, just do the validation on the description and inputs
+        if not (self._model_exists() and str(self.code).startswith("s3://")):
+            self.code, inputs, description = parse_code(self.code)
+            if self.description is None:
+                self.description = description
+            if len(self.inputs) == 0:
+                self.inputs = inputs
+            for input in self.inputs:
+                input.validate()
+        else:
+            logging.info("Utility Model Already Exists, skipping code validation")
+
         assert description is not None or self.description is not None, "Utility Model Error: Model description is required"
-        if self.description is None:
-            self.description = description
-        if len(self.inputs) == 0:
-            self.inputs = inputs
-        for input in self.inputs:
-            input.validate()
         assert self.name and self.name.strip() != "", "Name is required"
         assert self.description and self.description.strip() != "", "Description is required"
         assert self.code and self.code.strip() != "", "Code is required"
 
+    def _model_exists(self):
+        if self.id is None or self.id == "":
+            return False
+        url = urljoin(self.backend_url, f"sdk/models/{self.id}")
+        headers = {"Authorization": f"Token {self.api_key}", "Content-Type": "application/json"}
+        logging.info(f"Start service for GET Model  - {url} - {headers}")
+        r = _request_with_retry("get", url, headers=headers)
+        if r.status_code != 200:
+            raise Exception()
+        return True
+
     def to_dict(self):
         return {
             "name": self.name,

aixplain/modules/model/utils.py

@@ -102,7 +102,6 @@ def parse_code(code: Union[Text, Callable]) -> Tuple[Text, List, Text]:
         str_code = requests.get(code).text
     else:
         str_code = code
-
     # assert str_code has a main function
     if "def main(" not in str_code:
         raise Exception("Utility Model Error: Code must have a main function")