From bc53c8ae428698aef554c947e976c9d48b41988c Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 23 Apr 2024 10:12:52 -0600 Subject: [PATCH] [Security solution] Remove extra spacing from message (#181367) --- .../chat_send/use_chat_send.test.tsx | 4 +- .../impl/assistant/prompt/helpers.test.ts | 58 +++++++++++++++++-- .../impl/assistant/prompt/helpers.ts | 13 +++-- 3 files changed, 63 insertions(+), 12 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx index 73caae8ed6fc3..6b3b2f8f06efb 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx @@ -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}` ); }); }); @@ -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 () => { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt/helpers.test.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt/helpers.test.ts index 97840d777feac..33b33f83e6581 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt/helpers.test.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt/helpers.test.ts @@ -94,9 +94,7 @@ User prompt text`); selectedSystemPrompt: undefined, // <-- no system prompt }); - expect(message.content).toEqual(` - -CONTEXT: + expect(message.content).toEqual(`CONTEXT: """ alert data """ @@ -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 = { @@ -210,9 +258,7 @@ User prompt text`); selectedSystemPrompt: mockSystemPrompt, }); - expect(message.content).toEqual(` - -User prompt text`); + expect(message.content).toEqual(`User prompt text`); }); }); }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt/helpers.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt/helpers.ts index 5d334f64f0e11..448ab6aa895b1 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt/helpers.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt/helpers.ts @@ -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,