Update index.js

index.js

@@ -39,6 +39,7 @@ const CONFIG = {
     'mistral-large': 'bedrock-mistral.mistral-large-2407-v1-0'
   },
   DEFAULT_HEADERS: {
+    "Host": "partyrock.aws",
     "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
     "Connection": "keep-alive",
     "Accept": "text/event-stream",
@@ -101,7 +102,6 @@ class TokenManager {
     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};cwr_s=${Tokens[currentIndex].cwr_s};cwr_u=${sessionId}`;
     CONFIG.DEFAULT_HEADERS.referer = Tokens[currentIndex].refreshUrl;
-    CONFIG.DEFAULT_HEADERS["User-Agent"] = await Utils.getRandomUserAgent();
   }
 
   async updateTokens(response, isWaf = false) {
@@ -111,9 +111,12 @@ class TokenManager {
         Tokens[currentIndex].aws_waf_token = wafToken;
         await this.updateCacheTokens();
         this.updateRedisTokens();
+        currentIndex = (currentIndex + 1) % Tokens.length;
+        console.log("成功提取 aws-waf-token");
       } else {
         currentIndex = (currentIndex + 1) % Tokens.length;
         await this.updateCacheTokens();
+        console.log("提取aws-waf-token失败");
       }
     } else {
       const newCsrfToken = response.headers.get('anti-csrftoken-a2z');
@@ -149,6 +152,7 @@ class Utils {
     }
   }
   static async extractWaf() {
+    puppeteer.use(StealthPlugin())
     const browser = await puppeteer.launch({
       headless: true,
       args: [
@@ -165,7 +169,7 @@ class Utils {
         cookie: `pr_refresh_token=${Tokens[currentIndex].pr_refresh_token};idToken=${Tokens[currentIndex].idToken};aws-waf-token=${Tokens[currentIndex].aws_waf_token};cwr_s=${Tokens[currentIndex].cwr_s};cwr_u=${Utils.uuidv4()}`
       });
       await page.setUserAgent(
-        await Utils.getRandomUserAgent()
+        CONFIG.DEFAULT_HEADERS["User-Agent"]
       )
       await page.goto(Tokens[currentIndex].refreshUrl, {
         waitUntil: 'networkidle2',
@@ -184,11 +188,9 @@ class Utils {
       )?.value;
 
       if (awsWafToken) {
-        console.log("成功提取 aws-waf-token");
         await browser.close();
         return awsWafToken;
       } else {
-        console.log("提取aws-waf-token失败");
         await browser.close();
         return null;
       }
@@ -242,7 +244,6 @@ async function initializeService() {
   console.log('服务初始化中...');
   tokenManager = new TokenManager();
   redisClient = new RedisClient();
-  puppeteer.use(StealthPlugin())
   let index = 0;
   while (true) {
     console.log(index, '开始检测是否有缓存');
@@ -515,17 +516,7 @@ app.get('/hf/v1/models', (req, res) => {
     }))
   });
 });
-app.post('/hf/delete/redis', (req, res) => {
-  res.json({
-    object: "list",
-    data: Object.keys(CONFIG.MODELS).map((model, index) => ({
-      id: model,
-      object: "model",
-      created: Math.floor(Date.now() / 1000),
-      owned_by: "partyrock",
-    }))
-  });
-});
+
 app.post('/hf/v1/chat/completions', async (req, res) => {
   var reqStatus = 500;
   try {
@@ -536,80 +527,52 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
     await tokenManager.updateCacheTokens();
     const apiClient = new ApiClient(req.body.model);
     const requestPayload = await apiClient.transformMessages(req.body);
-    let retryCount = 0;
-
-    while (retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
-      try {
-        console.log("开始请求");
-        //发送请求
-        var response = await fetch(`https://partyrock.aws/stream/getCompletion`, {
-          method: "POST",
-          headers: {
-            ...CONFIG.DEFAULT_HEADERS
-          },
-          body: JSON.stringify(requestPayload)
-        });
-        reqStatus = response.status;
-        switch (reqStatus) {
-          case 200:
-            console.log("请求成功");
-            // 异步更新token
-            tokenManager.updateTokens(response)
-            // 处理响应
-            if (req.body.stream) {
-              await ResponseHandler.handleStreamResponse(response, req.body.model, res);
-            } else {
-              await ResponseHandler.handleNormalResponse(response, req.body.model, res);
-            }
-            return;
-          case 202:
-            console.log("请求受限，更新WAF");
-            await tokenManager.updateTokens(response, true);
-            break;
-          case 405:
-            console.log("人机验证");
-            await tokenManager.updateTokens(response, true);//常识获取waf，然后返回错误提示。
-            res.status(405).json({
-              error: {
-                message: "人机验证，请重新请求，如果多次失败，请重新更换token",
-                type: 'server_error',
-                param: null,
-                code: null
-              }
-            });
-            return;
-          case 400:
-            console.log("信息过期，请求失败");
-            await tokenManager.updateTokens(response);
-            res.status(400).json({
-              error: {
-                message: "已刷新验证信息，请重新请求",
-                type: 'server_error',
-                param: null,
-                code: null
-              }
-            });
-            return;
-          case 403:
-            console.log("请求被阻止");
-            throw new Error(`请求失败! status: ${response.statusText}，请重新请求，如果多次失败，请重新更换token`);
-          default:
-            throw new Error(`请求失败! status: ${response.status}`);
-        }
-        console.log("请求失败，重试中...");
-        retryCount++;
-        if (retryCount >= CONFIG.RETRY.MAX_ATTEMPTS) {
-          throw new Error(`上游服务请求失败! status: ${response.status}`);
-        }
-        await new Promise(resolve => setTimeout(resolve, CONFIG.RETRY.DELAY_BASE * retryCount));
-
-      } catch (error) {
-        retryCount++;
-        if (retryCount >= CONFIG.RETRY.MAX_ATTEMPTS) {
-          throw error;
-        }
-        await new Promise(resolve => setTimeout(resolve, CONFIG.RETRY.DELAY_BASE * retryCount));
+    try {
+      console.log("开始请求");
+      //发送请求
+      var response = await fetch(`https://partyrock.aws/stream/getCompletion`, {
+        method: "POST",
+        headers: {
+          ...CONFIG.DEFAULT_HEADERS
+        },
+        body: JSON.stringify(requestPayload)
+      });
+      reqStatus = response.status;
+      switch (reqStatus) {
+        case 200:
+          console.log("请求成功");
+          // 异步更新token
+          tokenManager.updateTokens(response)
+          // 处理响应
+          if (req.body.stream) {
+            await ResponseHandler.handleStreamResponse(response, req.body.model, res);
+          } else {
+            await ResponseHandler.handleNormalResponse(response, req.body.model, res);
+          }
+          return;
+        case 202:
+          console.log("请求受限，更新WAF");
+          await tokenManager.updateTokens(response, true);
+          throw new Error(`请求失败! status: ${response.statusText}，已刷新验证信息，请重新请求`);
+        case 405:
+          console.log("人机验证");
+          await tokenManager.updateTokens(response, true);//尝试获取waf，然后返回错误提示。
+          throw new Error(`请求失败! status: ${response.statusText}，人机验证，请重新请求，如果多次失败，请重新更换token`);
+        case 400:
+          console.log("信息过期，请求失败");
+          await tokenManager.updateTokens(response);
+          throw new Error(`请求失败! status: ${response.statusText}，已刷新验证信息，请重新请求`);
+        case 403:
+          console.log("请求被阻止");
+          console.log(JSON.stringify(CONFIG.DEFAULT_HEADERS, null, 2));
+          //CONFIG.DEFAULT_HEADERS["User-Agent"] = await Utils.getRandomUserAgent();
+          console.log(response.headers);
+          throw new Error(`请求失败! status: ${response.statusText}，请重新请求，如果多次失败，请重新更换token`);
+        default:
+          throw new Error(`请求失败! status: ${response.status}`);
       }
+    } catch (error) {
+      throw new Error(`请求失败! status: ${response.status}`);
     }
   } catch (error) {
     res.status(reqStatus).json({