rag-gateway/app/service/v2/chat.py

import json

from Log import logger
from app.config.agent_base_url import RG_CHAT_DIALOG
from app.config.config import settings
from app.config.const import max_chunk_size
from app.models.v2.session_model import ChatSessionDao
from app.service.v2.app_driver.chat_dialog import ChatDialog


async def service_chat_dialog(db, chat_id:str, question: str, session_id: str, user_id):
    token = "ragflow-YzMzE1NDRjYzMyZjExZWY5ZjkxMDI0Mm"
    url = settings.fwr_base_url+RG_CHAT_DIALOG.format(chat_id)
    chat = ChatDialog(token)
    request_data = {
        "question": question,
        "stream": True,
        "session_id": session_id
    }
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f"Bearer {token}"
    }
    try:
        await ChatSessionDao(db).update_or_insert_by_id(
            session_id=session_id,
            name=question[:255],
            agent_id=chat_id,
            agent_type=1,
            tenant_id=user_id,
            message={"role": "user", "content": question},
            conversation_id=session_id,
            event_type="message"
        )
    except Exception as e:
        logger.error(e)
    try:
        message = {"role": "assistant","answer":"", "reference": {}}
        async for ans in chat.chat_completions(url, request_data, headers):
            if ans.get("code", None) == 102:
                error = ans.get("message", "请输入你的问题！")
                data = {"answer":error}
                event = "message"
            else:
                if isinstance(ans.get("data"), bool) and ans.get("data") is True:
                    data = {}
                    event = "message_end"
                else:
                    data = ans.get("data", {})
                    message = ans.get("data", {})
                    event = "message"
            message_str = "data: " + json.dumps({"event": event, "data": data}, ensure_ascii=False) + "\n\n"
            for i in range(0, len(message_str), max_chunk_size):
                chunk = message_str[i:i + max_chunk_size]
                # print(chunk)
                yield chunk  # 发送分块消息
        await ChatSessionDao(db).update_session_by_id(
            session_id=session_id,
            session=None,
            message=message
        )
    except Exception as e:
        logger.error(e)
        yield "data: " + json.dumps({"message": "message",
                                    "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
                                   ensure_ascii=False) + "\n\n"

        yield "data: " + json.dumps({"message": "message_end",
                                    "data": {}},
                                   ensure_ascii=False) + "\n\n"