Update app.js

app.js

@@ -1,293 +1,236 @@
-const express = require("express");
-const { io } = require("socket.io-client");
-const { v4: uuidv4 } = require("uuid");
-const { ProxyAgent } = require("proxy-agent");
-const agent = new ProxyAgent();
+const express = require('express');
+const fetch = require('node-fetch');
+const zlib = require('zlib');
+const { promisify } = require('util');
+const stream = require('stream');
+const bodyParser = require('body-parser');
+
+const gunzip = promisify(zlib.gunzip);
+const pipeline = promisify(stream.pipeline);
+
+const PROJECT_ID = process.env.PROJECT_ID;
+const CLIENT_ID = process.env.CLIENT_ID;
+const CLIENT_SECRET = process.env.CLIENT_SECRET;
+const REFRESH_TOKEN = process.env.REFRESH_TOKEN;
+const API_KEY = process.env.API_KEY;
+
+const TOKEN_URL = 'https://www.googleapis.com/oauth2/v4/token';
+
+let tokenCache = {
+  accessToken: '',
+  expiry: 0,
+  refreshPromise: null
+};
 
-const app = express();
-const port = process.env.PORT || 8081;
+function logRequest(req, status, message) {
+  const timestamp = new Date().toISOString();
+  const method = req.method;
+  const url = req.originalUrl;
+  const ip = req.ip;
+  console.log(`[${timestamp}] ${method} ${url} - Status: ${status}, IP: ${ip}, Message: ${message}`);
+}
 
-// 从环境变量中获取 accessToken
-const accessToken = process.env.PPLX_KEY;
+async function getAccessToken() {
+  const now = Date.now() / 1000;
 
-console.log(`Server starting. Access token: ${accessToken ? "Set" : "Not set"}`);
+  if (tokenCache.accessToken && now < tokenCache.expiry - 120) {
+    return tokenCache.accessToken;
+  }
 
-// 添加一个中间件来解析 JSON 请求体
-app.use(express.json());
+  if (tokenCache.refreshPromise) {
+    await tokenCache.refreshPromise;
+    return tokenCache.accessToken;
+  }
 
-// 添加一个中间件来记录所有请求
-app.use((req, res, next) => {
-    console.log(`[${new Date().toISOString()}] ${req.method} request received for ${req.path}`);
-    console.log(`Headers: ${JSON.stringify(req.headers)}`);
-    next();
-});
+  tokenCache.refreshPromise = (async () => {
+    try {
+      const response = await fetch(TOKEN_URL, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json'
+        },
+        body: JSON.stringify({
+          client_id: CLIENT_ID,
+          client_secret: CLIENT_SECRET,
+          refresh_token: REFRESH_TOKEN,
+          grant_type: 'refresh_token'
+        })
+      });
+
+      const data = await response.json();
+      
+      tokenCache.accessToken = data.access_token;
+      tokenCache.expiry = now + data.expires_in;
+    } finally {
+      tokenCache.refreshPromise = null;
+    }
+  })();
 
-var opts = {
-    agent: agent,
-    auth: {
-        jwt: "anonymous-ask-user",
-    },
-    reconnection: false,
-    transports: ["websocket"],
-    path: "/socket.io",
-    hostname: "www.perplexity.ai",
-    secure: true,
-    port: "443",
-    extraHeaders: {
-        Cookie: process.env.PPLX_COOKIE,
-        "User-Agent": process.env.USER_AGENT,
-        Accept: "*/*",
-        priority: "u=1, i",
-        Referer: "https://www.perplexity.ai/",
-    },
-};
+  await tokenCache.refreshPromise;
+  return tokenCache.accessToken;
+}
 
-app.post("/v1/messages", (req, res) => {
-    console.log("[POST /v1/messages] Processing request");
-    
-    // 严格验证客户端的 Authorization 头
-    const clientAuth = req.headers['authorization'];
-    console.log(`[POST /v1/messages] Received Authorization header: ${clientAuth}`);
-    console.log(`[POST /v1/messages] Expected Authorization: Bearer ${accessToken}`);
+function getLocation() {
+  const currentSeconds = new Date().getSeconds();
+  return currentSeconds < 30 ? 'europe-west1' : 'us-east5';
+}
+
+function constructApiUrl(location, model) {
+  return `https://${location}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${location}/publishers/anthropic/models/${model}:streamRawPredict`;
+}
+
+function formatModelName(model) {
+  if (model === 'claude-3-5-sonnet-20240620') {
+    return 'claude-3-5-sonnet@20240620';
+  }
+  return model;
+}
+
+async function handleRequest(req, res) {
+  if (req.method === 'OPTIONS') {
+    handleOptions(res);
+    logRequest(req, 204, 'CORS preflight request');
+    return;
+  }
+
+  const apiKey = req.headers['x-api-key'];
+  if (apiKey !== API_KEY) {
+    res.status(403).json({
+      type: "error",
+      error: {
+        type: "permission_error",
+        message: "Your API key does not have permission to use the specified resource."
+      }
+    });
+    logRequest(req, 403, 'Invalid API key');
+    return;
+  }
+
+  const accessToken = await getAccessToken();
+  const location = getLocation();
+
+  let requestBody = req.body;
+  
+  let model = requestBody.model || 'claude-3-5-sonnet@20240620';
+  model = formatModelName(model);
+  
+  const apiUrl = constructApiUrl(location, model);
+
+  if (requestBody.anthropic_version) {
+    delete requestBody.anthropic_version;
+  }
+  
+  if (requestBody.model) {
+    delete requestBody.model;
+  }
+  
+  requestBody.anthropic_version = "vertex-2023-10-16";
+
+  try {
+    const response = await fetch(apiUrl, {
+      method: 'POST',
+      headers: {
+        'Authorization': `Bearer ${accessToken}`,
+        'Content-Type': 'application/json; charset=utf-8',
+        'Accept-Encoding': 'gzip, deflate'
+      },
+      body: JSON.stringify(requestBody),
+      compress: false
+    });
+
+    res.status(response.status);
     
-    if (!clientAuth || clientAuth !== `Bearer ${accessToken}`) {
-        console.log("[POST /v1/messages] Authorization failed. Sending 401 Unauthorized response.");
-        res.status(401).json({ error: "Unauthorized" });
-        return;
+    for (const [key, value] of response.headers.entries()) {
+      if (key.toLowerCase() !== 'content-encoding' && key.toLowerCase() !== 'content-length') {
+        res.setHeader(key, value);
+      }
     }
-    
-    console.log("[POST /v1/messages] Authorization successful. Proceeding with request processing.");
 
-    // 使用解析后的 JSON 体
-    const jsonBody = req.body;
-    console.log(`[POST /v1/messages] Received request body: ${JSON.stringify(jsonBody)}`);
+    res.setHeader('Access-Control-Allow-Origin', '*');
+    res.setHeader('Access-Control-Allow-Methods', 'POST, GET, OPTIONS');
+    res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization, x-api-key, anthropic-version, model');
 
-    res.setHeader("Content-Type", "text/event-stream;charset=utf-8");
-    try {
-        if (jsonBody.stream == false) {
-            console.log("[POST /v1/messages] Stream is false. Sending non-streaming response.");
-            res.send(
-                JSON.stringify({
-                    id: uuidv4(),
-                    content: [
-                        {
-                            text: "Please turn on streaming.",
-                        },
-                        {
-                            id: "string",
-                            name: "string",
-                            input: {},
-                        },
-                    ],
-                    model: "string",
-                    stop_reason: "end_turn",
-                    stop_sequence: "string",
-                    usage: {
-                        input_tokens: 0,
-                        output_tokens: 0,
-                    },
-                })
-            );
-        } else if (jsonBody.stream == true) {
-            console.log("[POST /v1/messages] Stream is true. Processing streaming request.");
-            // 计算用户消息长度
-            let userMessage = [{ question: "", answer: "" }];
-            let lastUpdate = true;
-            if (jsonBody.system) {
-                // 把系统消息加入messages的首条
-                jsonBody.messages.unshift({ role: "system", content: jsonBody.system });
-            }
-            console.log(`[POST /v1/messages] Processed messages: ${JSON.stringify(jsonBody.messages)}`);
-            jsonBody.messages.forEach((msg) => {
-                if (msg.role == "system" || msg.role == "user") {
-                    if (lastUpdate) {
-                        userMessage[userMessage.length - 1].question += msg.content + "\n";
-                    } else if (userMessage[userMessage.length - 1].question == "") {
-                        userMessage[userMessage.length - 1].question += msg.content + "\n";
-                    } else {
-                        userMessage.push({ question: msg.content + "\n", answer: "" });
-                    }
-                    lastUpdate = true;
-                } else if (msg.role == "assistant") {
-                    if (!lastUpdate) {
-                        userMessage[userMessage.length - 1].answer += msg.content + "\n";
-                    } else if (userMessage[userMessage.length - 1].answer == "") {
-                        userMessage[userMessage.length - 1].answer += msg.content + "\n";
-                    } else {
-                        userMessage.push({ question: "", answer: msg.content + "\n" });
-                    }
-                    lastUpdate = false;
-                }
-            });
-            // user message to plaintext
-            let previousMessages = jsonBody.messages
-                .map((msg) => {
-                    return msg.content
-                })
-                .join("\n\n");
-
-            console.log(`[POST /v1/messages] Previous messages: ${previousMessages}`);
-
-            let msgid = uuidv4();
-            // send message start
-            res.write(
-                createEvent("message_start", {
-                    type: "message_start",
-                    message: {
-                        id: msgid,
-                        type: "message",
-                        role: "assistant",
-                        content: [],
-                        model: "claude-3-opus-20240229",
-                        stop_reason: null,
-                        stop_sequence: null,
-                        usage: { input_tokens: 8, output_tokens: 1 },
-                    },
-                })
-            );
-            res.write(createEvent("content_block_start", { type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }));
-            res.write(createEvent("ping", { type: "ping" }));
-
-            console.log("[POST /v1/messages] Initiating WebSocket connection to Perplexity.ai");
-            // proxy response
-            var socket = io("wss://www.perplexity.ai/", opts);
-
-            socket.on("connect", function () {
-                console.log("[POST /v1/messages] WebSocket connected to Perplexity.ai");
-                socket
-                    .emitWithAck("perplexity_ask", previousMessages, {
-                        "version": "2.9",
-                        "source": "default",
-                        "attachments": [],
-                        "language": "en-GB",
-                        "timezone": "Europe/London",
-                        "search_focus": "writing",
-                        "frontend_uuid": uuidv4(),
-                        "mode": "concise",
-                        "is_related_query": false,
-                        "is_default_related_query": false,
-                        "visitor_id": uuidv4(),
-                        "frontend_context_uuid": uuidv4(),
-                        "prompt_source": "user",
-                        "query_source": "home"
-                    })
-                    .then((response) => {
-                        console.log(`[POST /v1/messages] Received response from Perplexity.ai: ${JSON.stringify(response)}`);
-                        res.write(createEvent("content_block_stop", { type: "content_block_stop", index: 0 }));
-                        res.write(
-                            createEvent("message_delta", {
-                                type: "message_delta",
-                                delta: { stop_reason: "end_turn", stop_sequence: null },
-                                usage: { output_tokens: 12 },
-                            })
-                        );
-                        res.write(createEvent("message_stop", { type: "message_stop" }));
-
-                        res.end();
-                    }).catch((error) => {
-                        if(error.message != "socket has been disconnected"){
-                            console.log(`[POST /v1/messages] Error in WebSocket communication: ${error}`);
-                        }
-                    });
-            });
-            socket.onAny((event, ...args) => {
-                console.log(`[POST /v1/messages] Received WebSocket event: ${event}`);
-            });
-            socket.on("query_progress", (data) => {
-                if(data.text){
-                    var text = JSON.parse(data.text)
-                    var chunk = text.chunks[text.chunks.length - 1];
-                    if(chunk){
-                        console.log(`[POST /v1/messages] Received chunk: ${chunk}`);
-                        chunkJSON = JSON.stringify({
-                            type: "content_block_delta",
-                            index: 0,
-                            delta: { type: "text_delta", text: chunk },
-                        });
-                        res.write(createEvent("content_block_delta", chunkJSON));
-                    }
-                }
-            });
-            socket.on("disconnect", function () {
-                console.log("[POST /v1/messages] WebSocket disconnected from Perplexity.ai");
-            });
-            socket.on("error", (error) => {
-                console.log(`[POST /v1/messages] WebSocket error: ${error}`);
-                chunkJSON = JSON.stringify({
-                    type: "content_block_delta",
-                    index: 0,
-                    delta: { type: "text_delta", text: "Error occured while fetching output 输出时出现错误\nPlease refer to the log for more information 请查看日志以获取更多信息" },
-                });
-                res.write(createEvent("content_block_delta", chunkJSON));
-                res.write(createEvent("content_block_stop", { type: "content_block_stop", index: 0 }));
-                res.write(
-                    createEvent("message_delta", {
-                        type: "message_delta",
-                        delta: { stop_reason: "end_turn", stop_sequence: null },
-                        usage: { output_tokens: 12 },
-                    })
-                );
-                res.write(createEvent("message_stop", { type: "message_stop" }));
-
-                res.end();
-            });
-            socket.on("connect_error", function (error) {
-                console.log(`[POST /v1/messages] WebSocket connection error: ${error}`);
-                chunkJSON = JSON.stringify({
-                    type: "content_block_delta",
-                    index: 0,
-                    delta: { type: "text_delta", text: "Failed to connect to the Perplexity.ai 连接到Perplexity失败\nPlease refer to the log for more information 请查看日志以获取更多信息" },
-                });
-                res.write(createEvent("content_block_delta", chunkJSON));
-                res.write(createEvent("content_block_stop", { type: "content_block_stop", index: 0 }));
-                res.write(
-                    createEvent("message_delta", {
-                        type: "message_delta",
-                        delta: { stop_reason: "end_turn", stop_sequence: null },
-                        usage: { output_tokens: 12 },
-                    })
-                );
-                res.write(createEvent("message_stop", { type: "message_stop" }));
-
-                res.end();
-            });
-            res.on("close", function () {
-                console.log("[POST /v1/messages] Client closed connection");
-                socket.disconnect();
-            });
-        } else {
-            console.log("[POST /v1/messages] Invalid request: stream is neither true nor false");
-            throw new Error("Invalid request");
+    const contentType = response.headers.get('content-type');
+    const contentEncoding = response.headers.get('content-encoding');
+
+    if (contentType && contentType.includes('text/event-stream')) {
+      // 处理 SSE
+      res.setHeader('Content-Type', 'text/event-stream');
+      res.setHeader('Cache-Control', 'no-cache');
+      res.setHeader('Connection', 'keep-alive');
+
+      const gunzipStream = zlib.createGunzip();
+      const transformStream = new stream.Transform({
+        transform(chunk, encoding, callback) {
+          this.push(chunk);
+          callback();
+        }
+      });
+
+      pipeline(response.body, gunzipStream, transformStream)
+        .then(() => {
+          console.log('Stream processing completed');
+        })
+        .catch((err) => {
+          console.error('Stream processing error:', err);
+        });
+
+      transformStream.pipe(res);
+    } else {
+      // 非流式响应的处理
+      const buffer = await response.buffer();
+
+      let data;
+      if (contentEncoding === 'gzip') {
+        try {
+          data = await gunzip(buffer);
+        } catch (error) {
+          console.error('Gunzip error:', error);
+          throw new Error('Failed to decompress the response');
         }
-    } catch (e) {
-        console.log(`[POST /v1/messages] Error in request processing: ${e}`);
-        res.write(JSON.stringify({ error: e.message }));
-        res.end();
-        return;
+      } else {
+        data = buffer;
+      }
+
+      res.send(data);
     }
-});
 
-// 处理 /ai/v1/messages 路由
-app.post("/ai/v1/messages", (req, res) => {
-    console.log("[POST /ai/v1/messages] Received request, forwarding to /v1/messages");
-    app.handle(req, res, req.next);
-});
+    logRequest(req, response.status, `Request forwarded successfully for model: ${model}`);
+  } catch (error) {
+    console.error('Request error:', error);
+    res.status(500).json({
+      type: "error",
+      error: {
+        type: "internal_server_error",
+        message: "An unexpected error occurred while processing your request."
+      }
+    });
+    logRequest(req, 500, `Error: ${error.message}`);
+  }
+}
 
-// handle other
-app.use((req, res, next) => {
-    console.log(`[${new Date().toISOString()}] Received request for ${req.path} - returning 404`);
-    res.status(404).send("Not Found");
-});
+function handleOptions(res) {
+  res.status(204);
+  res.setHeader('Access-Control-Allow-Origin', '*');
+  res.setHeader('Access-Control-Allow-Methods', 'POST, GET, OPTIONS');
+  res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization, x-api-key, anthropic-version, model');
+  res.end();
+}
+
+const app = express();
 
-app.listen(port, () => {
-    console.log(`[${new Date().toISOString()}] Perplexity proxy listening on port ${port}`);
+// 增加 body-parser 的限制
+app.use(bodyParser.json({ limit: '50mb' }));
+app.use(bodyParser.urlencoded({ limit: '50mb', extended: true }));
+
+// 根路由处理
+app.get('/', (req, res) => {
+  res.status(200).send('GCP VertexAI For Claude Proxy');
 });
 
-// eventStream util
-function createEvent(event, data) {
-    // if data is object, stringify it
-    if (typeof data === "object") {
-        data = JSON.stringify(data);
-    }
-    return `event: ${event}\ndata: ${data}\n\n`;
-}
+app.all('/ai/v1/messages', handleRequest);
+
+const PORT = 8080;
+app.listen(PORT, () => {
+  console.log(`Server is running on port ${PORT}`);
+});