Compare commits
2 commits
94992743e8
...
dc91830fae
Author | SHA1 | Date | |
---|---|---|---|
dc91830fae | |||
552143e345 |
1 changed files with 15 additions and 31 deletions
|
@ -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,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue