diff --git a/aixplain/enums/asset_status.py b/aixplain/enums/asset_status.py index 9274001f..994212fb 100644 --- a/aixplain/enums/asset_status.py +++ b/aixplain/enums/asset_status.py @@ -26,6 +26,7 @@ class AssetStatus(Text, Enum): + DRAFT = "draft" HIDDEN = "hidden" SCHEDULED = "scheduled" ONBOARDING = "onboarding" diff --git a/aixplain/factories/agent_factory/__init__.py b/aixplain/factories/agent_factory/__init__.py index 3e66884e..c56d1fd8 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 @@ -165,37 +139,42 @@ 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"} + 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: """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 9192f1d4..e5e73dc4 100644 --- a/aixplain/factories/agent_factory/utils.py +++ b/aixplain/factories/agent_factory/utils.py @@ -12,20 +12,22 @@ 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 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["function"]) if tool["function"] is not None else None, - supplier=tool["supplier"], + function=Function(tool.get("function", None)), + supplier=supplier, version=tool["version"], model=tool["assetId"], description=tool.get("description", ""), @@ -34,37 +36,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..3f65b4b0 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,67 +48,73 @@ 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 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, - } - - 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 @@ -119,33 +123,34 @@ 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) -> 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,14 +158,18 @@ 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}") - r = _request_with_retry("get", url, headers=headers) - resp = r.json() + logging.info(f"Start service for GET Team Agent - {url} - {headers}") + 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: 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..da859a43 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,24 @@ def build_team_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> TeamAgent: """Instantiate a new team agent in the platform.""" - agents = payload["agents"] - for i, agent in enumerate(agents): + 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[i] = agent + agents.append(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=True if payload["plannerId"] is not None else False, api_key=api_key, status=AssetStatus(payload["status"]), ) diff --git a/aixplain/modules/agent/__init__.py b/aixplain/modules/agent/__init__.py index 3f892723..8d7391af 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,18 @@ def run_async( response["error"] = msg return response + 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.value["code"] if isinstance(self.supplier, Supplier) else self.supplier, + "version": self.version, + "llmId": self.llm_id, + "status": self.status.value, + } + def delete(self) -> None: """Delete Agent service""" try: @@ -259,4 +291,34 @@ def delete(self) -> None: except ValueError: message = f"Agent Deletion Error (HTTP {r.status_code}): There was an error in deleting the agent." logging.error(message) - raise Exception(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 404ed8d7..628377a3 100644 --- a/aixplain/modules/agent/tool/model_tool.py +++ b/aixplain/modules/agent/tool/model_tool.py @@ -79,6 +79,26 @@ def __init__( self.model = model self.function = function + 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": supplier, + "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/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 diff --git a/aixplain/modules/team_agent/__init__.py b/aixplain/modules/team_agent/__init__.py index 86321489..2f8b5c3b 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,68 @@ def delete(self) -> None: ) logging.error(message) raise Exception(f"{message}") + + 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) + ], + "links": [], + "description": self.description, + "llmId": self.llm_id, + "supervisorId": self.llm_id, + "plannerId": self.llm_id if self.use_mentalist_and_inspector else None, + "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() diff --git a/tests/functional/agent/agent_functional_test.py b/tests/functional/agent/agent_functional_test.py index 478b23f3..5b247728 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, TeamAgentFactory +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 @@ -56,14 +59,15 @@ def test_end2end(run_input_map, delete_agents_and_team_agents): 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"])) @@ -72,6 +76,11 @@ def test_end2end(run_input_map, delete_agents_and_team_agents): 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"]) @@ -91,6 +100,43 @@ def test_list_agents(): assert type(agents_result) is list +def test_update_draft_agent(run_input_map): + for team in TeamAgentFactory.list()["results"]: + team.delete() + + 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/functional/team_agent/team_agent_functional_test.py b/tests/functional/team_agent/team_agent_functional_test.py index c28b01da..44ea5dbc 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" @@ -59,14 +61,15 @@ def test_end2end(run_input_map, delete_agents_and_team_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 + 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"])) @@ -74,6 +77,7 @@ def test_end2end(run_input_map, delete_agents_and_team_agents): 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( @@ -85,6 +89,9 @@ def test_end2end(run_input_map, delete_agents_and_team_agents): ) 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"]) @@ -99,6 +106,51 @@ def test_end2end(run_input_map, delete_agents_and_team_agents): 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( diff --git a/tests/unit/agent_test.py b/tests/unit/agent_test.py index 43d0d0a2..1b4fd929 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": [ @@ -145,3 +152,77 @@ def test_create_agent(): assert agent.llm_id == ref_response["llmId"] assert agent.tools[0].function.value == ref_response["assets"][0]["function"] assert agent.tools[0].description == ref_response["assets"][0]["description"] + 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..56564b73 100644 --- a/tests/unit/team_agent_test.py +++ b/tests/unit/team_agent_test.py @@ -1,8 +1,12 @@ import pytest import requests_mock -from aixplain.modules import TeamAgent +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(): @@ -71,3 +75,122 @@ 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" + + +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"]