| | |
|---|
| | | 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 |
|---|
| | |
|---|
| | | 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") |
|---|
| | | SessionService(db).create_session( |
|---|
| | | session_id=chat_id, |
|---|
| | | name=question, |
|---|
| | | agent_id=agent_id, |
|---|
| | | agent_type=AgentType.BASIC |
|---|
| | | ) |
|---|
| | | 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 |
|---|
| | | |
|---|
| | | service = BasicService(base_url=settings.basic_base_url) |
|---|
| | | async for result in service.excel_talk(question, chat_id): |
|---|
| | | try: |
|---|
| | | if result[:5] == "data:": |
|---|
| | | # 如果是，则截取掉前5个字符，并去除首尾空白符 |
|---|
| | | text = result[5:].strip() |
|---|
| | | else: |
|---|
| | | # 否则，保持原样 |
|---|
| | | text = result |
|---|
| | | if agent.agent_type == "questionTalk": |
|---|
| | | async for result in service.questions_talk(question, chat_id): |
|---|
| | | try: |
|---|
| | | data = json.loads(text) |
|---|
| | | output = data.get("output", "") |
|---|
| | | excel_name = data.get("excel_name", "") |
|---|
| | | image_name = data.get("excel_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} |
|---|
| | | 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}") |
|---|
| | | 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: |
|---|
