Update app.js

app.js

@@ -1,236 +1,293 @@
-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 express = require("express");
+const { io } = require("socket.io-client");
+const { v4: uuidv4 } = require("uuid");
+const { ProxyAgent } = require("proxy-agent");
+const agent = new ProxyAgent();
 
-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}`);
-}
+const app = express();
+const port = process.env.PORT || 8081;
 
-async function getAccessToken() {
-  const now = Date.now() / 1000;
+// 从环境变量中获取 accessToken
+const accessToken = process.env.PPLX_KEY;
 
-  if (tokenCache.accessToken && now < tokenCache.expiry - 120) {
-    return tokenCache.accessToken;
-  }
+console.log(`Server starting. Access token: ${accessToken ? "Set" : "Not set"}`);
 
-  if (tokenCache.refreshPromise) {
-    await tokenCache.refreshPromise;
-    return tokenCache.accessToken;
-  }
+// 添加一个中间件来解析 JSON 请求体
+app.use(express.json());
 
-  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;
-    }
-  })();
-
-  await tokenCache.refreshPromise;
-  return tokenCache.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`;
-}
+// 添加一个中间件来记录所有请求
+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();
+});
 
-function formatModelName(model) {
-  if (model === 'claude-3-5-sonnet-20240620') {
-    return 'claude-3-5-sonnet@20240620';
-  }
-  return model;
-}
+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/",
+    },
+};
 
-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);
+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}`);
     
-    for (const [key, value] of response.headers.entries()) {
-      if (key.toLowerCase() !== 'content-encoding' && key.toLowerCase() !== 'content-length') {
-        res.setHeader(key, value);
-      }
+    if (!clientAuth || clientAuth !== `Bearer ${accessToken}`) {
+        console.log("[POST /v1/messages] Authorization failed. Sending 401 Unauthorized response.");
+        res.status(401).json({ error: "Unauthorized" });
+        return;
     }
+    
+    console.log("[POST /v1/messages] Authorization successful. Proceeding with request processing.");
 
-    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');
-
-    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');
+    // 使用解析后的 JSON 体
+    const jsonBody = req.body;
+    console.log(`[POST /v1/messages] Received request body: ${JSON.stringify(jsonBody)}`);
 
-      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');
+    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");
         }
-      } else {
-        data = buffer;
-      }
-
-      res.send(data);
+    } catch (e) {
+        console.log(`[POST /v1/messages] Error in request processing: ${e}`);
+        res.write(JSON.stringify({ error: e.message }));
+        res.end();
+        return;
     }
+});
 
-    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}`);
-  }
-}
-
-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();
-
-// 增加 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');
+// 处理 /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);
 });
 
-app.all('/ai/v1/messages', handleRequest);
+// 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");
+});
 
-const PORT = 8080;
-app.listen(PORT, () => {
-  console.log(`Server is running on port ${PORT}`);
+app.listen(port, () => {
+    console.log(`[${new Date().toISOString()}] Perplexity proxy listening on port ${port}`);
 });
+
+// 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`;
+}