Refactor Google authentication functions for Vertex AI

- Exported `getVertexAIAuth`, `generateJWTToken`, and `getAccessToken` functions from google.js for better modularity.
- Removed redundant JWT token generation and access token retrieval logic from chat-completions.js, utilizing the newly exported functions instead.
- Improved code organization and maintainability by centralizing authentication logic.
This commit is contained in:
InterestingDarkness
2025-05-27 00:31:43 +08:00
parent 7fcd40c829
commit 1c048a6c79
2 changed files with 4 additions and 116 deletions

View File

@@ -1,6 +1,5 @@
import process from 'node:process'; import process from 'node:process';
import util from 'node:util'; import util from 'node:util';
import crypto from 'node:crypto';
import express from 'express'; import express from 'express';
import fetch from 'node-fetch'; import fetch from 'node-fetch';
@@ -44,6 +43,7 @@ import {
webTokenizers, webTokenizers,
getWebTokenizer, getWebTokenizer,
} from '../tokenizers.js'; } from '../tokenizers.js';
import { getVertexAIAuth } from '../google.js';
const API_OPENAI = 'https://api.openai.com/v1'; const API_OPENAI = 'https://api.openai.com/v1';
const API_CLAUDE = 'https://api.anthropic.com/v1'; const API_CLAUDE = 'https://api.anthropic.com/v1';
@@ -61,123 +61,11 @@ const API_DEEPSEEK = 'https://api.deepseek.com/beta';
const API_XAI = 'https://api.x.ai/v1'; const API_XAI = 'https://api.x.ai/v1';
const API_POLLINATIONS = 'https://text.pollinations.ai/openai'; const API_POLLINATIONS = 'https://text.pollinations.ai/openai';
/**
* Generates a JWT token for Google Cloud authentication using service account credentials.
* @param {object} serviceAccount Service account JSON object
* @returns {Promise<string>} JWT token
*/
async function generateJWTToken(serviceAccount) {
const now = Math.floor(Date.now() / 1000);
const expiry = now + 3600; // 1 hour
const header = {
alg: 'RS256',
typ: 'JWT',
};
const payload = {
iss: serviceAccount.client_email,
scope: 'https://www.googleapis.com/auth/cloud-platform',
aud: 'https://oauth2.googleapis.com/token',
iat: now,
exp: expiry,
};
const headerBase64 = Buffer.from(JSON.stringify(header)).toString('base64url');
const payloadBase64 = Buffer.from(JSON.stringify(payload)).toString('base64url');
const signatureInput = `${headerBase64}.${payloadBase64}`;
// Create signature using private key
const sign = crypto.createSign('RSA-SHA256');
sign.update(signatureInput);
const signature = sign.sign(serviceAccount.private_key, 'base64url');
return `${signatureInput}.${signature}`;
}
/**
* Gets an access token from Google OAuth2 using JWT assertion.
* @param {string} jwtToken JWT token
* @returns {Promise<string>} Access token
*/
async function getAccessToken(jwtToken) {
const response = await fetch('https://oauth2.googleapis.com/token', {
method: 'POST',
headers: {
'Content-Type': 'application/x-www-form-urlencoded',
},
body: new URLSearchParams({
grant_type: 'urn:ietf:params:oauth:grant-type:jwt-bearer',
assertion: jwtToken,
}),
});
if (!response.ok) {
throw new Error(`Failed to get access token: ${response.status} ${response.statusText}`);
}
const data = await response.json();
return data.access_token;
}
/**
* Gets authentication for Vertex AI - either API key or service account token.
* @param {object} request Express request
* @returns {Promise<{authHeader: string, authType: string}>} Authentication header and type
*/
async function getVertexAIAuth(request) {
// Get the authentication mode from frontend
const authMode = request.body.vertexai_auth_mode || 'express';
// Check if using reverse proxy
if (request.body.reverse_proxy) {
return {
authHeader: `Bearer ${request.body.proxy_password}`,
authType: 'proxy',
};
}
if (authMode === 'express') {
// Express mode: use API key
const apiKey = readSecret(request.user.directories, SECRET_KEYS.VERTEXAI);
if (apiKey) {
return {
authHeader: `Bearer ${apiKey}`,
authType: 'express',
};
}
throw new Error('API key is required for Vertex AI Express mode');
} else if (authMode === 'full') {
// Full mode: use service account JSON
// First try to read from backend secret storage
let serviceAccountJson = readSecret(request.user.directories, SECRET_KEYS.VERTEXAI_SERVICE_ACCOUNT);
// If not found in secrets, fall back to request body (for backward compatibility)
if (!serviceAccountJson) {
serviceAccountJson = request.body.vertexai_service_account_json;
}
if (serviceAccountJson) {
try {
const serviceAccount = JSON.parse(serviceAccountJson);
const jwtToken = await generateJWTToken(serviceAccount);
const accessToken = await getAccessToken(jwtToken);
return {
authHeader: `Bearer ${accessToken}`,
authType: 'full',
};
} catch (error) {
console.error('Failed to authenticate with service account:', error);
throw new Error(`Service account authentication failed: ${error.message}`);
}
}
throw new Error('Service Account JSON is required for Vertex AI Full mode');
}
throw new Error(`Unsupported Vertex AI authentication mode: ${authMode}`);
}
/** /**
* Applies a post-processing step to the generated messages. * Applies a post-processing step to the generated messages.

View File

@@ -11,7 +11,7 @@ const API_MAKERSUITE = 'https://generativelanguage.googleapis.com';
const API_VERTEX_AI = 'https://us-central1-aiplatform.googleapis.com'; const API_VERTEX_AI = 'https://us-central1-aiplatform.googleapis.com';
// Vertex AI authentication helper functions // Vertex AI authentication helper functions
async function getVertexAIAuth(request) { export async function getVertexAIAuth(request) {
const authMode = request.body.vertexai_auth_mode || 'express'; const authMode = request.body.vertexai_auth_mode || 'express';
if (request.body.reverse_proxy) { if (request.body.reverse_proxy) {
@@ -64,7 +64,7 @@ async function getVertexAIAuth(request) {
* @param {object} serviceAccount Service account JSON object * @param {object} serviceAccount Service account JSON object
* @returns {Promise<string>} JWT token * @returns {Promise<string>} JWT token
*/ */
async function generateJWTToken(serviceAccount) { export async function generateJWTToken(serviceAccount) {
const now = Math.floor(Date.now() / 1000); const now = Math.floor(Date.now() / 1000);
const expiry = now + 3600; // 1 hour const expiry = now + 3600; // 1 hour
@@ -93,7 +93,7 @@ async function generateJWTToken(serviceAccount) {
return `${signatureInput}.${signature}`; return `${signatureInput}.${signature}`;
} }
async function getAccessToken(jwtToken) { export async function getAccessToken(jwtToken) {
const response = await fetch('https://oauth2.googleapis.com/token', { const response = await fetch('https://oauth2.googleapis.com/token', {
method: 'POST', method: 'POST',
headers: { 'Content-Type': 'application/x-www-form-urlencoded' }, headers: { 'Content-Type': 'application/x-www-form-urlencoded' },