Update index.js

index.js

@@ -4,6 +4,7 @@ import cors from 'cors';
 import dotenv from 'dotenv';
 import puppeteer from 'puppeteer-extra'
 import StealthPlugin from 'puppeteer-extra-plugin-stealth'
+import UserAgent from 'user-agents';
 
 dotenv.config();
 
@@ -20,7 +21,7 @@ const CONFIG = {
     RedisToken: process.env.RedisToken
   },
   RETRY: {
-    MAX_ATTEMPTS: 2,
+    MAX_ATTEMPTS: 1,
     DELAY_BASE: 1000
   },
   SERVER: {
@@ -100,6 +101,7 @@ 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) {
@@ -113,7 +115,7 @@ class TokenManager {
         currentIndex = (currentIndex + 1) % Tokens.length;
         await this.updateCacheTokens();
       }
-    }else{
+    } else {
       const newCsrfToken = response.headers.get('anti-csrftoken-a2z');
       const cookies = response.headers.get('set-cookie');
       if (newCsrfToken && cookies) {
@@ -133,6 +135,19 @@ class TokenManager {
 
 
 class Utils {
+  static async getRandomUserAgent() {
+    try {
+      let type = ["Win32", "MacIntel", "Linux x86_64"]
+      const userAgent = new UserAgent({ platform: type[Math.floor(Math.random() * type.length)] });
+      return userAgent.random().toString();
+    } catch (error) {
+      let type = [
+        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.0 Safari/605.1.15"
+      ]
+      return type[Math.floor(Math.random() * type.length)]
+    }
+  }
   static async extractWaf() {
     const browser = await puppeteer.launch({
       headless: true,
@@ -150,9 +165,9 @@ 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(
-        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.0 Safari/605.1.15"
+        await Utils.getRandomUserAgent()
       )
-      await page.goto(Tokens[currentIndex].refreshUrl, { 
+      await page.goto(Tokens[currentIndex].refreshUrl, {
         waitUntil: 'networkidle2',
         timeout: 30000
       });
@@ -345,7 +360,7 @@ class ApiClient {
       temperature = 1;
     }
     const extractPartyRockId = url => url.match(/https:\/\/partyrock\.aws\/u\/[^/]+\/([^/]+)/)?.[1];
-    console.log("当前请求的是",CONFIG.DEFAULT_HEADERS.referer);
+    console.log("当前请求的是", CONFIG.DEFAULT_HEADERS.referer);
 
     const requestPayload = {
       "messages": mergedMessages,
@@ -413,10 +428,9 @@ class ResponseHandler {
         buffer = lines.pop() || '';
 
         for (const line of lines) {
-          if (!line) continue;
-          const trimmedLine = line.trim();
-          if (trimmedLine && trimmedLine.startsWith('data: ')) {
-            const data = trimmedLine.substring(6);
+          if (!line.trim()) continue;
+          if (line.startsWith('data: ')) {
+            const data = line.substring(6);
             if (!data) continue;
             if (data == "[DONE]") {
               res.write('data: [DONE]\n\n');
@@ -501,7 +515,17 @@ 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 {
@@ -524,7 +548,7 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
             ...CONFIG.DEFAULT_HEADERS
           },
           body: JSON.stringify(requestPayload)
-        });     
+        });
         reqStatus = response.status;
         switch (reqStatus) {
           case 200:
@@ -547,13 +571,28 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
             await tokenManager.updateTokens(response, true);//常识获取waf，然后返回错误提示。
             res.status(405).json({
               error: {
-                message: "人机验证",
+                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}`);
         }
@@ -586,7 +625,7 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
 
 
 app.use((req, res) => {
-  res.status(404).send("API服务运行正常，请使用正确请求路径");
+  res.status(404).send("API服务运行正常，，请使用正确请求路径");
 });
 
 // 启动服务器