Upload index.js

index.js

@@ -0,0 +1,474 @@
+import fetch from 'node-fetch';
+import express from 'express';
+import cors from 'cors';
+import dotenv from 'dotenv';
+
+dotenv.config();
+
+const Tokens =[];
+let tokenManager;
+const CONFIG = {
+  API: {
+    BASE_URL: "https://partyrock.aws",
+    API_KEY: process.env.API_KEY || "sk-123456"//自定义你自己的认证密钥，记得修改
+  },
+  RETRY: {
+    MAX_ATTEMPTS: 1,
+    DELAY_BASE: 1000
+  },
+  SERVER: {
+    PORT: process.env.PORT || 3000,
+    BODY_LIMIT: '5mb'
+  },
+  MODELS: {
+    'claude-3-5-haiku-20241022': 'bedrock-anthropic.claude-3-5-haiku',
+    'claude-3-5-sonnet-20241022': 'bedrock-anthropic.claude-3-5-sonnet-v2-0',
+    'nova-lite-v1-0': 'bedrock-amazon.nova-lite-v1-0',
+    'nova-pro-v1-0': 'bedrock-amazon.nova-pro-v1-0',
+    'llama3-1-7b': 'bedrock-meta.llama3-1-8b-instruct-v1',
+    'llama3-1-70b': 'bedrock-meta.llama3-1-70b-instruct-v1',
+    'mistral-small': 'bedrock-mistral.mistral-small-2402-v1-0',
+    'mistral-large': 'bedrock-mistral.mistral-large-2407-v1-0'
+  },
+  DEFAULT_HEADERS: {
+    "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",
+    "Accept-Encoding": "gzip, deflate, br, zstd",
+    "Content-Type": "application/json",
+    "anti-csrftoken-a2z": "",
+    "sec-ch-ua-mobile": "?0",
+    "origin": "https://partyrock.aws",
+    "sec-fetch-site": "same-origin",
+    "sec-fetch-mode": "cors",
+    "sec-fetch-dest": "empty",
+    "sec-ch-ua-platform": "\"Windows\"",
+    "sec-ch-ua": "\"Google Chrome\";v=\"131\", \"Chromium\";v=\"131\", \"Not_A Brand\";v=\"24\"",
+    "referer": "",
+    "Cookie": "",
+    "accept-language": "zh-CN,zh;q=0.9",
+    "priority": "u=1, i"
+  }
+};
+
+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 updateTokens(response) {
+    const newCsrfToken = response.headers.get('anti-csrftoken-a2z');
+    if (newCsrfToken) {
+      this.cache.tokens.csrfToken = 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];
+      }
+    }
+
+    await this.updateTokenIndex();
+  }
+}
+class Utils {
+  static async extractTokens(cookieString) {
+    const tokens = {};
+    const cookiePairs = cookieString.split(';').map(pair => pair.trim());
+
+    cookiePairs.forEach(pair => {
+      const splitIndex = pair.indexOf('=');
+      const key = pair.slice(0, splitIndex).trim();
+      const value = pair.slice(splitIndex + 1).trim();
+
+      tokens[key] = value;
+    });
+
+    return tokens;
+  }
+  // 获取数组中的随机元素
+  static getRandomElement(arr) {
+    return arr[Math.floor(Math.random() * arr.length)];
+  }
+
+  // 生成UUID
+  static uuidv4() {
+    return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, function (c) {
+      const r = (Math.random() * 16) | 0;
+      const v = c === 'x' ? r : (r & 0x3) | 0x8;
+      return v.toString(16);
+    });
+  }
+
+  // 生成随机十六进制字符串
+  static generateRandomHexString(length) {
+    let result = '';
+    const characters = '0123456789ABCDEF';
+    for (let i = 0; i < length; i++) {
+      result += characters.charAt(Math.floor(Math.random() * characters.length));
+    }
+    return result;
+  }
+}
+async function initializeService() {
+  console.log('服务初始化中...');
+  // 遍历所有可能的 token 组
+  let index = 0;
+  while (true) {
+    const refreshUrl = process.env[`AUTH_TOKENS_${index}_REFRESH_URL`];
+    const anti_csrftoken_a2z = process.env[`AUTH_TOKENS_${index}_ANTI_CSRF_TOKEN`];
+    const cookie = process.env[`AUTH_TOKENS_${index}_COOKIE`];
+    if (!refreshUrl && !anti_csrftoken_a2z && !cookie) {
+      break;
+    }
+    const cookies = await Utils.extractTokens(cookie);
+    // 只有当所有属性都存在时才添加 token
+    if (refreshUrl && anti_csrftoken_a2z && cookie) {
+      Tokens.push({
+        refreshUrl,
+        anti_csrftoken_a2z,
+        pr_refresh_token: cookies["pr_refresh_token"],
+        aws_waf_token: cookies["aws-waf-token"],
+        idToken: cookies["idToken"]
+      });
+    }
+    index++;
+  }
+  tokenManager = new TokenManager();
+}
+
+await initializeService();
+
+
+
+class ApiClient {
+  constructor(modelId) {
+    console.log(modelId);
+    if (!CONFIG.MODELS[modelId]) {
+      throw new Error(`不支持的模型: ${modelId}`);
+    }
+    this.modelId = CONFIG.MODELS[modelId];
+  }
+
+  processMessageContent(content) {
+    if (typeof content === 'string') return content;
+
+    if (Array.isArray(content)) {
+      return content
+        .map(item => item.text)
+        .join('\n');
+    }
+
+    if (typeof content === 'object') return content.text || null;
+    return null;
+  }
+
+  //合并相同role的消息
+  async transformMessages(request) {
+    const mergedMessages = await request.messages.reduce(async (accPromise, current) => {
+      const acc = await accPromise;
+      const lastMessage = acc[acc.length - 1];
+      if (lastMessage && lastMessage.role == "system") {
+        lastMessage.role = "user"
+      }
+      if (current && current.role == "system") {
+        current.role = "user"
+      }
+      const currentContent = this.processMessageContent(current.content);
+
+      if (currentContent === null) return acc;
+
+      if (lastMessage && current && (lastMessage.role == current.role)) {
+        const lastContent = this.processMessageContent(lastMessage.content);
+        if (lastContent !== null) {
+          lastMessage.content = [
+            {
+              "text": `${lastContent}\r\n${currentContent}`
+            }
+          ];
+          return acc;
+        }
+      }
+      current.content = [
+        {
+          "text": currentContent
+        }
+      ]
+      acc.push(current);
+      return acc;
+    }, Promise.resolve([]));
+    // 处理请求参数
+    let topP = request.top_p || 0.5;
+    let temperature = request.temperature || 0.95;
+    if (topP >= 1) {
+      topP = 1;
+    }
+    if (temperature >= 1) {
+      temperature = 1;
+    }
+    const extractPartyRockId = url => url.match(/https:\/\/partyrock\.aws\/u\/[^/]+\/([^/]+)/)?.[1];
+    console.log(CONFIG.DEFAULT_HEADERS.referer);
+
+    const requestPayload = {
+      "messages": mergedMessages,
+      "modelName": this.modelId,
+      "context": {
+        "type": "chat-widget",
+        "appId": extractPartyRockId(CONFIG.DEFAULT_HEADERS.referer)
+      },
+      "options": {
+        "temperature": temperature,
+        "topP": topP
+      },
+      "apiVersion": 3
+    }
+    console.log(JSON.stringify(requestPayload, null, 2));
+    return requestPayload;
+  }
+}
+class MessageProcessor {
+  static createChatResponse(message, model, isStream = false) {
+    const baseResponse = {
+      id: `chatcmpl-${Utils.uuidv4()}`,
+      created: Math.floor(Date.now() / 1000),
+      model: model
+    };
+
+    if (isStream) {
+      return {
+        ...baseResponse,
+        object: 'chat.completion.chunk',
+        choices: [{
+          index: 0,
+          delta: { content: message }
+        }]
+      };
+    }
+
+    return {
+      ...baseResponse,
+      object: 'chat.completion',
+      choices: [{
+        index: 0,
+        message: {
+          role: 'assistant',
+          content: message
+        },
+        finish_reason: 'stop'
+      }],
+      usage: null
+    };
+  }
+}
+class ResponseHandler {
+  static async handleStreamResponse(response, model, res) {
+    res.setHeader('Content-Type', 'text/event-stream');
+    res.setHeader('Cache-Control', 'no-cache');
+    res.setHeader('Connection', 'keep-alive');
+
+    try {
+      const stream = response.body;
+      let buffer = '';
+      let decoder = new TextDecoder('utf-8');
+      stream.on('data', (chunk) => {
+        buffer += decoder.decode(chunk, { stream: true });
+        const lines = buffer.split('\n');
+        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 (!data) continue;
+            try {
+              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`);
+              }
+            } catch (error) {
+              console.error('JSON解析错误:', error);
+            }
+          }
+        }
+      });
+      stream.on('end', () => {
+        res.write('data: [DONE]\n\n');
+        res.end();
+      });
+      stream.on('error', (error) => {
+        console.error('流处理错误:', error);
+        res.write('data: [DONE]\n\n');
+        res.end();
+      });
+
+    } catch (error) {
+      console.error('处理响应错误:', error);
+      res.write('data: [DONE]\n\n');
+      res.end();
+    }
+  }
+
+  static async handleNormalResponse(response, model, res) {
+    const text = await response.text();
+    const lines = text.split("\n");
+    let fullResponse = '';
+
+    for (let line of lines) {
+      line = line.trim();
+      if (line) {
+        if (line.startsWith('data: ')) {
+          let data = line.substring(6);
+          if (data === '[DONE]') break;
+          try {
+            let json = JSON.parse(data)
+            if (json?.text) {
+              fullResponse += json.text;
+            }
+          } catch (error) {
+            console.log("json解析错误");
+            continue
+          }
+        }
+      }
+    }
+    const responseData = MessageProcessor.createChatResponse(fullResponse, model);
+    res.json(responseData);
+  }
+}
+// Express 应用设置
+const app = express();
+app.use(express.json({ limit: CONFIG.SERVER.BODY_LIMIT }));
+app.use(express.urlencoded({ extended: true, limit: CONFIG.SERVER.BODY_LIMIT }));
+
+app.use(cors({
+  origin: '*',
+  methods: ['GET', 'POST', 'OPTIONS'],
+  allowedHeaders: ['*']
+}));
+// 路由处理
+app.get('/v1/models', (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('/v1/chat/completions', async (req, res) => {
+  try {
+    const authToken = req.headers.authorization?.replace('Bearer ', '');
+    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));
+
+    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)
+        });
+        await tokenManager.updateTokens(response);
+        if (response.status == 200) {
+          console.log("请求成功");
+          break; // 如果请求成功，跳出重试循环
+        }
+        if (response.status != 200) {
+          console.error(JSON.stringify(await response.text(), null, 2));
+        }
+        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));
+      }
+    }
+    if (req.body.stream) {
+      await ResponseHandler.handleStreamResponse(response, req.body.model, res);
+    } else {
+      await ResponseHandler.handleNormalResponse(response, req.body.model, res);
+    }
+
+  } catch (error) {
+    res.status(500).json({
+      error: {
+        message: error.message,
+        type: 'server_error',
+        param: null,
+        code: error.code || null
+      }
+    });
+  }
+});
+
+
+app.use((req, res) => {
+  res.status(404).json({ message: "API服务运行正常" });
+});
+
+// 启动服务器
+app.listen(CONFIG.SERVER.PORT, () => {
+  console.log(`服务器运行在端口 ${CONFIG.SERVER.PORT} `);
+});