diff --git a/app/client/platforms/openai.ts b/app/client/platforms/openai.ts index 8ea864692..196d76902 100644 --- a/app/client/platforms/openai.ts +++ b/app/client/platforms/openai.ts @@ -18,6 +18,14 @@ import { prettyObject } from "@/app/utils/format"; import { getClientConfig } from "@/app/config/client"; import { makeAzurePath } from "@/app/azure"; +//# for AdEx custom usage tracking (calls per model per user per datekey) +// import { NextApiRequest, NextApiResponse } from 'next'; +import { getServerAuthSession } from '../../auth'; // Adjust the import path as necessary +import { incrementAPICallCount } from '../../utils/cloud/redisRestClient'; +// app\utils\cloud\redisRestClient.ts +// app\client\platforms\openai.ts +// app\auth.ts + export interface OpenAIListModelResponse { object: string; data: Array<{ @@ -94,6 +102,30 @@ export class ChatGPTApi implements LLMApi { console.log("[Request] openai payload: ", requestPayload); + + + // export default async function handler(req: NextApiRequest, res: NextApiResponse) { + // Retrieve the session using getServerAuthSession + const session = await getServerAuthSession()(req); + + if (session?.user?.email) { + // Now you have the user's email from the session + const userEmail = session.user.email; + const modelIdentifier = modelConfig.mode; + const dateKey = new Date().toISOString().slice(0, 7); // "YYYY-MM" + console.log("API Call: ",token.email, modelIdentifier); + + // Use the userEmail to increment the API call count + await incrementAPICallCount(userEmail, modelIdentifier, dateKey); + + // ... rest of your API route logic ... + } else { + // Handle cases where the session or email is not available + res.status(401).json({ error: 'Unauthorized' }); + } + // } + + const shouldStream = !!options.config.stream; const controller = new AbortController(); options.onController?.(controller); diff --git a/app/utils/cloud/redisRestClient.ts b/app/utils/cloud/redisRestClient.ts index ffc1a85cb..6e51c51b9 100644 --- a/app/utils/cloud/redisRestClient.ts +++ b/app/utils/cloud/redisRestClient.ts @@ -68,6 +68,31 @@ export const incrementSessionRefreshCount = async ( } }; +export const incrementAPICallCount = async ( + email: string | undefined, + modelIdentifier: string, + dateKey: string +): Promise => { + if (!email || !modelIdentifier) { + console.error('Email or model identifier is undefined, cannot increment API call count.'); + return; // This is fine for a function returning Promise + } + + const redis = getRedisClient(); + if (!redis) { + console.error('Redis client is not initialized.'); + return; // This is fine for a function returning Promise + } + + const key = `api_calls:${email}:${modelIdentifier}`; + + try { + await redis.hincrby(key, dateKey, 1); + } catch (error) { + console.error('Failed to increment API call count in Redis via Upstash', error); + } +}; + export const incrementTokenCounts = async ( email: string | undefined,