Merge to prod

aixplain/factories/__init__.py

@@ -21,6 +21,7 @@
 """
 from .asset_factory import AssetFactory
 from .agent_factory import AgentFactory
+from .team_agent_factory import TeamAgentFactory
 from .benchmark_factory import BenchmarkFactory
 from .corpus_factory import CorpusFactory
 from .data_factory import DataFactory

aixplain/factories/agent_factory/__init__.py

@@ -34,7 +34,7 @@
 from aixplain.utils import config
 from typing import Dict, List, Optional, Text, Union
 
-from aixplain.factories.agent_factory.utils import build_agent, validate_llm
+from aixplain.factories.agent_factory.utils import build_agent, validate_llm, validate_name
 from aixplain.utils.file_utils import _request_with_retry
 from urllib.parse import urljoin
 
@@ -44,36 +44,31 @@ class AgentFactory:
     def create(
         cls,
         name: Text,
-        llm_id: Text,
+        description: Text,
+        llm_id: Text = "669a63646eb56306647e1091",
         tools: List[Tool] = [],
-        description: Text = "",
         api_key: Text = config.TEAM_API_KEY,
         supplier: Union[Dict, Text, Supplier, int] = "aiXplain",
         version: Optional[Text] = None,
-        use_mentalist_and_inspector: bool = False,
     ) -> Agent:
         """Create a new agent in the platform.
 
         Args:
             name (Text): name of the agent
-            llm_id (Text): aiXplain ID of the large language model to be used as agent.
+            description (Text): description of the agent role.
+            llm_id (Text, optional): aiXplain ID of the large language model to be used as agent. Defaults to "669a63646eb56306647e1091" (GPT-4o mini).
             tools (List[Tool], optional): list of tool for the agent. Defaults to [].
-            description (Text, optional): description of the agent role. Defaults to "".
             api_key (Text, optional): team/user API key. Defaults to config.TEAM_API_KEY.
             supplier (Union[Dict, Text, Supplier, int], optional): owner of the agent. Defaults to "aiXplain".
             version (Optional[Text], optional): version of the agent. Defaults to None.
-            use_mentalist_and_inspector (bool, optional): flag to enable mentalist and inspector agents (which only works when a supervisor is enabled). Defaults to False.
 
         Returns:
             Agent: created Agent
         """
+        validate_name(name)
         # validate LLM ID
         validate_llm(llm_id)
 
-        orchestrator_llm_id, mentalist_and_inspector_llm_id = llm_id, None
-        if use_mentalist_and_inspector is True:
-            mentalist_and_inspector_llm_id = llm_id
-
         try:
             agent = None
             url = urljoin(config.BACKEND_URL, "sdk/agents")
@@ -117,8 +112,6 @@ def create(
                 "supplier": supplier,
                 "version": version,
                 "llmId": llm_id,
-                "supervisorId": orchestrator_llm_id,
-                "plannerId": mentalist_and_inspector_llm_id,
             }
 
             logging.info(f"Start service for POST Create Agent  - {url} - {headers} - {json.dumps(payload)}")
@@ -128,13 +121,13 @@ def create(
                 agent = build_agent(payload=response, api_key=api_key)
             else:
                 error = r.json()
-                error_msg = "Agent Onboarding Error: Please contant the administrators."
+                error_msg = "Agent Onboarding Error: Please contact the administrators."
                 if "message" in error:
                     msg = error["message"]
                     if error["message"] == "err.name_already_exists":
                         msg = "Agent name already exists."
                     elif error["message"] == "err.asset_is_not_available":
-                        msg = "Some the tools are not available."
+                        msg = "Some tools are not available."
                     error_msg = f"Agent Onboarding Error (HTTP {r.status_code}): {msg}"
                 logging.exception(error_msg)
                 raise Exception(error_msg)
@@ -190,7 +183,7 @@ def list(cls) -> Dict:
                     agents.append(build_agent(agent))
                 return {"results": agents, "page_total": page_total, "page_number": 0, "total": total}
             else:
-                error_msg = "Agent Listing Error: Please contant the administrators."
+                error_msg = "Agent Listing Error: Please contact the administrators."
                 if "message" in resp:
                     msg = resp["message"]
                     error_msg = f"Agent Listing Error (HTTP {r.status_code}): {msg}"
@@ -214,7 +207,7 @@ def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent:
         if 200 <= r.status_code < 300:
             return build_agent(resp)
         else:
-            msg = "Please contant the administrators."
+            msg = "Please contact the administrators."
             if "message" in resp:
                 msg = resp["message"]
             error_msg = f"Agent Get Error (HTTP {r.status_code}): {msg}"

aixplain/factories/agent_factory/utils.py

@@ -7,6 +7,8 @@
 from typing import Dict, Text
 from urllib.parse import urljoin
 
+GPT_4o_ID = "6646261c6eb563165658bbb1"
+
 
 def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
     """Instantiate a new agent in the platform."""
@@ -41,7 +43,7 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
         supplier=payload["teamId"] if "teamId" in payload else None,
         version=payload["version"] if "version" in payload else None,
         cost=payload["cost"] if "cost" in payload else None,
-        llm_id=payload["llmId"] if "llmId" in payload else "6646261c6eb563165658bbb1",
+        llm_id=payload["llmId"] if "llmId" in payload else GPT_4o_ID,
         api_key=api_key,
         status=AssetStatus(payload["status"]),
     )
@@ -57,3 +59,11 @@ def validate_llm(model_id: Text) -> None:
         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 '{model_id}' not found.")
+
+
+def validate_name(name: Text) -> None:
+    import re
+
+    assert (
+        re.match("^[a-zA-Z0-9 ]*$", name) is not None
+    ), "Agent Creation Error: Agent name must not contain special characters."

aixplain/factories/pipeline_factory/__init__.py

@@ -290,7 +290,7 @@ def create(
                     pipeline = json.load(f)
 
             for i, node in enumerate(pipeline["nodes"]):
-                if "functionType" in node and node["functionType"] == "AI":
+                if "functionType" in node:
                     pipeline["nodes"][i]["functionType"] = pipeline["nodes"][i]["functionType"].lower()
             # prepare payload
             payload = {

aixplain/factories/pipeline_factory/utils.py

@@ -7,7 +7,8 @@
 from aixplain.modules.pipeline.designer import (
     Input,
     Output,
-    AssetNode,
+    BareAsset,
+    BareMetric,
     Decision,
     Router,
     Route,
@@ -36,14 +37,16 @@ def build_from_response(response: Dict, load_architecture: bool = False) -> Pipe
         try:
             # instantiating nodes
             for node_json in response["nodes"]:
-                print(node_json)
                 if node_json["type"].lower() == "input":
                     node = Input(
                         data=node_json["data"] if "data" in node_json else None,
                         data_types=[DataType(dt) for dt in node_json["dataType"]],
                     )
                 elif node_json["type"].lower() == "asset":
-                    node = AssetNode(asset_id=node_json["assetId"])
+                    if node_json["functionType"] == "metric":
+                        node = BareMetric(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":
@@ -53,7 +56,7 @@ def build_from_response(response: Dict, load_architecture: bool = False) -> Pipe
                 elif node_json["type"].lower() == "router":
                     node = Router(routes=[Route(**route) for route in node_json["routes"]])
                 elif node_json["type"].lower() == "script":
-                    node = Script(fileId=node_json["fileId"])
+                    node = Script(fileId=node_json["fileId"], fileMetadata=node_json["fileMetadata"])
                 elif node_json["type"].lower() == "output":
                     node = Output()
 

aixplain/factories/script_factory.py

@@ -8,15 +8,12 @@
 
 
 class ScriptFactory:
-
     @classmethod
     def upload_script(cls, script_path: str) -> Tuple[str, str]:
         try:
             url = f"{config.BACKEND_URL}/sdk/pipelines/script"
             headers = {"Authorization": f"Token {config.TEAM_API_KEY}"}
-            r = requests.post(
-                url, headers=headers, files={"file": open(script_path, "rb")}
-            )
+            r = requests.post(url, headers=headers, files={"file": open(script_path, "rb")})
             if 200 <= r.status_code < 300:
                 response = r.json()
             else:
@@ -26,6 +23,6 @@ def upload_script(cls, script_path: str) -> Tuple[str, str]:
 
         # get metadata info
         fname = os.path.splitext(os.path.basename(script_path))[0]
-        file_size_kb = int(os.path.getsize(script_path) / 1024)
-        metadata = json.dumps({"name": fname, "size": file_size_kb})
+        file_size = int(os.path.getsize(script_path))
+        metadata = json.dumps({"name": fname, "size": file_size})
         return response["fileId"], metadata

aixplain/factories/team_agent_factory/__init__.py

@@ -0,0 +1,166 @@
+__author__ = "lucaspavanelli"
+
+"""
+Copyright 2024 The aiXplain SDK authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Author: Thiago Castro Ferreira and Lucas Pavanelli
+Date: August 15th 2024
+Description:
+    TeamAgent Factory Class
+"""
+
+import json
+import logging
+
+from aixplain.enums.supplier import Supplier
+from aixplain.factories.agent_factory import AgentFactory
+from aixplain.factories.agent_factory.utils import validate_llm, validate_name
+from aixplain.modules.agent import Agent
+from aixplain.modules.team_agent import TeamAgent
+from aixplain.utils import config
+from aixplain.factories.team_agent_factory.utils import build_team_agent
+from aixplain.utils.file_utils import _request_with_retry
+from typing import Dict, List, Optional, Text, Union
+from urllib.parse import urljoin
+
+
+class TeamAgentFactory:
+    @classmethod
+    def create(
+        cls,
+        name: Text,
+        agents: List[Union[Text, Agent]],
+        llm_id: Text = "669a63646eb56306647e1091",
+        description: Text = "",
+        api_key: Text = config.TEAM_API_KEY,
+        supplier: Union[Dict, Text, Supplier, int] = "aiXplain",
+        version: Optional[Text] = None,
+        use_mentalist_and_inspector: bool = True,
+    ) -> TeamAgent:
+        """Create a new team agent in the platform."""
+        validate_name(name)
+        # validate LLM ID
+        validate_llm(llm_id)
+        assert len(agents) > 0, "TeamAgent Onboarding Error: At least one agent must be provided."
+        for agent in agents:
+            if isinstance(agent, Text) is True:
+                try:
+                    agent = AgentFactory.get(agent)
+                except Exception:
+                    raise Exception(f"TeamAgent Onboarding Error: Agent {agent} does not exist.")
+            else:
+                assert isinstance(agent, Agent), "TeamAgent Onboarding Error: Agents must be instances of Agent class"
+
+        mentalist_and_inspector_llm_id = None
+        if use_mentalist_and_inspector is True:
+            mentalist_and_inspector_llm_id = llm_id
+        try:
+            team_agent = None
+            url = urljoin(config.BACKEND_URL, "sdk/agent-communities")
+            headers = {"x-api-key": api_key}
+
+            if isinstance(supplier, dict):
+                supplier = supplier["code"]
+            elif isinstance(supplier, Supplier):
+                supplier = supplier.value["code"]
+
+            agent_list = []
+            for idx, agent in enumerate(agents):
+                agent_list.append({"assetId": agent.id, "number": idx, "type": "AGENT", "label": "AGENT"})
+
+            payload = {
+                "name": name,
+                "agents": agent_list,
+                "links": [],
+                "description": description,
+                "llmId": llm_id,
+                "supervisorId": llm_id,
+                "plannerId": mentalist_and_inspector_llm_id,
+                "supplier": supplier,
+                "version": version,
+            }
+
+            logging.info(f"Start service for POST Create TeamAgent  - {url} - {headers} - {json.dumps(payload)}")
+            r = _request_with_retry("post", url, headers=headers, json=payload)
+            if 200 <= r.status_code < 300:
+                response = r.json()
+                team_agent = build_team_agent(payload=response, api_key=api_key)
+            else:
+                error = r.json()
+                error_msg = "TeamAgent Onboarding Error: Please contact the administrators."
+                if "message" in error:
+                    msg = error["message"]
+                    if error["message"] == "err.name_already_exists":
+                        msg = "TeamAgent name already exists."
+                    elif error["message"] == "err.asset_is_not_available":
+                        msg = "Some tools are not available."
+                    error_msg = f"TeamAgent Onboarding Error (HTTP {r.status_code}): {msg}"
+                logging.exception(error_msg)
+                raise Exception(error_msg)
+        except Exception as e:
+            raise Exception(e)
+        return team_agent
+
+    @classmethod
+    def list(cls) -> Dict:
+        """List all agents available in the platform."""
+        url = urljoin(config.BACKEND_URL, "sdk/agent-communities")
+        headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"}
+
+        payload = {}
+        logging.info(f"Start service for GET List Agents - {url} - {headers} - {json.dumps(payload)}")
+        try:
+            r = _request_with_retry("get", url, headers=headers)
+            resp = r.json()
+
+            if 200 <= r.status_code < 300:
+                agents, page_total, total = [], 0, 0
+                results = resp
+                page_total = len(results)
+                total = len(results)
+                logging.info(f"Response for GET List Agents - Page Total: {page_total} / Total: {total}")
+                for agent in results:
+                    agents.append(build_team_agent(agent))
+                return {"results": agents, "page_total": page_total, "page_number": 0, "total": total}
+            else:
+                error_msg = "Agent Listing Error: Please contact the administrators."
+                if "message" in resp:
+                    msg = resp["message"]
+                    error_msg = f"Agent Listing Error (HTTP {r.status_code}): {msg}"
+                logging.exception(error_msg)
+                raise Exception(error_msg)
+        except Exception as e:
+            raise Exception(e)
+
+    @classmethod
+    def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent:
+        """Get agent by id."""
+        url = urljoin(config.BACKEND_URL, f"sdk/agent-communities/{agent_id}")
+        if config.AIXPLAIN_API_KEY != "":
+            headers = {"x-aixplain-key": f"{config.AIXPLAIN_API_KEY}", "Content-Type": "application/json"}
+        else:
+            api_key = api_key if api_key is not None else config.TEAM_API_KEY
+            headers = {"x-api-key": api_key, "Content-Type": "application/json"}
+        logging.info(f"Start service for GET Agent  - {url} - {headers}")
+        r = _request_with_retry("get", url, headers=headers)
+        resp = r.json()
+        if 200 <= r.status_code < 300:
+            return build_team_agent(resp)
+        else:
+            msg = "Please contact the administrators."
+            if "message" in resp:
+                msg = resp["message"]
+            error_msg = f"Agent Get Error (HTTP {r.status_code}): {msg}"
+            raise Exception(error_msg)