Bläddra i källkod

增加ping機制

zhaoqingang 1 år sedan
förälder
incheckning
48e33c7fbd
3 ändrade filer med 27 tillägg och 12 borttagningar
  1. 3 1
      app/config/agent_base_url.py
  2. 14 3
      app/service/v2/app_driver/chat_data.py
  3. 10 8
      app/task/sync_account_token.py

+ 3 - 1
app/config/agent_base_url.py

@@ -5,10 +5,12 @@ RG_CHAT_SESSIONS= "/api/v1/chats/{}/sessions"
 RG_APP_NEW_TOKEN= "/v1/system/new_token"
 RG_APP_TOKEN_LIST= "/v1/system/token_list"
 RG_USER_LOGIN= "/v1/user/login"
+RG_PING= "/v1/system/version"
 
 ### ----------
 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"
+DF_USER_LOGIN= "/console/api/login"
+DF_PING = "/console/api/workspaces"

+ 14 - 3
app/service/v2/app_driver/chat_data.py

@@ -47,6 +47,12 @@ class ChatBaseApply(ChatBase):
         password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt(password)
         return password
 
+    @staticmethod
+    async def get_chat_headers(token):
+        return {
+            'Content-Type': 'application/json',
+            'Authorization': token
+        }
 
 
 
@@ -69,9 +75,14 @@ if __name__ == "__main__":
             'Content-Type': 'application/json',
             'Authorization': f"Bearer {token}"
         }
-        ans = await chat.chat_parameters(url, params, headers)
-        print(ans)
-
+        # ans = await chat.chat_parameters(url, params, headers)
+        # print(ans)
+
+        ping_url = "http://192.168.20.119:13002/console/api/workspaces"
+        user_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoiNzBhNWVhZWYtYzU4Yi00YTE0LTliOGQtYWE2MzY0ZmE1ZDE3IiwiZXhwIjoxNzM5MjU2MDMyLCJpc3MiOiJTRUxGX0hPU1RFRCIsInN1YiI6IkNvbnNvbGUgQVBJIFBhc3Nwb3J0In0.tlxdRIVL2Ewgb0V4M5pj6v_U2yyvQxtCDrbvSxLjUBs"
+        token = "Bearer {}"
+        res = await chat.chat_ping(ping_url, {}, await chat.get_chat_headers(token.format(user_token)))
+        print(res)
 
     import asyncio
 

+ 10 - 8
app/task/sync_account_token.py

@@ -5,9 +5,9 @@ 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.agent_base_url import RG_USER_LOGIN, DF_USER_LOGIN, RG_PING, DF_PING
 from app.config.config import settings
-from app.config.const import chat_server, workflow_server
+from app.config.const import chat_server, workflow_server, http_200
 from app.models import UserTokenModel
 from app.models.app_token_model import AppToken
 from app.models.base_model import SessionLocal
@@ -29,19 +29,21 @@ async def sync_token():
     :return:
     """
     app_data = [{"token_id": chat_server, "url": f"{settings.fwr_base_url}{RG_USER_LOGIN}",
-                 "data": {"email": "", "password": ""}, "is_crypt": True},
+                 "data": {"email": "", "password": ""}, "is_crypt": True, "ping_url": f"{settings.fwr_base_url}{RG_PING}",
+                 "token": "{}"},
                 {"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}]
+                 "data": {"email": "", "password": "", "remember_me": True, "language": "zh-Hans"}, "is_crypt": False
+                 , "ping_url": f"{settings.dify_base_url}{DF_PING}", "token": "Bearer {}"}]
 
-    async def sync_token_chat(token_id, url, data, is_crypt):
+    async def sync_token_chat(token_id, url, data, is_crypt, ping_url, token):
         db = SessionLocal()
         # pdb = PostgresqlSessionLocal()
         current_time = datetime.now() - timedelta(hours=12)
         try:
             user_token = db.query(UserTokenModel).filter(UserTokenModel.id == token_id).first()
+            chat = ChatBaseApply()
+            if user_token and (user_token.updated_at < current_time or not user_token.access_token or await chat.chat_ping(ping_url, {}, await chat.get_chat_headers(token.format(user_token.access_token))) != http_200):
 
-            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))
@@ -61,7 +63,7 @@ async def sync_token():
 
     tasks = []
     for app in app_data:
-        tasks.append(asyncio.create_task(sync_token_chat(app["token_id"], app["url"], app["data"], app["is_crypt"])))
+        tasks.append(asyncio.create_task(sync_token_chat(app["token_id"], app["url"], app["data"], app["is_crypt"], app["ping_url"], app["token"])))
     done, pending = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)