From e078028f1a8da34f3cc2cb5095b8e103a996f553 Mon Sep 17 00:00:00 2001 From: zhaoqingang <zhaoqg0118@163.com> Date: 星期五, 11 四月 2025 10:16:42 +0800 Subject: [PATCH] 首页智能体中心过滤掉自研类型 --- app/task/fetch_agent.py | 258 ++++++++++++++++++++++++++++++++++++++++----------- 1 files changed, 202 insertions(+), 56 deletions(-) diff --git a/app/task/fetch_agent.py b/app/task/fetch_agent.py index ea9a2e6..295b7b0 100644 --- a/app/task/fetch_agent.py +++ b/app/task/fetch_agent.py @@ -1,17 +1,20 @@ +import json +import os from pickle import PROTO from typing import Dict, List, Tuple from sqlalchemy import create_engine, Column, String, Integer, Text from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import sessionmaker, Session from app.config.config import settings -from app.config.const import RAGFLOW, BISHENG, DIFY -from app.models import KnowledgeModel +from app.config.const import RAGFLOW, BISHENG, DIFY, ENV_CONF_PATH, Dialog_STATSU_DELETE, Dialog_STATSU_ON +from app.models import KnowledgeModel, ComplexChatDao from app.models.dialog_model import DialogModel from app.models.user_model import UserAppModel from app.models.agent_model import AgentModel from app.models.base_model import SessionLocal, Base +from app.models.resource_model import ResourceModel, ResourceTypeModel from app.service.v2.app_register import AppRegisterDao # 鍒涘缓鏁版嵁搴撳紩鎿庡拰浼氳瘽宸ュ巶 @@ -40,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): @@ -49,6 +53,7 @@ status = Column(String(16), nullable=False) description = Column(Text, nullable=False) tenant_id = Column(String(36), nullable=False) + mode = Column(String(36), nullable=False) class RgKnowledge(Base): @@ -60,6 +65,16 @@ description = Column(Text) # 璇存槑 status = Column(String(1)) # 鐘舵�� doc_num = Column(Integer) # 鏂囨。 + embd_id = Column(String(128)) # 鏂囨。 + + +class RgUserTenant(Base): + __tablename__ = 'user_tenant' + id = Column(String(36), primary_key=True) # id + tenant_id = Column(String(32)) # 鍚嶇О + user_id = Column(String(32)) + role = Column(String(32)) # 鍒涘缓浜篿d + # 瑙f瀽鍚嶅瓧 def parse_names(names_str: str) -> List[str]: @@ -150,7 +165,7 @@ ('basic_excel_talk', 6, '鏅鸿兘鏁版嵁', 'BASIC', 'excelTalk'), ('basic_question_talk', 7, '鍑洪缁勫嵎', 'BASIC', 'questionTalk'), ('9d75142a-66eb-4e23-b7d4-03efe4584915', 8, '灏忔暟缁樺浘', 'DIFY', 'imageTalk'), - ('basic_paper_talk', 9, '鏂囨。鍑哄嵎', 'BASIC', 'paperTalk'), + ('2f6ddf93-7ba6-4b2d-b991-d96421404600', 9, '鏂囨。鍑哄嵎', 'DIFY', 'paperTalk'), ('basic_report_clean', 10, '鏂囨。鎶ュ憡', 'DIFY', 'reportWorkflow') ] @@ -197,9 +212,9 @@ finally: db.close() -def get_rag_user_id(db, tenant_id, app_type): - user = db.query(UserAppModel).filter(UserAppModel.app_type==app_type, UserAppModel.app_id==tenant_id).first() +def get_rag_user_id(db, tenant_id, app_type): + user = db.query(UserAppModel).filter(UserAppModel.app_type == app_type, UserAppModel.app_id == tenant_id).first() if user: return user.user_id return tenant_id @@ -210,75 +225,99 @@ try: if names: query = db.query(Flow.id, Flow.name, Flow.description, Flow.status, Flow.user_id) \ - .filter(Flow.name.in_(names)) + .filter(Flow.name.in_(names), Flow.status == "1") else: - query = db.query(Flow.id, Flow.name, Flow.description, Flow.status, Flow.user_id) + query = db.query(Flow.id, Flow.name, Flow.description, Flow.status, Flow.user_id).filter(Flow.status == "1") results = query.all() # print(f"Executing query: {query}") # 鏍煎紡鍖杋d涓篣UID - 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] + formatted_results = [ + {"id": row[0], "name": row[1], "description": row[2], "status": row[3], "user_id": str(row[4]), + "mode": "agent-dialog"} for row in results] return formatted_results finally: db.close() -def get_data_from_ragflow_v2(names: List[str]) -> List[Dict]: + +def get_data_from_ragflow_v2(base_db, names: List[str], tenant_id) -> List[Dict]: db = SessionRagflow() + para = { + "user_input_form": [], + "retriever_resource": { + "enabled": True + }, + "file_upload": { + "enabled": False + } + } try: + chat_ids = ComplexChatDao(base_db).get_complex_chat_ids() + # print(chat_ids) if names: query = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id) \ - .filter( Dialog.name.in_(names)) + .filter(Dialog.name.in_(names), Dialog.status == "1") else: - query = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id) + 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": 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] + {"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, "kb_ids": row[5]} for row in results if row[0] not in chat_ids] return formatted_results finally: db.close() -def get_data_from_dify_v2(names: List[str]) -> List[Dict]: +def get_data_from_dy_v2(base_db, names: List[str]) -> List[Dict]: db = SessionDify() try: + chat_ids = ComplexChatDao(base_db).get_complex_chat_ids() + # print(chat_ids) if names: - query = db.query(DfApps.id, DfApps.name, DfApps.description, DfApps.status, DfApps.tenant_id) \ - .filter( DfApps.name.in_(names)) + query = db.query(DfApps.id, DfApps.name, DfApps.description, DfApps.status, DfApps.tenant_id, DfApps.mode) \ + .filter(DfApps.name.in_(names)) else: - query = db.query(DfApps.id, DfApps.name, DfApps.description, DfApps.status, DfApps.tenant_id) + query = db.query(DfApps.id, DfApps.name, DfApps.description, DfApps.status, DfApps.tenant_id, DfApps.mode) results = query.all() formatted_results = [ {"id": str(row[0]), "name": row[1], "description": row[2], "status": "1", - "user_id": str(row[4])} for row in results] + "user_id": str(row[4]), "mode": row[5], "parameters": {}} for row in results if str(row[0]) not in chat_ids] return formatted_results finally: db.close() - -def update_ids_in_local_v2(data: List[Dict], dialog_type:str): +def update_ids_in_local_v2(data: List[Dict], dialog_type: str): db = SessionLocal() agent_id_list = [] - type_dict = {"1": RAGFLOW,"2": BISHENG,"4": DIFY} + type_dict = {"1": RAGFLOW, "2": BISHENG, "4": DIFY} try: for row in data: agent_id_list.append(row["id"]) 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.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"], 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"], 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) 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).update({"status": "2"}) + # print(dialog.id) + db.query(DialogModel).filter_by(id=dialog.id).update({"status": Dialog_STATSU_DELETE}) db.commit() except IntegrityError: db.rollback() @@ -287,12 +326,12 @@ db.close() - -def get_data_from_ragflow_knowledge(): +def get_data_from_ragflow_knowledge(tenant_id): db = SessionRagflow() try: - results = db.query(RgKnowledge.id, RgKnowledge.name, RgKnowledge.description, RgKnowledge.status, RgKnowledge.tenant_id, RgKnowledge.doc_num, RgKnowledge.permission).all() + results = db.query(RgKnowledge.id, RgKnowledge.name, RgKnowledge.description, RgKnowledge.status, + RgKnowledge.tenant_id, RgKnowledge.doc_num, RgKnowledge.permission).filter(RgKnowledge.tenant_id==tenant_id).all() formatted_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] @@ -300,27 +339,30 @@ finally: db.close() + def sync_agents_v2(): db = SessionLocal() try: app_register = AppRegisterDao(db).get_apps() for app in app_register: - if app["id"] == RAGFLOW: - ragflow_data = get_data_from_ragflow_v2([]) - if ragflow_data: - update_ids_in_local_v2(ragflow_data, "1") - elif app["id"] == BISHENG: - bisheng_data = get_data_from_bisheng_v2([]) - if bisheng_data: - update_ids_in_local_v2(bisheng_data, "2") - elif app["id"] == DIFY: - dify_data = get_data_from_dify_v2([]) - if dify_data: - update_ids_in_local_v2(dify_data, "4") - print("Agents synchronized successfully") + try: + if app["id"] == RAGFLOW: + ragflow_data = get_data_from_ragflow_v2(db, [], app["name"]) + if ragflow_data: + update_ids_in_local_v2(ragflow_data, "1") + elif app["id"] == DIFY: + dify_data = get_data_from_dy_v2(db, []) + if dify_data: + update_ids_in_local_v2(dify_data, "4") + except Exception as e: + print(f"Failed to sync agents: {str(e)}") + print("v2 Agents synchronized successfully") except Exception as e: - print(f"Failed to sync agents: {str(e)}") + print(f"v2 Failed to sync agents: {str(e)}") + finally: + db.close() + def update_ids_in_local_knowledge(data, klg_type): type_dict = {"1": RAGFLOW, "2": BISHENG, "4": DIFY} @@ -334,16 +376,17 @@ 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"] + 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"]) + 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(): + for dialog in db.query(KnowledgeModel).filter_by(knowledge_type=klg_type).all(): if dialog.id not in agent_id_list: db.query(KnowledgeModel).filter_by(id=dialog.id).delete() db.commit() @@ -353,13 +396,27 @@ 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() + row = db.query(RgKnowledge.id, RgKnowledge.name, RgKnowledge.description, RgKnowledge.status, + RgKnowledge.tenant_id, RgKnowledge.doc_num, RgKnowledge.permission, RgKnowledge.embd_id).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 {} + "user_id": str(row[4]), "doc_num": row[5], "permission": row[6], "embd_id": row[7]} if row else {} + finally: + db.close() + + +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, 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]), "kb_ids": row[5]} if row else {} finally: db.close() @@ -371,7 +428,7 @@ app_register = AppRegisterDao(db).get_apps() for app in app_register: if app["id"] == RAGFLOW: - ragflow_data = get_data_from_ragflow_knowledge() + ragflow_data = get_data_from_ragflow_knowledge(app["name"]) if ragflow_data: update_ids_in_local_knowledge(ragflow_data, "1") # elif app["id"] == BISHENG: @@ -380,11 +437,100 @@ # elif app["id"] == DIFY: # dify_data = get_data_from_dify_v2([]) # update_ids_in_local_v2(dify_data, "4") - print("Agents synchronized successfully") + print("sync knowledge successfully") except Exception as e: - print(f"Failed to sync agents: {str(e)}") + print(f"Failed to sync knowledge: {str(e)}") + finally: + db.close() + + +def update_ragflow_user_tenant(user_id: str): + db = SessionRagflow() + try: + if user_id: + db.query(RgUserTenant).filter(RgUserTenant.user_id == user_id, RgUserTenant.role == "invite").update( + {"role": "normal"}) + db.query(RgUserTenant).filter(RgUserTenant.tenant_id == user_id, RgUserTenant.role == "invite").update( + {"role": "normal"}) + else: + db.query(RgUserTenant).filter(RgUserTenant.role == "invite").update({"role": "normal"}) + db.commit() + finally: + db.close() + + +def import_type_table(session: Session, node: dict, parent=None): + resource_type = ResourceTypeModel( + id=node['id'], + name=node['name'], + description=node.get('description') + ) + if parent: + resource_type.parent = parent + session.add(resource_type) + session.commit() + + +def import_tree(session: Session, node: dict, parent=None): + resource = ResourceModel( + id=node['id'], + name=node['name'], + url=node['url'], + path=node.get('path'), + perms=node['perms'], + description=node.get('description'), + icon=node.get('icon'), + seq=node['seq'], + target=node.get('target'), + canbdeeleted=node.get('canbdeeleted'), + resource_type_id=node['resource_type_id'], + resource_id=node.get('resource_id'), + status=node['status'], + hidden=node.get('hidden') + ) + if parent: + resource.parent = parent + session.add(resource) + if 'children' in node: + for child in node['children']: + import_tree(session, child, parent=resource) + session.commit() + + +def sync_resources_from_json(): + db = SessionLocal() + try: + if db.query(ResourceTypeModel).count() == 0: + with open(os.path.join(ENV_CONF_PATH, "resource_type.json"), 'r', encoding='utf-8') as file: + type_json_data = json.load(file) + + db.query(ResourceTypeModel).delete() + db.commit() + + for node in type_json_data: + import_type_table(db, node) + print("add resourceType record successfully") + else: + print("sync resourcesType successfully") + if db.query(ResourceModel).count() == 0: + with open(os.path.join(ENV_CONF_PATH, "resource.json"), 'r', encoding='utf-8') as file: + json_data = json.load(file) + + db.query(ResourceModel).delete() + db.commit() + + for node in json_data: + import_tree(db, node) + print("add resources record successfully") + else: + print("sync resources successfully") + except Exception as e: + print(f"Failed to sync resources or resource type: {str(e)}") + finally: + db.close() if __name__ == "__main__": - a = get_data_from_dify_v2([]) - print(a) + # a = get_data_from_dify_v2([]) + # print(a) + update_ragflow_user_tenant("") -- Gitblit v1.8.0