828 streamline change chat model request (#883)

backend/src/controller/modelController.ts

@@ -2,7 +2,11 @@ import { Response } from 'express';
 
 import { OpenAiConfigureModelRequest } from '@src/models/api/OpenAiConfigureModelRequest';
 import { OpenAiSetModelRequest } from '@src/models/api/OpenAiSetModelRequest';
-import { MODEL_CONFIG_ID, modelConfigIds } from '@src/models/chat';
+import {
+ MODEL_CONFIG_ID,
+ chatModelIds,
+ modelConfigIds,
+} from '@src/models/chat';
 import { ChatInfoMessage } from '@src/models/chatMessage';
 import { LEVEL_NAMES } from '@src/models/level';
 import { pushMessageToHistory } from '@src/utils/chat';
@@ -13,14 +17,32 @@ function handleSetModel(req: OpenAiSetModelRequest, res: Response) {
  const { model } = req.body;
 
  if (model === undefined) {
- res.status(400).send();
- } else {
-  const configuration =
- req.body.configuration ?? req.session.chatModel.configuration;
-  req.session.chatModel = { id: model, configuration };
- console.debug('GPT model set:', JSON.stringify(req.session.chatModel));
- res.status(200).send();
+ sendErrorResponse(res, 400, 'Missing model');
+  return;
+ }
+
+ if (!chatModelIds.includes(model)) {
+ sendErrorResponse(res, 400, 'Invalid model');
+ return;
  }
+
+ const configuration =
+ req.body.configuration ?? req.session.chatModel.configuration;
+ req.session.chatModel = { id: model, configuration };
+ console.debug('GPT model set:', JSON.stringify(req.session.chatModel));
+
+ const chatInfoMessage = {
+ infoMessage: `changed model to ${model}`,
+ chatMessageType: 'GENERIC_INFO',
+ } as ChatInfoMessage;
+ // for now, the chat model only changes in the sandbox level
+ req.session.levelState[LEVEL_NAMES.SANDBOX].chatHistory =
+ pushMessageToHistory(
+ req.session.levelState[LEVEL_NAMES.SANDBOX].chatHistory,
+ chatInfoMessage
+ );
+
+ res.send({ chatInfoMessage });
 }
 
 function handleConfigureModel(req: OpenAiConfigureModelRequest, res: Response) {
@@ -59,13 +81,14 @@ function handleConfigureModel(req: OpenAiConfigureModelRequest, res: Response) {
  infoMessage: `changed ${configId} to ${value}`,
  chatMessageType: 'GENERIC_INFO',
  } as ChatInfoMessage;
+ // for now, the chat model only changes in the sandbox level
  req.session.levelState[LEVEL_NAMES.SANDBOX].chatHistory =
  pushMessageToHistory(
  req.session.levelState[LEVEL_NAMES.SANDBOX].chatHistory,
  chatInfoMessage
  );
 
- res.status(200).send({ chatInfoMessage });
+ res.send({ chatInfoMessage });
 }
 
 export { handleSetModel, handleConfigureModel };

backend/src/langchain.ts

@@ -4,7 +4,7 @@ import { OpenAI } from 'langchain/llms/openai';
 import { PromptTemplate } from 'langchain/prompts';
 
 import { getDocumentVectors } from './document';
-import { CHAT_MODELS } from './models/chat';
+import { CHAT_MODEL_ID } from './models/chat';
 import { PromptEvaluationChainReply, QaChainReply } from './models/langchain';
 import { LEVEL_NAMES } from './models/level';
 import { getOpenAIKey, getValidOpenAIModels } from './openai';
@@ -30,10 +30,8 @@ function makePromptTemplate(
  return PromptTemplate.fromTemplate(fullPrompt);
 }
 
-function getChatModel() {
- return getValidOpenAIModels().includes(CHAT_MODELS.GPT_4)
- ? CHAT_MODELS.GPT_4
- : CHAT_MODELS.GPT_3_5_TURBO;
+function getChatModel(): CHAT_MODEL_ID {
+ return getValidOpenAIModels().includes('gpt-4') ? 'gpt-4' : 'gpt-3.5-turbo';
 }
 
 function initQAModel(level: LEVEL_NAMES, Prompt: string) {

backend/src/models/api/OpenAiSetModelRequest.ts

@@ -1,12 +1,15 @@
 import { Request } from 'express';
 
-import { CHAT_MODELS, ChatModelConfigurations } from '@src/models/chat';
+import { CHAT_MODEL_ID, ChatModelConfigurations } from '@src/models/chat';
+import { ChatMessage } from '@src/models/chatMessage';
 
 export type OpenAiSetModelRequest = Request<
  never,
- never,
  {
- model?: CHAT_MODELS;
+ chatInfoMessage: ChatMessage;
+ },
+ {
+ model?: CHAT_MODEL_ID;
  configuration?: ChatModelConfigurations;
  },
  never

backend/src/models/chat.ts

@@ -7,18 +7,17 @@ import { ChatInfoMessage, ChatMessage } from './chatMessage';
 import { DEFENCE_ID } from './defence';
 import { EmailInfo } from './email';
 
-enum CHAT_MODELS {
- GPT_4_TURBO = 'gpt-4-1106-preview',
- GPT_4 = 'gpt-4',
- GPT_4_0613 = 'gpt-4-0613',
- GPT_3_5_TURBO = 'gpt-3.5-turbo',
- GPT_3_5_TURBO_0613 = 'gpt-3.5-turbo-0613',
- GPT_3_5_TURBO_16K = 'gpt-3.5-turbo-16k',
- GPT_3_5_TURBO_16K_0613 = 'gpt-3.5-turbo-16k-0613',
-}
+const chatModelIds = [
+ 'gpt-4-1106-preview',
+ 'gpt-4',
+ 'gpt-4-0613',
+ 'gpt-3.5-turbo',
+] as const;
+
+type CHAT_MODEL_ID = (typeof chatModelIds)[number];
 
 type ChatModel = {
- id: CHAT_MODELS;
+ id: CHAT_MODEL_ID;
  configuration: ChatModelConfigurations;
 };
 
@@ -106,7 +105,7 @@ interface LevelHandlerResponse {
 }
 
 const defaultChatModel: ChatModel = {
- id: CHAT_MODELS.GPT_3_5_TURBO,
+ id: 'gpt-3.5-turbo',
  configuration: {
  temperature: 1,
  topP: 1,
@@ -131,4 +130,5 @@ export type {
  SingleDefenceReport,
  MODEL_CONFIG_ID,
 };
-export { CHAT_MODELS, defaultChatModel, modelConfigIds };
+export { defaultChatModel, modelConfigIds, chatModelIds };
+export type { CHAT_MODEL_ID };

backend/src/openai.ts

@@ -9,12 +9,13 @@ import { getQAPromptFromConfig } from './defence';
 import { sendEmail } from './email';
 import { queryDocuments } from './langchain';
 import {
- CHAT_MODELS,
+ CHAT_MODEL_ID,
  ChatGptReply,
  ChatModel,
  ChatResponse,
  FunctionCallResponse,
  ToolCallResponse,
+ chatModelIds,
 } from './models/chat';
 import { ChatMessage } from './models/chatMessage';
 import { QaLlmDefence } from './models/defence';
@@ -84,10 +85,10 @@ const chatGptTools: ChatCompletionTool[] = [
 
 // list of valid chat models for the api key
 const validOpenAiModels = (() => {
- let validModels: string[] = [];
+ let validModels: CHAT_MODEL_ID[] = [];
  return {
  get: () => validModels,
- set: (models: string[]) => {
+ set: (models: CHAT_MODEL_ID[]) => {
  validModels = models;
  },
  };
@@ -117,8 +118,8 @@ async function getValidModelsFromOpenAI() {
 
  // get the model ids that are supported by our app. Non-chat models like Dall-e and whisper are not supported.
  const validModels = models.data
- .map((model) => model.id)
- .filter((id) => Object.values(CHAT_MODELS).includes(id as CHAT_MODELS))
+ .map((model) => model.id as CHAT_MODEL_ID)
+ .filter((id) => chatModelIds.includes(id))
  .sort();
 
  validOpenAiModels.set(validModels);
@@ -283,7 +284,7 @@ async function chatGptChatCompletion(
 
 function getChatCompletionsInContextWindow(
  chatHistory: ChatMessage[],
- gptModel: CHAT_MODELS
+ gptModel: CHAT_MODEL_ID
 ): ChatCompletionMessageParam[] {
  const completions = chatHistory
  .map((chatMessage) =>

backend/src/utils/token.ts

@@ -4,18 +4,17 @@ import {
 } from 'openai/resources/chat/completions';
 import { promptTokensEstimate, stringTokens } from 'openai-chat-tokens';
 
-import { CHAT_MODELS } from '@src/models/chat';
+import { CHAT_MODEL_ID } from '@src/models/chat';
 import { chatGptTools } from '@src/openai';
 
 // The size of each model's context window in number of tokens. https://platform.openai.com/docs/models
-const chatModelMaxTokens = {
- [CHAT_MODELS.GPT_4_TURBO]: 128000,
- [CHAT_MODELS.GPT_4]: 8192,
- [CHAT_MODELS.GPT_4_0613]: 8192,
- [CHAT_MODELS.GPT_3_5_TURBO]: 4097,
- [CHAT_MODELS.GPT_3_5_TURBO_0613]: 4097,
- [CHAT_MODELS.GPT_3_5_TURBO_16K]: 16385,
- [CHAT_MODELS.GPT_3_5_TURBO_16K_0613]: 16385,
+const chatModelMaxTokens: {
+ [key in CHAT_MODEL_ID]: number;
+} = {
+ 'gpt-4': 8192,
+ 'gpt-4-1106-preview': 128000,
+ 'gpt-4-0613': 8192,
+ 'gpt-3.5-turbo': 16385,
 };
 
 const TOKENS_PER_TOOL_CALL = 4;

backend/test/integration/openai.test.ts

@@ -1,6 +1,6 @@
 import { expect, jest, test, describe } from '@jest/globals';
 
-import { CHAT_MODELS, ChatModel } from '@src/models/chat';
+import { ChatModel } from '@src/models/chat';
 import { ChatMessage } from '@src/models/chatMessage';
 import { chatGptSendMessage } from '@src/openai';
 
@@ -54,7 +54,7 @@ describe('OpenAI Integration Tests', () => {
  },
  ];
  const chatModel: ChatModel = {
- id: CHAT_MODELS.GPT_4,
+ id: 'gpt-4',
  configuration: {
  temperature: 1,
  topP: 1,

backend/test/unit/controller/modelController.test.ts

@@ -1,8 +1,12 @@
 import { expect, jest, test, describe } from '@jest/globals';
 import { Response } from 'express';
 
-import { handleConfigureModel } from '@src/controller/modelController';
+import {
+ handleConfigureModel,
+ handleSetModel,
+} from '@src/controller/modelController';
 import { OpenAiConfigureModelRequest } from '@src/models/api/OpenAiConfigureModelRequest';
+import { OpenAiSetModelRequest } from '@src/models/api/OpenAiSetModelRequest';
 import { modelConfigIds } from '@src/models/chat';
 import { ChatMessage } from '@src/models/chatMessage';
 import { LEVEL_NAMES, LevelState } from '@src/models/level';
@@ -14,6 +18,64 @@ function responseMock() {
  } as unknown as Response;
 }
 
+describe('handleSetModel', () => {
+ test('WHEN passed sensible parameters THEN sets model AND adds info message to chat history AND responds with info message', () => {
+ const req = {
+ body: {
+ model: 'gpt-4',
+ },
+ session: {
+ chatModel: {
+ id: 'gpt-3.5-turbo',
+ },
+ levelState: [{}, {}, {}, { chatHistory: [] } as unknown as LevelState],
+ },
+ } as OpenAiSetModelRequest;
+ const res = responseMock();
+
+ handleSetModel(req, res);
+
+ expect(req.session.chatModel.id).toBe('gpt-4');
+
+ const expectedInfoMessage = {
+ infoMessage: 'changed model to gpt-4',
+ chatMessageType: 'GENERIC_INFO',
+ } as ChatMessage;
+ expect(
+ req.session.levelState[LEVEL_NAMES.SANDBOX].chatHistory.at(-1)
+ ).toEqual(expectedInfoMessage);
+ expect(res.send).toHaveBeenCalledWith({
+ chatInfoMessage: expectedInfoMessage,
+ });
+ });
+
+ test('WHEN missing model THEN does not set model', () => {
+ const req = {
+ body: {},
+ } as OpenAiSetModelRequest;
+ const res = responseMock();
+
+ handleSetModel(req, res);
+
+ expect(res.status).toHaveBeenCalledWith(400);
+ expect(res.send).toHaveBeenCalledWith('Missing model');
+ });
+
+ test('WHEN model is invalid THEN does not set model', () => {
+ const req = {
+ body: {
+ model: 'invalid model',
+ },
+ } as unknown as OpenAiSetModelRequest;
+ const res = responseMock();
+
+ handleSetModel(req, res);
+
+ expect(res.status).toHaveBeenCalledWith(400);
+ expect(res.send).toHaveBeenCalledWith('Invalid model');
+ });
+});
+
 describe('handleConfigureModel', () => {
  test('WHEN passed sensible parameters THEN configures model AND adds info message to chat history AND responds with info message', () => {
  const req = {
@@ -37,7 +99,6 @@ describe('handleConfigureModel', () => {
 
  handleConfigureModel(req, res);
 
- expect(res.status).toHaveBeenCalledWith(200);
  expect(req.session.chatModel.configuration.topP).toBe(0.5);
 
  const expectedInfoMessage = {

backend/test/unit/openai.test.ts

@@ -42,12 +42,7 @@ describe('getValidModelsFromOpenAI', () => {
  { id: 'da-vinci-1' },
  { id: 'da-vinci-2' },
  ];
- const expectedValidModels = [
- 'gpt-3.5-turbo',
- 'gpt-3.5-turbo-0613',
- 'gpt-4',
- 'gpt-4-0613',
- ];
+ const expectedValidModels = ['gpt-3.5-turbo', 'gpt-4', 'gpt-4-0613'];
 
  mockListFn.mockResolvedValueOnce({
  data: mockModelList,

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 { ChatMessage, ChatModel } from '@src/models/chat';
+import { CHAT_MODEL_ID, ChatMessage, ChatModel } from '@src/models/chat';
 import {
  DEFENCE_ID,
  DefenceConfigItem,
@@ -24,13 +24,12 @@ function ControlPanel({
  resetDefenceConfiguration,
  setDefenceConfiguration,
  openDocumentViewer,
- addInfoMessage,
  addChatMessage,
 }: {
  currentLevel: LEVEL_NAMES;
  defences: Defence[];
  chatModel?: ChatModel;
- setChatModelId: (modelId: string) => void;
+ setChatModelId: (modelId: CHAT_MODEL_ID) => void;
  chatModelOptions: string[];
  toggleDefence: (defence: Defence) => void;
  resetDefenceConfiguration: (
@@ -42,7 +41,6 @@ function ControlPanel({
  config: DefenceConfigItem[]
  ) => Promise<boolean>;
  openDocumentViewer: () => void;
- addInfoMessage: (message: string) => void;
  addChatMessage: (chatMessage: ChatMessage) => void;
 }) {
  const configurableDefences =
@@ -101,7 +99,6 @@ function ControlPanel({
  chatModel={chatModel}
  setChatModelId={setChatModelId}
  chatModelOptions={chatModelOptions}
- addInfoMessage={addInfoMessage}
  addChatMessage={addChatMessage}
  />
  )}