From 8df7b789f540b79893c14dac0ddd9296f5f1fdb8 Mon Sep 17 00:00:00 2001 From: Kibana Machine <42973632+kibanamachine@users.noreply.github.com> Date: Sat, 5 Oct 2024 01:29:13 +1000 Subject: [PATCH] [8.x] [Security Assistant] Vertex chat model (#193032) (#194949) # Backport This will backport the following commits from `main` to `8.x`: - [[Security Assistant] Vertex chat model (#193032)](https://github.com/elastic/kibana/pull/193032) ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport) Co-authored-by: Steph Milovic --- package.json | 7 +- x-pack/packages/kbn-langchain/server/index.ts | 2 + .../chat_vertex/chat_vertex.test.ts | 199 ++++++++++++ .../chat_vertex/chat_vertex.ts | 187 +++++++++++ .../language_models/chat_vertex/connection.ts | 102 ++++++ .../language_models/chat_vertex/index.ts | 8 + .../server/language_models/gemini_chat.ts | 282 +---------------- .../server/language_models/index.ts | 1 + .../kbn-langchain/server/utils/gemini.ts | 297 ++++++++++++++++++ .../nodes/generate_chat_title.ts | 2 +- .../nodes/run_agent.ts | 3 + .../nodes/translations.ts | 11 +- .../graphs/default_assistant_graph/prompts.ts | 17 +- .../elastic_assistant/server/routes/utils.ts | 4 +- .../plugins/elastic_assistant/server/types.ts | 2 + .../esql_language_knowledge_base_tool.ts | 2 +- .../nl_to_esql_tool.ts | 7 +- .../server/connector_types/gemini/gemini.ts | 4 +- yarn.lock | 124 ++++++-- 19 files changed, 941 insertions(+), 320 deletions(-) create mode 100644 x-pack/packages/kbn-langchain/server/language_models/chat_vertex/chat_vertex.test.ts create mode 100644 x-pack/packages/kbn-langchain/server/language_models/chat_vertex/chat_vertex.ts create mode 100644 x-pack/packages/kbn-langchain/server/language_models/chat_vertex/connection.ts create mode 100644 x-pack/packages/kbn-langchain/server/language_models/chat_vertex/index.ts diff --git a/package.json b/package.json index 59101da738873..1b8b8845f98cf 100644 --- a/package.json +++ b/package.json @@ -82,6 +82,7 @@ "**/@bazel/typescript/protobufjs": "6.11.4", "**/@hello-pangea/dnd": "16.6.0", "**/@langchain/core": "^0.2.18", + "**/@langchain/google-common": "^0.1.1", "**/@types/node": "20.10.5", "**/@typescript-eslint/utils": "5.62.0", "**/chokidar": "^3.5.3", @@ -999,7 +1000,9 @@ "@kbn/zod-helpers": "link:packages/kbn-zod-helpers", "@langchain/community": "0.2.18", "@langchain/core": "^0.2.18", - "@langchain/google-genai": "^0.0.23", + "@langchain/google-common": "^0.1.1", + "@langchain/google-genai": "^0.1.0", + "@langchain/google-vertexai": "^0.1.0", "@langchain/langgraph": "0.0.34", "@langchain/openai": "^0.1.3", "@langtrase/trace-attributes": "^3.0.8", @@ -1148,7 +1151,7 @@ "jsts": "^1.6.2", "kea": "^2.6.0", "langchain": "^0.2.11", - "langsmith": "^0.1.39", + "langsmith": "^0.1.55", "launchdarkly-js-client-sdk": "^3.4.0", "load-json-file": "^6.2.0", "lodash": "^4.17.21", diff --git a/x-pack/packages/kbn-langchain/server/index.ts b/x-pack/packages/kbn-langchain/server/index.ts index 7f5691c8f9907..4ffe3aec864d6 100644 --- a/x-pack/packages/kbn-langchain/server/index.ts +++ b/x-pack/packages/kbn-langchain/server/index.ts @@ -10,6 +10,7 @@ import { ActionsClientChatOpenAI } from './language_models/chat_openai'; import { ActionsClientLlm } from './language_models/llm'; import { ActionsClientSimpleChatModel } from './language_models/simple_chat_model'; import { ActionsClientGeminiChatModel } from './language_models/gemini_chat'; +import { ActionsClientChatVertexAI } from './language_models/chat_vertex'; import { parseBedrockStream } from './utils/bedrock'; import { parseGeminiResponse } from './utils/gemini'; import { getDefaultArguments } from './language_models/constants'; @@ -20,6 +21,7 @@ export { getDefaultArguments, ActionsClientBedrockChatModel, ActionsClientChatOpenAI, + ActionsClientChatVertexAI, ActionsClientGeminiChatModel, ActionsClientLlm, ActionsClientSimpleChatModel, diff --git a/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/chat_vertex.test.ts b/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/chat_vertex.test.ts new file mode 100644 index 0000000000000..37506922ff69b --- /dev/null +++ b/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/chat_vertex.test.ts @@ -0,0 +1,199 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { PassThrough } from 'stream'; +import { loggerMock } from '@kbn/logging-mocks'; +import { actionsClientMock } from '@kbn/actions-plugin/server/actions_client/actions_client.mock'; + +import { BaseMessage, HumanMessage, SystemMessage } from '@langchain/core/messages'; +import { ActionsClientChatVertexAI } from './chat_vertex'; +import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager'; + +const connectorId = 'mock-connector-id'; + +const mockExecute = jest.fn(); +const actionsClient = actionsClientMock.create(); + +const mockLogger = loggerMock.create(); + +const mockStreamExecute = jest.fn().mockImplementation(() => { + const passThrough = new PassThrough(); + + // Write the data chunks to the stream + setTimeout(() => { + passThrough.write( + Buffer.from( + `data: {"candidates": [{"content": {"role": "model","parts": [{"text": "token1"}]}}],"modelVersion": "gemini-1.5-pro-001"}` + ) + ); + }); + setTimeout(() => { + passThrough.write( + Buffer.from( + `data: {"candidates": [{"content": {"role": "model","parts": [{"text": "token2"}]}}],"modelVersion": "gemini-1.5-pro-001"}` + ) + ); + }); + setTimeout(() => { + passThrough.write( + Buffer.from( + `data: {"candidates": [{"content": {"role": "model","parts": [{"text": "token3"}]}}],"modelVersion": "gemini-1.5-pro-001"}` + ) + ); + // End the stream + passThrough.end(); + }); + + return { + data: passThrough, // PassThrough stream will act as the async iterator + status: 'ok', + }; +}); + +const callMessages = [ + new SystemMessage('Answer the following questions truthfully and as best you can.'), + new HumanMessage('Question: Do you know my name?\n\n'), +] as unknown as BaseMessage[]; + +const callOptions = { + stop: ['\n'], + recursionLimit: 0, + /** Maximum number of parallel calls to make. */ + maxConcurrency: 0, +}; +const handleLLMNewToken = jest.fn(); +const callRunManager = { + handleLLMNewToken, +} as unknown as CallbackManagerForLLMRun; +const onFailedAttempt = jest.fn(); +const defaultArgs = { + actionsClient, + connectorId, + logger: mockLogger, + streaming: false, + maxRetries: 0, + onFailedAttempt, +}; + +const testMessage = 'Yes, your name is Andrew. How can I assist you further, Andrew?'; + +export const mockActionResponse = { + candidates: [ + { + content: { + role: 'model', + parts: [ + { + text: testMessage, + }, + ], + }, + finishReason: 'STOP', + }, + ], + usageMetadata: { input_tokens: 4, output_tokens: 10, total_tokens: 14 }, +}; + +describe('ActionsClientChatVertexAI', () => { + beforeEach(() => { + jest.clearAllMocks(); + actionsClient.execute.mockImplementation( + jest.fn().mockImplementation(() => ({ + data: mockActionResponse, + status: 'ok', + })) + ); + mockExecute.mockImplementation(() => ({ + data: mockActionResponse, + status: 'ok', + })); + }); + + describe('_generate streaming: false', () => { + it('returns the expected content when _generate is invoked', async () => { + const actionsClientChatVertexAI = new ActionsClientChatVertexAI(defaultArgs); + + const result = await actionsClientChatVertexAI._generate( + callMessages, + callOptions, + callRunManager + ); + const subAction = actionsClient.execute.mock.calls[0][0].params.subAction; + expect(subAction).toEqual('invokeAIRaw'); + + expect(result.generations[0].text).toEqual(testMessage); + }); + + it('rejects with the expected error when the action result status is error', async () => { + const hasErrorStatus = jest.fn().mockImplementation(() => { + throw new Error( + 'ActionsClientChatVertexAI: action result status is error: action-result-message - action-result-service-message' + ); + }); + + actionsClient.execute.mockRejectedValueOnce(hasErrorStatus); + + const actionsClientChatVertexAI = new ActionsClientChatVertexAI({ + ...defaultArgs, + actionsClient, + }); + + await expect( + actionsClientChatVertexAI._generate(callMessages, callOptions, callRunManager) + ).rejects.toThrowError(); + expect(onFailedAttempt).toHaveBeenCalled(); + }); + + it('rejects with the expected error the message has invalid content', async () => { + actionsClient.execute.mockImplementation( + jest.fn().mockResolvedValue({ + data: { + Bad: true, + finishReason: 'badness', + }, + status: 'ok', + }) + ); + + const actionsClientChatVertexAI = new ActionsClientChatVertexAI(defaultArgs); + + await expect( + actionsClientChatVertexAI._generate(callMessages, callOptions, callRunManager) + ).rejects.toThrowError("Cannot read properties of undefined (reading 'text')"); + }); + }); + + describe('*_streamResponseChunks', () => { + it('iterates over gemini chunks', async () => { + actionsClient.execute.mockImplementationOnce(mockStreamExecute); + + const actionsClientChatVertexAI = new ActionsClientChatVertexAI({ + ...defaultArgs, + actionsClient, + streaming: true, + }); + + const gen = actionsClientChatVertexAI._streamResponseChunks( + callMessages, + callOptions, + callRunManager + ); + + const chunks = []; + + for await (const chunk of gen) { + chunks.push(chunk); + } + + expect(chunks.map((c) => c.text)).toEqual(['token1', 'token2', 'token3']); + expect(handleLLMNewToken).toHaveBeenCalledTimes(3); + expect(handleLLMNewToken).toHaveBeenCalledWith('token1'); + expect(handleLLMNewToken).toHaveBeenCalledWith('token2'); + expect(handleLLMNewToken).toHaveBeenCalledWith('token3'); + }); + }); +}); diff --git a/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/chat_vertex.ts b/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/chat_vertex.ts new file mode 100644 index 0000000000000..5627abe717291 --- /dev/null +++ b/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/chat_vertex.ts @@ -0,0 +1,187 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { EnhancedGenerateContentResponse } from '@google/generative-ai'; +import { ActionsClient } from '@kbn/actions-plugin/server'; +import { PublicMethodsOf } from '@kbn/utility-types'; +import { BaseMessage, UsageMetadata } from '@langchain/core/messages'; +import { ChatGenerationChunk } from '@langchain/core/outputs'; +import { ChatVertexAI } from '@langchain/google-vertexai'; +import { get } from 'lodash/fp'; +import { Readable } from 'stream'; + +import { Logger } from '@kbn/logging'; +import { BaseChatModelParams } from '@langchain/core/language_models/chat_models'; +import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager'; +import { GeminiPartText } from '@langchain/google-common/dist/types'; +import { + convertResponseBadFinishReasonToErrorMsg, + convertResponseContentToChatGenerationChunk, +} from '../../utils/gemini'; +import { ActionsClientChatConnection } from './connection'; + +const DEFAULT_GEMINI_TEMPERATURE = 0; +export interface CustomChatModelInput extends BaseChatModelParams { + actionsClient: PublicMethodsOf; + connectorId: string; + logger: Logger; + streaming: boolean; + temperature?: number; + signal?: AbortSignal; + model?: string; + maxTokens?: number; +} + +export class ActionsClientChatVertexAI extends ChatVertexAI { + #actionsClient: PublicMethodsOf; + #connectorId: string; + #model?: string; + constructor({ actionsClient, connectorId, ...props }: CustomChatModelInput) { + super({ + ...props, + maxOutputTokens: props.maxTokens ?? 2048, + temperature: props.temperature ?? DEFAULT_GEMINI_TEMPERATURE, + }); + // LangChain needs model to be defined for logging purposes + this.model = props.model ?? this.model; + // If model is not specified by consumer, the connector will define it so do not pass + // a LangChain default to the actionsClient + this.#model = props.model; + this.#actionsClient = actionsClient; + this.#connectorId = connectorId; + const client = this.buildClient(props); + this.connection = new ActionsClientChatConnection( + { + ...this, + }, + this.caller, + client, + false, + actionsClient, + connectorId + ); + } + + buildConnection() { + // prevent ChatVertexAI from overwriting our this.connection defined in super + } + + async *_streamResponseChunks( + messages: BaseMessage[], + options: this['ParsedCallOptions'], + runManager?: CallbackManagerForLLMRun + ): AsyncGenerator { + const parameters = this.invocationParams(options); + const data = await this.connection.formatData(messages, parameters); + const stream = await this.caller.callWithOptions({ signal: options?.signal }, async () => { + const systemPart: GeminiPartText | undefined = data?.systemInstruction + ?.parts?.[0] as unknown as GeminiPartText; + const systemInstruction = systemPart?.text.length + ? { systemInstruction: systemPart?.text } + : {}; + const requestBody = { + actionId: this.#connectorId, + params: { + subAction: 'invokeStream', + subActionParams: { + model: this.#model, + messages: data?.contents, + tools: data?.tools, + temperature: this.temperature, + ...systemInstruction, + }, + }, + }; + + const actionResult = await this.#actionsClient.execute(requestBody); + + if (actionResult.status === 'error') { + throw new Error( + `ActionsClientChatVertexAI: action result status is error: ${actionResult?.message} - ${actionResult?.serviceMessage}` + ); + } + + const readable = get('data', actionResult) as Readable; + + if (typeof readable?.read !== 'function') { + throw new Error('Action result status is error: result is not streamable'); + } + return readable; + }); + let usageMetadata: UsageMetadata | undefined; + let index = 0; + let partialStreamChunk = ''; + for await (const rawStreamChunk of stream) { + const streamChunk = rawStreamChunk.toString(); + const nextChunk = `${partialStreamChunk + streamChunk}`; + + let parsedStreamChunk: EnhancedGenerateContentResponse | null = null; + try { + parsedStreamChunk = JSON.parse(nextChunk.replaceAll('data: ', '').replaceAll('\r\n', '')); + partialStreamChunk = ''; + } catch (_) { + partialStreamChunk += nextChunk; + } + + if (parsedStreamChunk !== null && !parsedStreamChunk.candidates?.[0]?.finishReason) { + const response = { + ...parsedStreamChunk, + functionCalls: () => + parsedStreamChunk?.candidates?.[0]?.content.parts[0].functionCall + ? [parsedStreamChunk.candidates?.[0]?.content.parts[0].functionCall] + : [], + }; + + if ( + 'usageMetadata' in response && + this.streamUsage !== false && + options.streamUsage !== false + ) { + const genAIUsageMetadata = response.usageMetadata as { + promptTokenCount: number; + candidatesTokenCount: number; + totalTokenCount: number; + }; + if (!usageMetadata) { + usageMetadata = { + input_tokens: genAIUsageMetadata.promptTokenCount, + output_tokens: genAIUsageMetadata.candidatesTokenCount, + total_tokens: genAIUsageMetadata.totalTokenCount, + }; + } else { + // Under the hood, LangChain combines the prompt tokens. Google returns the updated + // total each time, so we need to find the difference between the tokens. + const outputTokenDiff = + genAIUsageMetadata.candidatesTokenCount - usageMetadata.output_tokens; + usageMetadata = { + input_tokens: 0, + output_tokens: outputTokenDiff, + total_tokens: outputTokenDiff, + }; + } + } + + const chunk = convertResponseContentToChatGenerationChunk(response, { + usageMetadata, + index, + }); + index += 1; + + if (chunk) { + yield chunk; + await runManager?.handleLLMNewToken(chunk.text ?? ''); + } + } else if (parsedStreamChunk) { + // handle bad finish reason + const errorMessage = convertResponseBadFinishReasonToErrorMsg(parsedStreamChunk); + if (errorMessage != null) { + throw new Error(errorMessage); + } + } + } + } +} diff --git a/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/connection.ts b/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/connection.ts new file mode 100644 index 0000000000000..0340d71b438db --- /dev/null +++ b/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/connection.ts @@ -0,0 +1,102 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { + ChatConnection, + GoogleAbstractedClient, + GoogleAIBaseLLMInput, + GoogleLLMResponse, +} from '@langchain/google-common'; +import { ActionsClient } from '@kbn/actions-plugin/server'; +import { PublicMethodsOf } from '@kbn/utility-types'; +import { EnhancedGenerateContentResponse } from '@google/generative-ai'; +import { AsyncCaller } from '@langchain/core/utils/async_caller'; +import { convertResponseBadFinishReasonToErrorMsg } from '../../utils/gemini'; + +// only implements non-streaming requests +// stream is handled by ActionsClientChatVertexAI.*_streamResponseChunks +export class ActionsClientChatConnection extends ChatConnection { + actionsClient: PublicMethodsOf; + connectorId: string; + #model?: string; + temperature: number; + caller: AsyncCaller; + constructor( + fields: GoogleAIBaseLLMInput, + caller: AsyncCaller, + client: GoogleAbstractedClient, + _streaming: boolean, // defaulting to false in the super + actionsClient: PublicMethodsOf, + connectorId: string + ) { + super(fields, caller, client, false); + this.actionsClient = actionsClient; + this.connectorId = connectorId; + this.caller = caller; + this.#model = fields.model; + this.temperature = fields.temperature ?? 0; + } + + async _request( + // TODO better types here + data: { + contents: unknown; + tools: unknown[]; + systemInstruction?: { parts: [{ text: string }] }; + }, + options: { signal?: AbortSignal } + ) { + const systemInstruction = data?.systemInstruction?.parts?.[0]?.text.length + ? { systemInstruction: data?.systemInstruction?.parts?.[0]?.text } + : {}; + return this.caller.callWithOptions({ signal: options?.signal }, async () => { + try { + const requestBody = { + actionId: this.connectorId, + params: { + subAction: 'invokeAIRaw', + subActionParams: { + model: this.#model, + messages: data?.contents, + tools: data?.tools, + temperature: this.temperature, + ...systemInstruction, + }, + }, + }; + const actionResult = (await this.actionsClient.execute(requestBody)) as { + status: string; + data: EnhancedGenerateContentResponse; + message?: string; + serviceMessage?: string; + }; + + if (actionResult.status === 'error') { + throw new Error( + `ActionsClientChatVertexAI: action result status is error: ${actionResult?.message} - ${actionResult?.serviceMessage}` + ); + } + + if (actionResult.data.candidates && actionResult.data.candidates.length > 0) { + // handle bad finish reason + const errorMessage = convertResponseBadFinishReasonToErrorMsg(actionResult.data); + if (errorMessage != null) { + throw new Error(errorMessage); + } + } + return actionResult as unknown as GoogleLLMResponse; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } catch (e: any) { + // TODO: Improve error handling + if (e.message?.includes('400 Bad Request')) { + e.status = 400; + } + throw e; + } + }); + } +} diff --git a/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/index.ts b/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/index.ts new file mode 100644 index 0000000000000..54197f7e2cb08 --- /dev/null +++ b/x-pack/packages/kbn-langchain/server/language_models/chat_vertex/index.ts @@ -0,0 +1,8 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +export * from './chat_vertex'; diff --git a/x-pack/packages/kbn-langchain/server/language_models/gemini_chat.ts b/x-pack/packages/kbn-langchain/server/language_models/gemini_chat.ts index fb98232359340..197360c2f06e6 100644 --- a/x-pack/packages/kbn-langchain/server/language_models/gemini_chat.ts +++ b/x-pack/packages/kbn-langchain/server/language_models/gemini_chat.ts @@ -8,34 +8,24 @@ import { Content, EnhancedGenerateContentResponse, - FunctionCallPart, - FunctionResponsePart, GenerateContentRequest, GenerateContentResult, - InlineDataPart, - POSSIBLE_ROLES, - Part, - TextPart, - FinishReason, - SafetyRating, } from '@google/generative-ai'; import { ActionsClient } from '@kbn/actions-plugin/server'; import { PublicMethodsOf } from '@kbn/utility-types'; import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager'; -import { ToolCallChunk } from '@langchain/core/dist/messages/tool'; -import { - AIMessageChunk, - BaseMessage, - ChatMessage, - isBaseMessage, - UsageMetadata, -} from '@langchain/core/messages'; +import { BaseMessage, UsageMetadata } from '@langchain/core/messages'; import { ChatGenerationChunk } from '@langchain/core/outputs'; import { ChatGoogleGenerativeAI } from '@langchain/google-genai'; import { Logger } from '@kbn/logging'; import { BaseChatModelParams } from '@langchain/core/language_models/chat_models'; import { get } from 'lodash/fp'; import { Readable } from 'stream'; +import { + convertBaseMessagesToContent, + convertResponseBadFinishReasonToErrorMsg, + convertResponseContentToChatGenerationChunk, +} from '../utils/gemini'; const DEFAULT_GEMINI_TEMPERATURE = 0; export interface CustomChatModelInput extends BaseChatModelParams { @@ -48,12 +38,6 @@ export interface CustomChatModelInput extends BaseChatModelParams { maxTokens?: number; } -// not sure why these properties are not on the type, as they are on the data -interface SafetyReason extends SafetyRating { - blocked: boolean; - severity: string; -} - export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI { #actionsClient: PublicMethodsOf; #connectorId: string; @@ -265,257 +249,3 @@ export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI { } } } - -export function convertResponseContentToChatGenerationChunk( - response: EnhancedGenerateContentResponse, - extra: { - usageMetadata?: UsageMetadata | undefined; - index: number; - } -): ChatGenerationChunk | null { - if (!response.candidates || response.candidates.length === 0) { - return null; - } - const functionCalls = response.functionCalls(); - const [candidate] = response.candidates; - const { content, ...generationInfo } = candidate; - const text = content?.parts[0]?.text ?? ''; - - const toolCallChunks: ToolCallChunk[] = []; - if (functionCalls) { - toolCallChunks.push( - ...functionCalls.map((fc) => ({ - ...fc, - args: JSON.stringify(fc.args), - index: extra.index, - type: 'tool_call_chunk' as const, - })) - ); - } - return new ChatGenerationChunk({ - text, - message: new AIMessageChunk({ - content: text, - name: !content ? undefined : content.role, - tool_call_chunks: toolCallChunks, - // Each chunk can have unique "generationInfo", and merging strategy is unclear, - // so leave blank for now. - additional_kwargs: {}, - usage_metadata: extra.usageMetadata, - }), - generationInfo, - }); -} - -export function convertAuthorToRole(author: string): (typeof POSSIBLE_ROLES)[number] { - switch (author) { - /** - * Note: Gemini currently is not supporting system messages - * we will convert them to human messages and merge with following - * */ - case 'ai': - case 'model': // getMessageAuthor returns message.name. code ex.: return message.name ?? type; - return 'model'; - case 'system': - case 'human': - return 'user'; - case 'tool': - case 'function': - return 'function'; - default: - throw new Error(`Unknown / unsupported author: ${author}`); - } -} -export function convertBaseMessagesToContent(messages: BaseMessage[], isMultimodalModel: boolean) { - return messages.reduce<{ - content: Content[]; - mergeWithPreviousContent: boolean; - }>( - (acc, message, index) => { - if (!isBaseMessage(message)) { - throw new Error('Unsupported message input'); - } - const author = getMessageAuthor(message); - if (author === 'system' && index !== 0) { - throw new Error('System message should be the first one'); - } - const role = convertAuthorToRole(author); - const parts = convertMessageContentToParts(message, isMultimodalModel); - - if (acc.mergeWithPreviousContent) { - const prevContent = acc.content[acc.content.length - 1]; - if (!prevContent) { - throw new Error( - 'There was a problem parsing your system message. Please try a prompt without one.' - ); - } - prevContent.parts.push(...parts); - - return { - mergeWithPreviousContent: false, - content: acc.content, - }; - } - let actualRole = role; - if (actualRole === 'function') { - // GenerativeAI API will throw an error if the role is not "user" or "model." - actualRole = 'user'; - } - const content: Content = { - role: actualRole, - parts, - }; - return { - mergeWithPreviousContent: author === 'system', - content: [...acc.content, content], - }; - }, - { content: [], mergeWithPreviousContent: false } - ).content; -} - -export function convertMessageContentToParts( - message: BaseMessage, - isMultimodalModel: boolean -): Part[] { - if (typeof message.content === 'string' && message.content !== '') { - return [{ text: message.content }]; - } - - let functionCalls: FunctionCallPart[] = []; - let functionResponses: FunctionResponsePart[] = []; - let messageParts: Part[] = []; - - if ( - 'tool_calls' in message && - Array.isArray(message.tool_calls) && - message.tool_calls.length > 0 - ) { - functionCalls = message.tool_calls.map((tc) => ({ - functionCall: { - name: tc.name, - args: tc.args, - }, - })); - } else if (message._getType() === 'tool' && message.name && message.content) { - functionResponses = [ - { - functionResponse: { - name: message.name, - response: message.content, - }, - }, - ]; - } else if (Array.isArray(message.content)) { - messageParts = message.content.map((c) => { - if (c.type === 'text') { - return { - text: c.text, - } as TextPart; - } - - if (c.type === 'image_url') { - if (!isMultimodalModel) { - throw new Error(`This model does not support images`); - } - let source; - if (typeof c.image_url === 'string') { - source = c.image_url; - } else if (typeof c.image_url === 'object' && 'url' in c.image_url) { - source = c.image_url.url; - } else { - throw new Error('Please provide image as base64 encoded data URL'); - } - const [dm, data] = source.split(','); - if (!dm.startsWith('data:')) { - throw new Error('Please provide image as base64 encoded data URL'); - } - - const [mimeType, encoding] = dm.replace(/^data:/, '').split(';'); - if (encoding !== 'base64') { - throw new Error('Please provide image as base64 encoded data URL'); - } - - return { - inlineData: { - data, - mimeType, - }, - } as InlineDataPart; - } else if (c.type === 'media') { - return messageContentMedia(c); - } else if (c.type === 'tool_use') { - return { - functionCall: { - name: c.name, - args: c.input, - }, - } as FunctionCallPart; - } - throw new Error(`Unknown content type ${(c as { type: string }).type}`); - }); - } - - return [...messageParts, ...functionCalls, ...functionResponses]; -} - -export function getMessageAuthor(message: BaseMessage) { - const type = message._getType(); - if (ChatMessage.isInstance(message)) { - return message.role; - } - if (type === 'tool') { - return type; - } - return message.name ?? type; -} - -// will be removed once FileDataPart is supported in @langchain/google-genai -function messageContentMedia(content: Record): InlineDataPart { - if ('mimeType' in content && 'data' in content) { - return { - inlineData: { - mimeType: content.mimeType, - data: content.data, - }, - } as InlineDataPart; - } - throw new Error('Invalid media content'); -} - -const badFinishReasons = [FinishReason.RECITATION, FinishReason.SAFETY]; -function hadBadFinishReason(candidate: { finishReason?: FinishReason }) { - return !!candidate.finishReason && badFinishReasons.includes(candidate.finishReason); -} - -export function convertResponseBadFinishReasonToErrorMsg( - response: EnhancedGenerateContentResponse -): string | null { - if (response.candidates && response.candidates.length > 0) { - const candidate = response.candidates[0]; - if (hadBadFinishReason(candidate)) { - if ( - candidate.finishReason === FinishReason.SAFETY && - candidate.safetyRatings && - (candidate.safetyRatings?.length ?? 0) > 0 - ) { - const safetyReasons = getSafetyReasons(candidate.safetyRatings as SafetyReason[]); - return `ActionsClientGeminiChatModel: action result status is error. Candidate was blocked due to ${candidate.finishReason} - ${safetyReasons}`; - } else { - return `ActionsClientGeminiChatModel: action result status is error. Candidate was blocked due to ${candidate.finishReason}`; - } - } - } - return null; -} - -const getSafetyReasons = (safetyRatings: SafetyReason[]) => { - const reasons = safetyRatings.filter((t: SafetyReason) => t.blocked); - return reasons.reduce( - (acc: string, t: SafetyReason, i: number) => - `${acc.length ? `${acc} ` : ''}${t.category}: ${t.severity}${ - i < reasons.length - 1 ? ',' : '' - }`, - '' - ); -}; diff --git a/x-pack/packages/kbn-langchain/server/language_models/index.ts b/x-pack/packages/kbn-langchain/server/language_models/index.ts index f5415079cbc11..b3c5053f11701 100644 --- a/x-pack/packages/kbn-langchain/server/language_models/index.ts +++ b/x-pack/packages/kbn-langchain/server/language_models/index.ts @@ -8,5 +8,6 @@ export { ActionsClientBedrockChatModel } from './bedrock_chat'; export { ActionsClientChatOpenAI } from './chat_openai'; export { ActionsClientGeminiChatModel } from './gemini_chat'; +export { ActionsClientChatVertexAI } from './chat_vertex'; export { ActionsClientLlm } from './llm'; export { ActionsClientSimpleChatModel } from './simple_chat_model'; diff --git a/x-pack/packages/kbn-langchain/server/utils/gemini.ts b/x-pack/packages/kbn-langchain/server/utils/gemini.ts index 7d08923083831..9cc975d8c1f2b 100644 --- a/x-pack/packages/kbn-langchain/server/utils/gemini.ts +++ b/x-pack/packages/kbn-langchain/server/utils/gemini.ts @@ -6,9 +6,306 @@ */ import { Logger } from '@kbn/core/server'; +import { + Content, + EnhancedGenerateContentResponse, + FinishReason, + FunctionCallPart, + FunctionResponsePart, + InlineDataPart, + Part, + POSSIBLE_ROLES, + SafetyRating, + TextPart, +} from '@google/generative-ai'; +import { + AIMessageChunk, + BaseMessage, + ChatMessage, + isBaseMessage, + UsageMetadata, +} from '@langchain/core/messages'; +import { ChatGenerationChunk } from '@langchain/core/outputs'; +import { ToolCallChunk } from '@langchain/core/dist/messages/tool'; import { Readable } from 'stream'; import { StreamParser } from './types'; +export function convertResponseContentToChatGenerationChunk( + response: EnhancedGenerateContentResponse, + extra: { + usageMetadata?: UsageMetadata | undefined; + index: number; + } +): ChatGenerationChunk | null { + if (!response.candidates || response.candidates.length === 0) { + return null; + } + const functionCalls = response.functionCalls(); + const [candidate] = response.candidates; + const { content, ...generationInfo } = candidate; + const text = content?.parts[0]?.text ?? ''; + + const toolCallChunks: ToolCallChunk[] = []; + if (functionCalls) { + toolCallChunks.push( + ...functionCalls.map((fc) => ({ + ...fc, + args: JSON.stringify(fc.args), + index: extra.index, + type: 'tool_call_chunk' as const, + })) + ); + } + return new ChatGenerationChunk({ + text, + message: new AIMessageChunk({ + content: text, + name: !content ? undefined : content.role, + tool_call_chunks: toolCallChunks, + // Each chunk can have unique "generationInfo", and merging strategy is unclear, + // so leave blank for now. + additional_kwargs: {}, + usage_metadata: extra.usageMetadata, + }), + generationInfo, + }); +} + +export function convertAuthorToRole(author: string): (typeof POSSIBLE_ROLES)[number] { + switch (author) { + /** + * Note: Gemini currently is not supporting system messages + * we will convert them to human messages and merge with following + * */ + case 'ai': + case 'model': // getMessageAuthor returns message.name. code ex.: return message.name ?? type; + return 'model'; + case 'system': + case 'human': + return 'user'; + case 'tool': + case 'function': + return 'function'; + default: + throw new Error(`Unknown / unsupported author: ${author}`); + } +} +export function convertBaseMessagesToContent(messages: BaseMessage[], isMultimodalModel: boolean) { + return messages.reduce<{ + content: Content[]; + mergeWithPreviousContent: boolean; + }>( + (acc, message, index) => { + if (!isBaseMessage(message)) { + throw new Error('Unsupported message input'); + } + const author = getMessageAuthor(message); + if (author === 'system' && index !== 0) { + throw new Error('System message should be the first one'); + } + const role = convertAuthorToRole(author); + const parts = convertMessageContentToParts(message, isMultimodalModel); + + if (acc.mergeWithPreviousContent) { + const prevContent = acc.content[acc.content.length - 1]; + if (!prevContent) { + throw new Error( + 'There was a problem parsing your system message. Please try a prompt without one.' + ); + } + prevContent.parts.push(...parts); + + return { + mergeWithPreviousContent: false, + content: acc.content, + }; + } + let actualRole = role; + if (actualRole === 'function') { + // GenerativeAI API will throw an error if the role is not "user" or "model." + actualRole = 'user'; + } + const content: Content = { + role: actualRole, + parts, + }; + return { + mergeWithPreviousContent: author === 'system', + content: [...acc.content, content], + }; + }, + { content: [], mergeWithPreviousContent: false } + ).content; +} + +export function convertMessageContentToParts( + message: BaseMessage, + isMultimodalModel: boolean +): Part[] { + if (typeof message.content === 'string' && message.content !== '') { + return [{ text: message.content }]; + } + + let functionCalls: FunctionCallPart[] = []; + let functionResponses: FunctionResponsePart[] = []; + let messageParts: Part[] = []; + + if ( + 'tool_calls' in message && + Array.isArray(message.tool_calls) && + message.tool_calls.length > 0 + ) { + functionCalls = message.tool_calls.map((tc) => ({ + functionCall: { + name: tc.name, + args: tc.args, + }, + })); + } else if (message._getType() === 'tool' && message.name && message.content) { + functionResponses = [ + { + functionResponse: { + name: message.name, + response: message.content, + }, + }, + ]; + } else if (Array.isArray(message.content)) { + messageParts = message.content.map((c) => { + if (c.type === 'text') { + return { + text: c.text, + } as TextPart; + } + + if (c.type === 'image_url') { + if (!isMultimodalModel) { + throw new Error(`This model does not support images`); + } + let source; + if (typeof c.image_url === 'string') { + source = c.image_url; + } else if (typeof c.image_url === 'object' && 'url' in c.image_url) { + source = c.image_url.url; + } else { + throw new Error('Please provide image as base64 encoded data URL'); + } + const [dm, data] = source.split(','); + if (!dm.startsWith('data:')) { + throw new Error('Please provide image as base64 encoded data URL'); + } + + const [mimeType, encoding] = dm.replace(/^data:/, '').split(';'); + if (encoding !== 'base64') { + throw new Error('Please provide image as base64 encoded data URL'); + } + + return { + inlineData: { + data, + mimeType, + }, + } as InlineDataPart; + } else if (c.type === 'media') { + return messageContentMedia(c); + } else if (c.type === 'tool_use') { + return { + functionCall: { + name: c.name, + args: c.input, + }, + } as FunctionCallPart; + } + throw new Error(`Unknown content type ${(c as { type: string }).type}`); + }); + } + + return [...messageParts, ...functionCalls, ...functionResponses]; +} + +export function getMessageAuthor(message: BaseMessage) { + const type = message._getType(); + if (ChatMessage.isInstance(message)) { + return message.role; + } + if (type === 'tool') { + return type; + } + return message.name ?? type; +} + +// will be removed once FileDataPart is supported in @langchain/google-genai +function messageContentMedia(content: Record): InlineDataPart { + if ('mimeType' in content && 'data' in content) { + return { + inlineData: { + mimeType: content.mimeType, + data: content.data, + }, + } as InlineDataPart; + } + throw new Error('Invalid media content'); +} + +// TODO Google's TS library is behind the API +// remove this enum once the library is updated +// https://github.com/google-gemini/generative-ai-js/pull/270 +enum FinishReasonMore { + BLOCKLIST = 'BLOCKLIST', + PROHIBITED_CONTENT = 'PROHIBITED_CONTENT', + SPII = 'SPII', + MALFORMED_FUNCTION_CALL = 'MALFORMED_FUNCTION_CALL', +} +const badFinishReasons = [ + FinishReason.RECITATION, + FinishReason.SAFETY, + FinishReasonMore.BLOCKLIST, + FinishReasonMore.PROHIBITED_CONTENT, + FinishReasonMore.SPII, + FinishReasonMore.MALFORMED_FUNCTION_CALL, +]; +function hadBadFinishReason(candidate: { finishReason?: FinishReason }) { + return !!candidate.finishReason && badFinishReasons.includes(candidate.finishReason); +} + +export function convertResponseBadFinishReasonToErrorMsg( + response: EnhancedGenerateContentResponse +): string | null { + if (response.candidates && response.candidates.length > 0) { + const candidate = response.candidates[0]; + if (hadBadFinishReason(candidate)) { + if ( + candidate.finishReason === FinishReason.SAFETY && + candidate.safetyRatings && + (candidate.safetyRatings?.length ?? 0) > 0 + ) { + const safetyReasons = getSafetyReasons(candidate.safetyRatings as SafetyReason[]); + return `Gemini Utils: action result status is error. Candidate was blocked due to ${candidate.finishReason} - ${safetyReasons}`; + } else { + return `Gemini Utils: action result status is error. Candidate was blocked due to ${candidate.finishReason}`; + } + } + } + return null; +} + +// not sure why these properties are not on the type, as they are on the data +interface SafetyReason extends SafetyRating { + blocked: boolean; + severity: string; +} + +const getSafetyReasons = (safetyRatings: SafetyReason[]) => { + const reasons = safetyRatings.filter((t: SafetyReason) => t.blocked); + return reasons.reduce( + (acc: string, t: SafetyReason, i: number) => + `${acc.length ? `${acc} ` : ''}${t.category}: ${t.severity}${ + i < reasons.length - 1 ? ',' : '' + }`, + '' + ); +}; + export const parseGeminiStreamAsAsyncIterator = async function* ( stream: Readable, logger: Logger, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts index dcef2ae6345f4..47a36ddf844b0 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts @@ -28,7 +28,7 @@ export const GENERATE_CHAT_TITLE_PROMPT = (responseLanguage: string, llmType?: s ? 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: + `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. Generate a relevant conversation title for the human's message in plain text. Make sure the title is formatted for the user, without using quotes or markdown. The title should clearly reflect the content of the message and be appropriate for a list of conversations. Please create the title in ${responseLanguage}. Respond only with the title. 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 diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts index 36c15aa44445d..2d076f6bd1472 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts @@ -7,6 +7,7 @@ import { RunnableConfig } from '@langchain/core/runnables'; import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; +import { formatLatestUserMessage } from '../prompts'; import { AgentState, NodeParamsBase } from '../types'; import { NodeType } from '../constants'; @@ -37,6 +38,8 @@ export async function runAgent({ const agentOutcome = await agentRunnable.withConfig({ tags: [AGENT_NODE_TAG] }).invoke( { ...state, + // prepend any user prompt (gemini) + input: formatLatestUserMessage(state.input, state.llmType), chat_history: state.messages, // TODO: Message de-dupe with ...state spread }, config diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts index ae8e3c18c2217..9eedce48ba69d 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts @@ -10,8 +10,11 @@ const YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT = const IF_YOU_DONT_KNOW_THE_ANSWER = 'Do not answer questions unrelated to Elastic Security.'; export const DEFAULT_SYSTEM_PROMPT = `${YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT} ${IF_YOU_DONT_KNOW_THE_ANSWER}`; - -export const GEMINI_SYSTEM_PROMPT = - `ALWAYS use the provided tools, as they have access to the latest data and syntax.` + - "The final response is the only output the user sees and should be a complete answer to the user's question. Do not leave out important tool output. The final response should never be empty. Don't forget to use tools."; +// system prompt from @afirstenberg +const BASE_GEMINI_PROMPT = + 'You are an assistant that is an expert at using tools and Elastic Security, doing your best to use these tools to answer questions or follow instructions. It is very important to use tools to answer the question or follow the instructions rather than coming up with your own answer. Tool calls are good. Sometimes you may need to make several tool calls to accomplish the task or get an answer to the question that was asked. Use as many tool calls as necessary.'; +const KB_CATCH = + 'If the knowledge base tool gives empty results, do your best to answer the question from the perspective of an expert security analyst.'; +export const GEMINI_SYSTEM_PROMPT = `${BASE_GEMINI_PROMPT} ${KB_CATCH}`; export const BEDROCK_SYSTEM_PROMPT = `Use tools as often as possible, as they have access to the latest data and syntax. Always return value from ESQLKnowledgeBaseTool as is. Never return tags in the response, but make sure to include tags content in the response. Do not reflect on the quality of the returned search results in your response.`; +export const GEMINI_USER_PROMPT = `Now, always using the tools at your disposal, step by step, come up with a response to this request:\n\n`; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts index eb52c227421fc..4a7b1fd46ccb8 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts @@ -10,6 +10,7 @@ import { BEDROCK_SYSTEM_PROMPT, DEFAULT_SYSTEM_PROMPT, GEMINI_SYSTEM_PROMPT, + GEMINI_USER_PROMPT, } from './nodes/translations'; export const formatPrompt = (prompt: string, additionalPrompt?: string) => @@ -23,7 +24,8 @@ export const formatPrompt = (prompt: string, additionalPrompt?: string) => export const systemPrompts = { openai: DEFAULT_SYSTEM_PROMPT, bedrock: `${DEFAULT_SYSTEM_PROMPT} ${BEDROCK_SYSTEM_PROMPT}`, - gemini: `${DEFAULT_SYSTEM_PROMPT} ${GEMINI_SYSTEM_PROMPT}`, + // The default prompt overwhelms gemini, do not prepend + gemini: GEMINI_SYSTEM_PROMPT, structuredChat: `Respond to the human as helpfully and accurately as possible. You have access to the following tools: {tools} @@ -98,3 +100,16 @@ export const formatPromptStructured = (prompt: string, additionalPrompt?: string ]); export const structuredChatAgentPrompt = formatPromptStructured(systemPrompts.structuredChat); + +/** + * If Gemini is the llmType, + * Adds a user prompt for the latest message in a conversation + * @param prompt + * @param llmType + */ +export const formatLatestUserMessage = (prompt: string, llmType?: string): string => { + if (llmType === 'gemini') { + return `${GEMINI_USER_PROMPT}${prompt}`; + } + return prompt; +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/utils.ts b/x-pack/plugins/elastic_assistant/server/routes/utils.ts index e163526d996ae..651a809e1a56e 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/utils.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/utils.ts @@ -17,7 +17,7 @@ import { ActionsClientChatOpenAI, ActionsClientBedrockChatModel, ActionsClientSimpleChatModel, - ActionsClientGeminiChatModel, + ActionsClientChatVertexAI, } from '@kbn/langchain/server'; import { CustomHttpRequestError } from './custom_http_request_error'; @@ -187,5 +187,5 @@ export const getLlmClass = (llmType?: string, bedrockChatEnabled?: boolean) => : llmType === 'bedrock' && bedrockChatEnabled ? ActionsClientBedrockChatModel : llmType === 'gemini' && bedrockChatEnabled - ? ActionsClientGeminiChatModel + ? ActionsClientChatVertexAI : ActionsClientSimpleChatModel; diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 58ee4a9dd4545..af8d019539a66 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -39,6 +39,7 @@ import { import { ActionsClientBedrockChatModel, ActionsClientChatOpenAI, + ActionsClientChatVertexAI, ActionsClientGeminiChatModel, ActionsClientLlm, ActionsClientSimpleChatModel, @@ -230,6 +231,7 @@ export type AssistantToolLlm = | ActionsClientBedrockChatModel | ActionsClientChatOpenAI | ActionsClientGeminiChatModel + | ActionsClientChatVertexAI | ActionsClientSimpleChatModel; export interface AssistantToolParams { diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts index 908d69ab0802e..6bf116c28719a 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts @@ -17,7 +17,7 @@ import { APP_UI_ID } from '../../../../common'; const toolDetails = { description: - 'Call this for knowledge on how to build an ESQL query, or answer questions about the ES|QL query language. Input must always be the query on a single line, with no other text. Your answer will be parsed as JSON, so never use quotes within the output and instead use backticks. Do not add any additional text to describe your output.', + 'Call this for knowledge on how to build an ESQL query, or answer questions about the ES|QL query language. Input must always be the user query on a single line, with no other text. Your answer will be parsed as JSON, so never use quotes within the output and instead use backticks. Do not add any additional text to describe your output.', id: 'esql-knowledge-base-tool', name: 'ESQLKnowledgeBaseTool', }; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/nl_to_esql_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/nl_to_esql_tool.ts index c313d877cc55b..a26d16607ac46 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/nl_to_esql_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/nl_to_esql_tool.ts @@ -20,16 +20,11 @@ const toolDetails = { id: 'nl-to-esql-tool', name: TOOL_NAME, description: `You MUST use the "${TOOL_NAME}" function when the user wants to: - - run any arbitrary query - breakdown or filter ES|QL queries that are displayed on the current page - convert queries from another language to ES|QL - asks general questions about ES|QL - DO NOT UNDER ANY CIRCUMSTANCES generate ES|QL queries or explain anything about the ES|QL query language yourself. - DO NOT UNDER ANY CIRCUMSTANCES try to correct an ES|QL query yourself - always use the "${TOOL_NAME}" function for this. - - Even if the "${TOOL_NAME}" function was used before that, follow it up with the "${TOOL_NAME}" function. If a query fails, do not attempt to correct it yourself. Again you should call the "${TOOL_NAME}" function, - even if it has been called before.`, + ALWAYS use this tool to generate ES|QL queries or explain anything about the ES|QL query language rather than coming up with your own answer.`, }; export const NL_TO_ESQL_TOOL: AssistantTool = { diff --git a/x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.ts b/x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.ts index 895dfe66d6de4..10aedb52bc6b0 100644 --- a/x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.ts +++ b/x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.ts @@ -395,9 +395,7 @@ const formatGeminiPayload = ({ temperature, maxOutputTokens: DEFAULT_TOKEN_LIMIT, }, - ...(systemInstruction - ? { system_instruction: { role: 'user', parts: [{ text: systemInstruction }] } } - : {}), + ...(systemInstruction ? { system_instruction: { parts: [{ text: systemInstruction }] } } : {}), ...(toolConfig ? { tool_config: { diff --git a/yarn.lock b/yarn.lock index d7f3fcbf7f72d..4fc1022f5a005 100644 --- a/yarn.lock +++ b/yarn.lock @@ -7312,17 +7312,16 @@ zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>=0.2.11 <0.3.0", "@langchain/core@>=0.2.16 <0.3.0", "@langchain/core@>=0.2.20 <0.3.0", "@langchain/core@>=0.2.5 <0.3.0", "@langchain/core@^0.2.18", "@langchain/core@~0.2.11": - version "0.2.18" - resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.18.tgz#1ac4f307fa217ab3555c9634147a6c4ad9826092" - integrity sha512-ru542BwNcsnDfjTeDbIkFIchwa54ctHZR+kVrC8U9NPS9/36iM8p8ruprOV7Zccj/oxtLE5UpEhV+9MZhVcFlA== +"@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>=0.2.11 <0.3.0", "@langchain/core@>=0.2.20 <0.3.0", "@langchain/core@>=0.2.5 <0.3.0", "@langchain/core@^0.2.18", "@langchain/core@~0.2.11": + version "0.2.32" + resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.32.tgz#a5dfbc49f8b6c15c8082763b93aeae8f9f4ca1a0" + integrity sha512-S27M+9Qou2qtcLfFGEvANkJ/zHq5XApeQsR6Q4I7C6v9x07eoYr558h6vVy6WQmKcksgbCIJ854ikwp173wBjA== dependencies: ansi-styles "^5.0.0" camelcase "6" decamelize "1.2.0" js-tiktoken "^1.0.12" - langsmith "~0.1.39" - ml-distance "^4.0.0" + langsmith "^0.1.43" mustache "^4.2.0" p-queue "^6.6.2" p-retry "4" @@ -7330,15 +7329,37 @@ zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/google-genai@^0.0.23": - version "0.0.23" - resolved "https://registry.yarnpkg.com/@langchain/google-genai/-/google-genai-0.0.23.tgz#e73af501bc1df4c7642b531759b82dc3eb7ae459" - integrity sha512-MTSCJEoKsfU1inz0PWvAjITdNFM4s41uvBCwLpcgx3jWJIEisczFD82x86ahYqJlb2fD6tohYSaCH/4tKAdkXA== +"@langchain/google-common@^0.1.1", "@langchain/google-common@~0.1.0": + version "0.1.1" + resolved "https://registry.yarnpkg.com/@langchain/google-common/-/google-common-0.1.1.tgz#7f8730e3dfb2920487dece5cbe738c13266b5e7b" + integrity sha512-oT/6lBev/Ufkp1dJbOTJ2S7xD9c+w9CqnqKqFOSxuZJbM4G8hzJtt7PDBOGfamIwtQP8dR7ORKXs1sCl+f5Tig== + dependencies: + uuid "^10.0.0" + zod-to-json-schema "^3.22.4" + +"@langchain/google-gauth@~0.1.0": + version "0.1.0" + resolved "https://registry.yarnpkg.com/@langchain/google-gauth/-/google-gauth-0.1.0.tgz#ea44941afede32a1d14fd18a6c6b500c7c7a53eb" + integrity sha512-0kps1NmaNiSl4n3lRw+7xsyhrEfIxNqBjih0kNYWPjLg55f9I9+QAlz7F1Sz/628HF1WQLFLQcBQA4geGzvenQ== + dependencies: + "@langchain/google-common" "~0.1.0" + google-auth-library "^8.9.0" + +"@langchain/google-genai@^0.1.0": + version "0.1.0" + resolved "https://registry.yarnpkg.com/@langchain/google-genai/-/google-genai-0.1.0.tgz#89552873210d72a5834de20fcbef3e6753283344" + integrity sha512-6rIba77zJVMj+048tLfkCBrkFbfAMiT+AfLEsu5s+CFoFmXMiI/dbKeDL4vhUWrJVb9uL4ZZyrnl0nKxyEKYgA== dependencies: "@google/generative-ai" "^0.7.0" - "@langchain/core" ">=0.2.16 <0.3.0" zod-to-json-schema "^3.22.4" +"@langchain/google-vertexai@^0.1.0": + version "0.1.0" + resolved "https://registry.yarnpkg.com/@langchain/google-vertexai/-/google-vertexai-0.1.0.tgz#e8ac6ae0bbdb6364d579def171ccfc1a26fd0cf2" + integrity sha512-xTi5NvNGSLQl/7OTsj4QTT0DkNbZ7cYDrEB0HqpZOwo6I5dulh/h2payGVQ6hdXj7Yyv78dRc5FdQSbyHui/WQ== + dependencies: + "@langchain/google-gauth" "~0.1.0" + "@langchain/langgraph@0.0.34": version "0.0.34" resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.34.tgz#1504c29ce524d08d6f076c34e0623c6de1f1246c" @@ -11509,16 +11530,16 @@ resolved "https://registry.yarnpkg.com/@types/use-sync-external-store/-/use-sync-external-store-0.0.3.tgz#b6725d5f4af24ace33b36fafd295136e75509f43" integrity sha512-EwmlvuaxPNej9+T4v5AuBPJa2x2UOJVdjCtDHgcDqitUeOtjnJKJ+apYjVcAoBEMjKW1VVFGZLUb5+qqa09XFA== +"@types/uuid@^10.0.0": + version "10.0.0" + resolved "https://registry.yarnpkg.com/@types/uuid/-/uuid-10.0.0.tgz#e9c07fe50da0f53dc24970cca94d619ff03f6f6d" + integrity sha512-7gqG38EyHgyP1S+7+xomFtL+ZNHcKv6DwNaCZmJmo1vgMugyF3TCnXVg4t1uk89mLNwnLtnY3TpOpCOyp1/xHQ== + "@types/uuid@^9.0.0": version "9.0.0" resolved "https://registry.yarnpkg.com/@types/uuid/-/uuid-9.0.0.tgz#53ef263e5239728b56096b0a869595135b7952d2" integrity sha512-kr90f+ERiQtKWMz5rP32ltJ/BtULDI5RVO0uavn1HQUOwjx0R1h0rnDYNL0CepF1zL5bSY6FISAfd9tOdDhU5Q== -"@types/uuid@^9.0.1": - version "9.0.2" - resolved "https://registry.yarnpkg.com/@types/uuid/-/uuid-9.0.2.tgz#ede1d1b1e451548d44919dc226253e32a6952c4b" - integrity sha512-kNnC1GFBLuhImSnV7w4njQkUiJi0ZXUycu1rUaouPqiKlXkh77JKgdRnTAp1x5eBwcIwbtI+3otwzuIDEuDoxQ== - "@types/vinyl-fs@*", "@types/vinyl-fs@^3.0.2": version "3.0.2" resolved "https://registry.yarnpkg.com/@types/vinyl-fs/-/vinyl-fs-3.0.2.tgz#cbaef5160ad7695483af0aa1b4fe67f166c18feb" @@ -12760,7 +12781,7 @@ arrify@^1.0.1: resolved "https://registry.yarnpkg.com/arrify/-/arrify-1.0.1.tgz#898508da2226f380df904728456849c1501a4b0d" integrity sha1-iYUI2iIm84DfkEcoRWhJwVAaSw0= -arrify@^2.0.1: +arrify@^2.0.0, arrify@^2.0.1: version "2.0.1" resolved "https://registry.yarnpkg.com/arrify/-/arrify-2.0.1.tgz#c9655e9331e0abcd588d2a7cad7e9956f66701fa" integrity sha512-3duEwti880xqi4eAMN8AyR4a0ByT90zoYdLlevfrvU43vb0YZwZVfxOgxWrLXXXpyugL0hNZc9G6BiB5B3nUug== @@ -18019,6 +18040,11 @@ fast-stream-to-buffer@^1.0.0: dependencies: end-of-stream "^1.4.1" +fast-text-encoding@^1.0.0: + version "1.0.6" + resolved "https://registry.yarnpkg.com/fast-text-encoding/-/fast-text-encoding-1.0.6.tgz#0aa25f7f638222e3396d72bf936afcf1d42d6867" + integrity sha512-VhXlQgj9ioXCqGstD37E/HBeqEGV/qOD/kmbVG8h5xKBYvM1L3lR1Zn4555cQ8GkYbJa8aJSipLPndE1k6zK2w== + fastest-levenshtein@^1.0.12: version "1.0.12" resolved "https://registry.yarnpkg.com/fastest-levenshtein/-/fastest-levenshtein-1.0.12.tgz#9990f7d3a88cc5a9ffd1f1745745251700d497e2" @@ -18691,6 +18717,16 @@ gauge@^3.0.0: strip-ansi "^6.0.1" wide-align "^1.1.2" +gaxios@^5.0.0, gaxios@^5.0.1: + version "5.1.3" + resolved "https://registry.yarnpkg.com/gaxios/-/gaxios-5.1.3.tgz#f7fa92da0fe197c846441e5ead2573d4979e9013" + integrity sha512-95hVgBRgEIRQQQHIbnxBXeHbW4TqFk4ZDJW7wmVtvYar72FdhRIo1UGOLS2eRAKCPEdPBWu+M7+A33D9CdX9rA== + dependencies: + extend "^3.0.2" + https-proxy-agent "^5.0.0" + is-stream "^2.0.0" + node-fetch "^2.6.9" + gaxios@^6.0.0, gaxios@^6.1.1: version "6.6.0" resolved "https://registry.yarnpkg.com/gaxios/-/gaxios-6.6.0.tgz#af8242fff0bbb82a682840d5feaa91b6a1c58be4" @@ -18702,6 +18738,14 @@ gaxios@^6.0.0, gaxios@^6.1.1: node-fetch "^2.6.9" uuid "^9.0.1" +gcp-metadata@^5.3.0: + version "5.3.0" + resolved "https://registry.yarnpkg.com/gcp-metadata/-/gcp-metadata-5.3.0.tgz#6f45eb473d0cb47d15001476b48b663744d25408" + integrity sha512-FNTkdNEnBdlqF2oatizolQqNANMrcqJt6AAYt99B3y1aLLC8Hc5IOBb+ZnnzllodEEf6xMBp6wRcBbc16fa65w== + dependencies: + gaxios "^5.0.0" + json-bigint "^1.0.0" + gcp-metadata@^6.1.0: version "6.1.0" resolved "https://registry.yarnpkg.com/gcp-metadata/-/gcp-metadata-6.1.0.tgz#9b0dd2b2445258e7597f2024332d20611cbd6b8c" @@ -19149,6 +19193,21 @@ gonzales-pe@^4.3.0: dependencies: minimist "^1.2.5" +google-auth-library@^8.9.0: + version "8.9.0" + resolved "https://registry.yarnpkg.com/google-auth-library/-/google-auth-library-8.9.0.tgz#15a271eb2ec35d43b81deb72211bd61b1ef14dd0" + integrity sha512-f7aQCJODJFmYWN6PeNKzgvy9LI2tYmXnzpNDHEjG5sDNPgGb2FXQyTBnXeSH+PAtpKESFD+LmHw3Ox3mN7e1Fg== + dependencies: + arrify "^2.0.0" + base64-js "^1.3.0" + ecdsa-sig-formatter "^1.0.11" + fast-text-encoding "^1.0.0" + gaxios "^5.0.0" + gcp-metadata "^5.3.0" + gtoken "^6.1.0" + jws "^4.0.0" + lru-cache "^6.0.0" + google-auth-library@^9.10.0: version "9.10.0" resolved "https://registry.yarnpkg.com/google-auth-library/-/google-auth-library-9.10.0.tgz#c9fb940923f7ff2569d61982ee1748578c0bbfd4" @@ -19161,6 +19220,13 @@ google-auth-library@^9.10.0: gtoken "^7.0.0" jws "^4.0.0" +google-p12-pem@^4.0.0: + version "4.0.1" + resolved "https://registry.yarnpkg.com/google-p12-pem/-/google-p12-pem-4.0.1.tgz#82841798253c65b7dc2a4e5fe9df141db670172a" + integrity sha512-WPkN4yGtz05WZ5EhtlxNDWPhC4JIic6G8ePitwUWy4l+XPVYec+a0j0Ts47PDtW59y3RwAhUd9/h9ZZ63px6RQ== + dependencies: + node-forge "^1.3.1" + google-protobuf@^3.6.1: version "3.19.4" resolved "https://registry.yarnpkg.com/google-protobuf/-/google-protobuf-3.19.4.tgz#8d32c3e34be9250956f28c0fb90955d13f311888" @@ -19226,6 +19292,15 @@ graphql@^16.6.0: resolved "https://registry.yarnpkg.com/graphql/-/graphql-16.8.1.tgz#1930a965bef1170603702acdb68aedd3f3cf6f07" integrity sha512-59LZHPdGZVh695Ud9lRzPBVTtlX9ZCV150Er2W43ro37wVof0ctenSaskPPjN7lVTIN8mSZt8PHUNKZuNQUuxw== +gtoken@^6.1.0: + version "6.1.2" + resolved "https://registry.yarnpkg.com/gtoken/-/gtoken-6.1.2.tgz#aeb7bdb019ff4c3ba3ac100bbe7b6e74dce0e8bc" + integrity sha512-4ccGpzz7YAr7lxrT2neugmXQ3hP9ho2gcaityLVkiUecAiwiy60Ii8gRbZeOsXV19fYaRjgBSshs8kXw+NKCPQ== + dependencies: + gaxios "^5.0.1" + google-p12-pem "^4.0.0" + jws "^4.0.0" + gtoken@^7.0.0: version "7.1.0" resolved "https://registry.yarnpkg.com/gtoken/-/gtoken-7.1.0.tgz#d61b4ebd10132222817f7222b1e6064bd463fc26" @@ -19915,7 +19990,7 @@ https-browserify@^1.0.0: resolved "https://registry.yarnpkg.com/https-browserify/-/https-browserify-1.0.0.tgz#ec06c10e0a34c0f2faf199f7fd7fc78fffd03c73" integrity sha1-7AbBDgo0wPL68Zn3/X/Hj//QPHM= -https-proxy-agent@^5.0.1: +https-proxy-agent@^5.0.0, https-proxy-agent@^5.0.1: version "5.0.1" resolved "https://registry.yarnpkg.com/https-proxy-agent/-/https-proxy-agent-5.0.1.tgz#c59ef224a04fe8b754f3db0063a25ea30d0005d6" integrity sha512-dFcAjpTQFgoLMzC2VwU+C/CbS7uRL0lWmxDITmqm7C+7F0Odmj6s9l6alZc6AELXhrnggM2CeWSXHGOdX2YtwA== @@ -22093,16 +22168,17 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.39, langsmith@~0.1.30, langsmith@~0.1.39: - version "0.1.39" - resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.39.tgz#cc99f1828a9c0f5ba24bec6b0121edc44e8d282d" - integrity sha512-K2/qbc96JhrZbSL74RbZ0DBOpTB9Mxicu8RQrZ88Xsp1bH2O3+y5EdcvC0g/1YzQWQhcQ4peknCA24c3VTNiNA== +langsmith@^0.1.43, langsmith@^0.1.55, langsmith@~0.1.30: + version "0.1.55" + resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.55.tgz#bdbb8015a28093f4a248c0ee9b8937731c5baa93" + integrity sha512-6NVtI04UUnIY59I/imOX02FG/QMGfqStu8tiJtyyreKMv2GAN0EE9Z5Ap1wzOe6v8ukEcV3NwEO2LYOPwup1PQ== dependencies: - "@types/uuid" "^9.0.1" + "@types/uuid" "^10.0.0" commander "^10.0.1" p-queue "^6.6.2" p-retry "4" - uuid "^9.0.0" + semver "^7.6.3" + uuid "^10.0.0" language-subtag-registry@~0.3.2: version "0.3.21"