BBIT-Kai
76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
# import db.postgres as pg
|
|
import uuid
|
|
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
|
|
|
|
|
|
# 对话历史记录
|
|
@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"))
|
|
|
|
|
|
@chatRouter.post("/chatForBot")
|
|
def chat(req: ChatRequest, user_id: UUID = Depends(get_user_id_from_token)):
|
|
if not user_id:
|
|
return {"error": "userId is required"}
|
|
if not req.aiId:
|
|
return {"error": "aiId is required"}
|
|
session_name = get_title(req.userInput)
|
|
# 如果没有 sessionId 就新建
|
|
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")
|
|
else:
|
|
is_new_session = False
|
|
pg.update_session_updated_at(req.sessionId)
|
|
|
|
# 插入用户消息
|
|
pg.insert_message(req.sessionId, False, req.userInput)
|
|
|
|
# 调用 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,
|
|
)
|
|
else:
|
|
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,
|
|
}
|
|
)
|