zhaoqingang
2025-02-25 383c21560fcb7012cd2e5b15d559e46d038c87b0
app/api/chat.py
@@ -11,7 +11,8 @@
from app.api import get_current_user_websocket
from app.config.config import settings
from app.config.const import IMAGE_TO_TEXT, DOCUMENT_TO_REPORT, DOCUMENT_TO_CLEANING, DOCUMENT_IA_QUESTIONS, \
    DOCUMENT_TO_REPORT_TITLE, DOCUMENT_TO_TITLE, DOCUMENT_TO_PAPER, DOCUMENT_IA_QUESTIONS_DS
    DOCUMENT_TO_REPORT_TITLE, DOCUMENT_TO_TITLE, DOCUMENT_TO_PAPER, DOCUMENT_IA_QUESTIONS_DS, \
    DOCUMENT_IA_QUESTIONS_EQUIPMENT
from app.models import MenuCapacityModel
from app.models.agent_model import AgentModel, AgentType
from app.models.base_model import get_db
@@ -39,6 +40,7 @@
    print(f"Client {agent_id} connected")
    agent = db.query(MenuCapacityModel).filter(MenuCapacityModel.chat_id == agent_id).first()
    if not agent:
        print("Agent not found")
        agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
        agent_type = agent.agent_type
        chat_type = agent.type
@@ -808,10 +810,11 @@
                                        result = {"message": f"内部错误: {e2}", "type": "close"}
                                        await websocket.send_json(result)
                                        print(f"Error process message of ragflow: {e2}")
                elif chat_type == "documentIa" or chat_type == "documentIaDs":
                elif chat_type == "documentIa" or chat_type == "documentIaDs" or chat_type == "documentIaEq":
                    token_dict = {
                        "documentIa": DOCUMENT_IA_QUESTIONS,
                        "documentIaDs": DOCUMENT_IA_QUESTIONS_DS,
                        "documentIaEq": DOCUMENT_IA_QUESTIONS_EQUIPMENT,
                    }
                    token = DfTokenDao(db).get_token_by_id(token_dict[chat_type])
                    # print(token)