Skip to content

Commit

Permalink
[Security solution] Remove extra spacing from message (#181367)
Browse files Browse the repository at this point in the history
  • Loading branch information
stephmilovic authored Apr 23, 2024
1 parent e45a26e commit bc53c8a
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ describe('use chat send', () => {
expect(sendMessage).toHaveBeenCalled();
const appendMessageSend = sendMessage.mock.calls[0][0].message;
expect(appendMessageSend).toEqual(
`You are a helpful, expert assistant who answers questions about Elastic Security. Do not answer questions unrelated to Elastic Security.\nIf you answer a question related to KQL or EQL, it should be immediately usable within an Elastic Security timeline; please always format the output correctly with back ticks. Any answer provided for Query DSL should also be usable in a security timeline. This means you should only ever include the "filter" portion of the query.\nUse the following context to answer questions:\n\n\n\n${promptText}`
`You are a helpful, expert assistant who answers questions about Elastic Security. Do not answer questions unrelated to Elastic Security.\nIf you answer a question related to KQL or EQL, it should be immediately usable within an Elastic Security timeline; please always format the output correctly with back ticks. Any answer provided for Query DSL should also be usable in a security timeline. This means you should only ever include the "filter" portion of the query.\nUse the following context to answer questions:\n\n${promptText}`
);
});
});
Expand All @@ -124,7 +124,7 @@ describe('use chat send', () => {
await waitFor(() => {
expect(sendMessage).toHaveBeenCalled();
const messages = setCurrentConversation.mock.calls[0][0].messages;
expect(messages[messages.length - 1].content).toEqual(`\n\n${promptText}`);
expect(messages[messages.length - 1].content).toEqual(promptText);
});
});
it('handleRegenerateResponse removes the last message of the conversation, resends the convo to GenAI, and appends the message received', async () => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -94,9 +94,7 @@ User prompt text`);
selectedSystemPrompt: undefined, // <-- no system prompt
});

expect(message.content).toEqual(`
CONTEXT:
expect(message.content).toEqual(`CONTEXT:
"""
alert data
"""
Expand Down Expand Up @@ -148,6 +146,56 @@ User prompt text`);

expect(Date.parse(message.timestamp)).not.toBeNaN();
});
it('should return the correct combined message for a new chat without prompt context', () => {
const result = getCombinedMessage({
currentReplacements: {},
isNewChat: true,
promptText: 'User prompt text',
selectedSystemPrompt: mockSystemPrompt,
selectedPromptContexts: {},
});

expect(result.content).toEqual(
`You are a helpful, expert assistant who answers questions about Elastic Security.\n\nUser prompt text`
);
});

it('should return the correct combined message for a new chat without system context and multiple selectedPromptContext', () => {
const result = getCombinedMessage({
currentReplacements: {},
isNewChat: true,
promptText: 'User prompt text',
selectedPromptContexts: {
context1: {
promptContextId: 'context1',
rawData: 'This is raw data for context 1',
replacements: {},
},
context2: {
promptContextId: 'context2',
rawData: 'This is raw data for context 2',
replacements: {},
},
},
selectedSystemPrompt: { ...mockSystemPrompt, content: '' },
});

expect(result.content).toEqual(
`CONTEXT:\n\"\"\"\nThis is raw data for context 1\n\"\"\"\n,CONTEXT:\n\"\"\"\nThis is raw data for context 2\n\"\"\"\n\nUser prompt text`
);
});

it('should remove extra spaces when there is no prompt content or system prompt', () => {
const result = getCombinedMessage({
currentReplacements: {},
isNewChat: true,
promptText: 'User prompt text',
selectedPromptContexts: {},
selectedSystemPrompt: { ...mockSystemPrompt, content: '' },
});

expect(result.content).toEqual(`User prompt text`);
});

describe('when there is data to anonymize', () => {
const mockPromptContextWithDataToAnonymize: SelectedPromptContext = {
Expand Down Expand Up @@ -210,9 +258,7 @@ User prompt text`);
selectedSystemPrompt: mockSystemPrompt,
});

expect(message.content).toEqual(`
User prompt text`);
expect(message.content).toEqual(`User prompt text`);
});
});
});
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,13 +71,18 @@ export function getCombinedMessage({
rawData: selectedPromptContexts[id].rawData,
});

return `${SYSTEM_PROMPT_CONTEXT_NON_I18N(promptContextData)}`;
return `${SYSTEM_PROMPT_CONTEXT_NON_I18N(promptContextData)}\n`;
});

const content = `${
isNewChat && selectedSystemPrompt && selectedSystemPrompt.content.length > 0
? `${selectedSystemPrompt?.content ?? ''}\n\n`
: ''
}${promptContextsContent.length > 0 ? `${promptContextsContent}\n` : ''}${promptText}`;

return {
content: `${
isNewChat ? `${selectedSystemPrompt?.content ?? ''}\n\n` : ''
}${promptContextsContent}\n\n${promptText}`,
// trim ensures any extra \n and other whitespace is removed
content: content.trim(),
role: 'user', // we are combining the system and user messages into one message
timestamp: new Date().toLocaleString(),
replacements,
Expand Down

0 comments on commit bc53c8a

Please sign in to comment.