diff --git a/src/pages/api/usage.ts b/src/pages/api/usage.ts index c904370..f77302b 100644 --- a/src/pages/api/usage.ts +++ b/src/pages/api/usage.ts @@ -2,6 +2,7 @@ import { NextApiRequest, NextApiResponse } from "next"; import { getServerSession } from "next-auth/next"; import { authOptions } from "./auth/[...nextauth]"; import { getSubscriptionByEmail } from "./utils/subscription"; +import { getModel } from "@/utils"; import { addUsage, getCurrentMonthUsage } from "./utils/usage"; import { getEndUser } from "./auth/end-user"; import { Quota } from "@/types"; @@ -24,6 +25,7 @@ const handler = async (req: NextApiRequest, res: NextApiResponse) => { if (req.method === "GET") { usage = await getCurrentMonthUsage(endUser); } else if (req.method === "POST") { + const model = getModel((req.headers["x-openai-model"] as string) || ""); usage = await addUsage(endUser); } diff --git a/src/pages/api/utils/usage.ts b/src/pages/api/utils/usage.ts index 3d01fda..c8fa0e1 100644 --- a/src/pages/api/utils/usage.ts +++ b/src/pages/api/utils/usage.ts @@ -23,7 +23,7 @@ export const getCurrentMonthUsage = async (endUser: string): Promise => }; // We coerce individual usage to the begining of the day to reduce the usage records. -export const addUsage = async (endUser: string): Promise => { +export const addUsage = async (endUser: string, addition: number): Promise => { const now = new Date(); const today = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate())); const usage = await prisma.usage.findFirst({ @@ -35,7 +35,7 @@ export const addUsage = async (endUser: string): Promise => { let newUsage = 0; if (usage) { - newUsage = usage.count + 1; + newUsage = usage.count + addition; await prisma.usage.update({ where: { id: usage.id, @@ -45,12 +45,12 @@ export const addUsage = async (endUser: string): Promise => { }, }); } else { - newUsage = 1; + newUsage = addition; await prisma.usage.create({ data: { endUser: endUser, createdAt: today, - count: 1, + count: newUsage, }, }); } diff --git a/src/utils/model.ts b/src/utils/model.ts index 20067d5..5e99909 100644 --- a/src/utils/model.ts +++ b/src/utils/model.ts @@ -4,6 +4,7 @@ const gpt35 = { frequency_penalty: 0.0, presence_penalty: 0.0, max_token: 4000, + cost_per_call: 1, }; const gpt4 = { @@ -12,6 +13,7 @@ const gpt4 = { frequency_penalty: 0.0, presence_penalty: 0.0, max_token: 8000, + cost_per_call: 10, }; export const models = [gpt35, gpt4];