diff --git a/aixplain/factories/agent_factory/__init__.py b/aixplain/factories/agent_factory/__init__.py index 0caf0814..a23eb529 100644 --- a/aixplain/factories/agent_factory/__init__.py +++ b/aixplain/factories/agent_factory/__init__.py @@ -98,10 +98,10 @@ def create( from aixplain.utils.llm_utils import get_llm_instance if llm is None and llm_id is not None: - llm = get_llm_instance(llm_id, api_key=api_key) + llm = get_llm_instance(llm_id, api_key=api_key, use_cache=True) elif llm is None: # Use default GPT-4o if no LLM specified - llm = get_llm_instance("669a63646eb56306647e1091", api_key=api_key) + llm = get_llm_instance("669a63646eb56306647e1091", api_key=api_key, use_cache=True) if output_format == OutputFormat.JSON: assert expected_output is not None and ( @@ -152,7 +152,7 @@ def create( } if llm is not None: - llm = get_llm_instance(llm, api_key=api_key) if isinstance(llm, str) else llm + llm = get_llm_instance(llm, api_key=api_key, use_cache=True) if isinstance(llm, str) else llm payload["tools"].append( { "type": "llm", @@ -519,11 +519,12 @@ def list(cls) -> Dict: raise Exception(error_msg) @classmethod - def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent: - """Retrieve an agent by its ID. + def get(cls, agent_id: Optional[Text] = None, name: Optional[Text] = None, api_key: Optional[Text] = None) -> Agent: + """Retrieve an agent by its ID or name. Args: - agent_id (Text): ID of the agent to retrieve. + agent_id (Optional[Text], optional): ID of the agent to retrieve. + name (Optional[Text], optional): Name of the agent to retrieve. api_key (Optional[Text], optional): API key for authentication. Defaults to None, using the configured TEAM_API_KEY. @@ -532,14 +533,23 @@ def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent: Raises: Exception: If the agent cannot be retrieved or doesn't exist. + ValueError: If neither agent_id nor name is provided, or if both are provided. """ from aixplain.factories.agent_factory.utils import build_agent - url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent_id}") + # Validate that exactly one parameter is provided + if not (agent_id or name) or (agent_id and name): + raise ValueError("Must provide exactly one of 'agent_id' or 'name'") + + # Construct URL based on parameter type + if agent_id: + url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent_id}") + else: # name is provided + url = urljoin(config.BACKEND_URL, f"sdk/agents/by-name/{name}") api_key = api_key if api_key is not None else config.TEAM_API_KEY headers = {"x-api-key": api_key, "Content-Type": "application/json"} - logging.info(f"Start service for GET Agent - {url} - {headers}") + logging.info(f"Start service for GET Agent - {url} - {headers}") r = _request_with_retry("get", url, headers=headers) resp = r.json() if 200 <= r.status_code < 300: diff --git a/aixplain/factories/agent_factory/utils.py b/aixplain/factories/agent_factory/utils.py index da73f21c..8b101b6f 100644 --- a/aixplain/factories/agent_factory/utils.py +++ b/aixplain/factories/agent_factory/utils.py @@ -140,7 +140,7 @@ def build_llm(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> LLM: for tool in payload["tools"]: if tool["type"] == "llm" and tool["description"] == "main": - llm = get_llm_instance(payload["llmId"], api_key=api_key) + llm = get_llm_instance(payload["llmId"], api_key=api_key, use_cache=True) # Set parameters from the tool if "parameters" in tool: # Apply all parameters directly to the LLM properties @@ -191,18 +191,34 @@ def build_agent(payload: Dict, tools: List[Tool] = None, api_key: Text = config. payload_tools = tools if payload_tools is None: payload_tools = [] - for tool in tools_dict: + # Use parallel tool building with ThreadPoolExecutor for better performance + from concurrent.futures import ThreadPoolExecutor, as_completed + + def build_tool_safe(tool_data): + """Build a single tool with error handling""" try: - payload_tools.append(build_tool(tool)) + return build_tool(tool_data) except (ValueError, AssertionError) as e: logging.warning(str(e)) - continue + return None except Exception: logging.warning( - f"Tool {tool['assetId']} is not available. Make sure it exists or you have access to it. " + f"Tool {tool_data['assetId']} is not available. Make sure it exists or you have access to it. " "If you think this is an error, please contact the administrators." ) - continue + return None + + # Build all tools in parallel (only if there are tools to build) + if len(tools_dict) > 0: + with ThreadPoolExecutor(max_workers=min(len(tools_dict), 10)) as executor: + # Submit all tool build tasks + future_to_tool = {executor.submit(build_tool_safe, tool): tool for tool in tools_dict} + + # Collect results as they complete + for future in as_completed(future_to_tool): + tool_result = future.result() + if tool_result is not None: + payload_tools.append(tool_result) llm = build_llm(payload, api_key) diff --git a/aixplain/factories/team_agent_factory/__init__.py b/aixplain/factories/team_agent_factory/__init__.py index ba8d1119..da1ea6f5 100644 --- a/aixplain/factories/team_agent_factory/__init__.py +++ b/aixplain/factories/team_agent_factory/__init__.py @@ -334,14 +334,15 @@ def list(cls) -> Dict: raise Exception(error_msg) @classmethod - def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> TeamAgent: - """Retrieve a team agent by its ID. + def get(cls, agent_id: Optional[Text] = None, name: Optional[Text] = None, api_key: Optional[Text] = None) -> TeamAgent: + """Retrieve a team agent by its ID or name. This method fetches a specific team agent from the platform using its - unique identifier. + unique identifier or name. Args: - agent_id (Text): Unique identifier of the team agent to retrieve. + agent_id (Optional[Text], optional): Unique identifier of the team agent to retrieve. + name (Optional[Text], optional): Name of the team agent to retrieve. api_key (Optional[Text], optional): API key for authentication. Defaults to None, using the configured TEAM_API_KEY. @@ -350,15 +351,25 @@ def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> TeamAgent: Raises: Exception: If: - - Team agent ID is invalid + - Team agent ID/name is invalid - Authentication fails - Service is unavailable - Other API errors occur + ValueError: If neither agent_id nor name is provided, or if both are provided. """ - url = urljoin(config.BACKEND_URL, f"sdk/agent-communities/{agent_id}") + # Validate that exactly one parameter is provided + if not (agent_id or name) or (agent_id and name): + raise ValueError("Must provide exactly one of 'agent_id' or 'name'") + + # Construct URL based on parameter type + if agent_id: + url = urljoin(config.BACKEND_URL, f"sdk/agent-communities/{agent_id}") + else: # name is provided + url = urljoin(config.BACKEND_URL, f"sdk/agent-communities/by-name/{name}") + api_key = api_key if api_key is not None else config.TEAM_API_KEY headers = {"x-api-key": api_key, "Content-Type": "application/json"} - logging.info(f"Start service for GET Team Agent - {url} - {headers}") + logging.info(f"Start service for GET Team Agent - {url} - {headers}") try: r = _request_with_retry("get", url, headers=headers) resp = r.json() diff --git a/aixplain/factories/team_agent_factory/utils.py b/aixplain/factories/team_agent_factory/utils.py index 773ad59b..edbef6d4 100644 --- a/aixplain/factories/team_agent_factory/utils.py +++ b/aixplain/factories/team_agent_factory/utils.py @@ -59,15 +59,31 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text = payload_agents = agents if payload_agents is None: payload_agents = [] - for i, agent in enumerate(agents_dict): + # Use parallel agent fetching with ThreadPoolExecutor for better performance + from concurrent.futures import ThreadPoolExecutor, as_completed + + def fetch_agent(agent_data): + """Fetch a single agent by ID with error handling""" try: - payload_agents.append(AgentFactory.get(agent["assetId"])) - except Exception: + return AgentFactory.get(agent_data["assetId"]) + except Exception as e: logging.warning( - f"Agent {agent['assetId']} not found. Make sure it exists or you have access to it. " - "If you think this is an error, please contact the administrators." + f"Agent {agent_data['assetId']} not found. Make sure it exists or you have access to it. " + "If you think this is an error, please contact the administrators. Error: {e}" ) - continue + return None + + # Fetch all agents in parallel (only if there are agents to fetch) + if len(agents_dict) > 0: + with ThreadPoolExecutor(max_workers=min(len(agents_dict), 10)) as executor: + # Submit all agent fetch tasks + future_to_agent = {executor.submit(fetch_agent, agent): agent for agent in agents_dict} + + # Collect results as they complete + for future in as_completed(future_to_agent): + agent_result = future.result() + if agent_result is not None: + payload_agents.append(agent_result) # Ensure custom classes are instantiated: for compatibility with backend return format inspectors = [] @@ -90,6 +106,15 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text = # Get LLMs from tools if present supervisor_llm = None mentalist_llm = None + + # Cache for models to avoid duplicate fetching of the same model ID + model_cache = {} + + def get_cached_model(model_id: str) -> any: + """Get model from cache or fetch if not cached""" + if model_id not in model_cache: + model_cache[model_id] = ModelFactory.get(model_id, api_key=api_key, use_cache=True) + return model_cache[model_id] # First check if we have direct LLM objects in the payload if "supervisor_llm" in payload: @@ -100,14 +125,8 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text = elif "tools" in payload: for tool in payload["tools"]: if tool["type"] == "llm": - try: - llm = ModelFactory.get(payload["llmId"], api_key=api_key) - except Exception: - logging.warning( - f"LLM {payload['llmId']} not found. Make sure it exists or you have access to it. " - "If you think this is an error, please contact the administrators." - ) - continue + # Use cached model fetching to avoid duplicate API calls + llm = get_cached_model(payload["llmId"]) # Set parameters from the tool if "parameters" in tool: # Apply all parameters directly to the LLM properties @@ -258,7 +277,7 @@ def build_team_agent_from_yaml(yaml_code: str, llm_id: str, api_key: str, team_i team_name = system_data.get("name", "") team_description = system_data.get("description", "") team_instructions = system_data.get("instructions", "") - llm = ModelFactory.get(llm_id) + llm = ModelFactory.get(llm_id, use_cache=True) # Create agent mapping by name for easier task assignment agents_mapping = {} agent_objs = [] diff --git a/aixplain/modules/agent/__init__.py b/aixplain/modules/agent/__init__.py index 68007914..e59ade7f 100644 --- a/aixplain/modules/agent/__init__.py +++ b/aixplain/modules/agent/__init__.py @@ -175,7 +175,7 @@ def _validate(self) -> None: re.match(r"^[a-zA-Z0-9 \-\(\)]*$", self.name) is not None ), "Agent Creation Error: Agent name contains invalid characters. Only alphanumeric characters, spaces, hyphens, and brackets are allowed." - llm = get_llm_instance(self.llm_id, api_key=self.api_key) + llm = get_llm_instance(self.llm_id, api_key=self.api_key, use_cache=True) assert llm.function == Function.TEXT_GENERATION, "Large Language Model must be a text generation model." diff --git a/aixplain/modules/agent/tool/model_tool.py b/aixplain/modules/agent/tool/model_tool.py index da455188..ab3ae1af 100644 --- a/aixplain/modules/agent/tool/model_tool.py +++ b/aixplain/modules/agent/tool/model_tool.py @@ -247,7 +247,7 @@ def _get_model(self, model_id: Text = None): from aixplain.factories.model_factory import ModelFactory model_id = model_id or self.model - return ModelFactory.get(model_id, api_key=self.api_key) + return ModelFactory.get(model_id, api_key=self.api_key, use_cache=True) def validate_parameters(self, received_parameters: Optional[List[Dict]] = None) -> Optional[List[Dict]]: """Validates and formats the parameters for the tool. diff --git a/aixplain/modules/team_agent/__init__.py b/aixplain/modules/team_agent/__init__.py index b5669ece..fcc5df7d 100644 --- a/aixplain/modules/team_agent/__init__.py +++ b/aixplain/modules/team_agent/__init__.py @@ -693,7 +693,7 @@ def _validate(self) -> None: ), "Team Agent Creation Error: Team name contains invalid characters. Only alphanumeric characters, spaces, hyphens, and brackets are allowed." try: - llm = get_llm_instance(self.llm_id) + llm = get_llm_instance(self.llm_id, use_cache=True) 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.") diff --git a/aixplain/utils/llm_utils.py b/aixplain/utils/llm_utils.py index d9897380..3860f225 100644 --- a/aixplain/utils/llm_utils.py +++ b/aixplain/utils/llm_utils.py @@ -6,12 +6,14 @@ def get_llm_instance( llm_id: Text, api_key: Optional[Text] = None, + use_cache: bool = True, ) -> LLM: """Get an LLM instance with specific configuration. Args: llm_id (Text): ID of the LLM model to use. api_key (Optional[Text], optional): API key to use. Defaults to None. + use_cache (bool, optional): Whether to use caching for model retrieval. Defaults to True. Returns: LLM: Configured LLM instance. @@ -20,7 +22,7 @@ def get_llm_instance( Exception: If the LLM model with the given ID is not found. """ try: - llm = ModelFactory.get(llm_id, api_key=api_key) + llm = ModelFactory.get(llm_id, api_key=api_key, use_cache=use_cache) return llm except Exception: raise Exception(f"Large Language Model with ID '{llm_id}' not found.") diff --git a/aixplain/v2/agent.py b/aixplain/v2/agent.py index a493f723..9f837b61 100644 --- a/aixplain/v2/agent.py +++ b/aixplain/v2/agent.py @@ -71,10 +71,10 @@ def list(cls, **kwargs: Unpack[BareListParams]) -> Page["Agent"]: return AgentFactory.list(**kwargs) @classmethod - def get(cls, id: str, **kwargs: Unpack[BareGetParams]) -> "Agent": + def get(cls, id: Optional[str] = None, name: Optional[str] = None, **kwargs: Unpack[BareGetParams]) -> "Agent": from aixplain.factories import AgentFactory - return AgentFactory.get(agent_id=id) + return AgentFactory.get(agent_id=id, name=name) @classmethod def create(cls, *args, **kwargs: Unpack[AgentCreateParams]) -> "Agent":