From 461c8e1bbcc3918a5536eb977bcfd99bac400fd9 Mon Sep 17 00:00:00 2001
From: mykonos-ibiza <222371740+mykonos-ibiza@users.noreply.github.com>
Date: Sat, 9 Aug 2025 22:56:05 +0530
Subject: [PATCH] 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.
---
 backend/agent/api.py | 47 +++++++++++++++++++++++++++-----------------
 1 file changed, 29 insertions(+), 18 deletions(-)

diff --git a/backend/agent/api.py b/backend/agent/api.py
index fc7b27b8..cbbb6c7e 100644
--- a/backend/agent/api.py
+++ b/backend/agent/api.py
@@ -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: