| | |
|---|
| | | import asyncio |
|---|
| | | import websockets |
|---|
| | | from sqlalchemy.orm import Session |
|---|
| | | |
|---|
| | | from Log import logger |
|---|
| | | from app.api import get_current_user_websocket |
|---|
| | | from app.config.config import settings |
|---|
| | | from app.models.agent_model import AgentModel, AgentType |
|---|
| | | from app.models.base_model import get_db |
|---|
| | | from app.models.user_model import UserModel |
|---|
| | | from app.service.dialog import update_session_history |
|---|
| | | from app.service.basic import BasicService |
|---|
| | | from app.service.ragflow import RagflowService |
|---|
| | | from app.service.token import get_bisheng_token, get_ragflow_token |
|---|
| | | from app.service.service_token import get_bisheng_token, get_ragflow_token |
|---|
| | | from app.service.session import SessionService |
|---|
| | | |
|---|
| | | router = APIRouter() |
|---|
| | | |
|---|
| | | # 存储客户端 WebSocket 连接 |
|---|
| | | client_websockets = {} |
|---|
| | | |
|---|
| | | |
|---|
| | | # 中间层WebSocket 服务器，接收客户端的连接 |
|---|
| | |
|---|
| | | chat_id: str, |
|---|
| | | current_user: UserModel = Depends(get_current_user_websocket), |
|---|
| | | db: Session = Depends(get_db)): |
|---|
| | | tasks = [] |
|---|
| | | await websocket.accept() |
|---|
| | | print(f"Client {agent_id} connected") |
|---|
| | | |
|---|
| | | if agent_id == "0": |
|---|
| | | agent_id = settings.bisheng_agent_id |
|---|
| | | elif agent_id == "1": |
|---|
| | | agent_id = settings.ragflow_agent_id |
|---|
| | | chat_id = settings.ragflow_chat_id |
|---|
| | | agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first() |
|---|
| | | if not agent: |
|---|
| | | ret = {"message": "Agent not found", "type": "close"} |
|---|
| | | await websocket.send_json(ret) |
|---|
| | | return |
|---|
| | | agent_type = agent.agent_type |
|---|
| | | if chat_id == "" or chat_id == "0": |
|---|
| | | ret = {"message": "Chat ID not found", "type": "close"} |
|---|
| | | await websocket.send_json(ret) |
|---|
| | | return |
|---|
| | | |
|---|
| | | if chat_id == "0": |
|---|
| | | chat_id = uuid.uuid4().hex |
|---|
| | | |
|---|
| | | client_websockets[chat_id] = websocket |
|---|
| | | if agent_id == settings.ragflow_agent_id: |
|---|
| | | ragflow_service = RagflowService(settings.ragflow_base_url) |
|---|
| | | if agent_type == AgentType.RAGFLOW: |
|---|
| | | ragflow_service = RagflowService(settings.fwr_base_url) |
|---|
| | | token = get_ragflow_token(db, current_user.id) |
|---|
| | | try: |
|---|
| | | async def forward_to_ragflow(): |
|---|
| | | while True: |
|---|
| | | is_new = False |
|---|
| | | message = await websocket.receive_json() |
|---|
| | | print(f"Received from client {chat_id}: {message}") |
|---|
| | | async for rag_response in ragflow_service.chat(token, chat_id, message["chatHistory"]): |
|---|
| | | chat_history = message.get('chatHistory', []) |
|---|
| | | message["role"] = "user" |
|---|
| | | if len(chat_history) == 0: |
|---|
| | | chat_history = await ragflow_service.get_session_history(token, chat_id) |
|---|
| | | if len(chat_history) == 0: |
|---|
| | | is_new = True |
|---|
| | | chat_history = await ragflow_service.set_session(token, agent_id, |
|---|
| | | message, chat_id, True) |
|---|
| | | # print("chat_history------------------------", chat_history) |
|---|
| | | if len(chat_history) == 0: |
|---|
| | | result = {"message": "内部错误：创建会话失败", "type": "close"} |
|---|
| | | await websocket.send_json(result) |
|---|
| | | await websocket.close() |
|---|
| | | return |
|---|
| | | else: |
|---|
| | | chat_history.append({ |
|---|
| | | "content": message["message"], |
|---|
| | | "doc_ids": message.get("doc_ids", []), |
|---|
| | | "role": "user" |
|---|
| | | }) |
|---|
| | | complete_response = "" |
|---|
| | | async for rag_response in ragflow_service.chat(token, chat_id, chat_history): |
|---|
| | | try: |
|---|
| | | print(f"Received from ragflow: {rag_response}") |
|---|
| | | json_str = rag_response[5:].strip() |
|---|
| | | json_data = json.loads(json_str) |
|---|
| | | if json_data.get("data") is not True: |
|---|
| | | answer = json_data.get("data", {}).get("answer", "") |
|---|
| | | result = {"message": answer, "type": "stream"} |
|---|
| | | if rag_response[:5] == "data:": |
|---|
| | | # 如果是，则截取掉前5个字符，并去除首尾空白符 |
|---|
| | | text = rag_response[5:].strip() |
|---|
| | | else: |
|---|
| | | result = {"message": "", "type": "close"} |
|---|
| | | # 否则，保持原样 |
|---|
| | | text = rag_response |
|---|
| | | complete_response += text |
|---|
| | | try: |
|---|
| | | json_data = json.loads(complete_response) |
|---|
| | | data = json_data.get("data") |
|---|
| | | if data is True: # 完成输出 |
|---|
| | | result = {"message": "", "type": "close"} |
|---|
| | | elif data is None: # 发生错误 |
|---|
| | | answer = json_data.get("retmsg", json_data.get("retcode")) |
|---|
| | | result = {"message": "内部错误：" + answer, "type": "message"} |
|---|
| | | else: # 正常输出 |
|---|
| | | answer = data.get("answer", "") |
|---|
| | | reference = data.get("reference", {}) |
|---|
| | | result = {"message": answer, "type": "message", "reference": reference} |
|---|
| | | await websocket.send_json(result) |
|---|
| | | complete_response = "" |
|---|
| | | except json.JSONDecodeError as e: |
|---|
| | | print(f"Error decoding JSON: {e}") |
|---|
| | | # print(f"Response text: {text}") |
|---|
| | | except Exception as e2: |
|---|
| | | result = {"message": f"内部错误： {e2}", "type": "close"} |
|---|
| | | await websocket.send_json(result) |
|---|
| | | print(f"Forwarded to client {chat_id}: {result}") |
|---|
| | | except Exception as e: |
|---|
| | | print(f"Error forwarding message to ragflow: {e}") |
|---|
| | | print(f"Error process message of ragflow: {e2}") |
|---|
| | | dialog_chat_history = await ragflow_service.get_session_history(token, chat_id, 1) |
|---|
| | | await update_session_history(db, dialog_chat_history, current_user.id, is_new) |
|---|
| | | |
|---|
| | | # 启动任务处理客户端消息 |
|---|
| | | tasks = [ |
|---|
| | | asyncio.create_task(forward_to_ragflow()) |
|---|
| | | ] |
|---|
| | | await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) |
|---|
| | | except WebSocketDisconnect: |
|---|
| | | print(f"Client {chat_id} disconnected") |
|---|
| | | finally: |
|---|
| | | del client_websockets[chat_id] |
|---|
| | | except WebSocketDisconnect as e1: |
|---|
| | | print(f"Client {chat_id} disconnected: {e1}") |
|---|
| | | await websocket.close() |
|---|
| | | except Exception as e: |
|---|
| | | print(f"Exception occurred: {e}") |
|---|
| | | |
|---|
| | | else: |
|---|
| | | finally: |
|---|
| | | print("Cleaning up resources of ragflow") |
|---|
| | | # 取消所有任务 |
|---|
| | | for task in tasks: |
|---|
| | | if not task.done(): |
|---|
| | | task.cancel() |
|---|
| | | try: |
|---|
| | | await task |
|---|
| | | except asyncio.CancelledError: |
|---|
| | | pass |
|---|
| | | |
|---|
| | | elif agent_type == AgentType.BISHENG: |
|---|
| | | token = get_bisheng_token(db, current_user.id) |
|---|
| | | service_uri = f"{settings.bisheng_websocket_url}/api/v1/assistant/chat/{agent_id}?t=&chat_id={chat_id}" |
|---|
| | | service_uri = f"{settings.sgb_websocket_url}/api/v1/assistant/chat/{agent_id}?t=&chat_id={chat_id}" |
|---|
| | | headers = {'cookie': f"access_token_cookie={token};"} |
|---|
| | | |
|---|
| | | async with websockets.connect(service_uri, extra_headers=headers) as service_websocket: |
|---|
| | |
|---|
| | | async def forward_to_service(): |
|---|
| | | while True: |
|---|
| | | message = await websocket.receive_json() |
|---|
| | | print(f"Received from client {chat_id}: {message}") |
|---|
| | | print(f"Received from client, {chat_id}: {message}") |
|---|
| | | # 添加 'agent_id' 和 'chat_id' 字段 |
|---|
| | | message['flow_id'] = agent_id |
|---|
| | | message['chat_id'] = chat_id |
|---|
| | |
|---|
| | | async def forward_to_client(): |
|---|
| | | while True: |
|---|
| | | message = await service_websocket.recv() |
|---|
| | | print(f"Received from service S: {message}") |
|---|
| | | await websocket.send_text(message) |
|---|
| | | print(f"Forwarded to client {chat_id}: {message}") |
|---|
| | | print(f"Received from bisheng: {message}") |
|---|
| | | data = json.loads(message) |
|---|
| | | if data["type"] == "close" or data["type"] == "stream" or data["type"] == "end_cover": |
|---|
| | | if data["type"] == "close": |
|---|
| | | t = "close" |
|---|
| | | else: |
|---|
| | | t = "stream" |
|---|
| | | result = {"message": data["message"], "type": t} |
|---|
| | | await websocket.send_json(result) |
|---|
| | | print(f"Forwarded to client, {chat_id}: {result}") |
|---|
| | | |
|---|
| | | # 启动两个任务，分别处理客户端和服务端的消息 |
|---|
| | | tasks = [ |
|---|
| | |
|---|
| | | except asyncio.CancelledError: |
|---|
| | | pass |
|---|
| | | |
|---|
| | | except WebSocketDisconnect: |
|---|
| | | print(f"Client {chat_id} disconnected") |
|---|
| | | except WebSocketDisconnect as e: |
|---|
| | | print(f"WebSocket connection closed with code {e.code}: {e.reason}") |
|---|
| | | await websocket.close() |
|---|
| | | await service_websocket.close() |
|---|
| | | except Exception as e: |
|---|
| | | print(f"Exception occurred: {e}") |
|---|
| | | |
|---|
| | | finally: |
|---|
| | | del client_websockets[chat_id] |
|---|
| | | print("Cleaning up resources of bisheng") |
|---|
| | | # 取消所有任务 |
|---|
| | | for task in tasks: |
|---|
| | | if not task.done(): |
|---|
| | | task.cancel() |
|---|
| | | try: |
|---|
| | | await task |
|---|
| | | except asyncio.CancelledError: |
|---|
| | | pass |
|---|
| | | elif agent_type == AgentType.BASIC: |
|---|
| | | try: |
|---|
| | | service = BasicService(base_url=settings.basic_base_url) |
|---|
| | | while True: |
|---|
| | | # 接收前端消息 |
|---|
| | | message = await websocket.receive_json() |
|---|
| | | question = message.get("message") |
|---|
| | | try: |
|---|
| | | SessionService(db).create_session( |
|---|
| | | session_id=chat_id, |
|---|
| | | name=question, |
|---|
| | | agent_id=agent_id, |
|---|
| | | agent_type=AgentType.BASIC, |
|---|
| | | user_id=current_user.id |
|---|
| | | ) |
|---|
| | | except Exception as e: |
|---|
| | | logger.error(e) |
|---|
| | | if not question: |
|---|
| | | await websocket.send_json({"message": "Invalid request", "type": "error"}) |
|---|
| | | continue |
|---|
| | | if agent.agent_type == "questionTalk": |
|---|
| | | async for result in service.questions_talk(question, chat_id): |
|---|
| | | try: |
|---|
| | | if result[:5] == "data:": |
|---|
| | | # 如果是，则截取掉前5个字符，并去除首尾空白符 |
|---|
| | | text = result[5:].strip() |
|---|
| | | else: |
|---|
| | | # 否则，保持原样 |
|---|
| | | text = result |
|---|
| | | try: |
|---|
| | | data = json.loads(text) |
|---|
| | | output = data.get("output", "") |
|---|
| | | file_name = data.get("filename", "") |
|---|
| | | |
|---|
| | | |
|---|
| | | excel_url = None |
|---|
| | | if file_name: |
|---|
| | | excel_url = f"/api/files/download/?agent_id=basic_question_talk&file_id={file_name}&file_type=word" |
|---|
| | | result = {"message": output, "type": "message", "file_url": excel_url} |
|---|
| | | try: |
|---|
| | | SessionService(db).update_session(session_id=chat_id, is_incr=1, |
|---|
| | | message={"role": "assistant", "content": result}) |
|---|
| | | except Exception as e: |
|---|
| | | logger.error(e) |
|---|
| | | await websocket.send_json(result | data) |
|---|
| | | except json.JSONDecodeError as e: |
|---|
| | | print(f"Error decoding JSON: {e}") |
|---|
| | | # print(f"Response text: {text}") |
|---|
| | | except Exception as e2: |
|---|
| | | result = {"message": f"内部错误： {e2}", "type": "close"} |
|---|
| | | await websocket.send_json(result) |
|---|
| | | print(f"Error process message of basic agent: {e2}") |
|---|
| | | else: |
|---|
| | | async for result in service.excel_talk(question, chat_id): |
|---|
| | | try: |
|---|
| | | if result[:5] == "data:": |
|---|
| | | # 如果是，则截取掉前5个字符，并去除首尾空白符 |
|---|
| | | text = result[5:].strip() |
|---|
| | | else: |
|---|
| | | # 否则，保持原样 |
|---|
| | | text = result |
|---|
| | | try: |
|---|
| | | data = json.loads(text) |
|---|
| | | output = data.get("output", "") |
|---|
| | | excel_name = data.get("excel_name", "") |
|---|
| | | image_name = data.get("image_name", "") |
|---|
| | | excel_url = None |
|---|
| | | image_url = None |
|---|
| | | if excel_name: |
|---|
| | | excel_url = f"/api/files/download/?agent_id=basic_excel_talk&file_id={excel_name}&file_type=excel" |
|---|
| | | if image_name: |
|---|
| | | image_url = f"/api/files/download/?agent_id=basic_excel_talk&file_id={image_name}&file_type=image" |
|---|
| | | result = {"message": output, "type": "message", "excel_url": excel_url, "image_url": image_url} |
|---|
| | | try: |
|---|
| | | SessionService(db).update_session(session_id=chat_id, is_incr=1, |
|---|
| | | message={"role": "assistant", "content": result}) |
|---|
| | | except Exception as e: |
|---|
| | | logger.error(e) |
|---|
| | | await websocket.send_json(result | data) |
|---|
| | | except json.JSONDecodeError as e: |
|---|
| | | print(f"Error decoding JSON: {e}") |
|---|
| | | # print(f"Response text: {text}") |
|---|
| | | except Exception as e2: |
|---|
| | | result = {"message": f"内部错误： {e2}", "type": "close"} |
|---|
| | | await websocket.send_json(result) |
|---|
| | | print(f"Error process message of basic agent: {e2}") |
|---|
| | | except Exception as e: |
|---|
| | | await websocket.send_json({"message": str(e), "type": "error"}) |
|---|
| | | finally: |
|---|
| | | await websocket.close() |
|---|
| | | print(f"Client {agent_id} disconnected") |
|---|
| | | else: |
|---|
| | | ret = {"message": "Agent not found", "type": "close"} |
|---|
| | | await websocket.send_json(ret) |
|---|
