Stop sending full defences on low levels

backend/src/controller/chatController.ts

@@ -22,7 +22,7 @@ import {
 	ChatInfoMessage,
 	chatInfoMessageTypes,
 } from '@src/models/chatMessage';
-import { Defence } from '@src/models/defence';
+import { DEFENCE_ID, Defence, QaLlmDefence } from '@src/models/defence';
 import { EmailInfo } from '@src/models/email';
 import { LEVEL_NAMES } from '@src/models/level';
 import { chatGptSendMessage } from '@src/openai';
@@ -95,8 +95,7 @@ async function handleChatWithoutDefenceDetection(
 	chatResponse: ChatHttpResponse,
 	currentLevel: LEVEL_NAMES,
 	chatModel: ChatModel,
-	chatHistory: ChatMessage[],
-	defences: Defence[]
+	chatHistory: ChatMessage[]
 ): Promise<LevelHandlerResponse> {
 	console.log(`User message: '${message}'`);
 
@@ -107,7 +106,6 @@ async function handleChatWithoutDefenceDetection(
 
 	const openAiReply = await chatGptSendMessage(
 		updatedChatHistory,
-		defences,
 		chatModel,
 		currentLevel
 	);
@@ -150,11 +148,15 @@ async function handleChatWithDefenceDetection(
 		}'`
 	);
 
+	const qaLlmDefence = defences.find(
+		(defence) => defence.id === DEFENCE_ID.QA_LLM
+	) as QaLlmDefence | undefined;
+
 	const openAiReplyPromise = chatGptSendMessage(
 		chatHistoryWithNewUserMessages,
-		defences,
 		chatModel,
-		currentLevel
+		currentLevel,
+		qaLlmDefence
 	);
 
 	// run input defence detection and chatGPT concurrently
@@ -258,8 +260,7 @@ async function handleChatToGPT(req: OpenAiChatRequest, res: Response) {
 				initChatResponse,
 				currentLevel,
 				chatModel,
-				currentChatHistory,
-				defences
+				currentChatHistory
 			);
 		} else {
 			levelResult = await handleChatWithDefenceDetection(

backend/src/models/defence.ts

@@ -23,11 +23,24 @@ type DefenceConfigItem = {
 	value: string;
 };
 
-type Defence = {
-	id: DEFENCE_ID;
+type QaLlmDefence = {
+	id: DEFENCE_ID.QA_LLM;
 	config: DefenceConfigItem[];
 	isActive: boolean;
 };
 
+type Defence =
+	| {
+			id: DEFENCE_ID;
+			config: DefenceConfigItem[];
+			isActive: boolean;
+	  }
+	| QaLlmDefence;
+
 export { DEFENCE_ID };
-export type { Defence, DefenceConfigItem, DEFENCE_CONFIG_ITEM_ID };
+export type {
+	Defence,
+	DefenceConfigItem,
+	DEFENCE_CONFIG_ITEM_ID,
+	QaLlmDefence,
+};

backend/src/openai.ts

@@ -5,7 +5,7 @@ import {
 	ChatCompletionMessageToolCall,
 } from 'openai/resources/chat/completions';
 
-import { isDefenceActive, getQAPromptFromConfig } from './defence';
+import { getQAPromptFromConfig } from './defence';
 import { sendEmail } from './email';
 import { queryDocuments } from './langchain';
 import {
@@ -17,7 +17,7 @@ import {
 	ToolCallResponse,
 } from './models/chat';
 import { ChatMessage } from './models/chatMessage';
-import { DEFENCE_ID, Defence } from './models/defence';
+import { QaLlmDefence } from './models/defence';
 import { EmailResponse } from './models/email';
 import { LEVEL_NAMES } from './models/level';
 import {
@@ -141,14 +141,14 @@ function isChatGptFunction(functionName: string) {
 async function handleAskQuestionFunction(
 	functionCallArgs: string | undefined,
 	currentLevel: LEVEL_NAMES,
-	defences: Defence[]
+	qaLlmDefence?: QaLlmDefence
 ) {
 	if (functionCallArgs) {
 		const params = JSON.parse(functionCallArgs) as FunctionAskQuestionParams;
 		console.debug(`Asking question: ${params.question}`);
 		// if asking a question, call the queryDocuments
-		const configQAPrompt = isDefenceActive(DEFENCE_ID.QA_LLM, defences)
-			? getQAPromptFromConfig(defences)
+		const configQAPrompt = qaLlmDefence?.isActive
+			? getQAPromptFromConfig([qaLlmDefence])
 			: '';
 		return await queryDocuments(params.question, configQAPrompt, currentLevel);
 	} else {
@@ -191,11 +191,11 @@ function handleSendEmailFunction(
 }
 
 async function chatGptCallFunction(
-	defences: Defence[],
 	toolCallId: string,
 	functionCall: ChatCompletionMessageToolCall.Function,
 	// default to sandbox
-	currentLevel: LEVEL_NAMES = LEVEL_NAMES.SANDBOX
+	currentLevel: LEVEL_NAMES = LEVEL_NAMES.SANDBOX,
+	qaLlmDefence?: QaLlmDefence
 ): Promise<FunctionCallResponse> {
 	const functionName = functionCall.name;
 	let functionReply = '';
@@ -220,7 +220,7 @@ async function chatGptCallFunction(
 			functionReply = await handleAskQuestionFunction(
 				functionCall.arguments,
 				currentLevel,
-				defences
+				qaLlmDefence
 			);
 		}
 	} else {
@@ -326,18 +326,18 @@ function getChatCompletionsInContextWindow(
 async function performToolCalls(
 	toolCalls: ChatCompletionMessageToolCall[],
 	chatHistory: ChatMessage[],
-	defences: Defence[],
-	currentLevel: LEVEL_NAMES
+	currentLevel: LEVEL_NAMES,
+	qaLlmDefence?: QaLlmDefence
 ): Promise<ToolCallResponse> {
 	for (const toolCall of toolCalls) {
 		// only tool type supported by openai is function
 		// eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
 		if (toolCall.type === 'function') {
 			const functionCallReply = await chatGptCallFunction(
-				defences,
 				toolCall.id,
 				toolCall.function,
-				currentLevel
+				currentLevel,
+				qaLlmDefence
 			);
 
 			// We assume only one function call in toolCalls, and so we return after getting function reply.
@@ -358,9 +358,9 @@ async function performToolCalls(
 
 async function getFinalReplyAfterAllToolCalls(
 	chatHistory: ChatMessage[],
-	defences: Defence[],
 	chatModel: ChatModel,
-	currentLevel: LEVEL_NAMES
+	currentLevel: LEVEL_NAMES,
+	qaLlmDefence?: QaLlmDefence
 ) {
 	let updatedChatHistory = [...chatHistory];
 	const sentEmails = [];
@@ -385,8 +385,8 @@ async function getFinalReplyAfterAllToolCalls(
 			const toolCallReply = await performToolCalls(
 				gptReply.completion.tool_calls,
 				updatedChatHistory,
-				defences,
-				currentLevel
+				currentLevel,
+				qaLlmDefence
 			);
 
 			updatedChatHistory = toolCallReply.chatHistory;
@@ -408,17 +408,17 @@ async function getFinalReplyAfterAllToolCalls(
 
 async function chatGptSendMessage(
 	chatHistory: ChatMessage[],
-	defences: Defence[],
 	chatModel: ChatModel,
-	currentLevel: LEVEL_NAMES = LEVEL_NAMES.SANDBOX
+	currentLevel: LEVEL_NAMES = LEVEL_NAMES.SANDBOX,
+	qaLlmDefence?: QaLlmDefence
 ) {
 	// this method just calls getFinalReplyAfterAllToolCalls then reformats the output. Does it need to exist?
 
 	const finalToolCallResponse = await getFinalReplyAfterAllToolCalls(
 		chatHistory,
-		defences,
 		chatModel,
-		currentLevel
+		currentLevel,
+		qaLlmDefence
 	);
 
 	const chatResponse: ChatResponse = {

backend/test/integration/openai.test.ts

@@ -1,9 +1,7 @@
 import { expect, jest, test, describe } from '@jest/globals';
 
-import { defaultDefences } from '@src/defaultDefences';
 import { CHAT_MODELS, ChatModel } from '@src/models/chat';
 import { ChatMessage } from '@src/models/chatMessage';
-import { Defence } from '@src/models/defence';
 import { chatGptSendMessage } from '@src/openai';
 
 const mockCreateChatCompletion =
@@ -55,7 +53,6 @@ describe('OpenAI Integration Tests', () => {
 				},
 			},
 		];
-		const defences: Defence[] = defaultDefences;
 		const chatModel: ChatModel = {
 			id: CHAT_MODELS.GPT_4,
 			configuration: {
@@ -68,11 +65,7 @@ describe('OpenAI Integration Tests', () => {
 
 		mockCreateChatCompletion.mockResolvedValueOnce(chatResponseAssistant('Hi'));
 
-		const reply = await chatGptSendMessage(
-			chatHistoryWithMessage,
-			defences,
-			chatModel
-		);
+		const reply = await chatGptSendMessage(chatHistoryWithMessage, chatModel);
 
 		expect(reply).toBeDefined();
 		expect(reply.chatResponse.completion).toBeDefined();

backend/test/unit/controller/chatController.test.ts

@@ -16,7 +16,7 @@ import {
 	MessageTransformation,
 } from '@src/models/chat';
 import { ChatMessage } from '@src/models/chatMessage';
-import { DEFENCE_ID, Defence } from '@src/models/defence';
+import { DEFENCE_ID, Defence, QaLlmDefence } from '@src/models/defence';
 import { EmailInfo } from '@src/models/email';
 import { LEVEL_NAMES, LevelState } from '@src/models/level';
 import { chatGptSendMessage } from '@src/openai';
@@ -497,7 +497,6 @@ describe('handleChatToGPT unit tests', () => {
 
 			expect(mockChatGptSendMessage).toHaveBeenCalledWith(
 				[...existingHistory, newUserChatMessage],
-				[],
 				mockChatModel,
 				LEVEL_NAMES.LEVEL_1
 			);
@@ -558,10 +557,18 @@ describe('handleChatToGPT unit tests', () => {
 				},
 			} as ChatMessage;
 
+			const qaLllmDefence = {
+				id: DEFENCE_ID.QA_LLM,
+				isActive: true,
+				config: [{ id: 'PROMPT', value: 'query them documents!' }],
+			} as QaLlmDefence;
+
 			const req = openAiChatRequestMock(
 				'send an email to [email protected] saying hi',
 				LEVEL_NAMES.SANDBOX,
-				existingHistory
+				existingHistory,
+				[],
+				[qaLllmDefence]
 			);
 			const res = responseMock();
 
@@ -590,9 +597,9 @@ describe('handleChatToGPT unit tests', () => {
 
 			expect(mockChatGptSendMessage).toHaveBeenCalledWith(
 				[...existingHistory, newUserChatMessage],
-				[],
 				mockChatModel,
-				LEVEL_NAMES.SANDBOX
+				LEVEL_NAMES.SANDBOX,
+				qaLllmDefence
 			);
 
 			expect(res.send).toHaveBeenCalledWith({
@@ -655,10 +662,18 @@ describe('handleChatToGPT unit tests', () => {
 				},
 			} as ChatMessage;
 
+			const qaLllmDefence = {
+				id: DEFENCE_ID.QA_LLM,
+				isActive: true,
+				config: [{ id: 'PROMPT', value: 'query them documents!' }],
+			} as QaLlmDefence;
+
 			const req = openAiChatRequestMock(
 				'hello bot',
 				LEVEL_NAMES.SANDBOX,
-				existingHistory
+				existingHistory,
+				[],
+				[qaLllmDefence]
 			);
 			const res = responseMock();
 
@@ -690,9 +705,9 @@ describe('handleChatToGPT unit tests', () => {
 
 			expect(mockChatGptSendMessage).toHaveBeenCalledWith(
 				[...existingHistory, ...newTransformationChatMessages],
-				[],
 				mockChatModel,
-				LEVEL_NAMES.SANDBOX
+				LEVEL_NAMES.SANDBOX,
+				qaLllmDefence
 			);
 
 			expect(res.send).toHaveBeenCalledWith({