TOKEN同步功能

zhaoqingang

2025-02-06 bcc63761bdc4c1604c9275a3c5cdf8a483ad5611

TOKEN同步功能

 app/api/auth.py

@@ -9,13 +9,13 @@
from app.api import Response, pwd_context, get_current_user
from app.api.user import reset_user_pwd
from app.config.config import settings
from app.config.const import RAGFLOW, BISHENG, DIFY
from app.config.const import RAGFLOW, BISHENG, DIFY, chat_server, workflow_server
from app.models.app_token_model import AppToken
from app.models.base_model import get_db
from app.models.postgresql_base_model import get_pdb
from app.models.token_model import upsert_token, update_token
from app.models.user import UserCreate, LoginData
from app.models.user_model import UserModel, UserAppModel
from app.models.user_model import UserModel, UserAppModel, UserTokenModel
from app.service.auth import authenticate_user, create_access_token, is_valid_password, save_register_user, \
    update_user_token, UserAppDao, update_user_info, password_rsa
from app.service.bisheng import BishengService
@@ -143,15 +143,18 @@

    # await update_token(db, user.id, access_token, token_dict)
    # await update_user_token(db, user.id, token_dict)
    '''
    token_dict = {}
    for app in  [{"app": RAGFLOW, "token_id": chat_server},
                {"app": DIFY,"token_id": workflow_server}]:
        user_token = db.query(UserTokenModel).filter(UserTokenModel.id == app["token_id"]).first()
        if user_token and user_token.access_token:
            token_dict[app["app"]] = user_token.access_token
    result = await pdb.execute(select(AppToken).where(AppToken.id == user.id))
    db_app_token = result.scalars().first()
    if isinstance(access_token, bytes):
        access_token = access_token.decode()
    if not db_app_token:
        app_token_str = json.dumps(token_dict)
        # print(app_token_str)
        app_token = AppToken(id=user.id, token=access_token, app_token=app_token_str)
        app_token = AppToken(id=user.id, token=access_token, app_token=json.dumps(token_dict))
        pdb.add(app_token)
        await pdb.commit()
        await pdb.refresh(app_token)
@@ -160,7 +163,7 @@
        db_app_token.app_token = json.dumps(token_dict)
        await pdb.commit()
        await pdb.refresh(db_app_token)
    '''

    return Response(code=200, msg="Login successful", data={
        "access_token": access_token,
        "token_type": "bearer",

 app/api/v2/chat.py

@@ -19,21 +19,21 @@

    chat_info = await get_chat_info(db, chatId)
    if not chat_info:
        error_msg = json.dumps({"message": smart_message_error, "error": "**ERROR**: parameter exception", "status": http_400})
        error_msg = json.dumps({"message": smart_message_error, "error": "\n**ERROR**: parameter exception", "status": http_400})
        return StreamingResponse(f"data: {error_msg}\n\n",
                                 media_type="text/event-stream")
    session_id = dialog.sessionId
    if chat_info.mode == dialog_chat:
        if not dialog.query:
            error_msg = json.dumps(
                {"message": smart_message_error, "error": "**ERROR**: question cannot be empty.", "status": http_400})
                {"message": smart_message_error, "error": "\n**ERROR**: question cannot be empty.", "status": http_400})
            return StreamingResponse(f"data: {error_msg}\n\n",
                                     media_type="text/event-stream")
        if not session_id:
            session = await service_chat_sessions(db, chatId, dialog.query)
            if not session or session.get("code") != 0:
                error_msg = json.dumps(
                    {"message": smart_message_error, "error": "**ERROR**: chat agent error", "status": http_500})
                    {"message": smart_message_error, "error": "\n**ERROR**: chat agent error", "status": http_500})
                return StreamingResponse(f"data: {error_msg}\n\n",
                                         media_type="text/event-stream")
            session_id = session.get("data", {}).get("id")
@@ -47,7 +47,7 @@
        return StreamingResponse(service_chat_basic(db, chatId, dialog.question, dialog_chat.sessionId, 1),
                                 media_type="text/event-stream")
    else:
        error_msg = json.dumps({"message": smart_message_error, "error": "**ERROR**: chat agent error", "status": http_500})
        error_msg = json.dumps({"message": smart_message_error, "error": "\n**ERROR**: chat agent error", "status": http_500})
        return StreamingResponse(f"data: {error_msg}\n\n",
                                 media_type="text/event-stream")


 app/config/agent_base_url.py

@@ -3,10 +3,12 @@
RG_CHAT_DIALOG= "/api/v1/chats/{}/completions"
RG_CHAT_SESSIONS= "/api/v1/chats/{}/sessions"
RG_APP_NEW_TOKEN= "/v1/system/new_token"
RG_APP_token_LIST= "/v1/system/token_list"
RG_APP_TOKEN_LIST= "/v1/system/token_list"
RG_USER_LOGIN= "/v1/user/login"

### ----------
DF_CHAT_AGENT= "/v1/chat-messages"
DF_CHAT_WORKFLOW= "/v1/workflows/run"
DF_CHAT_PARAMETERS= "/v1/parameters"
DF_CHAT_API_KEY= "/console/api/apps/{}/api-keys"
DF_USER_LOGIN= "/console/api/login"

 app/config/env_conf/account.yaml

@@ -1,7 +1,7 @@
chat_server:
  account: user@example.com
  password: secret
  account: zhao1@example.com
  password: gAAAAABnpFLtotY2OIRH12BJh4MzMgn5Zil7-DVpIeuqlFwvr0g6g_n4ULogn-LNhCbtk6cCDkzZlqAHvBSX2e_zf7AsoyzbiQ==

workflow_server:
  account: basic@mail.com
  password: Basic2024
  account: zqg@mail.com
  password: gAAAAABnpCfh1IAlSg1wBn7xPjbiucm-YAyzrKygw4rQiRg267ZMrhIolE38n94E8GgOmWpFfQx3s79zdD6yUrIqOLa2YeW7eA==

 app/config/env_conf/config.yaml

@@ -1,10 +1,10 @@
secret_key: your-secret-key
sgb_base_url: http://192.168.20.119:13001
sgb_websocket_url: ws://192.168.20.119:13001
fwr_base_url: http://192.168.20.119:11080
fwr_base_url: http://192.168.20.116:11080
database_url: mysql+pymysql://root:infini_rag_flow@192.168.20.119:5455/rag_basic
sgb_db_url: mysql+pymysql://root:1234@192.168.20.119:13306/bisheng
fwr_db_url: mysql+pymysql://root:infini_rag_flow@192.168.20.119:15455/rag_flow
fwr_db_url: mysql+pymysql://root:infini_rag_flow@192.168.20.116:15455/rag_flow
PUBLIC_KEY: |
  -----BEGIN PUBLIC KEY-----
  MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArq9XTUSeYr2+N1h3Afl/z8Dse/2yD0ZGrKwx+EEEcdsBLca9Ynmx3nIB5obmLlSfmskLpBo0UACBmB5rEjBp2Q2f3AG3Hjd4B+gNCG6BDaawuDlgANIhGnaTLrIqWrrcm4EMzJOnAOI1fgzJRsOOUEfaS318Eq9OVO3apEyCCt0lOQK6PuksduOjVxtltDav+guVAA068NrPYmRNabVKRNLJpL8w4D44sfth5RvZ3q9t+6RTArpEtc5sh5ChzvqPOzKGMXW83C95TxmXqpbK6olN4RevSfVjEAgCydH6HN6OhtOQEcnrU97r9H0iZOWwbw3pVrZiUkuRD1R56Wzs2wIDAQAB
@@ -15,11 +15,11 @@
PASSWORD_KEY: VKinqB-8XMrwCLLrcf_PyHyo12_4PVKvWzaHjNFions=
basic_base_url: http://192.168.20.231:8000
basic_paper_url: http://192.168.20.231:8000
dify_base_url: http://192.168.20.116
dify_base_url: http://192.168.20.119:13002
dify_api_token: app-YmOAMDsPpDDlqryMHnc9TzTO
postgresql_database_url: postgresql+asyncpg://kong:kongpass@192.168.20.119:5432/kong
dify_workflow_clean: app-OpF0drPu0XcgqcekQpT4FA8a
dify_workflow_report: app-0MAkdFWqh9zxwmU69O0BFU1s
dify_database_url: postgresql+psycopg2://postgres:difyai123456@192.168.20.116:15432/dify
dify_database_url: postgresql+psycopg2://postgres:difyai123456@192.168.20.119:15432/dify



 app/config/env_conf/default_agent_conf.json

@@ -108,6 +108,15 @@
      "description": "文档出卷",
      "icon": "intellFrame4",
      "agentType": "paperTalk",
      "parameters": {
        "retriever_resource": {
          "enabled": true
        },
        "user_input_form": [],
        "file_upload":{
          "enabled": false
        }
      },
      "dialogType": "3",
      "mode": "agent-basic"
    }

 app/models/postgresql_base_model.py

@@ -14,4 +14,8 @@

async def get_pdb() -> AsyncSession:
    async with PostgresqlSessionLocal() as session:
        # yield session
        try:
        yield session
        finally:
            session.close()

 app/service/v2/chat.py

@@ -119,7 +119,7 @@
        logger.error(e)
        try:
            yield "data: " + json.dumps({"message": smart_message_error,
                                         "error": "**ERROR**: " + str(e), "status": http_500},
                                         "error": "\n**ERROR**: " + str(e), "status": http_500},
                                        ensure_ascii=False) + "\n\n"
        except:
            ...
@@ -199,7 +199,7 @@
        logger.error(e)
        try:
            yield "data: " + json.dumps({"message": smart_message_error,
                                         "error": "**ERROR**: " + str(e), "status": http_500},
                                         "error": "\n**ERROR**: " + str(e), "status": http_500},
                                        ensure_ascii=False) + "\n\n"
        except:
            ...

 app/service/v2/initialize_data.py

@@ -135,6 +135,7 @@
                dialog.name = agent["name"]
                dialog.description = agent["description"]
                dialog.icon = agent["icon"]
                dialog.mode = agent["mode"]
                dialog.parameters = json.dumps(agent["parameters"])
                db.commit()
            except Exception as e:

 app/task/fetch_agent.py

@@ -8,7 +8,7 @@
from sqlalchemy.orm import sessionmaker, Session

from app.config.config import settings
from app.config.const import RAGFLOW, BISHENG, DIFY, ENV_CONF_PATH
from app.config.const import RAGFLOW, BISHENG, DIFY, ENV_CONF_PATH, Dialog_STATSU_DELETE, Dialog_STATSU_ON
from app.models import KnowledgeModel
from app.models.dialog_model import DialogModel
from app.models.user_model import UserAppModel
@@ -296,6 +296,10 @@
                existing_agent.name = row["name"]
                existing_agent.description = row["description"]
                existing_agent.mode = row["mode"]
                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"],
@@ -306,7 +310,7 @@
        for dialog in db.query(DialogModel).filter_by(dialog_type=dialog_type).all():
            if dialog.id not in agent_id_list:
                # print(dialog.id)
                db.query(DialogModel).filter_by(id=dialog.id).update({"status": "2"})
                db.query(DialogModel).filter_by(id=dialog.id).update({"status": Dialog_STATSU_DELETE})
                db.commit()
    except IntegrityError:
        db.rollback()

 app/task/sync_account_token.py

@@ -1,10 +1,26 @@
import asyncio
import json
import time
from datetime import datetime, timedelta

from sqlalchemy.future import select

from app.config.agent_base_url import RG_USER_LOGIN, DF_USER_LOGIN
from app.config.config import settings
from app.config.const import chat_server, workflow_server
from app.models import UserTokenModel
from app.models.app_token_model import AppToken
from app.models.base_model import SessionLocal
from app.models.postgresql_base_model import get_pdb, PostgresqlSessionLocal
from app.service.ragflow import RagflowService
from app.service.v2.app_driver.chat_data import ChatBaseApply
from app.utils.password_handle import password_decrypted



def sync_token():
async def sync_token():
    """
    1.获取到user_token表中的账号
    2.判断是否超过12小时，否则获取新的token
    2.判断是否超过12小时，超过则获取新的token
    3.未超过12小时，则测试token有效性，chat_ping接口，返回401，则重新获取token
    4.token获取：login
    df:/console/api/workspaces
@@ -12,3 +28,52 @@
    5.跟新本地token和kong网关token
    :return:
    """
    app_data = [{"token_id": chat_server, "url": f"{settings.fwr_base_url}{RG_USER_LOGIN}",
                 "data": {"email": "", "password": ""}, "is_crypt": True},
                {"token_id": workflow_server, "url": f"{settings.dify_base_url}{DF_USER_LOGIN}",
                 "data": {"email": "", "password": "", "remember_me": True, "language": "zh-Hans"}, "is_crypt": False}]

    async def sync_token_chat(token_id, url, data, is_crypt):
        db = SessionLocal()
        # pdb = PostgresqlSessionLocal()
        current_time = datetime.now() - timedelta(hours=12)
        try:
            user_token = db.query(UserTokenModel).filter(UserTokenModel.id == token_id).first()

            if user_token and (user_token.updated_at < current_time or not user_token.access_token):
                chat = ChatBaseApply()
                data["email"] = user_token.account
                if is_crypt:
                    data["password"] = await chat.password_encrypt(await password_decrypted(user_token.password))
                else:
                    data["password"] = await password_decrypted(user_token.password)
                res = await chat.chat_login(url, data, {'Content-Type': 'application/json'})
                if res:
                    access_token = res["data"]["access_token"]
                    user_token.access_token = access_token
                    user_token.updated_at = datetime.now()
                    db.commit()
        except Exception as e:
            print(e)
        finally:
            db.close()
            # await pdb.close()

    tasks = []
    for app in app_data:
        tasks.append(asyncio.create_task(sync_token_chat(app["token_id"], app["url"], app["data"], app["is_crypt"])))
    done, pending = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)



def start_sync_token_task():
    asyncio.run(sync_token())


if __name__ == "__main__":



    start_sync_token_task()



 main.py

@@ -29,7 +29,7 @@
from app.task.fetch_agent import sync_agents, initialize_agents, sync_agents_v2, sync_knowledge, \
    sync_resources_from_json
from app.init_config.init_run_data import sync_default_data
from app.task.sync_account_token import sync_token
from app.task.sync_account_token import sync_token, start_sync_token_task

init_db()

@@ -70,7 +70,7 @@
# 创建调度器
scheduler = BackgroundScheduler()
scheduler.add_job(sync_agents_v2, 'interval', minutes=60, id="sync_resource_data")
scheduler.add_job(sync_token, 'interval', minutes=5, id="sync_token_1")
scheduler.add_job(start_sync_token_task, 'interval', minutes=5, id="sync_token_1")
# scheduler.add_job(delete_file_after_delay, 'interval', minutes=10, id="delete_file_after_delay")
scheduler.start()