Skip to content

Commit

Permalink
741 transformed messages not showing correctly (#800)
Browse files Browse the repository at this point in the history 
* combines two separate bits of logic for winning level in processChatResponse

* renames increamentNumCompletedLevels to updateNumCompletedLevels

* renames ChatHistoryMessage to ChatMessageDTO

* refactors getChatHistory

* further refactors getChatHistory to remove immutability

* refactors makeChatMessageFromDTO

* removed some outdated comments

* adds reminder comment and transformedMessage as propery to chatHistoryMessage

* sets transformed message inn chat history

* adds ability to retrieve transformed message from the dto

* add the transformed info message in backend rather than frontend

* adds the transformedMessageInfo to the chat response so it can be shown in the frontend

* combine message transformation objects into one object

* tidies random sequence transformation test

* finalise random sequence transformation test

* tidy up xml tagging transformation test

* tidy up xml tagging transformation test with escaping

* removes unnecessary test and reorders

* moves no transformation into transformation test block and removes unused stuff from file

* moves transform message tests into separate test file

* remove isTriggered from defence object

* complete message transformation test

* use undefined instead of null for transofrmed messages

* updates test

* removes isTriggered from test to make it pass

* implements undefined tricks
  • Loading branch information
@pmarsh-scottlogic @chriswilty
pmarsh-scottlogic authored and chriswilty committed Apr 8, 2024
1 parent fc153a9 commit 2b64c6c
Show file tree
Hide file tree
Showing 18 changed files with 328 additions and 171 deletions.
35 changes: 16 additions & 19 deletions backend/src/controller/chatController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ import { Response } from 'express';
import {
transformMessage,
detectTriggeredInputDefences,
combineTransformedMessage,
detectTriggeredOutputDefences,
} from '@src/defence';
import { OpenAiAddHistoryRequest } from '@src/models/api/OpenAiAddHistoryRequest';
Expand All @@ -17,6 +16,7 @@ import {
ChatHttpResponse,
ChatModel,
LevelHandlerResponse,
MessageTransformation,
defaultChatModel,
} from '@src/models/chat';
import { Defence } from '@src/models/defence';
Expand Down Expand Up @@ -46,28 +46,30 @@ function combineChatDefenceReports(

function createNewUserMessages(
message: string,
transformedMessage: string | null
messageTransformation?: MessageTransformation
): ChatHistoryMessage[] {
if (transformedMessage) {
// if message has been transformed
if (messageTransformation) {
return [
// original message
{
completion: null,
chatMessageType: CHAT_MESSAGE_TYPE.USER,
infoMessage: message,
},
// transformed message
{
completion: null,
chatMessageType: CHAT_MESSAGE_TYPE.INFO,
infoMessage: messageTransformation.transformedMessageInfo,
},
{
completion: {
role: 'user',
content: transformedMessage,
content: messageTransformation.transformedMessageCombined,
},
chatMessageType: CHAT_MESSAGE_TYPE.USER_TRANSFORMED,
transformedMessage: messageTransformation.transformedMessage,
},
];
} else {
// not transformed, so just return the original message
return [
{
completion: {
Expand All @@ -88,7 +90,7 @@ async function handleChatWithoutDefenceDetection(
chatHistory: ChatHistoryMessage[],
defences: Defence[]
): Promise<LevelHandlerResponse> {
const updatedChatHistory = createNewUserMessages(message, null).reduce(
const updatedChatHistory = createNewUserMessages(message).reduce(
pushMessageToHistory,
chatHistory
);
Expand Down Expand Up @@ -123,28 +125,22 @@ async function handleChatWithDefenceDetection(
chatHistory: ChatHistoryMessage[],
defences: Defence[]
): Promise<LevelHandlerResponse> {
// transform the message according to active defences
const transformedMessage = transformMessage(message, defences);
const transformedMessageCombined = transformedMessage
? combineTransformedMessage(transformedMessage)
: null;
const messageTransformation = transformMessage(message, defences);
const chatHistoryWithNewUserMessages = createNewUserMessages(
message,
transformedMessageCombined ?? null
messageTransformation
).reduce(pushMessageToHistory, chatHistory);

// detect defences on input message
const triggeredInputDefencesPromise = detectTriggeredInputDefences(
message,
defences
);

// get the chatGPT reply
const openAiReplyPromise = chatGptSendMessage(
chatHistoryWithNewUserMessages,
defences,
chatModel,
transformedMessageCombined ?? message,
messageTransformation?.transformedMessageCombined ?? message,
currentLevel
);

Expand Down Expand Up @@ -178,10 +174,11 @@ async function handleChatWithDefenceDetection(
defenceReport: combinedDefenceReport,
openAIErrorMessage: openAiReply.chatResponse.openAIErrorMessage,
reply: !combinedDefenceReport.isBlocked && botReply ? botReply : '',
transformedMessage: transformedMessage ?? undefined,
transformedMessage: messageTransformation?.transformedMessage,
wonLevel:
openAiReply.chatResponse.wonLevel && !combinedDefenceReport.isBlocked,
sentEmails: combinedDefenceReport.isBlocked ? [] : openAiReply.sentEmails,
transformedMessageInfo: messageTransformation?.transformedMessageInfo,
};
return {
chatResponse: updatedChatResponse,
Expand Down
1 change: 0 additions & 1 deletion backend/src/defaultDefences.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ function createDefence(id: DEFENCE_ID, config: DefenceConfigItem[]): Defence {
id,
config,
isActive: false,
isTriggered: false,
};
}

Expand Down
23 changes: 16 additions & 7 deletions backend/src/defence.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import { defaultDefences } from './defaultDefences';
import { queryPromptEvaluationModel } from './langchain';
import {
ChatDefenceReport,
MessageTransformation,
SingleDefenceReport,
TransformedChatMessage,
} from './models/chat';
Expand Down Expand Up @@ -253,26 +254,34 @@ function combineTransformedMessage(transformedMessage: TransformedChatMessage) {
function transformMessage(
message: string,
defences: Defence[]
): TransformedChatMessage | null {
): MessageTransformation | undefined {
const transformedMessage = isDefenceActive(DEFENCE_ID.XML_TAGGING, defences)
? transformXmlTagging(message, defences)
: isDefenceActive(DEFENCE_ID.RANDOM_SEQUENCE_ENCLOSURE, defences)
? transformRandomSequenceEnclosure(message, defences)
: isDefenceActive(DEFENCE_ID.INSTRUCTION, defences)
? transformInstructionDefence(message, defences)
: null;
: undefined;

if (!transformedMessage) {
console.debug('No defences applied. Message unchanged.');
return null;
return;
}

const transformedMessageCombined =
combineTransformedMessage(transformedMessage);

const transformedMessageInfo =
`${transformedMessage.transformationName} enabled, your message has been transformed`.toLocaleLowerCase();

console.debug(
`Defences applied. Transformed message: ${combineTransformedMessage(
transformedMessage
)}`
`Defences applied. Transformed message: ${transformedMessageCombined}`
);
return transformedMessage;
return {
transformedMessage,
transformedMessageCombined,
transformedMessageInfo,
};
}

// detects triggered defences in original message and blocks the message if necessary
Expand Down
9 changes: 9 additions & 0 deletions backend/src/models/chat.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,12 @@ interface TransformedChatMessage {
transformationName: string;
}

interface MessageTransformation {
transformedMessage: TransformedChatMessage;
transformedMessageInfo: string;
transformedMessageCombined: string;
}

interface ChatHttpResponse {
reply: string;
defenceReport: ChatDefenceReport;
Expand All @@ -112,6 +118,7 @@ interface ChatHttpResponse {
isError: boolean;
openAIErrorMessage: string | null;
sentEmails: EmailInfo[];
transformedMessageInfo?: string;
}

interface LevelHandlerResponse {
Expand All @@ -123,6 +130,7 @@ interface ChatHistoryMessage {
completion: ChatCompletionMessageParam | null;
chatMessageType: CHAT_MESSAGE_TYPE;
infoMessage?: string | null;
transformedMessage?: TransformedChatMessage;
}

// default settings for chat model
Expand All @@ -148,6 +156,7 @@ export type {
TransformedChatMessage,
FunctionCallResponse,
ToolCallResponse,
MessageTransformation,
};
export {
CHAT_MODELS,
Expand Down
1 change: 0 additions & 1 deletion backend/src/models/defence.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ type Defence = {
id: DEFENCE_ID;
config: DefenceConfigItem[];
isActive: boolean;
isTriggered: boolean;
};

export { DEFENCE_ID };
Expand Down
113 changes: 112 additions & 1 deletion backend/test/unit/controller/chatController.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import {
handleClearChatHistory,
handleGetChatHistory,
} from '@src/controller/chatController';
import { detectTriggeredInputDefences } from '@src/defence';
import { detectTriggeredInputDefences, transformMessage } from '@src/defence';
import { OpenAiAddHistoryRequest } from '@src/models/api/OpenAiAddHistoryRequest';
import { OpenAiChatRequest } from '@src/models/api/OpenAiChatRequest';
import { OpenAiClearRequest } from '@src/models/api/OpenAiClearRequest';
Expand All @@ -18,6 +18,7 @@ import {
ChatHistoryMessage,
ChatModel,
ChatResponse,
MessageTransformation,
} from '@src/models/chat';
import { DEFENCE_ID, Defence } from '@src/models/defence';
import { EmailInfo } from '@src/models/email';
Expand Down Expand Up @@ -54,6 +55,9 @@ const mockDetectTriggeredDefences =
detectTriggeredInputDefences as jest.MockedFunction<
typeof detectTriggeredInputDefences
>;
const mockTransformMessage = transformMessage as jest.MockedFunction<
typeof transformMessage
>;

function responseMock() {
return {
Expand Down Expand Up @@ -542,6 +546,113 @@ describe('handleChatToGPT unit tests', () => {
];
expect(history).toEqual(expectedHistory);
});

test('Given sandbox AND message transformation defence active WHEN message sent THEN send reply AND session chat history is updated', async () => {
const transformedMessage = {
preMessage: '[pre message] ',
message: 'hello bot',
postMessage: '[post message]',
transformationName: 'one of the transformation defences',
};
const newTransformationChatHistoryMessages = [
{
completion: null,
chatMessageType: CHAT_MESSAGE_TYPE.USER,
infoMessage: 'hello bot',
},
{
completion: null,
chatMessageType: CHAT_MESSAGE_TYPE.INFO,
infoMessage: 'your message has been transformed by a defence',
},
{
completion: {
role: 'user',
content: '[pre message] hello bot [post message]',
},
chatMessageType: CHAT_MESSAGE_TYPE.USER_TRANSFORMED,
transformedMessage,
},
] as ChatHistoryMessage[];

const newBotChatHistoryMessage = {
chatMessageType: CHAT_MESSAGE_TYPE.BOT,
completion: {
role: 'assistant',
content: 'hello user',
},
} as ChatHistoryMessage;

const req = openAiChatRequestMock(
'hello bot',
LEVEL_NAMES.SANDBOX,
existingHistory
);
const res = responseMock();

mockChatGptSendMessage.mockResolvedValueOnce({
chatResponse: {
completion: { content: 'hello user', role: 'assistant' },
wonLevel: true,
openAIErrorMessage: null,
},
chatHistory: [
...existingHistory,
...newTransformationChatHistoryMessages,
],
sentEmails: [] as EmailInfo[],
});

mockTransformMessage.mockReturnValueOnce({
transformedMessage,
transformedMessageCombined: '[pre message] hello bot [post message]',
transformedMessageInfo:
'your message has been transformed by a defence',
} as MessageTransformation);

mockDetectTriggeredDefences.mockResolvedValueOnce({
blockedReason: null,
isBlocked: false,
alertedDefences: [],
triggeredDefences: [], // do these get updated when the message is transformed?
} as ChatDefenceReport);

await handleChatToGPT(req, res);

expect(mockChatGptSendMessage).toHaveBeenCalledWith(
[...existingHistory, ...newTransformationChatHistoryMessages],
[],
mockChatModel,
'[pre message] hello bot [post message]',
LEVEL_NAMES.SANDBOX
);

expect(res.send).toHaveBeenCalledWith({
reply: 'hello user',
defenceReport: {
blockedReason: '',
isBlocked: false,
alertedDefences: [],
triggeredDefences: [],
},
wonLevel: true,
isError: false,
sentEmails: [],
openAIErrorMessage: null,
transformedMessage,
transformedMessageInfo:
'your message has been transformed by a defence',
});

const history =
req.session.levelState[LEVEL_NAMES.SANDBOX.valueOf()].chatHistory;
const expectedHistory = [
...existingHistory,
...newTransformationChatHistoryMessages,
newBotChatHistoryMessage,
];
expect(history).toEqual(expectedHistory);
});
});
});

Expand Down
Loading

0 comments on commit 2b64c6c

Please sign in to comment.