zack3d/simple-proxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add temporary user tokens (khanon/oai-reverse-proxy!42)

Browse Source
This commit is contained in:
khanon
2023-09-09 22:21:38 +00:00
parent 5728e235dc
commit 2a453ab657
27 changed files with 758 additions and 259 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/shared/users/schema.ts
+26 -18
View File
@@ -1,5 +1,6 @@
import { ZodType, z } from "zod";
import type { ModelFamily } from "../models";
import { makeOptionalPropsNullable } from "../utils";
export const tokenCountsSchema: ZodType<UserTokenCounts> = z
.object({
@@ -15,44 +16,51 @@ export const tokenCountsSchema: ZodType<UserTokenCounts> = z
export const UserSchema = z
.object({
/** The user's personal access token. */
/** User's personal access token. */
token: z.string(),
/** The IP addresses the user has connected from. */
/** IP addresses the user has connected from. */
ip: z.array(z.string()),
/** The user's nickname. */
nickname: z.string().max(80).nullish(),
/** User's nickname. */
nickname: z.string().max(80).optional(),
/**
* The user's privilege level.
* - `normal`: Default role. Subject to usual rate limits and quotas.
* - `special`: Special role. Higher quotas and exempt from
* auto-ban/lockout.
**/
type: z.enum(["normal", "special"]),
/** The number of prompts the user has made. */
type: z.enum(["normal", "special", "temporary"]),
/** Number of prompts the user has made. */
promptCount: z.number(),
/**
* @deprecated Use `tokenCounts` instead.
* Never used; retained for backwards compatibility.
*/
tokenCount: z.any().optional(),
/** The number of tokens the user has consumed, by model family. */
/** Number of tokens the user has consumed, by model family. */
tokenCounts: tokenCountsSchema,
/** The maximum number of tokens the user can consume, by model family. */
/** Maximum number of tokens the user can consume, by model family. */
tokenLimits: tokenCountsSchema,
/** The time at which the user was created. */
/** Time at which the user was created. */
createdAt: z.number(),
/** The time at which the user last connected. */
lastUsedAt: z.number().nullish(),
/** The time at which the user was disabled, if applicable. */
disabledAt: z.number().nullish(),
/** The reason for which the user was disabled, if applicable. */
disabledReason: z.string().nullish(),
/** Time at which the user last connected. */
lastUsedAt: z.number().optional(),
/** Time at which the user was disabled, if applicable. */
disabledAt: z.number().optional(),
/** Reason for which the user was disabled, if applicable. */
disabledReason: z.string().optional(),
/** Time at which the user will expire and be disabled (for temp users). */
expiresAt: z.number().optional(),
})
.strict();
export const UserPartialSchema = UserSchema.partial().extend({
token: z.string(),
});
/**
* Variant of `UserSchema` which allows for partial updates, and makes any
* optional properties on the base schema nullable. Null values are used to
* indicate that the property should be deleted from the user object.
*/
export const UserPartialSchema = makeOptionalPropsNullable(UserSchema)
.partial()
.extend({ token: z.string() });
// gpt4-32k was added after the initial release, so this tries to allow for
// data imported from older versions of the app which may be missing the
src/shared/users/user-store.ts
+88 -23
View File
@@ -22,6 +22,7 @@ const MAX_IPS_PER_USER = config.maxIpsPerUser;
const users: Map<string, User> = new Map();
const usersToFlush = new Set<string>();
let quotaRefreshJob: schedule.Job | null = null;
let userCleanupJob: schedule.Job | null = null;
export async function init() {
log.info({ store: config.gatekeeperStore }, "Initializing user store...");
@@ -29,16 +30,8 @@ export async function init() {
await initFirebase();
}
if (config.quotaRefreshPeriod) {
quotaRefreshJob = schedule.scheduleJob(getRefreshCrontab(), () => {
for (const user of users.values()) {
refreshQuota(user.token);
}
log.info(
{ users: users.size, nextRefresh: quotaRefreshJob!.nextInvocation() },
"Token quotas refreshed."
);
});
const crontab = getRefreshCrontab();
quotaRefreshJob = schedule.scheduleJob(crontab, refreshAllQuotas);
if (!quotaRefreshJob) {
throw new Error(
"Unable to schedule quota refresh. Is QUOTA_REFRESH_PERIOD set correctly?"
@@ -49,26 +42,42 @@ export async function init() {
"Scheduled token quota refresh."
);
}
userCleanupJob = schedule.scheduleJob("* * * * *", cleanupExpiredTokens);
log.info("User store initialized.");
}
export function getNextQuotaRefresh() {
if (!quotaRefreshJob) return "never (manual refresh only)";
return quotaRefreshJob.nextInvocation().getTime();
}
/** Creates a new user and returns their token. */
export function createUser() {
/**
* Creates a new user and returns their token. Optionally accepts parameters
* for setting an expiry date and/or token limits for temporary users.
**/
export function createUser(createOptions?: {
type?: User["type"];
expiresAt?: number;
tokenLimits?: User["tokenLimits"];
}) {
const token = uuid();
users.set(token, {
const newUser: User = {
token,
ip: [],
type: "normal",
promptCount: 0,
tokenCounts: { turbo: 0, gpt4: 0, "gpt4-32k": 0, claude: 0 },
tokenLimits: { ...config.tokenQuota },
tokenLimits: createOptions?.tokenLimits ?? { ...config.tokenQuota },
createdAt: Date.now(),
});
};
if (createOptions?.type === "temporary") {
Object.assign(newUser, {
type: "temporary",
expiresAt: createOptions.expiresAt,
});
} else {
Object.assign(newUser, { type: createOptions?.type ?? "normal" });
}
users.set(token, newUser);
usersToFlush.add(token);
return token;
}
@@ -114,6 +123,14 @@ export function upsertUser(user: UserUpdate) {
}
}
// TODO: Write firebase migration to backfill gpt4-32k token counts
if (updates.tokenCounts) {
updates.tokenCounts["gpt4-32k"] ??= 0;
}
if (updates.tokenLimits) {
updates.tokenLimits["gpt4-32k"] ??= 0;
}
users.set(user.token, Object.assign(existing, updates));
usersToFlush.add(user.token);
@@ -161,7 +178,7 @@ export function authenticate(token: string, ip: string) {
const ipLimit =
user.type === "special" || !MAX_IPS_PER_USER ? Infinity : MAX_IPS_PER_USER;
if (user.ip.length > ipLimit) {
disableUser(token, "Too many IP addresses associated with this token.");
disableUser(token, "IP address limit exceeded.");
return;
}
@@ -223,6 +240,48 @@ export function disableUser(token: string, reason?: string) {
usersToFlush.add(token);
}
export function getNextQuotaRefresh() {
if (!quotaRefreshJob) return "never (manual refresh only)";
return quotaRefreshJob.nextInvocation().getTime();
}
/**
* Cleans up expired temporary tokens by disabling tokens past their access
* expiry date and permanently deleting tokens three days after their access
* expiry date.
*/
function cleanupExpiredTokens() {
const now = Date.now();
let disabled = 0;
let deleted = 0;
for (const user of users.values()) {
if (user.type !== "temporary") continue;
if (user.expiresAt && user.expiresAt < now && !user.disabledAt) {
disableUser(user.token, "Temporary token expired.");
disabled++;
}
if (user.disabledAt && user.disabledAt + 72 * 60 * 60 * 1000 < now) {
users.delete(user.token);
usersToFlush.add(user.token);
deleted++;
}
}
log.debug({ disabled, deleted }, "Expired tokens cleaned up.");
}
function refreshAllQuotas() {
let count = 0;
for (const user of users.values()) {
if (user.type === "temporary") continue;
refreshQuota(user.token);
count++;
}
log.info(
{ refreshed: count, nextRefresh: quotaRefreshJob!.nextInvocation() },
"Token quotas refreshed."
);
}
// TODO: Firebase persistence is pretend right now and just polls the in-memory
// store to sync it with Firebase when it changes. Will refactor to abstract
// persistence layer later so we can support multiple stores.
@@ -253,10 +312,12 @@ async function flushUsers() {
const db = admin.database(app);
const usersRef = db.ref("users");
const updates: Record<string, User> = {};
const deletions = [];
for (const token of usersToFlush) {
const user = users.get(token);
if (!user) {
deletions.push(token);
continue;
}
updates[token] = user;
@@ -264,13 +325,17 @@ async function flushUsers() {
usersToFlush.clear();
const numUpdates = Object.keys(updates).length;
const numUpdates = Object.keys(updates).length + deletions.length;
if (numUpdates === 0) {
return;
}
await usersRef.update(updates);
log.info({ users: Object.keys(updates).length }, "Flushed users to Firebase");
await Promise.all(deletions.map((token) => usersRef.child(token).remove()));
log.info(
{ users: Object.keys(updates).length, deletions: deletions.length },
"Flushed changes to Firebase"
);
}
// TODO: use key-management/models.ts for family mapping