Explorar o código

Merge remote-tracking branch 'origin/master'

xuyonghao hai 1 ano
pai
achega
40408073a2

+ 65 - 0
alembic/versions/75d247bc0248_knowledgebases_tables_add.py

@@ -0,0 +1,65 @@
+"""knowledgebases tables add 
+
+Revision ID: 75d247bc0248
+Revises: 16efd20cce44
+Create Date: 2024-12-17 19:38:44.939447
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision: str = '75d247bc0248'
+down_revision: Union[str, None] = '16efd20cce44'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('knowledgebases',
+    sa.Column('id', sa.String(length=36), nullable=False),
+    sa.Column('name', sa.String(length=128), nullable=True),
+    sa.Column('create_date', sa.DateTime(), nullable=True),
+    sa.Column('update_date', sa.DateTime(), nullable=True),
+    sa.Column('permission', sa.String(length=32), nullable=True),
+    sa.Column('tenant_id', sa.String(length=32), nullable=True),
+    sa.Column('description', sa.Text(), nullable=True),
+    sa.Column('status', sa.String(length=1), nullable=True),
+    sa.Column('documents', sa.Integer(), nullable=True),
+    sa.Column('icon', sa.String(length=128), nullable=True),
+    sa.Column('knowledge_type', sa.String(length=1), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.drop_table('knowledgebase')
+    op.drop_constraint('group_Knowledge_ibfk_2', 'group_Knowledge', type_='foreignkey')
+    op.create_foreign_key(None, 'group_Knowledge', 'knowledgebases', ['knowledge_id'], ['id'], ondelete='CASCADE')
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, 'group_Knowledge', type_='foreignkey')
+    op.create_foreign_key('group_Knowledge_ibfk_2', 'group_Knowledge', 'knowledgebase', ['knowledge_id'], ['id'], ondelete='CASCADE')
+    op.create_table('knowledgebase',
+    sa.Column('id', mysql.VARCHAR(length=36), nullable=False),
+    sa.Column('name', mysql.VARCHAR(length=128), nullable=True),
+    sa.Column('create_date', mysql.DATETIME(), nullable=True),
+    sa.Column('update_date', mysql.DATETIME(), nullable=True),
+    sa.Column('tenant_id', mysql.VARCHAR(length=32), nullable=True),
+    sa.Column('description', mysql.TEXT(), nullable=True),
+    sa.Column('status', mysql.VARCHAR(length=1), nullable=True),
+    sa.Column('documents', mysql.INTEGER(), autoincrement=False, nullable=True),
+    sa.Column('icon', mysql.VARCHAR(length=128), nullable=True),
+    sa.Column('knowledge_type', mysql.VARCHAR(length=1), nullable=True),
+    sa.Column('permission', mysql.VARCHAR(length=32), nullable=True),
+    sa.PrimaryKeyConstraint('id'),
+    mysql_collate='utf8mb4_0900_ai_ci',
+    mysql_default_charset='utf8mb4',
+    mysql_engine='InnoDB'
+    )
+    op.drop_table('knowledgebases')
+    # ### end Alembic commands ###

+ 1 - 1
app/api/chat.py

@@ -622,7 +622,7 @@ async def handle_client(websocket: WebSocket,
                                     elif data.get("event") == "message":  # "event": "message_end"
                                          # 正常输出
                                         answer = data.get("answer", "")
-                                        result = {"message": answer, "type": "message"}
+                                        result = {"message": answer, "type": "stream"}
                                     elif data.get("event") == "workflow_finished":
                                         answer = data.get("data", "")
                                         if isinstance(answer, str):

+ 11 - 2
app/api/dialog.py

@@ -5,7 +5,8 @@ from app.api import Response, get_current_user, ResponseList
 from app.models.dialog_model import dialogData, dialogDataUpdate
 from app.models.base_model import get_db
 from app.models.user_model import UserModel
-from app.service.dialog import get_dialog_list, create_dialog_service, update_dialog_status_service
+from app.service.dialog import get_dialog_list, create_dialog_service, update_dialog_status_service, \
+    delete_dialog_service
 
 dialog_router = APIRouter()
 
@@ -46,4 +47,12 @@ async def change_dialog_status(dialog: dialogDataUpdate, current_user: UserModel
     is_create = await update_dialog_status_service(db, dialog.id, dialog.status)
     if not is_create:
         return Response(code=500, msg="dialog update failure", data={})
-    return Response(code=200, msg="dialog update success", data={})
+    return Response(code=200, msg="dialog update success", data={})
+
+
+@dialog_router.get("/delete", response_model=Response)
+async def sync_knowledge_api(dialogId: str, current_user: UserModel = Depends(get_current_user), db=Depends(get_db)):
+    is_delete = await delete_dialog_service(db, dialogId)
+    if not is_delete:
+        return Response(code=500, msg="knowledge delete failure", data={})
+    return Response(code=200, msg="knowledge delete success", data={})

+ 21 - 4
app/api/knowledge.py

@@ -5,7 +5,8 @@ from app.api import Response, get_current_user, ResponseList
 from app.models import klgParameter
 from app.models.base_model import get_db
 from app.models.user_model import UserModel
-from app.service.knowledge import get_knowledge_list, create_knowledge_service
+from app.service.knowledge import get_knowledge_list, create_knowledge_service, sync_knowledge_service, \
+    delete_knowledge_service
 from typing import Optional
 knowledge_router = APIRouter()
 
@@ -30,8 +31,8 @@ async def knowledge_list(
 async def create_knowledge_api(klg: klgParameter, current_user: UserModel = Depends(get_current_user), db=Depends(get_db)):
     is_create = await create_knowledge_service(db, klg.id, klg.name, klg.description, klg.icon, klg.klgType, current_user.id)
     if not is_create:
-        return Response(code=500, msg="role create failure", data={})
-    return Response(code=200, msg="role create success", data={})
+        return Response(code=500, msg="role knowledge failure", data={})
+    return Response(code=200, msg="role knowledge success", data={})
 
 
 @knowledge_router.get("/update", response_model=Response)
@@ -39,4 +40,20 @@ async def change_knowledge_api(knowledgeId: str, current_user: UserModel = Depen
     # is_create = await update_dialog_status_service(db, dialog.id, dialog.status)
     # if not is_create:
     #     return Response(code=500, msg="dialog update failure", data={})
-    return Response(code=200, msg="dialog update success", data={})
+    return Response(code=200, msg="knowledge update success", data={})
+
+
+@knowledge_router.get("/sync", response_model=Response)
+async def sync_knowledge_api(knowledgeId: str, current_user: UserModel = Depends(get_current_user), db=Depends(get_db)):
+    is_create = await sync_knowledge_service(db, knowledgeId)
+    if not is_create:
+        return Response(code=500, msg="knowledge update failure", data={})
+    return Response(code=200, msg="knowledge update success", data={})
+
+
+@knowledge_router.get("/delete", response_model=Response)
+async def sync_knowledge_api(knowledgeId: str, current_user: UserModel = Depends(get_current_user), db=Depends(get_db)):
+    is_delete = await delete_knowledge_service(db, knowledgeId)
+    if not is_delete:
+        return Response(code=500, msg="knowledge delete failure", data={})
+    return Response(code=200, msg="knowledge delete success", data={})

+ 2 - 6
app/api/user.py

@@ -1,5 +1,6 @@
 from fastapi import APIRouter, Depends
 from app.api import Response, pwd_context, get_current_user, ResponseList
+from app.config.const import USER_STATSU_DELETE
 from app.models.base_model import get_db
 from app.models.user import PageParameter, UserStatus, UserInfo, LoginData, UserPassword, UserDept
 from app.models.user_model import UserModel
@@ -25,7 +26,7 @@ async def add_user(user: UserInfo, current_user: UserModel = Depends(get_current
     if user.pwd:
         if not is_valid_password(user.pwd):
             return Response(code=400, msg="The password must be at least 8 and contain both numbers and letters")
-    db_user = db.query(UserModel).filter(UserModel.username == user.userName).first()
+    db_user = db.query(UserModel).filter(UserModel.username == user.userName, UserModel.status != USER_STATSU_DELETE).first()
     if db_user:
         return Response(code=200, msg="user already created")
     pwd = user.pwd
@@ -40,14 +41,9 @@ async def add_user(user: UserInfo, current_user: UserModel = Depends(get_current
 
 @user_router.put("/edit_user", response_model=Response)
 async def edit_user(user: UserInfo, current_user: UserModel = Depends(get_current_user), db=Depends(get_db)):
-    # if not user.userName:
-    #     return Response(code=400, msg="The userName cannot be empty!")
     user_info = db.query(UserModel).filter(UserModel.id == user.userId).first()
     if not user_info:
         return Response(code=200, msg="user does not exist")
-    # db_user = db.query(UserModel).filter(UserModel.username == user.userName).first()
-    # if db_user and db_user.id != user.userId:
-    #     return Response(code=200, msg="user already created")
     is_edit = await edit_user_data(db, user.userId, user.email, user.phone, user.loginName, user.roles, user.groups)
     if not is_edit:
         return Response(code=500, msg="user edit failure", data={})

+ 8 - 1
app/config/const.py

@@ -15,4 +15,11 @@ DIFY = "dify_app"
 BASIC_ID = 3
 RAGFLOW_ID = 1
 BISHENG_ID = 2
-DIFY_ID = 4
+DIFY_ID = 4
+
+###  ------------------status--------------------
+USER_STATSU_DELETE = "2"
+USER_STATSU_ON = "1"
+USER_STATSU_OFF = "0"
+
+Dialog_STATSU_DELETE = "2"

+ 2 - 2
app/models/group_model.py

@@ -9,9 +9,9 @@ from sqlalchemy.orm import relationship, backref
 from app.models.base_model import Base
 
 
-group_knowledge_table = Table('group_Knowledge', Base.metadata
+group_knowledge_table = Table('group_knowledge', Base.metadata
                            , Column('group_id', Integer, ForeignKey('group.id', ondelete='CASCADE'))
-                           , Column('knowledge_id', String(32), ForeignKey('knowledgebase.id', ondelete='CASCADE')))
+                           , Column('knowledge_id', String(32), ForeignKey('knowledgebases.id', ondelete='CASCADE')))
 
 group_dialog_table = Table('group_dialogs', Base.metadata
                            , Column('group_id', Integer, ForeignKey('group.id', ondelete='CASCADE'))

+ 1 - 1
app/models/knowledge_model.py

@@ -9,7 +9,7 @@ from sqlalchemy.orm import relationship, backref
 from app.models.base_model import Base
 
 class KnowledgeModel(Base):
-    __tablename__ = 'knowledgebase'
+    __tablename__ = 'knowledgebases'
     __mapper_args__ = {
         # "order_by": 'SEQ'
     }

+ 3 - 1
app/models/resource_model.py

@@ -25,7 +25,7 @@ class ResourceModel(Base):
     resource_id = Column(String(36), ForeignKey('resource.id'))
     parent = relationship('ResourceModel', remote_side=[id], backref='resources', uselist=False)
     children = relationship('ResourceModel')
-    status = Column(String(10))
+    status = Column(String(10), default="1")
     hidden = Column(Integer, default=0)
     def get_id(self):
         return str(self.ID)
@@ -56,6 +56,8 @@ class ResourceModel(Base):
             'menuId': self.id,
             'menuName': self.name,
             'menuType': self.resource_type_id,
+            'parentId': self.get_pid(),
+            'parentName': self.get_pName(),
             'children': [res.to_tree_select_json() for res in self.children]
         }
 

+ 2 - 2
app/service/auth.py

@@ -11,7 +11,7 @@ from sqlalchemy.orm import Session
 
 from Log import logger
 from app.config.config import settings
-from app.config.const import RAGFLOW, BISHENG, DIFY
+from app.config.const import RAGFLOW, BISHENG, DIFY, USER_STATSU_DELETE
 from app.models import RoleModel, GroupModel, TokenModel
 from app.models.user_model import UserModel, UserAppModel
 # from app.service.ragflow import RagflowService
@@ -34,7 +34,7 @@ def get_password_hash(password):
 
 
 def authenticate_user(db, username: str, password: str):
-    user = db.query(UserModel).filter(UserModel.username == username).first()
+    user = db.query(UserModel).filter(UserModel.username == username, UserModel.status != USER_STATSU_DELETE).first()
     if not user:
         return False
     if not verify_password(password, user.hashed_password):

+ 13 - 1
app/service/dialog.py

@@ -2,6 +2,7 @@ from datetime import datetime
 
 from sqlalchemy import or_
 
+from app.config.const import Dialog_STATSU_DELETE
 from app.models import KnowledgeModel, GroupModel, DialogModel, ConversationModel, group_dialog_table, LabelWorkerModel, \
     LabelModel
 from app.models.user_model import UserModel
@@ -12,7 +13,7 @@ async def get_dialog_list(db, user_id, keyword, label, status, page_size, page_i
     user = db.query(UserModel).filter(UserModel.id == user_id).first()
     if user is None:
         return {"rows": []}
-    query = db.query(DialogModel)
+    query = db.query(DialogModel).filter(DialogModel.status != Dialog_STATSU_DELETE)
     id_list = []
     if label:
         id_list = [i.object_id for i in db.query(LabelWorkerModel).filter(LabelWorkerModel.label_id==label).all()]
@@ -111,4 +112,15 @@ async def update_dialog_status_service(db, dialog_id, status):
         logger.error(e)
         db.rollback()
         return False
+    return True
+
+
+async def delete_dialog_service(db, dialog_id):
+    try:
+        db.query(DialogModel).filter_by(id=dialog_id).update({"status":Dialog_STATSU_DELETE, "update_date": datetime.now()})
+        db.commit()
+    except Exception as e:
+        logger.error(e)
+        db.rollback()
+        return False
     return True

+ 35 - 1
app/service/knowledge.py

@@ -3,6 +3,7 @@ from sqlalchemy import or_
 from app.models import KnowledgeModel, group_knowledge_table
 from app.models.user_model import UserModel
 from Log import logger
+from app.task.fetch_agent import get_one_from_ragflow_knowledge
 
 
 async def get_knowledge_list(db, user_id, keyword, page_size, page_index, status, knowledge, location):
@@ -19,7 +20,7 @@ async def get_knowledge_list(db, user_id, keyword, page_size, page_index, status
         query = query.filter(KnowledgeModel.name.like('%{}%'.format(keyword)))
     total = query.count()
     if page_size:
-        query = query.limit(page_size).offset((page_index - 1) * page_size)
+        query = query.order_by(KnowledgeModel.update_date.desc()).limit(page_size).offset((page_index - 1) * page_size)
     return {"total":  total, "rows": [kld.to_json() for kld in query.all()]}
 
 
@@ -35,3 +36,36 @@ async def create_knowledge_service(db, klg_id, name, description, icon, klg_type
         return False
     return True
 
+
+async def sync_knowledge_service(db, klg_id):
+    knowledge = db.query(KnowledgeModel).filter(KnowledgeModel.id == klg_id).first()
+    if knowledge:
+        try:
+            app_knowledge = get_one_from_ragflow_knowledge(klg_id)
+            if app_knowledge:
+                knowledge.name = app_knowledge["name"]
+                knowledge.description = app_knowledge["description"]
+                knowledge.permission = app_knowledge["permission"]
+                knowledge.documents = app_knowledge["doc_num"]
+                knowledge.status = app_knowledge["status"]
+                db.add(knowledge)
+                db.commit()
+                db.refresh(knowledge)
+        except Exception as e:
+            logger.error(e)
+            db.rollback()
+            return False
+        return True
+
+
+async def delete_knowledge_service(db, klg_id):
+
+    try:
+        db.query(KnowledgeModel).filter(KnowledgeModel.id==klg_id).delete()
+        db.commit()
+    except Exception as e:
+        logger.error(e)
+        db.rollback()
+        return False
+    return True
+

+ 7 - 6
app/service/user.py

@@ -4,7 +4,7 @@ from datetime import datetime
 from app.api import pwd_context
 from app.api.dialog import dialog_list
 from app.config.config import settings
-from app.config.const import RAGFLOW, BISHENG, DIFY
+from app.config.const import RAGFLOW, BISHENG, DIFY, USER_STATSU_DELETE
 from app.models import RoleModel, GroupModel, AgentType, role_resource_table, DialogModel, OrganizationModel
 from app.models.menu_model import WebMenuModel, MenuCapacityModel
 from app.models.user_model import UserModel, UserAppModel
@@ -19,7 +19,7 @@ from app.utils.password_handle import generate_password
 
 
 async def get_user_list(db, page_index: int, page_size: int, keyword: str, role_key: str, user_id):
-    query = db.query(UserModel).filter(UserModel.permission != "admin")
+    query = db.query(UserModel).filter(UserModel.permission != "admin", UserModel.status != USER_STATSU_DELETE)
     # if role_key != "admin":
     #     query.filter(UserModel.creator==user_id)
     if keyword:
@@ -32,7 +32,7 @@ async def get_user_list(db, page_index: int, page_size: int, keyword: str, role_
 
 async def edit_user_status(db, status: str, user_id: int):
     try:
-        db.query(UserModel).filter(UserModel.id == user_id).update({"status": status})
+        db.query(UserModel).filter(UserModel.id == user_id, UserModel.status != USER_STATSU_DELETE).update({"status": status})
         db.commit()
     except Exception as e:
         logger.error(e)
@@ -43,7 +43,7 @@ async def edit_user_status(db, status: str, user_id: int):
 
 async def delete_user_data(db, user_id: str):
     try:
-        db.query(UserModel).filter(UserModel.id == user_id, UserModel.permission != "admin").delete()
+        db.query(UserModel).filter(UserModel.id == user_id, UserModel.permission != "admin").update({"status": USER_STATSU_DELETE})
         db.commit()
     except Exception as e:
         logger.error(e)
@@ -165,7 +165,7 @@ async def get_user_info(db, user_id):
     dept = []
     role_set = set()
     dept_set = set()
-    user = db.query(UserModel).filter_by(id=user_id).first()
+    user = db.query(UserModel).filter(UserModel.id==user_id,UserModel.status !=USER_STATSU_DELETE).first()
     await role_resource(role_set, roles, permissions, user.roles)
     for ogt in user.organizations:
         dept.append(ogt.to_json())
@@ -241,11 +241,12 @@ async def get_user_menus(db, user_id):
     dialog_list = []
     menu_dict = {}
     res = []
+    # print(user_id)
     user = db.query(UserModel).filter_by(id=user_id).first()
     for group in user.groups:
         for dialog in group.dialogs:
             dialog_list.append(dialog.id)
-    print(dialog_list)
+    # print(dialog_list)
     menu_list = db.query(WebMenuModel.id, WebMenuModel.title, WebMenuModel.describe, WebMenuModel.icon, WebMenuModel.desc,WebMenuModel.rank,
                          WebMenuModel.img, MenuCapacityModel.capacity_id, MenuCapacityModel.capacity_type, MenuCapacityModel.chat_id.label("agentId")).outerjoin(
         MenuCapacityModel, WebMenuModel.id == MenuCapacityModel.menu_id).outerjoin(

+ 71 - 15
app/task/fetch_agent.py

@@ -7,6 +7,7 @@ from sqlalchemy.orm import sessionmaker
 
 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
@@ -50,6 +51,16 @@ class DfApps(Base):
     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]:
     return [name.strip() for name in names_str.split(',')]
@@ -199,14 +210,14 @@ def get_data_from_bisheng_v2(names: List[str]) -> List[Dict]:
     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 @@ def get_data_from_ragflow_v2(names: List[str]) -> List[Dict]:
 
         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 @@ def update_ids_in_local_v2(data: List[Dict], dialog_type:str):
             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 @@ def get_data_from_ragflow_knowledge():
     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 @@ def sync_agents_v2():
         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 @@ def sync_knowledge():
         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")

+ 7 - 5
main.py

@@ -1,5 +1,5 @@
 from contextlib import asynccontextmanager
-# from apscheduler.schedulers.background import BackgroundScheduler
+from apscheduler.schedulers.background import BackgroundScheduler
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.staticfiles import StaticFiles
@@ -22,7 +22,8 @@ from app.api.resource import menu_router
 from app.api.user import user_router
 from app.api.group import group_router
 from app.api.role import role_router
-from app.task.fetch_agent import sync_agents, initialize_agents, sync_agents_v2
+# from app.models.base_model import init_db
+from app.task.fetch_agent import sync_agents, initialize_agents, sync_agents_v2, sync_knowledge
 from app.init_config.init_run_data import sync_default_data
 
 
@@ -36,6 +37,7 @@ async def lifespan(app: FastAPI):
     await sync_default_data()
 
     sync_agents_v2()
+    sync_knowledge()
     yield
     # 在应用关闭时执行清理操作(如果需要)
     pass
@@ -57,9 +59,9 @@ app = FastAPI(
 #     allow_headers=["*"],  # 允许所有头部
 # )
 # 创建调度器
-# scheduler = BackgroundScheduler()
-# scheduler.add_job(sync_resource, 'interval', minutes=1, id="sync_resource_data")
-# scheduler.start()
+scheduler = BackgroundScheduler()
+scheduler.add_job(sync_agents_v2, 'interval', minutes=60, id="sync_resource_data")
+scheduler.start()
 
 
 app.include_router(auth_router, prefix='/api/auth', tags=["auth"])