Fix bug duplicated first message in context

This commit is contained in:
Logan Yang 2023-04-08 23:28:34 -07:00
parent 02d91ed3da
commit b0936e16ab
3 changed files with 67 additions and 30 deletions

View file

@ -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<ChatProps> = ({ 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();

View file

@ -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);
};

View file

@ -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
]);
});
});