daita-agents 0.1.0__py3-none-any.whl

daita/core/relay.py

@@ -0,0 +1,695 @@
+"""
+Simplified Relay System for Daita Agents.
+
+Provides simple message passing between agents in a workflow.
+Focus on essential communication functionality without complex features.
+
+Updated to pass only result data between agents for cleaner interfaces.
+
+Example:
+    ```python
+    from daita.core.relay import RelayManager
+    
+    # Initialize relay manager
+    relay = RelayManager()
+    
+    # Agent publishes data (full response gets stored, only result gets forwarded)
+    await relay.publish("data_channel", {
+        "status": "success", 
+        "result": {"processed": True},
+        "agent_id": "agent_123"
+    })
+    
+    # Another agent subscribes and receives only: {"processed": True}
+    async def handle_data(result_data):
+        print(f"Got result: {result_data}")
+    
+    await relay.subscribe("data_channel", handle_data)
+    ```
+"""
+import asyncio
+import logging
+import random
+import time
+import uuid
+from typing import Dict, Any, Optional, List, Callable, Union
+from collections import deque
+from dataclasses import dataclass, field
+from enum import Enum
+import weakref
+
+try:
+    from ..core.exceptions import DaitaError, AcknowledgmentTimeoutError
+except ImportError:
+    # Fallback for direct execution or testing
+    from core.exceptions import DaitaError, AcknowledgmentTimeoutError
+
+logger = logging.getLogger(__name__)
+
+class RelayError(DaitaError):
+    """Exception raised for relay-related errors."""
+    pass
+
+class MessageStatus(str, Enum):
+    """Status of a message in the relay system."""
+    PENDING = "pending"
+    ACKNOWLEDGED = "acknowledged"
+    FAILED = "failed"
+    TIMEOUT = "timeout"
+
+@dataclass
+class ReliableMessage:
+    """A message with acknowledgment tracking."""
+    id: str
+    channel: str
+    data: Any
+    publisher: Optional[str] = None
+    timestamp: float = field(default_factory=time.time)
+    status: MessageStatus = MessageStatus.PENDING
+    ack_timeout: float = 30.0  # Default 30 second timeout
+    error: Optional[str] = None
+    attempts: int = 0
+    max_attempts: int = 3
+
+class RelayManager:
+    """
+    Simple relay manager for agent-to-agent communication.
+
+    Handles message publishing and subscription with minimal complexity.
+    Now passes only result data to maintain clean interfaces.
+    """
+
+    def __init__(
+        self,
+        max_messages_per_channel: int = 10,
+        enable_reliability: bool = False,
+        default_ack_timeout: float = 30.0
+    ):
+        """
+        Initialize the relay manager.
+
+        Args:
+            max_messages_per_channel: Maximum messages to keep per channel (reduced to 10 for memory efficiency)
+            enable_reliability: Enable message acknowledgments and reliability features
+            default_ack_timeout: Default acknowledgment timeout in seconds
+        """
+        self.max_messages_per_channel = max_messages_per_channel
+        self.enable_reliability = enable_reliability
+        self.default_ack_timeout = default_ack_timeout
+
+        # Channel storage: channel_name -> deque of result data only
+        self.channels: Dict[str, deque] = {}
+
+        # Subscribers: channel_name -> set of callbacks
+        self.subscribers: Dict[str, weakref.WeakSet] = {}
+
+        # Per-channel locks to prevent race conditions between publish and subscribe
+        self._channel_locks: Dict[str, asyncio.Lock] = {}
+
+        # Subscriber error tracking
+        self.subscriber_errors: deque = deque(maxlen=100)
+
+        # Reliability features (when enabled)
+        self.pending_messages: Dict[str, ReliableMessage] = {} if enable_reliability else {}
+        self.message_timeouts: Dict[str, asyncio.Task] = {} if enable_reliability else {}
+
+        # Running state
+        self._running = False
+
+        logger.debug(f"RelayManager initialized (reliability: {enable_reliability})")
+    
+    async def start(self) -> None:
+        """Start the relay manager."""
+        if self._running:
+            return
+
+        self._running = True
+        logger.info("RelayManager started")
+
+    async def stop(self) -> None:
+        """Stop the relay manager and cleanup."""
+        if not self._running:
+            return
+
+        self._running = False
+
+        # Clear all channels and subscribers
+        self.channels.clear()
+        self.subscribers.clear()
+
+        logger.info("RelayManager stopped")
+    
+    def _ensure_channel(self, channel: str) -> None:
+        """Ensure channel exists."""
+        if channel not in self.channels:
+            self.channels[channel] = deque(maxlen=self.max_messages_per_channel)
+            self.subscribers[channel] = weakref.WeakSet()
+            self._channel_locks[channel] = asyncio.Lock()
+    
+    async def publish(
+        self,
+        channel: str,
+        agent_response: Dict[str, Any],
+        publisher: Optional[str] = None,
+        require_ack: Optional[bool] = None
+    ) -> Optional[str]:
+        """
+        Publish data to a channel.
+
+        Args:
+            channel: Channel name
+            agent_response: Full agent response (we extract 'result' field)
+            publisher: Optional publisher identifier
+            require_ack: Whether this message requires acknowledgment (overrides global setting)
+
+        Returns:
+            Message ID if reliability is enabled, None otherwise
+        """
+        if not self._running:
+            await self.start()
+
+        # Extract just the result data
+        result_data = agent_response.get('result')
+
+        if result_data is None:
+            logger.warning(f"No 'result' field found in agent response for channel '{channel}'")
+            result_data = agent_response  # Fallback to full response if no result field
+
+        self._ensure_channel(channel)
+
+        # Determine if we need reliability
+        needs_reliability = require_ack if require_ack is not None else self.enable_reliability
+
+        if needs_reliability:
+            return await self._publish_reliable(channel, result_data, agent_response, publisher)
+        else:
+            return await self._publish_fire_and_forget(channel, result_data, agent_response, publisher)
+    
+    async def _publish_fire_and_forget(
+        self,
+        channel: str,
+        result_data: Any,
+        agent_response: Dict[str, Any],
+        publisher: Optional[str]
+    ) -> None:
+        """Publish message without reliability features (original behavior)."""
+        # Use per-channel lock to make publish atomic
+        async with self._channel_locks[channel]:
+            # Create message with result data only
+            message = {
+                'data': result_data,
+                'publisher': publisher,
+                'timestamp': time.time()
+            }
+
+            # Store result message
+            self.channels[channel].append(message)
+
+            # Notify subscribers with just the result data (while holding lock)
+            await self._notify_subscribers(channel, result_data)
+
+        logger.debug(f"Published result to channel '{channel}' from {publisher}")
+        return None
+    
+    async def _publish_reliable(
+        self,
+        channel: str,
+        result_data: Any,
+        agent_response: Dict[str, Any],
+        publisher: Optional[str]
+    ) -> str:
+        """Publish message with reliability features."""
+        # Create reliable message
+        message_id = uuid.uuid4().hex
+        reliable_message = ReliableMessage(
+            id=message_id,
+            channel=channel,
+            data=result_data,
+            publisher=publisher,
+            ack_timeout=self.default_ack_timeout
+        )
+
+        # Store pending message
+        self.pending_messages[message_id] = reliable_message
+
+        # Use per-channel lock to make publish atomic
+        async with self._channel_locks[channel]:
+            # Create message for channel storage
+            message = {
+                'id': message_id,
+                'data': result_data,
+                'publisher': publisher,
+                'timestamp': time.time(),
+                'requires_ack': True
+            }
+
+            # Store result message
+            self.channels[channel].append(message)
+
+            # Set up timeout task
+            timeout_task = asyncio.create_task(
+                self._handle_message_timeout(message_id, reliable_message.ack_timeout)
+            )
+            timeout_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
+            self.message_timeouts[message_id] = timeout_task
+
+            # Notify subscribers with message ID for acknowledgment (while holding lock)
+            await self._notify_subscribers_reliable(channel, result_data, message_id)
+
+        logger.debug(f"Published reliable message {message_id} to channel '{channel}' from {publisher}")
+        return message_id
+
+    async def _notify_subscribers(self, channel: str, result_data: Any) -> None:
+        """Notify all subscribers of a channel with result data only."""
+        if channel not in self.subscribers:
+            return
+        
+        # Get snapshot of subscribers to avoid modification during iteration
+        subscriber_list = list(self.subscribers[channel])
+        
+        if not subscriber_list:
+            return
+        
+        # Notify all subscribers concurrently
+        tasks = []
+        for subscriber in subscriber_list:
+            task = asyncio.create_task(self._call_subscriber(subscriber, result_data))
+            task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
+            tasks.append(task)
+        
+        # Wait for all notifications to complete
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+    
+    async def _notify_subscribers_reliable(self, channel: str, result_data: Any, message_id: str) -> None:
+        """Notify all subscribers of a reliable message."""
+        if channel not in self.subscribers:
+            return
+        
+        # Get snapshot of subscribers to avoid modification during iteration
+        subscriber_list = list(self.subscribers[channel])
+        
+        if not subscriber_list:
+            return
+        
+        # Notify all subscribers concurrently with message ID
+        tasks = []
+        for subscriber in subscriber_list:
+            task = asyncio.create_task(
+                self._call_subscriber_reliable(subscriber, result_data, message_id)
+            )
+            task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
+            tasks.append(task)
+        
+        # Wait for all notifications to complete
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+    
+    async def _call_subscriber_reliable(self, callback: Callable, result_data: Any, message_id: str) -> None:
+        """Safely call a subscriber callback for reliable message."""
+        try:
+            if asyncio.iscoroutinefunction(callback):
+                # Check if callback supports message_id parameter
+                import inspect
+                sig = inspect.signature(callback)
+                if 'message_id' in sig.parameters:
+                    await callback(result_data, message_id=message_id)
+                else:
+                    await callback(result_data)
+            else:
+                # Run sync callback in thread pool
+                loop = asyncio.get_event_loop()
+                await loop.run_in_executor(None, callback, result_data)
+        except Exception as e:
+            logger.error(f"Error in reliable subscriber callback: {str(e)}")
+            # NACK the message on callback error
+            await self.nack_message(message_id, str(e))
+    
+    async def _call_subscriber(self, callback: Callable, result_data: Any) -> None:
+        """Safely call a subscriber callback with result data."""
+        try:
+            if asyncio.iscoroutinefunction(callback):
+                await callback(result_data)
+            else:
+                # Run sync callback in thread pool
+                loop = asyncio.get_event_loop()
+                await loop.run_in_executor(None, callback, result_data)
+        except Exception as e:
+            error_info = {
+                'callback': str(callback),
+                'error': str(e),
+                'error_type': type(e).__name__,
+                'timestamp': time.time(),
+                'data_preview': str(result_data)[:100]
+            }
+            self.subscriber_errors.append(error_info)
+            logger.error(f"Subscriber callback failed: {e}")
+            # Don't re-raise - we want other subscribers to continue
+    
+    async def subscribe(self, channel: str, callback: Callable) -> None:
+        """
+        Subscribe to a channel.
+
+        Args:
+            channel: Channel name
+            callback: Callback function to receive result data
+        """
+        self._ensure_channel(channel)
+        self.subscribers[channel].add(callback)
+
+        logger.debug(f"Subscribed to channel '{channel}'")
+    
+    def unsubscribe(self, channel: str, callback: Callable) -> bool:
+        """
+        Unsubscribe from a channel.
+        
+        Args:
+            channel: Channel name
+            callback: Callback to remove
+            
+        Returns:
+            True if callback was removed
+        """
+        if channel not in self.subscribers:
+            return False
+        
+        try:
+            self.subscribers[channel].remove(callback)
+            logger.debug(f"Unsubscribed from channel '{channel}'")
+            return True
+        except KeyError:
+            return False
+    
+    async def get_latest(self, channel: str, count: int = 1) -> List[Any]:
+        """
+        Get latest result data from a channel.
+
+        Args:
+            channel: Channel name
+            count: Number of messages to retrieve
+
+        Returns:
+            List of result data (newest first)
+        """
+        if channel not in self.channels:
+            return []
+
+        # Get latest messages
+        messages = list(self.channels[channel])
+        latest = messages[-count:] if count < len(messages) else messages
+        latest.reverse()  # Newest first
+
+        # Return just the result data
+        return [msg['data'] for msg in latest]
+
+    def list_channels(self) -> List[str]:
+        """List all channels."""
+        return list(self.channels.keys())
+    
+    def clear_channel(self, channel: str) -> bool:
+        """
+        Clear all messages from a channel.
+        
+        Args:
+            channel: Channel name
+            
+        Returns:
+            True if channel was cleared
+        """
+        if channel not in self.channels:
+            return False
+
+        self.channels[channel].clear()
+
+        logger.debug(f"Cleared channel '{channel}'")
+        return True
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get relay manager statistics."""
+        total_messages = sum(len(ch) for ch in self.channels.values())
+        total_subscribers = sum(len(subs) for subs in self.subscribers.values())
+
+        stats = {
+            'running': self._running,
+            'total_channels': len(self.channels),
+            'total_messages': total_messages,
+            'total_subscribers': total_subscribers,
+            'channels': list(self.channels.keys()),
+            'subscriber_errors_count': len(self.subscriber_errors)
+        }
+
+        # Add per-channel error breakdown
+        if self.subscriber_errors:
+            stats['errors_by_channel'] = {}
+            for channel in self.channels.keys():
+                channel_errors = sum(
+                    1 for err in self.subscriber_errors
+                    if channel in err.get('callback', '')
+                )
+                if channel_errors > 0:
+                    stats['errors_by_channel'][channel] = channel_errors
+
+        # Add reliability stats if enabled
+        if self.enable_reliability:
+            stats.update({
+                'reliability_enabled': True,
+                'pending_messages': len(self.pending_messages),
+                'active_timeouts': len(self.message_timeouts)
+            })
+        else:
+            stats['reliability_enabled'] = False
+
+        return stats
+    
+    # Reliability methods (acknowledgments and timeouts)
+    
+    async def ack_message(self, message_id: str) -> bool:
+        """
+        Acknowledge a message as successfully processed.
+        
+        Args:
+            message_id: ID of the message to acknowledge
+            
+        Returns:
+            True if message was acknowledged, False if not found
+        """
+        if not self.enable_reliability or message_id not in self.pending_messages:
+            return False
+        
+        # Update message status
+        message = self.pending_messages[message_id]
+        message.status = MessageStatus.ACKNOWLEDGED
+        
+        # Cancel timeout task
+        if message_id in self.message_timeouts:
+            timeout_task = self.message_timeouts.pop(message_id)
+            if not timeout_task.done():
+                timeout_task.cancel()
+        
+        # Remove from pending
+        del self.pending_messages[message_id]
+        
+        logger.debug(f"Acknowledged message {message_id}")
+        return True
+    
+    async def nack_message(self, message_id: str, error: str) -> bool:
+        """
+        Negative acknowledge a message (processing failed).
+        
+        Args:
+            message_id: ID of the message to NACK
+            error: Error message
+            
+        Returns:
+            True if message was NACKed, False if not found
+        """
+        if not self.enable_reliability or message_id not in self.pending_messages:
+            return False
+        
+        # Update message status
+        message = self.pending_messages[message_id]
+        message.status = MessageStatus.FAILED
+        message.error = error
+        message.attempts += 1
+        
+        # Cancel timeout task
+        if message_id in self.message_timeouts:
+            timeout_task = self.message_timeouts.pop(message_id)
+            if not timeout_task.done():
+                timeout_task.cancel()
+        
+        # Check if we should retry
+        if message.attempts < message.max_attempts:
+            logger.warning(f"Message {message_id} failed (attempt {message.attempts}/{message.max_attempts}): {error}")
+            await self._schedule_retry(message)
+        else:
+            logger.error(f"Message {message_id} failed permanently after {message.attempts} attempts: {error}")
+            # Remove from pending messages
+            del self.pending_messages[message_id]
+        
+        return True
+    
+    async def _handle_message_timeout(self, message_id: str, timeout_duration: float) -> None:
+        """Handle message acknowledgment timeout."""
+        try:
+            await asyncio.sleep(timeout_duration)
+            
+            # Check if message is still pending
+            if message_id in self.pending_messages:
+                message = self.pending_messages[message_id]
+                message.status = MessageStatus.TIMEOUT
+                message.attempts += 1
+                
+                logger.warning(f"Message {message_id} timed out after {timeout_duration}s")
+                
+                # Check if we should retry
+                if message.attempts < message.max_attempts:
+                    await self._schedule_retry(message)
+                else:
+                    logger.error(f"Message {message_id} timed out permanently after {message.attempts} attempts")
+                    # Remove from pending messages
+                    del self.pending_messages[message_id]
+                
+                # Remove timeout task
+                self.message_timeouts.pop(message_id, None)
+        
+        except asyncio.CancelledError:
+            # Timeout was cancelled (message was acknowledged)
+            pass
+        except Exception as e:
+            logger.error(f"Error handling timeout for message {message_id}: {e}")
+    
+    def _calculate_retry_delay(self, attempt: int, base_delay: float = 1.0) -> float:
+        """Calculate retry delay with exponential backoff and jitter."""
+        # Exponential backoff: 1s, 2s, 4s, 8s, ...
+        delay = base_delay * (2 ** attempt)
+
+        # Cap at 60 seconds
+        delay = min(delay, 60.0)
+
+        # Add ±50% jitter to prevent thundering herd
+        jitter_factor = 0.5 + random.random()  # Random between 0.5 and 1.5
+        delay *= jitter_factor
+
+        return delay
+
+    async def _schedule_retry(self, message: ReliableMessage) -> None:
+        """Schedule a retry for a failed or timed out message."""
+        # Calculate exponential backoff delay with jitter
+        delay = self._calculate_retry_delay(message.attempts - 1)
+
+        logger.debug(f"Scheduling retry for message {message.id} in {delay:.1f}s (attempt {message.attempts + 1})")
+        
+        # Reset message status for retry
+        message.status = MessageStatus.PENDING
+        message.timestamp = time.time()  # Update timestamp for retry
+        
+        # Schedule the retry task
+        retry_task = asyncio.create_task(self._execute_retry(message, delay))
+        retry_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
+        # Use a different key to avoid conflicts with timeout tasks
+        self.message_timeouts[f"{message.id}_retry"] = retry_task
+    
+    async def _execute_retry(self, message: ReliableMessage, delay: float) -> None:
+        """Execute a message retry after the specified delay."""
+        try:
+            await asyncio.sleep(delay)
+            
+            # Check if message is still in pending (not cancelled)
+            if message.id in self.pending_messages:
+                logger.debug(f"Retrying message {message.id} (attempt {message.attempts + 1})")
+                
+                # Republish the message by re-adding it to the channel
+                self._ensure_channel(message.channel)
+                
+                # Create new message entry in channel
+                retry_message = {
+                    'id': message.id,
+                    'data': message.data,
+                    'publisher': message.publisher,
+                    'timestamp': message.timestamp,
+                    'requires_ack': True
+                }
+                
+                # Add to channel
+                self.channels[message.channel].append(retry_message)
+                
+                # Set up new timeout task for the retry
+                timeout_task = asyncio.create_task(
+                    self._handle_message_timeout(message.id, message.ack_timeout)
+                )
+                timeout_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
+                self.message_timeouts[message.id] = timeout_task
+                
+                # Notify subscribers again
+                await self._notify_subscribers_reliable(message.channel, message.data, message.id)
+            
+            # Clean up retry task reference
+            self.message_timeouts.pop(f"{message.id}_retry", None)
+        
+        except asyncio.CancelledError:
+            # Retry was cancelled
+            pass
+        except Exception as e:
+            logger.error(f"Error executing retry for message {message.id}: {e}")
+    
+    def get_pending_messages(self) -> List[Dict[str, Any]]:
+        """Get list of pending messages waiting for acknowledgment."""
+        if not self.enable_reliability:
+            return []
+        
+        return [
+            {
+                'id': msg.id,
+                'channel': msg.channel,
+                'publisher': msg.publisher,
+                'status': msg.status.value,
+                'timestamp': msg.timestamp,
+                'attempts': msg.attempts,
+                'max_attempts': msg.max_attempts,
+                'error': msg.error,
+                'age': time.time() - msg.timestamp
+            }
+            for msg in self.pending_messages.values()
+        ]
+
+    def get_subscriber_errors(self, limit: int = 10) -> List[Dict[str, Any]]:
+        """Get recent subscriber errors for debugging."""
+        errors = list(self.subscriber_errors)
+        return errors[-limit:] if limit < len(errors) else errors
+
+    # Context manager support
+    async def __aenter__(self) -> "RelayManager":
+        """Async context manager entry."""
+        await self.start()
+        return self
+    
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Async context manager exit."""
+        await self.stop()
+
+# Global relay manager instance
+_global_relay = None
+
+def get_global_relay() -> RelayManager:
+    """Get the global relay manager instance."""
+    global _global_relay
+    if _global_relay is None:
+        _global_relay = RelayManager()
+    return _global_relay
+
+# Convenience functions
+async def publish(channel: str, agent_response: Dict[str, Any], publisher: Optional[str] = None) -> None:
+    """Publish using the global relay manager."""
+    relay = get_global_relay()
+    await relay.publish(channel, agent_response, publisher)
+
+async def subscribe(channel: str, callback: Callable) -> None:
+    """Subscribe using the global relay manager."""
+    relay = get_global_relay()
+    await relay.subscribe(channel, callback)
+
+async def get_latest(channel: str, count: int = 1) -> List[Any]:
+    """Get latest result data using the global relay manager."""
+    relay = get_global_relay()
+    return await relay.get_latest(channel, count)