|
|
const redisClient = require('../models/redis') |
|
|
const { v4: uuidv4 } = require('uuid') |
|
|
const crypto = require('crypto') |
|
|
const config = require('../../config/config') |
|
|
const logger = require('../utils/logger') |
|
|
|
|
|
|
|
|
const ALGORITHM = 'aes-256-cbc' |
|
|
const IV_LENGTH = 16 |
|
|
|
|
|
|
|
|
const ENCRYPTION_SALT = config.security?.azureOpenaiSalt || 'azure-openai-account-default-salt' |
|
|
|
|
|
class EncryptionKeyManager { |
|
|
constructor() { |
|
|
this.keyCache = new Map() |
|
|
this.keyRotationInterval = 24 * 60 * 60 * 1000 |
|
|
} |
|
|
|
|
|
getKey(version = 'current') { |
|
|
const cached = this.keyCache.get(version) |
|
|
if (cached && Date.now() - cached.timestamp < this.keyRotationInterval) { |
|
|
return cached.key |
|
|
} |
|
|
|
|
|
|
|
|
const key = crypto.scryptSync(config.security.encryptionKey, ENCRYPTION_SALT, 32) |
|
|
this.keyCache.set(version, { |
|
|
key, |
|
|
timestamp: Date.now() |
|
|
}) |
|
|
|
|
|
logger.debug('🔑 Azure OpenAI encryption key generated/refreshed') |
|
|
return key |
|
|
} |
|
|
|
|
|
|
|
|
cleanup() { |
|
|
const now = Date.now() |
|
|
for (const [version, cached] of this.keyCache.entries()) { |
|
|
if (now - cached.timestamp > this.keyRotationInterval) { |
|
|
this.keyCache.delete(version) |
|
|
} |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
const encryptionKeyManager = new EncryptionKeyManager() |
|
|
|
|
|
|
|
|
setInterval( |
|
|
() => { |
|
|
encryptionKeyManager.cleanup() |
|
|
}, |
|
|
60 * 60 * 1000 |
|
|
) |
|
|
|
|
|
|
|
|
function generateEncryptionKey() { |
|
|
return encryptionKeyManager.getKey() |
|
|
} |
|
|
|
|
|
|
|
|
const AZURE_OPENAI_ACCOUNT_KEY_PREFIX = 'azure_openai:account:' |
|
|
const SHARED_AZURE_OPENAI_ACCOUNTS_KEY = 'shared_azure_openai_accounts' |
|
|
const ACCOUNT_SESSION_MAPPING_PREFIX = 'azure_openai_session_account_mapping:' |
|
|
|
|
|
|
|
|
function encrypt(text) { |
|
|
if (!text) { |
|
|
return '' |
|
|
} |
|
|
const key = generateEncryptionKey() |
|
|
const iv = crypto.randomBytes(IV_LENGTH) |
|
|
const cipher = crypto.createCipheriv(ALGORITHM, key, iv) |
|
|
let encrypted = cipher.update(text) |
|
|
encrypted = Buffer.concat([encrypted, cipher.final()]) |
|
|
return `${iv.toString('hex')}:${encrypted.toString('hex')}` |
|
|
} |
|
|
|
|
|
|
|
|
function decrypt(text) { |
|
|
if (!text) { |
|
|
return '' |
|
|
} |
|
|
|
|
|
try { |
|
|
const key = generateEncryptionKey() |
|
|
|
|
|
const ivHex = text.substring(0, 32) |
|
|
const encryptedHex = text.substring(33) |
|
|
|
|
|
if (ivHex.length !== 32 || !encryptedHex) { |
|
|
throw new Error('Invalid encrypted text format') |
|
|
} |
|
|
|
|
|
const iv = Buffer.from(ivHex, 'hex') |
|
|
const encryptedText = Buffer.from(encryptedHex, 'hex') |
|
|
const decipher = crypto.createDecipheriv(ALGORITHM, key, iv) |
|
|
let decrypted = decipher.update(encryptedText) |
|
|
decrypted = Buffer.concat([decrypted, decipher.final()]) |
|
|
const result = decrypted.toString() |
|
|
|
|
|
return result |
|
|
} catch (error) { |
|
|
logger.error('Azure OpenAI decryption error:', error.message) |
|
|
return '' |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function createAccount(accountData) { |
|
|
const accountId = uuidv4() |
|
|
const now = new Date().toISOString() |
|
|
|
|
|
const account = { |
|
|
id: accountId, |
|
|
name: accountData.name, |
|
|
description: accountData.description || '', |
|
|
accountType: accountData.accountType || 'shared', |
|
|
groupId: accountData.groupId || null, |
|
|
priority: accountData.priority || 50, |
|
|
|
|
|
azureEndpoint: accountData.azureEndpoint || '', |
|
|
apiVersion: accountData.apiVersion || '2024-02-01', |
|
|
deploymentName: accountData.deploymentName || 'gpt-4', |
|
|
apiKey: encrypt(accountData.apiKey || ''), |
|
|
|
|
|
supportedModels: JSON.stringify( |
|
|
accountData.supportedModels || ['gpt-4', 'gpt-4-turbo', 'gpt-35-turbo', 'gpt-35-turbo-16k'] |
|
|
), |
|
|
|
|
|
|
|
|
|
|
|
subscriptionExpiresAt: accountData.subscriptionExpiresAt || null, |
|
|
|
|
|
|
|
|
isActive: accountData.isActive !== false ? 'true' : 'false', |
|
|
status: 'active', |
|
|
schedulable: accountData.schedulable !== false ? 'true' : 'false', |
|
|
createdAt: now, |
|
|
updatedAt: now |
|
|
} |
|
|
|
|
|
|
|
|
if (accountData.proxy) { |
|
|
account.proxy = |
|
|
typeof accountData.proxy === 'string' ? accountData.proxy : JSON.stringify(accountData.proxy) |
|
|
} |
|
|
|
|
|
const client = redisClient.getClientSafe() |
|
|
await client.hset(`${AZURE_OPENAI_ACCOUNT_KEY_PREFIX}${accountId}`, account) |
|
|
|
|
|
|
|
|
if (account.accountType === 'shared') { |
|
|
await client.sadd(SHARED_AZURE_OPENAI_ACCOUNTS_KEY, accountId) |
|
|
} |
|
|
|
|
|
logger.info(`Created Azure OpenAI account: ${accountId}`) |
|
|
return account |
|
|
} |
|
|
|
|
|
|
|
|
async function getAccount(accountId) { |
|
|
const client = redisClient.getClientSafe() |
|
|
const accountData = await client.hgetall(`${AZURE_OPENAI_ACCOUNT_KEY_PREFIX}${accountId}`) |
|
|
|
|
|
if (!accountData || Object.keys(accountData).length === 0) { |
|
|
return null |
|
|
} |
|
|
|
|
|
|
|
|
if (accountData.apiKey) { |
|
|
accountData.apiKey = decrypt(accountData.apiKey) |
|
|
} |
|
|
|
|
|
|
|
|
if (accountData.proxy && typeof accountData.proxy === 'string') { |
|
|
try { |
|
|
accountData.proxy = JSON.parse(accountData.proxy) |
|
|
} catch (e) { |
|
|
accountData.proxy = null |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if (accountData.supportedModels && typeof accountData.supportedModels === 'string') { |
|
|
try { |
|
|
accountData.supportedModels = JSON.parse(accountData.supportedModels) |
|
|
} catch (e) { |
|
|
accountData.supportedModels = ['gpt-4', 'gpt-35-turbo'] |
|
|
} |
|
|
} |
|
|
|
|
|
return accountData |
|
|
} |
|
|
|
|
|
|
|
|
async function updateAccount(accountId, updates) { |
|
|
const existingAccount = await getAccount(accountId) |
|
|
if (!existingAccount) { |
|
|
throw new Error('Account not found') |
|
|
} |
|
|
|
|
|
updates.updatedAt = new Date().toISOString() |
|
|
|
|
|
|
|
|
if (updates.apiKey) { |
|
|
updates.apiKey = encrypt(updates.apiKey) |
|
|
} |
|
|
|
|
|
|
|
|
if (updates.proxy) { |
|
|
updates.proxy = |
|
|
typeof updates.proxy === 'string' ? updates.proxy : JSON.stringify(updates.proxy) |
|
|
} |
|
|
|
|
|
|
|
|
if (updates.supportedModels) { |
|
|
updates.supportedModels = |
|
|
typeof updates.supportedModels === 'string' |
|
|
? updates.supportedModels |
|
|
: JSON.stringify(updates.supportedModels) |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
if (updates.subscriptionExpiresAt !== undefined) { |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
const client = redisClient.getClientSafe() |
|
|
if (updates.accountType && updates.accountType !== existingAccount.accountType) { |
|
|
if (updates.accountType === 'shared') { |
|
|
await client.sadd(SHARED_AZURE_OPENAI_ACCOUNTS_KEY, accountId) |
|
|
} else { |
|
|
await client.srem(SHARED_AZURE_OPENAI_ACCOUNTS_KEY, accountId) |
|
|
} |
|
|
} |
|
|
|
|
|
await client.hset(`${AZURE_OPENAI_ACCOUNT_KEY_PREFIX}${accountId}`, updates) |
|
|
|
|
|
logger.info(`Updated Azure OpenAI account: ${accountId}`) |
|
|
|
|
|
|
|
|
const updatedAccount = { ...existingAccount, ...updates } |
|
|
|
|
|
|
|
|
if (updatedAccount.proxy && typeof updatedAccount.proxy === 'string') { |
|
|
try { |
|
|
updatedAccount.proxy = JSON.parse(updatedAccount.proxy) |
|
|
} catch (e) { |
|
|
updatedAccount.proxy = null |
|
|
} |
|
|
} |
|
|
|
|
|
return updatedAccount |
|
|
} |
|
|
|
|
|
|
|
|
async function deleteAccount(accountId) { |
|
|
|
|
|
const accountGroupService = require('./accountGroupService') |
|
|
await accountGroupService.removeAccountFromAllGroups(accountId) |
|
|
|
|
|
const client = redisClient.getClientSafe() |
|
|
const accountKey = `${AZURE_OPENAI_ACCOUNT_KEY_PREFIX}${accountId}` |
|
|
|
|
|
|
|
|
await client.del(accountKey) |
|
|
|
|
|
|
|
|
await client.srem(SHARED_AZURE_OPENAI_ACCOUNTS_KEY, accountId) |
|
|
|
|
|
logger.info(`Deleted Azure OpenAI account: ${accountId}`) |
|
|
return true |
|
|
} |
|
|
|
|
|
|
|
|
async function getAllAccounts() { |
|
|
const client = redisClient.getClientSafe() |
|
|
const keys = await client.keys(`${AZURE_OPENAI_ACCOUNT_KEY_PREFIX}*`) |
|
|
|
|
|
if (!keys || keys.length === 0) { |
|
|
return [] |
|
|
} |
|
|
|
|
|
const accounts = [] |
|
|
for (const key of keys) { |
|
|
const accountData = await client.hgetall(key) |
|
|
if (accountData && Object.keys(accountData).length > 0) { |
|
|
|
|
|
delete accountData.apiKey |
|
|
|
|
|
|
|
|
if (accountData.proxy && typeof accountData.proxy === 'string') { |
|
|
try { |
|
|
accountData.proxy = JSON.parse(accountData.proxy) |
|
|
} catch (e) { |
|
|
accountData.proxy = null |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if (accountData.supportedModels && typeof accountData.supportedModels === 'string') { |
|
|
try { |
|
|
accountData.supportedModels = JSON.parse(accountData.supportedModels) |
|
|
} catch (e) { |
|
|
accountData.supportedModels = ['gpt-4', 'gpt-35-turbo'] |
|
|
} |
|
|
} |
|
|
|
|
|
accounts.push({ |
|
|
...accountData, |
|
|
isActive: accountData.isActive === 'true', |
|
|
schedulable: accountData.schedulable !== 'false', |
|
|
|
|
|
|
|
|
expiresAt: accountData.subscriptionExpiresAt || null, |
|
|
platform: 'azure-openai' |
|
|
}) |
|
|
} |
|
|
} |
|
|
|
|
|
return accounts |
|
|
} |
|
|
|
|
|
|
|
|
async function getSharedAccounts() { |
|
|
const client = redisClient.getClientSafe() |
|
|
const accountIds = await client.smembers(SHARED_AZURE_OPENAI_ACCOUNTS_KEY) |
|
|
|
|
|
if (!accountIds || accountIds.length === 0) { |
|
|
return [] |
|
|
} |
|
|
|
|
|
const accounts = [] |
|
|
for (const accountId of accountIds) { |
|
|
const account = await getAccount(accountId) |
|
|
if (account && account.isActive === 'true') { |
|
|
accounts.push(account) |
|
|
} |
|
|
} |
|
|
|
|
|
return accounts |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
function isSubscriptionExpired(account) { |
|
|
if (!account.subscriptionExpiresAt) { |
|
|
return false |
|
|
} |
|
|
const expiryDate = new Date(account.subscriptionExpiresAt) |
|
|
return expiryDate <= new Date() |
|
|
} |
|
|
|
|
|
|
|
|
async function selectAvailableAccount(sessionId = null) { |
|
|
|
|
|
if (sessionId) { |
|
|
const client = redisClient.getClientSafe() |
|
|
const mappingKey = `${ACCOUNT_SESSION_MAPPING_PREFIX}${sessionId}` |
|
|
const accountId = await client.get(mappingKey) |
|
|
|
|
|
if (accountId) { |
|
|
const account = await getAccount(accountId) |
|
|
if (account && account.isActive === 'true' && account.schedulable === 'true') { |
|
|
logger.debug(`Reusing Azure OpenAI account ${accountId} for session ${sessionId}`) |
|
|
return account |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
const sharedAccounts = await getSharedAccounts() |
|
|
|
|
|
|
|
|
const availableAccounts = sharedAccounts.filter((acc) => { |
|
|
|
|
|
if (isSubscriptionExpired(acc)) { |
|
|
logger.debug( |
|
|
`⏰ Skipping expired Azure OpenAI account: ${acc.name}, expired at ${acc.subscriptionExpiresAt}` |
|
|
) |
|
|
return false |
|
|
} |
|
|
|
|
|
return acc.isActive === 'true' && acc.schedulable === 'true' |
|
|
}) |
|
|
|
|
|
if (availableAccounts.length === 0) { |
|
|
throw new Error('No available Azure OpenAI accounts') |
|
|
} |
|
|
|
|
|
|
|
|
availableAccounts.sort((a, b) => (b.priority || 50) - (a.priority || 50)) |
|
|
const selectedAccount = availableAccounts[0] |
|
|
|
|
|
|
|
|
if (sessionId && selectedAccount) { |
|
|
const client = redisClient.getClientSafe() |
|
|
const mappingKey = `${ACCOUNT_SESSION_MAPPING_PREFIX}${sessionId}` |
|
|
await client.setex(mappingKey, 3600, selectedAccount.id) |
|
|
} |
|
|
|
|
|
logger.debug(`Selected Azure OpenAI account: ${selectedAccount.id}`) |
|
|
return selectedAccount |
|
|
} |
|
|
|
|
|
|
|
|
async function updateAccountUsage(accountId, tokens) { |
|
|
const client = redisClient.getClientSafe() |
|
|
const now = new Date().toISOString() |
|
|
|
|
|
|
|
|
await client.hincrby(`${AZURE_OPENAI_ACCOUNT_KEY_PREFIX}${accountId}`, 'totalTokensUsed', tokens) |
|
|
await client.hset(`${AZURE_OPENAI_ACCOUNT_KEY_PREFIX}${accountId}`, 'lastUsedAt', now) |
|
|
|
|
|
logger.debug(`Updated Azure OpenAI account ${accountId} usage: ${tokens} tokens`) |
|
|
} |
|
|
|
|
|
|
|
|
async function healthCheckAccount(accountId) { |
|
|
try { |
|
|
const account = await getAccount(accountId) |
|
|
if (!account) { |
|
|
return { id: accountId, status: 'error', message: 'Account not found' } |
|
|
} |
|
|
|
|
|
|
|
|
if (!account.azureEndpoint || !account.apiKey || !account.deploymentName) { |
|
|
return { |
|
|
id: accountId, |
|
|
status: 'error', |
|
|
message: 'Incomplete configuration' |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
return { |
|
|
id: accountId, |
|
|
status: 'healthy', |
|
|
message: 'Account is configured correctly' |
|
|
} |
|
|
} catch (error) { |
|
|
logger.error(`Health check failed for Azure OpenAI account ${accountId}:`, error) |
|
|
return { |
|
|
id: accountId, |
|
|
status: 'error', |
|
|
message: error.message |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function performHealthChecks() { |
|
|
const accounts = await getAllAccounts() |
|
|
const results = [] |
|
|
|
|
|
for (const account of accounts) { |
|
|
const result = await healthCheckAccount(account.id) |
|
|
results.push(result) |
|
|
} |
|
|
|
|
|
return results |
|
|
} |
|
|
|
|
|
|
|
|
async function toggleSchedulable(accountId) { |
|
|
const account = await getAccount(accountId) |
|
|
if (!account) { |
|
|
throw new Error('Account not found') |
|
|
} |
|
|
|
|
|
const newSchedulable = account.schedulable === 'true' ? 'false' : 'true' |
|
|
await updateAccount(accountId, { schedulable: newSchedulable }) |
|
|
|
|
|
return { |
|
|
id: accountId, |
|
|
schedulable: newSchedulable === 'true' |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function migrateApiKeysForAzureSupport() { |
|
|
const client = redisClient.getClientSafe() |
|
|
const apiKeyIds = await client.smembers('api_keys') |
|
|
|
|
|
let migratedCount = 0 |
|
|
for (const keyId of apiKeyIds) { |
|
|
const keyData = await client.hgetall(`api_key:${keyId}`) |
|
|
if (keyData && !keyData.azureOpenaiAccountId) { |
|
|
|
|
|
await client.hset(`api_key:${keyId}`, 'azureOpenaiAccountId', '') |
|
|
migratedCount++ |
|
|
} |
|
|
} |
|
|
|
|
|
logger.info(`Migrated ${migratedCount} API keys for Azure OpenAI support`) |
|
|
return migratedCount |
|
|
} |
|
|
|
|
|
module.exports = { |
|
|
createAccount, |
|
|
getAccount, |
|
|
updateAccount, |
|
|
deleteAccount, |
|
|
getAllAccounts, |
|
|
getSharedAccounts, |
|
|
selectAvailableAccount, |
|
|
updateAccountUsage, |
|
|
healthCheckAccount, |
|
|
performHealthChecks, |
|
|
toggleSchedulable, |
|
|
migrateApiKeysForAzureSupport, |
|
|
encrypt, |
|
|
decrypt |
|
|
} |
|
|
|