Support message history compression 2

memfreeme · Jan 10, 2025 · 5ec8dd6 · 5ec8dd6
1 parent 4aee9a5
commit 5ec8dd6
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 23 deletions.
diff --git a/frontend/app/api/search/route.ts b/frontend/app/api/search/route.ts
@@ -44,8 +44,6 @@ export async function POST(req: NextRequest) {
 
     let { model, source, messages, profile, isSearch, questionLanguage, answerLanguage, summary } = await req.json();
 
-    console.log('model', model, 'source', source, 'messages', messages.length, 'userId', userId, 'isSearch', isSearch, 'summary', summary);
-
     if (isProModel(model) && !isPro) {
         return NextResponse.json(
             {
@@ -71,7 +69,7 @@ export async function POST(req: NextRequest) {
             async start(controller) {
                 switch (source) {
                     case SearchCategory.O1: {
-                        await o1Answer(isSearch, messages, isPro, userId, profile, streamController(controller), model);
+                        await o1Answer(isSearch, messages, isPro, userId, profile, summary, streamController(controller), model);
                         break;
                     }
                     case SearchCategory.CHAT: {

diff --git a/frontend/components/search/search-window.tsx b/frontend/components/search/search-window.tsx
@@ -66,7 +66,6 @@ export default function SearchWindow({ id, initialMessages, user, isReadOnly = f
     // monitorMemoryUsage();
 
     const { incrementSearchCount, canSearch } = useSearchLimit();
-    const { isCompressHistory } = useSearchState();
 
     const sendMessage = useCallback(
         async (question?: string, attachments?: string[], messageIdToUpdate?: string) => {
@@ -122,11 +121,11 @@ export default function SearchWindow({ id, initialMessages, user, isReadOnly = f
             }
 
             const waitForCompression = async () => {
-                if (!isCompressHistory) return;
+                if (!useSearchState.getState().isCompressHistory) return;
 
                 return new Promise<void>((resolve) => {
                     const checkCompressionStatus = () => {
-                        if (!isCompressHistory) {
+                        if (!useSearchState.getState().isCompressHistory) {
                             resolve();
                         } else {
                             console.log('Waiting for compression to finish...');

diff --git a/frontend/hooks/use-compress-history.ts b/frontend/hooks/use-compress-history.ts
@@ -8,25 +8,27 @@ export function useCompressHistory() {
     const { isCompressHistory, setIsCompressHistory } = useSearchState();
 
     const compressHistoryMessages = async () => {
-        console.log('compressHistoryMessages');
         if (isCompressHistory) return;
         if (!activeSearch?.messages) return;
 
         const messages = activeSearch.messages;
         const totalMessages = messages.length;
         if (totalMessages < 4) return;
 
+        const totalLength = messages.reduce((acc, message) => {
+            return acc + (message.content?.length || 0);
+        }, 0);
+
+        if (totalLength < 4000) return;
+
         const model = useConfigStore.getState().model;
         if (!isProModel(model)) {
             return;
         }
 
-        console.log('compressHistoryMessages totalMessages:', totalMessages);
-
         try {
             const compressIndex = activeSearch.lastCompressIndex || 0;
             const newMessagesToCompress = messages.slice(compressIndex);
-            console.log('compressHistoryMessages newMessagesToCompress:', newMessagesToCompress, compressIndex);
             if (newMessagesToCompress.length < 4 || newMessagesToCompress.length > 6) {
                 return;
             }

diff --git a/frontend/lib/llm/utils.ts b/frontend/lib/llm/utils.ts
@@ -5,12 +5,13 @@ export function getHistoryMessages(isPro: boolean, messages: any[], summary?: st
     const sliceNum = isPro ? -7 : -3;
     const slicedMessages = messages?.slice(sliceNum);
     if (summary) {
+        console.log('summary', summary, 'isPro', isPro);
         return [
             {
                 content: summary,
                 role: 'system',
             },
-            ...slicedMessages.slice(-2),
+            ...slicedMessages.slice(-4),
         ];
     }
     return slicedMessages;

diff --git a/frontend/lib/tools/compress-history.ts b/frontend/lib/tools/compress-history.ts
@@ -6,25 +6,28 @@ import { GPT_4o_MIMI } from '@/lib/model';
 import { Message } from '@/lib/types';
 import { generateText } from 'ai';
 
+const systemPrompt = `You're an assistant who's good at extracting key takeaways from conversations and summarizing them. Please summarize according to the user's needs. Create a new summary from the messages.The summary needs to maintain the original language.`;
+
 export async function compressHistory(messages: Message[], previousSummary?: string): Promise<string> {
     if (messages.length < 4) {
         return '';
     }
 
     try {
-        console.log('compressHistory messages:', messages);
-        console.log('compressHistory previousSummary:', previousSummary);
         console.time('compressHistory');
 
-        const systemPrompt = `You're an assistant who's good at extracting key takeaways from conversations and summarizing them. Please summarize according to the user's needs. ${
-            previousSummary
-                ? 'Please incorporate the previous summary with new messages to create an updated comprehensive summary.'
-                : 'Create a new summary from the messages.'
-        }`;
+        const allMessages = previousSummary
+            ? [
+                  {
+                      role: 'user',
+                      content: previousSummary,
+                  },
+                  ...messages,
+              ]
+            : messages;
 
-        const userPrompt = previousSummary
-            ? `Previous Summary: ${previousSummary}\n\nNew Messages: ${formatHistoryMessages(messages)}\n\nPlease create an updated summary incorporating both the previous summary and new messages. Limit to 400 tokens.`
-            : `${formatHistoryMessages(messages)}\nPlease summarize the above conversation and retain key information. Limit to 400 tokens.`;
+        const userPrompt = `${formatHistoryMessages(allMessages as Message[])}
+Please summarize the above conversation and retain key information. The summarized content will be used as context for subsequent prompts, and should be limited to 400 tokens.`;
 
         const { text } = await generateText({
             model: getLLM(GPT_4o_MIMI),

diff --git a/frontend/lib/tools/o1-answer.ts b/frontend/lib/tools/o1-answer.ts
@@ -3,7 +3,7 @@ import 'server-only';
 import { incSearchCount } from '@/lib/db';
 import { getLLM } from '@/lib/llm/llm';
 import { DirectAnswerPrompt } from '@/lib/llm/prompt';
-import { getHistory, streamResponse } from '@/lib/llm/utils';
+import { getHistoryMessages, streamResponse } from '@/lib/llm/utils';
 import { logError } from '@/lib/log';
 import { GPT_4o_MIMI } from '@/lib/model';
 import { getSearchEngine } from '@/lib/search/search';
@@ -20,6 +20,7 @@ export async function o1Answer(
     isPro: boolean,
     userId: string,
     profile?: string,
+    summary?: string,
     onStream?: (...args: any[]) => void,
     model = GPT_4o_MIMI,
     source = SearchCategory.ALL,
@@ -32,7 +33,7 @@ export async function o1Answer(
         let images: ImageSource[] = [];
         let videos: VideoSource[] = [];
 
-        let history = getHistory(isPro, messages);
+        let history = getHistoryMessages(isPro, messages, summary);
 
         let imageFetchPromise;
         let videoFetchPromise;