mirror of
https://github.com/SillyTavern/SillyTavern.git
synced 2025-06-05 21:59:27 +02:00
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:
@@ -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.
|
||||||
|
@@ -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' },
|
||||||
|
Reference in New Issue
Block a user