Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 18 additions & 8 deletions aixplain/factories/agent_factory/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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.

Expand All @@ -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:
Expand Down
28 changes: 22 additions & 6 deletions aixplain/factories/agent_factory/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down
25 changes: 18 additions & 7 deletions aixplain/factories/team_agent_factory/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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()
Expand Down
49 changes: 34 additions & 15 deletions aixplain/factories/team_agent_factory/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Expand All @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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 = []
Expand Down
2 changes: 1 addition & 1 deletion aixplain/modules/agent/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."

Expand Down
2 changes: 1 addition & 1 deletion aixplain/modules/agent/tool/model_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion aixplain/modules/team_agent/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.")
Expand Down
4 changes: 3 additions & 1 deletion aixplain/utils/llm_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.")
4 changes: 2 additions & 2 deletions aixplain/v2/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down