diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2500c2c8..4ad3e5bc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -8,18 +8,26 @@ repos: pass_filenames: false types: [python] always_run: true - + - repo: https://github.com/psf/black - rev: 22.10.0 + rev: 25.1.0 hooks: - id: black language_version: python3 args: # arguments to configure black - --line-length=128 - + - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v2.0.0 # Use the latest version + rev: v5.0.0 # Use the latest version + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-merge-conflict + - id: check-added-large-files + + - repo: https://github.com/pycqa/flake8 + rev: 7.2.0 hooks: - id: flake8 args: # arguments to configure flake8 - - --ignore=E402,E501,E203,W503 \ No newline at end of file + - --ignore=E402,E501,E203,W503 diff --git a/aixplain/enums/__init__.py b/aixplain/enums/__init__.py index 17308467..5c3a71f5 100644 --- a/aixplain/enums/__init__.py +++ b/aixplain/enums/__init__.py @@ -20,4 +20,5 @@ from .asset_status import AssetStatus from .index_stores import IndexStores from .function_type import FunctionType +from .evolve_type import EvolveType from .code_interpreter import CodeInterpreterModel diff --git a/aixplain/enums/evolve_type.py b/aixplain/enums/evolve_type.py new file mode 100644 index 00000000..555fdb53 --- /dev/null +++ b/aixplain/enums/evolve_type.py @@ -0,0 +1,6 @@ +from enum import Enum + + +class EvolveType(str, Enum): + TEAM_TUNING = "team_tuning" + INSTRUCTION_TUNING = "instruction_tuning" diff --git a/aixplain/factories/agent_factory/__init__.py b/aixplain/factories/agent_factory/__init__.py index a7662b81..5c90d887 100644 --- a/aixplain/factories/agent_factory/__init__.py +++ b/aixplain/factories/agent_factory/__init__.py @@ -81,7 +81,7 @@ def create( Args: name (Text): name of the agent - description (Text): description of the agent role. + description (Text): description of the agent instructions. instructions (Text): instructions of the agent. llm (Optional[Union[LLM, Text]], optional): LLM instance to use as an object or as an ID. llm_id (Optional[Text], optional): ID of LLM to use if no LLM instance provided. Defaults to None. diff --git a/aixplain/factories/team_agent_factory/utils.py b/aixplain/factories/team_agent_factory/utils.py index ee0f9d2d..db9f5d7b 100644 --- a/aixplain/factories/team_agent_factory/utils.py +++ b/aixplain/factories/team_agent_factory/utils.py @@ -1,12 +1,14 @@ __author__ = "lucaspavanelli" import logging -from typing import Dict, Text, List +from typing import Dict, Text, List, Optional from urllib.parse import urljoin import aixplain.utils.config as config from aixplain.enums.asset_status import AssetStatus from aixplain.modules.agent import Agent +from aixplain.modules.agent.agent_task import AgentTask +from aixplain.modules.agent.tool.model_tool import ModelTool from aixplain.modules.team_agent import TeamAgent, InspectorTarget from aixplain.modules.team_agent.inspector import Inspector from aixplain.factories.agent_factory import AgentFactory @@ -29,7 +31,7 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text = - name: Team agent name - agents: List of agent configurations - description: Optional description - - role: Optional instructions + - instructions: Optional instructions - teamId: Optional supplier information - version: Optional version - cost: Optional cost information @@ -112,7 +114,10 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text = # Convert parameters list to dictionary format expected by ModelParameters params_dict = {} for param in tool["parameters"]: - params_dict[param["name"]] = {"required": False, "value": param["value"]} + params_dict[param["name"]] = { + "required": False, + "value": param["value"], + } # Create ModelParameters and set it on the LLM llm.model_params = ModelParameters(params_dict) @@ -158,5 +163,158 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text = if task_dependency: team_agent.agents[idx].tasks[i].dependencies[j] = task_dependency else: - raise Exception(f"Team Agent Creation Error: Task dependency not found - {dependency}") + team_agent.agents[idx].tasks[i].dependencies[j] = None + return team_agent + + +def parse_tool_from_yaml(tool: str) -> ModelTool: + from aixplain.enums import Function + + tool_name = tool.strip() + if tool_name == "translation": + return ModelTool( + function=Function.TRANSLATION, + ) + elif tool_name == "speech-recognition": + return ModelTool( + function=Function.SPEECH_RECOGNITION, + ) + elif tool_name == "text-to-speech": + return ModelTool( + function=Function.SPEECH_SYNTHESIS, + ) + elif tool_name == "llm": + return ModelTool(function=Function.TEXT_GENERATION) + elif tool_name == "serper_search": + return ModelTool(model="65c51c556eb563350f6e1bb1") + elif tool.strip() == "website_search": + return ModelTool(model="6736411cf127849667606689") + elif tool.strip() == "website_scrape": + return ModelTool(model="6748e4746eb5633559668a15") + elif tool.strip() == "website_crawl": + return ModelTool(model="6748d4cff12784b6014324e2") + else: + raise Exception(f"Tool {tool} in yaml not found.") + + +import yaml + + +def is_yaml_formatted(text): + """ + Check if a string is valid YAML format with additional validation. + + Args: + text (str): The string to check + + Returns: + bool: True if valid YAML, False otherwise + """ + if not text or not isinstance(text, str): + return False + + # Strip whitespace + text = text.strip() + + # Empty string is valid YAML + if not text: + return True + + try: + parsed = yaml.safe_load(text) + + # If it's just a plain string without YAML structure, + # we might want to consider it as non-YAML + # This is optional depending on your requirements + if isinstance(parsed, str) and "\n" not in text and ":" not in text: + return False + + return True + except yaml.YAMLError: + return False + + +def build_team_agent_from_yaml(yaml_code: str, llm_id: str, api_key: str, team_id: Optional[str] = None) -> TeamAgent: + import yaml + from aixplain.factories import AgentFactory, TeamAgentFactory + + # check if it is a yaml or just as string + if not is_yaml_formatted(yaml_code): + return None + team_config = yaml.safe_load(yaml_code) + + agents_data = team_config["agents"] + tasks_data = team_config.get("tasks", []) + system_data = team_config["system"] if "system" in team_config else {"query": "", "name": "Test Team"} + team_name = system_data["name"] + + # Create agent mapping by name for easier task assignment + agents_mapping = {} + agent_objs = [] + + # Parse agents + for agent_entry in agents_data: + agent_name = list(agent_entry.keys())[0] + agent_info = agent_entry[agent_name] + agent_instructions = agent_info["instructions"] + agent_goal = agent_info["goal"] + agent_backstory = agent_info["backstory"] + + description = f"You are an expert {agent_instructions}. {agent_backstory} Your primary goal is to {agent_goal}. Use your expertise to ensure the success of your tasks." + agent_name = agent_name.replace("_", " ") + agent_name = f"{agent_name} agent" if not agent_name.endswith(" agent") else agent_name + agent_obj = Agent( + id="", + name=agent_name, + description=description, + instructions=description, + tasks=[], # Tasks will be assigned later + tools=[parse_tool_from_yaml(tool) for tool in agent_info.get("tools", []) if tool != "language_model"], + llmId=llm_id, + ) + agents_mapping[agent_name] = agent_obj + agent_objs.append(agent_obj) + + # Parse tasks and assign them to the corresponding agents + for task in tasks_data: + for task_name, task_info in task.items(): + description = task_info["description"] + expected_output = task_info["expected_output"] + dependencies = task_info.get("dependencies", []) + agent_name = task_info["agent"] + agent_name = agent_name.replace("_", " ") + agent_name = f"{agent_name} agent" if not agent_name.endswith(" agent") else agent_name + + task_obj = AgentTask( + name=task_name, + description=description, + expected_output=expected_output, + dependencies=dependencies, + ) + + # Assign the task to the corresponding agent + if agent_name in agents_mapping: + agent = agents_mapping[agent_name] + agent.tasks.append(task_obj) + else: + raise Exception(f"Agent '{agent_name}' referenced in tasks not found.") + + for i, agent in enumerate(agent_objs): + agent_objs[i] = AgentFactory.create( + name=agent.name, + description=agent.instructions, + instructions=agent.instructions, + tools=agent.tools, + llm_id=llm_id, + tasks=agent.tasks, + ) + + return TeamAgentFactory.create( + name=team_name, + agents=agent_objs, + llm_id=llm_id, + api_key=api_key, + use_mentalist=True, + inspectors=[], + ) diff --git a/aixplain/modules/agent/__init__.py b/aixplain/modules/agent/__init__.py index 48a19d21..af32bdb3 100644 --- a/aixplain/modules/agent/__init__.py +++ b/aixplain/modules/agent/__init__.py @@ -29,6 +29,7 @@ from aixplain.utils.file_utils import _request_with_retry from aixplain.enums import Function, Supplier, AssetStatus, StorageType, ResponseStatus +from aixplain.enums.evolve_type import EvolveType from aixplain.modules.model import Model from aixplain.modules.agent.agent_task import WorkflowTask, AgentTask from aixplain.modules.agent.output_format import OutputFormat @@ -37,7 +38,8 @@ from aixplain.modules.agent.agent_response_data import AgentResponseData from aixplain.modules.agent.utils import process_variables, validate_history from pydantic import BaseModel -from typing import Dict, List, Text, Optional, Union +from typing import Dict, List, Text, Optional, Union, Any +from aixplain.modules.agent.evolve_param import EvolveParam, validate_evolve_param from urllib.parse import urljoin from aixplain.modules.model.llm_model import LLM @@ -384,6 +386,7 @@ def run_async( max_iterations: int = 10, output_format: Optional[OutputFormat] = None, expected_output: Optional[Union[BaseModel, Text, dict]] = None, + evolve: Union[Dict[str, Any], EvolveParam, None] = None, ) -> AgentResponse: """Runs asynchronously an agent call. @@ -399,6 +402,8 @@ def run_async( max_iterations (int, optional): maximum number of iterations between the agent and the tools. Defaults to 10. output_format (OutputFormat, optional): response format. If not provided, uses the format set during initialization. expected_output (Union[BaseModel, Text, dict], optional): expected output. Defaults to None. + output_format (ResponseFormat, optional): response format. Defaults to TEXT. + evolve (Union[Dict[str, Any], EvolveParam, None], optional): evolve the agent configuration. Can be a dictionary, EvolveParam instance, or None. Returns: dict: polling URL in response """ @@ -415,8 +420,9 @@ def run_async( from aixplain.factories.file_factory import FileFactory - if not self.is_valid: - raise Exception("Agent is not valid. Please validate the agent before running.") + # Validate and normalize evolve parameters using the base model + evolve_param = validate_evolve_param(evolve) + evolve_dict = evolve_param.to_dict() if output_format == OutputFormat.JSON: assert expected_output is not None and ( @@ -478,6 +484,7 @@ def run_async( "outputFormat": output_format, "expectedOutput": expected_output, }, + "evolve": json.dumps(evolve_dict), } payload.update(parameters) payload = json.dumps(payload) @@ -594,7 +601,7 @@ def from_dict(cls, data: Dict) -> "Agent": id=data["id"], name=data["name"], description=data["description"], - instructions=data.get("role"), + instructions=data.get("instructions"), tools=tools, llm_id=data.get("llmId", "6646261c6eb563165658bbb1"), llm=llm, @@ -736,3 +743,128 @@ def __repr__(self) -> str: str: A string in the format "Agent: (id=)". """ return f"Agent: {self.name} (id={self.id})" + + def evolve_async( + self, + evolve_type: Union[EvolveType, str] = EvolveType.TEAM_TUNING, + max_successful_generations: int = 3, + max_failed_generation_retries: int = 3, + recursion_limit: int = 50, + max_non_improving_generations: Optional[int] = 2, + evolver_llm: Optional[Union[Text, LLM]] = None, + ) -> AgentResponse: + """Asynchronously evolve the Agent and return a polling URL in the AgentResponse. + + Args: + evolve_type (Union[EvolveType, str]): Type of evolution (TEAM_TUNING or INSTRUCTION_TUNING). Defaults to TEAM_TUNING. + max_successful_generations (int): Maximum number of successful generations to evolve. Defaults to 3. + max_failed_generation_retries (int): Maximum retry attempts for failed generations. Defaults to 3. + recursion_limit (int): Limit for recursive operations. Defaults to 50. + max_non_improving_generations (Optional[int]): Stop condition parameter for non-improving generations. Defaults to 2, can be None. + evolver_llm (Optional[Union[Text, LLM]]): LLM to use for evolution. Can be an LLM ID string or LLM object. Defaults to None. + + Returns: + AgentResponse: Response containing polling URL and status. + """ + from aixplain.utils.evolve_utils import create_evolver_llm_dict + + query = "" + + # Create EvolveParam from individual parameters + evolve_parameters = EvolveParam( + to_evolve=True, + evolve_type=evolve_type, + max_generations=max_successful_generations, + max_retries=max_failed_generation_retries, + recursion_limit=recursion_limit, + max_iterations_without_improvement=max_non_improving_generations, + evolver_llm=create_evolver_llm_dict(evolver_llm), + ) + + return self.run_async(query=query, evolve=evolve_parameters) + + def evolve( + self, + evolve_type: Union[EvolveType, str] = EvolveType.TEAM_TUNING, + max_successful_generations: int = 3, + max_failed_generation_retries: int = 3, + recursion_limit: int = 50, + max_non_improving_generations: Optional[int] = 2, + evolver_llm: Optional[Union[Text, LLM]] = None, + ) -> AgentResponse: + """Synchronously evolve the Agent and poll for the result. + + Args: + evolve_type (Union[EvolveType, str]): Type of evolution (TEAM_TUNING or INSTRUCTION_TUNING). Defaults to TEAM_TUNING. + max_successful_generations (int): Maximum number of successful generations to evolve. Defaults to 3. + max_failed_generation_retries (int): Maximum retry attempts for failed generations. Defaults to 3. + recursion_limit (int): Limit for recursive operations. Defaults to 50. + max_non_improving_generations (Optional[int]): Stop condition parameter for non-improving generations. Defaults to 2, can be None. + evolver_llm (Optional[Union[Text, LLM]]): LLM to use for evolution. Can be an LLM ID string or LLM object. Defaults to None. + + Returns: + AgentResponse: Final response from the evolution process. + """ + from aixplain.utils.evolve_utils import from_yaml, create_evolver_llm_dict + + # Create EvolveParam from individual parameters + evolve_parameters = EvolveParam( + to_evolve=True, + evolve_type=evolve_type, + max_generations=max_successful_generations, + max_retries=max_failed_generation_retries, + recursion_limit=recursion_limit, + max_iterations_without_improvement=max_non_improving_generations, + evolver_llm=create_evolver_llm_dict(evolver_llm), + ) + + start = time.time() + try: + logging.info(f"Evolve started with parameters: {evolve_parameters}") + logging.info("It might take a while...") + response = self.evolve_async( + evolve_type=evolve_type, + max_successful_generations=max_successful_generations, + max_failed_generation_retries=max_failed_generation_retries, + recursion_limit=recursion_limit, + max_non_improving_generations=max_non_improving_generations, + evolver_llm=evolver_llm, + ) + if response["status"] == ResponseStatus.FAILED: + end = time.time() + response["elapsed_time"] = end - start + return response + poll_url = response["url"] + end = time.time() + result = self.sync_poll(poll_url, name="evolve_process", timeout=600) + result_data = result.data + + if "current_code" in result_data and result_data["current_code"] is not None: + if evolve_parameters.evolve_type == EvolveType.TEAM_TUNING: + result_data["evolved_agent"] = from_yaml( + result_data["current_code"], + self.llm_id, + ) + elif evolve_parameters.evolve_type == EvolveType.INSTRUCTION_TUNING: + self.instructions = result_data["current_code"] + self.update() + result_data["evolved_agent"] = self + else: + raise ValueError( + "evolve_parameters.evolve_type must be one of the following: TEAM_TUNING, INSTRUCTION_TUNING" + ) + return AgentResponse( + status=ResponseStatus.SUCCESS, + completed=True, + data=result_data, + used_credits=getattr(result, "used_credits", 0.0), + run_time=getattr(result, "run_time", end - start), + ) + except Exception as e: + logging.error(f"Agent Evolve: Error in evolving: {e}") + end = time.time() + return AgentResponse( + status=ResponseStatus.FAILED, + completed=False, + error_message="No response from the service.", + ) diff --git a/aixplain/modules/agent/agent_response.py b/aixplain/modules/agent/agent_response.py index fc35e072..d10cff29 100644 --- a/aixplain/modules/agent/agent_response.py +++ b/aixplain/modules/agent/agent_response.py @@ -1,8 +1,11 @@ from aixplain.enums import ResponseStatus -from typing import Any, Dict, Optional, Text, Union, List +from typing import Any, Dict, Optional, Text, Union, List, TYPE_CHECKING from aixplain.modules.agent.agent_response_data import AgentResponseData from aixplain.modules.model.response import ModelResponse +if TYPE_CHECKING: + from aixplain.modules.team_agent.evolver_response_data import EvolverResponseData + class AgentResponse(ModelResponse): """A response object for agent execution results. @@ -25,7 +28,7 @@ class AgentResponse(ModelResponse): def __init__( self, status: ResponseStatus = ResponseStatus.FAILED, - data: Optional[AgentResponseData] = None, + data: Optional[Union[AgentResponseData, "EvolverResponseData"]] = None, details: Optional[Union[Dict, List]] = {}, completed: bool = False, error_message: Text = "", diff --git a/aixplain/modules/agent/agent_response_data.py b/aixplain/modules/agent/agent_response_data.py index 4de10d12..b97b16c1 100644 --- a/aixplain/modules/agent/agent_response_data.py +++ b/aixplain/modules/agent/agent_response_data.py @@ -95,6 +95,9 @@ def to_dict(self) -> Dict[str, Any]: "critiques": self.critiques, } + def get(self, key: str, default: Optional[Any] = None) -> Any: + return getattr(self, key, default) + def __getitem__(self, key: str) -> Any: """Get an attribute value using dictionary-style access. diff --git a/aixplain/modules/agent/agent_task.py b/aixplain/modules/agent/agent_task.py index 74db8517..fa2ea66c 100644 --- a/aixplain/modules/agent/agent_task.py +++ b/aixplain/modules/agent/agent_task.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Text, Union +from typing import List, Text, Union class WorkflowTask: diff --git a/aixplain/modules/agent/evolve_param.py b/aixplain/modules/agent/evolve_param.py new file mode 100644 index 00000000..c622f90b --- /dev/null +++ b/aixplain/modules/agent/evolve_param.py @@ -0,0 +1,256 @@ +__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: aiXplain Team +Date: December 2024 +Description: + EvolveParam Base Model Class for Agent and TeamAgent evolve functionality +""" +from aixplain.enums import EvolveType +from dataclasses import dataclass, field +from typing import Any, Dict, Optional, Union + + +@dataclass +class EvolveParam: + """Base model for evolve parameters used in Agent and TeamAgent evolution. + + Attributes: + to_evolve (bool): Whether to enable evolution. Defaults to False. + evolve_type (Optional[EvolveType]): Type of evolve. + max_generations (int): Maximum number of generations. + max_retries (int): Maximum number of retries. + recursion_limit (int): Maximum number of recursion. + max_iterations_without_improvement (int): Maximum number of iterations without improvement. + evolver_llm (Optional[Dict[str, Any]]): Evolver LLM configuration with all parameters. + additional_params (Optional[Dict[str, Any]]): Additional parameters. + """ + + to_evolve: bool = False + evolve_type: Optional[EvolveType] = EvolveType.TEAM_TUNING + max_generations: int = 3 + max_retries: int = 3 + recursion_limit: int = 50 + max_iterations_without_improvement: Optional[int] = 2 + evolver_llm: Optional[Dict[str, Any]] = None + additional_params: Optional[Dict[str, Any]] = field(default_factory=dict) + + def __post_init__(self): + """Validate parameters after initialization.""" + self.validate() + + def validate(self) -> None: + """Validate evolve parameters. + + Raises: + ValueError: If any parameter is invalid. + """ + if self.evolve_type is not None: + if isinstance(self.evolve_type, str): + # Convert string to EvolveType + try: + self.evolve_type = EvolveType(self.evolve_type) + except ValueError: + raise ValueError( + f"evolve_type '{self.evolve_type}' is not a valid EvolveType. Valid values are: {list(EvolveType)}" + ) + elif not isinstance(self.evolve_type, EvolveType): + raise ValueError("evolve_type must be a valid EvolveType or string") + if self.additional_params is not None: + if not isinstance(self.additional_params, dict): + raise ValueError("additional_params must be a dictionary") + + if self.max_generations is not None: + if not isinstance(self.max_generations, int): + raise ValueError("max_generations must be an integer") + if self.max_generations <= 0: + raise ValueError("max_generations must be positive") + + if self.max_retries is not None: + if not isinstance(self.max_retries, int): + raise ValueError("max_retries must be an integer") + if self.max_retries <= 0: + raise ValueError("max_retries must be positive") + + if self.recursion_limit is not None: + if not isinstance(self.recursion_limit, int): + raise ValueError("recursion_limit must be an integer") + if self.recursion_limit <= 0: + raise ValueError("recursion_limit must be positive") + + if self.max_iterations_without_improvement is not None: + if not isinstance(self.max_iterations_without_improvement, int): + raise ValueError("max_iterations_without_improvement must be an integer or None") + if self.max_iterations_without_improvement <= 0: + raise ValueError("max_iterations_without_improvement must be positive or None") + + @classmethod + def from_dict(cls, data: Union[Dict[str, Any], None]) -> "EvolveParam": + """Create EvolveParam instance from dictionary. + + Args: + data (Union[Dict[str, Any], None]): Dictionary containing evolve parameters. + + Returns: + EvolveParam: Instance with parameters set from dictionary. + + Raises: + ValueError: If data format is invalid. + """ + if data is None: + return cls() + + if not isinstance(data, dict): + raise ValueError("evolve parameter must be a dictionary or None") + + # Extract known parameters + known_params = { + "to_evolve": data.get("toEvolve", data.get("to_evolve", False)), + "evolve_type": data.get("evolve_type"), + "max_generations": data.get("max_generations"), + "max_retries": data.get("max_retries"), + "recursion_limit": data.get("recursion_limit"), + "max_iterations_without_improvement": data.get("max_iterations_without_improvement"), + "evolver_llm": data.get("evolver_llm"), + "additional_params": data.get("additional_params"), + } + + # Remove None values + known_params = {k: v for k, v in known_params.items() if v is not None} + + # Collect additional parameters + additional_params = { + k: v + for k, v in data.items() + if k + not in [ + "toEvolve", + "to_evolve", + "evolve_type", + "max_generations", + "max_retries", + "recursion_limit", + "max_iterations_without_improvement", + "evolver_llm", + "additional_params", + ] + } + + return cls(additional_params=additional_params, **known_params) + + def to_dict(self) -> Dict[str, Any]: + """Convert EvolveParam instance to dictionary for API calls. + + Returns: + Dict[str, Any]: Dictionary representation with API-compatible keys. + """ + result = { + "toEvolve": self.to_evolve, + } + + # Add optional parameters if they are set + if self.evolve_type is not None: + result["evolve_type"] = self.evolve_type + if self.max_generations is not None: + result["max_generations"] = self.max_generations + if self.max_retries is not None: + result["max_retries"] = self.max_retries + if self.recursion_limit is not None: + result["recursion_limit"] = self.recursion_limit + # Always include max_iterations_without_improvement, even if None + result["max_iterations_without_improvement"] = self.max_iterations_without_improvement + if self.evolver_llm is not None: + result["evolver_llm"] = self.evolver_llm + if self.additional_params is not None: + result.update(self.additional_params) + + return result + + def merge(self, other: Union[Dict[str, Any], "EvolveParam"]) -> "EvolveParam": + """Merge this EvolveParam with another set of parameters. + + Args: + other (Union[Dict[str, Any], EvolveParam]): Other parameters to merge. + + Returns: + EvolveParam: New instance with merged parameters. + """ + if isinstance(other, dict): + other = EvolveParam.from_dict(other) + elif not isinstance(other, EvolveParam): + raise ValueError("other must be a dictionary or EvolveParam instance") + + # Create merged parameters + merged_additional = {**self.additional_params, **other.additional_params} + + return EvolveParam( + to_evolve=other.to_evolve if other.to_evolve else self.to_evolve, + evolve_type=(other.evolve_type if other.evolve_type is not None else self.evolve_type), + max_generations=(other.max_generations if other.max_generations is not None else self.max_generations), + max_retries=(other.max_retries if other.max_retries is not None else self.max_retries), + recursion_limit=(other.recursion_limit if other.recursion_limit is not None else self.recursion_limit), + max_iterations_without_improvement=( + other.max_iterations_without_improvement + if other.max_iterations_without_improvement is not None + else self.max_iterations_without_improvement + ), + evolver_llm=(other.evolver_llm if other.evolver_llm is not None else self.evolver_llm), + additional_params=merged_additional, + ) + + def __repr__(self) -> str: + return ( + f"EvolveParam(" + f"to_evolve={self.to_evolve}, " + f"evolve_type={self.evolve_type}, " + f"max_generations={self.max_generations}, " + f"max_retries={self.max_retries}, " + f"recursion_limit={self.recursion_limit}, " + f"max_iterations_without_improvement={self.max_iterations_without_improvement}, " + f"evolver_llm={self.evolver_llm}, " + f"additional_params={self.additional_params})" + ) + + +def validate_evolve_param( + evolve_param: Union[Dict[str, Any], EvolveParam, None], +) -> EvolveParam: + """Utility function to validate and convert evolve parameters. + + Args: + evolve_param (Union[Dict[str, Any], EvolveParam, None]): Input evolve parameters. + + Returns: + EvolveParam: Validated EvolveParam instance. + + Raises: + ValueError: If parameters are invalid. + """ + if evolve_param is None: + return EvolveParam() + + if isinstance(evolve_param, EvolveParam): + evolve_param.validate() + return evolve_param + + if isinstance(evolve_param, dict): + # Check for required toEvolve key for backward compatibility + if "toEvolve" not in evolve_param and "to_evolve" not in evolve_param: + raise ValueError("evolve parameter must contain 'toEvolve' key") + return EvolveParam.from_dict(evolve_param) + + raise ValueError("evolve parameter must be a dictionary, EvolveParam instance, or None") diff --git a/aixplain/modules/team_agent/__init__.py b/aixplain/modules/team_agent/__init__.py index 497664c7..077026ac 100644 --- a/aixplain/modules/team_agent/__init__.py +++ b/aixplain/modules/team_agent/__init__.py @@ -27,7 +27,7 @@ import traceback import re from enum import Enum -from typing import Dict, List, Text, Optional, Union +from typing import Dict, List, Text, Optional, Union, Any from urllib.parse import urljoin from datetime import datetime @@ -36,12 +36,15 @@ from aixplain.enums.supplier import Supplier from aixplain.enums.asset_status import AssetStatus from aixplain.enums.storage_type import StorageType +from aixplain.enums.evolve_type import EvolveType from aixplain.modules.model import Model from aixplain.modules.agent import Agent, OutputFormat from aixplain.modules.agent.agent_response import AgentResponse from aixplain.modules.agent.agent_response_data import AgentResponseData +from aixplain.modules.agent.evolve_param import EvolveParam, validate_evolve_param from aixplain.modules.agent.utils import process_variables, validate_history from aixplain.modules.team_agent.inspector import Inspector +from aixplain.modules.team_agent.evolver_response_data import EvolverResponseData from aixplain.utils import config from aixplain.utils.request_utils import _request_with_retry from aixplain.modules.model.llm_model import LLM @@ -123,6 +126,7 @@ def __init__( self.agents = agents self.llm_id = llm_id self.llm = llm + self.api_key = api_key self.use_mentalist = use_mentalist self.inspectors = inspectors self.inspector_targets = inspector_targets @@ -213,7 +217,7 @@ def run( output_format (OutputFormat, optional): response format. If not provided, uses the format set during initialization. expected_output (Union[BaseModel, Text, dict], optional): expected output. Defaults to None. Returns: - Dict: parsed output from model + AgentResponse: parsed output from model """ start = time.time() result_data = {} @@ -283,6 +287,7 @@ def run_async( max_iterations: int = 30, output_format: Optional[OutputFormat] = None, expected_output: Optional[Union[BaseModel, Text, dict]] = None, + evolve: Union[Dict[str, Any], EvolveParam, None] = None, ) -> AgentResponse: """Runs asynchronously a Team Agent call. @@ -298,8 +303,9 @@ def run_async( max_iterations (int, optional): maximum number of iterations between the agents. Defaults to 30. output_format (OutputFormat, optional): response format. If not provided, uses the format set during initialization. expected_output (Union[BaseModel, Text, dict], optional): expected output. Defaults to None. + evolve (Union[Dict[str, Any], EvolveParam, None], optional): evolve the team agent configuration. Can be a dictionary, EvolveParam instance, or None. Returns: - dict: polling URL in response + AgentResponse: polling URL in response """ if session_id is not None and history is not None: raise ValueError("Provide either `session_id` or `history`, not both.") @@ -313,6 +319,10 @@ def run_async( from aixplain.factories.file_factory import FileFactory + # Validate and normalize evolve parameters using the base model + evolve_param = validate_evolve_param(evolve) + evolve_dict = evolve_param.to_dict() + if not self.is_valid: raise Exception("Team Agent is not valid. Please validate the team agent before running.") @@ -371,6 +381,7 @@ def run_async( "outputFormat": output_format, "expectedOutput": expected_output, }, + "evolve": json.dumps(evolve_dict), } payload.update(parameters) payload = json.dumps(payload) @@ -384,7 +395,7 @@ def run_async( logging.info(f"Result of request for {name} - {r.status_code} - {resp}") poll_url = resp["data"] - return AgentResponse( + response = AgentResponse( status=ResponseStatus.IN_PROGRESS, url=poll_url, data=AgentResponseData(input=input_data), @@ -394,10 +405,70 @@ def run_async( except Exception: msg = f"Error in request for {name} - {traceback.format_exc()}" logging.error(f"Team Agent Run Async: Error in running for {name}: {resp}") - return AgentResponse( - status=ResponseStatus.FAILED, - error=msg, + if resp is not None: + response = AgentResponse( + status=ResponseStatus.FAILED, + error=msg, + ) + return response + + def poll(self, poll_url: Text, name: Text = "model_process") -> AgentResponse: + used_credits, run_time = 0.0, 0.0 + resp, error_message, status = None, None, ResponseStatus.SUCCESS + headers = {"x-api-key": self.api_key, "Content-Type": "application/json"} + r = _request_with_retry("get", poll_url, headers=headers) + try: + resp = r.json() + if resp["completed"] is True: + status = ResponseStatus(resp.get("status", "FAILED")) + if "error_message" in resp or "supplierError" in resp: + status = ResponseStatus.FAILED + error_message = resp.get("error_message") + else: + status = ResponseStatus.IN_PROGRESS + logging.debug(f"Single Poll for Team Agent: Status of polling for {name}: {resp}") + + resp_data = resp.get("data") or {} + used_credits = resp_data.get("usedCredits", 0.0) + run_time = resp_data.get("runTime", 0.0) + evolve_type = resp_data.get("evolve_type", EvolveType.TEAM_TUNING.value) + if "evolved_agent" in resp_data and status == ResponseStatus.SUCCESS: + if evolve_type == EvolveType.INSTRUCTION_TUNING.value: + # return this class as it is but replace its description and instructions + evolved_agent = self + current_code = resp_data.get("current_code", "") + evolved_agent.description = current_code + evolved_agent.instructions = current_code + evolved_agent.update() + resp_data["evolved_agent"] = evolved_agent + else: + resp_data = EvolverResponseData.from_dict(resp_data, llm_id=self.llm_id, api_key=self.api_key) + else: + resp_data = AgentResponseData( + input=resp_data.get("input"), + output=resp_data.get("output"), + session_id=resp_data.get("session_id"), + intermediate_steps=resp_data.get("intermediate_steps"), + execution_stats=resp_data.get("executionStats"), + ) + except Exception as e: + import traceback + + logging.error(f"Single Poll for Team Agent: Error of polling for {name}: {e}, traceback: {traceback.format_exc()}") + status = ResponseStatus.FAILED + error_message = str(e) + finally: + response = AgentResponse( + status=status, + data=resp_data, + details=resp.get("details", {}), + completed=resp.get("completed", False), + used_credits=used_credits, + run_time=run_time, + usage=resp.get("usage", None), + error_message=error_message, ) + return response def delete(self) -> None: """Delete Corpus service""" @@ -480,7 +551,7 @@ def to_dict(self) -> Dict: - supplier (str): The supplier code - version (str): The version number - status (str): The current status - - role (str): The team agent's instructions + - instructions (str): The team agent's instructions """ if self.use_mentalist: planner_id = self.mentalist_llm.id if self.mentalist_llm else self.llm_id @@ -493,11 +564,11 @@ def to_dict(self) -> Dict: "links": [], "description": self.description, "llmId": self.llm.id if self.llm else self.llm_id, - "supervisorId": self.supervisor_llm.id if self.supervisor_llm else self.llm_id, + "supervisorId": (self.supervisor_llm.id if self.supervisor_llm else self.llm_id), "plannerId": planner_id, "inspectors": [inspector.model_dump(by_alias=True) for inspector in self.inspectors], "inspectorTargets": [target.value for target in self.inspector_targets], - "supplier": self.supplier.value["code"] if isinstance(self.supplier, Supplier) else self.supplier, + "supplier": (self.supplier.value["code"] if isinstance(self.supplier, Supplier) else self.supplier), "version": self.version, "status": self.status.value, "instructions": self.instructions, @@ -604,7 +675,7 @@ def from_dict(cls, data: Dict) -> "TeamAgent": version=data.get("version"), use_mentalist=use_mentalist, status=status, - instructions=data.get("role"), + instructions=data.get("instructions"), inspectors=inspectors, inspector_targets=inspector_targets, output_format=OutputFormat(data.get("outputFormat", OutputFormat.TEXT)), @@ -729,3 +800,130 @@ def __repr__(self): str: A string in the format "TeamAgent: (id=)". """ return f"TeamAgent: {self.name} (id={self.id})" + + def evolve_async( + self, + evolve_type: Union[EvolveType, str] = EvolveType.TEAM_TUNING, + max_successful_generations: int = 3, + max_failed_generation_retries: int = 3, + recursion_limit: int = 50, + max_non_improving_generations: Optional[int] = 2, + evolver_llm: Optional[Union[Text, LLM]] = None, + ) -> AgentResponse: + """Asynchronously evolve the Team Agent and return a polling URL in the AgentResponse. + + Args: + evolve_type (Union[EvolveType, str]): Type of evolution (TEAM_TUNING or INSTRUCTION_TUNING). Defaults to TEAM_TUNING. + max_successful_generations (int): Maximum number of successful generations to evolve. Defaults to 3. + max_failed_generation_retries (int): Maximum retry attempts for failed generations. Defaults to 3. + recursion_limit (int): Limit for recursive operations. Defaults to 50. + max_non_improving_generations (Optional[int]): Stop condition parameter for non-improving generations. Defaults to 2, can be None. + evolver_llm (Optional[Union[Text, LLM]]): LLM to use for evolution. Can be an LLM ID string or LLM object. Defaults to None. + + Returns: + AgentResponse: Response containing polling URL and status. + """ + from aixplain.utils.evolve_utils import create_evolver_llm_dict + + query = "" + + # Create EvolveParam from individual parameters (map new names to old keys) + evolve_parameters = EvolveParam( + to_evolve=True, + evolve_type=evolve_type, + max_generations=max_successful_generations, + max_retries=max_failed_generation_retries, + recursion_limit=recursion_limit, + max_iterations_without_improvement=max_non_improving_generations, + evolver_llm=create_evolver_llm_dict(evolver_llm), + ) + + response = self.run_async(query=query, evolve=evolve_parameters) + return response + + def evolve( + self, + evolve_type: Union[EvolveType, str] = EvolveType.TEAM_TUNING, + max_successful_generations: int = 3, + max_failed_generation_retries: int = 3, + recursion_limit: int = 50, + max_non_improving_generations: Optional[int] = 2, + evolver_llm: Optional[Union[Text, LLM]] = None, + ) -> AgentResponse: + """Synchronously evolve the Team Agent and poll for the result. + + Args: + evolve_type (Union[EvolveType, str]): Type of evolution (TEAM_TUNING or INSTRUCTION_TUNING). Defaults to TEAM_TUNING. + max_successful_generations (int): Maximum number of successful generations to evolve. Defaults to 3. + max_failed_generation_retries (int): Maximum retry attempts for failed generations. Defaults to 3. + recursion_limit (int): Limit for recursive operations. Defaults to 50. + max_non_improving_generations (Optional[int]): Stop condition parameter for non-improving generations. Defaults to 2, can be None. + evolver_llm (Optional[Union[Text, LLM]]): LLM to use for evolution. Can be an LLM ID string or LLM object. Defaults to None. + + Returns: + AgentResponse: Final response from the evolution process. + """ + from aixplain.enums import EvolveType + from aixplain.utils.evolve_utils import from_yaml, create_evolver_llm_dict + + # Create EvolveParam from individual parameters (map new names to old keys) + evolve_parameters = EvolveParam( + to_evolve=True, + evolve_type=evolve_type, + max_generations=max_successful_generations, + max_retries=max_failed_generation_retries, + recursion_limit=recursion_limit, + max_iterations_without_improvement=max_non_improving_generations, + evolver_llm=create_evolver_llm_dict(evolver_llm), + ) + start = time.time() + try: + logging.info(f"Evolve started with parameters: {evolve_parameters}") + logging.info("It might take a while...") + response = self.evolve_async( + evolve_type=evolve_type, + max_successful_generations=max_successful_generations, + max_failed_generation_retries=max_failed_generation_retries, + recursion_limit=recursion_limit, + max_non_improving_generations=max_non_improving_generations, + evolver_llm=evolver_llm, + ) + if response["status"] == ResponseStatus.FAILED: + end = time.time() + response["elapsed_time"] = end - start + return response + poll_url = response["url"] + end = time.time() + result = self.sync_poll(poll_url, name="evolve_process", timeout=600) + result_data = result.data + current_code = result_data.get("current_code") if isinstance(result_data, dict) else result_data.current_code + if current_code is not None: + if evolve_parameters.evolve_type == EvolveType.TEAM_TUNING: + result_data["evolved_agent"] = from_yaml( + result_data["current_code"], + self.llm_id, + ) + elif evolve_parameters.evolve_type == EvolveType.INSTRUCTION_TUNING: + self.instructions = result_data["current_code"] + self.description = result_data["current_code"] + self.update() + result_data["evolved_agent"] = self + else: + raise ValueError( + "evolve_parameters.evolve_type must be one of the following: TEAM_TUNING, INSTRUCTION_TUNING" + ) + return AgentResponse( + status=ResponseStatus.SUCCESS, + completed=True, + data=result_data, + used_credits=getattr(result, "used_credits", 0.0), + run_time=getattr(result, "run_time", end - start), + ) + except Exception as e: + logging.error(f"Team Agent Evolve: Error in evolving: {e}") + end = time.time() + return AgentResponse( + status=ResponseStatus.FAILED, + completed=False, + error_message="No response from the service.", + ) diff --git a/aixplain/modules/team_agent/evolver_response_data.py b/aixplain/modules/team_agent/evolver_response_data.py new file mode 100644 index 00000000..cd2499b5 --- /dev/null +++ b/aixplain/modules/team_agent/evolver_response_data.py @@ -0,0 +1,71 @@ +from typing import Any, Dict, List, Text, TYPE_CHECKING + +if TYPE_CHECKING: + from aixplain.modules.team_agent import TeamAgent + + +class EvolverResponseData: + def __init__( + self, + evolved_agent: "TeamAgent", + current_code: Text, + evaluation_report: Text, + comparison_report: Text, + criteria: Text, + archive: List[Text], + current_output: Text = "", + ) -> None: + self.evolved_agent = evolved_agent + self.current_code = current_code + self.evaluation_report = evaluation_report + self.comparison_report = comparison_report + self.criteria = criteria + self.archive = archive + self.current_output = current_output + + @classmethod + def from_dict(cls, data: Dict[str, Any], llm_id: Text, api_key: Text) -> "EvolverResponseData": + from aixplain.factories.team_agent_factory.utils import build_team_agent_from_yaml + + yaml_code = data.get("current_code", "") + evolved_team_agent = build_team_agent_from_yaml(yaml_code=yaml_code, llm_id=llm_id, api_key=api_key) + + return cls( + evolved_agent=evolved_team_agent, + current_code=yaml_code, + evaluation_report=data.get("evaluation_report", ""), + comparison_report=data.get("comparison_report", ""), + criteria=data.get("criteria", ""), + archive=data.get("archive", []), + current_output=data.get("current_output", ""), + ) + + def to_dict(self) -> Dict[str, Any]: + return { + "evolved_agent": self.evolved_agent, + "current_code": self.current_code, + "evaluation_report": self.evaluation_report, + "comparison_report": self.comparison_report, + "criteria": self.criteria, + "archive": self.archive, + "current_output": self.current_output, + } + + def __getitem__(self, key: str) -> Any: + return getattr(self, key, None) + + def __setitem__(self, key: str, value: Any) -> None: + if hasattr(self, key): + setattr(self, key, value) + else: + raise KeyError(f"{key} is not a valid attribute of {self.__class__.__name__}") + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(" + f"evolved_agent='{self.evolved_agent}', " + f"evaluation_report='{self.evaluation_report}', " + f"comparison_report='{self.comparison_report}', " + f"criteria='{self.criteria}', " + f"archive='{self.archive}', " + ) diff --git a/aixplain/utils/evolve_utils.py b/aixplain/utils/evolve_utils.py new file mode 100644 index 00000000..8e1b1c92 --- /dev/null +++ b/aixplain/utils/evolve_utils.py @@ -0,0 +1,203 @@ +import yaml +from typing import Union, Dict, Any, Optional, Text +from aixplain.enums import Function, Supplier +from aixplain.modules.agent import Agent +from aixplain.modules.team_agent import TeamAgent +from aixplain.modules.model.llm_model import LLM +from aixplain.factories import AgentFactory, TeamAgentFactory + + +def create_evolver_llm_dict(evolver_llm: Optional[Union[Text, LLM]]) -> Optional[Dict[str, Any]]: + """Create a dictionary representation of an evolver LLM for evolution parameters. + + Args: + evolver_llm: Either an LLM ID string or an LLM object instance. + + Returns: + Dictionary with LLM information if evolver_llm is provided, None otherwise. + """ + if evolver_llm is None: + return None + + if isinstance(evolver_llm, LLM): + return { + "id": evolver_llm.id, + "name": evolver_llm.name, + "description": evolver_llm.description, + "supplier": evolver_llm.supplier, + "version": evolver_llm.version, + "function": evolver_llm.function, + "parameters": (evolver_llm.get_parameters().to_list() if evolver_llm.get_parameters() else None), + "temperature": getattr(evolver_llm, "temperature", None), + } + else: + return {"id": evolver_llm} + + +def generate_tools_from_str(tools: list) -> list: + obj_tools = [] + for tool in tools: + tool_name = tool.strip() + if tool_name == "translation": + obj_tools.append( + AgentFactory.create_model_tool( + function=Function.TRANSLATION, + supplier=Supplier.GOOGLE, + ) + ) + elif tool_name == "speech-recognition": + obj_tools.append( + AgentFactory.create_model_tool( + function=Function.SPEECH_RECOGNITION, + supplier=Supplier.GOOGLE, + ) + ) + elif tool_name == "text-to-speech": + obj_tools.append( + AgentFactory.create_model_tool( + function=Function.SPEECH_SYNTHESIS, + supplier=Supplier.GOOGLE, + ) + ) + elif tool_name == "serper_search": + obj_tools.append(AgentFactory.create_model_tool(model="65c51c556eb563350f6e1bb1")) + elif tool_name == "website_search": + obj_tools.append(AgentFactory.create_model_tool(model="6736411cf127849667606689")) + elif tool_name == "website_scrape": + obj_tools.append(AgentFactory.create_model_tool(model="6748e4746eb5633559668a15")) + else: + continue + return obj_tools + + +def from_yaml( + yaml_code: str, + llm_id: str, +) -> Union[TeamAgent, Agent]: + team_config = yaml.safe_load(yaml_code) + + agents_data = team_config["agents"] + tasks_data = team_config.get("tasks", []) + system_data = team_config["system"] if "system" in team_config else {"query": ""} + expected_output = system_data.get("expected_output") + team_name = system_data.get("name") + team_description = system_data.get("description", "") + team_instructions = system_data.get("instructions", "") + + # Create agent mapping by name for easier task assignment + agents_mapping = {} + agent_objs = [] + + # Parse agents + for agent_entry in agents_data: + # Handle different YAML structures + if isinstance(agent_entry, dict): + # Case 1: Standard structure - {agent_name: {instructions: ..., goal: ..., backstory: ...}} + for agent_name, agent_info in agent_entry.items(): + # Check if agent_info is a list (malformed YAML structure) + if isinstance(agent_info, list): + # Handle malformed structure where agent_info is a list + # This happens when YAML has unquoted keys that create nested structures + continue + elif isinstance(agent_info, dict): + agent_instructions = agent_info["instructions"] + agent_goal = agent_info["goal"] + agent_backstory = agent_info["backstory"] + + description = f"## ROLE\n{agent_instructions}\n\n## GOAL\n{agent_goal}\n\n## BACKSTORY\n{agent_backstory}" + agent_obj = AgentFactory.create( + name=agent_name.replace("_", " "), + description=description, + tasks=[], # Tasks will be assigned later + tools=generate_tools_from_str(agent_info.get("tools", [])), + llm_id=llm_id, + ) + agents_mapping[agent_name] = agent_obj + agent_objs.append(agent_obj) + elif isinstance(agent_entry, list): + # Case 2: Handle list structure (alternative YAML format) + for item in agent_entry: + if isinstance(item, dict): + for agent_name, agent_info in item.items(): + if isinstance(agent_info, dict): + agent_instructions = agent_info["instructions"] + agent_goal = agent_info["goal"] + agent_backstory = agent_info["backstory"] + + description = ( + f"## ROLE\n{agent_instructions}\n\n## GOAL\n{agent_goal}\n\n## BACKSTORY\n{agent_backstory}" + ) + agent_tools = generate_tools_from_str(agent_info.get("tools", [])) + agent_obj = AgentFactory.create( + name=agent_name.replace("_", " "), + description=description, + tools=agent_tools, + tasks=[], + llm_id=llm_id, + ) + agents_mapping[agent_name] = agent_obj + agent_objs.append(agent_obj) + + # Parse tasks and assign them to the corresponding agents + for task_entry in tasks_data: + # Handle different YAML structures for tasks + if isinstance(task_entry, dict): + # Case 1: Standard structure - {task_name: {description: ..., expected_output: ..., agent: ...}} + for task_name, task_info in task_entry.items(): + # Check if task_info is a list (malformed YAML structure) + if isinstance(task_info, list): + # Handle malformed structure where task_info is a list + continue + elif isinstance(task_info, dict): + description = task_info["description"] + expected_output = task_info["expected_output"] + dependencies = task_info.get("dependencies", []) + agent_name = task_info["agent"] + task_obj = AgentFactory.create_task( + name=task_name.replace("_", " "), + description=description, + expected_output=expected_output, + dependencies=dependencies, + ) + + # Assign the task to the corresponding agent + if agent_name in agents_mapping: + agent = agents_mapping[agent_name] + agent.tasks.append(task_obj) + else: + raise ValueError(f"Agent '{agent_name}' referenced in tasks not found.") + elif isinstance(task_entry, list): + # Case 2: Handle list structure (alternative YAML format) + for item in task_entry: + if isinstance(item, dict): + for task_name, task_info in item.items(): + if isinstance(task_info, dict): + description = task_info["description"] + expected_output = task_info["expected_output"] + dependencies = task_info.get("dependencies", []) + agent_name = task_info["agent"] + + task_obj = AgentFactory.create_task( + name=task_name.replace("_", " "), + description=description, + expected_output=expected_output, + dependencies=dependencies, + ) + + # Assign the task to the corresponding agent + if agent_name in agents_mapping: + agent = agents_mapping[agent_name] + agent.tasks.append(task_obj) + else: + raise ValueError(f"Agent '{agent_name}' referenced in tasks not found.") + + team = TeamAgentFactory.create( + name=team_name, + description=team_description, + instructions=team_instructions, + agents=agent_objs, + llm_id=llm_id, + use_mentalist=True, + inspectors=[], + ) + return team diff --git a/tests/functional/team_agent/evolver_test.py b/tests/functional/team_agent/evolver_test.py new file mode 100644 index 00000000..4edf0c39 --- /dev/null +++ b/tests/functional/team_agent/evolver_test.py @@ -0,0 +1,138 @@ +import pytest +from aixplain.enums.function import Function +from aixplain.enums.supplier import Supplier +from aixplain.enums import ResponseStatus +from aixplain.factories.agent_factory import AgentFactory +from aixplain.factories.team_agent_factory import TeamAgentFactory +import time + + +team_dict = { + "team_agent_name": "Test Text Speech Team", + "llm_id": "6646261c6eb563165658bbb1", + "llm_name": "GPT4o", + "query": "Translate this text into Portuguese: 'This is a test'. Translate to pt and synthesize in audio", + "description": "You are a text translation and speech synthesizing agent. You will be provided a text in the source language and expected to translate and synthesize in the target language.", + "agents": [ + { + "agent_name": "Text Translation agent", + "llm_id": "6646261c6eb563165658bbb1", + "llm_name": "GPT4o", + "description": "## ROLE\nText Translator\n\n## GOAL\nTranslate the text supplied into the users desired language.\n\n## BACKSTORY\nYou are a text translation agent. You will be provided a text in the source language and expected to translate in the target language.", + "tasks": [ + { + "name": "Text translation", + "description": "Translate a text from source language (English) to target language (Portuguese)", + "expected_output": "target language text", + } + ], + "model_tools": [{"function": "translation", "supplier": "AWS"}], + }, + { + "agent_name": "Test Speech Synthesis agent", + "llm_id": "6646261c6eb563165658bbb1", + "llm_name": "GPT4o", + "description": "## ROLE\nSpeech Synthesizer\n\n## GOAL\nTranscribe the translated text into speech.\n\n## BACKSTORY\nYou are a speech synthesizing agent. You will be provided a text to synthesize into audio and return the audio link.", + "tasks": [ + { + "name": "Speech synthesis", + "description": "Synthesize a text from text to speech", + "expected_output": "audio link of the synthesized text", + "dependencies": ["Text translation"], + } + ], + "model_tools": [{"function": "speech_synthesis", "supplier": "Google"}], + }, + ], +} + + +def parse_tools(tools_info): + tools = [] + for tool in tools_info: + function_enum = Function[tool["function"].upper().replace(" ", "_")] + supplier_enum = Supplier[tool["supplier"].upper().replace(" ", "_")] + tools.append(AgentFactory.create_model_tool(function=function_enum, supplier=supplier_enum)) + return tools + + +def build_team_agent_from_json(team_config: dict): + agents_data = team_config["agents"] + tasks_data = team_config.get("tasks", []) + + agent_objs = [] + for agent_entry in agents_data: + agent_name = agent_entry["agent_name"] + agent_description = agent_entry["description"] + agent_llm_id = agent_entry.get("llm_id", None) + + agent_tasks = [] + for task in tasks_data: + task_name = task.get("task_name", "") + task_info = task + + if agent_name == task_info["agent"]: + task_obj = AgentFactory.create_task( + name=task_name.replace("_", " "), + description=task_info.get("description", ""), + expected_output=task_info.get("expected_output", ""), + dependencies=[t.replace("_", " ") for t in task_info.get("dependencies", [])], + ) + agent_tasks.append(task_obj) + + if "model_tools" in agent_entry: + agent_tools = parse_tools(agent_entry["model_tools"]) + else: + agent_tools = [] + + agent_obj = AgentFactory.create( + name=agent_name.replace("_", " "), + description=agent_description, + tools=agent_tools, + tasks=agent_tasks, + llm_id=agent_llm_id, + ) + agent_objs.append(agent_obj) + + return TeamAgentFactory.create( + name=team_config["team_agent_name"], + agents=agent_objs, + description=team_config["description"], + llm_id=team_config.get("llm_id", None), + inspectors=[], + use_mentalist=True, + ) + + +@pytest.fixture +def team_agent(): + return build_team_agent_from_json(team_dict) + + +def test_evolver_output(team_agent): + response = team_agent.evolve() + poll_url = response["url"] + result = team_agent.poll(poll_url) + + while result.status == ResponseStatus.IN_PROGRESS: + time.sleep(30) + result = team_agent.poll(poll_url) + + assert "system" in result["data"]["evolved_agent"]["name"].lower(), "System should be in the system name" + assert result["status"] == ResponseStatus.SUCCESS, "Final result should have a 'SUCCESS' status" + assert "evolved_agent" in result["data"], "Data should contain 'evolved_agent'" + assert "evaluation_report" in result["data"], "Data should contain 'evaluation_report'" + assert "criteria" in result["data"], "Data should contain 'criteria'" + assert "archive" in result["data"], "Data should contain 'archive'" + + +def test_evolver_with_custom_llm_id(team_agent): + """Test evolver functionality with custom LLM ID""" + custom_llm_id = "6646261c6eb563165658bbb1" # GPT-4o ID + + # Test with evolver_llm parameter + response = team_agent.evolve_async(evolver_llm=custom_llm_id) + + assert response is not None + assert "url" in response or response.get("url") is not None + assert response["status"] == ResponseStatus.IN_PROGRESS diff --git a/tests/functional/team_agent/team_agent_functional_test.py b/tests/functional/team_agent/team_agent_functional_test.py index b51044ea..e8352e1d 100644 --- a/tests/functional/team_agent/team_agent_functional_test.py +++ b/tests/functional/team_agent/team_agent_functional_test.py @@ -291,7 +291,7 @@ def test_team_agent_with_instructions(delete_agents_and_team_agents): instructions="Use only 'Agent 2' to solve the tasks.", llm_id="6646261c6eb563165658bbb1", use_mentalist=True, - use_inspector=False, + inspectors=[], ) response = team_agent.run(data="Translate 'cat' to Portuguese") @@ -408,7 +408,7 @@ class Response(BaseModel): description="Team agent", llm_id="6646261c6eb563165658bbb1", use_mentalist=False, - use_inspector=False, + inspectors=[], ) # Run the team agent @@ -484,7 +484,7 @@ def test_team_agent_with_slack_connector(): description="Team agent", llm_id="6646261c6eb563165658bbb1", use_mentalist=False, - use_inspector=False, + inspectors=[], ) response = team_agent.run( diff --git a/tests/unit/agent/evolve_param_test.py b/tests/unit/agent/evolve_param_test.py new file mode 100644 index 00000000..b108c146 --- /dev/null +++ b/tests/unit/agent/evolve_param_test.py @@ -0,0 +1,187 @@ +""" +Unit tests for EvolveParam base model functionality +""" + +import pytest +from aixplain.modules.agent.evolve_param import ( + EvolveParam, + EvolveType, + validate_evolve_param, +) + + +class TestEvolveParam: + """Test class for EvolveParam functionality""" + + def test_default_initialization(self): + """Test EvolveParam default initialization""" + default_param = EvolveParam() + + assert default_param is not None + assert default_param.to_evolve is False + assert default_param.evolve_type == EvolveType.TEAM_TUNING + assert default_param.max_generations == 3 + assert default_param.max_retries == 3 + assert default_param.recursion_limit == 50 + assert default_param.max_iterations_without_improvement == 2 + assert default_param.evolver_llm is None + assert default_param.additional_params == {} + + # Test to_dict method + result_dict = default_param.to_dict() + assert isinstance(result_dict, dict) + assert "toEvolve" in result_dict + + def test_custom_initialization(self): + """Test EvolveParam custom initialization""" + custom_param = EvolveParam( + to_evolve=True, + max_generations=5, + max_retries=2, + recursion_limit=30, + max_iterations_without_improvement=4, + evolve_type=EvolveType.TEAM_TUNING, + evolver_llm={"id": "test_llm_id", "name": "Test LLM"}, + additional_params={"customParam": "custom_value"}, + ) + + assert custom_param.to_evolve is True + assert custom_param.max_generations == 5 + assert custom_param.max_retries == 2 + assert custom_param.recursion_limit == 30 + assert custom_param.max_iterations_without_improvement == 4 + assert custom_param.evolve_type == EvolveType.TEAM_TUNING + assert custom_param.evolver_llm == {"id": "test_llm_id", "name": "Test LLM"} + assert custom_param.additional_params == {"customParam": "custom_value"} + + # Test to_dict method + result_dict = custom_param.to_dict() + assert result_dict["toEvolve"] is True + assert result_dict["max_generations"] == 5 + assert result_dict["max_retries"] == 2 + assert result_dict["recursion_limit"] == 30 + assert result_dict["max_iterations_without_improvement"] == 4 + assert result_dict["evolve_type"] == EvolveType.TEAM_TUNING + assert result_dict["evolver_llm"] == {"id": "test_llm_id", "name": "Test LLM"} + assert result_dict["customParam"] == "custom_value" + + def test_from_dict_with_api_format(self): + """Test EvolveParam from_dict() with API format""" + api_dict = { + "toEvolve": True, + "max_generations": 10, + "max_retries": 4, + "recursion_limit": 40, + "max_iterations_without_improvement": 5, + "evolve_type": EvolveType.TEAM_TUNING, + "evolver_llm": {"id": "api_llm_id", "name": "API LLM"}, + "customParam": "custom_value", + } + + from_dict_param = EvolveParam.from_dict(api_dict) + + assert from_dict_param.to_evolve is True + assert from_dict_param.max_generations == 10 + assert from_dict_param.max_retries == 4 + assert from_dict_param.recursion_limit == 40 + assert from_dict_param.max_iterations_without_improvement == 5 + assert from_dict_param.evolve_type == EvolveType.TEAM_TUNING + assert from_dict_param.evolver_llm == {"id": "api_llm_id", "name": "API LLM"} + + # Test round-trip conversion + result_dict = from_dict_param.to_dict() + assert result_dict["toEvolve"] is True + assert result_dict["max_generations"] == 10 + assert result_dict["max_retries"] == 4 + assert result_dict["recursion_limit"] == 40 + assert result_dict["max_iterations_without_improvement"] == 5 + + def test_validate_evolve_param_with_none(self): + """Test validate_evolve_param() with None input""" + validated_none = validate_evolve_param(None) + + assert validated_none is not None + assert isinstance(validated_none, EvolveParam) + assert validated_none.to_evolve is False + + result_dict = validated_none.to_dict() + assert "toEvolve" in result_dict + + def test_validate_evolve_param_with_dict(self): + """Test validate_evolve_param() with dictionary input""" + input_dict = {"toEvolve": True, "max_generations": 5} + validated_dict = validate_evolve_param(input_dict) + + assert isinstance(validated_dict, EvolveParam) + assert validated_dict.to_evolve is True + assert validated_dict.max_generations == 5 + + result_dict = validated_dict.to_dict() + assert result_dict["toEvolve"] is True + assert result_dict["max_generations"] == 5 + + def test_validate_evolve_param_with_instance(self): + """Test validate_evolve_param() with EvolveParam instance""" + custom_param = EvolveParam( + to_evolve=True, + max_generations=5, + max_retries=2, + recursion_limit=30, + max_iterations_without_improvement=4, + evolve_type=EvolveType.TEAM_TUNING, + evolver_llm={"id": "instance_llm_id"}, + additional_params={"customParam": "custom_value"}, + ) + + validated_instance = validate_evolve_param(custom_param) + + assert validated_instance is custom_param # Should return the same instance + assert validated_instance.to_evolve is True + assert validated_instance.max_generations == 5 + assert validated_instance.max_retries == 2 + + def test_invalid_max_generations_raises_error(self): + """Test that invalid max_generations raises ValueError""" + with pytest.raises(ValueError, match="max_generations must be positive"): + EvolveParam(max_generations=0) # max_generations <= 0 should fail + + def test_validate_evolve_param_missing_to_evolve_key(self): + """Test that missing toEvolve key raises ValueError""" + with pytest.raises(ValueError, match="evolve parameter must contain 'toEvolve' key"): + validate_evolve_param({"no_to_evolve": True}) # Missing toEvolve key + + def test_evolve_type_enum_values(self): + """Test that EvolveType enum values work correctly""" + param_team_tuning = EvolveParam(evolve_type=EvolveType.TEAM_TUNING) + + assert param_team_tuning.evolve_type == EvolveType.TEAM_TUNING + + # Test in to_dict conversion + dict_team_tuning = param_team_tuning.to_dict() + + assert "evolve_type" in dict_team_tuning + + def test_invalid_additional_params_type(self): + """Test that invalid additional_params type raises ValueError""" + with pytest.raises(ValueError, match="additional_params must be a dictionary"): + EvolveParam(additional_params="not a dict") + + def test_merge_with_dict(self): + """Test merging with a dictionary""" + base_param = EvolveParam(to_evolve=False, max_generations=3, additional_params={"base": "value"}) + merge_dict = { + "toEvolve": True, + "max_generations": 5, + "evolver_llm": {"id": "merged_llm_id"}, + "customParam": "custom_value", + } + + merged = base_param.merge(merge_dict) + + assert merged.to_evolve is True + assert merged.max_generations == 5 + assert merged.evolver_llm == {"id": "merged_llm_id"} + assert merged.additional_params == { + "base": "value", + "customParam": "custom_value", + } diff --git a/tests/unit/agent/test_agent_evolve.py b/tests/unit/agent/test_agent_evolve.py new file mode 100644 index 00000000..1aa6cd70 --- /dev/null +++ b/tests/unit/agent/test_agent_evolve.py @@ -0,0 +1,214 @@ +""" +Unit tests for Agent evolve functionality with evolver_llm parameter +""" + +import pytest +from unittest.mock import Mock, patch +from aixplain.modules.agent import Agent +from aixplain.modules.model.llm_model import LLM +from aixplain.modules.agent.evolve_param import EvolveParam +from aixplain.enums import EvolveType, Function, Supplier, ResponseStatus +from aixplain.modules.agent.agent_response import AgentResponse +from aixplain.modules.agent.agent_response_data import AgentResponseData + + +class TestAgentEvolve: + """Test class for Agent evolve functionality""" + + @pytest.fixture + def mock_agent(self): + """Create a mock Agent for testing""" + agent = Mock(spec=Agent) + agent.id = "test_agent_id" + agent.name = "Test Agent" + agent.api_key = "test_api_key" + return agent + + @pytest.fixture + def mock_llm(self): + """Create a mock LLM for testing""" + llm = Mock(spec=LLM) + llm.id = "test_llm_id" + llm.name = "Test LLM" + llm.description = "Test LLM Description" + llm.supplier = Supplier.OPENAI + llm.version = "1.0.0" + llm.function = Function.TEXT_GENERATION + llm.temperature = 0.7 + + # Mock get_parameters + mock_params = Mock() + mock_params.to_list.return_value = [{"name": "temperature", "type": "float"}] + llm.get_parameters.return_value = mock_params + + return llm + + def test_evolve_async_with_evolver_llm_string(self, mock_agent): + """Test evolve_async with evolver_llm as string ID""" + from aixplain.modules.agent import Agent + + # Create a real Agent instance but mock its methods + agent = Agent( + id="test_agent_id", + name="Test Agent", + description="Test Description", + instructions="Test Instructions", + tools=[], + llm_id="6646261c6eb563165658bbb1", + ) + + # Mock the run_async method + mock_response = AgentResponse( + status=ResponseStatus.IN_PROGRESS, + url="http://test-poll-url.com", + data=AgentResponseData(input="test input"), + run_time=0.0, + used_credits=0.0, + ) + + with patch.object(agent, "run_async", return_value=mock_response) as mock_run_async: + result = agent.evolve_async(evolver_llm="custom_llm_id_123") + + # Verify run_async was called with correct evolve parameter + mock_run_async.assert_called_once() + call_args = mock_run_async.call_args + + # Check that evolve parameter contains evolver_llm + evolve_param = call_args[1]["evolve"] + assert isinstance(evolve_param, EvolveParam) + assert evolve_param.evolver_llm == {"id": "custom_llm_id_123"} + + assert result == mock_response + + def test_evolve_async_with_evolver_llm_object(self, mock_agent, mock_llm): + """Test evolve_async with evolver_llm as LLM object""" + from aixplain.modules.agent import Agent + + # Create a real Agent instance but mock its methods + agent = Agent( + id="test_agent_id", + name="Test Agent", + description="Test Description", + instructions="Test Instructions", + tools=[], + llm_id="6646261c6eb563165658bbb1", + ) + + # Mock the run_async method + mock_response = AgentResponse( + status=ResponseStatus.IN_PROGRESS, + url="http://test-poll-url.com", + data=AgentResponseData(input="test input"), + run_time=0.0, + used_credits=0.0, + ) + + with patch.object(agent, "run_async", return_value=mock_response) as mock_run_async: + result = agent.evolve_async(evolver_llm=mock_llm) + + # Verify run_async was called with correct evolve parameter + mock_run_async.assert_called_once() + call_args = mock_run_async.call_args + + # Check that evolve parameter contains evolver_llm + evolve_param = call_args[1]["evolve"] + assert isinstance(evolve_param, EvolveParam) + + expected_llm_dict = { + "id": "test_llm_id", + "name": "Test LLM", + "description": "Test LLM Description", + "supplier": Supplier.OPENAI, + "version": "1.0.0", + "function": Function.TEXT_GENERATION, + "parameters": [{"name": "temperature", "type": "float"}], + "temperature": 0.7, + } + assert evolve_param.evolver_llm == expected_llm_dict + + assert result == mock_response + + def test_evolve_async_without_evolver_llm(self, mock_agent): + """Test evolve_async without evolver_llm parameter""" + from aixplain.modules.agent import Agent + + # Create a real Agent instance but mock its methods + agent = Agent( + id="test_agent_id", + name="Test Agent", + description="Test Description", + instructions="Test Instructions", + tools=[], + llm_id="6646261c6eb563165658bbb1", + ) + + # Mock the run_async method + mock_response = AgentResponse( + status=ResponseStatus.IN_PROGRESS, + url="http://test-poll-url.com", + data=AgentResponseData(input="test input"), + run_time=0.0, + used_credits=0.0, + ) + + with patch.object(agent, "run_async", return_value=mock_response) as mock_run_async: + result = agent.evolve_async() + + # Verify run_async was called with correct evolve parameter + mock_run_async.assert_called_once() + call_args = mock_run_async.call_args + + # Check that evolve parameter has evolver_llm as None + evolve_param = call_args[1]["evolve"] + assert isinstance(evolve_param, EvolveParam) + assert evolve_param.evolver_llm is None + + assert result == mock_response + + def test_evolve_with_custom_parameters(self, mock_agent): + """Test evolve with custom parameters including evolver_llm""" + from aixplain.modules.agent import Agent + + # Create a real Agent instance but mock its methods + agent = Agent( + id="test_agent_id", + name="Test Agent", + description="Test Description", + instructions="Test Instructions", + tools=[], + llm_id="6646261c6eb563165658bbb1", + ) + + with patch.object(agent, "evolve_async") as mock_evolve_async, patch.object(agent, "sync_poll") as mock_sync_poll: + + # Mock evolve_async response + mock_evolve_async.return_value = {"status": ResponseStatus.IN_PROGRESS, "url": "http://test-poll-url.com"} + + # Mock sync_poll response + mock_result = Mock() + mock_result.data = {"current_code": "test code", "evolved_agent": "evolved_agent_data"} + mock_sync_poll.return_value = mock_result + + result = agent.evolve( + evolve_type=EvolveType.TEAM_TUNING, + max_successful_generations=5, + max_failed_generation_retries=3, + recursion_limit=40, + max_non_improving_generations=3, + evolver_llm="custom_evolver_llm_id", + ) + + # Verify evolve_async was called with correct parameters + mock_evolve_async.assert_called_once_with( + evolve_type=EvolveType.TEAM_TUNING, + max_successful_generations=5, + max_failed_generation_retries=3, + recursion_limit=40, + max_non_improving_generations=3, + evolver_llm="custom_evolver_llm_id", + ) + + # Verify sync_poll was called + mock_sync_poll.assert_called_once_with("http://test-poll-url.com", name="evolve_process", timeout=600) + + assert result is not None diff --git a/tests/unit/agent/test_evolver_llm_utils.py b/tests/unit/agent/test_evolver_llm_utils.py new file mode 100644 index 00000000..dfcb4ec7 --- /dev/null +++ b/tests/unit/agent/test_evolver_llm_utils.py @@ -0,0 +1,129 @@ +""" +Unit tests for evolver LLM utility functions +""" + +from unittest.mock import Mock +from aixplain.utils.evolve_utils import create_evolver_llm_dict +from aixplain.modules.model.llm_model import LLM +from aixplain.enums import Function, Supplier + + +class TestCreateEvolverLLMDict: + """Test class for create_evolver_llm_dict functionality""" + + def test_create_evolver_llm_dict_with_none(self): + """Test create_evolver_llm_dict with None input""" + result = create_evolver_llm_dict(None) + assert result is None + + def test_create_evolver_llm_dict_with_string_id(self): + """Test create_evolver_llm_dict with LLM ID string""" + llm_id = "test_llm_id_123" + result = create_evolver_llm_dict(llm_id) + + expected = {"id": llm_id} + assert result == expected + + def test_create_evolver_llm_dict_with_llm_object(self): + """Test create_evolver_llm_dict with LLM object""" + # Create a mock LLM object + mock_llm = Mock(spec=LLM) + mock_llm.id = "llm_id_456" + mock_llm.name = "Test LLM Model" + mock_llm.description = "A test LLM model for unit testing" + mock_llm.supplier = Supplier.OPENAI + mock_llm.version = "1.0.0" + mock_llm.function = Function.TEXT_GENERATION + mock_llm.temperature = 0.7 + + # Mock the get_parameters method + mock_parameters = Mock() + mock_parameters.to_list.return_value = [ + {"name": "max_tokens", "type": "integer", "default": 2048}, + {"name": "temperature", "type": "float", "default": 0.7}, + ] + mock_llm.get_parameters.return_value = mock_parameters + + result = create_evolver_llm_dict(mock_llm) + + expected = { + "id": "llm_id_456", + "name": "Test LLM Model", + "description": "A test LLM model for unit testing", + "supplier": Supplier.OPENAI, + "version": "1.0.0", + "function": Function.TEXT_GENERATION, + "parameters": [ + {"name": "max_tokens", "type": "integer", "default": 2048}, + {"name": "temperature", "type": "float", "default": 0.7}, + ], + "temperature": 0.7, + } + assert result == expected + + def test_create_evolver_llm_dict_with_llm_object_no_parameters(self): + """Test create_evolver_llm_dict with LLM object that has no parameters""" + # Create a mock LLM object + mock_llm = Mock(spec=LLM) + mock_llm.id = "llm_id_789" + mock_llm.name = "Simple LLM" + mock_llm.description = "A simple LLM without parameters" + mock_llm.supplier = Supplier.OPENAI + mock_llm.version = "2.0.0" + mock_llm.function = Function.TEXT_GENERATION + mock_llm.temperature = 0.5 + + # Mock get_parameters to return None + mock_llm.get_parameters.return_value = None + + result = create_evolver_llm_dict(mock_llm) + + expected = { + "id": "llm_id_789", + "name": "Simple LLM", + "description": "A simple LLM without parameters", + "supplier": Supplier.OPENAI, + "version": "2.0.0", + "function": Function.TEXT_GENERATION, + "parameters": None, + "temperature": 0.5, + } + assert result == expected + + def test_create_evolver_llm_dict_with_llm_object_no_temperature(self): + """Test create_evolver_llm_dict with LLM object that has no temperature attribute""" + # Create a mock LLM object without temperature + mock_llm = Mock(spec=LLM) + mock_llm.id = "llm_id_999" + mock_llm.name = "No Temp LLM" + mock_llm.description = "LLM without temperature" + mock_llm.supplier = Supplier.GOOGLE + mock_llm.version = "3.0.0" + mock_llm.function = Function.TEXT_GENERATION + + # Remove temperature attribute + del mock_llm.temperature + + # Mock get_parameters to return None + mock_llm.get_parameters.return_value = None + + result = create_evolver_llm_dict(mock_llm) + + expected = { + "id": "llm_id_999", + "name": "No Temp LLM", + "description": "LLM without temperature", + "supplier": Supplier.GOOGLE, + "version": "3.0.0", + "function": Function.TEXT_GENERATION, + "parameters": None, + "temperature": None, + } + assert result == expected + + def test_create_evolver_llm_dict_with_empty_string(self): + """Test create_evolver_llm_dict with empty string""" + result = create_evolver_llm_dict("") + + expected = {"id": ""} + assert result == expected