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:
mykonos-ibiza 2025-08-09 22:56:05 +05:30
parent 9de8e0e093
commit 461c8e1bbc
1 changed files with 29 additions and 18 deletions

View File

@ -1328,13 +1328,10 @@ async def get_agents(
# Default to created_at # Default to created_at
query = query.order("created_at", desc=(sort_order == "desc")) query = query.order("created_at", desc=(sort_order == "desc"))
# Execute query to get total count first # Get paginated data and total count in one request
count_result = await query.execute()
total_count = count_result.count
# Now get the actual data with pagination
query = query.range(offset, offset + limit - 1) query = query.range(offset, offset + limit - 1)
agents_result = await query.execute() agents_result = await query.execute()
total_count = agents_result.count if agents_result.count is not None else 0
if not agents_result.data: if not agents_result.data:
logger.info(f"No agents found for user: {user_id}") logger.info(f"No agents found for user: {user_id}")
@ -1352,21 +1349,35 @@ async def get_agents(
agents_data = agents_result.data agents_data = agents_result.data
# First, fetch version data for all agents to ensure we have correct tool info # 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 = {} agent_version_map = {}
for agent in agents_data: version_ids = list({agent['current_version_id'] for agent in agents_data if agent.get('current_version_id')})
if agent.get('current_version_id'): if version_ids:
try: try:
version_service = await _get_version_service() 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( for row in (versions_result.data or []):
agent_id=agent['agent_id'], config = row.get('config') or {}
version_id=agent['current_version_id'], tools = config.get('tools') or {}
user_id=user_id version_dict = {
) 'version_id': row['version_id'],
version_dict = version_obj.to_dict() 'agent_id': row['agent_id'],
agent_version_map[agent['agent_id']] = version_dict 'version_number': row['version_number'],
except Exception as e: 'version_name': row['version_name'],
logger.warning(f"Failed to get version data for agent {agent['agent_id']}: {e}") '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 # Apply tool-based filters using version data
if has_mcp_tools is not None or has_agentpress_tools is not None or tools: if has_mcp_tools is not None or has_agentpress_tools is not None or tools: