| import aio_pika |
| from typing import Optional |
| from api.chat.chat_api import ChatAPI |
| from .base import ChatQueueBase |
|
|
| class RabbitMQQueue(ChatQueueBase): |
| def __init__(self, url: str): |
| self.url = url |
| self.connection = None |
| self.channel = None |
| self.queue_name = "chat_queue" |
| |
| async def connect(self): |
| if not self.connection: |
| self.connection = await aio_pika.connect_robust(self.url) |
| self.channel = await self.connection.channel() |
| await self.channel.declare_queue(self.queue_name) |
| |
| async def add(self, api_key: str) -> None: |
| await self.connect() |
| message = aio_pika.Message(body=api_key.encode()) |
| await self.channel.default_exchange.publish( |
| message, |
| routing_key=self.queue_name |
| ) |
| |
| async def get(self) -> Optional[ChatAPI]: |
| await self.connect() |
| message = await self.channel.get(self.queue_name, no_ack=True) |
| if message: |
| api_key = message.body.decode() |
| chat = ChatAPI(api_key=api_key) |
| |
| await self.add(api_key) |
| return chat |
| return None |
| |
| async def length(self) -> int: |
| await self.connect() |
| queue = await self.channel.declare_queue(self.queue_name) |
| return queue.declaration_result.message_count |
