Separate prompt and completion tokens

Author: backwardsmap

prisma/migrations/20250103193523_token_stats_multiple_types/migration.sql

@@ -0,0 +1,10 @@
+/*
+  Warnings:
+
+  - Added the required column `completionTokens` to the `token_stats` table without a default value. This is not possible if the table is not empty.
+  - Added the required column `promptTokens` to the `token_stats` table without a default value. This is not possible if the table is not empty.
+
+*/
+-- AlterTable
+ALTER TABLE "token_stats" ADD COLUMN     "completionTokens" INTEGER NOT NULL,
+ADD COLUMN     "promptTokens" INTEGER NOT NULL;

prisma/schema.prisma

@@ -67,6 +67,8 @@ model TokenStat {
   id                  String   @id @default(cuid())
   userId              String
   messageIds          String[]
+  promptTokens        Int
+  completionTokens    Int
   totalTokens         Int
   createdAt           DateTime @default(now())
   updatedAt           DateTime @updatedAt

src/app/api/chat/route.ts

@@ -165,13 +165,25 @@ export async function POST(req: Request) {
           });
 
           // Save the token stats
-          if (messages && newUserMessage && !isNaN(usage.totalTokens)) {
+          if (
+            messages &&
+            newUserMessage &&
+            !isNaN(usage.promptTokens) &&
+            !isNaN(usage.completionTokens) &&
+            !isNaN(usage.totalTokens)
+          ) {
             const messageIds = newUserMessage
               .concat(messages)
               .map((message) => message.id);
-            const totalTokens = usage.totalTokens;
-
-            await dbCreateTokenStat({ userId, messageIds, totalTokens });
+            const { promptTokens, completionTokens, totalTokens } = usage;
+
+            await dbCreateTokenStat({
+              userId,
+              messageIds,
+              promptTokens,
+              completionTokens,
+              totalTokens,
+            });
           }
 
           revalidatePath('/api/conversations');

src/server/db/queries.ts

@@ -167,17 +167,23 @@ export async function dbCreateTokenStat({
   userId,
   messageIds,
   totalTokens,
+  promptTokens,
+  completionTokens,
 }: {
   userId: string;
   messageIds: string[];
   totalTokens: number;
+  promptTokens: number;
+  completionTokens: number;
 }) {
   try {
     return await prisma.tokenStat.create({
       data: {
         userId,
         messageIds,
         totalTokens,
+        promptTokens,
+        completionTokens,
       },
     });
   } catch (error) {