loads keys on startup

This commit is contained in:
nai-degen 2023-04-08 06:02:59 -05:00 committed by nai-degen
parent ba4aca7608
commit 66b8b6a5d0
7 changed files with 50 additions and 9 deletions

20
.env.example Normal file
View File

@ -0,0 +1,20 @@
# Copy this file to .env and fill in the values.
# Uncomment the following line and replace the value with your own secret key
# to control access to the proxy server
# PROXY_KEY=your-secret-key
# Set your OpenAI API key below.
OPENAI_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
# You can also set a base-64 encoded JSON array of keys matching this schema:
# Trial keys will be prioritized. GPT4-enabled keys are necessary to use GPT-4.
# For example:
# [
# { "key": "your-openai-key-1", "isTrial": true, "isGpt4": false },
# { "key": "your-openai-key-2", "isTrial": false, "isGpt4": false },
# { "key": "your-openai-key-3", "isTrial": false, "isGpt4": true }
# ]
# Encoded in base-64, this would look like:
# OPENAI_KEYS=WwogeyAia2V5IjogInlvdXItb3BlbmFpLWtleS0xIiwgImlzVHJpYWwiOiB0cnVlLCAiaXNHcHQ0IjogZmFsc2UgfSwKIHsgImtleSI6ICJ5b3VyLW9wZW5haS1rZXktMiIsICJpc1RyaWFsIjogZmFsc2UsICJpc0dwdDQiOiBmYWxzZSB9LAogeyAia2V5IjogInlvdXItb3BlbmFpLWtleS0zIiwgImlzVHJpYWwiOiBmYWxzZSwgImlzR3B0NCI6IHRydWUgfQpd

1
.gitignore vendored
View File

@ -1,3 +1,4 @@
.vscode .vscode
.env
build build
node_modules node_modules

9
package-lock.json generated
View File

@ -10,6 +10,7 @@
"license": "MIT", "license": "MIT",
"dependencies": { "dependencies": {
"cors": "^2.8.5", "cors": "^2.8.5",
"dotenv": "^16.0.3",
"express": "^4.18.2", "express": "^4.18.2",
"http-proxy-middleware": "^2.0.6", "http-proxy-middleware": "^2.0.6",
"pino-http": "^8.3.3" "pino-http": "^8.3.3"
@ -495,6 +496,14 @@
"node": ">=0.3.1" "node": ">=0.3.1"
} }
}, },
"node_modules/dotenv": {
"version": "16.0.3",
"resolved": "https://registry.npmjs.org/dotenv/-/dotenv-16.0.3.tgz",
"integrity": "sha512-7GO6HghkA5fYG9TYnNxi14/7K9f5occMlp3zXAuSxn7CKCxt9xbNWG7yF8hTCSUchlfWSe3uLmlPfigevRItzQ==",
"engines": {
"node": ">=12"
}
},
"node_modules/ee-first": { "node_modules/ee-first": {
"version": "1.1.1", "version": "1.1.1",
"resolved": "https://registry.npmjs.org/ee-first/-/ee-first-1.1.1.tgz", "resolved": "https://registry.npmjs.org/ee-first/-/ee-first-1.1.1.tgz",

View File

@ -11,6 +11,7 @@
"license": "MIT", "license": "MIT",
"dependencies": { "dependencies": {
"cors": "^2.8.5", "cors": "^2.8.5",
"dotenv": "^16.0.3",
"express": "^4.18.2", "express": "^4.18.2",
"http-proxy-middleware": "^2.0.6", "http-proxy-middleware": "^2.0.6",
"pino-http": "^8.3.3" "pino-http": "^8.3.3"

View File

@ -36,13 +36,15 @@ const keyPool: Key[] = [];
function init() { function init() {
const keyString = process.env.OPENAI_KEY; const keyString = process.env.OPENAI_KEY;
if (!keyString) { if (!keyString?.trim()) {
throw new Error("OPENAI_KEY environment variable is not set"); throw new Error("OPENAI_KEY environment variable is not set");
} }
let keyList: KeySchema[]; let keyList: KeySchema[];
try { try {
keyList = JSON.parse(Buffer.from(keyString, "base64").toString()); const decoded = Buffer.from(keyString, "base64").toString();
keyList = JSON.parse(decoded) as KeySchema[];
} catch (err) { } catch (err) {
console.log("Key is not base64-encoded JSON, assuming it's a bare key");
// We don't actually know if bare keys are paid/GPT-4 so we assume they are // We don't actually know if bare keys are paid/GPT-4 so we assume they are
keyList = [{ key: keyString, isTrial: false, isGpt4: true }]; keyList = [{ key: keyString, isTrial: false, isGpt4: true }];
} }
@ -63,7 +65,7 @@ function init() {
}; };
keyPool.push(newKey); keyPool.push(newKey);
logger.info("Key added", { key: newKey.hash }); logger.info({ key: newKey.hash }, "Key added");
} }
// TODO: check each key's usage upon startup. // TODO: check each key's usage upon startup.
} }
@ -79,7 +81,7 @@ function disable(key: Key) {
const keyFromPool = keyPool.find((k) => k.key === key.key)!; const keyFromPool = keyPool.find((k) => k.key === key.key)!;
if (keyFromPool.isDisabled) return; if (keyFromPool.isDisabled) return;
keyFromPool.isDisabled = true; keyFromPool.isDisabled = true;
logger.warn("Key disabled", { key: key.hash }); logger.warn({ key: key.hash }, "Key disabled");
} }
function anyAvailable() { function anyAvailable() {
@ -104,13 +106,13 @@ function get(model: string) {
// Prioritize trial keys // Prioritize trial keys
const trialKeys = availableKeys.filter((key) => key.isTrial); const trialKeys = availableKeys.filter((key) => key.isTrial);
if (trialKeys.length > 0) { if (trialKeys.length > 0) {
logger.info("Using trial key", { key: trialKeys[0].hash }); logger.info({ key: trialKeys[0].hash }, "Using trial key");
return trialKeys[0]; return trialKeys[0];
} }
// Otherwise, return the oldest key // Otherwise, return the oldest key
const oldestKey = availableKeys.sort((a, b) => a.lastUsed - b.lastUsed)[0]; const oldestKey = availableKeys.sort((a, b) => a.lastUsed - b.lastUsed)[0];
logger.info("Using key", { key: oldestKey.hash }); logger.info({ key: oldestKey.hash }, "Assigning key to request.");
oldestKey.lastUsed = Date.now(); oldestKey.lastUsed = Date.now();
return { ...oldestKey }; return { ...oldestKey };
} }

View File

@ -45,7 +45,7 @@ const handleResponse = (
try { try {
errorPayload = JSON.parse(body); errorPayload = JSON.parse(body);
} catch (err) { } catch (err) {
logger.error(errorPayload.error, { error: err }); logger.error({ error: err }, errorPayload.error);
res.status(statusCode).json(errorPayload); res.status(statusCode).json(errorPayload);
return; return;
} }
@ -71,8 +71,8 @@ const handleResponse = (
errorPayload.proxy_note = message; errorPayload.proxy_note = message;
} else { } else {
logger.warn( logger.warn(
`OpenAI rate limit exceeded or model overloaded. Keyhash ${req.key?.hash}`, { errorCode: errorPayload.error?.type },
{ errorCode: errorPayload.error?.type } `OpenAI rate limit exceeded or model overloaded. Keyhash ${req.key?.hash}`
); );
} }
} }

View File

@ -1,8 +1,11 @@
import dotenv from "dotenv";
dotenv.config();
import express from "express"; import express from "express";
import cors from "cors"; import cors from "cors";
import pinoHttp from "pino-http"; import pinoHttp from "pino-http";
import { logger } from "./logger"; import { logger } from "./logger";
import { keys } from "./keys";
import { proxy } from "./proxy"; import { proxy } from "./proxy";
const PORT = process.env.PORT || 7860; const PORT = process.env.PORT || 7860;
@ -18,6 +21,10 @@ app.use(
app.use("/", proxy); app.use("/", proxy);
app.use((_req: unknown, res: express.Response) => {
res.status(404).json({ error: "Not found" });
});
app.use((err: any, _req: unknown, res: express.Response, _next: unknown) => { app.use((err: any, _req: unknown, res: express.Response, _next: unknown) => {
if (err.status) { if (err.status) {
res.status(err.status).json({ error: err.message }); res.status(err.status).json({ error: err.message });
@ -29,4 +36,5 @@ app.use((err: any, _req: unknown, res: express.Response, _next: unknown) => {
app.listen(PORT, () => { app.listen(PORT, () => {
logger.info(`Server listening on port ${PORT}`); logger.info(`Server listening on port ${PORT}`);
keys.init();
}); });