707 refactor of chatgptchatcompletion

backend/src/controller/chatController.ts

@@ -17,7 +17,7 @@ import {
 	defaultChatModel,
 } from '@src/models/chat';
 import { LEVEL_NAMES } from '@src/models/level';
-import { chatGptSendMessage } from '@src/openai';
+import { chatGptSendMessage, setSystemRoleInChatHistory } from '@src/openai';
 
 import { handleChatError } from './handleError';
 
@@ -29,7 +29,6 @@ async function handleLowLevelChat(
 	currentLevel: LEVEL_NAMES,
 	chatModel: ChatModel
 ) {
-	// get the chatGPT reply
 	const openAiReply = await chatGptSendMessage(
 		req.session.levelState[currentLevel].chatHistory,
 		req.session.levelState[currentLevel].defences,
@@ -177,13 +176,19 @@ async function handleChatToGPT(req: OpenAiChatRequest, res: Response) {
 			? req.session.chatModel
 			: defaultChatModel;
 
+	req.session.levelState[currentLevel].chatHistory = setSystemRoleInChatHistory(
+		currentLevel,
+		req.session.levelState[currentLevel].defences,
+		req.session.levelState[currentLevel].chatHistory
+	);
+
 	// record the history before chat completion called
 	const chatHistoryBefore = [
 		...req.session.levelState[currentLevel].chatHistory,
 	];
 	try {
-		// skip defence detection / blocking for levels 1 and 2 - sets chatResponse obj
 		if (currentLevel < LEVEL_NAMES.LEVEL_3) {
+			// skip defence detection / blocking for levels 1 and 2 - sets chatResponse obj
 			await handleLowLevelChat(
 				req,
 				message,

backend/src/openai.ts

@@ -218,50 +218,55 @@ async function chatGptCallFunction(
 	};
 }
 
-async function chatGptChatCompletion(
-	chatResponse: ChatResponse,
-	chatHistory: ChatHistoryMessage[],
+function setSystemRoleInChatHistory(
+	currentLevel: LEVEL_NAMES,
 	defences: Defence[],
-	chatModel: ChatModel,
-	openai: OpenAI,
-	// default to sandbox
-	// eslint-disable-next-line @typescript-eslint/no-unused-vars
-	currentLevel: LEVEL_NAMES = LEVEL_NAMES.SANDBOX
+	chatHistory: ChatHistoryMessage[]
 ) {
-	// check if we need to set a system role
-	// system role is always active on levels
-	if (
-		currentLevel !== LEVEL_NAMES.SANDBOX ||
-		isDefenceActive(DEFENCE_ID.SYSTEM_ROLE, defences)
-	) {
+	const systemRoleNeededInChatHistory =
+		currentLevel === LEVEL_NAMES.SANDBOX
+			? isDefenceActive(DEFENCE_ID.SYSTEM_ROLE, defences)
+			: true;
+
+	if (systemRoleNeededInChatHistory) {
 		const completionConfig: ChatCompletionSystemMessageParam = {
 			role: 'system',
 			content: getSystemRole(defences, currentLevel),
 		};
 
-		// check to see if there's already a system role
-		const systemRole = chatHistory.find(
+		const existingSystemRole = chatHistory.find(
 			(message) => message.completion?.role === 'system'
 		);
-		if (!systemRole) {
-			// add the system role to the start of the chat history
-			chatHistory.unshift({
-				completion: completionConfig,
-				chatMessageType: CHAT_MESSAGE_TYPE.SYSTEM,
-			});
+		if (!existingSystemRole) {
+			return [
+				{
+					completion: completionConfig,
+					chatMessageType: CHAT_MESSAGE_TYPE.SYSTEM,
+				},
+				...chatHistory,
+			];
 		} else {
-			// replace with the latest system role
-			systemRole.completion = completionConfig;
+			return chatHistory.map((message) => {
+				if (message.completion?.role === 'system') {
+					return { ...existingSystemRole, completion: completionConfig };
+				} else {
+					return message;
+				}
+			});
 		}
 	} else {
-		// remove the system role from the chat history
-		while (
-			chatHistory.length > 0 &&
-			chatHistory[0].completion?.role === 'system'
-		) {
-			chatHistory.shift();
-		}
+		return chatHistory.filter(
+			(message) => message.completion?.role !== 'system'
+		);
 	}
+}
+
+async function chatGptChatCompletion(
+	chatResponse: ChatResponse,
+	chatHistory: ChatHistoryMessage[],
+	chatModel: ChatModel,
+	openai: OpenAI
+) {
 	console.debug('Talking to model: ', JSON.stringify(chatModel));
 
 	// get start time
@@ -449,10 +454,8 @@ async function getFinalReplyAfterAllToolCalls(
 	let reply = await chatGptChatCompletion(
 		chatResponse,
 		chatHistory,
-		defences,
 		chatModel,
-		openai,
-		currentLevel
+		openai
 	);
 
 	// check if GPT wanted to call a tool
@@ -477,10 +480,8 @@ async function getFinalReplyAfterAllToolCalls(
 		reply = await chatGptChatCompletion(
 			chatResponse,
 			chatHistory,
-			defences,
 			chatModel,
-			openai,
-			currentLevel
+			openai
 		);
 	}
 
@@ -551,4 +552,5 @@ export {
 	chatGptSendMessage,
 	getOpenAIKey,
 	getValidModelsFromOpenAI,
+	setSystemRoleInChatHistory,
 };

backend/test/integration/chatController.test.ts

@@ -8,6 +8,7 @@ import { ChatHistoryMessage, ChatModel } from '@src/models/chat';
 import { Defence } from '@src/models/defence';
 import { EmailInfo } from '@src/models/email';
 import { LEVEL_NAMES, LevelState } from '@src/models/level';
+import * as openaiModule from '@src/openai';
 
 declare module 'express-session' {
 	interface Session {
@@ -152,12 +153,17 @@ describe('handleChatToGPT integration tests', () => {
 		const req = openAiChatRequestMock('Hello chatbot', LEVEL_NAMES.LEVEL_1);
 		const res = responseMock();
 
+		const mockedSetSystemRoleInChatHistory = jest.spyOn(
+			openaiModule,
+			'setSystemRoleInChatHistory'
+		);
 		mockCreateChatCompletion.mockResolvedValueOnce(
 			chatResponseAssistant('Howdy human!')
 		);
 
 		await handleChatToGPT(req, res);
 
+		expect(mockedSetSystemRoleInChatHistory).toHaveBeenCalled();
 		expect(res.send).toHaveBeenCalledWith({
 			reply: 'Howdy human!',
 			defenceReport: {