Development 2 Test

.pre-commit-config.yaml

@@ -21,5 +21,5 @@ repos:
     rev: v2.0.0 # Use the latest version
     hooks:
       - id: flake8
-        args: # arguments to configure black
-          - --ignore=E402,E501
+        args: # arguments to configure flake8
+          - --ignore=E402,E501,E203

aixplain/enums/status.py

@@ -0,0 +1,8 @@
+from enum import Enum
+from typing import Text
+
+
+class Status(Text, Enum):
+    FAILED = "failed"
+    IN_PROGRESS = "in_progress"
+    SUCCESS = "success"

aixplain/factories/agent_factory/__init__.py

@@ -121,8 +121,8 @@ def create(
             "status": "draft",
             "tasks": [task.to_dict() for task in tasks],
         }
-        agent = build_agent(payload=payload, api_key=api_key)
-        agent.validate()
+        agent = build_agent(payload=payload, tools=tools, api_key=api_key)
+        agent.validate(raise_exception=True)
         response = "Unspecified error"
         try:
             logging.debug(f"Start service for POST Create Agent  - {url} - {headers} - {json.dumps(agent.to_dict())}")
@@ -132,7 +132,7 @@ def create(
             raise Exception("Agent Onboarding Error: Please contact the administrators.")
 
         if 200 <= r.status_code < 300:
-            agent = build_agent(payload=response, api_key=api_key)
+            agent = build_agent(payload=response, tools=tools, api_key=api_key)
         else:
             error_msg = f"Agent Onboarding Error: {response}"
             if "message" in response:

aixplain/factories/agent_factory/utils.py

@@ -1,69 +1,93 @@
 __author__ = "thiagocastroferreira"
 
+import logging
 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
+from aixplain.modules.agent.tool import Tool
 from aixplain.modules.agent.agent_task import AgentTask
 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.agent.tool.sql_tool import SQLTool
-from typing import Dict, Text
+from typing import Dict, Text, List
 from urllib.parse import urljoin
 
 GPT_4o_ID = "6646261c6eb563165658bbb1"
 
 
-def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
+def build_tool(tool: Dict):
+    """Build a tool from a dictionary.
+
+    Args:
+        tool (Dict): Tool dictionary.
+
+    Returns:
+        Tool: Tool object.
+    """
+    if tool["type"] == "model":
+        supplier = "aixplain"
+        for supplier_ in Supplier:
+            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)),
+            supplier=supplier,
+            version=tool["version"],
+            model=tool["assetId"],
+            description=tool.get("description", ""),
+            parameters=tool.get("parameters", None),
+        )
+    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()
+    elif tool["type"] == "sql":
+        parameters = {parameter["name"]: parameter["value"] for parameter in tool.get("parameters", [])}
+        database = parameters.get("database")
+        schema = parameters.get("schema")
+        tables = parameters.get("tables", None)
+        tables = tables.split(",") if tables is not None else None
+        enable_commit = parameters.get("enable_commit", False)
+        tool = SQLTool(
+            description=tool["description"], database=database, schema=schema, tables=tables, enable_commit=enable_commit
+        )
+    else:
+        raise Exception("Agent Creation Error: Tool type not supported.")
+
+    return tool
+
+
+def build_agent(payload: Dict, tools: List[Tool] = None, api_key: Text = config.TEAM_API_KEY) -> Agent:
     """Instantiate a new agent in the platform."""
     tools_dict = payload["assets"]
-    tools = []
-    for tool in tools_dict:
-        if tool["type"] == "model":
-            supplier = "aixplain"
-            for supplier_ in Supplier:
-                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)),
-                supplier=supplier,
-                version=tool["version"],
-                model=tool["assetId"],
-                description=tool.get("description", ""),
-                parameters=tool.get("parameters", None),
-            )
-        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()
-        elif tool["type"] == "sql":
-            parameters = {parameter["name"]: parameter["value"] for parameter in tool.get("parameters", [])}
-            database = parameters.get("database")
-            schema = parameters.get("schema")
-            tables = parameters.get("tables", None)
-            tables = tables.split(",") if tables is not None else None
-            enable_commit = parameters.get("enable_commit", False)
-            tool = SQLTool(
-                description=tool["description"], database=database, schema=schema, tables=tables, enable_commit=enable_commit
-            )
-        else:
-            raise Exception("Agent Creation Error: Tool type not supported.")
-        tools.append(tool)
+    payload_tools = tools
+    if payload_tools is None:
+        payload_tools = []
+        for tool in tools_dict:
+            try:
+                payload_tools.append(build_tool(tool))
+            except Exception:
+                logging.warning(
+                    f"Tool {tool['assetId']} is not available. Make sure it exists or you have access to it. "
+                    "If you think this is an error, please contact the administrators."
+                )
+                continue
 
     agent = Agent(
         id=payload["id"] if "id" in payload else "",
         name=payload.get("name", ""),
-        tools=tools,
+        tools=payload_tools,
         description=payload.get("description", ""),
         instructions=payload.get("role", ""),
         supplier=payload.get("teamId", None),

aixplain/factories/team_agent_factory/__init__.py

@@ -45,27 +45,38 @@ def create(
         api_key: Text = config.TEAM_API_KEY,
         supplier: Union[Dict, Text, Supplier, int] = "aiXplain",
         version: Optional[Text] = None,
+        use_mentalist: bool = True,
         use_inspector: bool = True,
-        use_mentalist_and_inspector: bool = True,
+        use_mentalist_and_inspector: bool = False,  # TODO: remove this
     ) -> TeamAgent:
         """Create a new team agent in the platform."""
         assert len(agents) > 0, "TeamAgent Onboarding Error: At least one agent must be provided."
+        agent_list = []
         for agent in agents:
             if isinstance(agent, Text) is True:
                 try:
                     from aixplain.factories.agent_factory import AgentFactory
 
-                    agent = AgentFactory.get(agent)
+                    agent_obj = AgentFactory.get(agent)
                 except Exception:
                     raise Exception(f"TeamAgent Onboarding Error: Agent {agent} does not exist.")
             else:
                 from aixplain.modules.agent import Agent
 
+                agent_obj = agent
+
                 assert isinstance(agent, Agent), "TeamAgent Onboarding Error: Agents must be instances of Agent class"
-
-        mentalist_and_inspector_llm_id = None
-        if use_inspector or use_mentalist_and_inspector:
-            mentalist_and_inspector_llm_id = llm_id
+            agent_list.append(agent_obj)
+
+        if use_inspector and not use_mentalist:
+            raise Exception("TeamAgent Onboarding Error: To use the Inspector agent, you must enable Mentalist.")
+
+        if use_mentalist_and_inspector:
+            mentalist_llm_id = llm_id
+            inspector_llm_id = llm_id
+        else:
+            mentalist_llm_id = llm_id if use_mentalist else None
+            inspector_llm_id = llm_id if use_inspector else None
 
         team_agent = None
         url = urljoin(config.BACKEND_URL, "sdk/agent-communities")
@@ -76,26 +87,26 @@ def create(
         elif isinstance(supplier, Supplier):
             supplier = supplier.value["code"]
 
-        agent_list = []
+        agent_payload_list = []
         for idx, agent in enumerate(agents):
-            agent_list.append({"assetId": agent.id, "number": idx, "type": "AGENT", "label": "AGENT"})
+            agent_payload_list.append({"assetId": agent.id, "number": idx, "type": "AGENT", "label": "AGENT"})
 
         payload = {
             "name": name,
-            "agents": agent_list,
+            "agents": agent_payload_list,
             "links": [],
             "description": description,
             "llmId": llm_id,
             "supervisorId": llm_id,
-            "plannerId": mentalist_and_inspector_llm_id,
-            "inspectorId": mentalist_and_inspector_llm_id,
+            "plannerId": mentalist_llm_id,
+            "inspectorId": inspector_llm_id,
             "supplier": supplier,
             "version": version,
             "status": "draft",
         }
 
-        team_agent = build_team_agent(payload=payload, api_key=api_key)
-        team_agent.validate()
+        team_agent = build_team_agent(payload=payload, agents=agent_list, api_key=api_key)
+        team_agent.validate(raise_exception=True)
         response = "Unspecified error"
         try:
             logging.debug(f"Start service for POST Create TeamAgent  - {url} - {headers} - {json.dumps(payload)}")
@@ -105,7 +116,7 @@ def create(
             raise Exception(e)
 
         if 200 <= r.status_code < 300:
-            team_agent = build_team_agent(payload=response, api_key=api_key)
+            team_agent = build_team_agent(payload=response, agents=agent_list, api_key=api_key)
         else:
             error_msg = f"{response}"
             if "message" in response:

aixplain/factories/team_agent_factory/utils.py

@@ -1,34 +1,45 @@
 __author__ = "lucaspavanelli"
 
+import logging
 import aixplain.utils.config as config
 from aixplain.enums.asset_status import AssetStatus
+from aixplain.modules.agent import Agent
 from aixplain.modules.team_agent import TeamAgent
-from typing import Dict, Text
+from typing import Dict, Text, List
 from urllib.parse import urljoin
 
 GPT_4o_ID = "6646261c6eb563165658bbb1"
 
 
-def build_team_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> TeamAgent:
+def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text = config.TEAM_API_KEY) -> TeamAgent:
     """Instantiate a new team agent in the platform."""
     from aixplain.factories.agent_factory import AgentFactory
 
     agents_dict = payload["agents"]
-    agents = []
-    for i, agent in enumerate(agents_dict):
-        agent = AgentFactory.get(agent["assetId"])
-        agents.append(agent)
+    payload_agents = agents
+    if payload_agents is None:
+        payload_agents = []
+        for i, agent in enumerate(agents_dict):
+            try:
+                payload_agents.append(AgentFactory.get(agent["assetId"]))
+            except Exception:
+                logging.warning(
+                    f"Agent {agent['assetId']} not found. Make sure it exists or you have access to it. "
+                    "If you think this is an error, please contact the administrators."
+                )
+                continue
 
     team_agent = TeamAgent(
         id=payload.get("id", ""),
         name=payload.get("name", ""),
-        agents=agents,
+        agents=payload_agents,
         description=payload.get("description", ""),
         supplier=payload.get("teamId", None),
         version=payload.get("version", None),
         cost=payload.get("cost", None),
         llm_id=payload.get("llmId", GPT_4o_ID),
-        use_mentalist_and_inspector=True if payload["plannerId"] is not None else False,
+        use_mentalist=True if payload.get("plannerId", None) is not None else False,
+        use_inspector=True if payload.get("inspectorId", None) is not None else False,
         api_key=api_key,
         status=AssetStatus(payload["status"]),
     )