anu151105/agentic-browser
2
1#!/usr/bin/env python32"""3Agent Orchestrator module for the Enhanced AI Agentic Browser Agent Architecture.4 5This module serves as the central coordinator for all layers in the architecture,6managing task execution flow and communication between components.7"""8 9import asyncio10import json11import logging12import os13import time14import uuid15from typing import Dict, List, Optional, Any16 17from fastapi import WebSocket18 19from src.perception.multimodal_processor import MultimodalProcessor20from src.browser_control.browser_controller import BrowserController21from src.action_execution.action_executor import ActionExecutor, APIInteractionModule22from src.planning.task_planner import TaskPlanner23from src.memory.continuous_memory import ContinuousMemory24from src.user_interaction.hybrid_executor import HybridExecutor25from src.a2a_protocol.agent_communication import A2AProtocol26from src.security.ethical_guardian import EthicalGuardian27from src.monitoring.metrics_collector import MetricsCollector28 29# Configure logging30logging.basicConfig(31 level=logging.INFO,32 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",33)34logger = logging.getLogger(__name__)35 36class AgentOrchestrator:37 """38 Orchestrates all layers of the AI Agent Architecture.39 40 This class manages the lifecycle of tasks, coordinates between different 41 layers, and handles communication with external systems.42 """43 44 def __init__(self):45 """Initialize the Agent Orchestrator."""46 self.perception = None47 self.browser_control = None48 self.action_executor = None49 self.api_module = None50 self.task_planner = None51 self.memory = None52 self.hybrid_executor = None53 self.a2a_protocol = None54 self.ethical_guardian = None55 self.metrics = None56 57 # Task tracking58 self.tasks = {} # task_id -> task_info59 self.websockets = {} # task_id -> List[WebSocket]60 61 logger.info("Agent Orchestrator instance created")62 63 @classmethod64 async def initialize(cls):65 """66 Initialize all components of the architecture.67 68 Returns:69 AgentOrchestrator: Initialized orchestrator instance70 """71 instance = cls()72 73 # Initialize components in parallel for efficiency74 await asyncio.gather(75 instance._init_perception(),76 instance._init_browser_control(),77 instance._init_action_execution(),78 instance._init_planning(),79 instance._init_memory(),80 instance._init_user_interaction(),81 instance._init_a2a_protocol(),82 instance._init_security(),83 instance._init_monitoring(),84 )85 86 logger.info("All components initialized successfully")87 return instance88 89 async def _init_perception(self):90 """Initialize the Perception & Understanding Layer."""91 self.perception = MultimodalProcessor()92 await self.perception.initialize()93 logger.info("Perception & Understanding Layer initialized")94 95 async def _init_browser_control(self):96 """Initialize the Browser Control Layer."""97 self.browser_control = BrowserController()98 await self.browser_control.initialize()99 logger.info("Browser Control Layer initialized")100 101 async def _init_action_execution(self):102 """Initialize the Action Execution Layer."""103 self.api_module = APIInteractionModule()104 self.action_executor = ActionExecutor(self.browser_control, self.api_module)105 await self.action_executor.initialize()106 logger.info("Action Execution Layer initialized")107 108 async def _init_planning(self):109 """Initialize the Planning & Reasoning Layer."""110 self.task_planner = TaskPlanner()111 await self.task_planner.initialize()112 logger.info("Planning & Reasoning Layer initialized")113 114 async def _init_memory(self):115 """Initialize the Memory & Learning Layer."""116 self.memory = ContinuousMemory()117 await self.memory.initialize()118 logger.info("Memory & Learning Layer initialized")119 120 async def _init_user_interaction(self):121 """Initialize the User Interaction Layer."""122 self.hybrid_executor = HybridExecutor(self.action_executor)123 await self.hybrid_executor.initialize()124 logger.info("User Interaction Layer initialized")125 126 async def _init_a2a_protocol(self):127 """Initialize the Agent-to-Agent Protocol."""128 self.a2a_protocol = A2AProtocol()129 await self.a2a_protocol.initialize()130 logger.info("Agent-to-Agent Protocol initialized")131 132 async def _init_security(self):133 """Initialize the Security & Ethics components."""134 self.ethical_guardian = EthicalGuardian()135 await self.ethical_guardian.initialize()136 logger.info("Security & Ethics components initialized")137 138 async def _init_monitoring(self):139 """Initialize the Monitoring & Analytics components."""140 self.metrics = MetricsCollector()141 await self.metrics.initialize()142 logger.info("Monitoring & Analytics components initialized")143 144 async def create_task(self, task_config: Dict) -> str:145 """146 Create a new task with the given configuration.147 148 Args:149 task_config: Dictionary containing task configuration150 151 Returns:152 str: Task ID153 """154 task_id = str(uuid.uuid4())155 156 # Validate the task with ethical guardian157 is_valid, reason = await self.ethical_guardian.validate_task(task_config["task_description"])158 if not is_valid:159 self.tasks[task_id] = {160 "status": "rejected",161 "error": f"Task rejected due to ethical concerns: {reason}",162 "config": task_config,163 "result": None,164 "created_at": time.time()165 }166 return task_id167 168 # Create and store task169 self.tasks[task_id] = {170 "status": "created",171 "config": task_config,172 "result": None,173 "created_at": time.time()174 }175 176 logger.info(f"Created task {task_id}: {task_config['task_description']}")177 self.metrics.record_task_created()178 179 return task_id180 181 async def execute_task(self, task_id: str):182 """183 Execute a task with the given ID.184 185 Args:186 task_id: ID of the task to execute187 """188 if task_id not in self.tasks:189 logger.error(f"Task {task_id} not found")190 return191 192 task = self.tasks[task_id]193 task["status"] = "running"194 await self._notify_task_update(task_id, {"status": "running"})195 196 try:197 # Record start time for metrics198 start_time = time.time()199 200 # Retrieve task configuration201 task_config = task["config"]202 human_assisted = task_config.get("human_assisted", False)203 204 # 1. Planning phase205 await self._notify_task_update(task_id, {"status": "planning"})206 plan = await self.task_planner.decompose_task(207 task_config["task_description"],208 urls=task_config.get("urls", [])209 )210 211 # 2. Check similar past tasks in memory212 similar_tasks = await self.memory.retrieve_similar_tasks(task_config["task_description"])213 214 # 3. Execute plan steps215 results = []216 for step_idx, step in enumerate(plan["steps"]):217 # Update task status with current step218 step_status = f"Executing step {step_idx+1}/{len(plan['steps'])}: {step['description']}"219 await self._notify_task_update(task_id, {"status": step_status})220 221 # Process current web page if browser is initialized222 if self.browser_control.is_page_loaded():223 page_data = await self.browser_control.get_page_data()224 understanding = await self.perception.analyze_page(225 page_data["screenshot"], 226 page_data["dom_text"], 227 step["description"]228 )229 230 # Execute the step based on operation mode231 if human_assisted:232 step_result = await self.hybrid_executor.execute_task(233 step, 234 human_assist=True235 )236 else:237 step_result = await self.action_executor.execute_action(step)238 239 # Record results240 results.append({241 "step": step["description"],242 "result": step_result243 })244 245 # Check for cancellation after each step246 if task["status"] == "cancelling":247 task["status"] = "cancelled"248 await self._notify_task_update(task_id, {"status": "cancelled"})249 return250 251 # 4. Store experience in memory252 await self.memory.store_experience(253 task_config["task_description"],254 plan["steps"],255 {256 "success": True,257 "results": results258 }259 )260 261 # 5. Update task status to completed262 task["status"] = "completed"263 task["result"] = {264 "steps_executed": len(results),265 "results": results,266 "execution_time": time.time() - start_time267 }268 269 # Record metrics270 self.metrics.record_task_completed(time.time() - start_time)271 272 await self._notify_task_update(task_id, {273 "status": "completed",274 "result": task["result"]275 })276 277 logger.info(f"Task {task_id} completed successfully")278 279 except Exception as e:280 # Handle errors and update task status281 logger.exception(f"Error executing task {task_id}: {str(e)}")282 task["status"] = "failed"283 task["error"] = str(e)284 285 # Record metrics286 self.metrics.record_task_failed()287 288 await self._notify_task_update(task_id, {289 "status": "failed",290 "error": str(e)291 })292 293 async def cancel_task(self, task_id: str) -> bool:294 """295 Cancel a running task.296 297 Args:298 task_id: ID of the task to cancel299 300 Returns:301 bool: True if task was found and cancellation was initiated, False otherwise302 """303 if task_id not in self.tasks:304 return False305 306 task = self.tasks[task_id]307 if task["status"] == "running":308 task["status"] = "cancelling"309 await self._notify_task_update(task_id, {"status": "cancelling"})310 logger.info(f"Initiated cancellation of task {task_id}")311 return True312 313 return False314 315 async def get_task_status(self, task_id: str) -> Optional[Dict]:316 """317 Get the current status of a task.318 319 Args:320 task_id: ID of the task321 322 Returns:323 Optional[Dict]: Task status and results if found, None otherwise324 """325 if task_id not in self.tasks:326 return None327 328 task = self.tasks[task_id]329 return {330 "task_id": task_id,331 "status": task["status"],332 "result": task.get("result"),333 "error": task.get("error")334 }335 336 async def register_websocket(self, task_id: str, websocket: WebSocket):337 """338 Register a WebSocket connection for real-time task updates.339 340 Args:341 task_id: ID of the task to subscribe to342 websocket: WebSocket connection343 """344 if task_id not in self.websockets:345 self.websockets[task_id] = []346 347 self.websockets[task_id].append(websocket)348 349 # Send initial status350 status = await self.get_task_status(task_id)351 if status:352 await websocket.send_text(json.dumps(status))353 354 async def unregister_websocket(self, task_id: str, websocket: WebSocket):355 """356 Unregister a WebSocket connection.357 358 Args:359 task_id: ID of the task360 websocket: WebSocket connection to unregister361 """362 if task_id in self.websockets:363 try:364 self.websockets[task_id].remove(websocket)365 except ValueError:366 pass367 368 async def _notify_task_update(self, task_id: str, update: Dict):369 """370 Send task updates to all registered WebSocket connections.371 372 Args:373 task_id: ID of the task374 update: Update data to send375 """376 if task_id not in self.websockets:377 return378 379 full_update = {"task_id": task_id, **update}380 381 dead_sockets = []382 for websocket in self.websockets[task_id]:383 try:384 await websocket.send_text(json.dumps(full_update))385 except Exception:386 dead_sockets.append(websocket)387 388 # Clean up dead connections389 for dead in dead_sockets:390 try:391 self.websockets[task_id].remove(dead)392 except ValueError:393 pass394 395 async def shutdown(self):396 """Clean up resources when the application is shutting down."""397 # Cancel all running tasks398 for task_id, task in self.tasks.items():399 if task["status"] == "running":400 task["status"] = "cancelled"401 402 # Close all websocket connections403 for task_id, sockets in self.websockets.items():404 for socket in sockets:405 try:406 await socket.close()407 except Exception:408 pass409 410 # Shutdown all components411 await asyncio.gather(412 self.browser_control.shutdown(),413 self.action_executor.shutdown(),414 self.memory.shutdown(),415 self.metrics.shutdown(),416 )417 418 logger.info("Agent Orchestrator shut down successfully")419