From deaaf771005065d9951a6df4df3e8baade2acca2 Mon Sep 17 00:00:00 2001
From: zhaoqingang <zhaoqg0118@163.com>
Date: 星期五, 15 十一月 2024 15:02:18 +0800
Subject: [PATCH] 增加数据同步机制

---
 app/models/app_model.py    |   31 ++++++++++
 app/models/user_model.py   |    4 
 requirements.txt           |    0 
 app/config/config.py       |    1 
 app/api/sync_data.py       |   26 ++++++++
 main.py                    |   14 ++++
 app/config/config.yaml     |    3 
 app/api/__init__.py        |    3 +
 app/task/sync_resources.py |   40 +++++++++++++
 9 files changed, 119 insertions(+), 3 deletions(-)

diff --git a/app/api/__init__.py b/app/api/__init__.py
index a96baa1..2e317ce 100644
--- a/app/api/__init__.py
+++ b/app/api/__init__.py
@@ -2,6 +2,7 @@
 from urllib.parse import urlencode
 
 import jwt
+from cryptography.fernet import Fernet
 from fastapi import FastAPI, Depends, HTTPException
 from fastapi.security import OAuth2PasswordBearer
 from passlib.context import CryptContext
@@ -11,11 +12,13 @@
 
 from app.models.user_model import UserModel
 from app.service.auth import SECRET_KEY, ALGORITHM
+from app.config.config import settings
 
 app = FastAPI()
 
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+cipher_suite = Fernet(settings.HASH_SUB_KEY)
 
 
 class Response(BaseModel):
diff --git a/app/api/sync_data.py b/app/api/sync_data.py
new file mode 100644
index 0000000..30bfbe3
--- /dev/null
+++ b/app/api/sync_data.py
@@ -0,0 +1,26 @@
+# coding:utf-8
+
+from fastapi import APIRouter, Depends
+from app.api import Response, get_current_user
+from app.models.base_model import get_db
+from app.models.user_model import UserModel
+from app.task.sync_resources import sync_knowledge, sync_dialog, sync_agent, sync_resource, sync_llm
+
+sync_router = APIRouter()
+
+
+@sync_router.get("/data", response_model=Response)
+async def user_group_list(resource_type:int, current_user: UserModel = Depends(get_current_user),
+                          db=Depends(get_db)):
+    if resource_type == 1:
+        await sync_knowledge()
+    elif resource_type == 2:
+        await sync_dialog()
+    elif resource_type == 3:
+        await sync_agent()
+    elif resource_type == 4:
+        await sync_llm()
+    else:
+        await sync_resource()
+
+    return Response(code=200, msg="", data={})
\ No newline at end of file
diff --git a/app/config/config.py b/app/config/config.py
index 2d7bdf2..25875d5 100644
--- a/app/config/config.py
+++ b/app/config/config.py
@@ -14,6 +14,7 @@
     fetch_fwr_agent: str = ''
     PUBLIC_KEY: str
     PRIVATE_KEY: str
+    HASH_SUB_KEY: str
 
     def __init__(self, **kwargs):
         # Check if all required fields are provided and set them
diff --git a/app/config/config.yaml b/app/config/config.yaml
index 9477b79..26ec7e9 100644
--- a/app/config/config.yaml
+++ b/app/config/config.yaml
@@ -11,4 +11,5 @@
   -----END PUBLIC KEY-----
 PRIVATE_KEY: str
 fetch_sgb_agent: 鎶ュ憡鐢熸垚
-fetch_fwr_agent: 鐭ヨ瘑闂瓟,鏂囨。鏅鸿兘,鏅鸿兘闂瓟
\ No newline at end of file
+fetch_fwr_agent: 鐭ヨ瘑闂瓟,鏂囨。鏅鸿兘,鏅鸿兘闂瓟
+HASH_SUB_KEY: bZQDT7xZ8KkGfdnwbbxy4Ovk1Ci8AHJh7NEC9mwnLhk=
\ No newline at end of file
diff --git a/app/models/app_model.py b/app/models/app_model.py
new file mode 100644
index 0000000..ba0a580
--- /dev/null
+++ b/app/models/app_model.py
@@ -0,0 +1,31 @@
+from datetime import datetime
+from enum import IntEnum
+from sqlalchemy import Column, String, Enum as SQLAlchemyEnum, Integer, BigInteger, DateTime, Text, Float, Boolean
+from app.models.base_model import Base
+
+class AppType(IntEnum):
+    BASIC = 0
+    RAGFLOW = 1
+    BISHENG = 2
+
+
+
+class AppRegisterModel(Base):
+    __tablename__ = "app_register"
+    id = Column(Integer, primary_key=True)
+    name = Column(String(255))
+    app_type = Column(Integer, nullable=False)
+    status = Column(Integer, nullable=False, default=1)
+    created_at = Column(DateTime, default=datetime.now())
+    updated_at = Column(DateTime, default=datetime.now(), onupdate=datetime.now())
+
+
+    # to_dict 鏂规硶
+    def to_dict(self):
+        return {
+            'id': self.id,
+            'name': self.name,
+            'agent_type': self.agent_type,
+            'type': self.type
+        }
+
diff --git a/app/models/user_model.py b/app/models/user_model.py
index 96d2d32..2f87e11 100644
--- a/app/models/user_model.py
+++ b/app/models/user_model.py
@@ -28,8 +28,8 @@
     phone = Column(String(255), nullable=False, default="")
     email = Column(String(255), nullable=False, default="")
     description = Column(String(255), nullable=False, default="")
-    ragflow_id = Column(String(32), unique=True, index=True)
-    bisheng_id = Column(Integer, unique=True, index=True)
+    ragflow_id = Column(String(32))
+    bisheng_id = Column(Integer)
     login_name = Column(String(100))
     status = Column(String(10), nullable=False, default="1")
     creator = Column(String(36))
diff --git a/app/task/sync_resources.py b/app/task/sync_resources.py
new file mode 100644
index 0000000..5fddf33
--- /dev/null
+++ b/app/task/sync_resources.py
@@ -0,0 +1,40 @@
+from fastapi import Depends
+
+from Log import logger
+from app.config.config import settings
+from app.models import AgentType
+from app.models.app_model import AppRegisterModel
+from app.models.base_model import get_db
+from app.service.bisheng import BishengService
+from app.service.ragflow import RagflowService
+
+
+async def sync_resource():
+    await sync_knowledge()
+    await sync_dialog()
+    await sync_agent()
+    await sync_llm()
+
+async def sync_knowledge(db=Depends(get_db)):
+
+    token = ""
+
+    register_app = db.query(AppRegisterModel).filter(AppRegisterModel.status.__eq__(1)).all()
+    for rapp in register_app:
+        if rapp.app_type == AgentType.RAGFLOW:
+            token = ""
+            ragflow_service = RagflowService(settings.fwr_base_url)
+        elif rapp.app_type ==AgentType.BISHENG:
+            token = ""
+            bisheng_service = BishengService(settings.sgb_base_url)
+        else:
+            logger.error("娉ㄥ唽鏈煡搴旂敤锛歿}".format(rapp.app_type))
+
+async def sync_dialog():
+    ...
+
+async def sync_agent():
+    ...
+
+async def sync_llm():
+    ...
\ No newline at end of file
diff --git a/main.py b/main.py
index 0197ea4..78ffd91 100644
--- a/main.py
+++ b/main.py
@@ -1,4 +1,5 @@
 from contextlib import asynccontextmanager
+from apscheduler.schedulers.background import BackgroundScheduler
 from fastapi import FastAPI
 from app.api.auth import router as auth_router
 from app.api.canvas import canvas_router
@@ -11,11 +12,13 @@
 from app.api.organization import dept_router
 from app.api.report import router as report_router
 from app.api.resource import menu_router
+from app.api.sync_data import sync_router
 from app.api.user import user_router
 from app.api.group import group_router
 from app.api.role import role_router
 from app.models.base_model import init_db
 from app.task.fetch_agent import sync_agents, initialize_agents
+from app.task.sync_resources import sync_resource
 
 
 # 浣跨敤 Lifespan 浜嬩欢澶勭悊绋嬪簭
@@ -38,6 +41,16 @@
 )
 
 
+# 鍒涘缓璋冨害鍣�
+scheduler = BackgroundScheduler()
+scheduler.add_job(sync_resource,minutes=1, id="sync_resource_data")
+scheduler.start()
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    print("Shutdown event handler")
+    scheduler.shutdown()
+
 app.include_router(auth_router, prefix='/api/auth', tags=["auth"])
 app.include_router(chat_router, prefix='/api/chat', tags=["chat"])
 app.include_router(agent_router, prefix='/api/agent', tags=["agent"])
@@ -52,6 +65,7 @@
 app.include_router(knowledge_router, prefix='/api/knowledge', tags=["knowledge"])
 app.include_router(dialog_router, prefix='/api/dialog', tags=["dialog"])
 app.include_router(canvas_router, prefix='/api/canvas', tags=["canvas"])
+app.include_router(sync_router, prefix='/api/sync', tags=["sync"])
 
 if __name__ == "__main__":
     import uvicorn
diff --git a/requirements.txt b/requirements.txt
index a1e863b..a83ed48 100644
--- a/requirements.txt
+++ b/requirements.txt
Binary files differ

--
Gitblit v1.8.0