741 fix backend chat history message types

backend/src/controller/chatController.ts

@@ -5,20 +5,23 @@ import {
 	detectTriggeredInputDefences,
 	detectTriggeredOutputDefences,
 } from '@src/defence';
-import { OpenAiAddHistoryRequest } from '@src/models/api/OpenAiAddHistoryRequest';
+import { OpenAiAddInfoToChatHistoryRequest } from '@src/models/api/OpenAiAddInfoToChatHistoryRequest';
 import { OpenAiChatRequest } from '@src/models/api/OpenAiChatRequest';
 import { OpenAiClearRequest } from '@src/models/api/OpenAiClearRequest';
 import { OpenAiGetHistoryRequest } from '@src/models/api/OpenAiGetHistoryRequest';
 import {
-	CHAT_MESSAGE_TYPE,
 	ChatDefenceReport,
-	ChatHistoryMessage,
 	ChatHttpResponse,
 	ChatModel,
 	LevelHandlerResponse,
 	MessageTransformation,
 	defaultChatModel,
 } from '@src/models/chat';
+import {
+	ChatMessage,
+	ChatInfoMessage,
+	chatInfoMessageType,
+} from '@src/models/chatMessage';
 import { Defence } from '@src/models/defence';
 import { EmailInfo } from '@src/models/email';
 import { LEVEL_NAMES } from '@src/models/level';
@@ -47,25 +50,23 @@ function combineChatDefenceReports(
 function createNewUserMessages(
 	message: string,
 	messageTransformation?: MessageTransformation
-): ChatHistoryMessage[] {
+): ChatMessage[] {
 	if (messageTransformation) {
 		return [
 			{
-				completion: null,
-				chatMessageType: CHAT_MESSAGE_TYPE.USER,
+				chatMessageType: 'USER',
 				infoMessage: message,
 			},
 			{
-				completion: null,
-				chatMessageType: CHAT_MESSAGE_TYPE.INFO,
+				chatMessageType: 'GENERIC_INFO',
 				infoMessage: messageTransformation.transformedMessageInfo,
 			},
 			{
 				completion: {
 					role: 'user',
 					content: messageTransformation.transformedMessageCombined,
 				},
-				chatMessageType: CHAT_MESSAGE_TYPE.USER_TRANSFORMED,
+				chatMessageType: 'USER_TRANSFORMED',
 				transformedMessage: messageTransformation.transformedMessage,
 			},
 		];
@@ -76,7 +77,7 @@ function createNewUserMessages(
 					role: 'user',
 					content: message,
 				},
-				chatMessageType: CHAT_MESSAGE_TYPE.USER,
+				chatMessageType: 'USER',
 			},
 		];
 	}
@@ -87,7 +88,7 @@ async function handleChatWithoutDefenceDetection(
 	chatResponse: ChatHttpResponse,
 	currentLevel: LEVEL_NAMES,
 	chatModel: ChatModel,
-	chatHistory: ChatHistoryMessage[],
+	chatHistory: ChatMessage[],
 	defences: Defence[]
 ): Promise<LevelHandlerResponse> {
 	const updatedChatHistory = createNewUserMessages(message).reduce(
@@ -122,7 +123,7 @@ async function handleChatWithDefenceDetection(
 	chatResponse: ChatHttpResponse,
 	currentLevel: LEVEL_NAMES,
 	chatModel: ChatModel,
-	chatHistory: ChatHistoryMessage[],
+	chatHistory: ChatMessage[],
 	defences: Defence[]
 ): Promise<LevelHandlerResponse> {
 	const messageTransformation = transformMessage(message, defences);
@@ -162,11 +163,7 @@ async function handleChatWithDefenceDetection(
 
 	// if blocked, restore original chat history and add user message to chat history without completion
 	const updatedChatHistory = combinedDefenceReport.isBlocked
-		? pushMessageToHistory(chatHistory, {
-				completion: null,
-				chatMessageType: CHAT_MESSAGE_TYPE.USER,
-				infoMessage: message,
-		  })
+		? chatHistoryWithNewUserMessages
 		: openAiReply.chatHistory;
 
 	const updatedChatResponse: ChatHttpResponse = {
@@ -284,9 +281,10 @@ async function handleChatToGPT(req: OpenAiChatRequest, res: Response) {
 	if (updatedChatResponse.defenceReport.isBlocked) {
 		// chatReponse.reply is empty if blocked
 		updatedChatHistory = pushMessageToHistory(updatedChatHistory, {
-			completion: null,
-			chatMessageType: CHAT_MESSAGE_TYPE.BOT_BLOCKED,
-			infoMessage: updatedChatResponse.defenceReport.blockedReason,
+			chatMessageType: 'BOT_BLOCKED',
+			infoMessage:
+				updatedChatResponse.defenceReport.blockedReason ??
+				'block reason unknown',
 		});
 	} else if (updatedChatResponse.openAIErrorMessage) {
 		const errorMsg = simplifyOpenAIErrorMessage(
@@ -307,13 +305,12 @@ async function handleChatToGPT(req: OpenAiChatRequest, res: Response) {
 		handleChatError(res, updatedChatResponse, errorMsg, 500);
 		return;
 	} else {
-		// add bot message to chat history
 		updatedChatHistory = pushMessageToHistory(updatedChatHistory, {
 			completion: {
 				role: 'assistant',
 				content: updatedChatResponse.reply,
 			},
-			chatMessageType: CHAT_MESSAGE_TYPE.BOT,
+			chatMessageType: 'BOT',
 		});
 	}
 
@@ -339,13 +336,12 @@ function simplifyOpenAIErrorMessage(openAIErrorMessage: string) {
 }
 
 function addErrorToChatHistory(
-	chatHistory: ChatHistoryMessage[],
+	chatHistory: ChatMessage[],
 	errorMessage: string
-): ChatHistoryMessage[] {
+): ChatMessage[] {
 	console.error(errorMessage);
 	return pushMessageToHistory(chatHistory, {
-		completion: null,
-		chatMessageType: CHAT_MESSAGE_TYPE.ERROR_MSG,
+		chatMessageType: 'ERROR_MSG',
 		infoMessage: errorMessage,
 	});
 }
@@ -360,23 +356,26 @@ function handleGetChatHistory(req: OpenAiGetHistoryRequest, res: Response) {
 	}
 }
 
-function handleAddToChatHistory(req: OpenAiAddHistoryRequest, res: Response) {
-	const infoMessage = req.body.message;
+function handleAddInfoToChatHistory(
+	req: OpenAiAddInfoToChatHistoryRequest,
+	res: Response
+) {
+	const infoMessage = req.body.infoMessage;
 	const chatMessageType = req.body.chatMessageType;
 	const level = req.body.level;
 	if (
 		infoMessage &&
 		chatMessageType &&
+		chatInfoMessageType.includes(chatMessageType) &&
 		level !== undefined &&
 		level >= LEVEL_NAMES.LEVEL_1
 	) {
 		req.session.levelState[level].chatHistory = pushMessageToHistory(
 			req.session.levelState[level].chatHistory,
 			{
-				completion: null,
 				chatMessageType,
 				infoMessage,
-			}
+			} as ChatInfoMessage
 		);
 		res.send();
 	} else {
@@ -400,6 +399,6 @@ function handleClearChatHistory(req: OpenAiClearRequest, res: Response) {
 export {
 	handleChatToGPT,
 	handleGetChatHistory,
-	handleAddToChatHistory,
+	handleAddInfoToChatHistory as handleAddInfoToChatHistory,
 	handleClearChatHistory,
 };

backend/src/models/api/OpenAiAddInfoToChatHistoryRequest.ts

@@ -0,0 +1,16 @@
+import { Request } from 'express';
+
+import { CHAT_INFO_MESSAGE_TYPE } from '@src/models/chatMessage';
+import { LEVEL_NAMES } from '@src/models/level';
+
+export type OpenAiAddInfoToChatHistoryRequest = Request<
+	never,
+	never,
+	{
+		chatMessageType?: CHAT_INFO_MESSAGE_TYPE;
+		infoMessage?: string;
+		level?: LEVEL_NAMES;
+	},
+	never,
+	never
+>;

backend/src/models/api/OpenAiGetHistoryRequest.ts

@@ -1,10 +1,10 @@
 import { Request } from 'express';
 
-import { ChatHistoryMessage } from '@src/models/chat';
+import { ChatMessage } from '@src/models/chatMessage';
 
 export type OpenAiGetHistoryRequest = Request<
 	never,
-	ChatHistoryMessage[] | string,
+	ChatMessage[] | string,
 	never,
 	{
 		level?: string;

backend/src/models/chat.ts

@@ -1,8 +1,9 @@
 import {
-	ChatCompletionMessage,
+	ChatCompletionAssistantMessageParam,
 	ChatCompletionMessageParam,
 } from 'openai/resources/chat/completions';
 
+import { ChatMessage } from './chatMessage';
 import { DEFENCE_ID } from './defence';
 import { EmailInfo } from './email';
 
@@ -16,21 +17,6 @@ enum CHAT_MODELS {
 	GPT_3_5_TURBO_16K_0613 = 'gpt-3.5-turbo-16k-0613',
 }
 
-enum CHAT_MESSAGE_TYPE {
-	BOT,
-	BOT_BLOCKED,
-	INFO,
-	USER,
-	USER_TRANSFORMED,
-	LEVEL_INFO,
-	DEFENCE_ALERTED,
-	DEFENCE_TRIGGERED,
-	SYSTEM,
-	FUNCTION_CALL,
-	ERROR_MSG,
-	RESET_LEVEL,
-}
-
 enum MODEL_CONFIG {
 	TEMPERATURE = 'temperature',
 	TOP_P = 'topP',
@@ -72,7 +58,7 @@ interface FunctionCallResponse {
 interface ToolCallResponse {
 	functionCallReply?: FunctionCallResponse;
 	chatResponse?: ChatResponse;
-	chatHistory: ChatHistoryMessage[];
+	chatHistory: ChatMessage[];
 }
 
 interface ChatAnswer {
@@ -92,8 +78,8 @@ interface ChatResponse {
 }
 
 interface ChatGptReply {
-	chatHistory: ChatHistoryMessage[];
-	completion: ChatCompletionMessage | null;
+	chatHistory: ChatMessage[];
+	completion: ChatCompletionAssistantMessageParam | null;
 	openAIErrorMessage: string | null;
 }
 
@@ -123,17 +109,9 @@ interface ChatHttpResponse {
 
 interface LevelHandlerResponse {
 	chatResponse: ChatHttpResponse;
-	chatHistory: ChatHistoryMessage[];
-}
-
-interface ChatHistoryMessage {
-	completion: ChatCompletionMessageParam | null;
-	chatMessageType: CHAT_MESSAGE_TYPE;
-	infoMessage?: string | null;
-	transformedMessage?: TransformedChatMessage;
+	chatHistory: ChatMessage[];
 }
 
-// default settings for chat model
 const defaultChatModel: ChatModel = {
 	id: CHAT_MODELS.GPT_3_5_TURBO,
 	configuration: {
@@ -154,11 +132,10 @@ export type {
 	ChatResponse,
 	LevelHandlerResponse,
 	ChatHttpResponse,
-	ChatHistoryMessage,
-	SingleDefenceReport,
 	TransformedChatMessage,
 	FunctionCallResponse,
 	ToolCallResponse,
 	MessageTransformation,
+	SingleDefenceReport,
 };
-export { CHAT_MODELS, CHAT_MESSAGE_TYPE, MODEL_CONFIG, defaultChatModel };
+export { CHAT_MODELS, MODEL_CONFIG, defaultChatModel };

backend/src/models/chatMessage.ts

@@ -0,0 +1,70 @@
+import {
+	ChatCompletionAssistantMessageParam,
+	ChatCompletionMessageParam,
+	ChatCompletionSystemMessageParam,
+	ChatCompletionUserMessageParam,
+} from 'openai/resources/chat/completions';
+
+import { TransformedChatMessage } from './chat';
+
+const chatInfoMessageType = [
+	'DEFENCE_ALERTED',
+	'DEFENCE_TRIGGERED',
+	'LEVEL_INFO',
+	'RESET_LEVEL',
+	'ERROR_MSG',
+	'BOT_BLOCKED',
+	'USER',
+	'GENERIC_INFO',
+] as const;
+
+type CHAT_INFO_MESSAGE_TYPE = (typeof chatInfoMessageType)[number];
+
+type ChatInfoMessage = {
+	chatMessageType: CHAT_INFO_MESSAGE_TYPE;
+	infoMessage: string;
+};
+
+type ChatFunctionCallMessage = {
+	completion: ChatCompletionMessageParam;
+	chatMessageType: 'FUNCTION_CALL';
+};
+
+type ChatSystemMessage = {
+	completion: ChatCompletionSystemMessageParam;
+	chatMessageType: 'SYSTEM';
+};
+
+type ChatBotMessage = {
+	completion: ChatCompletionAssistantMessageParam;
+	chatMessageType: 'BOT';
+};
+
+type ChatUserMessageAsCompletion = {
+	completion: ChatCompletionUserMessageParam;
+	chatMessageType: 'USER';
+};
+
+type ChatUserTransformedMessage = {
+	completion: ChatCompletionUserMessageParam;
+	chatMessageType: 'USER_TRANSFORMED';
+	transformedMessage: TransformedChatMessage;
+};
+
+type ChatCompletionMessage =
+	| ChatFunctionCallMessage
+	| ChatSystemMessage
+	| ChatBotMessage
+	| ChatUserMessageAsCompletion
+	| ChatUserTransformedMessage;
+
+type ChatMessage = ChatInfoMessage | ChatCompletionMessage;
+
+export type {
+	ChatMessage,
+	ChatSystemMessage,
+	ChatInfoMessage,
+	CHAT_INFO_MESSAGE_TYPE,
+};
+
+export { chatInfoMessageType };