Merge branch 'master' of http://192.168.5.5:10010/r/rag-gateway

zhaoqingang

2024-12-18 916d984d2628fd8d863183cf5ba9e5a0a7339871

 app/service/v2/initialize_data.py

@@ -1,8 +1,17 @@
import json

import time
from Log import logger
from app.models import MenuCapacityModel, WebMenuModel, GroupModel, RoleModel, DialogModel, UserModel
from app.config.const import DIFY
from app.models import MenuCapacityModel, WebMenuModel, GroupModel, RoleModel, DialogModel, UserModel, UserAppModel, \
    cipher_suite
from app.service.auth import UserAppDao
from app.service.bisheng import BishengService
from app.service.difyService import DifyService
from app.service.ragflow import RagflowService
from app.service.service_token import get_new_token
from app.service.v2.app_register import AppRegisterDao
from app.config.config import settings
from app.utils.password_handle import generate_password


async def dialog_menu_sync(db):
@@ -20,7 +29,9 @@
        # print(menu)
        dialog = menu.pop("dialog", [])
        for i in dialog:
            capacity = MenuCapacityModel(menu_id=menu["id"], capacity_id=i["id"], capacity_type=i["agentType"], chat_id=i["id"] if not i["chat_id"] else i["chat_id"], chat_type=i["chat_type"])
            capacity = MenuCapacityModel(menu_id=menu["id"], capacity_id=i["id"], capacity_type=i["agentType"],
                                         chat_id=i["id"] if not i["chat_id"] else i["chat_id"],
                                         chat_type=i["chat_type"])
            db.add(capacity)
        menu_obj = WebMenuModel(**menu)
        db.add(menu_obj)
@@ -73,7 +84,7 @@
        agent_list = agent_dict.get("basic", [])
    user = db.query(UserModel).filter_by(permission="admin").first()
    for agent in agent_list:
        dialog = db.query(DialogModel).filter(DialogModel.id==agent["id"]).first()
        dialog = db.query(DialogModel).filter(DialogModel.id == agent["id"]).first()
        if dialog:
            try:
                dialog.name = agent["name"]
@@ -92,4 +103,78 @@
                db.refresh(dialog)
            except Exception as e:
                print(e)
                db.rollback()
                db.rollback()


async def user_update_app(userid, db):
    user = db.query(UserModel).filter(UserModel.id == userid).first()
    if not user:
        raise Exception("User id not found")
    app_register = AppRegisterDao(db).get_apps()
    register_dict = {}
    token = ""
    app_password = await generate_password(10)
    crypt_password = UserAppModel.encrypted_password(app_password)
    for app in app_register:
        if app["id"] == 'ragflow_app':
            user_rag_app = db.query(UserAppModel).filter(UserAppModel.user_id == userid,
                                                         UserAppModel.app_type == 'ragflow_app').all()
            if not user_rag_app:
                service = RagflowService(settings.fwr_base_url)

                register_info = await register_app(service, app["id"], app_password, token)
                if register_info:
                    register_dict[app["id"]] = register_info
                app_name = register_info.get("name")
                app_id = register_info.get("id")
                app_email = register_info.get("email")
                await save_db(db, app_name, crypt_password, app_email, user.id, app_id, "ragflow_app")
        elif app["id"] == 'bisheng_app':
            user_bs_app = db.query(UserAppModel).filter(UserAppModel.user_id == userid,
                                                        UserAppModel.app_type == 'bisheng_app').all()
            if not user_bs_app:
                service = BishengService(settings.sgb_base_url)

                register_info = await register_app(service, app["id"], app_password, token)
                if register_info:
                    register_dict[app["id"]] = register_info
                app_name = register_info.get("name")
                app_id = register_info.get("id")
                app_email = register_info.get("email")
                await save_db(db, app_name, crypt_password, app_email, user.id, app_id, "bisheng_app")
        elif app["id"] == 'dify_app':
            user_df_app = db.query(UserAppModel).filter(UserAppModel.user_id == userid,
                                                        UserAppModel.app_type == 'dify_app').all()
            if not user_df_app:
                admin_user = db.query(UserModel).filter(UserModel.permission == "admin").first()
                token = await get_new_token(db, admin_user.id, DIFY)
                if not token:
                    print("用户注册获取dftoken失败！")
                service = DifyService(settings.dify_base_url)
                register_info = await register_app(service, app["id"], app_password, token)
                if register_info:
                    register_dict[app["id"]] = register_info
                app_name = register_info.get("name")
                app_id = register_info.get("id")
                app_email = register_info.get("email")
                await save_db(db, app_name, crypt_password, app_email, user.id, app_id, "dify_app")
        else:
            raise Exception("未知注册应用---")


async def register_app(service, app_id, app_password, token):
    name = app_id + str(int(time.time()))
    try:
        register_info = await service.register(name, app_password, token)
        return {"id": register_info.get("id"), "name": name, "email": register_info.get("email")}
    except Exception as e:
        print(f"Failed to register with {app_id}: {str(e)}")
        return None


async def save_db(db, username, password, email, user_id, app_id, app_type):
    user_app_dao = UserAppDao(db)
    user_id = await user_app_dao.insert_user_app_data(username, password, email, user_id, app_id, app_type)
    if not user_id:
        raise Exception("Failed to register with app")
    print({"msg": "User registered successfully", "userFlag": user_id})