app.js

const express = require('express');
const fetch = require('node-fetch');

const MODEL = 'claude-3-5-sonnet@20240620';

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
};

// 日志函数
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}`);
}

async function getAccessToken() {
  const now = Date.now() / 1000;

  if (tokenCache.accessToken && now < tokenCache.expiry - 120) {
    return tokenCache.accessToken;
  }

  if (tokenCache.refreshPromise) {
    await tokenCache.refreshPromise;
    return tokenCache.accessToken;
  }

  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) {
  return `https://${location}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${location}/publishers/anthropic/models/${MODEL}:streamRawPredict`;
}

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();
  const apiUrl = constructApiUrl(location);

  let requestBody = req.body;
  
  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'
      },
      body: JSON.stringify(requestBody)
    });

    res.status(response.status);
    response.headers.forEach((value, name) => {
      res.setHeader(name, value);
    });
    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');

    logRequest(req, response.status, 'Request forwarded successfully');
    response.body.pipe(res);
  } catch (error) {
    console.error('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();
app.use(express.json());

app.all('/ai/v1/messages', handleRequest);

const PORT = 8080;
app.listen(PORT, () => {
  console.log(`Server is running on port ${PORT}`);
});