Merge pull request #376 from cnoe-io/fix-aws

fix: make AWS agent to run in executor to prevent blocking

Author: suwhang-cisco

ai_platform_engineering/agents/aws/agent_aws/protocol_bindings/a2a_server/agent.py

@@ -15,43 +15,25 @@ class AWSAgent:
 
     def __init__(self):
         """Initialize the A2A AWS Agent."""
-        self._agent = None
+        logger.info("Initializing AWS Agent and MCP servers...")
+        # Initialize agent eagerly to download MCP packages at startup
+        self._agent = BaseAWSAgent()
+        logger.info("AWS Agent initialized successfully")
 
     async def _get_agent(self) -> BaseAWSAgent:
         """Get or create the agent instance."""
-        if self._agent is None:
-            self._agent = BaseAWSAgent()
         return self._agent
 
     async def stream(self, query: str, context_id: str = None) -> AsyncIterator[Dict[str, Any]]:
         """Stream response from the agent."""
         agent = await self._get_agent()
-        
+
         # Run the synchronous agent in an executor to avoid blocking
         loop = asyncio.get_event_loop()
         response = await loop.run_in_executor(None, agent.run_sync, query)
-        
-        # For now, we'll implement basic streaming by chunking the response
-        # The underlying agent doesn't support native streaming yet
-        words = response.split()
-        chunk_size = 10  # words per chunk
-        
-        for i in range(0, len(words), chunk_size):
-            chunk = " ".join(words[i:i+chunk_size])
-            if i + chunk_size < len(words):
-                chunk += " "
-            
-            # Send intermediate chunks
-            yield {
-                'content': chunk,
-                'is_task_complete': False,
-                'context_id': context_id
-            }
-            
-            # Small delay to simulate streaming
-            await asyncio.sleep(0.1)
-        
-        # Send final completion event
+
+        # Send final completion event with full response
+        # Don't send fake intermediate chunks - just send the complete response
         yield {
             'content': response,
             'is_task_complete': True,
@@ -60,7 +42,5 @@ async def stream(self, query: str, context_id: str = None) -> AsyncIterator[Dict
 
     def run_sync(self, query: str) -> str:
         """Run the agent synchronously."""
-        if self._agent is None:
-            self._agent = BaseAWSAgent()
         result = self._agent.chat(query)
         return result.get("answer", "No response generated")

ai_platform_engineering/agents/aws/agent_aws/protocol_bindings/a2a_server/agent_executor.py

@@ -13,70 +13,103 @@
     TaskStatus,
     TaskStatusUpdateEvent,
 )
-from a2a.utils import new_task, new_text_artifact
+from a2a.utils import new_agent_text_message, new_task, new_text_artifact
 
 
 logger = logging.getLogger(__name__)
 
 
 class AWSAgentExecutor(AgentExecutor):
     """A2A Agent Executor for AWS Agent."""
-    
+
     SUPPORTED_CONTENT_TYPES = ["text/plain"]
-    
+
     def __init__(self):
         """Initialize the AWS Agent Executor."""
         self.agent = AWSAgent()
         logger.info("AWS Agent Executor initialized")
-    
+
     @override
     async def execute(
         self,
         context: RequestContext,
         event_queue: EventQueue,
     ) -> None:
         """Execute the agent with the given context.
-        
+
         Args:
             context: Request context containing user input and task info
             event_queue: Event queue for publishing task updates
         """
         query = context.get_user_input()
         task = context.current_task
+        context_id = context.message.contextId if context.message else None
 
         if not context.message:
             raise Exception('No message provided')
 
         if not task:
             task = new_task(context.message)
             await event_queue.enqueue_event(task)
-        
+
         try:
-            # Use the run_sync method from the A2A wrapper
-            content = self.agent.run_sync(query)
-            
-            await event_queue.enqueue_event(
-                TaskArtifactUpdateEvent(
-                    append=False,
-                    contextId=task.contextId,
-                    taskId=task.id,
-                    lastChunk=True,
-                    artifact=new_text_artifact(
-                        name='current_result',
-                        description='Result of request to agent.',
-                        text=content
+            # Run agent and stream response
+            async for event in self.agent.stream(query, context_id):
+                if event['is_task_complete']:
+                    # Send artifact chunk that client can accumulate
+                    await event_queue.enqueue_event(
+                        TaskArtifactUpdateEvent(
+                            append=False,
+                            contextId=task.contextId,
+                            taskId=task.id,
+                            lastChunk=False,
+                            artifact=new_text_artifact(
+                                name='current_result',
+                                description='Result of request to agent.',
+                                text=event['content'],
+                            ),
+                        )
+                    )
+                    await event_queue.enqueue_event(
+                        TaskStatusUpdateEvent(
+                            status=TaskStatus(state=TaskState.completed),
+                            final=True,
+                            contextId=task.contextId,
+                            taskId=task.id,
+                        )
+                    )
+                elif event['require_user_input']:
+                    await event_queue.enqueue_event(
+                        TaskStatusUpdateEvent(
+                            status=TaskStatus(
+                                state=TaskState.input_required,
+                                message=new_agent_text_message(
+                                    event['content'],
+                                    task.contextId,
+                                    task.id,
+                                ),
+                            ),
+                            final=True,
+                            contextId=task.contextId,
+                            taskId=task.id,
+                        )
+                    )
+                else:
+                    await event_queue.enqueue_event(
+                        TaskStatusUpdateEvent(
+                            status=TaskStatus(
+                                state=TaskState.working,
+                                message=new_agent_text_message(
+                                    event['content'],
+                                    task.contextId,
+                                    task.id,
+                                ),
+                            ),
+                            final=False,
+                            contextId=task.contextId,
+                            taskId=task.id,
+                        )
                     )
-                )
-            )
-            await event_queue.enqueue_event(
-                TaskStatusUpdateEvent(
-                    status=TaskStatus(state=TaskState.completed),
-                    final=True,
-                    contextId=task.contextId,
-                    taskId=task.id,
-                )
-            )
-                
         except Exception as e:
             logger.error(f"Error executing agent: {e}")
             await event_queue.enqueue_event(
@@ -100,15 +133,15 @@ async def execute(
                     taskId=task.id,
                 )
             )
-    
+
     @override
     async def cancel(
         self,
         context: RequestContext,
         event_queue: EventQueue,
     ) -> None:
         """Cancel the current task execution.
-        
+
         Args:
             context: Request context
             event_queue: Event queue for publishing cancellation updates

docker-compose/docker-compose.rag-only.yaml

@@ -45,6 +45,7 @@ services:
     - ENABLE_WEBEX_AGENT=false
     - ENABLE_PETSTORE_AGENT=false
     - ENABLE_RAG=true
+    - RAG_AGENT_PORT=8099
     depends_on:
     - agent-rag-rag-only-p2p
     command: platform-engineer