agent-os-kernel 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

modules/amb/amb_core/adapters/aws_sqs_broker.py

@@ -0,0 +1,374 @@
+"""AWS SQS broker adapter for AMB.
+
+Amazon SQS provides fully managed message queues for microservices,
+distributed systems, and serverless applications.
+"""
+
+import asyncio
+import json
+import uuid
+from typing import Dict, List, Optional
+
+from amb_core.broker import BrokerAdapter, MessageHandler
+from amb_core.models import Message
+
+try:
+    import aioboto3
+    from botocore.exceptions import ClientError
+except ImportError:
+    raise ImportError(
+        "AWS SQS adapter requires 'aioboto3' package. "
+        "Install it with: pip install amb-core[aws]"
+    )
+
+
+class AWSSQSBroker(BrokerAdapter):
+    """
+    AWS SQS broker adapter.
+    
+    This adapter uses AWS SQS queues for reliable message delivery.
+    Supports both standard queues and FIFO queues for ordered processing.
+    
+    Example:
+        ```python
+        from amb_core.adapters import AWSSQSBroker
+        
+        broker = AWSSQSBroker(
+            region_name="us-east-1",
+            queue_url="https://sqs.us-east-1.amazonaws.com/123456789/my-queue"
+        )
+        await broker.connect()
+        
+        # Subscribe
+        async def handler(msg):
+            print(f"Received: {msg.payload}")
+        
+        await broker.subscribe("agent.tasks", handler)
+        
+        # Publish
+        await broker.publish(Message(topic="agent.tasks", payload={"task": "analyze"}))
+        ```
+    """
+
+    def __init__(
+        self,
+        region_name: str = "us-east-1",
+        queue_url: Optional[str] = None,
+        queue_name: str = "amb-messages",
+        use_fifo: bool = False,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None
+    ):
+        """
+        Initialize AWS SQS broker.
+        
+        Args:
+            region_name: AWS region
+            queue_url: Existing queue URL (optional)
+            queue_name: Queue name to create/use
+            use_fifo: Use FIFO queue for ordered delivery
+            aws_access_key_id: AWS access key (uses env vars if not provided)
+            aws_secret_access_key: AWS secret key (uses env vars if not provided)
+        """
+        self.region_name = region_name
+        self.queue_url = queue_url
+        self.queue_name = queue_name
+        self.use_fifo = use_fifo
+        self.aws_access_key_id = aws_access_key_id
+        self.aws_secret_access_key = aws_secret_access_key
+        
+        self._session = None
+        self._sqs = None
+        self._subscriptions: Dict[str, str] = {}  # subscription_id -> topic
+        self._handlers: Dict[str, MessageHandler] = {}
+        self._tasks: set = set()
+        self._running = False
+        self._topic_queues: Dict[str, str] = {}  # topic -> queue_url
+
+    async def connect(self) -> None:
+        """Connect to AWS SQS."""
+        self._session = aioboto3.Session(
+            aws_access_key_id=self.aws_access_key_id,
+            aws_secret_access_key=self.aws_secret_access_key,
+            region_name=self.region_name
+        )
+        
+        self._running = True
+        
+        # Get or create main queue
+        async with self._session.client('sqs') as sqs:
+            if not self.queue_url:
+                queue_name = f"{self.queue_name}.fifo" if self.use_fifo else self.queue_name
+                
+                try:
+                    # Try to get existing queue
+                    response = await sqs.get_queue_url(QueueName=queue_name)
+                    self.queue_url = response['QueueUrl']
+                except ClientError:
+                    # Create new queue
+                    attributes = {}
+                    if self.use_fifo:
+                        attributes['FifoQueue'] = 'true'
+                        attributes['ContentBasedDeduplication'] = 'true'
+                    
+                    response = await sqs.create_queue(
+                        QueueName=queue_name,
+                        Attributes=attributes
+                    )
+                    self.queue_url = response['QueueUrl']
+        
+        self._topic_queues['default'] = self.queue_url
+
+    async def disconnect(self) -> None:
+        """Disconnect from AWS SQS."""
+        self._running = False
+        
+        # Cancel all tasks
+        for task in self._tasks:
+            if not task.done():
+                task.cancel()
+        
+        if self._tasks:
+            await asyncio.gather(*self._tasks, return_exceptions=True)
+        
+        self._tasks.clear()
+        self._session = None
+
+    async def publish(self, message: Message, wait_for_confirmation: bool = False) -> Optional[str]:
+        """
+        Publish message to SQS queue.
+        
+        Args:
+            message: Message to publish
+            wait_for_confirmation: Wait for SQS acknowledgment (always true for SQS)
+        
+        Returns:
+            Message ID
+        """
+        if not self._session:
+            raise ConnectionError("Not connected to AWS SQS")
+        
+        async with self._session.client('sqs') as sqs:
+            # Serialize message
+            message_body = message.model_dump_json()
+            
+            # Build send parameters
+            params = {
+                'QueueUrl': self.queue_url,
+                'MessageBody': message_body,
+                'MessageAttributes': {
+                    'topic': {
+                        'DataType': 'String',
+                        'StringValue': message.topic
+                    },
+                    'source': {
+                        'DataType': 'String',
+                        'StringValue': message.source
+                    },
+                    'message_id': {
+                        'DataType': 'String',
+                        'StringValue': message.id
+                    }
+                }
+            }
+            
+            # FIFO queues require message group ID
+            if self.use_fifo:
+                params['MessageGroupId'] = message.topic.replace('/', '-')
+                params['MessageDeduplicationId'] = message.id
+            
+            # Send message
+            response = await sqs.send_message(**params)
+            
+            return response.get('MessageId', message.id)
+
+    async def subscribe(self, topic: str, handler: MessageHandler) -> str:
+        """
+        Subscribe to messages on a topic.
+        
+        Starts polling the SQS queue for messages matching the topic.
+        
+        Args:
+            topic: Topic to filter messages by
+            handler: Message handler
+        
+        Returns:
+            Subscription ID
+        """
+        subscription_id = str(uuid.uuid4())
+        
+        self._subscriptions[subscription_id] = topic
+        self._handlers[topic] = handler
+        
+        # Start polling task
+        task = asyncio.create_task(self._poll_task(subscription_id, topic, handler))
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+        
+        return subscription_id
+
+    async def _poll_task(self, subscription_id: str, topic: str, handler: MessageHandler):
+        """Poll SQS queue for messages."""
+        while self._running and subscription_id in self._subscriptions:
+            try:
+                async with self._session.client('sqs') as sqs:
+                    # Receive messages
+                    response = await sqs.receive_message(
+                        QueueUrl=self.queue_url,
+                        MaxNumberOfMessages=10,
+                        WaitTimeSeconds=20,  # Long polling
+                        MessageAttributeNames=['All']
+                    )
+                    
+                    messages = response.get('Messages', [])
+                    
+                    for sqs_message in messages:
+                        try:
+                            # Parse message
+                            body = sqs_message['Body']
+                            amb_message = Message.model_validate_json(body)
+                            
+                            # Check topic filter
+                            if amb_message.topic == topic or topic == "*":
+                                await handler(amb_message)
+                            
+                            # Delete message after processing
+                            await sqs.delete_message(
+                                QueueUrl=self.queue_url,
+                                ReceiptHandle=sqs_message['ReceiptHandle']
+                            )
+                            
+                        except Exception as e:
+                            print(f"Error handling SQS message: {e}")
+                            # Leave message in queue for retry
+                
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                print(f"Error polling SQS: {e}")
+                await asyncio.sleep(1.0)
+
+    async def unsubscribe(self, subscription_id: str) -> None:
+        """
+        Unsubscribe from messages.
+        
+        Args:
+            subscription_id: Subscription ID
+        """
+        if subscription_id in self._subscriptions:
+            topic = self._subscriptions[subscription_id]
+            del self._subscriptions[subscription_id]
+            
+            if topic in self._handlers:
+                del self._handlers[topic]
+
+    async def request(self, message: Message, timeout: float = 30.0) -> Message:
+        """
+        Send request and wait for response.
+        
+        Uses a temporary response queue for the request-response pattern.
+        
+        Args:
+            message: Request message
+            timeout: Timeout in seconds
+        
+        Returns:
+            Response message
+        """
+        if not self._session:
+            raise ConnectionError("Not connected to AWS SQS")
+        
+        # Generate correlation ID
+        if not message.correlation_id:
+            message.correlation_id = str(uuid.uuid4())
+        
+        # Create response queue
+        response_queue: asyncio.Queue = asyncio.Queue()
+        
+        async def response_handler(msg: Message):
+            if msg.correlation_id == message.correlation_id:
+                await response_queue.put(msg)
+        
+        # Subscribe to responses (using correlation ID as topic filter)
+        response_topic = f"response.{message.correlation_id}"
+        sub_id = await self.subscribe(response_topic, response_handler)
+        
+        try:
+            # Set reply-to
+            message.reply_to = response_topic
+            
+            # Publish request
+            await self.publish(message, wait_for_confirmation=False)
+            
+            # Wait for response
+            try:
+                response = await asyncio.wait_for(response_queue.get(), timeout=timeout)
+                return response
+            except asyncio.TimeoutError:
+                raise TimeoutError(f"No response received within {timeout} seconds")
+        
+        finally:
+            await self.unsubscribe(sub_id)
+
+    async def get_pending_messages(self, topic: str, limit: int = 10) -> List[Message]:
+        """
+        Peek at pending messages in the queue.
+        
+        Note: SQS doesn't support true peeking, so this uses visibility timeout
+        and then returns messages to the queue.
+        
+        Args:
+            topic: Topic to filter by
+            limit: Maximum messages to retrieve
+        
+        Returns:
+            List of messages
+        """
+        if not self._session:
+            raise ConnectionError("Not connected to AWS SQS")
+        
+        messages = []
+        
+        async with self._session.client('sqs') as sqs:
+            # Receive with short visibility timeout
+            response = await sqs.receive_message(
+                QueueUrl=self.queue_url,
+                MaxNumberOfMessages=min(limit, 10),
+                VisibilityTimeout=1,  # Short timeout so messages return to queue
+                MessageAttributeNames=['All']
+            )
+            
+            for sqs_message in response.get('Messages', []):
+                try:
+                    body = sqs_message['Body']
+                    amb_message = Message.model_validate_json(body)
+                    
+                    if amb_message.topic == topic or topic == "*":
+                        messages.append(amb_message)
+                except Exception:
+                    continue
+        
+        return messages
+
+    async def health_check(self) -> bool:
+        """Check if connected to AWS SQS."""
+        if not self._session:
+            return False
+        
+        try:
+            async with self._session.client('sqs') as sqs:
+                await sqs.get_queue_attributes(
+                    QueueUrl=self.queue_url,
+                    AttributeNames=['QueueArn']
+                )
+                return True
+        except Exception:
+            return False
+
+    async def purge_queue(self) -> None:
+        """Purge all messages from the queue."""
+        if not self._session:
+            raise ConnectionError("Not connected to AWS SQS")
+        
+        async with self._session.client('sqs') as sqs:
+            await sqs.purge_queue(QueueUrl=self.queue_url)

modules/amb/amb_core/adapters/azure_servicebus_broker.py

@@ -0,0 +1,338 @@
+"""Azure Service Bus broker adapter for AMB.
+
+Azure Service Bus provides enterprise-grade messaging for cloud applications
+with features like dead-letter queues, sessions, and transactions.
+"""
+
+import asyncio
+import json
+import uuid
+from typing import Dict, List, Optional
+
+from amb_core.broker import BrokerAdapter, MessageHandler
+from amb_core.models import Message
+
+try:
+    from azure.servicebus.aio import ServiceBusClient, ServiceBusSender, ServiceBusReceiver
+    from azure.servicebus import ServiceBusMessage
+except ImportError:
+    raise ImportError(
+        "Azure Service Bus adapter requires 'azure-servicebus' package. "
+        "Install it with: pip install amb-core[azure]"
+    )
+
+
+class AzureServiceBusBroker(BrokerAdapter):
+    """
+    Azure Service Bus broker adapter.
+    
+    This adapter uses Azure Service Bus topics and subscriptions for
+    enterprise-grade messaging with guaranteed delivery.
+    
+    Example:
+        ```python
+        from amb_core.adapters import AzureServiceBusBroker
+        
+        broker = AzureServiceBusBroker(
+            connection_string="Endpoint=sb://...",
+            topic_name="agent-messages"
+        )
+        await broker.connect()
+        
+        # Subscribe
+        async def handler(msg):
+            print(f"Received: {msg.payload}")
+        
+        sub_id = await broker.subscribe("agent.tasks", handler)
+        
+        # Publish
+        await broker.publish(Message(topic="agent.tasks", payload={"task": "analyze"}))
+        ```
+    """
+
+    def __init__(
+        self,
+        connection_string: str,
+        topic_name: str = "amb-messages",
+        subscription_name: str = "amb-subscription"
+    ):
+        """
+        Initialize Azure Service Bus broker.
+        
+        Args:
+            connection_string: Azure Service Bus connection string
+            topic_name: Service Bus topic name
+            subscription_name: Default subscription name
+        """
+        self.connection_string = connection_string
+        self.topic_name = topic_name
+        self.subscription_name = subscription_name
+        
+        self._client: Optional[ServiceBusClient] = None
+        self._sender: Optional[ServiceBusSender] = None
+        self._receivers: Dict[str, ServiceBusReceiver] = {}
+        self._subscriptions: Dict[str, str] = {}  # subscription_id -> topic
+        self._handlers: Dict[str, MessageHandler] = {}
+        self._tasks: set = set()
+        self._running = False
+
+    async def connect(self) -> None:
+        """Connect to Azure Service Bus."""
+        self._client = ServiceBusClient.from_connection_string(
+            conn_str=self.connection_string
+        )
+        
+        # Create sender for the topic
+        self._sender = self._client.get_topic_sender(topic_name=self.topic_name)
+        self._running = True
+
+    async def disconnect(self) -> None:
+        """Disconnect from Azure Service Bus."""
+        self._running = False
+        
+        # Cancel all tasks
+        for task in self._tasks:
+            if not task.done():
+                task.cancel()
+        
+        if self._tasks:
+            await asyncio.gather(*self._tasks, return_exceptions=True)
+        
+        self._tasks.clear()
+        
+        # Close receivers
+        for receiver in self._receivers.values():
+            await receiver.close()
+        
+        self._receivers.clear()
+        
+        # Close sender
+        if self._sender:
+            await self._sender.close()
+        
+        # Close client
+        if self._client:
+            await self._client.close()
+
+    async def publish(self, message: Message, wait_for_confirmation: bool = False) -> Optional[str]:
+        """
+        Publish message to Azure Service Bus topic.
+        
+        Args:
+            message: Message to publish
+            wait_for_confirmation: Wait for broker acknowledgment
+        
+        Returns:
+            Message ID
+        """
+        if not self._sender:
+            raise ConnectionError("Not connected to Azure Service Bus")
+        
+        # Create Service Bus message
+        sb_message = ServiceBusMessage(
+            body=message.model_dump_json(),
+            message_id=message.id,
+            correlation_id=message.correlation_id,
+            subject=message.topic,  # Use subject for filtering
+            application_properties={
+                "topic": message.topic,
+                "source": message.source,
+                "target": message.target or "",
+                "message_type": message.message_type.value if hasattr(message.message_type, 'value') else str(message.message_type)
+            }
+        )
+        
+        # Send message
+        await self._sender.send_messages(sb_message)
+        
+        return message.id
+
+    async def subscribe(self, topic: str, handler: MessageHandler) -> str:
+        """
+        Subscribe to messages on a topic.
+        
+        Creates a receiver for the subscription and starts listening.
+        
+        Args:
+            topic: Topic to filter messages by
+            handler: Message handler
+        
+        Returns:
+            Subscription ID
+        """
+        if not self._client:
+            raise ConnectionError("Not connected to Azure Service Bus")
+        
+        subscription_id = str(uuid.uuid4())
+        
+        # Create receiver for the subscription
+        receiver = self._client.get_subscription_receiver(
+            topic_name=self.topic_name,
+            subscription_name=self.subscription_name,
+            max_wait_time=5  # seconds
+        )
+        
+        self._receivers[subscription_id] = receiver
+        self._subscriptions[subscription_id] = topic
+        self._handlers[topic] = handler
+        
+        # Start receiver task
+        task = asyncio.create_task(self._receive_task(subscription_id, topic, handler))
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+        
+        return subscription_id
+
+    async def _receive_task(self, subscription_id: str, topic: str, handler: MessageHandler):
+        """Receive messages from Azure Service Bus."""
+        receiver = self._receivers.get(subscription_id)
+        if not receiver:
+            return
+        
+        while self._running:
+            try:
+                # Receive messages in batches
+                async with receiver:
+                    messages = await receiver.receive_messages(max_message_count=10, max_wait_time=5)
+                    
+                    for sb_message in messages:
+                        try:
+                            # Parse message
+                            body = str(sb_message)
+                            amb_message = Message.model_validate_json(body)
+                            
+                            # Check if message matches topic filter
+                            if amb_message.topic == topic or topic == "*":
+                                await handler(amb_message)
+                            
+                            # Complete the message
+                            await receiver.complete_message(sb_message)
+                            
+                        except Exception as e:
+                            print(f"Error handling message: {e}")
+                            # Dead-letter the message for investigation
+                            await receiver.dead_letter_message(
+                                sb_message,
+                                reason="ProcessingError",
+                                error_description=str(e)
+                            )
+                
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                print(f"Error in Service Bus receiver: {e}")
+                await asyncio.sleep(1.0)
+
+    async def unsubscribe(self, subscription_id: str) -> None:
+        """
+        Unsubscribe from messages.
+        
+        Args:
+            subscription_id: Subscription ID
+        """
+        if subscription_id in self._receivers:
+            receiver = self._receivers[subscription_id]
+            await receiver.close()
+            del self._receivers[subscription_id]
+        
+        if subscription_id in self._subscriptions:
+            topic = self._subscriptions[subscription_id]
+            del self._subscriptions[subscription_id]
+            if topic in self._handlers:
+                del self._handlers[topic]
+
+    async def request(self, message: Message, timeout: float = 30.0) -> Message:
+        """
+        Send request and wait for response.
+        
+        Uses correlation ID and reply-to for request-response pattern.
+        
+        Args:
+            message: Request message
+            timeout: Timeout in seconds
+        
+        Returns:
+            Response message
+        """
+        if not self._client or not self._sender:
+            raise ConnectionError("Not connected to Azure Service Bus")
+        
+        # Generate correlation ID
+        if not message.correlation_id:
+            message.correlation_id = str(uuid.uuid4())
+        
+        # Create response queue
+        response_queue: asyncio.Queue = asyncio.Queue()
+        
+        # Create temporary receiver for response
+        response_subscription = f"response-{message.correlation_id[:8]}"
+        
+        async def response_handler(msg: Message):
+            if msg.correlation_id == message.correlation_id:
+                await response_queue.put(msg)
+        
+        # Subscribe to responses
+        sub_id = await self.subscribe(response_subscription, response_handler)
+        
+        try:
+            # Set reply-to
+            message.reply_to = response_subscription
+            
+            # Publish request
+            await self.publish(message, wait_for_confirmation=False)
+            
+            # Wait for response
+            try:
+                response = await asyncio.wait_for(response_queue.get(), timeout=timeout)
+                return response
+            except asyncio.TimeoutError:
+                raise TimeoutError(f"No response received within {timeout} seconds")
+        
+        finally:
+            await self.unsubscribe(sub_id)
+
+    async def get_pending_messages(self, topic: str, limit: int = 10) -> List[Message]:
+        """
+        Peek at pending messages without consuming them.
+        
+        Args:
+            topic: Topic to filter by
+            limit: Maximum messages to retrieve
+        
+        Returns:
+            List of messages
+        """
+        if not self._client:
+            raise ConnectionError("Not connected to Azure Service Bus")
+        
+        messages = []
+        
+        # Create temporary receiver
+        receiver = self._client.get_subscription_receiver(
+            topic_name=self.topic_name,
+            subscription_name=self.subscription_name
+        )
+        
+        try:
+            async with receiver:
+                # Peek messages without completing them
+                peeked = await receiver.peek_messages(max_message_count=limit)
+                
+                for sb_message in peeked:
+                    try:
+                        body = str(sb_message)
+                        amb_message = Message.model_validate_json(body)
+                        
+                        if amb_message.topic == topic or topic == "*":
+                            messages.append(amb_message)
+                    except Exception:
+                        continue
+        
+        finally:
+            await receiver.close()
+        
+        return messages
+
+    async def health_check(self) -> bool:
+        """Check if connected to Azure Service Bus."""
+        return self._client is not None and self._sender is not None