Merge dev for v1.0.3 release (#919)

ScottLogic · Jul 16, 2024 · 9ba0435 · 9ba0435
2 parents 40bab47 + 58f0252
commit 9ba0435
Show file tree

Hide file tree

Showing 61 changed files with 2,388 additions and 2,132 deletions.
diff --git a/backend/package-lock.json b/backend/package-lock.json
diff --git a/backend/package.json b/backend/package.json
@@ -1,6 +1,6 @@
 {
 	"name": "prompt-injection-api",
-	"version": "1.0.0",
+	"version": "1.0.3",
 	"type": "module",
 	"scripts": {
 		"build": "tsc --noEmit",
@@ -28,20 +28,20 @@
 		"@types/cors": "^2.8.17",
 		"@types/express": "^4.17.21",
 		"@types/express-session": "^1.18.0",
-		"@types/node": "^20.12.7",
+		"@types/node": "^20.14.10",
 		"@types/supertest": "^6.0.2",
-		"@typescript-eslint/eslint-plugin": "^7.6.0",
-		"@typescript-eslint/parser": "^7.6.0",
+		"@typescript-eslint/eslint-plugin": "^7.16.0",
+		"@typescript-eslint/parser": "^7.16.0",
 		"cross-env": "^7.0.3",
 		"eslint": "^8.57.0",
 		"eslint-config-prettier": "^9.1.0",
 		"eslint-import-resolver-alias": "^1.1.2",
 		"eslint-plugin-import": "^2.29.1",
-		"eslint-plugin-jest": "^28.2.0",
-		"prettier": "^3.2.5",
+		"eslint-plugin-jest": "^28.6.0",
+		"prettier": "^3.3.3",
 		"supertest": "^7.0.0",
-		"ts-jest": "^29.1.2",
-		"tsx": "^4.7.2",
-		"typescript": "^5.4.5"
+		"ts-jest": "^29.2.2",
+		"tsx": "^4.16.2",
+		"typescript": "^5.5.3"
 	}
 }
diff --git a/backend/src/controller/defenceController.ts b/backend/src/controller/defenceController.ts
@@ -254,14 +254,16 @@ function handleResetDefenceConfigItem(
 		currentDefences
 	);
 
-	const updatedDefences = req.session.levelState[level].defences;
-	if (updatedDefences === undefined) {
+	const updatedDefence = req.session.levelState[level].defences.find(
+		(defence) => defence.id === defenceId
+	);
+	if (!updatedDefence) {
 		sendErrorResponse(res, 500, 'Something went whacky');
 		return;
 	}
-	const updatedDefenceConfig = updatedDefences
-		.find((defence) => defence.id === defenceId)
-		?.config.find((config) => config.id === configItemId);
+	const updatedDefenceConfig = updatedDefence.config.find(
+		(config) => config.id === configItemId
+	);
 
 	res.send(updatedDefenceConfig);
 }

diff --git a/backend/src/controller/documentController.ts b/backend/src/controller/documentController.ts
diff --git a/backend/src/controller/levelController.ts b/backend/src/controller/levelController.ts
@@ -1,27 +1,32 @@
 import { Response } from 'express';
 
-import { LevelGetRequest } from '@src/models/api/LevelGetRequest';
-import { LEVEL_NAMES, isValidLevel } from '@src/models/level';
+import { getSandboxDocumentMetas } from '@src/document';
+import {
+	LevelGetRequest,
+	LevelGetResponseBody,
+} from '@src/models/api/LevelGetRequest';
+import { LEVEL_NAMES } from '@src/models/level';
 
-function handleLoadLevel(req: LevelGetRequest, res: Response) {
-	const { level } = req.query;
+import { validateLevel } from './requestValidators';
 
-	if (level === undefined) {
-		res.status(400).send('Level not provided');
-		return;
-	}
+function handleLoadLevel(
+	req: LevelGetRequest,
+	res: Response<LevelGetResponseBody>
+) {
+	const level = validateLevel(res, req.query.level);
+	if (level === null) return;
 
-	if (!isValidLevel(level)) {
-		res.status(400).send('Invalid level');
-		return;
-	}
+	const chatModel =
+		level === LEVEL_NAMES.SANDBOX ? req.session.chatModel : undefined;
+	const availableDocs =
+		level === LEVEL_NAMES.SANDBOX ? getSandboxDocumentMetas() : undefined;
 
 	res.send({
 		emails: req.session.levelState[level].sentEmails,
 		chatHistory: req.session.levelState[level].chatHistory,
 		defences: req.session.levelState[level].defences,
-		chatModel:
-			level === LEVEL_NAMES.SANDBOX ? req.session.chatModel : undefined,
+		chatModel,
+		availableDocs,
 	});
 }
 

diff --git a/backend/src/controller/requestValidators.ts b/backend/src/controller/requestValidators.ts
@@ -0,0 +1,18 @@
+import { Response } from 'express';
+
+import { sendErrorResponse } from '@src/controller/handleError';
+import { isValidLevel, LEVEL_NAMES } from '@src/models/level';
+
+export function validateLevel(response: Response, level?: LEVEL_NAMES) {
+	if (level === undefined) {
+		sendErrorResponse(response, 400, 'Level not provided');
+		return null;
+	}
+
+	if (!isValidLevel(level)) {
+		sendErrorResponse(response, 400, 'Invalid level');
+		return null;
+	}
+
+	return level;
+}
diff --git a/backend/src/controller/resetController.ts b/backend/src/controller/resetController.ts
@@ -1,58 +1,59 @@
 import { Response } from 'express';
 
-import { LevelGetRequest } from '@src/models/api/LevelGetRequest';
-import { LevelResetRequest } from '@src/models/api/LevelResetRequest';
-import { defaultChatModel } from '@src/models/chat';
-import { ChatInfoMessage } from '@src/models/chatMessage';
+import { getSandboxDocumentMetas } from '@src/document';
 import {
-	LEVEL_NAMES,
-	getInitialLevelStates,
-	isValidLevel,
-} from '@src/models/level';
-
-function handleResetProgress(req: LevelGetRequest, res: Response) {
-	const { level } = req.query;
+	LevelResetRequest,
+	LevelResetResponseBody,
+} from '@src/models/api/LevelResetRequest';
+import {
+	ProgressResetRequest,
+	ProgressResetResponseBody,
+} from '@src/models/api/ProgressResetRequest';
+import { defaultChatModel } from '@src/models/chat';
+import { LEVEL_NAMES, getInitialLevelStates } from '@src/models/level';
 
-	if (level === undefined) {
-		res.status(400).send('Level not provided');
-		return;
-	}
+import { validateLevel } from './requestValidators';
 
-	if (!isValidLevel(level)) {
-		res.status(400).send('Invalid level');
-		return;
-	}
+function handleResetProgress(
+	req: ProgressResetRequest,
+	res: Response<ProgressResetResponseBody>
+) {
+	const level = validateLevel(res, req.query.level);
+	if (level === null) return;
 
-	console.debug('Resetting progress for all levels', req.session.levelState);
+	console.debug('Resetting progress for all levels');
 	req.session.levelState = getInitialLevelStates();
 	req.session.chatModel = defaultChatModel;
+
+	const chatModel =
+		level === LEVEL_NAMES.SANDBOX ? req.session.chatModel : undefined;
+	const availableDocs =
+		level === LEVEL_NAMES.SANDBOX ? getSandboxDocumentMetas() : undefined;
+
 	res.send({
 		emails: req.session.levelState[level].sentEmails,
 		chatHistory: req.session.levelState[level].chatHistory,
 		defences: req.session.levelState[level].defences,
-		chatModel:
-			level === LEVEL_NAMES.SANDBOX ? req.session.chatModel : undefined,
+		chatModel,
+		availableDocs,
 	});
 }
 
-function handleResetLevel(req: LevelResetRequest, res: Response) {
-	const { level: levelStr } = req.params;
-
-	const level = Number(levelStr) as LEVEL_NAMES;
-
-	if (!isValidLevel(level)) {
-		res.status(400).send('Invalid level');
-		return;
-	}
+function handleResetLevel(
+	req: LevelResetRequest,
+	res: Response<LevelResetResponseBody>
+) {
+	const level = validateLevel(res, Number(req.params.level) as LEVEL_NAMES);
+	if (level === null) return;
 
-	console.debug('Resetting progress for level ', level);
+	console.debug(`Resetting progress for level ${level}`);
 	req.session.levelState[level].chatHistory = [];
 	req.session.levelState[level].sentEmails = [];
 	res.send({
 		chatInfoMessage: {
 			infoMessage: `Level progress reset`,
 			chatMessageType: 'RESET_LEVEL',
-		} as ChatInfoMessage,
+		},
 	});
 }
 

diff --git a/backend/src/controller/startController.ts b/backend/src/controller/startController.ts
@@ -1,32 +1,31 @@
 import { Response } from 'express';
 
+import { getSandboxDocumentMetas } from '@src/document';
 import {
 	StartGetRequest,
-	StartResponse,
+	StartGetResponseBody,
 } from '@src/models/api/StartGetRequest';
-import { LEVEL_NAMES, isValidLevel } from '@src/models/level';
+import { LEVEL_NAMES } from '@src/models/level';
 import { getValidOpenAIModels } from '@src/openai';
 import {
 	systemRoleLevel1,
 	systemRoleLevel2,
 	systemRoleLevel3,
 } from '@src/promptTemplates';
 
-import { sendErrorResponse } from './handleError';
+import { validateLevel } from './requestValidators';
 
-function handleStart(req: StartGetRequest, res: Response) {
-	const { level } = req.query;
-
-	if (level === undefined) {
-		sendErrorResponse(res, 400, 'Level not provided');
-		return;
-	}
-
-	if (!isValidLevel(level)) {
-		sendErrorResponse(res, 400, 'Invalid level');
-		return;
-	}
+function handleStart(
+	req: StartGetRequest,
+	res: Response<StartGetResponseBody>
+) {
+	const level = validateLevel(res, req.query.level);
+	if (level === null) return;
 
+	const chatModel =
+		level === LEVEL_NAMES.SANDBOX ? req.session.chatModel : undefined;
+	const availableDocs =
+		level === LEVEL_NAMES.SANDBOX ? getSandboxDocumentMetas() : undefined;
 	const systemRoles = [
 		{ level: LEVEL_NAMES.LEVEL_1, systemRole: systemRoleLevel1 },
 		{ level: LEVEL_NAMES.LEVEL_2, systemRole: systemRoleLevel2 },
@@ -37,11 +36,11 @@ function handleStart(req: StartGetRequest, res: Response) {
 		emails: req.session.levelState[level].sentEmails,
 		chatHistory: req.session.levelState[level].chatHistory,
 		defences: req.session.levelState[level].defences,
+		chatModel,
+		availableDocs,
 		availableModels: getValidOpenAIModels(),
 		systemRoles,
-		chatModel:
-			level === LEVEL_NAMES.SANDBOX ? req.session.chatModel : undefined,
-	} as StartResponse);
+	});
 }
 
 export { handleStart };
diff --git a/backend/src/document.ts b/backend/src/document.ts
@@ -58,19 +58,16 @@ function getSandboxDocumentMetas() {
 	return [...getDocumentMetas('common'), ...getDocumentMetas('sandbox')];
 }
 
-function getDocumentMetas(folder: string) {
+function getDocumentMetas(folder: string): DocumentMeta[] {
 	const filepath = `resources/documents/${folder}`;
-	const documentMetas: DocumentMeta[] = [];
-
-	fs.readdirSync(filepath).forEach((file) => {
+	return fs.readdirSync(filepath).map((file) => {
 		const fileType = file.split('.').pop() ?? '';
-		documentMetas.push({
-			filename: file,
-			filetype: fileType === 'csv' ? 'text/csv' : fileType,
+		return {
+			fileName: file,
+			fileType: fileType === 'csv' ? 'text/csv' : fileType,
 			folder,
-		});
+		};
 	});
-	return documentMetas;
 }
 
 // store vectorised documents for each level as array

diff --git a/backend/src/langchain.ts b/backend/src/langchain.ts
@@ -30,7 +30,14 @@ function makePromptTemplate(
 }
 
 function getChatModel(): CHAT_MODEL_ID {
-	return getValidOpenAIModels().includes('gpt-4') ? 'gpt-4' : 'gpt-3.5-turbo';
+	const validModels = getValidOpenAIModels();
+	// GPT-4 is the most expensive model by a long way, avoid at all costs!
+	return (
+		validModels.find((model) => model === 'gpt-4o') ??
+		validModels.find((model) => model === 'gpt-4-turbo') ??
+		validModels.find((model) => model === 'gpt-3.5-turbo') ??
+		validModels[0]
+	);
 }
 
 function initQAModel(level: LEVEL_NAMES, Prompt: string) {

diff --git a/backend/src/models/api/LevelGetRequest.ts b/backend/src/models/api/LevelGetRequest.ts
@@ -3,17 +3,21 @@ import { Request } from 'express';
 import { ChatModel } from '@src/models/chat';
 import { ChatMessage } from '@src/models/chatMessage';
 import { Defence } from '@src/models/defence';
+import { DocumentMeta } from '@src/models/document';
 import { EmailInfo } from '@src/models/email';
 import { LEVEL_NAMES } from '@src/models/level';
 
+export type LevelGetResponseBody = {
+	emails: EmailInfo[];
+	chatHistory: ChatMessage[];
+	defences?: Defence[];
+	chatModel?: ChatModel;
+	availableDocs?: DocumentMeta[];
+};
+
 export type LevelGetRequest = Request<
 	never,
-	{
-		emails: EmailInfo[];
-		chatHistory: ChatMessage[];
-		defences: Defence[];
-		chatModel?: ChatModel;
-	},
+	LevelGetResponseBody,
 	never,
 	{
 		level?: LEVEL_NAMES;

diff --git a/backend/src/models/api/LevelResetRequest.ts b/backend/src/models/api/LevelResetRequest.ts
@@ -2,11 +2,13 @@ import { Request } from 'express';
 
 import { ChatMessage } from '@src/models/chatMessage';
 
+export type LevelResetResponseBody = {
+	chatInfoMessage?: ChatMessage;
+};
+
 export type LevelResetRequest = Request<
 	{ level: string },
-	{
-		chatInfoMessage?: ChatMessage;
-	},
+	LevelResetResponseBody,
 	never,
 	never
 >;