知识库和智能体同步

zhaoqingang

2024-12-18 c941b948f1fa8eff615985a83110a401ae426ab3

 app/task/fetch_agent.py

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

from app.config.config import settings
from app.config.const import RAGFLOW, BISHENG, DIFY
from app.models import KnowledgeModel
from app.models.dialog_model import DialogModel
from app.models.user_model import UserAppModel
from app.models.agent_model import AgentModel
@@ -49,6 +50,16 @@
    description = Column(Text, nullable=False)
    tenant_id = Column(String(36), nullable=False)


class RgKnowledge(Base):
    __tablename__ = 'knowledgebase'
    id = Column(String(36), primary_key=True)  # id
    name = Column(String(128))  # 名称
    permission = Column(String(32), default="me")
    tenant_id = Column(String(32))  # 创建人id
    description = Column(Text)  # 说明
    status = Column(String(1))  # 状态
    doc_num = Column(Integer)  # 文档

# 解析名字
def parse_names(names_str: str) -> List[str]:
@@ -199,14 +210,14 @@
    try:
        if names:
            query = db.query(Flow.id, Flow.name, Flow.description, Flow.status, Flow.user_id) \
                .filter(Flow.name.in_(names), Flow.status==2)
                .filter(Flow.name.in_(names))
        else:
            query = db.query(Flow.id, Flow.name, Flow.description, Flow.status, Flow.user_id).filter(Flow.status==2)
            query = db.query(Flow.id, Flow.name, Flow.description, Flow.status, Flow.user_id)

        results = query.all()
        # print(f"Executing query: {query}")
        # 格式化id为UUID
        formatted_results = [{"id":format_uuid(row[0]), "name": row[1], "description": row[2], "status": str(row[3]-1), "user_id": str(row[4])} for row in results]
        formatted_results = [{"id":format_uuid(row[0]), "name": row[1], "description": row[2], "status": "1" if row[3] ==2 else "0", "user_id": str(row[4])} for row in results]
        return formatted_results
    finally:
        db.close()
@@ -222,7 +233,7 @@

        results = query.all()
        formatted_results = [
            {"id": format_uuid(row[0]), "name": row[1], "description": row[2], "status": str(row[3]),
            {"id": format_uuid(row[0]), "name": row[1], "description": row[2], "status": str(row[3]) if row[3]  ==1 else "2",
             "user_id": str(row[4])} for row in results]
        return formatted_results
    finally:
@@ -258,15 +269,16 @@
            existing_agent = db.query(DialogModel).filter_by(id=row["id"]).first()
            if existing_agent:
                existing_agent.name = row["name"]
                existing_agent.status = row["status"]
                existing_agent.description = row["description"]
                existing_agent.tenant_id = get_rag_user_id(db, row["user_id"], type_dict[dialog_type])
                # existing_agent.tenant_id = get_rag_user_id(db, row["user_id"], type_dict[dialog_type])
            else:
                existing = DialogModel(id=row["id"], name=row["name"], description=row["description"], tenant_id=get_rag_user_id(db, row["user_id"], type_dict[dialog_type]), dialog_type=dialog_type)
                existing = DialogModel(id=row["id"], status=row["status"], name=row["name"], description=row["description"], tenant_id=get_rag_user_id(db, row["user_id"], type_dict[dialog_type]), dialog_type=dialog_type)
                db.add(existing)
        db.commit()
        for dialog in db.query(DialogModel).filter_by(dialog_type=dialog_type).all():
            if dialog.id not in agent_id_list:
                db.query(DialogModel).filter_by(id=dialog.id).delete()
                db.query(DialogModel).filter_by(id=dialog.id).update({"status": "2"})
                db.commit()
    except IntegrityError:
        db.rollback()
@@ -280,10 +292,10 @@
    db = SessionRagflow()
    try:

        results = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id).all()
        results = db.query(RgKnowledge.id, RgKnowledge.name, RgKnowledge.description, RgKnowledge.status, RgKnowledge.tenant_id, RgKnowledge.doc_num, RgKnowledge.permission).all()
        formatted_results = [
            {"id": format_uuid(row[0]), "name": row[1], "description": row[2], "status": str(row[3]),
             "user_id": str(row[4])} for row in results]
            {"id": row[0], "name": row[1], "description": row[2], "status": str(row[3]),
             "user_id": str(row[4]), "doc_num": row[5], "permission": row[6]} for row in results]
        return formatted_results
    finally:
        db.close()
@@ -296,17 +308,60 @@
        for app in app_register:
            if app["id"] == RAGFLOW:
                ragflow_data = get_data_from_ragflow_v2([])
                update_ids_in_local_v2(ragflow_data, "1")
                if ragflow_data:
                    update_ids_in_local_v2(ragflow_data, "1")
            elif app["id"] == BISHENG:
                bisheng_data = get_data_from_bisheng_v2([])
                update_ids_in_local_v2(bisheng_data, "2")
                if bisheng_data:
                    update_ids_in_local_v2(bisheng_data, "2")
            elif app["id"] == DIFY:
                dify_data = get_data_from_dify_v2([])
                update_ids_in_local_v2(dify_data, "4")
                if dify_data:
                    update_ids_in_local_v2(dify_data, "4")
        print("Agents synchronized successfully")
    except Exception as e:
        print(f"Failed to sync agents: {str(e)}")

def update_ids_in_local_knowledge(data, klg_type):
    type_dict = {"1": RAGFLOW, "2": BISHENG, "4": DIFY}
    db = SessionLocal()
    agent_id_list = []
    try:
        for row in data:
            agent_id_list.append(row["id"])
            existing_agent = db.query(KnowledgeModel).filter_by(id=row["id"]).first()
            if existing_agent:
                existing_agent.name = row["name"]
                existing_agent.description = row["description"]
                # existing_agent.tenant_id = get_rag_user_id(db, row["user_id"], type_dict[klg_type])
                existing_agent.permission =  row["permission"]
                existing_agent.documents =  row["doc_num"]
                existing_agent.status =  row["status"]
            else:
                existing = KnowledgeModel(id=row["id"], name=row["name"], description=row["description"],
                                       tenant_id=get_rag_user_id(db, row["user_id"], type_dict[klg_type]),status=row["status"],
                                       knowledge_type=1, permission=row["permission"], documents=row["doc_num"])
                db.add(existing)
        db.commit()
        for dialog in db.query(KnowledgeModel).filter_by(knowledge_type=type_dict[klg_type]).all():
            if dialog.id not in agent_id_list:
                db.query(KnowledgeModel).filter_by(id=dialog.id).delete()
                db.commit()
    except IntegrityError:
        db.rollback()
        raise
    finally:
        db.close()

def get_one_from_ragflow_knowledge(klg_id):
    db = SessionRagflow()
    try:

        row = db.query(RgKnowledge.id, RgKnowledge.name, RgKnowledge.description, RgKnowledge.status, RgKnowledge.tenant_id, RgKnowledge.doc_num, RgKnowledge.permission).filter(RgKnowledge.id==klg_id).first()
        return {"id": row[0], "name": row[1], "description": row[2], "status": str(row[3]),
             "user_id": str(row[4]), "doc_num": row[5], "permission": row[6]} if row else {}
    finally:
        db.close()


def sync_knowledge():
@@ -316,8 +371,9 @@
        app_register = AppRegisterDao(db).get_apps()
        for app in app_register:
            if app["id"] == RAGFLOW:
                ragflow_data = get_data_from_ragflow_knowledge([])
                update_ids_in_local_v2(ragflow_data, "1")
                ragflow_data = get_data_from_ragflow_knowledge()
                if ragflow_data:
                    update_ids_in_local_knowledge(ragflow_data, "1")
            # elif app["id"] == BISHENG:
            #     bisheng_data = get_data_from_bisheng_v2([])
            #     update_ids_in_local_v2(bisheng_data, "2")