From 80f4f18eabc7f88b8e9537c4cba21b4024ca182d Mon Sep 17 00:00:00 2001
From: Wroclaw <wroclaw223@outlook.com>
Date: Mon, 25 Sep 2023 09:52:22 +0200
Subject: [PATCH] quota/tokenCount: allow arbitrary multipliers for tokens

I didn't notice that OpenAI changed pricing again
---
 src/quota/tokenCount.ts | 40 ++++++++++++----------------------------
 1 file changed, 12 insertions(+), 28 deletions(-)

diff --git a/src/quota/tokenCount.ts b/src/quota/tokenCount.ts
index 4c24fbb..1756c85 100644
--- a/src/quota/tokenCount.ts
+++ b/src/quota/tokenCount.ts
@@ -12,16 +12,19 @@ import { Usage } from "@prisma/client";
 export default class tokenCount implements IQuota {
   defaultQuota: number;
   lookback: number;
-  considerInputTokensAsHalf: boolean;
+  requestTokenMultiplier: number;
+  responseTokenMultiplier: number;
 
   constructor(
     defaultQuota: number = 512 * 25,
     lookback: number = 1000 * 60 * 60 * 24,
-    considerInputTokensAsHalf: boolean = true,
+    requestTokenMultiplier: number = 1,
+    responseTokenMultiplier: number = 1,
   ) {
     this.defaultQuota = defaultQuota;
     this.lookback = lookback;
-    this.considerInputTokensAsHalf = considerInputTokensAsHalf;
+    this.requestTokenMultiplier = requestTokenMultiplier;
+    this.responseTokenMultiplier = responseTokenMultiplier;
   }
 
   private getUserQuota(id: string) {
@@ -55,9 +58,7 @@ export default class tokenCount implements IQuota {
     const usageResponse = usedTokens.usageResponse === null ? 0 : usedTokens.usageResponse;
 
     const usedUnits = (() => {
-      if (this.considerInputTokensAsHalf)
-        return usageResponse + usageRequest / 2;
-      return usageResponse + usageRequest;
+      return usageRequest * this.requestTokenMultiplier + usageResponse * this.responseTokenMultiplier;
     })();
 
     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
    */
   findNthUsage(user: string, requestTimestamp: number, unitCount: number) {
-    if (this.considerInputTokensAsHalf)
-      return database.$queryRaw<Array<Usage & {usage: number}>>`
-        SELECT t1.*, (
-            SELECT
-              SUM(usageResponse + usageRequest/2) AS usage
-            FROM \`usage\`
-            WHERE 
-              user = ${user} AND
-              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}>>`
+    return database.$queryRaw<Array<Usage & {usage: number | bigint}>>`
       SELECT t1.*, (
           SELECT
-            SUM(usageResponse + usageRequest) AS usage
+            SUM(
+              usageRequest * ${this.requestTokenMultiplier} +
+              usageResponse * ${this.responseTokenMultiplier}
+            ) AS usage
           FROM \`usage\`
           WHERE 
             user = ${user} AND