AI后端程序增加对RabbitMQ的使用

2025-10-09 14:11:37 +08:00
parent 3fb43c09f3
commit 4d38f5483e
12 changed files with 175 additions and 134 deletions
@@ -0,0 +1,69 @@
+# consumer.py
+import json
+import time
+
+import pika
+
+from config.rabbitMQ import *
+from models.AnalysisRequest import AnalysisRequest
+
+
+def mq_new_analysis(req: AnalysisRequest):
+    """将分析请求发送到 RabbitMQ 队列"""
+    credentials = pika.PlainCredentials(RABBIT_USER, RABBIT_PASSWORD)
+    connection = pika.BlockingConnection(
+        pika.ConnectionParameters(
+            host=RABBIT_HOST, virtual_host=RABBIT_VHOST, credentials=credentials
+        )
+    )
+    channel = connection.channel()
+    channel.queue_declare(queue=QUEUE_NAME, durable=True)
+
+    message = json.dumps(req.dict())
+    try:
+        channel.basic_publish(
+            exchange="",
+            routing_key=QUEUE_NAME,
+            body=message,
+            properties=pika.BasicProperties(delivery_mode=2),  # 持久化消息
+        )
+    finally:
+        channel.close()
+        connection.close()
+
+
+def mq_pull_analysis():
+    """
+    从队列拉取分析任务并处理
+    process_func: 一个函数，接收 AnalysisRequest 对象处理分析逻辑
+    """
+    credentials = pika.PlainCredentials(RABBIT_USER, RABBIT_PASSWORD)
+    connection = pika.BlockingConnection(
+        pika.ConnectionParameters(
+            host=RABBIT_HOST, virtual_host=RABBIT_VHOST, credentials=credentials
+        )
+    )
+    channel = connection.channel()
+    channel.queue_declare(queue=QUEUE_NAME, durable=True)
+    channel.basic_qos(prefetch_count=1)  # 保证一次只处理一条消息
+
+    def callback(ch, method, properties, body):
+        try:
+            data = json.loads(body)
+            req = AnalysisRequest(**data)
+            print(f"收到任务: {req}")
+            time.sleep(5)
+            print(f"完成任务: {req}")
+        except Exception as e:
+            print(f"处理任务出错: {e}")
+        finally:
+            ch.basic_ack(delivery_tag=method.delivery_tag)
+
+    channel.basic_consume(queue=QUEUE_NAME, on_message_callback=callback)
+    print("开启队列消费者")
+    print("等待分析任务...")
+    try:
+        channel.start_consuming()
+    finally:
+        channel.close()
+        connection.close()