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

[Security solution] Fix executor for SimpleChatModel (#186661)

Browse source
This commit is contained in:
Steph Milovic 2024-06-21 14:56:36 -06:00 • committed by GitHub
parent 781db4507d
commit 561c562724
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 11 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

18
x-pack/packages/kbn-langchain/server/language_models/simple_chat_model.ts
View file

@ -98,20 +98,13 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel {
if (!messages.length) {
throw new Error('No messages provided.');
}
const formattedMessages = [];
if (messages.length >= 2) {
messages.forEach((message, i) => {
if (typeof message.content !== 'string') {
throw new Error('Multimodal messages are not supported.');
}
formattedMessages.push(getMessageContentAndRole(message.content, message._getType()));
});
} else {
if (typeof messages[0].content !== 'string') {
const formattedMessages: Array<{ content: string; role: string }> = [];
messages.forEach((message, i) => {
if (typeof message.content !== 'string') {
throw new Error('Multimodal messages are not supported.');
}
formattedMessages.push(getMessageContentAndRole(messages[0].content));
}
formattedMessages.push(getMessageContentAndRole(message.content, message._getType()));
});
this.#logger.debug(
`ActionsClientSimpleChatModel#_call\ntraceId: ${
this.#traceId
@ -129,7 +122,6 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel {
},
},
};
// create an actions client from the authenticated request context:
const actionsClient = await this.#actions.getActionsClientWithRequest(this.#request);

2
x-pack/packages/kbn-langchain/server/utils/bedrock.ts
View file

@ -180,7 +180,7 @@ const prepareBedrockOutput = (responseBody: CompletionChunk, logger?: Logger): s
return responseBody.delta.text;
}
}
logger?.warn(`Failed to parse bedrock chunk ${JSON.stringify(responseBody)}`);
// ignore any chunks that do not include text output
return '';
};

6
x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts
View file

@ -17,6 +17,7 @@ import {
ActionsClientChatOpenAI,
ActionsClientSimpleChatModel,
} from '@kbn/langchain/server';
import { MessagesPlaceholder } from '@langchain/core/prompts';
import { AgentExecutor } from '../executors/types';
import { APMTracer } from '../tracers/apm_tracer';
import { AssistantToolParams } from '../../../types';
@ -126,7 +127,10 @@ export const callAgentExecutor: AgentExecutor<true | false> = async ({
returnIntermediateSteps: false,
agentArgs: {
// this is important to help LangChain correctly format tool input
humanMessageTemplate: `Question: {input}\n\n{agent_scratchpad}`,
humanMessageTemplate: `Remember, when you have enough information, always prefix your final JSON output with "Final Answer:"\n\nQuestion: {input}\n\n{agent_scratchpad}.`,
memoryPrompts: [new MessagesPlaceholder('chat_history')],
suffix:
'Begin! Reminder to ALWAYS use the above format, and to use tools if appropriate.',
},
});