优化token计数

app/api/logs/[...path]/route.ts

@@ -1,17 +1,7 @@
 import { NextRequest, NextResponse } from "next/server";
 import prisma from "@/lib/prisma";
 import { insertUser } from "@/lib/auth";
-
-// function cleanObjectKeys(input: { [key: string]: string }): {
-//   [key: string]: string;
-// } {
-//   const cleanedObj: { [key: string]: string } = {};
-//   Object.keys(input).forEach((key) => {
-//     cleanedObj[key] = input[key].trim();
-//   });
-//   console.log('========', input, cleanedObj)
-//   return cleanedObj;
-// }
+import { getTokenLength } from "@/app/utils/token";
 
 async function handle(
   req: NextRequest,
@@ -23,12 +13,22 @@ async function handle(
       await insertUser({ name: request_data?.userName });
     }
     // console.log("===========4", request_data);
+    if (request_data?.logEntry) {
+      const regex = /\[(.*)]/g;
+      const matchResponse = request_data.logEntry.match(regex);
+      if (matchResponse.length > 0) {
+        request_data.logToken = getTokenLength(matchResponse[0]);
+      }
+      // console.log('=======', request_data.logEntry, '=====', matchResponse);
+    }
+
     await prisma.logEntry.create({
       data: request_data,
     });
   } catch (e) {
+    console.log("[LOG]", e);
+
     return NextResponse.json({ status: 0 });
-    // console.log("[LOG]", e);
   }
 
   return NextResponse.json({ status: 1 });

app/app/(admin)/admin/usage-analysis.tsx

@@ -1,73 +0,0 @@
-// "use client";
-import { BarList, Bold, Card, Flex, Text, Title } from "@tremor/react";
-import prisma from "@/lib/prisma";
-import { estimateTokenLength } from "@/app/utils/token";
-
-function HandleLogData(todayLog: [{ userName: string; logEntry: string }]) {
-  const data1 = todayLog.map((log) => {
-    return {
-      name: log.userName ?? "unknown",
-      value: estimateTokenLength(log.logEntry ?? ""),
-    };
-  });
-
-  type Accumulator = {
-    [key: string]: number;
-  };
-  const data2 = data1.reduce<Accumulator>((acc, item) => {
-    if (acc[item.name]) {
-      acc[item.name] += item.value;
-    } else {
-      acc[item.name] = item.value;
-    }
-    return acc;
-  }, {});
-  const data3 = Object.entries(data2)
-    .map(([name, value]) => ({
-      name,
-      value,
-    }))
-    .sort((a, b) => {
-      return b.value - a.value;
-    });
-  return data3;
-}
-
-export default async function UsageAnalysis() {
-  // 今天日期的开始和结束
-  const startDate = new Date();
-  startDate.setHours(0, 0, 0, 0);
-
-  const endDate = new Date();
-  endDate.setHours(23, 59, 59, 999);
-  const todayLog = await prisma.logEntry.findMany({
-    where: {
-      createdAt: {
-        gte: startDate, // gte 表示 '大于等于'
-        lte: endDate, // lte 表示 '小于等于'
-      },
-    },
-    include: {
-      user: true,
-    },
-  });
-
-  // @ts-ignore
-  const log_data = HandleLogData(todayLog);
-
-  // @ts-ignore
-  return (
-    <Card className="max-w-lg">
-      <Title>Website Analytics</Title>
-      <Flex className="mt-4">
-        <Text>
-          <Bold>Source</Bold>
-        </Text>
-        <Text>
-          <Bold>Visits</Bold>
-        </Text>
-      </Flex>
-      <BarList data={log_data} className="mt-2" />
-    </Card>
-  );
-}

app/app/(admin)/admin/usage-by-model.tsx

@@ -2,17 +2,19 @@ import * as echarts from "echarts";
 import { EChartsOption } from "echarts";
 import dynamic from "next/dynamic";
 import prisma from "@/lib/prisma";
-import { estimateTokenLength } from "@/app/utils/token";
+// import { getTokenLength } from "@/app/utils/token";
 
 const UsageByModelChart = dynamic(() => import("./usage-by-model-chart"), {
   ssr: false,
 });
 
-function HandleLogData(todayLog: [{ userName: string; logEntry: string }]) {
+function HandleLogData(
+  todayLog: [{ userName: string; logEntry: string; logToken: number }],
+) {
   const data1 = todayLog.map((log) => {
     return {
       name: log.userName ?? "unknown",
-      value: estimateTokenLength(log.logEntry ?? ""),
+      value: log.logToken,
     };
   });
 

app/utils/token.ts

@@ -20,3 +20,22 @@ export function estimateTokenLength(input: string): number {
 
   return tokenLength;
 }
+
+// import { get_encoding } from "tiktoken";
+
+export function getTokenLength(input: string): number {
+  // const { get_encoding } = require( "tiktoken" );
+  // const encoding = get_encoding("cl100k_base");
+
+  const { Tiktoken } = require("tiktoken/lite");
+  const cl100k_base = require("tiktoken/encoders/cl100k_base.json");
+  const encoding = new Tiktoken(
+    cl100k_base.bpe_ranks,
+    cl100k_base.special_tokens,
+    cl100k_base.pat_str,
+  );
+  const tokenLength = encoding.encode(input).length;
+  // console.log('[TOKEN],=========', input, tokenLength)
+
+  return tokenLength;
+}

next.config.mjs

@@ -28,11 +28,16 @@ const nextConfig = {
     }
     // console.log('=======', config.optimization)
 
-
     config.resolve.fallback = {
       child_process: false,
     };
 
+    // tiktoken
+    config.experiments = {
+      asyncWebAssembly: true,
+      layers: true,
+    };
+
     return config;
   },
   output: mode,

package.json

@@ -17,9 +17,9 @@
   },
   "dependencies": {
     "@fortaine/fetch-event-source": "^3.0.6",
+    "@hello-pangea/dnd": "^16.5.0",
     "@next-auth/prisma-adapter": "^1.0.7",
     "@prisma/client": "^5.7.0",
-    "@hello-pangea/dnd": "^16.5.0",
     "@svgr/webpack": "^8.1.0",
     "@tremor/react": "^3.12.1",
     "@vercel/analytics": "^1.1.1",
@@ -46,6 +46,7 @@
     "sonner": "^1.2.0",
     "spark-md5": "^3.0.2",
     "tailwind-merge": "^2.0.0",
+    "tiktoken": "^1.0.11",
     "use-debounce": "^10.0.0",
     "zustand": "^4.3.8"
   },

prisma/schema.prisma

@@ -71,6 +71,7 @@ model LogEntry {
   userName  String?  @db.VarChar(50)
   createdAt DateTime @default(now())
   logEntry  String?  @db.Text
+  logToken  Int?     @default(0)
   user      User?    @relation(fields: [userName], references: [name], onDelete: NoAction)
 }