🧪 xp bnp: use larger model for summary when exceeding 12000 tokens

gmpetrov · Jun 22, 2023 · 476dcf6 · 476dcf6 · vercel · Jun 22, 2023
1 parent 8651186
commit 476dcf6
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/pages/api/xp/bnp/summary.ts b/pages/api/xp/bnp/summary.ts
@@ -1,5 +1,7 @@
+import { get_encoding } from '@dqbd/tiktoken';
 import { ConversationChannel, MessageFrom, Usage } from '@prisma/client';
 import cuid from 'cuid';
+import { TokenTextSplitter } from 'langchain/text_splitter';
 import { NextApiResponse } from 'next';
 
 import { AppNextApiRequest, ChatRequest } from '@app/types';
@@ -15,6 +17,8 @@ import prisma from '@app/utils/prisma-client';
 
 const handler = createAuthApiHandler();
 
+const MAX_TOKENS = 12000;
+
 export const XPBNPQuery = async (
   req: AppNextApiRequest,
   res: NextApiResponse
@@ -32,6 +36,7 @@ export const XPBNPQuery = async (
   const receivedDate = new Date();
 
   let content = '';
+  let nbTokens = 0;
 
   if (data.datasourceId || data.datastoreId) {
     if (!data.datasourceId || !data.datastoreId) {
@@ -68,6 +73,22 @@ export const XPBNPQuery = async (
 
     content = JSON.parse(s3Response.Body?.toString('utf-8') || '{}')
       ?.text as string;
+
+    const encoding = get_encoding('cl100k_base');
+    nbTokens = encoding.encode(content).length;
+    encoding.free();
+
+    if (nbTokens > MAX_TOKENS) {
+      const splitter = new TokenTextSplitter({
+        chunkSize: MAX_TOKENS,
+        chunkOverlap: 0,
+      });
+
+      const chunks = await splitter.splitText(content);
+      content = chunks[0];
+    }
+
+    console.log('nbTokens ------>', nbTokens);
   }
 
   if (data.streaming) {
@@ -90,6 +111,7 @@ export const XPBNPQuery = async (
     temperature: 0,
     query: data.query,
     stream: data.streaming ? streamData : undefined,
+    modelName: nbTokens > 3200 ? 'gpt-3.5-turbo-16k' : undefined,
   });
 
   await prisma.messageBNP.createMany({

diff --git a/utils/chat.ts b/utils/chat.ts
@@ -117,6 +117,7 @@ const chat = async ({
   stream,
   temperature,
   history,
+  modelName,
 }: {
   datastore?: Datastore;
   query: string;
@@ -125,6 +126,7 @@ const chat = async ({
   topK?: number;
   stream?: any;
   temperature?: number;
+  modelName?: string;
   history?: { from: MessageFrom; message: string }[];
 }) => {
   let results = [] as {
@@ -177,7 +179,7 @@ const chat = async ({
   }
 
   const model = new ChatOpenAI({
-    modelName: 'gpt-3.5-turbo-0613',
+    modelName: modelName || 'gpt-3.5-turbo-0613',
     temperature: temperature || 0,
     streaming: Boolean(stream),
     callbacks: [