Update index.js

index.js

@@ -7,6 +7,7 @@ dotenv.config();
 
 const Tokens =[];
 let tokenManager;
+let currentIndex = 0;
 const CONFIG = {
   API: {
     BASE_URL: "https://partyrock.aws",
@@ -52,55 +53,26 @@ const CONFIG = {
 };
 
 class TokenManager {
-  constructor() {
-    this.cache = {
-      currentIndex: 0,
-      tokens: {}
-    };
-  }
-
-  async getCurrentTokenIndex() {
-    return this.cache.currentIndex;
-  }
-
-  async updateTokenIndex() {
-    const nextIndex = (this.cache.currentIndex + 1) % Tokens.length;
-    this.cache.currentIndex = nextIndex;
-  }
-
-  async getCachedTokens() {
-    const currentIndex = await this.getCurrentTokenIndex();
-    const csrfToken = this.cache.tokens.csrfToken
-      ? this.cache.tokens.csrfToken
-      : Tokens[currentIndex]["anti_csrftoken_a2z"];
-
-    const idToken = this.cache.tokens.idToken
-      ? this.cache.tokens.idToken
-      : Tokens[currentIndex]["idToken"];
-    return {
-      csrfToken: csrfToken,
-      idToken: idToken,
-      pr_refresh_token: Tokens[currentIndex]["pr_refresh_token"],
-      aws_waf_token: Tokens[currentIndex]["aws_waf_token"],
-      refreshUrl: Tokens[currentIndex]["refreshUrl"]
-    };
+  async updateCacheTokens() {
+    CONFIG.DEFAULT_HEADERS["anti-csrftoken-a2z"] = Tokens[currentIndex]["anti_csrftoken_a2z"];
+    CONFIG.DEFAULT_HEADERS.Cookie = `idToken=${Tokens[currentIndex]["idToken"]}; pr_refresh_token=${Tokens[currentIndex]["pr_refresh_token"]};aws-waf-token=${Tokens[currentIndex]["aws_waf_token"]}`;
+    CONFIG.DEFAULT_HEADERS.referer = Tokens[currentIndex]["refreshUrl"];
   }
 
   async updateTokens(response) {
     const newCsrfToken = response.headers.get('anti-csrftoken-a2z');
     if (newCsrfToken) {
-      this.cache.tokens.csrfToken = newCsrfToken;
+      Tokens[currentIndex]["anti_csrftoken_a2z"] = newCsrfToken;
     }
 
     const cookies = response.headers.get('set-cookie');
     if (cookies) {
       const idTokenMatch = cookies.match(/idToken=([^;]+)/);
       if (idTokenMatch && idTokenMatch[1]) {
-        this.cache.tokens.idToken = idTokenMatch[1];
+        Tokens[currentIndex]["idToken"] = idTokenMatch[1];
       }
     }
-
-    await this.updateTokenIndex();
+    currentIndex = (currentIndex + 1) % Tokens.length;
   }
 }
 class Utils {
@@ -166,6 +138,7 @@ async function initializeService() {
     }
     index++;
   }
+  
   tokenManager = new TokenManager();
 }
 
@@ -254,7 +227,6 @@ class ApiClient {
       },
       "apiVersion": 3
     }
-    console.log(JSON.stringify(requestPayload, null, 2));
     return requestPayload;
   }
 }
@@ -317,7 +289,6 @@ class ResponseHandler {
               const json = JSON.parse(data);
               if (json?.text) {
                 var content = json.text;
-                console.log(content);
                 const responseData = MessageProcessor.createChatResponse(content, model, true);
                 res.write(`data: ${JSON.stringify(responseData)}\n\n`);
               }
@@ -400,11 +371,8 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
     if (authToken !== CONFIG.API.API_KEY) {
       return res.status(401).json({ error: "Unauthorized" });
     }
-    const { csrfToken, idToken, pr_refresh_token, aws_waf_token, refreshUrl } = await tokenManager.getCachedTokens();
-    CONFIG.DEFAULT_HEADERS["anti-csrftoken-a2z"] = csrfToken;
-    CONFIG.DEFAULT_HEADERS.Cookie = `idToken=${idToken}; pr_refresh_token=${pr_refresh_token};aws-waf-token=${aws_waf_token}`;
-    CONFIG.DEFAULT_HEADERS.referer = refreshUrl;
-    console.log(JSON.stringify(CONFIG.DEFAULT_HEADERS, null, 2));
+    await tokenManager.updateCacheTokens();
+
 
     const apiClient = new ApiClient(req.body.model);
     const requestPayload = await apiClient.transformMessages(req.body);