mirror of https://github.com/kortix-ai/suna.git
refactor: optimize agent retrieval by batching version data queries
- Combined total count retrieval and paginated data fetching into a single request for efficiency. - Replaced per-agent version fetching with a batched query to reduce API calls and improve performance. - Enhanced error handling for version data loading.
This commit is contained in:
parent
9de8e0e093
commit
461c8e1bbc
|
@ -1328,13 +1328,10 @@ async def get_agents(
|
|||
# Default to created_at
|
||||
query = query.order("created_at", desc=(sort_order == "desc"))
|
||||
|
||||
# Execute query to get total count first
|
||||
count_result = await query.execute()
|
||||
total_count = count_result.count
|
||||
|
||||
# Now get the actual data with pagination
|
||||
# Get paginated data and total count in one request
|
||||
query = query.range(offset, offset + limit - 1)
|
||||
agents_result = await query.execute()
|
||||
total_count = agents_result.count if agents_result.count is not None else 0
|
||||
|
||||
if not agents_result.data:
|
||||
logger.info(f"No agents found for user: {user_id}")
|
||||
|
@ -1352,21 +1349,35 @@ async def get_agents(
|
|||
agents_data = agents_result.data
|
||||
|
||||
# First, fetch version data for all agents to ensure we have correct tool info
|
||||
# Do this in a single batched query instead of per-agent service calls
|
||||
agent_version_map = {}
|
||||
for agent in agents_data:
|
||||
if agent.get('current_version_id'):
|
||||
try:
|
||||
version_service = await _get_version_service()
|
||||
version_ids = list({agent['current_version_id'] for agent in agents_data if agent.get('current_version_id')})
|
||||
if version_ids:
|
||||
try:
|
||||
versions_result = await client.table('agent_versions').select(
|
||||
'version_id, agent_id, version_number, version_name, is_active, created_at, updated_at, created_by, config'
|
||||
).in_('version_id', version_ids).execute()
|
||||
|
||||
version_obj = await version_service.get_version(
|
||||
agent_id=agent['agent_id'],
|
||||
version_id=agent['current_version_id'],
|
||||
user_id=user_id
|
||||
)
|
||||
version_dict = version_obj.to_dict()
|
||||
agent_version_map[agent['agent_id']] = version_dict
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to get version data for agent {agent['agent_id']}: {e}")
|
||||
for row in (versions_result.data or []):
|
||||
config = row.get('config') or {}
|
||||
tools = config.get('tools') or {}
|
||||
version_dict = {
|
||||
'version_id': row['version_id'],
|
||||
'agent_id': row['agent_id'],
|
||||
'version_number': row['version_number'],
|
||||
'version_name': row['version_name'],
|
||||
'system_prompt': config.get('system_prompt', ''),
|
||||
'configured_mcps': tools.get('mcp', []),
|
||||
'custom_mcps': tools.get('custom_mcp', []),
|
||||
'agentpress_tools': tools.get('agentpress', {}),
|
||||
'is_active': row.get('is_active', False),
|
||||
'created_at': row.get('created_at'),
|
||||
'updated_at': row.get('updated_at') or row.get('created_at'),
|
||||
'created_by': row.get('created_by'),
|
||||
}
|
||||
agent_version_map[row['agent_id']] = version_dict
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to batch load versions for agents: {e}")
|
||||
|
||||
# Apply tool-based filters using version data
|
||||
if has_mcp_tools is not None or has_agentpress_tools is not None or tools:
|
||||
|
|
Loading…
Reference in New Issue