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
|
|
|
async def sync_token():
|
"""
|
1.获取到user_token表中的账号
|
2.判断是否超过12小时，超过则获取新的token
|
3.未超过12小时，则测试token有效性，chat_ping接口，返回401，则重新获取token
|
4.token获取：login
|
df:/console/api/workspaces
|
rg:/v1/system/version
|
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()
|
