message_queue_client.py

"""
Message Queue Client Abstraction
Supports both Redis and RabbitMQ implementations
"""

import json
import uuid
import logging
from abc import ABC, abstractmethod
from typing import Dict, Any, Optional, Callable
from datetime import datetime
import redis
import pika
from pika.adapters.blocking_connection import BlockingChannel

logger = logging.getLogger(__name__)

class Message:
    """Standard message format for all queue communications"""

    def __init__(self, agent_name: str, session_id: str, payload: Dict[str, Any],
                 message_id: Optional[str] = None, retry_count: int = 0,
                 priority: str = "normal", ttl: int = 300):
        self.message_id = message_id or str(uuid.uuid4())
        self.timestamp = datetime.utcnow().isoformat()
        self.agent_name = agent_name
        self.session_id = session_id
        self.payload = payload
        self.metadata = {
            "retry_count": retry_count,
            "priority": priority,
            "ttl": ttl
        }

    def to_dict(self) -> Dict[str, Any]:
        return {
            "message_id": self.message_id,
            "timestamp": self.timestamp,
            "agent_name": self.agent_name,
            "session_id": self.session_id,
            "payload": self.payload,
            "metadata": self.metadata
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Message':
        msg = cls(
            agent_name=data["agent_name"],
            session_id=data["session_id"],
            payload=data["payload"],
            message_id=data.get("message_id"),
            retry_count=data["metadata"].get("retry_count", 0),
            priority=data["metadata"].get("priority", "normal"),
            ttl=data["metadata"].get("ttl", 300)
        )
        msg.timestamp = data.get("timestamp", datetime.utcnow().isoformat())
        return msg


class MessageQueueClient(ABC):
    """Abstract base class for message queue clients"""

    @abstractmethod
    def publish(self, topic: str, message: Message) -> None:
        """Publish a message to a topic/queue"""
        pass

    @abstractmethod
    def subscribe(self, topic: str, callback: Callable[[Message], None]) -> None:
        """Subscribe to a topic/queue with a callback"""
        pass

    @abstractmethod
    def close(self) -> None:
        """Close the connection"""
        pass


class RedisMessageQueue(MessageQueueClient):
    """Redis-based message queue implementation using pub/sub"""

    def __init__(self, host: str = "localhost", port: int = 6379, db: int = 0):
        self.redis_client = redis.Redis(host=host, port=port, db=db, decode_responses=True)
        self.subscriptions = {}
        self.pubsub = None

    def publish(self, topic: str, message: Message) -> None:
        """Publish message to Redis channel"""
        try:
            message_data = json.dumps(message.to_dict())
            self.redis_client.publish(topic, message_data)
            logger.info(f"Published message {message.message_id} to topic {topic}")
        except Exception as e:
            logger.error(f"Failed to publish message to {topic}: {e}")
            raise

    def subscribe(self, topic: str, callback: Callable[[Message], None]) -> None:
        """Subscribe to Redis channel"""
        if not self.pubsub:
            self.pubsub = self.redis_client.pubsub()

        self.subscriptions[topic] = callback
        self.pubsub.subscribe(**{topic: self._message_handler})
        logger.info(f"Subscribed to topic {topic}")

        # Start consuming messages in a separate thread if not already running
        if not hasattr(self, '_consumer_thread') or not self._consumer_thread.is_alive():
            self.start_consuming()

    def _message_handler(self, message):
        """Handle incoming Redis messages"""
        if message['type'] == 'message':
            try:
                data = json.loads(message['data'])
                msg = Message.from_dict(data)
                topic = message['channel']

                if topic in self.subscriptions:
                    # Get the callback function
                    callback = self.subscriptions[topic]
                    # Call the callback directly without asyncio
                    callback(msg)
            except Exception as e:
                logger.error(f"Error processing message: {e}")
                import traceback
                traceback.print_exc()

    def start_consuming(self):
        """Start consuming messages (run in separate thread)"""
        if self.pubsub:
            self._consumer_thread = self.pubsub.run_in_thread(sleep_time=0.001)

    def close(self) -> None:
        """Close Redis connection"""
        if self.pubsub:
            self.pubsub.close()
        self.redis_client.close()


class RabbitMQMessageQueue(MessageQueueClient):
    """RabbitMQ-based message queue implementation"""

    def __init__(self, host: str = "localhost", port: int = 5672,
                 username: str = "guest", password: str = "guest"):
        self.connection = pika.BlockingConnection(
            pika.ConnectionParameters(host=host, port=port,
                                    credentials=pika.PlainCredentials(username, password))
        )
        self.channel = self.connection.channel()
        self.consumers = []

    def publish(self, topic: str, message: Message) -> None:
        """Publish message to RabbitMQ exchange"""
        try:
            # Declare exchange (topic-based)
            self.channel.exchange_declare(exchange=topic, exchange_type='topic')

            message_data = json.dumps(message.to_dict())
            self.channel.basic_publish(
                exchange=topic,
                routing_key=f"{message.agent_name}.{message.session_id}",
                body=message_data.encode('utf-8'),
                properties=pika.BasicProperties(
                    delivery_mode=2,  # make message persistent
                    message_id=message.message_id,
                    timestamp=int(datetime.utcnow().timestamp())
                )
            )
            logger.info(f"Published message {message.message_id} to exchange {topic}")
        except Exception as e:
            logger.error(f"Failed to publish message to {topic}: {e}")
            raise

    def subscribe(self, topic: str, callback: Callable[[Message], None]) -> None:
        """Subscribe to RabbitMQ queue"""
        try:
            # Declare exchange and queue
            self.channel.exchange_declare(exchange=topic, exchange_type='topic')
            result = self.channel.queue_declare('', exclusive=True)
            queue_name = result.method.queue

            # Bind queue to exchange with routing key
            self.channel.queue_bind(exchange=topic, queue=queue_name,
                                  routing_key='#')  # Subscribe to all messages

            def rabbitmq_callback(ch: BlockingChannel, method, properties, body):
                try:
                    data = json.loads(body.decode('utf-8'))
                    msg = Message.from_dict(data)
                    callback(msg)
                    ch.basic_ack(delivery_tag=method.delivery_tag)
                except Exception as e:
                    logger.error(f"Error processing RabbitMQ message: {e}")
                    ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)

            self.channel.basic_consume(queue=queue_name, on_message_callback=rabbitmq_callback)
            self.consumers.append(queue_name)
            logger.info(f"Subscribed to topic {topic}")

        except Exception as e:
            logger.error(f"Failed to subscribe to {topic}: {e}")
            raise

    def start_consuming(self):
        """Start consuming messages (blocking)"""
        logger.info("Starting RabbitMQ message consumption")
        self.channel.start_consuming()

    def close(self) -> None:
        """Close RabbitMQ connection"""
        if self.channel and not self.channel.is_closed:
            self.channel.stop_consuming()
        if self.connection and not self.connection.is_closed:
            self.connection.close()


def create_message_queue_client(queue_type: str = "redis", **kwargs) -> MessageQueueClient:
    """Factory function to create message queue client"""
    if queue_type == "redis":
        return RedisMessageQueue(**kwargs)
    elif queue_type == "rabbitmq":
        return RabbitMQMessageQueue(**kwargs)
    else:
        raise ValueError(f"Unsupported queue type: {queue_type}")