From a1a2502728aebb7e7f785f3479feef8f2255ba97 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Mon, 28 Oct 2024 12:28:49 -0300 Subject: [PATCH 1/7] Update method for agents and teams --- aixplain/enums/asset_status.py | 34 +++--- aixplain/factories/agent_factory/__init__.py | 110 +++++++----------- aixplain/factories/agent_factory/utils.py | 41 ++----- .../factories/team_agent_factory/__init__.py | 17 +-- .../factories/team_agent_factory/utils.py | 19 +-- aixplain/modules/agent/__init__.py | 65 ++++++++++- aixplain/modules/agent/tool.py | 59 ---------- aixplain/modules/agent/tool/__init__.py | 7 ++ aixplain/modules/agent/tool/model_tool.py | 11 ++ aixplain/modules/agent/tool/pipeline_tool.py | 7 ++ aixplain/modules/team_agent/__init__.py | 66 +++++++++++ 11 files changed, 247 insertions(+), 189 deletions(-) delete mode 100644 aixplain/modules/agent/tool.py diff --git a/aixplain/enums/asset_status.py b/aixplain/enums/asset_status.py index 134af26e..994212fb 100644 --- a/aixplain/enums/asset_status.py +++ b/aixplain/enums/asset_status.py @@ -24,20 +24,22 @@ from enum import Enum from typing import Text + class AssetStatus(Text, Enum): - HIDDEN = 'hidden' - SCHEDULED = 'scheduled' - ONBOARDING = 'onboarding' - ONBOARDED = 'onboarded' - PENDING = 'pending' - FAILED = 'failed' - TRAINING = 'training' - REJECTED = 'rejected' - ENABLING = 'enabling' - DELETING = 'deleting' - DISABLED = 'disabled' - DELETED = 'deleted' - IN_PROGRESS = 'in_progress' - COMPLETED = 'completed' - CANCELING = 'canceling' - CANCELED = 'canceled' \ No newline at end of file + DRAFT = "draft" + HIDDEN = "hidden" + SCHEDULED = "scheduled" + ONBOARDING = "onboarding" + ONBOARDED = "onboarded" + PENDING = "pending" + FAILED = "failed" + TRAINING = "training" + REJECTED = "rejected" + ENABLING = "enabling" + DELETING = "deleting" + DISABLED = "disabled" + DELETED = "deleted" + IN_PROGRESS = "in_progress" + COMPLETED = "completed" + CANCELING = "canceling" + CANCELED = "canceled" diff --git a/aixplain/factories/agent_factory/__init__.py b/aixplain/factories/agent_factory/__init__.py index 2a16e191..2c9cf629 100644 --- a/aixplain/factories/agent_factory/__init__.py +++ b/aixplain/factories/agent_factory/__init__.py @@ -34,7 +34,6 @@ from aixplain.utils import config from typing import Dict, List, Optional, Text, Union -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 @@ -65,74 +64,49 @@ def create( Returns: Agent: created Agent """ - validate_name(name) - # validate LLM ID - validate_llm(llm_id) + from aixplain.factories.agent_factory.utils import build_agent + agent = None + url = urljoin(config.BACKEND_URL, "sdk/agents") + headers = {"x-api-key": api_key} + + if isinstance(supplier, dict): + supplier = supplier["code"] + elif isinstance(supplier, Supplier): + supplier = supplier.value["code"] + + payload = { + "name": name, + "assets": [tool.to_dict() for tool in tools], + "description": description, + "supplier": supplier, + "version": version, + "llmId": llm_id, + "status": "draft", + } + agent = build_agent(payload=payload, api_key=api_key) + agent.validate() + response = "Unspecified error" try: - agent = None - url = urljoin(config.BACKEND_URL, "sdk/agents") - headers = {"x-api-key": api_key} - - if isinstance(supplier, dict): - supplier = supplier["code"] - elif isinstance(supplier, Supplier): - supplier = supplier.value["code"] - - tool_payload = [] - for tool in tools: - if isinstance(tool, ModelTool): - tool.validate() - tool_payload.append( - { - "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): - tool.validate() - tool_payload.append( - { - "assetId": tool.pipeline, - "description": tool.description, - "type": "pipeline", - } - ) - else: - raise Exception("Agent Creation Error: Tool type not supported.") - - payload = { - "name": name, - "assets": tool_payload, - "description": description, - "supplier": supplier, - "version": version, - "llmId": llm_id, - } - - logging.info(f"Start service for POST Create Agent - {url} - {headers} - {json.dumps(payload)}") + logging.debug(f"Start service for POST Create Agent - {url} - {headers} - {json.dumps(payload)}") r = _request_with_retry("post", url, headers=headers, json=payload) - if 200 <= r.status_code < 300: - response = r.json() - agent = build_agent(payload=response, api_key=api_key) - else: - error = r.json() - 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 tools are not available." - error_msg = f"Agent Onboarding Error (HTTP {r.status_code}): {msg}" - logging.exception(error_msg) - raise Exception(error_msg) - except Exception as e: - raise Exception(e) + response = r.json() + except Exception: + raise Exception("Agent Onboarding Error: Please contact the administrators.") + + if 200 <= r.status_code < 300: + agent = build_agent(payload=response, api_key=api_key) + else: + error_msg = f"Agent Onboarding Error: {response}" + if "message" in response: + msg = response["message"] + if response["message"] == "err.name_already_exists": + msg = "Agent name already exists." + elif response["message"] == "err.asset_is_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) return agent @classmethod @@ -164,6 +138,8 @@ def create_pipeline_tool(cls, description: Text, pipeline: Union[Pipeline, Text] @classmethod def list(cls) -> Dict: """List all agents available in the platform.""" + from aixplain.factories.agent_factory.utils import build_agent + url = urljoin(config.BACKEND_URL, "sdk/agents") headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"} @@ -195,6 +171,8 @@ def list(cls) -> Dict: @classmethod def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent: """Get agent by id.""" + from aixplain.factories.agent_factory.utils import build_agent + url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent_id}") if config.AIXPLAIN_API_KEY != "": headers = {"x-aixplain-key": f"{config.AIXPLAIN_API_KEY}", "Content-Type": "application/json"} diff --git a/aixplain/factories/agent_factory/utils.py b/aixplain/factories/agent_factory/utils.py index d86982ef..99975b0d 100644 --- a/aixplain/factories/agent_factory/utils.py +++ b/aixplain/factories/agent_factory/utils.py @@ -12,8 +12,9 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent: """Instantiate a new agent in the platform.""" - tools = payload["assets"] - for i, tool in enumerate(tools): + tools_dict = payload["assets"] + tools = [] + for i, tool in enumerate(tools_dict): if tool["type"] == "model": for supplier in Supplier: if tool["supplier"] is not None and tool["supplier"].lower() in [ @@ -24,7 +25,7 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent: break tool = ModelTool( - function=Function(tool["function"]) if tool["function"] is not None else None, + function=Function(tool.get("function", None)), supplier=tool["supplier"], version=tool["version"], model=tool["assetId"], @@ -33,37 +34,19 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent: tool = PipelineTool(description=tool["description"], pipeline=tool["assetId"]) else: raise Exception("Agent Creation Error: Tool type not supported.") - tools[i] = tool + tools.append(tool) agent = Agent( - id=payload["id"], - name=payload["name"] if "name" in payload else "", + id=payload["id"] if "id" in payload else "", + name=payload.get("name", ""), tools=tools, - description=payload["description"] if "description" in payload else "", - 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 GPT_4o_ID, + 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), api_key=api_key, status=AssetStatus(payload["status"]), ) agent.url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent.id}/run") return agent - - -def validate_llm(model_id: Text) -> None: - from aixplain.factories.model_factory import ModelFactory - - try: - llm = ModelFactory.get(model_id) - 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." diff --git a/aixplain/factories/team_agent_factory/__init__.py b/aixplain/factories/team_agent_factory/__init__.py index 72d47c03..4923282c 100644 --- a/aixplain/factories/team_agent_factory/__init__.py +++ b/aixplain/factories/team_agent_factory/__init__.py @@ -25,8 +25,6 @@ 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 @@ -50,17 +48,18 @@ def create( 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: + from aixplain.factories.agent_factory import AgentFactory + agent = AgentFactory.get(agent) except Exception: raise Exception(f"TeamAgent Onboarding Error: Agent {agent} does not exist.") else: + from aixplain.modules.agent import Agent + assert isinstance(agent, Agent), "TeamAgent Onboarding Error: Agents must be instances of Agent class" mentalist_and_inspector_llm_id = None @@ -92,6 +91,8 @@ def create( "version": version, } + team_agent = build_team_agent(payload=payload, api_key=api_key) + team_agent.validate() 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: @@ -145,7 +146,7 @@ def list(cls) -> Dict: raise Exception(e) @classmethod - def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent: + def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> TeamAgent: """Get agent by id.""" url = urljoin(config.BACKEND_URL, f"sdk/agent-communities/{agent_id}") if config.AIXPLAIN_API_KEY != "": @@ -153,7 +154,7 @@ def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent: 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}") + logging.info(f"Start service for GET Team Agent - {url} - {headers}") r = _request_with_retry("get", url, headers=headers) resp = r.json() if 200 <= r.status_code < 300: @@ -162,5 +163,5 @@ def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent: msg = "Please contact the administrators." if "message" in resp: msg = resp["message"] - error_msg = f"Agent Get Error (HTTP {r.status_code}): {msg}" + error_msg = f"Team Agent Get Error (HTTP {r.status_code}): {msg}" raise Exception(error_msg) diff --git a/aixplain/factories/team_agent_factory/utils.py b/aixplain/factories/team_agent_factory/utils.py index 42fa5f6c..758f508e 100644 --- a/aixplain/factories/team_agent_factory/utils.py +++ b/aixplain/factories/team_agent_factory/utils.py @@ -3,7 +3,6 @@ import aixplain.utils.config as config from aixplain.enums.asset_status import AssetStatus from aixplain.modules.team_agent import TeamAgent -from aixplain.factories.agent_factory import AgentFactory from typing import Dict, Text from urllib.parse import urljoin @@ -12,21 +11,23 @@ def build_team_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> TeamAgent: """Instantiate a new team agent in the platform.""" + from aixplain.factories.agent_factory import AgentFactory + agents = payload["agents"] for i, agent in enumerate(agents): agent = AgentFactory.get(agent["assetId"]) agents[i] = agent team_agent = TeamAgent( - id=payload["id"], - name=payload["name"] if "name" in payload else "", + id=payload.get("id", ""), + name=payload.get("name", ""), agents=agents, - description=payload["description"] if "description" in payload else "", - 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 GPT_4o_ID, - use_mentalist_and_inspector=True if "plannerId" in payload and payload["plannerId"] is not None else False, + 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=payload.get("plannerId", False), api_key=api_key, status=AssetStatus(payload["status"]), ) diff --git a/aixplain/modules/agent/__init__.py b/aixplain/modules/agent/__init__.py index a7586c8b..d71dc264 100644 --- a/aixplain/modules/agent/__init__.py +++ b/aixplain/modules/agent/__init__.py @@ -22,10 +22,12 @@ """ import json import logging +import re import time import traceback from aixplain.utils.file_utils import _request_with_retry +from aixplain.enums.function import Function from aixplain.enums.supplier import Supplier from aixplain.enums.asset_status import AssetStatus from aixplain.enums.storage_type import StorageType @@ -66,7 +68,7 @@ def __init__( supplier: Union[Dict, Text, Supplier, int] = "aiXplain", version: Optional[Text] = None, cost: Optional[Dict] = None, - status: AssetStatus = AssetStatus.ONBOARDING, + status: AssetStatus = AssetStatus.DRAFT, **additional_info, ) -> None: """Create an Agent with the necessary information. @@ -91,9 +93,27 @@ def __init__( try: status = AssetStatus(status) except Exception: - status = AssetStatus.ONBOARDING + status = AssetStatus.DRAFT self.status = status + def validate(self) -> None: + """Validate the Agent.""" + from aixplain.factories.model_factory import ModelFactory + + # validate name + assert ( + re.match("^[a-zA-Z0-9 ]*$", self.name) is not None + ), "Agent Creation Error: Agent name must not contain special characters." + + try: + llm = ModelFactory.get(self.llm_id) + 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.") + + for tool in self.tools: + tool.validate() + def run( self, data: Optional[Union[Dict, Text]] = None, @@ -242,6 +262,17 @@ def run_async( response["error"] = msg return response + def to_dict(self) -> Dict: + return { + "name": self.name, + "assets": [tool.to_dict() for tool in self.tools], + "description": self.description, + "supplier": self.supplier, + "version": self.version, + "llmId": self.llm_id, + "status": self.status.value, + } + def delete(self) -> None: """Delete Agent service""" try: @@ -255,3 +286,33 @@ def delete(self) -> None: message = f"Agent Deletion Error (HTTP {r.status_code}): Make sure the agent exists and you are the owner." logging.error(message) raise Exception(f"{message}") + + def update(self) -> None: + """Update agent.""" + from aixplain.factories.agent_factory.utils import build_agent + + self.validate() + url = urljoin(config.BACKEND_URL, f"sdk/agents/{self.id}") + headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"} + + payload = self.to_dict() + + logging.debug(f"Start service for PUT Update Agent - {url} - {headers} - {json.dumps(payload)}") + resp = "No specified error." + try: + r = _request_with_retry("put", url, headers=headers, json=payload) + resp = r.json() + except Exception: + raise Exception("Agent Update Error: Please contact the administrators.") + + if 200 <= r.status_code < 300: + return build_agent(resp) + else: + error_msg = f"Agent Update Error (HTTP {r.status_code}): {resp}" + raise Exception(error_msg) + + def deploy(self) -> None: + assert self.status == AssetStatus.DRAFT, "Agent must be in draft status to be deployed." + assert self.status != AssetStatus.ONBOARDED, "Agent is already deployed." + self.status = AssetStatus.ONBOARDED + self.update() diff --git a/aixplain/modules/agent/tool.py b/aixplain/modules/agent/tool.py deleted file mode 100644 index 6651afe7..00000000 --- a/aixplain/modules/agent/tool.py +++ /dev/null @@ -1,59 +0,0 @@ -__author__ = "aiXplain" - -""" -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: Lucas Pavanelli and Thiago Castro Ferreira -Date: May 16th 2024 -Description: - Agentification Class -""" -from typing import Text, Optional - -from aixplain.enums.function import Function -from aixplain.enums.supplier import Supplier - - -class Tool: - """Specialized software or resource designed to assist the AI in executing specific tasks or functions based on user commands. - - Attributes: - name (Text): name of the tool - description (Text): descriptiion of the tool - function (Function): task that the tool performs - supplier (Optional[Union[Dict, Text, Supplier, int]], optional): Preferred supplier to perform the task. Defaults to None. - """ - - def __init__( - self, - name: Text, - description: Text, - function: Function, - supplier: Optional[Supplier] = None, - **additional_info, - ) -> None: - """Specialized software or resource designed to assist the AI in executing specific tasks or functions based on user commands. - - Args: - name (Text): name of the tool - description (Text): descriptiion of the tool - function (Function): task that the tool performs - supplier (Optional[Union[Dict, Text, Supplier, int]], optional): Preferred supplier to perform the task. Defaults to None. - """ - self.name = name - self.description = description - self.function = function - self.supplier = supplier - self.additional_info = additional_info diff --git a/aixplain/modules/agent/tool/__init__.py b/aixplain/modules/agent/tool/__init__.py index 9c7a7a09..01b44dfa 100644 --- a/aixplain/modules/agent/tool/__init__.py +++ b/aixplain/modules/agent/tool/__init__.py @@ -51,3 +51,10 @@ def __init__( self.description = description self.version = version self.additional_info = additional_info + + def to_dict(self): + """Converts the tool to a dictionary.""" + raise NotImplementedError + + def validate(self): + raise NotImplementedError diff --git a/aixplain/modules/agent/tool/model_tool.py b/aixplain/modules/agent/tool/model_tool.py index 3a84c45b..d651761b 100644 --- a/aixplain/modules/agent/tool/model_tool.py +++ b/aixplain/modules/agent/tool/model_tool.py @@ -77,6 +77,17 @@ def __init__( self.model = model self.function = function + def to_dict(self) -> Dict: + """Converts the tool to a dictionary.""" + return { + "function": self.function.value if self.function is not None else None, + "type": "model", + "description": self.description, + "supplier": self.supplier.value["code"] if self.supplier else None, + "version": self.version if self.version else None, + "assetId": self.model, + } + def validate(self) -> Model: from aixplain.factories.model_factory import ModelFactory diff --git a/aixplain/modules/agent/tool/pipeline_tool.py b/aixplain/modules/agent/tool/pipeline_tool.py index fa8394ea..9ea7a5fb 100644 --- a/aixplain/modules/agent/tool/pipeline_tool.py +++ b/aixplain/modules/agent/tool/pipeline_tool.py @@ -51,6 +51,13 @@ def __init__( pipeline = pipeline.id self.pipeline = pipeline + def to_dict(self): + return { + "assetId": self.pipeline, + "description": self.description, + "type": "pipeline", + } + def validate(self): from aixplain.factories.pipeline_factory import PipelineFactory diff --git a/aixplain/modules/team_agent/__init__.py b/aixplain/modules/team_agent/__init__.py index 86321489..44ee26ba 100644 --- a/aixplain/modules/team_agent/__init__.py +++ b/aixplain/modules/team_agent/__init__.py @@ -25,8 +25,10 @@ import logging import time import traceback +import re from aixplain.utils.file_utils import _request_with_retry +from aixplain.enums.function import Function from aixplain.enums.supplier import Supplier from aixplain.enums.asset_status import AssetStatus from aixplain.enums.storage_type import StorageType @@ -259,3 +261,67 @@ def delete(self) -> None: ) logging.error(message) raise Exception(f"{message}") + + def to_dict(self) -> Dict: + return { + "name": self.name, + "agents": [ + {"assetId": agent.id, "number": idx, "type": "AGENT", "label": "AGENT"} for idx, agent in enumerate(self.agents) + ], + "links": [], + "description": self.description, + "llmId": self.llm_id, + "supervisorId": self.llm_id, + "plannerId": self.mentalist_and_inspector_llm_id, + "supplier": self.supplier, + "version": self.version, + } + + def validate(self) -> None: + """Validate the Team.""" + from aixplain.factories.model_factory import ModelFactory + + # validate name + assert ( + re.match("^[a-zA-Z0-9 ]*$", self.name) is not None + ), "Team Agent Creation Error: Team name must not contain special characters." + + try: + llm = ModelFactory.get(self.llm_id) + 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.") + + for agent in self.agents: + agent.validate() + + def update(self) -> None: + """Update the Team Agent.""" + from aixplain.factories.team_agent_factory.utils import build_team_agent + + self.validate() + url = urljoin(config.BACKEND_URL, f"sdk/agent-communities/{self.id}") + headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"} + + payload = self.to_dict() + + logging.debug(f"Start service for PUT Update Team Agent - {url} - {headers} - {json.dumps(payload)}") + resp = "No specified error." + try: + r = _request_with_retry("put", url, headers=headers, json=payload) + resp = r.json() + except Exception: + raise Exception("Team Agent Update Error: Please contact the administrators.") + + if 200 <= r.status_code < 300: + return build_team_agent(resp) + else: + error_msg = f"Team Agent Update Error (HTTP {r.status_code}): {resp}" + raise Exception(error_msg) + + def deploy(self) -> None: + """Deploy the Team Agent.""" + assert self.status == AssetStatus.DRAFT, "Team Agent Deployment Error: Team Agent must be in draft status." + assert self.status != AssetStatus.ONBOARDED, "Team Agent Deployment Error: Team Agent must be onboarded." + self.status = AssetStatus.ONBOARDED + self.update() From 97438cbca4745c2c3eefd2f305b24e6dfee4cb3a Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Mon, 28 Oct 2024 16:01:52 -0300 Subject: [PATCH 2/7] Agent Functional Test --- aixplain/factories/agent_factory/utils.py | 13 +-- aixplain/modules/agent/__init__.py | 3 +- aixplain/modules/agent/tool/model_tool.py | 11 ++- aixplain/modules/team_agent/__init__.py | 3 +- .../functional/agent/agent_functional_test.py | 47 +++++++++- tests/unit/agent_test.py | 87 ++++++++++++++++++- tests/unit/team_agent_test.py | 35 +++++++- 7 files changed, 184 insertions(+), 15 deletions(-) diff --git a/aixplain/factories/agent_factory/utils.py b/aixplain/factories/agent_factory/utils.py index 99975b0d..19318fb3 100644 --- a/aixplain/factories/agent_factory/utils.py +++ b/aixplain/factories/agent_factory/utils.py @@ -14,19 +14,20 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent: """Instantiate a new agent in the platform.""" tools_dict = payload["assets"] tools = [] - for i, tool in enumerate(tools_dict): + for tool in tools_dict: if tool["type"] == "model": - for supplier in Supplier: + 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_.value["code"].lower(), + supplier_.value["name"].lower(), ]: - tool["supplier"] = supplier + supplier = supplier_ break tool = ModelTool( function=Function(tool.get("function", None)), - supplier=tool["supplier"], + supplier=supplier, version=tool["version"], model=tool["assetId"], ) diff --git a/aixplain/modules/agent/__init__.py b/aixplain/modules/agent/__init__.py index d71dc264..78da1269 100644 --- a/aixplain/modules/agent/__init__.py +++ b/aixplain/modules/agent/__init__.py @@ -264,10 +264,11 @@ def run_async( def to_dict(self) -> Dict: return { + "id": self.id, "name": self.name, "assets": [tool.to_dict() for tool in self.tools], "description": self.description, - "supplier": self.supplier, + "supplier": self.supplier.value["code"] if isinstance(self.supplier, Supplier) else self.supplier, "version": self.version, "llmId": self.llm_id, "status": self.status.value, diff --git a/aixplain/modules/agent/tool/model_tool.py b/aixplain/modules/agent/tool/model_tool.py index d651761b..bb980b94 100644 --- a/aixplain/modules/agent/tool/model_tool.py +++ b/aixplain/modules/agent/tool/model_tool.py @@ -79,11 +79,20 @@ def __init__( def to_dict(self) -> Dict: """Converts the tool to a dictionary.""" + supplier = self.supplier + if supplier is not None: + if isinstance(supplier, dict): + supplier = supplier["code"] + elif isinstance(supplier, Supplier): + supplier = supplier.value["code"] + else: + supplier = str(supplier) + return { "function": self.function.value if self.function is not None else None, "type": "model", "description": self.description, - "supplier": self.supplier.value["code"] if self.supplier else None, + "supplier": supplier, "version": self.version if self.version else None, "assetId": self.model, } diff --git a/aixplain/modules/team_agent/__init__.py b/aixplain/modules/team_agent/__init__.py index 44ee26ba..2f8b5c3b 100644 --- a/aixplain/modules/team_agent/__init__.py +++ b/aixplain/modules/team_agent/__init__.py @@ -264,6 +264,7 @@ def delete(self) -> None: def to_dict(self) -> Dict: return { + "id": self.id, "name": self.name, "agents": [ {"assetId": agent.id, "number": idx, "type": "AGENT", "label": "AGENT"} for idx, agent in enumerate(self.agents) @@ -272,7 +273,7 @@ def to_dict(self) -> Dict: "description": self.description, "llmId": self.llm_id, "supervisorId": self.llm_id, - "plannerId": self.mentalist_and_inspector_llm_id, + "plannerId": self.llm_id if self.use_mentalist_and_inspector else None, "supplier": self.supplier, "version": self.version, } diff --git a/tests/functional/agent/agent_functional_test.py b/tests/functional/agent/agent_functional_test.py index 648f4f28..3fd6f032 100644 --- a/tests/functional/agent/agent_functional_test.py +++ b/tests/functional/agent/agent_functional_test.py @@ -15,13 +15,16 @@ See the License for the specific language governing permissions and limitations under the License. """ +import copy import json from dotenv import load_dotenv load_dotenv() from aixplain.factories import AgentFactory +from aixplain.enums.asset_status import AssetStatus from aixplain.enums.function import Function from aixplain.enums.supplier import Supplier +from uuid import uuid4 import pytest @@ -44,14 +47,15 @@ def test_end2end(run_input_map): tools = [] if "model_tools" in run_input_map: for tool in run_input_map["model_tools"]: + tool_ = copy.copy(tool) for supplier in Supplier: if tool["supplier"] is not None and tool["supplier"].lower() in [ supplier.value["code"].lower(), supplier.value["name"].lower(), ]: - tool["supplier"] = supplier + tool_["supplier"] = supplier break - tools.append(AgentFactory.create_model_tool(**tool)) + tools.append(AgentFactory.create_model_tool(**tool_)) if "pipeline_tools" in run_input_map: for tool in run_input_map["pipeline_tools"]: tools.append(AgentFactory.create_pipeline_tool(pipeline=tool["pipeline_id"], description=tool["description"])) @@ -60,6 +64,11 @@ def test_end2end(run_input_map): name=run_input_map["agent_name"], description=run_input_map["agent_name"], llm_id=run_input_map["llm_id"], tools=tools ) assert agent is not None + assert agent.status == AssetStatus.DRAFT + # deploy agent + agent.deploy() + assert agent.status == AssetStatus.ONBOARDED + agent = AgentFactory.get(agent.id) assert agent is not None response = agent.run(data=run_input_map["query"]) @@ -79,6 +88,40 @@ def test_list_agents(): assert type(agents_result) is list +def test_update_draft_agent(run_input_map): + for agent in AgentFactory.list()["results"]: + agent.delete() + + tools = [] + if "model_tools" in run_input_map: + for tool in run_input_map["model_tools"]: + tool_ = copy.copy(tool) + for supplier in Supplier: + if tool["supplier"] is not None and tool["supplier"].lower() in [ + supplier.value["code"].lower(), + supplier.value["name"].lower(), + ]: + tool_["supplier"] = supplier + break + tools.append(AgentFactory.create_model_tool(**tool_)) + if "pipeline_tools" in run_input_map: + for tool in run_input_map["pipeline_tools"]: + tools.append(AgentFactory.create_pipeline_tool(pipeline=tool["pipeline_id"], description=tool["description"])) + + agent = AgentFactory.create( + name=run_input_map["agent_name"], description=run_input_map["agent_name"], llm_id=run_input_map["llm_id"], tools=tools + ) + + agent_name = str(uuid4()).replace("-", "") + agent.name = agent_name + agent.update() + + agent = AgentFactory.get(agent.id) + assert agent.name == agent_name + assert agent.status == AssetStatus.DRAFT + agent.delete() + + def test_fail_non_existent_llm(): with pytest.raises(Exception) as exc_info: AgentFactory.create( diff --git a/tests/unit/agent_test.py b/tests/unit/agent_test.py index 1be0682e..e4d35466 100644 --- a/tests/unit/agent_test.py +++ b/tests/unit/agent_test.py @@ -1,5 +1,6 @@ import pytest import requests_mock +from aixplain.enums.asset_status import AssetStatus from aixplain.modules import Agent from aixplain.utils import config from aixplain.factories import AgentFactory @@ -53,7 +54,7 @@ def test_fail_key_not_found(): assert str(exc_info.value) == "Key 'input2' not found in query." -def test_sucess_query_content(): +def test_success_query_content(): agent = Agent("123", "Test Agent", "Sample Description") with requests_mock.Mocker() as mock: url = agent.url @@ -83,6 +84,12 @@ def test_invalid_modeltool(): assert str(exc_info.value) == "Model Tool Unavailable. Make sure Model '309851793' exists or you have access to it." +def test_invalid_llm_id(): + with pytest.raises(Exception) as exc_info: + AgentFactory.create(name="Test", description="", tools=[], llm_id="123") + assert str(exc_info.value) == "Large Language Model with ID '123' not found." + + def test_invalid_agent_name(): with pytest.raises(Exception) as exc_info: AgentFactory.create(name="[Test]", description="", tools=[], llm_id="6646261c6eb563165658bbb1") @@ -102,7 +109,7 @@ def test_create_agent(): "description": "Test Agent Description", "teamId": "123", "version": "1.0", - "status": "onboarded", + "status": "draft", "llmId": "6646261c6eb563165658bbb1", "pricing": {"currency": "USD", "value": 0.0}, "assets": [ @@ -134,10 +141,84 @@ def test_create_agent(): name="Test Agent", description="Test Agent Description", llm_id="6646261c6eb563165658bbb1", - tools=[AgentFactory.create_model_tool(supplier=Supplier.OPENAI, function="text-generation")], + tools=[ModelTool(function="text-generation", supplier=Supplier.OPENAI)], ) assert agent.name == ref_response["name"] assert agent.description == ref_response["description"] assert agent.llm_id == ref_response["llmId"] assert agent.tools[0].function.value == ref_response["assets"][0]["function"] + assert agent.status == AssetStatus.DRAFT + + +def test_to_dict(): + agent = Agent( + id="", + name="Test Agent", + description="Test Agent Description", + llm_id="6646261c6eb563165658bbb1", + tools=[AgentFactory.create_model_tool(function="text-generation")], + ) + + agent_json = agent.to_dict() + assert agent_json["id"] == "" + assert agent_json["name"] == "Test Agent" + assert agent_json["description"] == "Test Agent Description" + assert agent_json["llmId"] == "6646261c6eb563165658bbb1" + assert agent_json["assets"][0]["function"] == "text-generation" + assert agent_json["assets"][0]["type"] == "model" + + +def test_update_success(): + agent = Agent( + id="123", + name="Test Agent", + description="Test Agent Description", + llm_id="6646261c6eb563165658bbb1", + tools=[AgentFactory.create_model_tool(function="text-generation")], + ) + + with requests_mock.Mocker() as mock: + url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent.id}") + headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"} + ref_response = { + "id": "123", + "name": "Test Agent", + "description": "Test Agent Description", + "teamId": "123", + "version": "1.0", + "status": "onboarded", + "llmId": "6646261c6eb563165658bbb1", + "pricing": {"currency": "USD", "value": 0.0}, + "assets": [ + { + "type": "model", + "supplier": "openai", + "version": "1.0", + "assetId": "6646261c6eb563165658bbb1", + "function": "text-generation", + } + ], + } + mock.put(url, headers=headers, json=ref_response) + + url = urljoin(config.BACKEND_URL, "sdk/models/6646261c6eb563165658bbb1") + model_ref_response = { + "id": "6646261c6eb563165658bbb1", + "name": "Test LLM", + "description": "Test LLM Description", + "function": {"id": "text-generation"}, + "supplier": "openai", + "version": {"id": "1.0"}, + "status": "onboarded", + "pricing": {"currency": "USD", "value": 0.0}, + } + mock.get(url, headers=headers, json=model_ref_response) + + agent.update() + + assert agent.id == ref_response["id"] + assert agent.name == ref_response["name"] + assert agent.description == ref_response["description"] + assert agent.llm_id == ref_response["llmId"] + assert agent.tools[0].function.value == ref_response["assets"][0]["function"] diff --git a/tests/unit/team_agent_test.py b/tests/unit/team_agent_test.py index fd738c04..731c2c15 100644 --- a/tests/unit/team_agent_test.py +++ b/tests/unit/team_agent_test.py @@ -1,6 +1,7 @@ import pytest import requests_mock -from aixplain.modules import TeamAgent +from aixplain.modules import Agent, TeamAgent +from aixplain.modules.agent import ModelTool from aixplain.factories import TeamAgentFactory from aixplain.utils import config @@ -71,3 +72,35 @@ def test_fail_number_agents(): TeamAgentFactory.create(name="Test Team Agent", agents=[]) assert str(exc_info.value) == "TeamAgent Onboarding Error: At least one agent must be provided." + + +def test_to_dict(): + team_agent = TeamAgent( + id="123", + name="Test Team Agent", + agents=[ + Agent( + id="", + name="Test Agent", + description="Test Agent Description", + llm_id="6646261c6eb563165658bbb1", + tools=[ModelTool(function="text-generation")], + ) + ], + description="Test Team Agent Description", + llm_id="6646261c6eb563165658bbb1", + use_mentalist_and_inspector=False, + ) + + team_agent_dict = team_agent.to_dict() + assert team_agent_dict["id"] == "123" + assert team_agent_dict["name"] == "Test Team Agent" + assert team_agent_dict["description"] == "Test Team Agent Description" + assert team_agent_dict["llmId"] == "6646261c6eb563165658bbb1" + assert team_agent_dict["supervisorId"] == "6646261c6eb563165658bbb1" + assert team_agent_dict["plannerId"] is None + assert len(team_agent_dict["agents"]) == 1 + assert team_agent_dict["agents"][0]["assetId"] == "" + assert team_agent_dict["agents"][0]["number"] == 0 + assert team_agent_dict["agents"][0]["type"] == "AGENT" + assert team_agent_dict["agents"][0]["label"] == "AGENT" From 42de8654412dea3e5b2702f07fd284a3bef33e2c Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Mon, 28 Oct 2024 16:23:25 -0300 Subject: [PATCH 3/7] Team Agent Function tests --- .../factories/team_agent_factory/__init__.py | 89 ++++++++++--------- .../factories/team_agent_factory/utils.py | 7 +- .../team_agent/team_agent_functional_test.py | 60 ++++++++++++- 3 files changed, 107 insertions(+), 49 deletions(-) diff --git a/aixplain/factories/team_agent_factory/__init__.py b/aixplain/factories/team_agent_factory/__init__.py index 4923282c..be9721db 100644 --- a/aixplain/factories/team_agent_factory/__init__.py +++ b/aixplain/factories/team_agent_factory/__init__.py @@ -65,53 +65,56 @@ def create( mentalist_and_inspector_llm_id = None if use_mentalist_and_inspector is True: mentalist_and_inspector_llm_id = llm_id + + 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, + "status": "draft", + } + + team_agent = build_team_agent(payload=payload, api_key=api_key) + team_agent.validate() + response = "Unspecified error" 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, - } - - team_agent = build_team_agent(payload=payload, api_key=api_key) - team_agent.validate() - logging.info(f"Start service for POST Create TeamAgent - {url} - {headers} - {json.dumps(payload)}") + logging.debug(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) + response = r.json() except Exception as e: raise Exception(e) + + if 200 <= r.status_code < 300: + team_agent = build_team_agent(payload=response, api_key=api_key) + else: + error_msg = f"{response}" + if "message" in response: + msg = response["message"] + if response["message"] == "err.name_already_exists": + msg = "TeamAgent name already exists." + elif response["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) return team_agent @classmethod diff --git a/aixplain/factories/team_agent_factory/utils.py b/aixplain/factories/team_agent_factory/utils.py index 758f508e..d31cf5c1 100644 --- a/aixplain/factories/team_agent_factory/utils.py +++ b/aixplain/factories/team_agent_factory/utils.py @@ -13,10 +13,11 @@ def build_team_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Team """Instantiate a new team agent in the platform.""" from aixplain.factories.agent_factory import AgentFactory - agents = payload["agents"] - for i, agent in enumerate(agents): + agents_dict = payload["agents"] + agents = [] + for i, agent in enumerate(agents_dict): agent = AgentFactory.get(agent["assetId"]) - agents[i] = agent + agents.append(agent) team_agent = TeamAgent( id=payload.get("id", ""), diff --git a/tests/functional/team_agent/team_agent_functional_test.py b/tests/functional/team_agent/team_agent_functional_test.py index c7e50b68..66924fc5 100644 --- a/tests/functional/team_agent/team_agent_functional_test.py +++ b/tests/functional/team_agent/team_agent_functional_test.py @@ -20,9 +20,11 @@ load_dotenv() from aixplain.factories import AgentFactory, TeamAgentFactory +from aixplain.enums.asset_status import AssetStatus from aixplain.enums.function import Function from aixplain.enums.supplier import Supplier - +from copy import copy +from uuid import uuid4 import pytest RUN_FILE = "tests/functional/team_agent/data/team_agent_test_end2end.json" @@ -38,6 +40,8 @@ def run_input_map(request): def test_end2end(run_input_map): + for team in TeamAgentFactory.list()["results"]: + team.delete() for agent in AgentFactory.list()["results"]: agent.delete() @@ -46,14 +50,15 @@ def test_end2end(run_input_map): tools = [] if "model_tools" in agent: for tool in agent["model_tools"]: + tool_ = copy(tool) for supplier in Supplier: if tool["supplier"] is not None and tool["supplier"].lower() in [ supplier.value["code"].lower(), supplier.value["name"].lower(), ]: - tool["supplier"] = supplier + tool_["supplier"] = supplier break - tools.append(AgentFactory.create_model_tool(**tool)) + tools.append(AgentFactory.create_model_tool(**tool_)) if "pipeline_tools" in agent: for tool in agent["pipeline_tools"]: tools.append(AgentFactory.create_pipeline_tool(pipeline=tool["pipeline_id"], description=tool["description"])) @@ -61,6 +66,7 @@ def test_end2end(run_input_map): agent = AgentFactory.create( name=agent["agent_name"], description=agent["agent_name"], llm_id=agent["llm_id"], tools=tools ) + agent.deploy() agents.append(agent) team_agent = TeamAgentFactory.create( @@ -72,6 +78,9 @@ def test_end2end(run_input_map): ) assert team_agent is not None + assert team_agent.status == AssetStatus.DRAFT + # deploy team agent + team_agent.deploy() team_agent = TeamAgentFactory.get(team_agent.id) assert team_agent is not None response = team_agent.run(data=run_input_map["query"]) @@ -86,6 +95,51 @@ def test_end2end(run_input_map): team_agent.delete() +def test_draft_team_agent_update(run_input_map): + for team in TeamAgentFactory.list()["results"]: + team.delete() + for agent in AgentFactory.list()["results"]: + agent.delete() + + agents = [] + for agent in run_input_map["agents"]: + tools = [] + if "model_tools" in agent: + for tool in agent["model_tools"]: + tool_ = copy(tool) + for supplier in Supplier: + if tool["supplier"] is not None and tool["supplier"].lower() in [ + supplier.value["code"].lower(), + supplier.value["name"].lower(), + ]: + tool_["supplier"] = supplier + break + tools.append(AgentFactory.create_model_tool(**tool_)) + if "pipeline_tools" in agent: + for tool in agent["pipeline_tools"]: + tools.append(AgentFactory.create_pipeline_tool(pipeline=tool["pipeline_id"], description=tool["description"])) + + agent = AgentFactory.create( + name=agent["agent_name"], description=agent["agent_name"], llm_id=agent["llm_id"], tools=tools + ) + agents.append(agent) + + team_agent = TeamAgentFactory.create( + name=run_input_map["team_agent_name"], + agents=agents, + description=run_input_map["team_agent_name"], + llm_id=run_input_map["llm_id"], + use_mentalist_and_inspector=True, + ) + + team_agent_name = str(uuid4()).replace("-", "") + team_agent.name = team_agent_name + team_agent.update() + team_agent = TeamAgentFactory.get(team_agent.id) + assert team_agent.name == team_agent_name + assert team_agent.status == AssetStatus.DRAFT + + def test_fail_non_existent_llm(): with pytest.raises(Exception) as exc_info: AgentFactory.create( From dc8887938ab6f02d70f3bdc48ef05b23e9c931c8 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Mon, 28 Oct 2024 16:45:05 -0300 Subject: [PATCH 4/7] Team Agent Unit Tests --- .../factories/team_agent_factory/utils.py | 2 +- tests/unit/team_agent_test.py | 90 +++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/aixplain/factories/team_agent_factory/utils.py b/aixplain/factories/team_agent_factory/utils.py index d31cf5c1..da859a43 100644 --- a/aixplain/factories/team_agent_factory/utils.py +++ b/aixplain/factories/team_agent_factory/utils.py @@ -28,7 +28,7 @@ def build_team_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Team version=payload.get("version", None), cost=payload.get("cost", None), llm_id=payload.get("llmId", GPT_4o_ID), - use_mentalist_and_inspector=payload.get("plannerId", False), + use_mentalist_and_inspector=True if payload["plannerId"] is not None else False, api_key=api_key, status=AssetStatus(payload["status"]), ) diff --git a/tests/unit/team_agent_test.py b/tests/unit/team_agent_test.py index 731c2c15..56564b73 100644 --- a/tests/unit/team_agent_test.py +++ b/tests/unit/team_agent_test.py @@ -1,9 +1,12 @@ import pytest import requests_mock +from aixplain.enums.asset_status import AssetStatus from aixplain.modules import Agent, TeamAgent from aixplain.modules.agent import ModelTool from aixplain.factories import TeamAgentFactory +from aixplain.factories import AgentFactory from aixplain.utils import config +from urllib.parse import urljoin def test_fail_no_data_query(): @@ -104,3 +107,90 @@ def test_to_dict(): assert team_agent_dict["agents"][0]["number"] == 0 assert team_agent_dict["agents"][0]["type"] == "AGENT" assert team_agent_dict["agents"][0]["label"] == "AGENT" + + +def test_create_team_agent(): + with requests_mock.Mocker() as mock: + headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"} + # MOCK GET LLM + url = urljoin(config.BACKEND_URL, "sdk/models/6646261c6eb563165658bbb1") + model_ref_response = { + "id": "6646261c6eb563165658bbb1", + "name": "Test LLM", + "description": "Test LLM Description", + "function": {"id": "text-generation"}, + "supplier": "openai", + "version": {"id": "1.0"}, + "status": "onboarded", + "pricing": {"currency": "USD", "value": 0.0}, + } + mock.get(url, headers=headers, json=model_ref_response) + + # AGENT MOCK CREATION + url = urljoin(config.BACKEND_URL, "sdk/agents") + ref_response = { + "id": "123", + "name": "Test Agent", + "description": "Test Agent Description", + "teamId": "123", + "version": "1.0", + "status": "draft", + "llmId": "6646261c6eb563165658bbb1", + "pricing": {"currency": "USD", "value": 0.0}, + "assets": [ + { + "type": "model", + "supplier": "openai", + "version": "1.0", + "assetId": "6646261c6eb563165658bbb1", + "function": "text-generation", + } + ], + } + mock.post(url, headers=headers, json=ref_response) + + agent = AgentFactory.create( + name="Test Agent", + description="Test Agent Description", + llm_id="6646261c6eb563165658bbb1", + tools=[ModelTool(model="6646261c6eb563165658bbb1")], + ) + + # AGENT MOCK GET + url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent.id}") + mock.get(url, headers=headers, json=ref_response) + + # TEAM MOCK CREATION + url = urljoin(config.BACKEND_URL, "sdk/agent-communities") + team_ref_response = { + "id": "team_agent_123", + "name": "TEST Multi agent", + "status": "draft", + "teamId": 645, + "description": "TEST Multi agent", + "llmId": "6646261c6eb563165658bbb1", + "assets": [], + "agents": [{"assetId": "123", "type": "AGENT", "number": 0, "label": "AGENT"}], + "links": [], + "plannerId": "6646261c6eb563165658bbb1", + "supervisorId": "6646261c6eb563165658bbb1", + "createdAt": "2024-10-28T19:30:25.344Z", + "updatedAt": "2024-10-28T19:30:25.344Z", + } + mock.post(url, headers=headers, json=team_ref_response) + + team_agent = TeamAgentFactory.create( + name="TEST Multi agent", + description="TEST Multi agent", + use_mentalist_and_inspector=True, + llm_id="6646261c6eb563165658bbb1", + agents=[agent], + ) + assert team_agent.id is not None + assert team_agent.name == team_ref_response["name"] + assert team_agent.description == team_ref_response["description"] + assert team_agent.llm_id == team_ref_response["llmId"] + assert team_agent.use_mentalist_and_inspector is True + assert team_agent.status == AssetStatus.DRAFT + assert len(team_agent.agents) == 1 + assert team_agent.agents[0].id == team_ref_response["agents"][0]["assetId"] From 8456f9076852a43a0c118f53f7eb5015f3ed7df4 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Fri, 1 Nov 2024 16:05:16 -0300 Subject: [PATCH 5/7] Improvements in tests and listing agents --- aixplain/factories/agent_factory/__init__.py | 37 ++++++++++--------- .../functional/agent/agent_functional_test.py | 8 +++- 2 files changed, 26 insertions(+), 19 deletions(-) diff --git a/aixplain/factories/agent_factory/__init__.py b/aixplain/factories/agent_factory/__init__.py index 2c9cf629..ece88c82 100644 --- a/aixplain/factories/agent_factory/__init__.py +++ b/aixplain/factories/agent_factory/__init__.py @@ -143,30 +143,31 @@ def list(cls) -> Dict: url = urljoin(config.BACKEND_URL, "sdk/agents") headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"} + resp = {} 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() + except Exception: + raise Exception("Agent Listing Error: Please contact the administrators.") - 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_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) + 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_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 isinstance(resp, dict) and "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) @classmethod def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent: diff --git a/tests/functional/agent/agent_functional_test.py b/tests/functional/agent/agent_functional_test.py index 3fd6f032..50befd03 100644 --- a/tests/functional/agent/agent_functional_test.py +++ b/tests/functional/agent/agent_functional_test.py @@ -20,7 +20,7 @@ from dotenv import load_dotenv load_dotenv() -from aixplain.factories import AgentFactory +from aixplain.factories import AgentFactory, TeamAgentFactory from aixplain.enums.asset_status import AssetStatus from aixplain.enums.function import Function from aixplain.enums.supplier import Supplier @@ -41,6 +41,9 @@ def run_input_map(request): def test_end2end(run_input_map): + for team in TeamAgentFactory.list()["results"]: + team.delete() + for agent in AgentFactory.list()["results"]: agent.delete() @@ -89,6 +92,9 @@ def test_list_agents(): def test_update_draft_agent(run_input_map): + for team in TeamAgentFactory.list()["results"]: + team.delete() + for agent in AgentFactory.list()["results"]: agent.delete() From 13f96fa80612fac4107f2f58eedf8b54d48a774f Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Fri, 1 Nov 2024 16:17:01 -0300 Subject: [PATCH 6/7] Refactoring listing teams method --- .../factories/team_agent_factory/__init__.py | 45 ++++++++++--------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/aixplain/factories/team_agent_factory/__init__.py b/aixplain/factories/team_agent_factory/__init__.py index be9721db..3f65b4b0 100644 --- a/aixplain/factories/team_agent_factory/__init__.py +++ b/aixplain/factories/team_agent_factory/__init__.py @@ -123,30 +123,31 @@ def list(cls) -> Dict: url = urljoin(config.BACKEND_URL, "sdk/agent-communities") headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"} + resp = {} 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() + except Exception: + raise Exception("Team Agent Listing Error: Please contact the administrators.") - 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) + 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 isinstance(resp, dict) and "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) @classmethod def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> TeamAgent: @@ -158,8 +159,12 @@ def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> TeamAgent: 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 Team Agent - {url} - {headers}") - r = _request_with_retry("get", url, headers=headers) - resp = r.json() + try: + r = _request_with_retry("get", url, headers=headers) + resp = r.json() + except Exception: + raise Exception("Team Agent Get Error: Please contact the administrators.") + if 200 <= r.status_code < 300: return build_team_agent(resp) else: From 5ca212a218311da0b84685a89a09932baf4c986d Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Mon, 4 Nov 2024 11:50:43 -0300 Subject: [PATCH 7/7] Add contain field in model response --- aixplain/modules/model/response.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/aixplain/modules/model/response.py b/aixplain/modules/model/response.py index 902b9987..42ed09a4 100644 --- a/aixplain/modules/model/response.py +++ b/aixplain/modules/model/response.py @@ -68,3 +68,10 @@ def __repr__(self) -> str: if self.additional_fields: fields.extend([f"{k}={repr(v)}" for k, v in self.additional_fields.items()]) return f"ModelResponse({', '.join(fields)})" + + def __contains__(self, key: Text) -> bool: + try: + self[key] + return True + except KeyError: + return False