agent_os_kernel 3.1.0__py3-none-any.whl

amb_core/memory_broker.py

@@ -0,0 +1,410 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""In-memory broker adapter for testing and simple use cases."""
+
+import asyncio
+import heapq
+import uuid
+from collections import defaultdict
+from typing import Dict, List, Optional, Set, Tuple
+
+from amb_core.broker import BrokerAdapter, MessageHandler
+from amb_core.models import Message, MessagePriority
+
+# Priority mapping for heapq (lower number = higher priority)
+PRIORITY_ORDER = {
+    MessagePriority.CRITICAL: 0,
+    MessagePriority.URGENT: 1,
+    MessagePriority.HIGH: 2,
+    MessagePriority.NORMAL: 3,
+    MessagePriority.LOW: 4,
+    MessagePriority.BACKGROUND: 5,
+}
+
+
+class InMemoryBroker(BrokerAdapter):
+    """
+    In-memory broker implementation for testing and simple use cases.
+    
+    This broker stores messages in memory and uses anyio for async handling.
+    It's suitable for testing, development, and single-process applications.
+    
+    Features:
+    - Backpressure: Automatically slows down producers when consumers are overwhelmed
+    - Priority lanes: CRITICAL messages jump ahead of BACKGROUND messages
+    """
+
+    def __init__(
+        self,
+        max_queue_size: int = 1000,
+        backpressure_threshold: float = 0.8,
+        backpressure_delay: float = 0.01,
+        use_priority_delivery: bool = True
+    ):
+        """
+        Initialize the in-memory broker.
+        
+        Args:
+            max_queue_size: Maximum messages per topic before backpressure kicks in
+            backpressure_threshold: Queue fill percentage (0.0-1.0) that triggers backpressure
+            backpressure_delay: Delay in seconds when backpressure is active
+            use_priority_delivery: If True, deliver messages in priority order (slower but respects priority lanes)
+        """
+        self._connected = False
+        self._subscriptions: Dict[str, Dict[str, MessageHandler]] = defaultdict(dict)
+
+        # Priority queue: List of tuples (priority_value, counter, message)
+        self._message_queues: Dict[str, List[Tuple[int, int, Message]]] = defaultdict(list)
+        self._message_counter = 0  # For stable sorting within same priority
+
+        self._response_queues: Dict[str, asyncio.Queue] = {}
+        self._request_message_ids: Set[str] = set()  # Track request message IDs to avoid self-capture
+        self._tasks: Set[asyncio.Task] = set()
+
+        # Backpressure configuration
+        self._max_queue_size = max_queue_size
+        self._backpressure_threshold = backpressure_threshold
+        self._backpressure_delay = backpressure_delay
+        self._use_priority_delivery = use_priority_delivery
+
+        # Statistics for monitoring
+        self._backpressure_events: Dict[str, int] = defaultdict(int)
+
+        # Background processing tasks for priority delivery
+        self._delivery_tasks: Dict[str, asyncio.Task] = {}
+
+    async def connect(self) -> None:
+        """Establish connection (no-op for in-memory broker)."""
+        self._connected = True
+
+    async def disconnect(self) -> None:
+        """Close connection and cancel all tasks."""
+        self._connected = False
+
+        # Cancel all delivery tasks
+        for task in self._delivery_tasks.values():
+            if not task.done():
+                task.cancel()
+
+        if self._delivery_tasks:
+            await asyncio.gather(*self._delivery_tasks.values(), return_exceptions=True)
+
+        # Cancel all running 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._delivery_tasks.clear()
+        self._subscriptions.clear()
+        self._message_queues.clear()
+        self._response_queues.clear()
+        self._request_message_ids.clear()
+        self._backpressure_events.clear()
+
+    def _check_backpressure(self, topic: str) -> bool:
+        """
+        Check if backpressure should be applied for a topic.
+        
+        Args:
+            topic: The topic to check
+            
+        Returns:
+            True if backpressure should be applied
+        """
+        queue_size = len(self._message_queues[topic])
+        threshold = int(self._max_queue_size * self._backpressure_threshold)
+        return queue_size >= threshold
+
+    async def _apply_backpressure(self, topic: str) -> None:
+        """
+        Apply backpressure delay to slow down the producer.
+        
+        Args:
+            topic: The topic experiencing backpressure
+        """
+        self._backpressure_events[topic] += 1
+        await asyncio.sleep(self._backpressure_delay)
+
+    async def _priority_delivery_worker(self, topic: str) -> None:
+        """
+        Background worker that delivers messages in priority order.
+        
+        Args:
+            topic: The topic to process
+        """
+        try:
+            while self._connected:
+                # Check if there are messages to deliver
+                if not self._message_queues[topic]:
+                    await asyncio.sleep(0.001)  # Small delay to avoid busy waiting
+                    continue
+
+                # Get highest priority message
+                _, _, message = heapq.heappop(self._message_queues[topic])
+
+                # Check if this is a response message
+                is_response = (
+                    message.correlation_id
+                    and message.correlation_id in self._response_queues
+                    and message.id not in self._request_message_ids
+                )
+
+                # Deliver to handlers (unless it's a response message)
+                if not is_response:
+                    handlers = self._subscriptions.get(topic, {})
+                    for handler in handlers.values():
+                        task = asyncio.create_task(handler(message))
+                        self._tasks.add(task)
+                        task.add_done_callback(self._tasks.discard)
+
+                # Handle response messages
+                if is_response:
+                    await self._response_queues[message.correlation_id].put(message)
+
+        except asyncio.CancelledError:
+            # Worker cancelled during shutdown
+            pass
+
+    async def publish(self, message: Message, wait_for_confirmation: bool = False) -> Optional[str]:
+        """
+        Publish a message to all subscribers of the topic.
+        
+        Implements backpressure: If the queue is near capacity, the publisher
+        is automatically slowed down to prevent overwhelming consumers.
+        
+        Implements priority lanes: CRITICAL messages are processed before
+        BACKGROUND messages, allowing urgent tasks to jump the queue.
+        
+        Args:
+            message: The message to publish
+            wait_for_confirmation: If True, wait for handlers to process
+        
+        Returns:
+            Message ID
+        """
+        if not self._connected:
+            raise ConnectionError("Broker not connected")
+
+        topic = message.topic
+
+        # Apply backpressure if queue is getting full (reactive streams flow control)
+        if self._check_backpressure(topic):
+            await self._apply_backpressure(topic)
+
+        # Check if queue is at max capacity
+        if len(self._message_queues[topic]) >= self._max_queue_size:
+            # Drop oldest BACKGROUND message, or raise error if no BACKGROUND messages
+            dropped = self._drop_background_message(topic)
+            if not dropped:
+                raise RuntimeError(
+                    f"Queue for topic '{topic}' is full ({self._max_queue_size} messages). "
+                    "Producer is overwhelmed. Consider increasing queue size or adding more consumers."
+                )
+
+        # Add message to priority queue
+        priority_value = PRIORITY_ORDER.get(message.priority, 3)
+        self._message_counter += 1
+        heapq.heappush(
+            self._message_queues[topic],
+            (priority_value, self._message_counter, message)
+        )
+
+        # Start priority delivery worker if not already running
+        if self._use_priority_delivery and topic not in self._delivery_tasks:
+            task = asyncio.create_task(self._priority_delivery_worker(topic))
+            self._delivery_tasks[topic] = task
+
+        # Check if this is a response message for request-response pattern
+        is_response = (
+            message.correlation_id
+            and message.correlation_id in self._response_queues
+            and message.id not in self._request_message_ids
+        )
+
+        # For non-priority mode or when waiting for confirmation, deliver immediately
+        if not self._use_priority_delivery:
+            # Deliver to subscribers (skip if this is a response message)
+            if not is_response:
+                handlers = self._subscriptions.get(topic, {})
+
+                if wait_for_confirmation:
+                    # Wait for all handlers to complete
+                    tasks = []
+                    for handler in handlers.values():
+                        tasks.append(handler(message))
+
+                    if tasks:
+                        await asyncio.gather(*tasks, return_exceptions=True)
+                else:
+                    # Fire and forget - schedule handlers without waiting
+                    for handler in handlers.values():
+                        task = asyncio.create_task(handler(message))
+                        self._tasks.add(task)
+                        task.add_done_callback(self._tasks.discard)
+
+            # Handle request-response pattern
+            # Capture response messages in the response queue
+            if is_response:
+                await self._response_queues[message.correlation_id].put(message)
+
+        return message.id
+
+    def _drop_background_message(self, topic: str) -> bool:
+        """
+        Drop the oldest BACKGROUND priority message from the queue.
+        
+        Args:
+            topic: The topic to drop from
+            
+        Returns:
+            True if a message was dropped, False if no BACKGROUND messages found
+        """
+        queue = self._message_queues[topic]
+        background_priority = PRIORITY_ORDER[MessagePriority.BACKGROUND]
+
+        # Find index of oldest BACKGROUND message (largest counter value for background priority)
+        background_idx = None
+        max_counter = -1
+        
+        for i, (priority, counter, _msg) in enumerate(queue):
+            if priority == background_priority and counter > max_counter:
+                background_idx = i
+                max_counter = counter
+        
+        if background_idx is not None:
+            # Remove the item and re-heapify
+            queue[background_idx] = queue[-1]
+            queue.pop()
+            if background_idx < len(queue):
+                heapq.heapify(queue)
+            return True
+
+        return False
+
+    async def subscribe(self, topic: str, handler: MessageHandler) -> str:
+        """
+        Subscribe to a topic.
+        
+        Args:
+            topic: Topic to subscribe to
+            handler: Message handler function
+        
+        Returns:
+            Subscription ID
+        """
+        if not self._connected:
+            raise ConnectionError("Broker not connected")
+
+        subscription_id = str(uuid.uuid4())
+        self._subscriptions[topic][subscription_id] = handler
+        return subscription_id
+
+    async def unsubscribe(self, subscription_id: str) -> None:
+        """
+        Unsubscribe from a topic.
+        
+        Args:
+            subscription_id: The subscription ID
+        """
+        for topic_handlers in self._subscriptions.values():
+            if subscription_id in topic_handlers:
+                del topic_handlers[subscription_id]
+                return
+
+    async def request(self, message: Message, timeout: float = 30.0) -> Message:
+        """
+        Send a request and wait for response.
+        
+        Args:
+            message: Request message
+            timeout: Timeout in seconds
+        
+        Returns:
+            Response message
+        
+        Raises:
+            TimeoutError: If timeout exceeded
+        """
+        if not self._connected:
+            raise ConnectionError("Broker not connected")
+
+        # Generate correlation ID if not present
+        if not message.correlation_id:
+            message.correlation_id = str(uuid.uuid4())
+
+        # Set up reply queue
+        reply_queue: asyncio.Queue = asyncio.Queue()
+        self._response_queues[message.correlation_id] = reply_queue
+
+        # Mark this message ID as a request to avoid self-capture
+        self._request_message_ids.add(message.id)
+
+        try:
+            # Publish the request
+            await self.publish(message, wait_for_confirmation=False)
+
+            # Wait for response
+            try:
+                response = await asyncio.wait_for(reply_queue.get(), timeout=timeout)
+                return response
+            except asyncio.TimeoutError:
+                raise TimeoutError(f"No response received within {timeout} seconds")
+
+        finally:
+            # Clean up
+            if message.correlation_id in self._response_queues:
+                del self._response_queues[message.correlation_id]
+            if message.id in self._request_message_ids:
+                self._request_message_ids.discard(message.id)
+
+    async def get_pending_messages(self, topic: str, limit: int = 10) -> List[Message]:
+        """
+        Get pending messages from a topic (in priority order).
+        
+        Args:
+            topic: Topic to get messages from
+            limit: Maximum number of messages
+        
+        Returns:
+            List of messages (highest priority first)
+        """
+        if not self._connected:
+            raise ConnectionError("Broker not connected")
+
+        queue = self._message_queues.get(topic, [])
+
+        # Sort by priority and return top N messages
+        sorted_messages = sorted(queue, key=lambda x: (x[0], x[1]))
+        messages = [item[2] for item in sorted_messages[:limit]]
+
+        return messages
+
+    def get_backpressure_stats(self, topic: Optional[str] = None) -> Dict[str, int]:
+        """
+        Get backpressure statistics.
+        
+        Args:
+            topic: Optional topic to get stats for. If None, returns all topics.
+            
+        Returns:
+            Dictionary mapping topics to backpressure event counts
+        """
+        if topic:
+            return {topic: self._backpressure_events.get(topic, 0)}
+        return dict(self._backpressure_events)
+
+    def get_queue_size(self, topic: str) -> int:
+        """
+        Get current queue size for a topic.
+        
+        Args:
+            topic: The topic to check
+            
+        Returns:
+            Number of messages in the queue
+        """
+        return len(self._message_queues.get(topic, []))

amb_core/models.py

@@ -0,0 +1,141 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Core message models for AMB."""
+
+from enum import Enum, IntEnum
+from typing import Any, Dict, Optional
+from datetime import datetime, timezone
+from pydantic import BaseModel, Field, ConfigDict, field_validator
+
+
+class MessagePriority(IntEnum):
+    """
+    Message priority levels.
+    
+    Higher values indicate higher priority. Messages with higher priority
+    are processed before lower priority messages when queued.
+    """
+    BACKGROUND = 0  # Lowest priority for background tasks
+    LOW = 1
+    NORMAL = 5
+    HIGH = 8
+    URGENT = 10
+    CRITICAL = 15  # For system-critical messages like fraud alerts
+
+
+class Priority:
+    """
+    Convenience class for accessing priority levels.
+    
+    Example:
+        message = Message(payload=data, priority=Priority.HIGH)
+    """
+    BACKGROUND = MessagePriority.BACKGROUND
+    LOW = MessagePriority.LOW
+    NORMAL = MessagePriority.NORMAL
+    HIGH = MessagePriority.HIGH
+    URGENT = MessagePriority.URGENT
+    CRITICAL = MessagePriority.CRITICAL
+
+
+class MessageStatus(str, Enum):
+    """Status of a message in its lifecycle."""
+    PENDING = "pending"
+    DELIVERED = "delivered"
+    ACKNOWLEDGED = "acknowledged"
+    FAILED = "failed"
+    EXPIRED = "expired"
+    DLQ = "dlq"  # Moved to dead letter queue
+
+
+class Message(BaseModel):
+    """
+    Core message model for the Agent Message Bus.
+    
+    This model represents a message that can be sent through the bus.
+    It includes metadata for routing, tracking, handling, and distributed tracing.
+    
+    New in v0.2.0:
+        - trace_id: For distributed tracing across agents
+        - ttl_seconds: Alias for ttl for clearer API
+        - is_expired: Property to check if message has expired
+    """
+    
+    id: str = Field(..., description="Unique message identifier")
+    topic: str = Field(..., description="Message topic/channel")
+    payload: Dict[str, Any] = Field(default_factory=dict, description="Message payload")
+    priority: MessagePriority = Field(default=MessagePriority.NORMAL, description="Message priority")
+    
+    # Metadata
+    sender: Optional[str] = Field(None, description="Sender identifier")
+    correlation_id: Optional[str] = Field(None, description="Correlation ID for request-response patterns")
+    reply_to: Optional[str] = Field(None, description="Topic to reply to")
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), description="Message timestamp")
+    
+    # TTL and expiration
+    ttl: Optional[int] = Field(None, description="Time to live in seconds", alias="ttl_seconds")
+    
+    # Distributed tracing (AMB-004)
+    trace_id: Optional[str] = Field(None, description="Distributed trace ID for tracking message flow")
+    span_id: Optional[str] = Field(None, description="Span ID within the trace")
+    parent_span_id: Optional[str] = Field(None, description="Parent span ID for nested operations")
+    
+    # Additional metadata
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
+    
+    model_config = ConfigDict(
+        json_encoders={
+            datetime: lambda v: v.isoformat()
+        },
+        populate_by_name=True,  # Allow both 'ttl' and 'ttl_seconds'
+    )
+    
+    @field_validator('priority', mode='before')
+    @classmethod
+    def validate_priority(cls, v):
+        """Accept both int and MessagePriority."""
+        if isinstance(v, int):
+            return MessagePriority(v)
+        if isinstance(v, str):
+            return MessagePriority[v.upper()]
+        return v
+    
+    @property
+    def ttl_seconds(self) -> Optional[int]:
+        """Alias for ttl for clearer API."""
+        return self.ttl
+    
+    @property
+    def is_expired(self) -> bool:
+        """
+        Check if the message has expired based on TTL.
+        
+        Returns:
+            True if the message has exceeded its TTL, False otherwise
+        """
+        if self.ttl is None:
+            return False
+        
+        now = datetime.now(timezone.utc)
+        age_seconds = (now - self.timestamp).total_seconds()
+        return age_seconds > self.ttl
+    
+    @property
+    def age_seconds(self) -> float:
+        """Get the age of the message in seconds."""
+        now = datetime.now(timezone.utc)
+        return (now - self.timestamp).total_seconds()
+    
+    @property
+    def remaining_ttl(self) -> Optional[float]:
+        """
+        Get remaining TTL in seconds.
+        
+        Returns:
+            Remaining TTL, 0 if expired, None if no TTL set
+        """
+        if self.ttl is None:
+            return None
+        
+        remaining = self.ttl - self.age_seconds
+        return max(0, remaining)