From b0936e16ab5e261ae171cc9da8ce77824bd75210 Mon Sep 17 00:00:00 2001 From: Logan Yang Date: Sat, 8 Apr 2023 23:28:34 -0700 Subject: [PATCH] Fix bug duplicated first message in context --- src/components/Chat.tsx | 22 ++++++++------- src/utils.ts | 16 +++++------ tests/utils.test.ts | 59 +++++++++++++++++++++++++++++++++-------- 3 files changed, 67 insertions(+), 30 deletions(-) diff --git a/src/components/Chat.tsx b/src/components/Chat.tsx index a17c99f2..3d7a9f05 100644 --- a/src/components/Chat.tsx +++ b/src/components/Chat.tsx @@ -4,7 +4,7 @@ import { USER_SENDER, AI_SENDER } from '@/constants'; import { BotIcon, RefreshIcon, SaveAsNoteIcon, UseActiveNoteAsContextIcon } from '@/components/Icons'; -import { OpenAIStream } from '@/openAiStream'; +import { OpenAIStream, Role } from '@/openAiStream'; import ChatMessageComponent from '@/components/ChatMessageComponent'; import ReactMarkdown from '@/components/Markdown/MemoizedReactMarkdown'; import { getChatContext } from '@/utils'; @@ -69,20 +69,24 @@ const Chat: React.FC = ({ sharedState, apiKey, model }) => { setInputMessage(''); // The number of past messages to use as context for the AI - // Increase this number later as needed - const chatContext = getChatContext([...chatHistory, userMessage], 5); + // Use a even number. Increase this number later as needed + const chatContext = getChatContext(chatHistory, 4); // Use OpenAIStream to send message to AI and get a response try { const stream = await OpenAIStream( currentModel, apiKey, - chatContext.map((chatMessage) => { - return { - role: chatMessage.sender === USER_SENDER ? 'user' : 'assistant', - content: chatMessage.message, - }; - }), + [ + ...chatContext.map((chatMessage) => { + return { + role: chatMessage.sender === USER_SENDER + ? 'user' as Role : 'assistant' as Role, + content: chatMessage.message, + }; + }), + { role: 'user', content: userMessage.message }, + ], ); const reader = stream.getReader(); const decoder = new TextDecoder(); diff --git a/src/utils.ts b/src/utils.ts index 6e29163b..318f417e 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -2,22 +2,18 @@ import { ChatMessage } from '@/sharedState'; import { USER_SENDER } from '@/constants'; // Returns the last N messages from the chat history, -// last one being the new user message +// last one being the newest ai message export const getChatContext = (chatHistory: ChatMessage[], contextSize: number) => { if (chatHistory.length === 0) { return []; } - - const lastUserMessageIndex = chatHistory.slice().reverse().findIndex(msg => msg.sender === USER_SENDER); - - if (lastUserMessageIndex === -1) { - // No user messages found, return an empty array + const lastAiMessageIndex = chatHistory.slice().reverse().findIndex(msg => msg.sender !== USER_SENDER); + if (lastAiMessageIndex === -1) { + // No ai messages found, return an empty array return []; } - const lastIndex = chatHistory.length - 1 - lastUserMessageIndex; + const lastIndex = chatHistory.length - 1 - lastAiMessageIndex; const startIndex = Math.max(0, lastIndex - contextSize + 1); - - const chatContext = chatHistory.slice(startIndex, lastIndex + 1); - return chatContext; + return chatHistory.slice(startIndex, lastIndex + 1); }; diff --git a/tests/utils.test.ts b/tests/utils.test.ts index 00b40a10..2203559d 100644 --- a/tests/utils.test.ts +++ b/tests/utils.test.ts @@ -1,9 +1,18 @@ -// src/utils.test.ts import { getChatContext } from '@/utils'; import { ChatMessage } from '@/sharedState'; import { USER_SENDER, AI_SENDER } from '@/constants'; describe('getChatContext', () => { + const userMessage0: ChatMessage = { + message: 'Hey', + sender: USER_SENDER, + }; + + const aiMessage0: ChatMessage = { + message: 'Hi!', + sender: AI_SENDER, + }; + const userMessage1: ChatMessage = { message: 'Hello', sender: USER_SENDER, @@ -26,27 +35,55 @@ describe('getChatContext', () => { it('should return an empty context when chatHistory is empty', () => { const chatHistory: ChatMessage[] = []; - const context = getChatContext(chatHistory, 5); + const context = getChatContext(chatHistory, 4); + expect(context).toEqual([]); + }); + + it('should return the correct context with chatHistory length 1', () => { + const chatHistory: ChatMessage[] = [userMessage1]; + const context = getChatContext(chatHistory, 4); expect(context).toEqual([]); }); it('should return the correct context with chatHistory length 2', () => { const chatHistory: ChatMessage[] = [aiMessage1, userMessage2]; - const context = getChatContext(chatHistory, 5); - expect(context).toEqual([aiMessage1, userMessage2]); + const context = getChatContext(chatHistory, 4); + expect(context).toEqual([aiMessage1]); }); - it('should return the correct context with chatHistory length 3', () => { + it('should return the correct context with chatHistory length 3, user message is last', () => { const chatHistory: ChatMessage[] = [userMessage1, aiMessage1, userMessage2]; - const context = getChatContext(chatHistory, 5); - expect(context).toEqual([userMessage1, aiMessage1, userMessage2]); + const context = getChatContext(chatHistory, 4); + expect(context).toEqual([userMessage1, aiMessage1]); }); - it('should return the correct context with chatHistory length 4', () => { + it('should return the correct context with chatHistory length 3, ', () => { const chatHistory: ChatMessage[] = [ - aiMessage1, userMessage1, aiMessage2, userMessage2, + aiMessage0, userMessage1, aiMessage1 ]; - const context = getChatContext(chatHistory, 5); - expect(context).toEqual([aiMessage1, userMessage1, aiMessage2, userMessage2]); + const context = getChatContext(chatHistory, 4); + expect(context).toEqual([ + aiMessage0, userMessage1, aiMessage1 + ]); + }); + + it('should return the correct context with chatHistory length 5, n=2', () => { + const chatHistory: ChatMessage[] = [ + userMessage0, aiMessage0, userMessage1, aiMessage1, userMessage2, + ]; + const context = getChatContext(chatHistory, 2); + expect(context).toEqual([ + userMessage1, aiMessage1, + ]); + }); + + it('should return the correct context with chatHistory length 6', () => { + const chatHistory: ChatMessage[] = [ + userMessage0, aiMessage0, userMessage1, aiMessage1, userMessage2, aiMessage2, + ]; + const context = getChatContext(chatHistory, 4); + expect(context).toEqual([ + userMessage1, aiMessage1, userMessage2, aiMessage2 + ]); }); }); \ No newline at end of file