diff --git a/src/endpoints/backends/chat-completions.js b/src/endpoints/backends/chat-completions.js index 7c340ace8..98cdd8f00 100644 --- a/src/endpoints/backends/chat-completions.js +++ b/src/endpoints/backends/chat-completions.js @@ -1,6 +1,5 @@ import process from 'node:process'; import util from 'node:util'; -import crypto from 'node:crypto'; import express from 'express'; import fetch from 'node-fetch'; @@ -44,6 +43,7 @@ import { webTokenizers, getWebTokenizer, } from '../tokenizers.js'; +import { getVertexAIAuth } from '../google.js'; const API_OPENAI = 'https://api.openai.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_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} 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} 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. diff --git a/src/endpoints/google.js b/src/endpoints/google.js index a42e15af7..eb32faa61 100644 --- a/src/endpoints/google.js +++ b/src/endpoints/google.js @@ -11,7 +11,7 @@ const API_MAKERSUITE = 'https://generativelanguage.googleapis.com'; const API_VERTEX_AI = 'https://us-central1-aiplatform.googleapis.com'; // Vertex AI authentication helper functions -async function getVertexAIAuth(request) { +export async function getVertexAIAuth(request) { const authMode = request.body.vertexai_auth_mode || 'express'; if (request.body.reverse_proxy) { @@ -64,7 +64,7 @@ async function getVertexAIAuth(request) { * @param {object} serviceAccount Service account JSON object * @returns {Promise} JWT token */ -async function generateJWTToken(serviceAccount) { +export async function generateJWTToken(serviceAccount) { const now = Math.floor(Date.now() / 1000); const expiry = now + 3600; // 1 hour @@ -93,7 +93,7 @@ async function generateJWTToken(serviceAccount) { return `${signatureInput}.${signature}`; } -async function getAccessToken(jwtToken) { +export async function getAccessToken(jwtToken) { const response = await fetch('https://oauth2.googleapis.com/token', { method: 'POST', headers: { 'Content-Type': 'application/x-www-form-urlencoded' },