Compare commits

..

2 commits

Author SHA1 Message Date
dc91830fae quota/tokenCount: allow arbitrary multipliers for tokens
I didn't notice that OpenAI changed pricing again
2023-09-25 09:52:22 +02:00
552143e345 quota/tokenCount: Rename single variable name to something that has more sense
I forgot to rename it when I commited it last time ;-;
2023-09-25 09:30:52 +02:00

View file

@ -12,16 +12,19 @@ import { Usage } from "@prisma/client";
export default class tokenCount implements IQuota { export default class tokenCount implements IQuota {
defaultQuota: number; defaultQuota: number;
lookback: number; lookback: number;
considerInputTokensAsHalf: boolean; requestTokenMultiplier: number;
responseTokenMultiplier: number;
constructor( constructor(
defaultQuota: number = 512 * 25, defaultQuota: number = 512 * 25,
lookback: number = 1000 * 60 * 60 * 24, lookback: number = 1000 * 60 * 60 * 24,
considerInputTokensAsHalf: boolean = true, requestTokenMultiplier: number = 1,
responseTokenMultiplier: number = 1,
) { ) {
this.defaultQuota = defaultQuota; this.defaultQuota = defaultQuota;
this.lookback = lookback; this.lookback = lookback;
this.considerInputTokensAsHalf = considerInputTokensAsHalf; this.requestTokenMultiplier = requestTokenMultiplier;
this.responseTokenMultiplier = responseTokenMultiplier;
} }
private getUserQuota(id: string) { private getUserQuota(id: string) {
@ -55,9 +58,7 @@ export default class tokenCount implements IQuota {
const usageResponse = usedTokens.usageResponse === null ? 0 : usedTokens.usageResponse; const usageResponse = usedTokens.usageResponse === null ? 0 : usedTokens.usageResponse;
const usedUnits = (() => { const usedUnits = (() => {
if (this.considerInputTokensAsHalf) return usageRequest * this.requestTokenMultiplier + usageResponse * this.responseTokenMultiplier;
return usageResponse + usageRequest / 2;
return usageResponse + usageRequest;
})(); })();
if (userQuota?.vip) return this.createUserQuotaData(Infinity, usedUnits); if (userQuota?.vip) return this.createUserQuotaData(Infinity, usedUnits);
@ -74,30 +75,13 @@ export default class tokenCount implements IQuota {
* @returns promise of giving out the record * @returns promise of giving out the record
*/ */
findNthUsage(user: string, requestTimestamp: number, unitCount: number) { findNthUsage(user: string, requestTimestamp: number, unitCount: number) {
if (this.considerInputTokensAsHalf) return database.$queryRaw<Array<Usage & {usage: number | bigint}>>`
return database.$queryRaw<Array<Usage & {usage: number}>>`
SELECT t1.*, ( SELECT t1.*, (
SELECT SELECT
SUM(usageResponse + usageRequest/2) AS usage SUM(
FROM \`usage\` usageRequest * ${this.requestTokenMultiplier} +
WHERE usageResponse * ${this.responseTokenMultiplier}
user = ${user} AND ) AS usage
timestamp >= ${requestTimestamp - this.lookback} AND
timestamp <= t1.timestamp
) as usage
FROM
\`usage\` AS t1
WHERE
user = ${user} AND
timestamp >= ${requestTimestamp - this.lookback} AND
usage >= ${unitCount}
ORDER BY timestamp ASC
LIMIT 1
`;
return database.$queryRaw<Array<Usage & {usage: bigint}>>`
SELECT t1.*, (
SELECT
SUM(usageResponse + usageRequest) AS usage
FROM \`usage\` FROM \`usage\`
WHERE WHERE
user = ${user} AND user = ${user} AND
@ -121,14 +105,14 @@ export default class tokenCount implements IQuota {
): Promise<userQuotaRecoveryData> { ): Promise<userQuotaRecoveryData> {
const userId = typeof user ==="string" ? user : user.id; const userId = typeof user ==="string" ? user : user.id;
const [userQuota, renameMebecause] = await Promise.all([ const [userQuota, overUnitCountRecord] = await Promise.all([
this.checkUser(userId, request), this.checkUser(userId, request),
this.findNthUsage(userId, request.createdTimestamp, unitCount) this.findNthUsage(userId, request.createdTimestamp, unitCount)
]); ]);
return { return {
...userQuota, ...userQuota,
recoveryTimestamp: (renameMebecause.at(0)?.timestamp.valueOf() ?? Infinity) + this.lookback, recoveryTimestamp: (overUnitCountRecord.at(0)?.timestamp.valueOf() ?? Infinity) + this.lookback,
}; };
} }