完善权限系统

bbit_ai/app/routers/Chat.py

@@ -1,33 +1,33 @@
-from models.ChatRequest import ChatRequest
-from models.ChatWithReportRequest import ChatWithReportRequest
-from models.BaseResponse import BaseResponse
+# import db.postgres as pg
 import uuid
-import db.postgres as pg
-import db.sqlserver as sqlserver
-import uuid
-import threading
-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
+
 chatRouter = APIRouter()
 from agent.dataAgent import get_graph_output
 from llm.chatLLM import get_chat_response
 from llm.titleChain import get_title
-from llm.sqlLLM import get_sql_description_response,get_chat_sql_response,get_chat_sql_improve_response
-from models.SaveReportRequest import SaveReportRequest
-from models.AIProfilesRequest import AIProfilesRequest
-import json
+
+
 # 对话历史记录
 @chatRouter.get("/history")
 def getHistory(sessionId: str):
     return BaseResponse(data=pg.get_history(sessionId))
 
+
 # 对话列表
 @chatRouter.get("/sessionsForBot")
 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,'bot'))
+    return BaseResponse(data=pg.get_sessions(user_id, "bot"))
+
 
 @chatRouter.post("/chatForBot")
 def chat(req: ChatRequest, user_id: UUID = Depends(get_user_id_from_token)):
@@ -40,7 +40,7 @@ def chat(req: ChatRequest, user_id: UUID = Depends(get_user_id_from_token)):
     if not req.sessionId:
         is_new_session = True
         req.sessionId = str(uuid.uuid4())
-        pg.insert_session(user_id,req.aiId, req.sessionId, session_name,"bot")
+        pg.insert_session(user_id, req.aiId, req.sessionId, session_name, "bot")
     else:
         is_new_session = False
         pg.update_session_updated_at(req.sessionId)
@@ -51,11 +51,25 @@ def chat(req: ChatRequest, user_id: UUID = Depends(get_user_id_from_token)):
     # 调用 LLM
     if req.aiId == "9d157dd1-921b-c768-5b90-3e903b50f6f9":
         # 数据专家AI
-        answer = get_graph_output(aiRole=pg.get_ai_personality(req.aiId),history=pg.get_history(req.sessionId), userInput= req.userInput)
+        answer = get_graph_output(
+            aiRole=pg.get_ai_personality(req.aiId),
+            history=pg.get_history(req.sessionId),
+            userInput=req.userInput,
+        )
     else:
-        answer = get_chat_response(aiRole=pg.get_ai_personality(req.aiId),history=pg.get_history(req.sessionId), userInput= req.userInput).content
+        answer = get_chat_response(
+            aiRole=pg.get_ai_personality(req.aiId),
+            history=pg.get_history(req.sessionId),
+            userInput=req.userInput,
+        ).content
     # 插入 AI 回复
     pg.insert_message(req.sessionId, True, answer)
 
-    return BaseResponse(data={"session_name":session_name,"isNewSession":is_new_session,"content":answer,"sessionId": req.sessionId})
-
+    return BaseResponse(
+        data={
+            "session_name": session_name,
+            "isNewSession": is_new_session,
+            "content": answer,
+            "sessionId": req.sessionId,
+        }
+    )