Cleanup chat model types in frontend

ScottLogic · Jul 5, 2024 · adf24b3 · adf24b3
1 parent 7a908db
commit adf24b3
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 32 deletions.
diff --git a/frontend/src/components/ControlPanel/ControlPanel.tsx b/frontend/src/components/ControlPanel/ControlPanel.tsx
@@ -3,7 +3,7 @@ import DefenceBox from '@src/components/DefenceBox/DefenceBox';
 import ModelBox from '@src/components/ModelBox/ModelBox';
 import DetailElement from '@src/components/ThemedButtons/DetailElement';
 import ThemedButton from '@src/components/ThemedButtons/ThemedButton';
-import { CHAT_MODEL_ID, ChatMessage, ChatModel } from '@src/models/chat';
+import { ChatMessage, ChatModel } from '@src/models/chat';
 import {
 	DEFENCE_ID,
 	DefenceConfigItem,
@@ -29,7 +29,7 @@ function ControlPanel({
 	currentLevel: LEVEL_NAMES;
 	defences: Defence[];
 	chatModel?: ChatModel;
-	setChatModelId: (modelId: CHAT_MODEL_ID) => void;
+	setChatModelId: (modelId: ChatModel['id']) => void;
 	chatModelOptions: string[];
 	toggleDefence: (defence: Defence) => void;
 	resetDefenceConfiguration: (

diff --git a/frontend/src/components/MainComponent/MainBody.tsx b/frontend/src/components/MainComponent/MainBody.tsx
@@ -1,7 +1,7 @@
 import ChatBox from '@src/components/ChatBox/ChatBox';
 import ControlPanel from '@src/components/ControlPanel/ControlPanel';
 import EmailBox from '@src/components/EmailBox/EmailBox';
-import { CHAT_MODEL_ID, ChatMessage, ChatModel } from '@src/models/chat';
+import { ChatMessage, ChatModel } from '@src/models/chat';
 import {
 	DEFENCE_ID,
 	DefenceConfigItem,
@@ -36,7 +36,7 @@ function MainBody({
 	emails: EmailInfo[];
 	messages: ChatMessage[];
 	chatModel?: ChatModel;
-	setChatModelId: (modelId: CHAT_MODEL_ID) => void;
+	setChatModelId: (modelId: ChatModel['id']) => void;
 	chatModels: string[];
 	addChatMessage: (message: ChatMessage) => void;
 	addSentEmails: (emails: EmailInfo[]) => void;

diff --git a/frontend/src/components/MainComponent/MainComponent.tsx b/frontend/src/components/MainComponent/MainComponent.tsx
@@ -5,7 +5,7 @@ import HandbookOverlay from '@src/components/HandbookOverlay/HandbookOverlay';
 import LevelMissionInfoBanner from '@src/components/LevelMissionInfoBanner/LevelMissionInfoBanner';
 import ResetLevelOverlay from '@src/components/Overlay/ResetLevel';
 import ResetProgressOverlay from '@src/components/Overlay/ResetProgress';
-import { CHAT_MODEL_ID, ChatMessage, ChatModel } from '@src/models/chat';
+import { ChatMessage, ChatModel } from '@src/models/chat';
 import {
 	DEFENCE_ID,
 	DefenceConfigItem,
@@ -55,7 +55,7 @@ function MainComponent({
 	resetCompletedLevels: () => void;
 	setIsNewUser: (isNewUser: boolean) => void;
 }) {
-	const [MainBodyKey, setMainBodyKey] = useState<number>(0);
+	const [mainBodyKey, setMainBodyKey] = useState<number>(0);
 	const [defences, setDefences] = useState<Defence[]>(DEFAULT_DEFENCES);
 	const [messages, setMessages] = useState<ChatMessage[]>([]);
 	const [emails, setEmails] = useState<EmailInfo[]>([]);
@@ -184,9 +184,10 @@ function MainComponent({
 
 		setDefences(defences);
 
-		// we will only update the chatModel if it is defined in the backend response. It will only defined for the sandbox level.
+		// Will only update chatModel if included in the backend response
+		// (currently only for sandbox level)
 		setChatModel(chatModel);
-		setMainBodyKey(MainBodyKey + 1);
+		setMainBodyKey((value) => value + 1);
 	}
 
 	function addInfoMessage(message: string) {
@@ -290,7 +291,7 @@ function MainComponent({
 		}
 	}
 
-	function setChatModelId(modelId: CHAT_MODEL_ID) {
+	function setChatModelId(modelId: ChatModel['id']) {
 		if (!chatModel) {
 			console.error(
 				'You are trying to change the id of the chatModel but it has not been loaded yet'
@@ -355,7 +356,7 @@ function MainComponent({
 				/>
 			)}
 			<MainBody
-				key={MainBodyKey}
+				key={mainBodyKey}
 				currentLevel={currentLevel}
 				defences={defences}
 				chatModel={chatModel}

diff --git a/frontend/src/components/ModelBox/ModelBox.tsx b/frontend/src/components/ModelBox/ModelBox.tsx
@@ -1,4 +1,4 @@
-import { CHAT_MODEL_ID, ChatMessage, ChatModel } from '@src/models/chat';
+import { ChatMessage, ChatModel } from '@src/models/chat';
 
 import ModelConfiguration from './ModelConfiguration';
 import ModelSelection from './ModelSelection';
@@ -12,7 +12,7 @@ function ModelBox({
 	addChatMessage,
 }: {
 	chatModel?: ChatModel;
-	setChatModelId: (modelId: CHAT_MODEL_ID) => void;
+	setChatModelId: (modelId: ChatModel['id']) => void;
 	chatModelOptions: string[];
 	addChatMessage: (chatMessage: ChatMessage) => void;
 }) {

diff --git a/frontend/src/components/ModelBox/ModelSelection.tsx b/frontend/src/components/ModelBox/ModelSelection.tsx
@@ -2,7 +2,7 @@
 import { useEffect, useState } from 'react';
 
 import LoadingButton from '@src/components/ThemedButtons/LoadingButton';
-import { CHAT_MODEL_ID, ChatMessage, ChatModel } from '@src/models/chat';
+import { ChatMessage, ChatModel } from '@src/models/chat';
 import { chatService } from '@src/service';
 
 import './ModelSelection.css';
@@ -15,14 +15,14 @@ function ModelSelection({
 	addChatMessage,
 }: {
 	chatModel?: ChatModel;
-	setChatModelId: (modelId: CHAT_MODEL_ID) => void;
+	setChatModelId: (modelId: ChatModel['id']) => void;
 	chatModelOptions: string[];
 	addChatMessage: (message: ChatMessage) => void;
 }) {
 	// model currently selected in the dropdown
-	const [selectedModel, setSelectedModel] = useState<CHAT_MODEL_ID | undefined>(
-		undefined
-	);
+	const [selectedModel, setSelectedModel] = useState<
+		ChatModel['id'] | undefined
+	>(chatModel?.id);
 
 	const [errorChangingModel, setErrorChangingModel] = useState(false);
 
@@ -65,7 +65,7 @@ function ModelSelection({
 									aria-label="model-select"
 									value={selectedModel ?? 0} // default to the first model
 									onChange={(e) => {
-										setSelectedModel(e.target.value as CHAT_MODEL_ID);
+										setSelectedModel(e.target.value);
 									}}
 								>
 									{chatModelOptions.map((model) => (

diff --git a/frontend/src/models/chat.ts b/frontend/src/models/chat.ts
@@ -15,20 +15,8 @@ type CHAT_MESSAGE_TYPE =
 	| 'ERROR_MSG'
 	| 'RESET_LEVEL';
 
-const chatModelIds = [
-	'gpt-4-1106-preview',
-	'gpt-4',
-	'gpt-4-0613',
-	'gpt-3.5-turbo',
-	'gpt-3.5-turbo-0613',
-	'gpt-3.5-turbo-16k',
-	'gpt-3.5-turbo-16k-0613',
-] as const;
-
-type CHAT_MODEL_ID = (typeof chatModelIds)[number];
-
 type ChatModel = {
-	id: CHAT_MODEL_ID;
+	id: string;
 	configuration: ChatModelConfigurations;
 };
 
@@ -107,5 +95,4 @@ export type {
 	CustomChatModelConfiguration,
 	CHAT_MESSAGE_TYPE,
 	MODEL_CONFIG_ID,
-	CHAT_MODEL_ID,
 };