github-mirrors/kibana
1
0
Fork 0
mirror of https://github.com/elastic/kibana.git synced 2025-04-23 17:28:26 -04:00
Code Issues Projects Releases Packages Wiki Activity

[Security solution] Title prompt improvements (#190087)

Browse source
This commit is contained in:
Steph Milovic 2024-08-12 15:54:54 -06:00 • committed by GitHub
parent 7f8b565d2f
commit b037ea7720
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 38 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

1
x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts
View file

@ -129,6 +129,7 @@ export const getDefaultAssistantGraph = ({
generateChatTitle({
...nodeParams,
model: createLlmInstance(),
llmType,
state,
responseLanguage,
});

46
x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts
View file

@ -10,20 +10,45 @@ import { ChatPromptTemplate } from '@langchain/core/prompts';
import { BaseChatModel } from '@langchain/core/language_models/chat_models';
import { AgentState, NodeParamsBase } from '../types';
export const GENERATE_CHAT_TITLE_PROMPT = (responseLanguage: string) =>
ChatPromptTemplate.fromMessages([
[
'system',
`You are a helpful assistant for Elastic Security. Assume the following user message is the start of a conversation between you and a user; give this conversation a title based on the content below. DO NOT UNDER ANY CIRCUMSTANCES wrap this title in single or double quotes. This title is shown in a list of conversations to the user, so title it for the user, not for you. Please create the title in ${responseLanguage}. As an example, for the given MESSAGE, this is the TITLE:
export const GENERATE_CHAT_TITLE_PROMPT = (responseLanguage: string, llmType?: string) =>
llmType === 'bedrock'
? ChatPromptTemplate.fromMessages([
[
'system',
`You are a helpful assistant for Elastic Security. Assume the following user message is the start of a conversation between you and a user; give this conversation a title based on the content below. DO NOT UNDER ANY CIRCUMSTANCES wrap this title in single or double quotes. This title is shown in a list of conversations to the user, so title it for the user, not for you. Please create the title in ${responseLanguage}. Respond with the title only with no other text explaining your response. As an example, for the given MESSAGE, this is the TITLE:
MESSAGE: I am having trouble with the Elastic Security app.
TITLE: Troubleshooting Elastic Security app issues
`,
],
['human', '{input}'],
]);
],
['human', '{input}'],
])
: llmType === 'gemini'
? ChatPromptTemplate.fromMessages([
[
'system',
`You are a title generator for a helpful assistant for Elastic Security. Assume the following human message is the start of a conversation between you and a human; Do not respond to the human message, instead respond with conversation title relevant to the human's message. DO NOT UNDER ANY CIRCUMSTANCES use quotes or markdown in your response. This title is shown in a list of conversations to the human, so title it for the user, not for you. Please create the title in ${responseLanguage}. Respond with the title only with no other text explaining your response. As an example, for the given MESSAGE, this is the TITLE:
MESSAGE: I am having trouble with the Elastic Security app.
TITLE: Troubleshooting Elastic Security app issues
`,
],
['human', '{input}'],
])
: ChatPromptTemplate.fromMessages([
[
'system',
`You are a helpful assistant for Elastic Security. Assume the following user message is the start of a conversation between you and a user; give this conversation a title based on the content below. DO NOT UNDER ANY CIRCUMSTANCES wrap this title in single or double quotes. This title is shown in a list of conversations to the user, so title it for the user, not for you. Please create the title in ${responseLanguage}. As an example, for the given MESSAGE, this is the TITLE:
MESSAGE: I am having trouble with the Elastic Security app.
TITLE: Troubleshooting Elastic Security app issues
`,
],
['human', '{input}'],
]);
export interface GenerateChatTitleParams extends NodeParamsBase {
llmType?: string;
responseLanguage: string;
state: AgentState;
model: BaseChatModel;
@ -32,6 +57,7 @@ export interface GenerateChatTitleParams extends NodeParamsBase {
export const GENERATE_CHAT_TITLE_NODE = 'generateChatTitle';
export const generateChatTitle = async ({
llmType,
responseLanguage,
logger,
model,
@ -40,7 +66,9 @@ export const generateChatTitle = async ({
logger.debug(() => `Node state:\n ${JSON.stringify(state, null, 2)}`);
const outputParser = new StringOutputParser();
const graph = GENERATE_CHAT_TITLE_PROMPT(responseLanguage).pipe(model).pipe(outputParser);
const graph = GENERATE_CHAT_TITLE_PROMPT(responseLanguage, llmType)
.pipe(model)
.pipe(outputParser);
const chatTitle = await graph.invoke({
input: JSON.stringify(state.input, null, 2),