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

adds user-specific overrides for daily quota refresh

Browse Source
This commit is contained in:
nai-degen
2024-07-27 14:25:53 -05:00
parent f242777596
commit 2aa19e5b09
8 changed files with 92 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/admin/web/views/admin_index.ejs
+1 -1
View File
@@ -43,7 +43,7 @@
<legend>Bulk Quota Management</legend>
<p>
<button id="refresh-quotas" type="button" onclick="submitForm('resetQuotas')">Refresh All Quotas</button>
Resets all users' quotas to the values set in the <code>TOKEN_QUOTA_*</code> environment variables.
Immediately refreshes all users' quotas by the configured amounts.
</p>
<p>
<button id="clear-token-counts" type="button" onclick="submitForm('resetCounts')">Clear All Token Counts</button>
src/admin/web/views/admin_view-user.ejs
+18 -9
View File
@@ -101,6 +101,10 @@
<% ["nickname", "type", "disabledAt", "disabledReason", "maxIps", "adminNote"].forEach(function (key) { %>
<input type="hidden" name="<%- key %>" value="<%- user[key] %>" />
<% }); %>
<!-- tokenRefresh_ keys are dynamically generated -->
<% Object.entries(quota).forEach(([family]) => { %>
<input type="hidden" name="tokenRefresh_<%- family %>" value="<%- user.tokenRefresh[family] || quota[family] %>" />
<% }); %>
</form>
<h3>Quota Information</h3>
@@ -111,7 +115,7 @@
<button type="submit" class="btn btn-primary">Refresh Quotas for User</button>
</form>
<% } %>
<%- include("partials/shared_quota-info", { quota, user }) %>
<%- include("partials/shared_quota-info", { quota, user, showRefreshEdit: true }) %>
<p><a href="/admin/manage/list-users">Back to User List</a></p>
@@ -122,18 +126,25 @@
const token = a.dataset.token;
const field = a.dataset.field;
const existingValue = document.querySelector(`#current-values input[name=${field}]`).value;
let value = prompt(`Enter new value for '${field}'':`, existingValue);
let value = prompt(`Enter new value for '${field}':`, existingValue);
if (value !== null) {
if (value === "") {
value = null;
}
const payload = { _csrf: document.querySelector("meta[name=csrf-token]").getAttribute("content") };
if (field.startsWith("tokenRefresh_")) {
const family = field.slice("tokenRefresh_".length);
payload.tokenRefresh = { [family]: Number(value) };
} else {
payload[field] = value;
}
fetch(`/admin/manage/edit-user/${token}`, {
method: "POST",
credentials: "same-origin",
body: JSON.stringify({
[field]: value,
_csrf: document.querySelector("meta[name=csrf-token]").getAttribute("content"),
}),
body: JSON.stringify(payload),
headers: { "Content-Type": "application/json", Accept: "application/json" },
})
.then((res) => Promise.all([res.ok, res.json()]))
@@ -141,9 +152,7 @@
const url = new URL(window.location.href);
const params = new URLSearchParams();
if (!ok) {
params.set("flash", `error: ${json.error.message}`);
} else {
params.set("flash", `success: User's ${field} updated.`);
alert(`Failed to edit user: ${json.message}`);
}
url.search = params.toString();
window.location.assign(url);
src/shared/users/schema.ts
+2
View File
@@ -37,6 +37,8 @@ export const UserSchema = z
tokenCounts: tokenCountsSchema,
/** Maximum number of tokens the user can consume, by model family. */
tokenLimits: tokenCountsSchema,
/** User-specific token refresh amount, by model family. */
tokenRefresh: tokenCountsSchema,
/** Time at which the user was created. */
createdAt: z.number(),
/** Time at which the user last connected. */
src/shared/users/user-store.ts
+30 -12
View File
@@ -70,6 +70,7 @@ export function createUser(createOptions?: {
type?: User["type"];
expiresAt?: number;
tokenLimits?: User["tokenLimits"];
tokenRefresh?: User["tokenRefresh"];
}) {
const token = uuid();
const newUser: User = {
@@ -79,6 +80,7 @@ export function createUser(createOptions?: {
promptCount: 0,
tokenCounts: { ...INITIAL_TOKENS },
tokenLimits: createOptions?.tokenLimits ?? { ...config.tokenQuota },
tokenRefresh: createOptions?.tokenRefresh ?? { ...INITIAL_TOKENS },
createdAt: Date.now(),
meta: {},
};
@@ -123,6 +125,7 @@ export function upsertUser(user: UserUpdate) {
promptCount: 0,
tokenCounts: { ...INITIAL_TOKENS },
tokenLimits: { ...config.tokenQuota },
tokenRefresh: { ...INITIAL_TOKENS },
createdAt: Date.now(),
meta: {},
};
@@ -139,7 +142,6 @@ export function upsertUser(user: UserUpdate) {
}
}
// TODO: Write firebase migration to backfill new fields
if (updates.tokenCounts) {
for (const family of MODEL_FAMILIES) {
updates.tokenCounts[family] ??= 0;
@@ -150,6 +152,11 @@ export function upsertUser(user: UserUpdate) {
updates.tokenLimits[family] ??= 0;
}
}
if (updates.tokenRefresh) {
for (const family of MODEL_FAMILIES) {
updates.tokenRefresh[family] ??= 0;
}
}
users.set(user.token, Object.assign(existing, updates));
usersToFlush.add(user.token);
@@ -245,19 +252,30 @@ export function hasAvailableQuota({
return tokensConsumed < tokenLimit;
}
/**
* For the given user, sets token limits for each model family to the sum of the
* current count and the refresh amount, up to the default limit. If a quota is
* not specified for a model family, it is not touched.
*/
export function refreshQuota(token: string) {
const user = users.get(token);
if (!user) return;
const { tokenCounts, tokenLimits } = user;
const quotas = Object.entries(config.tokenQuota) as [ModelFamily, number][];
quotas
// If a quota is not configured, don't touch any existing limits a user may
// already have been assigned manually.
.filter(([, quota]) => quota > 0)
.forEach(
([model, quota]) =>
(tokenLimits[model] = (tokenCounts[model] ?? 0) + quota)
);
const { tokenQuota } = config;
const { tokenCounts, tokenLimits, tokenRefresh } = user;
// Get default quotas for each model family.
const defaultQuotas = Object.entries(tokenQuota) as [ModelFamily, number][];
// If any user-specific refresh quotas are present, override default quotas.
const userQuotas = defaultQuotas.map(([f, q]) => [
f,
(tokenRefresh[f] ?? 0) || q,
] as const /* narrow to tuple */);
userQuotas
// Ignore families with no global or user-specific refresh quota.
.filter(([, q]) => q > 0)
// Increase family token limit by the family's refresh amount.
.forEach(([f, q]) => (tokenLimits[f] = (tokenCounts[f] ?? 0) + q));
usersToFlush.add(token);
}
@@ -307,7 +325,7 @@ function cleanupExpiredTokens() {
user.meta.refreshable = config.captchaMode !== "none";
disabled++;
}
const purgeTimeout = config.powTokenPurgeHours * 60 * 60 * 1000;
const purgeTimeout = config.powTokenPurgeHours * 60 * 60 * 1000;
if (user.disabledAt && user.disabledAt + purgeTimeout < now) {
users.delete(user.token);
usersToFlush.add(user.token);
src/shared/views/partials/shared_header.ejs
+20 -9
View File
@@ -33,7 +33,7 @@
.pagination li a {
display: block;
padding: 0.5em 1em;
border-bottom: none;
border-bottom: none;
text-decoration: none;
}
.pagination li.active a {
@@ -71,20 +71,24 @@
td.actions:hover {
background-color: #e0e6f6;
}
tr > td,
tr > th {
border-right: 1px solid #dedede;
}
@media (max-width: 800px) {
body {
padding: 0.5em;
}
table.full-width {
width: 100%;
position: static;
left: auto;
right: auto;
margin-left: 0;
margin-right: 0;
}
table.full-width {
width: 100%;
position: static;
left: auto;
right: auto;
margin-left: 0;
margin-right: 0;
}
}
@media (prefers-color-scheme: dark) {
@@ -95,6 +99,13 @@
th.active {
background-color: #446;
}
td.actions:hover {
background-color: #446;
}
tr > td,
tr > th {
border-right: 1px solid #444;
}
}
</style>
</head>
src/shared/views/partials/shared_quota-info.ejs
+19 -4
View File
@@ -1,4 +1,6 @@
<p>Next refresh: <time><%- nextQuotaRefresh %></time></p>
<p>
Next refresh: <time><%- nextQuotaRefresh %></time>
</p>
<table class="striped">
<thead>
<tr>
@@ -9,7 +11,7 @@
<% } %>
<th scope="col">Limit</th>
<th scope="col">Remaining</th>
<th scope="col">Refresh Amount</th>
<th scope="col" colspan="<%= showRefreshEdit ? 2 : 1 %>">Refresh Amount</th>
</tr>
</thead>
<tbody>
@@ -19,7 +21,7 @@
<td><%- prettyTokens(user.tokenCounts[key]) %></td>
<% if (showTokenCosts) { %>
<td>$<%- tokenCost(key, user.tokenCounts[key]).toFixed(2) %></td>
<% } %>
<% } %>
<% if (!user.tokenLimits[key]) { %>
<td colspan="2" style="text-align: center">unlimited</td>
<% } else { %>
@@ -29,7 +31,20 @@
<% if (user.type === "temporary") { %>
<td>N/A</td>
<% } else { %>
<td><%- prettyTokens(quota[key]) %></td>
<td><%- prettyTokens(user.tokenRefresh[key] || quota[key]) %></td>
<% } %>
<% if (showRefreshEdit) { %>
<td class="actions">
<a
title="Edit"
id="edit-refresh"
href="#"
data-field="tokenRefresh_<%= key %>"
data-token="<%= user.token %>"
data-modelFamily="<%= key %>"
>✏️</a
>
</td>
<% } %>
</tr>
<% }) %>
src/user/web/views/user_lookup.ejs
+1 -1
View File
@@ -64,7 +64,7 @@
</table>
<h3>Quota Information</h3>
<%- include("partials/shared_quota-info", { quota, user }) %>
<%- include("partials/shared_quota-info", { quota, user, showRefreshEdit: false }) %>
<form id="edit-nickname-form" style="display: none" action="/user/edit-nickname" method="post">
<input type="hidden" name="_csrf" value="<%= csrfToken %>" />