新增用户增加rg数据同步

zhaoqingang

2024-12-26 5eff7ec9d00fa835a76c354bad7cc7f284d28b11

 app/task/fetch_agent.py

@@ -62,6 +62,15 @@
    status = Column(String(1))  # 状态
    doc_num = Column(Integer)  # 文档


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))  # 创建人id


# 解析名字
def parse_names(names_str: str) -> List[str]:
    return [name.strip() for name in names_str.split(',')]
@@ -228,9 +237,9 @@
    try:
        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).filter( Dialog.status == "1")

        results = query.all()
        formatted_results = [
@@ -271,7 +280,7 @@
            if existing_agent:
                existing_agent.name = row["name"]
                existing_agent.description = row["description"]
                existing_agent.status = row["status"]
                # existing_agent.status = row["status"]
                existing_agent.mode = row["mode"]
                # existing_agent.tenant_id = get_rag_user_id(db, row["user_id"], type_dict[dialog_type])
            else:
@@ -324,6 +333,8 @@
        print("v2 Agents synchronized successfully")
    except Exception as 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}
@@ -386,8 +397,24 @@
        print("sync knowledge successfully")
    except Exception as 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()


if __name__ == "__main__":
    a = get_data_from_dify_v2([])
    print(a)
    # a = get_data_from_dify_v2([])
    # print(a)
    update_ragflow_user_tenant("")