|
|
|
|
|
|
|
import { logger } from "./logger"; |
|
import crypto from "crypto"; |
|
|
|
|
|
type KeySchema = { |
|
|
|
key: string; |
|
|
|
isTrial?: boolean; |
|
|
|
isGpt4?: boolean; |
|
}; |
|
|
|
|
|
export type Key = KeySchema & { |
|
|
|
isDisabled?: boolean; |
|
|
|
softLimit?: number; |
|
|
|
hardLimit?: number; |
|
|
|
systemHardLimit?: number; |
|
|
|
usage?: number; |
|
|
|
promptCount: number; |
|
|
|
lastUsed: number; |
|
|
|
hash: string; |
|
}; |
|
|
|
const keyPool: Key[] = []; |
|
|
|
function init() { |
|
const keyString = process.env.OPENAI_KEY; |
|
if (!keyString?.trim()) { |
|
throw new Error("OPENAI_KEY environment variable is not set"); |
|
} |
|
let keyList: KeySchema[]; |
|
try { |
|
const decoded = Buffer.from(keyString, "base64").toString(); |
|
keyList = JSON.parse(decoded) as KeySchema[]; |
|
} catch (err) { |
|
logger.info("OPENAI_KEY is not base64-encoded JSON, assuming bare key"); |
|
|
|
keyList = [{ key: keyString, isTrial: false, isGpt4: true }]; |
|
} |
|
for (const key of keyList) { |
|
const newKey = { |
|
...key, |
|
isDisabled: false, |
|
softLimit: 0, |
|
hardLimit: 0, |
|
systemHardLimit: 0, |
|
usage: 0, |
|
lastUsed: 0, |
|
promptCount: 0, |
|
hash: crypto |
|
.createHash("sha256") |
|
.update(key.key) |
|
.digest("hex") |
|
.slice(0, 6), |
|
}; |
|
keyPool.push(newKey); |
|
|
|
logger.info({ key: newKey.hash }, "Key added"); |
|
} |
|
|
|
} |
|
|
|
function list() { |
|
return keyPool.map((key) => ({ |
|
...key, |
|
key: undefined, |
|
})); |
|
} |
|
|
|
function disable(key: Key) { |
|
const keyFromPool = keyPool.find((k) => k.key === key.key)!; |
|
if (keyFromPool.isDisabled) return; |
|
keyFromPool.isDisabled = true; |
|
logger.warn({ key: key.hash }, "Key disabled"); |
|
} |
|
|
|
function anyAvailable() { |
|
return keyPool.some((key) => !key.isDisabled); |
|
} |
|
|
|
function get(model: string) { |
|
const needsGpt4Key = model.startsWith("gpt-4"); |
|
const availableKeys = keyPool.filter( |
|
(key) => !key.isDisabled && (!needsGpt4Key || key.isGpt4) |
|
); |
|
if (availableKeys.length === 0) { |
|
let message = "No keys available. Please add more keys."; |
|
if (needsGpt4Key) { |
|
message = |
|
"No GPT-4 keys available. Please add more keys or use a non-GPT-4 model."; |
|
} |
|
logger.error(message); |
|
throw new Error(message); |
|
} |
|
|
|
|
|
const trialKeys = availableKeys.filter((key) => key.isTrial); |
|
if (trialKeys.length > 0) { |
|
logger.info({ key: trialKeys[0].hash }, "Using trial key"); |
|
trialKeys[0].lastUsed = Date.now(); |
|
return trialKeys[0]; |
|
} |
|
|
|
|
|
const oldestKey = availableKeys.sort((a, b) => a.lastUsed - b.lastUsed)[0]; |
|
logger.info({ key: oldestKey.hash }, "Assigning key to request."); |
|
oldestKey.lastUsed = Date.now(); |
|
return { ...oldestKey }; |
|
} |
|
|
|
function incrementPrompt(keyHash?: string) { |
|
if (!keyHash) return; |
|
const key = keyPool.find((k) => k.hash === keyHash)!; |
|
key.promptCount++; |
|
} |
|
|
|
export const keys = { init, list, get, anyAvailable, disable, incrementPrompt }; |
|
|