首页知识库对话-rg

zhaoqingang

2025-03-18 282a631b9ceee9a634ee1d93751a5254ed37ccef

 app/service/v2/chat.py

@@ -6,6 +6,7 @@

import fitz
from fastapi import HTTPException
from sqlalchemy import or_

from Log import logger
from app.config.agent_base_url import RG_CHAT_DIALOG, DF_CHAT_AGENT, DF_CHAT_PARAMETERS, RG_CHAT_SESSIONS, \
@@ -13,7 +14,7 @@
from app.config.config import settings
from app.config.const import *
from app.models import DialogModel, ApiTokenModel, UserTokenModel, ComplexChatSessionDao, ChatDataRequest, \
    ComplexChatDao
    ComplexChatDao, KnowledgeModel, UserModel
from app.models.v2.session_model import ChatSessionDao, ChatData
from app.service.v2.app_driver.chat_agent import ChatAgent
from app.service.v2.app_driver.chat_data import ChatBaseApply
@@ -87,17 +88,45 @@
        return ChatAgent(), url


async def service_chat_dialog(db, chat_id: str, question: str, session_id: str, user_id, mode: str):

async def get_user_kb(db, user_id: int, kb_ids: list) -> list:
    res = []
    user = db.query(UserModel).filter(UserModel.id == user_id).first()
    if user is None:
        return res
    query = db.query(KnowledgeModel)
    if user.permission != "admin":
        klg_list = [j.id for i in user.groups for j in i.knowledges]
        query = query.filter(or_(KnowledgeModel.id.in_(klg_list), KnowledgeModel.tenant_id == str(user_id)))
        kb_list= query.all()
        for kb in kb_list:
            if kb.id in kb_ids:
                if kb.permission == "team":
                    res.append(kb.id)
                elif kb.tenant_id == str(user_id):
                    res.append(kb.id)
        return res
    else:
        return kb_ids


async def service_chat_dialog(db, chat_id: str, question: str, session_id: str, user_id: int, mode: str, kb_ids: list):
    conversation_id = ""
    token = await get_chat_token(db, rg_api_token)
    url = settings.fwr_base_url + RG_CHAT_DIALOG.format(chat_id)
    kb_id = await get_user_kb(db, user_id, kb_ids)
    if not kb_id:
        yield "data: " + json.dumps({"message": smart_message_error,
                                     "error": "\n**ERROR**: The agent has no knowledge base to work with!", "status": http_400},
                                    ensure_ascii=False) + "\n\n"
        return
    chat = ChatDialog()
    session = await add_session_log(db, session_id, question, chat_id, user_id, mode, session_id, RG_TYPE)
    if session:
        conversation_id = session.conversation_id
    message = {"role": "assistant", "answer": "", "reference": {}}
    try:
        async for ans in chat.chat_completions(url, await chat.request_data(question, conversation_id),
        async for ans in chat.chat_completions(url, await chat.complex_request_data(question, kb_id, conversation_id),
                                               await chat.get_headers(token)):
            data = {}
            error = ""