[feature]Support RabbitMQ.

cli.py

@@ -0,0 +1,10 @@
+from consumer.asr import TranscribeConsumer
+
+
+def cli():
+    transcribe_consumer = TranscribeConsumer()
+    transcribe_consumer.consume_messages(transcribe_consumer.queue_name, transcribe_consumer.consume)
+
+
+if __name__ == '__main__':
+    cli()

consumer/asr.py

@@ -0,0 +1,104 @@
+import base64
+import io
+import logging
+import os
+import uuid
+from pathlib import Path
+from typing import Literal
+
+from faster_whisper import WhisperModel
+from pydantic import BaseModel, Field, ValidationError, model_validator
+
+from .base import BasicMessageReceiver, BasicMessageSender, Headers, Priority
+
+
+class TranscribeInputMessage(BaseModel):
+    uuid: str = Field(title='Request Unique Id.')
+    audio_file: str
+    language: Literal['en', 'zh',]
+    using_file_content: bool
+
+    @model_validator(mode='after')
+    def check_audio_file(self):
+        if self.using_file_content:
+            return self
+
+        if not Path(self.audio_file).exists():
+            raise FileNotFoundError(f'Audio file not exists.')
+        return self
+
+
+class TranscribeOutputMessage(BaseModel):
+    uuid: str
+    if_success: bool
+    msg: str
+    transcribed_text: str = Field(default='')
+
+
+class TranscribeConsumer(BasicMessageReceiver):
+
+    def __init__(self):
+        super().__init__()
+
+        self.exchange_name = 'transcribe'
+        self.queue_name = 'transcribe-input'
+        self.routing_key = 'transcribe-input'
+
+        self.setup_consume_parameters()
+        self.setup_message_sender()
+
+        model_size = os.environ.get('WHISPER-MODEL-SIZE', 'large-v3')
+        # Run on GPU with FP16
+        self.asr_model = WhisperModel(model_size, device='cuda', compute_type='float16')
+
+    def setup_consume_parameters(self):
+        self.declare_exchange(self.exchange_name)
+        self.declare_queue(self.queue_name, max_priority=-1)
+        self.bind_queue(self.exchange_name, self.queue_name, self.routing_key)
+
+    def setup_message_sender(self):
+        self.sender = BasicMessageSender()
+
+    def send_message(self, message: dict):
+        routing_key = 'transcribe-output'
+        headers = Headers(job_id=f'{uuid.uuid4()}', priority=Priority.NORMAL)
+        self.sender.send_message(
+            exchange_name=self.exchange_name,
+            routing_key=routing_key,
+            body=message,
+            headers=None
+        )
+
+    def send_success_message(self, uuid: str, transcribed_text):
+        message = TranscribeOutputMessage(uuid=uuid, if_success=True, msg='Transcribe finished.',
+                                          transcribed_text=transcribed_text)
+        self.send_message(message.model_dump())
+
+    def send_fail_message(self, uuid: str, error: str):
+        message = TranscribeOutputMessage(uuid=uuid, if_success=False, msg=error)
+        self.send_message(message.model_dump())
+
+    def consume(self, channel, method, properties, message):
+        body = self.decode_message(message)
+
+        try:
+            validated_message = TranscribeInputMessage.model_validate(body)
+
+            audio_file = validated_message.audio_file
+            if validated_message.using_file_content:
+                audio_file = io.BytesIO(base64.b64decode(validated_message.audio_file))
+
+            segments, _ = self.asr_model.transcribe(audio_file, language=validated_message.language)
+
+            transcribed_text = ''
+            for segment in segments:
+                transcribed_text = segment.text
+                break
+        except ValidationError as exc:
+            logging.exception('Consume message failed: \n message: %s\n\n exception info: %s', message, exc)
+            self.send_fail_message(body.get('uuid'), f'{exc}')
+        except Exception as exc:
+            logging.exception('Consume message failed: \n message: %s\n\n exception info: %s', message, exc)
+            self.send_fail_message(body.get('uuid'), f'{exc}')
+        else:
+            self.send_success_message(validated_message.uuid, transcribed_text)

consumer/base.py

@@ -0,0 +1,235 @@
+import asyncio
+import functools
+import json
+import logging
+import os
+import ssl
+import time
+from enum import Enum
+from typing import Dict, Optional, Literal
+
+import msgpack
+import pika
+from pika.exceptions import AMQPConnectionError
+from pydantic import BaseModel, field_validator
+
+logger = logging.getLogger(__name__)
+
+
+def sync(f):
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        return asyncio.get_event_loop().run_until_complete(f(*args, **kwargs))
+
+    return wrapper
+
+
+class Priority(Enum):
+    LOW = 1
+    NORMAL = 5
+    HIGH = 10
+
+
+class Headers(BaseModel):
+    job_id: str
+    priority: Priority
+    task_type: Optional[str] = None
+
+    @field_validator('priority', mode='before')
+    @classmethod
+    def _convert_priority(cls, value):
+        return Priority[value]
+
+
+class RabbitMQConfig(BaseModel):
+    host: str
+    port: int
+    username: str
+    password: str
+    protocol: str
+
+
+class BasicPikaClient:
+    def __init__(self):
+        self.username = os.environ.get('RABBITMQ_USER', '')
+        self.password = os.environ.get('RABBITMQ_PASSWD', '')
+        self.host = os.environ.get('RABBITMQ_HOST', 'localhost')
+        self.port = os.environ.get('RABBITMQ_PORT', 5672)
+        self.protocol = "amqp"
+
+        self._init_connection_parameters()
+        self._connect()
+
+    def _connect(self):
+        tries = 0
+        while True:
+            try:
+                self.connection = pika.BlockingConnection(self.parameters)
+                self.channel = self.connection.channel()
+                if self.connection.is_open:
+                    break
+            except (AMQPConnectionError, Exception) as e:
+                time.sleep(5)
+                tries += 1
+                if tries == 20:
+                    raise AMQPConnectionError(e)
+
+    def _init_connection_parameters(self):
+        if any([self.username, self.password]):
+            self.credentials = pika.PlainCredentials(self.username, self.password)
+            self.parameters = pika.ConnectionParameters(
+                host=self.host,
+                port=int(self.port),
+                virtual_host="/",
+                credentials=self.credentials,
+            )
+        else:
+            self.parameters = pika.ConnectionParameters(
+                self.host,
+                int(self.port),
+                "/",
+            )
+
+        if self.protocol == "amqps":
+            # SSL Context for TLS configuration of Amazon MQ for RabbitMQ
+            ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
+            ssl_context.set_ciphers("ECDHE+AESGCM:!ECDSA")
+            self.parameters.ssl_options = pika.SSLOptions(context=ssl_context)
+
+    def check_connection(self):
+        if not self.connection or self.connection.is_closed:
+            self._connect()
+
+    def close(self):
+        self.channel.close()
+        self.connection.close()
+
+    def declare_queue(
+            self, queue_name, exclusive: bool = False, max_priority: int = 10
+    ):
+        self.check_connection()
+        logger.debug(f"Trying to declare queue({queue_name})...")
+
+        self.channel.queue_declare(
+            queue=queue_name,
+            exclusive=exclusive,
+            durable=True,
+            arguments={"x-max-priority": max_priority} if max_priority > 0 else None
+        )
+
+    def declare_exchange(self, exchange_name: str, exchange_type: str = "direct"):
+        self.check_connection()
+        self.channel.exchange_declare(
+            exchange=exchange_name, exchange_type=exchange_type
+        )
+
+    def bind_queue(self, exchange_name: str, queue_name: str, routing_key: str):
+        self.check_connection()
+        self.channel.queue_bind(
+            exchange=exchange_name, queue=queue_name, routing_key=routing_key
+        )
+
+    def unbind_queue(self, exchange_name: str, queue_name: str, routing_key: str):
+        self.channel.queue_unbind(
+            queue=queue_name, exchange=exchange_name, routing_key=routing_key
+        )
+
+
+class BasicMessageSender(BasicPikaClient):
+    message_encoding_type: Literal['bytes', 'json'] = 'json'
+
+    def encode_message(self, body: Dict, encoding_type: str = "bytes"):
+        if encoding_type == "bytes":
+            return msgpack.packb(body)
+        elif encoding_type == "json":
+            return json.dumps(body)
+        else:
+            raise NotImplementedError
+
+    def send_message(
+            self,
+            exchange_name: str,
+            routing_key: str,
+            body: Dict,
+            headers: Optional[Headers],
+    ):
+        body = self.encode_message(body=body, encoding_type=self.message_encoding_type)
+
+        properties = pika.BasicProperties(delivery_mode=pika.spec.PERSISTENT_DELIVERY_MODE,
+                                          priority=headers.priority.value,
+                                          headers=headers.model_dump() if headers else None)
+        self.channel.basic_publish(
+            exchange=exchange_name,
+            routing_key=routing_key,
+            body=body,
+            properties=properties,
+        )
+        logger.debug(
+            f"Sent message. Exchange: {exchange_name}, Routing Key: {routing_key}, Body: {body[:128]}"
+        )
+
+
+class BasicMessageReceiver(BasicPikaClient):
+    def __init__(self):
+        super().__init__()
+        self.channel_tag = None
+
+    def decode_message(self, body):
+        if type(body) == bytes:
+            return msgpack.unpackb(body)
+        elif type(body) == str:
+            return json.loads(body)
+        else:
+            raise NotImplementedError
+
+    def get_message(self, queue_name: str, auto_ack: bool = False):
+        method_frame, header_frame, body = self.channel.basic_get(
+            queue=queue_name, auto_ack=auto_ack
+        )
+        if method_frame:
+            logger.debug(f"{method_frame}, {header_frame}, {body}")
+            return method_frame, header_frame, body
+        else:
+            logger.debug("No message returned")
+            return None
+
+    def consume_messages(self, queue, callback):
+        self.check_connection()
+        self.channel_tag = self.channel.basic_consume(
+            queue=queue, on_message_callback=callback, auto_ack=True
+        )
+        logger.debug(" [*] Waiting for messages. To exit press CTRL+C")
+        self.channel.start_consuming()
+
+    def cancel_consumer(self):
+        if self.channel_tag is not None:
+            self.channel.basic_cancel(self.channel_tag)
+            self.channel_tag = None
+        else:
+            logger.error("Do not cancel a non-existing job")
+
+
+class ExtractFaceFramesConsumer(BasicMessageReceiver):
+    @sync
+    async def consume(self, channel, method, properties, body):
+        body = self.decode_message(body=body)
+        file_content = await self._download_image(img_url=body["url"])
+        # consume message logic ...
+
+    async def _download_image(self, img_url):
+        # do some async stuff here
+        pass
+
+
+def create_consumer():
+    worker = ExtractFaceFramesConsumer()
+    worker.declare_queue(queue_name="myqueue")
+    worker.declare_exchange(exchange_name="myexchange")
+    worker.bind_queue(
+        exchange_name="myexchange", queue_name="myqueue", routing_key="randomkey"
+    )
+    worker.consume_messages(queue="myqueue", callback=worker.consume)
+
+
+if __name__ == "__main__":
+    create_consumer()

requirements.txt

@@ -3,3 +3,4 @@ fastapi==0.115.5
 python-multipart==0.0.17
 websockets==14.1
 uvicorn==0.32.1
+pika==1.3.2