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

bbit_ai/app/app.py

@@ -6,9 +6,12 @@ from routers.Chat import chatRouter
 from routers.Datasource import reportDataRouter
 from routers.F8 import f8Router
 from routers.Knowledge import knowledgeRouter
+from routers.RabbitMQ import rqRouter
 from routers.Report import reportRouter
 from routers.Service import serviceRouter
 from routers.Vision import visionRouter
+from service.Analyze import mq_pull_analysis
+from utils import MyUtils
 
 app = FastAPI(title="BBIT_AI")
 
@@ -34,9 +37,12 @@ routers = [
     serviceRouter,
     botRouter,
     visionRouter,
+    rqRouter,
 ]
 
 for r in routers:
     app.include_router(r, prefix="/llm", tags=["llm"])
 
 app.include_router(f8Router, prefix="/f8", tags=["f8"])
+
+MyUtils.async_new_task(mq_pull_analysis)

bbit_ai/app/config/rabbitMQ.py

@@ -0,0 +1,5 @@
+RABBIT_HOST = "10.10.10.9"
+RABBIT_VHOST = "/bbit_ai"
+RABBIT_USER = "bbit_ai"
+RABBIT_PASSWORD = "123456"
+QUEUE_NAME = "analysis_queue"

bbit_ai/app/models/AnalysisRequest.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+# 请求
+class AnalysisRequest(BaseModel):
+    data: str  # 你可以根据需要扩展字段

bbit_ai/app/routers/RabbitMQ.py

@@ -0,0 +1,12 @@
+from fastapi import APIRouter
+
+from models.AnalysisRequest import AnalysisRequest
+from service.Analyze import mq_new_analysis
+
+rqRouter = APIRouter()
+
+
+@rqRouter.post("/analyze")
+def send_analysis_request(req: AnalysisRequest):
+    mq_new_analysis(req)
+    return {"status": "queued"}

bbit_ai/app/routers/Service.py

@@ -1,22 +1,27 @@
-from models.ChatRequest import ChatRequest
-from models.BaseResponse import BaseResponse
 import uuid
-import db.postgres as pg
-import uuid
-from fastapi import APIRouter, Depends
 from uuid import UUID
+
+from fastapi import APIRouter, Depends
+
+import db.postgres as pg
 from config.security import get_user_id_from_token
+from models.BaseResponse import BaseResponse
+from models.ChatRequest import ChatRequest
+
 serviceRouter = APIRouter()
 from llm.titleChain import get_title
 from agent.serviceAgent import get_service_agent_reply
 from llm.memLLM import take_memory
 import utils.MyUtils as utils
+
+
 # 对话列表
 @serviceRouter.get("/sessionsForService")
 def getSessions(user_id: UUID = Depends(get_user_id_from_token)):
     if not user_id:
         return {"error": "userId is required"}
-    return BaseResponse(data=pg.get_sessions(user_id,'service'))
+    return BaseResponse(data=pg.get_sessions(user_id, "service"))
+
 
 # 对话
 @serviceRouter.post("/chatForService")
@@ -30,7 +35,7 @@ def chat(req: ChatRequest, user_id: UUID = Depends(get_user_id_from_token)):
     if not req.sessionId:
         isNewSession = True
         req.sessionId = str(uuid.uuid4())
-        pg.insert_session(user_id,req.aiId, req.sessionId, sessionName, "service")
+        pg.insert_session(user_id, req.aiId, req.sessionId, sessionName, "service")
     else:
         isNewSession = False
         pg.update_session_updated_at(req.sessionId)
@@ -38,10 +43,27 @@ def chat(req: ChatRequest, user_id: UUID = Depends(get_user_id_from_token)):
     # 插入用户消息
     pg.insert_message(req.sessionId, False, req.userInput)
 
-    answer = get_service_agent_reply(aiId=req.aiId,history=pg.get_history_with_time(req.sessionId,6), userInput= req.userInput,kn_bases=pg.get_ai_available_kn_bases(req.aiId))
+    answer = get_service_agent_reply(
+        aiId=req.aiId,
+        history=pg.get_history_with_time(req.sessionId, 6),
+        userInput=req.userInput,
+        kn_bases=pg.get_ai_available_kn_bases(req.aiId),
+    )
     # 插入 AI 回复
     pg.insert_message(req.sessionId, True, answer)
 
     # 异步执行：记忆判断
-    utils.async_db_task(take_memory,req.aiId,req.sessionId,user_id,)
-    return BaseResponse(data={"sessionName":sessionName,"isNewSession":isNewSession,"content":answer,"sessionId": req.sessionId})
+    utils.async_new_task(
+        take_memory,
+        req.aiId,
+        req.sessionId,
+        user_id,
+    )
+    return BaseResponse(
+        data={
+            "sessionName": sessionName,
+            "isNewSession": isNewSession,
+            "content": answer,
+            "sessionId": req.sessionId,
+        }
+    )

bbit_ai/app/service/Analyze.py

@@ -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()

bbit_ai/app/utils/MyUtils.py

@@ -1,24 +1,19 @@
+import asyncio
 import threading
+from datetime import datetime
+
+import pytz
 
 
 # 后台操作
-def async_db_task(func, *args, **kwargs):
-    """将数据库操作放到后台线程执行"""
+def async_new_task(func, *args, **kwargs):
     threading.Thread(target=func, args=args, kwargs=kwargs, daemon=True).start()
 
 
-import asyncio
-
-
 async def async_task(func, *args, **kwargs):
     return await asyncio.to_thread(func, *args, **kwargs)
 
 
-from datetime import datetime
-
-import pytz
-
-
 def format_datetime(dt: datetime, tz="Asia/Shanghai"):
     if dt.tzinfo is None:
         dt = pytz.UTC.localize(dt)

bbit_ai/docker/requirements.txt

@@ -20,3 +20,4 @@ minio==7.2.16
 pyzxing==1.1.1
 Pillow==11.3.0
 python-multipart==0.0.20
+pika==1.3.2