首页知识库对话-rg

zhaoqingang

2025-03-18 282a631b9ceee9a634ee1d93751a5254ed37ccef

首页知识库对话-rg

 app/api/v2/chat.py

@@ -42,7 +42,7 @@
            return StreamingResponse(f"data: {error_msg}\n\n",
                                     media_type="text/event-stream")
        session_id = session.get("data", {}).get("id")
    return StreamingResponse(service_chat_dialog(db, chatId, dialog.query, session_id, current_user.id, chat_info.mode),
    return StreamingResponse(service_chat_dialog(db, chatId, dialog.query, session_id, current_user.id, chat_info.mode, chat_info.get_kb_ids()),
                             media_type="text/event-stream")

@chat_router_v2.post("/agent/{chatId}/completions")

 app/api/v2/mindmap.py

@@ -14,7 +14,7 @@
from app.models.v2.chat import RetrievalRequest, ComplexChatDao
from app.models.v2.mindmap import MindmapRequest
from app.models.v2.session_model import ChatData
from app.service.v2.mindmap import service_chat_mindmap
from app.service.v2.mindmap import service_chat_mindmap, service_message_mindmap_parse

mind_map_router = APIRouter()

@@ -28,4 +28,13 @@
            return Response(code=500, msg="create failure", data={})
    else:
        return Response(code=500, msg="网络异常！failure", data={})
    return Response(code=200, msg="create success", data=data)


@mind_map_router.get("/{messageId}/parse", response_model=Response)
async def api_chat_mindmap(messageId: str, current_user: UserModel = Depends(get_current_user), db: Session = Depends(get_db)): #  current_user: UserModel = Depends(get_current_user)

    data = await service_message_mindmap_parse(db, messageId, current_user.id)
    if not data:
        return Response(code=500, msg="create failure", data={})
    return Response(code=200, msg="create success", data=data)

 app/models/dialog_model.py

@@ -1,3 +1,4 @@
import json
from datetime import datetime
from typing import Optional

@@ -24,6 +25,7 @@
    # agent_id = Column(String(36))
    mode = Column(String(36))
    parameters = Column(Text)
    kb_ids = Column(String(128))

    def get_id(self):
        return str(self.id)
@@ -43,6 +45,9 @@
            'mode': self.mode,
        }

    def get_kb_ids(self):
        return json.loads(self.kb_ids) if self.kb_ids else []


class ConversationModel(Base):
    __tablename__ = 'conversation'

 app/models/v2/chat.py

@@ -6,7 +6,7 @@
from sqlalchemy import Column, Integer, String, BigInteger, ForeignKey, DateTime, Text, TEXT
from sqlalchemy.orm import Session

from app.config.const import Dialog_STATSU_DELETE, Dialog_STATSU_ON
from app.config.const import Dialog_STATSU_DELETE, Dialog_STATSU_ON, complex_knowledge_chat
from app.models.base_model import Base
from app.utils.common import current_time

@@ -187,14 +187,20 @@
            query = {}
            if self.query:
                query = json.loads(self.query)
            return {

            res = {
                'id': self.id,
                'role': "assistant",
                'answer': self.content,
                'chat_mode': self.chat_mode,
                'node_list': json.loads(self.node_data) if self.node_data else [],
                "parentId": query.get("parentId")
                "parentId": query.get("parentId"),
                "isDeep": query.get("isDeep", 1),
            }
            if self.chat_mode == complex_knowledge_chat:
                res['reference'] = json.loads(self.node_data) if self.node_data else {}
            else:
                res['node_list'] = json.loads(self.node_data) if self.node_data else []
            return res


class ComplexChatSessionDao:

 app/service/dialog.py

@@ -245,6 +245,7 @@
            if app_dialog:
                dialog.name = app_dialog["name"]
                dialog.description = app_dialog["description"]
                dialog.kb_ids = app_dialog["kb_ids"]
                dialog.update_date = datetime.now()
                db.add(dialog)
                db.commit()

 app/service/knowledge.py

@@ -17,8 +17,8 @@
        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)))

    if location:
        query = query.filter(or_(KnowledgeModel.permission == "team", KnowledgeModel.tenant_id == str(user_id)))
        if location:
            query = query.filter(or_(KnowledgeModel.permission == "team", KnowledgeModel.tenant_id == str(user_id)))

    if keyword:
        query = query.filter(KnowledgeModel.name.like('%{}%'.format(keyword)))

 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 = ""

 app/service/v2/mindmap.py

@@ -1,10 +1,11 @@
import json
from Log import logger
from app.config.agent_base_url import DF_CHAT_AGENT
from app.config.agent_base_url import DF_CHAT_AGENT, RG_CHAT_DIALOG
from app.config.config import settings
from app.config.const import message_error, message_event, complex_knowledge_chat
from app.config.const import message_error, message_event, complex_knowledge_chat, rg_api_token, workflow_finished
from app.models import ComplexChatSessionDao, ChatData
from app.service.v2.app_driver.chat_agent import ChatAgent
from app.service.v2.app_driver.chat_dialog import ChatDialog
from app.service.v2.chat import get_chat_token


@@ -77,23 +78,41 @@
        if session.mindmap:
            inputs = {"is_deep": chat_request.get("isDeep", 1)}
            if session.chat_mode == complex_knowledge_chat:
                inputs["query_json"] = json.dumps(
                    {"query": chat_request.get("query", ""), "dataset_ids": chat_request.get("knowledgeId", [])})
            try:
                async for ans in chat.chat_completions(url,
                                                       await chat.complex_request_data(message, session.conversation_id,
                                                                               str(user_id), files=chat_request.get("files", []), inputs=inputs),
                                                       await chat.get_headers(token)):
                    if ans.get("event") == message_error:
                        return res
                    elif ans.get("event") == message_event:
                        mindmap_query += ans.get("answer", "")
                    else:
                        continue
                token = await get_chat_token(db, rg_api_token)
                # print(token)
                dialog_url = settings.fwr_base_url + RG_CHAT_DIALOG.format(session.chat_id)
                dialog_chat = ChatDialog()
                try:
                    async for ans in dialog_chat.chat_completions(dialog_url, await dialog_chat.complex_request_data(f"简要总结：{message}",
                                                                                                chat_request["knowledgeId"],
                                                                                                session.conversation_id),
                                                           await dialog_chat.get_headers(token)):
                        if ans.get("code", None) == 102:
                            return res
                        else:
                            if isinstance(ans.get("data"), bool) and ans.get("data") is True:
                                break
                            else:
                                data = ans.get("data", {})
                                mindmap_query = data.get("answer", "")
                except Exception as e:
                    logger.error(e)
            else:
                try:
                    async for ans in chat.chat_completions(url,
                                                           await chat.complex_request_data(message, session.conversation_id,
                                                                                   str(user_id), files=chat_request.get("files", []), inputs=inputs),
                                                           await chat.get_headers(token)):
                        if ans.get("event") == message_error:
                            return res
                        elif ans.get("event") == workflow_finished:
                            mindmap_query = ans.get("data", {}).get("outputs", {}).get("answer", "")
                        else:
                            continue

            except Exception as e:
                logger.error(e)
                return res
                except Exception as e:
                    logger.error(e)
                    return res
        else:
            mindmap_query = session.content
        # print("-----------------", mindmap_query)
@@ -107,6 +126,7 @@
                                                   await chat.complex_request_data(mindmap_query, "",
                                                                           str(user_id)),
                                                   await chat.get_headers(token)):
                # print(ans)
                if ans.get("event") == message_error:
                    return res
                elif ans.get("event") == message_event:
@@ -195,6 +215,19 @@
    return parent_list[:index]+new_node_list+parent_list[index+1:]


async def service_message_mindmap_parse(db, message_id, user_id):
    res = {}
    complex_log = ComplexChatSessionDao(db)
    session = await complex_log.get_session_by_id(message_id)

    if session.mindmap:
        try:
            res_str = await mindmap_join_str(session.mindmap)
            res["mindmap"] = res_str
        except Exception as e:
            logger.error(e)
    return res


if __name__ == '__main__':
    a = '{  "title": "全生命周期管理",  "items": [    {      "title": "设备规划与采购",      "items": [        {          "title": "需求分析与选型"    ,"items": [{"title": "rererer"}, {"title": "trtrtrtrt"}]    },        {          "title": "供应商选择与合同管理"        }      ]    },    {      "title": "设备安装与调试",      "items": [        {          "title": "安装规范"        },        {          "title": "调试测试"        }      ]    },    {      "title": "设备使用",      "items": [        {          "title": "操作培训"        },        {          "title": "操作规程与记录"        }      ]    },    {      "title": "设备维护与维修",      "items": [        {          "title": "定期维护"        },        {          "title": "故障诊断"        },        {          "title": "备件管理"        }      ]    },    {      "title": "设备更新与改造",      "items": [        {          "title": "技术评估"        },        {          "title": "更新计划"        },        {          "title": "改造方案"        }      ]    },    {      "title": "设备报废",      "items": [        {          "title": "报废评估"        },        {          "title": "报废处理"        }      ]    },    {      "title": "信息化管理",      "items": [        {          "title": "设备管理系统"        },        {          "title": "数据分析"        },        {          "title": "远程监控"        }      ]    },    {      "title": "安全管理",      "items": [        {          "title": "安全培训"        },        {          "title": "安全检查"        },        {          "title": "应急预案"        }      ]    },    {      "title": "环境保护",      "items": [        {          "title": "环保设备"        },        {          "title": "废物处理"        },        {          "title": "节能减排"        }      ]    },    {      "title": "具体实践案例",      "items": [        {          "title": "高压开关设备润滑脂选用研究"        },        {          "title": "环保型 C4 混气 GIS 设备运维技术研究"        }      ]    },    {      "title": "总结",      "items": [        {          "title": "提高运营效率和竞争力"        }      ]    }  ]}'

 app/task/fetch_agent.py

@@ -43,6 +43,7 @@
    status = Column(String(1), nullable=False)
    description = Column(String(255), nullable=False)
    tenant_id = Column(String(36), nullable=False)
    kb_ids = Column(String(128), nullable=False)


class DfApps(Base):
@@ -257,13 +258,13 @@
            query = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id) \
                .filter(Dialog.name.in_(names), Dialog.status == "1")
        else:
            query = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id).filter(
            query = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id, Dialog.kb_ids).filter(
                Dialog.status == "1", Dialog.tenant_id == tenant_id)

        results = query.all()
        formatted_results = [
            {"id": row[0], "name": row[1], "description": row[2], "status": "1" if row[3] == "1" else "2",
             "user_id": str(row[4]), "mode": "agent-dialog", "parameters": para} for row in results if row[0] not in chat_ids]
             "user_id": str(row[4]), "mode": "agent-dialog", "parameters": para, "kb_ids": row[5]} for row in results if row[0] not in chat_ids]
        return formatted_results
    finally:
        db.close()
@@ -301,13 +302,14 @@
                existing_agent.name = row["name"]
                existing_agent.description = row["description"]
                existing_agent.mode = row["mode"]
                existing_agent.kb_ids = row.get("kb_ids", "")
                if existing_agent.status == Dialog_STATSU_DELETE:
                    existing_agent.status = Dialog_STATSU_ON
                if row["parameters"]:
                    existing_agent.parameters = json.dumps(row["parameters"])
            else:
                existing = DialogModel(id=row["id"], status=row["status"], name=row["name"],
                                       description=row["description"],
                                       description=row["description"], kb_ids=row.get("kb_ids", ""),
                                       tenant_id=get_rag_user_id(db, row["user_id"], type_dict[dialog_type]),
                                       dialog_type=dialog_type, mode=row["mode"], parameters=json.dumps(row["parameters"]))
                db.add(existing)
@@ -411,10 +413,10 @@
def get_one_from_ragflow_dialog(dialog_id):
    db = SessionRagflow()
    try:
        row = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id) \
        row = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id, Dialog.kb_ids) \
            .filter(Dialog.id==dialog_id).first()
        return {"id": row[0], "name": row[1], "description": row[2], "status": str(row[3]),
                "user_id": str(row[4])} if row else {}
                "user_id": str(row[4]), "kb_ids": row[5]} if row else {}
    finally:
        db.close()