From 0b03e9883e40d14e381a10fbf1bb0dc99ba4abaf Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 30 May 2024 17:23:53 -0600 Subject: [PATCH 01/47] Migrates to LangSmith and adds KB Tools --- package.json | 13 +- .../knowledge_base/common_attributes.gen.ts | 4 +- .../common_attributes.schema.yaml | 4 +- .../create_knowledge_base_entry.ts | 2 +- .../knowledge_base/index.ts | 39 +++- .../helpers/get_msearch_query_body.ts | 10 +- .../execute_custom_llm_chain/index.ts | 9 +- .../server/lib/langchain/executors/types.ts | 14 +- .../graphs/default_assistant_graph/graph.ts | 120 +++++++++++ .../graphs/default_assistant_graph/helpers.ts | 189 ++++++++++++++++++ .../graphs/default_assistant_graph/index.ts | 145 ++++++++++++++ .../graphs/default_assistant_graph/prompts.ts | 55 +++++ .../graphs/default_assistant_graph/types.ts | 19 ++ .../lib/langchain/tracers/apm_tracer.ts | 2 +- .../routes/post_actions_connector_execute.ts | 17 +- .../plugins/elastic_assistant/server/types.ts | 2 + .../server/assistant/tools/index.ts | 4 + .../knowledge_base_retrieval_tool.ts | 53 +++++ .../knowledge_base_write_tool.ts | 64 ++++++ yarn.lock | 157 ++++++++++----- 20 files changed, 844 insertions(+), 78 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts diff --git a/package.json b/package.json index 096c76042f828..7b98d63f3b257 100644 --- a/package.json +++ b/package.json @@ -80,7 +80,7 @@ "resolutions": { "**/@bazel/typescript/protobufjs": "6.11.4", "**/@hello-pangea/dnd": "16.6.0", - "**/@langchain/core": "0.1.53", + "**/@langchain/core": "0.2.0", "**/@types/node": "20.10.5", "**/@typescript-eslint/utils": "5.62.0", "**/chokidar": "^3.5.3", @@ -913,9 +913,10 @@ "@kbn/watcher-plugin": "link:x-pack/plugins/watcher", "@kbn/xstate-utils": "link:packages/kbn-xstate-utils", "@kbn/zod-helpers": "link:packages/kbn-zod-helpers", - "@langchain/community": "^0.0.44", - "@langchain/core": "^0.1.53", - "@langchain/openai": "^0.0.25", + "@langchain/community": "^0.2.2", + "@langchain/core": "^0.2.0", + "@langchain/langgraph": "^0.0.20", + "@langchain/openai": "^0.0.33", "@loaders.gl/core": "^3.4.7", "@loaders.gl/json": "^3.4.7", "@loaders.gl/shapefile": "^3.4.7", @@ -1052,8 +1053,8 @@ "jsonwebtoken": "^9.0.2", "jsts": "^1.6.2", "kea": "^2.6.0", - "langchain": "^0.1.30", - "langsmith": "^0.1.14", + "langchain": "^0.2.2", + "langsmith": "^0.1.28", "launchdarkly-js-client-sdk": "^3.1.4", "launchdarkly-node-server-sdk": "^7.0.3", "load-json-file": "^6.2.0", diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts index 0d44cbe51e320..a4ca22cb03dbd 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts @@ -33,11 +33,11 @@ export const KnowledgeBaseEntryErrorSchema = z export type Metadata = z.infer; export const Metadata = z.object({ /** - * Knowledge Base resource name + * Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc */ kbResource: z.string(), /** - * Original text content source + * Source document name or filepath */ source: z.string(), /** diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml index a9a9794852953..63fcdf16b8fb1 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml @@ -32,10 +32,10 @@ components: properties: kbResource: type: string - description: Knowledge Base resource name + description: Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc source: type: string - description: Original text content source + description: Source document name or filepath required: type: boolean description: Whether or not this resource should always be included diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts index 5430bf597ebe7..90c0850c503fe 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts @@ -18,8 +18,8 @@ import { CreateKnowledgeBaseEntrySchema } from './types'; export interface CreateKnowledgeBaseEntryParams { esClient: ElasticsearchClient; - logger: Logger; knowledgeBaseIndex: string; + logger: Logger; spaceId: string; user: AuthenticatedUser; knowledgeBaseEntry: KnowledgeBaseEntryCreateProps; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 771ec35c07c51..d2dd1165f6894 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -13,13 +13,16 @@ import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import type { KibanaRequest } from '@kbn/core-http-server'; import type { Document } from 'langchain/document'; import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; -import { KnowledgeBaseEntryResponse } from '@kbn/elastic-assistant-common'; +import { + KnowledgeBaseEntryCreateProps, + KnowledgeBaseEntryResponse, +} from '@kbn/elastic-assistant-common'; import pRetry from 'p-retry'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader'; import { GetElser } from '../../types'; -import { transformToCreateSchema } from './create_knowledge_base_entry'; +import { createKnowledgeBaseEntry, transformToCreateSchema } from './create_knowledge_base_entry'; import { EsKnowledgeBaseEntrySchema } from './types'; import { transformESSearchToKnowledgeBaseEntry } from './transforms'; import { ESQL_DOCS_LOADED_QUERY } from '../../routes/knowledge_base/constants'; @@ -261,4 +264,36 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { return created?.data ? transformESSearchToKnowledgeBaseEntry(created?.data) : []; }; + + /** + * Creates a new Knowledge Base Entry. + * + * @param knowledgeBaseEntry + */ + public createKnowledgeBaseEntry = async ({ + knowledgeBaseEntry, + }: { + knowledgeBaseEntry: KnowledgeBaseEntryCreateProps; + }): Promise => { + const authenticatedUser = this.options.currentUser; + if (authenticatedUser == null) { + throw new Error( + 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' + ); + } + + this.options.logger.debug( + `Creating Knowledge Base Entry:\n ${JSON.stringify(knowledgeBaseEntry, null, 2)}` + ); + this.options.logger.debug(`kbIndex: ${this.indexTemplateAndPattern.alias}`); + const esClient = await this.options.elasticsearchClientPromise; + return createKnowledgeBaseEntry({ + esClient, + knowledgeBaseIndex: this.indexTemplateAndPattern.alias, + logger: this.options.logger, + spaceId: this.spaceId, + user: authenticatedUser, + knowledgeBaseEntry, + }); + }; } diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_msearch_query_body.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_msearch_query_body.ts index c93c3f2e30954..7780eab2c2d1d 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_msearch_query_body.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_msearch_query_body.ts @@ -58,10 +58,10 @@ export const getMsearchQueryBody = ({ query: vectorSearchQuery, size: vectorSearchQuerySize, }, - { index }, - { - query: termsSearchQuery, - size: termsSearchQuerySize, - }, + // { index }, + // { + // query: termsSearchQuery, + // size: termsSearchQuerySize, + // }, ], }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts index 8323712c50aa7..555336b301c31 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts @@ -38,6 +38,7 @@ export const callAgentExecutor: AgentExecutor = async ({ connectorId, esClient, esStore, + kbDataClient, langChainMessages, llmType, logger, @@ -89,12 +90,14 @@ export const callAgentExecutor: AgentExecutor = async ({ // Fetch any applicable tools that the source plugin may have registered const assistantToolParams: AssistantToolParams = { - anonymizationFields, alertsIndexPattern, - isEnabledKnowledgeBase, + anonymizationFields, chain, - llm, esClient, + isEnabledKnowledgeBase, + kbDataClient, + llm, + logger, modelExists, onNewReplacements, replacements, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index 8acd7f4fcdde2..56cceb3448a9d 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -17,6 +17,13 @@ import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/s import { ResponseBody } from '../types'; import type { AssistantTool } from '../../../types'; import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store'; +import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base'; + +export type OnLlmResponse = ( + content: string, + traceData?: Message['traceData'], + isError?: boolean +) => Promise; export interface AgentExecutorParams { abortSignal?: AbortSignal; @@ -28,17 +35,14 @@ export interface AgentExecutorParams { connectorId: string; esClient: ElasticsearchClient; esStore: ElasticsearchStore; + kbDataClient?: AIAssistantKnowledgeBaseDataClient | undefined; langChainMessages: BaseMessage[]; llmType?: string; logger: Logger; onNewReplacements?: (newReplacements: Replacements) => void; replacements: Replacements; isStream?: T; - onLlmResponse?: ( - content: string, - traceData?: Message['traceData'], - isError?: boolean - ) => Promise; + onLlmResponse?: OnLlmResponse; request: KibanaRequest; size?: number; traceOptions?: TraceOptions; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts new file mode 100644 index 0000000000000..e70df0b506d25 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts @@ -0,0 +1,120 @@ +/* + * 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 { ToolExecutor } from '@langchain/langgraph/prebuilt'; +import { RunnableConfig, RunnableLambda } from '@langchain/core/runnables'; +import { END, START, StateGraph, StateGraphArgs } from '@langchain/langgraph'; +import { AgentAction, AgentFinish, AgentStep } from '@langchain/core/agents'; +import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; +import { StructuredTool } from '@langchain/core/tools'; +import type { CompiledStateGraph } from '@langchain/langgraph/dist/graph/state'; +import type { Logger } from '@kbn/logging'; + +import { BaseMessage } from '@langchain/core/messages'; +import { AgentState } from './types'; + +export const DEFAULT_ASSISTANT_GRAPH_ID = 'Default Security Assistant Graph'; + +interface GetDefaultAssistantGraphParams { + agentRunnable: AgentRunnableSequence; + logger: Logger; + messages: BaseMessage[]; + tools: StructuredTool[]; +} + +/** + * Returns a compiled default assistant graph + */ +export const getDefaultAssistantGraph = ({ + agentRunnable, + logger, + messages, + tools, +}: GetDefaultAssistantGraphParams): CompiledStateGraph => { + try { + // Default graph state + const graphState: StateGraphArgs['channels'] = { + input: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + steps: { + value: (x: AgentStep[], y: AgentStep[]) => x.concat(y), + default: () => [], + }, + agentOutcome: { + value: ( + x: AgentAction | AgentFinish | undefined, + y?: AgentAction | AgentFinish | undefined + ) => y ?? x, + default: () => undefined, + }, + messages: { + value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y), + default: () => messages, + }, + }; + + // Create a tool executor + const toolExecutor = new ToolExecutor({ tools }); + + // Define logic that will be used to determine which conditional edge to go down + const shouldContinue = (state: AgentState) => { + logger.debug(`graph:shouldContinue:state\n${JSON.stringify(state, null, 2)}`); + if (state.agentOutcome && 'returnValues' in state.agentOutcome) { + return 'end'; + } + return 'continue'; + }; + + const runAgent = async (state: AgentState, config?: RunnableConfig) => { + logger.debug(`graph:runAgent:\nstate\n${JSON.stringify(state, null, 2)}`); + + const agentOutcome = await agentRunnable.invoke( + { + ...state, + chat_history: messages, + knowledge_history: 'The users favorite color is blue', // TODO: Plumb through initial retrieval + }, + config + ); + return { + agentOutcome, + }; + }; + + const executeTools = async (state: AgentState, config?: RunnableConfig) => { + logger.debug(`graph:executeTools:state\n${JSON.stringify(state, null, 2)}`); + const agentAction = state.agentOutcome; + if (!agentAction || 'returnValues' in agentAction) { + throw new Error('Agent has not been run yet'); + } + const out = await toolExecutor.invoke(agentAction, config); + return { + steps: [{ action: agentAction, observation: JSON.stringify(out, null, 2) }], + }; + }; + + // Create a new graph, with the default state from above + const workflow = new StateGraph({ channels: graphState }); + + // Define the nodes to cycle between + workflow.addNode('agent', new RunnableLambda({ func: runAgent })); + workflow.addNode('action', new RunnableLambda({ func: executeTools })); + + // Add conditional edge for determining if we shouldContinue + workflow.addConditionalEdges('agent', shouldContinue, { continue: 'action', end: END }); + + // Add edges for start, and between agent and action (action always followed by agent) + workflow.addEdge(START, 'agent'); + workflow.addEdge('action', 'agent'); + + return workflow.compile(); + } catch (e) { + throw new Error(`Unable to compile DefaultAssistantGraph\n${e}`); + } +}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts new file mode 100644 index 0000000000000..f2c318b32a2ab --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts @@ -0,0 +1,189 @@ +/* + * 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 agent, { Span } from 'elastic-apm-node'; +import type { Logger } from '@kbn/logging'; +import { streamFactory, StreamResponseWithHeaders } from '@kbn/ml-response-stream/server'; +import { transformError } from '@kbn/securitysolution-es-utils'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { ExecuteConnectorRequestBody, TraceData } from '@kbn/elastic-assistant-common'; +import { CompiledStateGraph } from '@langchain/langgraph/dist/graph/state'; +import { DEFAULT_ASSISTANT_GRAPH_ID } from './graph'; +import type { OnLlmResponse, TraceOptions } from '../../executors/types'; +import type { APMTracer } from '../../tracers/apm_tracer'; +import { withAssistantSpan } from '../../tracers/with_assistant_span'; + +interface StreamGraphParams { + apmTracer: APMTracer; + assistantGraph: CompiledStateGraph; + inputs: { input: string }; + logger: Logger; + onLlmResponse?: OnLlmResponse; + request: KibanaRequest; + traceOptions?: TraceOptions; +} + +/** + * Execute the graph in streaming mode + * + * @param apmTracer + * @param assistantGraph + * @param inputs + * @param logger + * @param onLlmResponse + * @param request + * @param traceOptions + */ +export const streamGraph = async ({ + apmTracer, + assistantGraph, + inputs, + logger, + onLlmResponse, + request, + traceOptions, +}: StreamGraphParams): Promise => { + let streamingSpan: Span | undefined; + if (agent.isStarted()) { + streamingSpan = agent.startSpan(`${DEFAULT_ASSISTANT_GRAPH_ID} (Streaming)`) ?? undefined; + } + const { + end: streamEnd, + push, + responseWithHeaders, + } = streamFactory<{ type: string; payload: string }>(request.headers, logger, false, false); + + let didEnd = false; + const handleStreamEnd = (finalResponse: string, isError = false) => { + if (onLlmResponse) { + onLlmResponse( + finalResponse, + { + transactionId: streamingSpan?.transaction?.ids?.['transaction.id'], + traceId: streamingSpan?.ids?.['trace.id'], + }, + isError + ).catch(() => {}); + } + streamEnd(); + didEnd = true; + if ((streamingSpan && !streamingSpan?.outcome) || streamingSpan?.outcome === 'unknown') { + streamingSpan.outcome = 'success'; + } + streamingSpan?.end(); + }; + + let finalMessage = ''; + const stream = assistantGraph.streamEvents(inputs, { + callbacks: [apmTracer, ...(traceOptions?.tracers ?? [])], + runName: DEFAULT_ASSISTANT_GRAPH_ID, + streamMode: 'values', + tags: traceOptions?.tags ?? [], + version: 'v1', + }); + + const processEvent = async () => { + try { + const { value, done } = await stream.next(); + if (done) return; + + const event = value; + if (event.event === 'on_llm_stream') { + const chunk = event.data?.chunk; + // TODO: For Bedrock streaming support, override `handleLLMNewToken` in callbacks, + // TODO: or maybe we can update ActionsClientSimpleChatModel to handle this `on_llm_stream` event + if (event.name === 'ActionsClientChatOpenAI') { + const msg = chunk.message; + + if (msg.tool_call_chunks && msg.tool_call_chunks.length > 0) { + } else if (!didEnd) { + if (msg.response_metadata.finish_reason === 'stop') { + handleStreamEnd(finalMessage); + } else { + push({ payload: msg.content, type: 'content' }); + finalMessage += msg.content; + } + } + } + } + + processEvent(); + } catch (err) { + // if I throw an error here, it crashes the server. Not sure how to get around that. + // If I put await on this function the error works properly, but when there is not an error + // it waits for the entire stream to complete before resolving + const error = transformError(err); + + if (error.message === 'AbortError') { + // user aborted the stream, we must end it manually here + return handleStreamEnd(finalMessage); + } + logger.error(`Error streaming from LangChain: ${error.message}`); + push({ payload: error.message, type: 'content' }); + handleStreamEnd(error.message, true); + } + }; + + // Start processing events, do not await! Return `responseWithHeaders` immediately + processEvent(); + + return responseWithHeaders; +}; + +interface InvokeGraphParams { + apmTracer: APMTracer; + assistantGraph: CompiledStateGraph; + inputs: { input: string }; + onLlmResponse?: OnLlmResponse; + traceOptions?: TraceOptions; +} +interface InvokeGraphResponse { + output: string; + traceData: TraceData; +} + +/** + * Execute the graph in non-streaming mode + * + * @param apmTracer + * @param assistantGraph + * @param inputs + * @param onLlmResponse + * @param traceOptions + */ +export const invokeGraph = async ({ + apmTracer, + assistantGraph, + inputs, + onLlmResponse, + traceOptions, +}: InvokeGraphParams): Promise => { + return withAssistantSpan(DEFAULT_ASSISTANT_GRAPH_ID, async (span) => { + let traceData: TraceData = {}; + if (span?.transaction?.ids['transaction.id'] != null && span?.ids['trace.id'] != null) { + traceData = { + // Transactions ID since this span is the parent + transactionId: span.transaction.ids['transaction.id'], + traceId: span.ids['trace.id'], + }; + span.addLabels({ evaluationId: traceOptions?.evaluationId }); + } + + const r = await assistantGraph.invoke(inputs, { + callbacks: [apmTracer, ...(traceOptions?.tracers ?? [])], + runName: DEFAULT_ASSISTANT_GRAPH_ID, + tags: traceOptions?.tags ?? [], + }); + const output = r.agentOutcome.returnValues.output; + + if (onLlmResponse) { + await onLlmResponse(output, traceData); + } + + return { output, traceData }; + }); +}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts new file mode 100644 index 0000000000000..750f0b4d971cb --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts @@ -0,0 +1,145 @@ +/* + * 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 { StructuredTool } from '@langchain/core/tools'; +import { RetrievalQAChain } from 'langchain/chains'; +import { + getDefaultArguments, + ActionsClientChatOpenAI, + ActionsClientSimpleChatModel, +} from '@kbn/langchain/server'; +import { createOpenAIFunctionsAgent, createStructuredChatAgent } from 'langchain/agents'; +import { AssistantToolParams } from '../../../../types'; +import { AgentExecutor } from '../../executors/types'; +import { openAIFunctionAgentPrompt, structuredChatAgentPrompt } from './prompts'; +import { APMTracer } from '../../tracers/apm_tracer'; +import { getDefaultAssistantGraph } from './graph'; +import { invokeGraph, streamGraph } from './helpers'; + +/** + * + * + */ +export const callAssistantGraph: AgentExecutor = async ({ + abortSignal, + actions, + alertsIndexPattern, + anonymizationFields, + isEnabledKnowledgeBase, + assistantTools = [], + connectorId, + esClient, + esStore, + kbDataClient, + langChainMessages, + llmType, + logger: parentLogger, + isStream = false, + onLlmResponse, + onNewReplacements, + replacements, + request, + size, + traceOptions, +}) => { + const logger = parentLogger.get('defaultAssistantGraph'); + const isOpenAI = llmType === 'openai'; + const llmClass = isOpenAI ? ActionsClientChatOpenAI : ActionsClientSimpleChatModel; + + const llm = new llmClass({ + actions, + connectorId, + request, + llmType, + logger, + // possible client model override, + // let this be undefined otherwise so the connector handles the model + model: request.body.model, + // ensure this is defined because we default to it in the language_models + // This is where the LangSmith logs (Metadata > Invocation Params) are set + temperature: getDefaultArguments(llmType).temperature, + signal: abortSignal, + streaming: isStream, + // prevents the agent from retrying on failure + // failure could be due to bad connector, we should deliver that result to the client asap + maxRetries: 0, + }); + const model = llm; + + const messages = langChainMessages.slice(0, -1); // all but the last message + const latestMessage = langChainMessages.slice(-1); // the last message + + const modelExists = await esStore.isModelInstalled(); + + // Create a chain that uses the ELSER backed ElasticsearchStore, override k=10 for esql query generation for now + const chain = RetrievalQAChain.fromLLM(model, esStore.asRetriever(10)); + + // Fetch any applicable tools that the source plugin may have registered + const assistantToolParams: AssistantToolParams = { + alertsIndexPattern, + anonymizationFields, + chain, + esClient, + isEnabledKnowledgeBase, + kbDataClient, + llm: model, + logger, + modelExists, + onNewReplacements, + replacements, + request, + size, + }; + + const tools: StructuredTool[] = assistantTools.flatMap( + (tool) => tool.getTool(assistantToolParams) ?? [] + ); + + const agentRunnable = isOpenAI + ? await createOpenAIFunctionsAgent({ + llm, + tools, + prompt: openAIFunctionAgentPrompt, + streamRunnable: isStream, + }) + : await createStructuredChatAgent({ + llm, + tools, + prompt: structuredChatAgentPrompt, + streamRunnable: isStream, + }); + + const apmTracer = new APMTracer({ projectName: traceOptions?.projectName ?? 'default' }, logger); + + const assistantGraph = getDefaultAssistantGraph({ agentRunnable, logger, messages, tools }); + const inputs = { input: latestMessage[0].content as string }; + + if (isStream) { + return streamGraph({ apmTracer, assistantGraph, inputs, logger, onLlmResponse, request }); + } + + const graphResponse = await invokeGraph({ + apmTracer, + assistantGraph, + inputs, + onLlmResponse, + traceOptions, + }); + + return { + body: { + connector_id: connectorId, + data: graphResponse.output, + trace_data: graphResponse.traceData, + replacements, + status: 'ok', + }, + headers: { + 'content-type': 'application/json', + }, + }; +}; 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 new file mode 100644 index 0000000000000..1b6d8f5601c05 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts @@ -0,0 +1,55 @@ +/* + * 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 { ChatPromptTemplate } from '@langchain/core/prompts'; + +export const openAIFunctionAgentPrompt = ChatPromptTemplate.fromMessages([ + [ + 'system', + 'You are a helpful assistant\n\nUse the below context as a sample of information about the user from their knowledge base:\n\n```{knowledge_history}```', + ], + ['placeholder', '{chat_history}'], + // ['human', '{knowledge_history}'], + ['human', '{input}'], + ['placeholder', '{agent_scratchpad}'], +]); + +export const structuredChatAgentPrompt = ChatPromptTemplate.fromMessages([ + [ + 'system', + 'Respond to the human as helpfully and accurately as possible. You have access to the following tools:\n\n' + + '{tools}\n\n' + + 'Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).\n\n' + + 'Valid "action" values: "Final Answer" or {tool_names}\n\n' + + 'Provide only ONE action per $JSON_BLOB, as shown:\n\n' + + '```\n\n' + + '{{\n\n' + + ' "action": $TOOL_NAME,\n\n' + + ' "action_input": $INPUT\n\n' + + '}}\n\n' + + '```\n\n' + + 'Follow this format:\n\n' + + 'Question: input question to answer\n\n' + + 'Thought: consider previous and subsequent steps\n\n' + + 'Action:\n\n' + + '```\n\n' + + '$JSON_BLOB\n\n' + + '```\n\n' + + 'Observation: action result\n\n' + + '... (repeat Thought/Action/Observation N times)\n\n' + + 'Thought: I know what to respond\n\n' + + 'Action:\n\n' + + '```\n\n' + + '{{\n\n' + + ' "action": "Final Answer",\n\n' + + ' "action_input": "Final response to human"\n\n' + + '}}\n\n' + + 'Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation', + ], + ['placeholder', '{chat_history}'], + ['human', '{input}\n\n{agent_scratchpad}\n(reminder to respond in a JSON blob no matter what)'], +]); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts new file mode 100644 index 0000000000000..6c7543f9e7e28 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts @@ -0,0 +1,19 @@ +/* + * 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 { BaseMessage } from '@langchain/core/messages'; +import { AgentAction, AgentFinish, AgentStep } from '@langchain/core/agents'; + +export interface AgentStateBase { + agentOutcome?: AgentAction | AgentFinish; + steps: AgentStep[]; +} + +export interface AgentState extends AgentStateBase { + input: string; + messages: BaseMessage[]; +} diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/apm_tracer.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/apm_tracer.ts index e5f91a379ed3d..ba4baef0433ad 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/apm_tracer.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/apm_tracer.ts @@ -38,7 +38,7 @@ export class APMTracer extends BaseTracer implements LangChainTracerFields { this.projectName = projectName ?? 'default'; this.exampleId = exampleId; - this.logger = logger; + this.logger = logger.get('apmTracer'); } protected async persistRun(_run: Run): Promise {} diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 2d53106bacf13..4e46ad16e7d2e 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -42,6 +42,7 @@ import { getLangSmithTracer } from './evaluate/utils'; import { EsAnonymizationFieldsSchema } from '../ai_assistant_data_clients/anonymization_fields/types'; import { transformESSearchToAnonymizationFields } from '../ai_assistant_data_clients/anonymization_fields/helpers'; import { ElasticsearchStore } from '../lib/langchain/elasticsearch_store/elasticsearch_store'; +import { callAssistantGraph } from '../lib/langchain/graphs/default_assistant_graph'; export const postActionsConnectorExecuteRoute = ( router: IRouter, @@ -325,7 +326,7 @@ export const postActionsConnectorExecuteRoute = ( }); // Create an ElasticsearchStore for KB interactions - // Setup with kbDataClient if `enableKnowledgeBaseByDefault` FF is enabled + // Setup with kbDataClient if `assistantKnowledgeBaseByDefault` FF is enabled const enableKnowledgeBaseByDefault = assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; const kbDataClient = enableKnowledgeBaseByDefault @@ -345,7 +346,8 @@ export const postActionsConnectorExecuteRoute = ( kbDataClient ); - const result: StreamResponseWithHeaders | StaticReturnType = await callAgentExecutor({ + // Shared executor params + const executorParams = { abortSignal, alertsIndexPattern: request.body.alertsIndexPattern, anonymizationFields: anonymizationFieldsRes @@ -358,6 +360,7 @@ export const postActionsConnectorExecuteRoute = ( esClient, esStore, isStream: request.body.subAction !== 'invokeAI', + kbDataClient, llmType: getLlmType(actionTypeId), langChainMessages, logger, @@ -374,7 +377,15 @@ export const postActionsConnectorExecuteRoute = ( logger, }), }, - }); + }; + + // New code path for LangGraph implementation, behind `assistantKnowledgeBaseByDefault` FF + let result: StreamResponseWithHeaders | StaticReturnType; + if (enableKnowledgeBaseByDefault) { + result = await callAssistantGraph(executorParams); + } else { + result = await callAgentExecutor(executorParams); + } telemetry.reportEvent(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, { actionTypeId, diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 3a392eaa8c256..f12bacde983df 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -212,8 +212,10 @@ export interface AssistantToolParams { isEnabledKnowledgeBase: boolean; chain?: RetrievalQAChain; esClient: ElasticsearchClient; + kbDataClient?: AIAssistantKnowledgeBaseDataClient; langChainTimeout?: number; llm?: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; + logger: Logger; modelExists: boolean; onNewReplacements?: (newReplacements: Replacements) => void; replacements?: Replacements; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index b99c1f6e0cd38..0e5ea3a8f69d1 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -11,10 +11,14 @@ import { ALERT_COUNTS_TOOL } from './alert_counts/alert_counts_tool'; import { ESQL_KNOWLEDGE_BASE_TOOL } from './esql_language_knowledge_base/esql_language_knowledge_base_tool'; import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool'; import { ATTACK_DISCOVERY_TOOL } from './attack_discovery/attack_discovery_tool'; +import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from './knowledge_base/knowledge_base_retrieval_tool'; +import { KNOWLEDGE_BASE_WRITE_TOOL } from './knowledge_base/knowledge_base_write_tool'; export const getAssistantTools = (): AssistantTool[] => [ ALERT_COUNTS_TOOL, ATTACK_DISCOVERY_TOOL, ESQL_KNOWLEDGE_BASE_TOOL, + KNOWLEDGE_BASE_RETRIEVAL_TOOL, + KNOWLEDGE_BASE_WRITE_TOOL, OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL, ]; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts new file mode 100644 index 0000000000000..c06898b1ce3b9 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts @@ -0,0 +1,53 @@ +/* + * 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 { DynamicStructuredTool } from '@langchain/core/tools'; +import { z } from 'zod'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { APP_UI_ID } from '../../../../common'; + +export type KnowledgeBaseRetrievalToolParams = AssistantToolParams; + +const toolDetails = { + description: + "Call this for fetching details from the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Call this function when the user asks for information about themself, like 'what is my favorite...' or 'using my saved....'. Input must always be the free-text query on a single line, with no other text. You are welcome to re-write the query to be a summary of items/things to search for in the knowledge base, as a vector search will be performed to return similar results when requested. If the results returned do not look relevant, disregard and tell the user you were unable to find the information they were looking for.", + id: 'knowledge-base-retrieval-tool', + name: 'KnowledgeBaseRetrievalTool', +}; +export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { + ...toolDetails, + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is KnowledgeBaseRetrievalToolParams => { + const { chain, isEnabledKnowledgeBase, modelExists } = params; + return isEnabledKnowledgeBase && modelExists && chain != null; + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + + const { chain, logger } = params as KnowledgeBaseRetrievalToolParams; + if (chain == null) return null; + + return new DynamicStructuredTool({ + name: toolDetails.name, + description: toolDetails.description, + schema: z.object({ + query: z.string().describe(`Summary of items/things to search for in the knowledge base`), + }), + func: async (input, _, cbManager) => { + logger.debug(`KnowledgeBaseRetrievalToolParams:input\n ${JSON.stringify(input, null, 2)}`); + const result = await chain.invoke( + { + query: input.query, + }, + cbManager + ); + return result.text; + }, + tags: ['knowledge-base'], + }); + }, +}; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts new file mode 100644 index 0000000000000..754955dd648ee --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts @@ -0,0 +1,64 @@ +/* + * 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 { DynamicStructuredTool } from '@langchain/core/tools'; +import { z } from 'zod'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import type { AIAssistantKnowledgeBaseDataClient } from '@kbn/elastic-assistant-plugin/server/ai_assistant_data_clients/knowledge_base'; +import type { KnowledgeBaseEntryCreateProps } from '@kbn/elastic-assistant-common'; +import { APP_UI_ID } from '../../../../common'; + +export interface KnowledgeBaseWriteToolParams extends AssistantToolParams { + kbDataClient: AIAssistantKnowledgeBaseDataClient; +} + +const toolDetails = { + description: + "Call this for writing details to the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Input will be the summarized knowledge base entry to store, with no other text.", + id: 'knowledge-base-write-tool', + name: 'KnowledgeBaseWriteTool', +}; +export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { + ...toolDetails, + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is KnowledgeBaseWriteToolParams => { + const { isEnabledKnowledgeBase, kbDataClient, modelExists } = params; + return isEnabledKnowledgeBase && modelExists && kbDataClient != null; + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + + const { kbDataClient, logger } = params as KnowledgeBaseWriteToolParams; + if (kbDataClient == null) return null; + + return new DynamicStructuredTool({ + name: toolDetails.name, + description: toolDetails.description, + schema: z.object({ + query: z.string().describe(`Summary of items/things to save in the knowledge base`), + }), + func: async (input, _, cbManager) => { + logger.debug(`KnowledgeBaseWriteToolParams:input\n ${JSON.stringify(input, null, 2)}`); + + const knowledgeBaseEntry: KnowledgeBaseEntryCreateProps = { + metadata: { kbResource: 'user', source: 'conversation', required: false }, + text: input.query, + }; + + logger.debug(`knowledgeBaseEntry\n ${JSON.stringify(knowledgeBaseEntry, null, 2)}`); + + const resp = await kbDataClient.createKnowledgeBaseEntry({ knowledgeBaseEntry }); + + if (resp == null) { + return "I'm sorry, but I was unable to add this entry to your knowledge base."; + } + return "I've successfully saved this entry to your knowledge base. You can ask me to recall this information at any time."; + }, + tags: ['knowledge-base'], + }); + }, +}; diff --git a/yarn.lock b/yarn.lock index 5d161d64b442e..7292b0fd0bcee 100644 --- a/yarn.lock +++ b/yarn.lock @@ -35,21 +35,6 @@ "@jridgewell/gen-mapping" "^0.1.0" "@jridgewell/trace-mapping" "^0.3.9" -"@anthropic-ai/sdk@^0.9.1": - version "0.9.1" - resolved "https://registry.yarnpkg.com/@anthropic-ai/sdk/-/sdk-0.9.1.tgz#b2d2b7bf05c90dce502c9a2e869066870f69ba88" - integrity sha512-wa1meQ2WSfoY8Uor3EdrJq0jTiZJoKoSii2ZVWRY1oN4Tlr5s59pADg9T79FTbPe1/se5c3pBeZgJL63wmuoBA== - dependencies: - "@types/node" "^18.11.18" - "@types/node-fetch" "^2.6.4" - abort-controller "^3.0.0" - agentkeepalive "^4.2.1" - digest-fetch "^1.3.0" - form-data-encoder "1.7.2" - formdata-node "^4.3.2" - node-fetch "^2.6.7" - web-streams-polyfill "^3.2.1" - "@apidevtools/json-schema-ref-parser@^9.0.6": version "9.0.9" resolved "https://registry.yarnpkg.com/@apidevtools/json-schema-ref-parser/-/json-schema-ref-parser-9.0.9.tgz#d720f9256e3609621280584f2b47ae165359268b" @@ -6743,48 +6728,68 @@ resolved "https://registry.yarnpkg.com/@kwsites/promise-deferred/-/promise-deferred-1.1.1.tgz#8ace5259254426ccef57f3175bc64ed7095ed919" integrity sha512-GaHYm+c0O9MjZRu0ongGBRbinu8gVAMd2UZjji6jVmqKtZluZnptXGWhz1E8j8D2HJ3f/yMxKAUC0b+57wncIw== -"@langchain/community@^0.0.44", "@langchain/community@~0.0.41": - version "0.0.44" - resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.0.44.tgz#b4f3453e3fd0b7a8c704fc35b004d7d738bd3416" - integrity sha512-II9Hz90jJmfWRICtxTg1auQWzFw0npqacWiiOpaxNhzs6rptdf56gyfC48Z6n1ii4R8FfAlfX6YxhOE7lGGKXg== +"@langchain/community@^0.2.2": + version "0.2.2" + resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.2.tgz#0ce7cd56ff8940fe73983f1853e165d334a2a446" + integrity sha512-TtlZnPBYt7Sujc1hAYvdZKUmV97wuF15O7b4nBX4lBfQeW38N0DwGbhqpitDbpaJqZ2s8DM4rjapECk0kIdAww== dependencies: - "@langchain/core" "~0.1.44" - "@langchain/openai" "~0.0.19" + "@langchain/core" "~0.2.0" + "@langchain/openai" "~0.0.28" + binary-extensions "^2.2.0" expr-eval "^2.0.2" flat "^5.0.2" + js-yaml "^4.1.0" + langchain "~0.2.0" langsmith "~0.1.1" uuid "^9.0.0" zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@0.1.53", "@langchain/core@^0.1.53", "@langchain/core@~0.1.44", "@langchain/core@~0.1.45": - version "0.1.53" - resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.1.53.tgz#40bf273b6d5e1426c60ce9cc259562fe656573f1" - integrity sha512-khfRTu2DSCNMPUmnKx7iH0TpEaunW/4BgR6STTteRRDd0NFtXGfAwUuY9sm0+EKi/XKhdAmpGnfLwSfNg5F0Qw== +"@langchain/core@0.2.0", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@^0.2.0", "@langchain/core@~0.2.0": + version "0.2.0" + resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.0.tgz#19c6374a5ad80daf8e14cb58582bc988109a1403" + integrity sha512-UbCJUp9eh2JXd9AW/vhPbTgtZoMgTqJgSan5Wf/EP27X8JM65lWdCOpJW+gHyBXvabbyrZz3/EGaptTUL5gutw== dependencies: ansi-styles "^5.0.0" camelcase "6" decamelize "1.2.0" - js-tiktoken "^1.0.8" + js-tiktoken "^1.0.12" langsmith "~0.1.7" ml-distance "^4.0.0" + mustache "^4.2.0" p-queue "^6.6.2" p-retry "4" uuid "^9.0.0" zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/openai@^0.0.25", "@langchain/openai@~0.0.19": - version "0.0.25" - resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.25.tgz#8332abea1e3acb9b1169f90636e518c0ee90622e" - integrity sha512-cD9xPDDXK2Cjs6yYg27BpdzBnQZvBb1yaNgMoGLWIT27UQVRyT96PLC1OVMQOmMmHaKDBCj/1bW4GQQgX7+d2Q== +"@langchain/langgraph@^0.0.20": + version "0.0.20" + resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.20.tgz#9229af1a79107916910fa65fe185bf66cda7736f" + integrity sha512-/byqz3WDbIQqaPDmC+Bo2n36LBpD42yj8wR7KiDZvrOIJSlMIoqwZeRkONEp9D7o61ZRaAMwoUJWriG8L9xdFg== dependencies: - "@langchain/core" "~0.1.45" - js-tiktoken "^1.0.7" - openai "^4.26.0" + "@langchain/core" ">0.1.61 <0.3.0" + uuid "^9.0.1" + +"@langchain/openai@^0.0.33", "@langchain/openai@~0.0.28": + version "0.0.33" + resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.33.tgz#af88d815ff0095018c879d3a1a5a32b2795b5c69" + integrity sha512-hTBo9y9bHtFvMT5ySBW7TrmKhLSA91iNahigeqAFBVrLmBDz+6rzzLFc1mpq6JEAR3fZKdaUXqso3nB23jIpTw== + dependencies: + "@langchain/core" ">0.1.56 <0.3.0" + js-tiktoken "^1.0.12" + openai "^4.41.1" zod "^3.22.4" zod-to-json-schema "^3.22.3" +"@langchain/textsplitters@~0.0.0": + version "0.0.2" + resolved "https://registry.yarnpkg.com/@langchain/textsplitters/-/textsplitters-0.0.2.tgz#500baa8341fb7fc86fca531a4192665a319504a3" + integrity sha512-6bQOuYHTGYlkgPY/8M5WPq4nnXZpEysGzRopQCYjg2WLcEoIPUMMrXsAaNNdvU3BOeMrhin8izvpDPD165hX6Q== + dependencies: + "@langchain/core" ">0.1.0 <0.3.0" + js-tiktoken "^1.0.12" + "@leichtgewicht/ip-codec@^2.0.1": version "2.0.4" resolved "https://registry.yarnpkg.com/@leichtgewicht/ip-codec/-/ip-codec-2.0.4.tgz#b2ac626d6cb9c8718ab459166d4bb405b8ffa78b" @@ -20881,10 +20886,10 @@ js-string-escape@^1.0.1: resolved "https://registry.yarnpkg.com/js-string-escape/-/js-string-escape-1.0.1.tgz#e2625badbc0d67c7533e9edc1068c587ae4137ef" integrity sha1-4mJbrbwNZ8dTPp7cEGjFh65BN+8= -js-tiktoken@^1.0.7, js-tiktoken@^1.0.8: - version "1.0.10" - resolved "https://registry.yarnpkg.com/js-tiktoken/-/js-tiktoken-1.0.10.tgz#2b343ec169399dcee8f9ef9807dbd4fafd3b30dc" - integrity sha512-ZoSxbGjvGyMT13x6ACo9ebhDha/0FHdKA+OsQcMOWcm1Zs7r90Rhk5lhERLzji+3rA7EKpXCgwXcM5fF3DMpdA== +js-tiktoken@^1.0.12: + version "1.0.12" + resolved "https://registry.yarnpkg.com/js-tiktoken/-/js-tiktoken-1.0.12.tgz#af0f5cf58e5e7318240d050c8413234019424211" + integrity sha512-L7wURW1fH9Qaext0VzaUDpFGVQgjkdE3Dgsy9/+yXyGEpBKnylTd0mU0bfbNkKDlXRb6TEsZkwuflu1B8uQbJQ== dependencies: base64-js "^1.5.1" @@ -21271,17 +21276,38 @@ kuler@^2.0.0: resolved "https://registry.yarnpkg.com/kuler/-/kuler-2.0.0.tgz#e2c570a3800388fb44407e851531c1d670b061b3" integrity sha512-Xq9nH7KlWZmXAtodXDDRE7vs6DU1gTU8zYDHDiWLSip45Egwq3plLHzPn27NgvzL2r1LMPC1vdqh98sQxtqj4A== -langchain@^0.1.30: - version "0.1.30" - resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.1.30.tgz#e1adb3f1849fcd5c596c668300afd5dc8cb37a97" - integrity sha512-5h/vNMmutQ98tbB0sPDlAileZVca6A2McFgGa3+D56Dm8mSSCzTQL2DngPA6h09DlKDpSr7+6PdFw5Hoj0ZDSw== +langchain@^0.2.2: + version "0.2.2" + resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.2.tgz#21605450458c77f022c88fdb7671bb82f4a9a17f" + integrity sha512-4tt2QuwW8AXdIL8CRkQeGOCoYYH3QbLHfQ09yD0iWLV1rwUYJ8mIYFAz/+u6CB8YNEyR/HI105s4xrxFQbWa9g== dependencies: - "@anthropic-ai/sdk" "^0.9.1" - "@langchain/community" "~0.0.41" - "@langchain/core" "~0.1.44" - "@langchain/openai" "~0.0.19" + "@langchain/core" "~0.2.0" + "@langchain/openai" "~0.0.28" + "@langchain/textsplitters" "~0.0.0" binary-extensions "^2.2.0" - js-tiktoken "^1.0.7" + js-tiktoken "^1.0.12" + js-yaml "^4.1.0" + jsonpointer "^5.0.1" + langchainhub "~0.0.8" + langsmith "~0.1.7" + ml-distance "^4.0.0" + openapi-types "^12.1.3" + p-retry "4" + uuid "^9.0.0" + yaml "^2.2.1" + zod "^3.22.4" + zod-to-json-schema "^3.22.3" + +langchain@~0.2.0: + version "0.2.1" + resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.1.tgz#eb21f76e27ab7220734d1632a85baae4d9ec935a" + integrity sha512-aCAsUwcmXjvhVsbAbR7NnzQ8jIjJPOx1EW4CmHX9Ggxp150EZYbLv7RJ5uBfj47hYUEFMfAqsCt524BwGnelng== + dependencies: + "@langchain/core" "~0.2.0" + "@langchain/openai" "~0.0.28" + "@langchain/textsplitters" "~0.0.0" + binary-extensions "^2.2.0" + js-tiktoken "^1.0.12" js-yaml "^4.1.0" jsonpointer "^5.0.1" langchainhub "~0.0.8" @@ -21299,7 +21325,18 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.14, langsmith@~0.1.1, langsmith@~0.1.7: +langsmith@^0.1.28: + version "0.1.28" + resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.28.tgz#fbe01352d0b993fd11d4085dd337b1cec17ef28d" + integrity sha512-IQUbo7I7rEE6QYBhrcgwqvlkcUsHlia0yTQpDwWdITw/VJx1f7gLPjNdbwWE+jvOZ4HcD7gCf2HR6zFXputu5A== + dependencies: + "@types/uuid" "^9.0.1" + commander "^10.0.1" + p-queue "^6.6.2" + p-retry "4" + uuid "^9.0.0" + +langsmith@~0.1.1, langsmith@~0.1.7: version "0.1.14" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.14.tgz#2b889dbcfb49547614df276a4a5a063092a1585d" integrity sha512-iEzQLLB7/0nRpAwNBAR7B7N64fyByg5UsNjSvLaCCkQ9AS68PSafjB8xQkyI8QXXrGjU1dEqDRoa8m4SUuRdUw== @@ -23064,6 +23101,11 @@ mustache@^2.3.2: resolved "https://registry.yarnpkg.com/mustache/-/mustache-2.3.2.tgz#a6d4d9c3f91d13359ab889a812954f9230a3d0c5" integrity sha512-KpMNwdQsYz3O/SBS1qJ/o3sqUJ5wSb8gb0pul8CO0S56b9Y2ALm8zCfsjPXsqGFfoNBkDwZuZIAjhsZI03gYVQ== +mustache@^4.2.0: + version "4.2.0" + resolved "https://registry.yarnpkg.com/mustache/-/mustache-4.2.0.tgz#e5892324d60a12ec9c2a73359edca52972bf6f64" + integrity sha512-71ippSywq5Yb7/tVYyGbkBggbU8H3u5Rz56fH60jGFgr8uHwxs+aSKeqmluIVzM0m0kB7xQjKS6qPfd0b2ZoqQ== + mutation-observer@^1.0.3: version "1.0.3" resolved "https://registry.yarnpkg.com/mutation-observer/-/mutation-observer-1.0.3.tgz#42e9222b101bca82e5ba9d5a7acf4a14c0f263d0" @@ -23848,7 +23890,7 @@ open@^8.0.9, open@^8.4.0, open@~8.4.0: is-docker "^2.1.1" is-wsl "^2.2.0" -openai@^4.24.1, openai@^4.26.0: +openai@^4.24.1: version "4.26.1" resolved "https://registry.yarnpkg.com/openai/-/openai-4.26.1.tgz#7b7c0225c09922445f68f3c4cdbd5775ed31108c" integrity sha512-DvWbjhWbappsFRatOWmu4Dp1/Q4RG9oOz6CfOSjy0/Drb8G+5iAiqWAO4PfpGIkhOOKtvvNfQri2SItl+U7LhQ== @@ -23863,6 +23905,20 @@ openai@^4.24.1, openai@^4.26.0: node-fetch "^2.6.7" web-streams-polyfill "^3.2.1" +openai@^4.41.1: + version "4.47.1" + resolved "https://registry.yarnpkg.com/openai/-/openai-4.47.1.tgz#1d23c7a8eb3d7bcdc69709cd905f4c9af0181dba" + integrity sha512-WWSxhC/69ZhYWxH/OBsLEirIjUcfpQ5+ihkXKp06hmeYXgBBIUCa9IptMzYx6NdkiOCsSGYCnTIsxaic3AjRCQ== + dependencies: + "@types/node" "^18.11.18" + "@types/node-fetch" "^2.6.4" + abort-controller "^3.0.0" + agentkeepalive "^4.2.1" + form-data-encoder "1.7.2" + formdata-node "^4.3.2" + node-fetch "^2.6.7" + web-streams-polyfill "^3.2.1" + openapi-sampler@^1.3.1: version "1.4.0" resolved "https://registry.yarnpkg.com/openapi-sampler/-/openapi-sampler-1.4.0.tgz#c133cad6250481f2ec7e48b16eb70062adb514c0" @@ -30613,6 +30669,11 @@ uuid@^8.0.0, uuid@^8.3.0, uuid@^8.3.2: resolved "https://registry.yarnpkg.com/uuid/-/uuid-8.3.2.tgz#80d5b5ced271bb9af6c445f21a1a04c606cefbe2" integrity sha512-+NYs2QeMWy+GWFOEm9xnn6HCDp0l7QBD7ml8zLUmJ+93Q5NF0NocErnwkTkXVFNiX3/fpC6afS8Dhb/gz7R7eg== +uuid@^9.0.1: + version "9.0.1" + resolved "https://registry.yarnpkg.com/uuid/-/uuid-9.0.1.tgz#e188d4c8853cc722220392c424cd637f32293f30" + integrity sha512-b+1eJOlsR9K8HJpow9Ok3fiWOWSIcIzXodvv0rQjVoOVNpWMpxf1wZNpt4y9h10odCNrqnYp1OBzRktckBe3sA== + v8-compile-cache-lib@^3.0.1: version "3.0.1" resolved "https://registry.yarnpkg.com/v8-compile-cache-lib/-/v8-compile-cache-lib-3.0.1.tgz#6336e8d71965cb3d35a1bbb7868445a7c05264bf" From de9d53064b7540c12e455598594ece5d2b17377c Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 30 May 2024 20:00:18 -0600 Subject: [PATCH 02/47] yarn.lock update --- yarn.lock | 71 ++++--------------------------------------------------- 1 file changed, 5 insertions(+), 66 deletions(-) diff --git a/yarn.lock b/yarn.lock index adf5f053ea7a5..53746c707bd64 100644 --- a/yarn.lock +++ b/yarn.lock @@ -12579,11 +12579,6 @@ bare-path@^2.0.0, bare-path@^2.1.0: dependencies: bare-os "^2.1.0" -base-64@^0.1.0: - version "0.1.0" - resolved "https://registry.yarnpkg.com/base-64/-/base-64-0.1.0.tgz#780a99c84e7d600260361511c4877613bf24f6bb" - integrity sha512-Y5gU45svrR5tI2Vt/X9GPd3L0HNIKzGu202EjxrXMpuc2V2CiKgemAbUUsqYmZJvPtCXoUKjNZwBJzsNScUbXA== - base64-js@1.3.1: version "1.3.1" resolved "https://registry.yarnpkg.com/base64-js/-/base64-js-1.3.1.tgz#58ece8cb75dd07e71ed08c736abc5fac4dbf8df1" @@ -15598,14 +15593,6 @@ diffie-hellman@^5.0.0: miller-rabin "^4.0.0" randombytes "^2.0.0" -digest-fetch@^1.3.0: - version "1.3.0" - resolved "https://registry.yarnpkg.com/digest-fetch/-/digest-fetch-1.3.0.tgz#898e69264d00012a23cf26e8a3e40320143fc661" - integrity sha512-CGJuv6iKNM7QyZlM2T3sPAdZWd/p9zQiRNS9G+9COUCwzWFTs0Xp8NF5iePx7wtvhDykReiRRrSeNb4oMmB8lA== - dependencies: - base-64 "^0.1.0" - md5 "^2.3.0" - dir-glob@^2.2.2: version "2.2.2" resolved "https://registry.yarnpkg.com/dir-glob/-/dir-glob-2.2.2.tgz#fa09f0694153c8918b18ba0deafae94769fc50c4" @@ -21280,7 +21267,7 @@ kuler@^2.0.0: resolved "https://registry.yarnpkg.com/kuler/-/kuler-2.0.0.tgz#e2c570a3800388fb44407e851531c1d670b061b3" integrity sha512-Xq9nH7KlWZmXAtodXDDRE7vs6DU1gTU8zYDHDiWLSip45Egwq3plLHzPn27NgvzL2r1LMPC1vdqh98sQxtqj4A== -langchain@^0.2.2: +langchain@^0.2.2, langchain@~0.2.0: version "0.2.2" resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.2.tgz#21605450458c77f022c88fdb7671bb82f4a9a17f" integrity sha512-4tt2QuwW8AXdIL8CRkQeGOCoYYH3QbLHfQ09yD0iWLV1rwUYJ8mIYFAz/+u6CB8YNEyR/HI105s4xrxFQbWa9g== @@ -21302,34 +21289,12 @@ langchain@^0.2.2: zod "^3.22.4" zod-to-json-schema "^3.22.3" -langchain@~0.2.0: - version "0.2.1" - resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.1.tgz#eb21f76e27ab7220734d1632a85baae4d9ec935a" - integrity sha512-aCAsUwcmXjvhVsbAbR7NnzQ8jIjJPOx1EW4CmHX9Ggxp150EZYbLv7RJ5uBfj47hYUEFMfAqsCt524BwGnelng== - dependencies: - "@langchain/core" "~0.2.0" - "@langchain/openai" "~0.0.28" - "@langchain/textsplitters" "~0.0.0" - binary-extensions "^2.2.0" - js-tiktoken "^1.0.12" - js-yaml "^4.1.0" - jsonpointer "^5.0.1" - langchainhub "~0.0.8" - langsmith "~0.1.7" - ml-distance "^4.0.0" - openapi-types "^12.1.3" - p-retry "4" - uuid "^9.0.0" - yaml "^2.2.1" - zod "^3.22.4" - zod-to-json-schema "^3.22.3" - langchainhub@~0.0.8: version "0.0.8" resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.28: +langsmith@^0.1.28, langsmith@~0.1.1, langsmith@~0.1.7: version "0.1.28" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.28.tgz#fbe01352d0b993fd11d4085dd337b1cec17ef28d" integrity sha512-IQUbo7I7rEE6QYBhrcgwqvlkcUsHlia0yTQpDwWdITw/VJx1f7gLPjNdbwWE+jvOZ4HcD7gCf2HR6zFXputu5A== @@ -21340,17 +21305,6 @@ langsmith@^0.1.28: p-retry "4" uuid "^9.0.0" -langsmith@~0.1.1, langsmith@~0.1.7: - version "0.1.14" - resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.14.tgz#2b889dbcfb49547614df276a4a5a063092a1585d" - integrity sha512-iEzQLLB7/0nRpAwNBAR7B7N64fyByg5UsNjSvLaCCkQ9AS68PSafjB8xQkyI8QXXrGjU1dEqDRoa8m4SUuRdUw== - dependencies: - "@types/uuid" "^9.0.1" - commander "^10.0.1" - p-queue "^6.6.2" - p-retry "4" - uuid "^9.0.0" - language-subtag-registry@~0.3.2: version "0.3.21" resolved "https://registry.yarnpkg.com/language-subtag-registry/-/language-subtag-registry-0.3.21.tgz#04ac218bea46f04cb039084602c6da9e788dd45a" @@ -23894,22 +23848,7 @@ open@^8.0.9, open@^8.4.0, open@~8.4.0: is-docker "^2.1.1" is-wsl "^2.2.0" -openai@^4.24.1: - version "4.26.1" - resolved "https://registry.yarnpkg.com/openai/-/openai-4.26.1.tgz#7b7c0225c09922445f68f3c4cdbd5775ed31108c" - integrity sha512-DvWbjhWbappsFRatOWmu4Dp1/Q4RG9oOz6CfOSjy0/Drb8G+5iAiqWAO4PfpGIkhOOKtvvNfQri2SItl+U7LhQ== - dependencies: - "@types/node" "^18.11.18" - "@types/node-fetch" "^2.6.4" - abort-controller "^3.0.0" - agentkeepalive "^4.2.1" - digest-fetch "^1.3.0" - form-data-encoder "1.7.2" - formdata-node "^4.3.2" - node-fetch "^2.6.7" - web-streams-polyfill "^3.2.1" - -openai@^4.41.1: +openai@^4.24.1, openai@^4.41.1: version "4.47.1" resolved "https://registry.yarnpkg.com/openai/-/openai-4.47.1.tgz#1d23c7a8eb3d7bcdc69709cd905f4c9af0181dba" integrity sha512-WWSxhC/69ZhYWxH/OBsLEirIjUcfpQ5+ihkXKp06hmeYXgBBIUCa9IptMzYx6NdkiOCsSGYCnTIsxaic3AjRCQ== @@ -30658,7 +30597,7 @@ uuid-browser@^3.1.0: resolved "https://registry.yarnpkg.com/uuid-browser/-/uuid-browser-3.1.0.tgz#0f05a40aef74f9e5951e20efbf44b11871e56410" integrity sha1-DwWkCu90+eWVHiDvv0SxGHHlZBA= -uuid@9.0.0, uuid@^9, uuid@^9.0.0: +uuid@9.0.0: version "9.0.0" resolved "https://registry.yarnpkg.com/uuid/-/uuid-9.0.0.tgz#592f550650024a38ceb0c562f2f6aa435761efb5" integrity sha512-MXcSTerfPa4uqyzStbRoTgt5XIe3x5+42+q1sDuy3R5MDk66URdLMOZe5aPX/SQd+kuYAh0FdP/pO28IkQyTeg== @@ -30673,7 +30612,7 @@ uuid@^8.0.0, uuid@^8.3.0, uuid@^8.3.2: resolved "https://registry.yarnpkg.com/uuid/-/uuid-8.3.2.tgz#80d5b5ced271bb9af6c445f21a1a04c606cefbe2" integrity sha512-+NYs2QeMWy+GWFOEm9xnn6HCDp0l7QBD7ml8zLUmJ+93Q5NF0NocErnwkTkXVFNiX3/fpC6afS8Dhb/gz7R7eg== -uuid@^9.0.1: +uuid@^9, uuid@^9.0.0, uuid@^9.0.1: version "9.0.1" resolved "https://registry.yarnpkg.com/uuid/-/uuid-9.0.1.tgz#e188d4c8853cc722220392c424cd637f32293f30" integrity sha512-b+1eJOlsR9K8HJpow9Ok3fiWOWSIcIzXodvv0rQjVoOVNpWMpxf1wZNpt4y9h10odCNrqnYp1OBzRktckBe3sA== From b862a291f1542369fc10192855194f3c9e637774 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 31 May 2024 17:39:16 -0600 Subject: [PATCH 03/47] Port title generation into langgraph node --- package.json | 14 ++--- .../execute_custom_llm_chain/index.ts | 2 - .../server/lib/langchain/executors/types.ts | 14 ++++- .../generate_chat_title.ts | 54 +++++++++++++++++ .../graphs/default_assistant_graph/graph.ts | 27 ++++++++- .../graphs/default_assistant_graph/index.ts | 18 ++++-- .../graphs/default_assistant_graph/prompts.ts | 6 +- .../graphs/default_assistant_graph/types.ts | 7 +++ .../routes/post_actions_connector_execute.ts | 18 +++++- yarn.lock | 60 +++++++++++++------ 10 files changed, 178 insertions(+), 42 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/generate_chat_title.ts diff --git a/package.json b/package.json index 23aae409cdcfe..f1fe4bdb88532 100644 --- a/package.json +++ b/package.json @@ -80,7 +80,7 @@ "resolutions": { "**/@bazel/typescript/protobufjs": "6.11.4", "**/@hello-pangea/dnd": "16.6.0", - "**/@langchain/core": "0.2.0", + "**/@langchain/core": "0.2.3", "**/@types/node": "20.10.5", "**/@typescript-eslint/utils": "5.62.0", "**/chokidar": "^3.5.3", @@ -913,10 +913,10 @@ "@kbn/watcher-plugin": "link:x-pack/plugins/watcher", "@kbn/xstate-utils": "link:packages/kbn-xstate-utils", "@kbn/zod-helpers": "link:packages/kbn-zod-helpers", - "@langchain/community": "^0.2.2", - "@langchain/core": "^0.2.0", - "@langchain/langgraph": "^0.0.20", - "@langchain/openai": "^0.0.33", + "@langchain/community": "^0.2.4", + "@langchain/core": "^0.2.3", + "@langchain/langgraph": "^0.0.21", + "@langchain/openai": "^0.0.34", "@loaders.gl/core": "^3.4.7", "@loaders.gl/json": "^3.4.7", "@loaders.gl/shapefile": "^3.4.7", @@ -1053,8 +1053,8 @@ "jsonwebtoken": "^9.0.2", "jsts": "^1.6.2", "kea": "^2.6.0", - "langchain": "^0.2.2", - "langsmith": "^0.1.28", + "langchain": "^0.2.3", + "langsmith": "^0.1.30", "launchdarkly-js-client-sdk": "^3.1.4", "launchdarkly-node-server-sdk": "^7.0.3", "load-json-file": "^6.2.0", diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts index 555336b301c31..7ea1e5fb3c9b9 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts @@ -38,7 +38,6 @@ export const callAgentExecutor: AgentExecutor = async ({ connectorId, esClient, esStore, - kbDataClient, langChainMessages, llmType, logger, @@ -95,7 +94,6 @@ export const callAgentExecutor: AgentExecutor = async ({ chain, esClient, isEnabledKnowledgeBase, - kbDataClient, llm, logger, modelExists, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index 56cceb3448a9d..bd07099e312b3 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -9,7 +9,7 @@ import { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/s import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { BaseMessage } from '@langchain/core/messages'; import { Logger } from '@kbn/logging'; -import { KibanaRequest, ResponseHeaders } from '@kbn/core-http-server'; +import { KibanaRequest, KibanaResponseFactory, ResponseHeaders } from '@kbn/core-http-server'; import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain'; import { ExecuteConnectorRequestBody, Message, Replacements } from '@kbn/elastic-assistant-common'; import { StreamResponseWithHeaders } from '@kbn/ml-response-stream/server'; @@ -18,6 +18,8 @@ import { ResponseBody } from '../types'; import type { AssistantTool } from '../../../types'; import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store'; import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base'; +import { AIAssistantConversationsDataClient } from '../../../ai_assistant_data_clients/conversations'; +import { AIAssistantDataClient } from '../../../ai_assistant_data_clients'; export type OnLlmResponse = ( content: string, @@ -25,6 +27,12 @@ export type OnLlmResponse = ( isError?: boolean ) => Promise; +export interface AssistantDataClients { + anonymizationFieldsDataClient?: AIAssistantDataClient; + conversationsDataClient?: AIAssistantConversationsDataClient; + kbDataClient?: AIAssistantKnowledgeBaseDataClient; +} + export interface AgentExecutorParams { abortSignal?: AbortSignal; alertsIndexPattern?: string; @@ -33,9 +41,10 @@ export interface AgentExecutorParams { isEnabledKnowledgeBase: boolean; assistantTools?: AssistantTool[]; connectorId: string; + conversationId?: string; + dataClients?: AssistantDataClients; esClient: ElasticsearchClient; esStore: ElasticsearchStore; - kbDataClient?: AIAssistantKnowledgeBaseDataClient | undefined; langChainMessages: BaseMessage[]; llmType?: string; logger: Logger; @@ -44,6 +53,7 @@ export interface AgentExecutorParams { isStream?: T; onLlmResponse?: OnLlmResponse; request: KibanaRequest; + response?: KibanaResponseFactory; size?: number; traceOptions?: TraceOptions; } diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/generate_chat_title.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/generate_chat_title.ts new file mode 100644 index 0000000000000..9cdf8fdb96943 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/generate_chat_title.ts @@ -0,0 +1,54 @@ +/* + * 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 { StringOutputParser } from '@langchain/core/output_parsers'; + +import { ChatPromptTemplate } from '@langchain/core/prompts'; +import { AgentState, NodeParamsBase } from './types'; +import { AIAssistantConversationsDataClient } from '../../../../ai_assistant_data_clients/conversations'; + +export const GENERATE_CHAT_TITLE_PROMPT = 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. 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 { + conversationsDataClient?: AIAssistantConversationsDataClient; + conversationId?: string; + state: AgentState; +} +export const generateChatTitle = async ({ + conversationsDataClient, + logger, + model, + state, +}: GenerateChatTitleParams) => { + logger.debug('node:generateChatTitle'); + logger.debug(`state:\n ${JSON.stringify(state, null, 2)}`); + if (state.messages.length !== 0) { + logger.debug('No need to generate chat title, messages already exist'); + return; + } + const outputParser = new StringOutputParser(); + const graph = GENERATE_CHAT_TITLE_PROMPT.pipe(model).pipe(outputParser); + + const chatTitle = await graph.invoke({ + input: JSON.stringify(state.input, null, 2), + }); + + logger.debug(`chatTitle: ${chatTitle}`); + + return { + chatTitle, + }; +}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts index e70df0b506d25..32ccf28e14da0 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts @@ -15,12 +15,18 @@ import type { CompiledStateGraph } from '@langchain/langgraph/dist/graph/state'; import type { Logger } from '@kbn/logging'; import { BaseMessage } from '@langchain/core/messages'; -import { AgentState } from './types'; +import { BaseChatModel } from '@langchain/core/language_models/chat_models'; +import { AgentState, NodeParamsBase } from './types'; +import { generateChatTitle } from './generate_chat_title'; +import { AssistantDataClients } from '../../executors/types'; export const DEFAULT_ASSISTANT_GRAPH_ID = 'Default Security Assistant Graph'; interface GetDefaultAssistantGraphParams { agentRunnable: AgentRunnableSequence; + dataClients?: AssistantDataClients; + conversationId?: string; + llm: BaseChatModel; logger: Logger; messages: BaseMessage[]; tools: StructuredTool[]; @@ -31,6 +37,9 @@ interface GetDefaultAssistantGraphParams { */ export const getDefaultAssistantGraph = ({ agentRunnable, + conversationId, + dataClients, + llm, logger, messages, tools, @@ -59,6 +68,11 @@ export const getDefaultAssistantGraph = ({ }, }; + const nodeParams: NodeParamsBase = { + model: llm, + logger, + }; + // Create a tool executor const toolExecutor = new ToolExecutor({ tools }); @@ -103,6 +117,14 @@ export const getDefaultAssistantGraph = ({ const workflow = new StateGraph({ channels: graphState }); // Define the nodes to cycle between + workflow.addNode('generateChatTitle', (state: AgentState) => + generateChatTitle({ + state, + conversationsDataClient: dataClients?.conversationsDataClient, + conversationId, + ...nodeParams, + }) + ); workflow.addNode('agent', new RunnableLambda({ func: runAgent })); workflow.addNode('action', new RunnableLambda({ func: executeTools })); @@ -110,7 +132,8 @@ export const getDefaultAssistantGraph = ({ workflow.addConditionalEdges('agent', shouldContinue, { continue: 'action', end: END }); // Add edges for start, and between agent and action (action always followed by agent) - workflow.addEdge(START, 'agent'); + workflow.addEdge(START, 'generateChatTitle'); + workflow.addEdge('generateChatTitle', 'agent'); workflow.addEdge('action', 'agent'); return workflow.compile(); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts index 750f0b4d971cb..1e40f6b2fe127 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts @@ -21,8 +21,7 @@ import { getDefaultAssistantGraph } from './graph'; import { invokeGraph, streamGraph } from './helpers'; /** - * - * + * Drop in replacement for the existing `callAgentExecutor` that uses LangGraph */ export const callAssistantGraph: AgentExecutor = async ({ abortSignal, @@ -32,9 +31,10 @@ export const callAssistantGraph: AgentExecutor = async ({ isEnabledKnowledgeBase, assistantTools = [], connectorId, + conversationId, + dataClients, esClient, esStore, - kbDataClient, langChainMessages, llmType, logger: parentLogger, @@ -85,7 +85,7 @@ export const callAssistantGraph: AgentExecutor = async ({ chain, esClient, isEnabledKnowledgeBase, - kbDataClient, + kbDataClient: dataClients?.kbDataClient, llm: model, logger, modelExists, @@ -115,7 +115,15 @@ export const callAssistantGraph: AgentExecutor = async ({ const apmTracer = new APMTracer({ projectName: traceOptions?.projectName ?? 'default' }, logger); - const assistantGraph = getDefaultAssistantGraph({ agentRunnable, logger, messages, tools }); + const assistantGraph = getDefaultAssistantGraph({ + agentRunnable, + conversationId, + dataClients, + llm, + logger, + messages, + tools, + }); const inputs = { input: latestMessage[0].content as string }; if (isStream) { 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 1b6d8f5601c05..7d8a78f7387ec 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 @@ -13,7 +13,6 @@ export const openAIFunctionAgentPrompt = ChatPromptTemplate.fromMessages([ 'You are a helpful assistant\n\nUse the below context as a sample of information about the user from their knowledge base:\n\n```{knowledge_history}```', ], ['placeholder', '{chat_history}'], - // ['human', '{knowledge_history}'], ['human', '{input}'], ['placeholder', '{agent_scratchpad}'], ]); @@ -51,5 +50,8 @@ export const structuredChatAgentPrompt = ChatPromptTemplate.fromMessages([ 'Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation', ], ['placeholder', '{chat_history}'], - ['human', '{input}\n\n{agent_scratchpad}\n(reminder to respond in a JSON blob no matter what)'], + [ + 'human', + 'Use the below context as a sample of information about the user from their knowledge base:\n\n```\n{knowledge_history}\n```\n\n{input}\n\n{agent_scratchpad}\n(reminder to respond in a JSON blob no matter what)', + ], ]); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts index 6c7543f9e7e28..1d19646fb6eb3 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts @@ -7,6 +7,8 @@ import { BaseMessage } from '@langchain/core/messages'; import { AgentAction, AgentFinish, AgentStep } from '@langchain/core/agents'; +import { BaseChatModel } from '@langchain/core/language_models/chat_models'; +import type { Logger } from '@kbn/logging'; export interface AgentStateBase { agentOutcome?: AgentAction | AgentFinish; @@ -17,3 +19,8 @@ export interface AgentState extends AgentStateBase { input: string; messages: BaseMessage[]; } + +export interface NodeParamsBase { + logger: Logger; + model: BaseChatModel; +} diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 4e46ad16e7d2e..1ab0ef84c2574 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -21,7 +21,11 @@ import { import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { i18n } from '@kbn/i18n'; import { getLlmType } from './utils'; -import { StaticReturnType } from '../lib/langchain/executors/types'; +import { + AgentExecutorParams, + AssistantDataClients, + StaticReturnType, +} from '../lib/langchain/executors/types'; import { INVOKE_ASSISTANT_ERROR_EVENT, INVOKE_ASSISTANT_SUCCESS_EVENT, @@ -346,8 +350,14 @@ export const postActionsConnectorExecuteRoute = ( kbDataClient ); + const dataClients: AssistantDataClients = { + anonymizationFieldsDataClient: anonymizationFieldsDataClient ?? undefined, + conversationsDataClient: conversationsDataClient ?? undefined, + kbDataClient, + }; + // Shared executor params - const executorParams = { + const executorParams: AgentExecutorParams = { abortSignal, alertsIndexPattern: request.body.alertsIndexPattern, anonymizationFields: anonymizationFieldsRes @@ -357,16 +367,18 @@ export const postActionsConnectorExecuteRoute = ( isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase ?? false, assistantTools, connectorId, + conversationId, + dataClients, esClient, esStore, isStream: request.body.subAction !== 'invokeAI', - kbDataClient, llmType: getLlmType(actionTypeId), langChainMessages, logger, onNewReplacements, onLlmResponse, request, + response, replacements: request.body.replacements, size: request.body.size, traceOptions: { diff --git a/yarn.lock b/yarn.lock index 53746c707bd64..b8f012f0d144e 100644 --- a/yarn.lock +++ b/yarn.lock @@ -6732,10 +6732,10 @@ resolved "https://registry.yarnpkg.com/@kwsites/promise-deferred/-/promise-deferred-1.1.1.tgz#8ace5259254426ccef57f3175bc64ed7095ed919" integrity sha512-GaHYm+c0O9MjZRu0ongGBRbinu8gVAMd2UZjji6jVmqKtZluZnptXGWhz1E8j8D2HJ3f/yMxKAUC0b+57wncIw== -"@langchain/community@^0.2.2": - version "0.2.2" - resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.2.tgz#0ce7cd56ff8940fe73983f1853e165d334a2a446" - integrity sha512-TtlZnPBYt7Sujc1hAYvdZKUmV97wuF15O7b4nBX4lBfQeW38N0DwGbhqpitDbpaJqZ2s8DM4rjapECk0kIdAww== +"@langchain/community@^0.2.4": + version "0.2.4" + resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.4.tgz#fb5feb4f4a01a1b33adfd28ce7126d0dedb3e6d1" + integrity sha512-rwrPNQLyIe84TPqPYbYOfDA4G/ba1rdj7OtZg63dQmxIvNDOmUCh4xIQac2iuRUnM3o4Ben0Faa9qz+V5oPgIA== dependencies: "@langchain/core" "~0.2.0" "@langchain/openai" "~0.0.28" @@ -6743,16 +6743,16 @@ expr-eval "^2.0.2" flat "^5.0.2" js-yaml "^4.1.0" - langchain "~0.2.0" + langchain "0.2.3" langsmith "~0.1.1" uuid "^9.0.0" zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@0.2.0", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@^0.2.0", "@langchain/core@~0.2.0": - version "0.2.0" - resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.0.tgz#19c6374a5ad80daf8e14cb58582bc988109a1403" - integrity sha512-UbCJUp9eh2JXd9AW/vhPbTgtZoMgTqJgSan5Wf/EP27X8JM65lWdCOpJW+gHyBXvabbyrZz3/EGaptTUL5gutw== +"@langchain/core@0.2.3", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@~0.2.0": + version "0.2.3" + resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.3.tgz#7faa82f92b0c7843506e827a38bfcbb60f009d13" + integrity sha512-mVuFHSLpPQ4yOHNXeoSA3LnmIMuFmUiit5rvbYcPZqM6SrB2zCNN2nD4Ty5+3H5X4tYItDoSqsTuUNUQySXRQw== dependencies: ansi-styles "^5.0.0" camelcase "6" @@ -6767,15 +6767,26 @@ zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/langgraph@^0.0.20": - version "0.0.20" - resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.20.tgz#9229af1a79107916910fa65fe185bf66cda7736f" - integrity sha512-/byqz3WDbIQqaPDmC+Bo2n36LBpD42yj8wR7KiDZvrOIJSlMIoqwZeRkONEp9D7o61ZRaAMwoUJWriG8L9xdFg== +"@langchain/langgraph@^0.0.21": + version "0.0.21" + resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.21.tgz#5037597a954abad9ed5f0a1742226f5fcf27e7d7" + integrity sha512-7jtVZFAwvxSbIribYNzGXYIRrsAXV7YF4u1Xcpd8MYNz8sD3h8+rpIOJcYF1AdFh6laajnz0Gv8abPBHHQ2QiQ== dependencies: "@langchain/core" ">0.1.61 <0.3.0" uuid "^9.0.1" -"@langchain/openai@^0.0.33", "@langchain/openai@~0.0.28": +"@langchain/openai@^0.0.34": + version "0.0.34" + resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.34.tgz#36c9bca0721ab9f7e5d40927e7c0429cacbd5b56" + integrity sha512-M+CW4oXle5fdoz2T2SwdOef8pl3/1XmUx1vjn2mXUVM/128aO0l23FMF0SNBsAbRV6P+p/TuzjodchJbi0Ht/A== + dependencies: + "@langchain/core" ">0.1.56 <0.3.0" + js-tiktoken "^1.0.12" + openai "^4.41.1" + zod "^3.22.4" + zod-to-json-schema "^3.22.3" + +"@langchain/openai@~0.0.28": version "0.0.33" resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.33.tgz#af88d815ff0095018c879d3a1a5a32b2795b5c69" integrity sha512-hTBo9y9bHtFvMT5ySBW7TrmKhLSA91iNahigeqAFBVrLmBDz+6rzzLFc1mpq6JEAR3fZKdaUXqso3nB23jIpTw== @@ -21267,10 +21278,10 @@ kuler@^2.0.0: resolved "https://registry.yarnpkg.com/kuler/-/kuler-2.0.0.tgz#e2c570a3800388fb44407e851531c1d670b061b3" integrity sha512-Xq9nH7KlWZmXAtodXDDRE7vs6DU1gTU8zYDHDiWLSip45Egwq3plLHzPn27NgvzL2r1LMPC1vdqh98sQxtqj4A== -langchain@^0.2.2, langchain@~0.2.0: - version "0.2.2" - resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.2.tgz#21605450458c77f022c88fdb7671bb82f4a9a17f" - integrity sha512-4tt2QuwW8AXdIL8CRkQeGOCoYYH3QbLHfQ09yD0iWLV1rwUYJ8mIYFAz/+u6CB8YNEyR/HI105s4xrxFQbWa9g== +langchain@0.2.3: + version "0.2.3" + resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.3.tgz#c14bb05cf871b21bd63b84b3ab89580b1d62539f" + integrity sha512-T9xR7zd+Nj0oXy6WoYKmZLy0DlQiDLFPGYWdOXDxy+AvqlujoPdVQgDSpdqiOHvAjezrByAoKxoHCz5XMwTP/Q== dependencies: "@langchain/core" "~0.2.0" "@langchain/openai" "~0.0.28" @@ -21294,7 +21305,18 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.28, langsmith@~0.1.1, langsmith@~0.1.7: +langsmith@^0.1.30: + version "0.1.30" + resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" + integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== + dependencies: + "@types/uuid" "^9.0.1" + commander "^10.0.1" + p-queue "^6.6.2" + p-retry "4" + uuid "^9.0.0" + +langsmith@~0.1.1, langsmith@~0.1.7: version "0.1.28" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.28.tgz#fbe01352d0b993fd11d4085dd337b1cec17ef28d" integrity sha512-IQUbo7I7rEE6QYBhrcgwqvlkcUsHlia0yTQpDwWdITw/VJx1f7gLPjNdbwWE+jvOZ4HcD7gCf2HR6zFXputu5A== From 64511427cec870974f8380ca6c4455d14932081e Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Mon, 3 Jun 2024 09:31:24 -0600 Subject: [PATCH 04/47] yarn.lock update --- yarn.lock | 217 +++++++++++++++++++++++++++++++++++------------------- 1 file changed, 141 insertions(+), 76 deletions(-) diff --git a/yarn.lock b/yarn.lock index 510b9154486a3..ff76fc98269a9 100644 --- a/yarn.lock +++ b/yarn.lock @@ -35,21 +35,6 @@ "@jridgewell/gen-mapping" "^0.1.0" "@jridgewell/trace-mapping" "^0.3.9" -"@anthropic-ai/sdk@^0.9.1": - version "0.9.1" - resolved "https://registry.yarnpkg.com/@anthropic-ai/sdk/-/sdk-0.9.1.tgz#b2d2b7bf05c90dce502c9a2e869066870f69ba88" - integrity sha512-wa1meQ2WSfoY8Uor3EdrJq0jTiZJoKoSii2ZVWRY1oN4Tlr5s59pADg9T79FTbPe1/se5c3pBeZgJL63wmuoBA== - dependencies: - "@types/node" "^18.11.18" - "@types/node-fetch" "^2.6.4" - abort-controller "^3.0.0" - agentkeepalive "^4.2.1" - digest-fetch "^1.3.0" - form-data-encoder "1.7.2" - formdata-node "^4.3.2" - node-fetch "^2.6.7" - web-streams-polyfill "^3.2.1" - "@apidevtools/json-schema-ref-parser@^9.0.6": version "9.0.9" resolved "https://registry.yarnpkg.com/@apidevtools/json-schema-ref-parser/-/json-schema-ref-parser-9.0.9.tgz#d720f9256e3609621280584f2b47ae165359268b" @@ -6880,48 +6865,97 @@ resolved "https://registry.yarnpkg.com/@kwsites/promise-deferred/-/promise-deferred-1.1.1.tgz#8ace5259254426ccef57f3175bc64ed7095ed919" integrity sha512-GaHYm+c0O9MjZRu0ongGBRbinu8gVAMd2UZjji6jVmqKtZluZnptXGWhz1E8j8D2HJ3f/yMxKAUC0b+57wncIw== -"@langchain/community@^0.0.44", "@langchain/community@~0.0.41": - version "0.0.44" - resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.0.44.tgz#b4f3453e3fd0b7a8c704fc35b004d7d738bd3416" - integrity sha512-II9Hz90jJmfWRICtxTg1auQWzFw0npqacWiiOpaxNhzs6rptdf56gyfC48Z6n1ii4R8FfAlfX6YxhOE7lGGKXg== +"@langchain/community@^0.2.4": + version "0.2.4" + resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.4.tgz#fb5feb4f4a01a1b33adfd28ce7126d0dedb3e6d1" + integrity sha512-rwrPNQLyIe84TPqPYbYOfDA4G/ba1rdj7OtZg63dQmxIvNDOmUCh4xIQac2iuRUnM3o4Ben0Faa9qz+V5oPgIA== dependencies: - "@langchain/core" "~0.1.44" - "@langchain/openai" "~0.0.19" + "@langchain/core" "~0.2.0" + "@langchain/openai" "~0.0.28" + binary-extensions "^2.2.0" expr-eval "^2.0.2" flat "^5.0.2" + js-yaml "^4.1.0" + langchain "0.2.3" langsmith "~0.1.1" uuid "^9.0.0" zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@0.1.53", "@langchain/core@^0.1.53", "@langchain/core@~0.1.44", "@langchain/core@~0.1.45": - version "0.1.53" - resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.1.53.tgz#40bf273b6d5e1426c60ce9cc259562fe656573f1" - integrity sha512-khfRTu2DSCNMPUmnKx7iH0TpEaunW/4BgR6STTteRRDd0NFtXGfAwUuY9sm0+EKi/XKhdAmpGnfLwSfNg5F0Qw== +"@langchain/core@0.2.3", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@>=0.2.5 <0.3.0", "@langchain/core@~0.2.0": + version "0.2.3" + resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.3.tgz#7faa82f92b0c7843506e827a38bfcbb60f009d13" + integrity sha512-mVuFHSLpPQ4yOHNXeoSA3LnmIMuFmUiit5rvbYcPZqM6SrB2zCNN2nD4Ty5+3H5X4tYItDoSqsTuUNUQySXRQw== dependencies: ansi-styles "^5.0.0" camelcase "6" decamelize "1.2.0" - js-tiktoken "^1.0.8" + js-tiktoken "^1.0.12" langsmith "~0.1.7" ml-distance "^4.0.0" + mustache "^4.2.0" + p-queue "^6.6.2" + p-retry "4" + uuid "^9.0.0" + zod "^3.22.4" + zod-to-json-schema "^3.22.3" + +"@langchain/core@^0.2.3": + version "0.2.5" + resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.5.tgz#4983469d04615f7f66ae573319939ed7301f6889" + integrity sha512-tMaKRFVewFn8crQwlbXGjT7hlMdX1yXHap1ebBx7Bb2C3C9AeZ+sXbX11m27yamypNlVVegwUcisw3YCaDkZJA== + dependencies: + ansi-styles "^5.0.0" + camelcase "6" + decamelize "1.2.0" + js-tiktoken "^1.0.12" + langsmith "~0.1.30" + ml-distance "^4.0.0" + mustache "^4.2.0" p-queue "^6.6.2" p-retry "4" uuid "^9.0.0" zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/openai@^0.0.25", "@langchain/openai@~0.0.19": - version "0.0.25" - resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.25.tgz#8332abea1e3acb9b1169f90636e518c0ee90622e" - integrity sha512-cD9xPDDXK2Cjs6yYg27BpdzBnQZvBb1yaNgMoGLWIT27UQVRyT96PLC1OVMQOmMmHaKDBCj/1bW4GQQgX7+d2Q== +"@langchain/langgraph@^0.0.21": + version "0.0.21" + resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.21.tgz#5037597a954abad9ed5f0a1742226f5fcf27e7d7" + integrity sha512-7jtVZFAwvxSbIribYNzGXYIRrsAXV7YF4u1Xcpd8MYNz8sD3h8+rpIOJcYF1AdFh6laajnz0Gv8abPBHHQ2QiQ== dependencies: - "@langchain/core" "~0.1.45" - js-tiktoken "^1.0.7" - openai "^4.26.0" + "@langchain/core" ">0.1.61 <0.3.0" + uuid "^9.0.1" + +"@langchain/openai@^0.0.34", "@langchain/openai@~0.0.28": + version "0.0.34" + resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.34.tgz#36c9bca0721ab9f7e5d40927e7c0429cacbd5b56" + integrity sha512-M+CW4oXle5fdoz2T2SwdOef8pl3/1XmUx1vjn2mXUVM/128aO0l23FMF0SNBsAbRV6P+p/TuzjodchJbi0Ht/A== + dependencies: + "@langchain/core" ">0.1.56 <0.3.0" + js-tiktoken "^1.0.12" + openai "^4.41.1" zod "^3.22.4" zod-to-json-schema "^3.22.3" +"@langchain/openai@~0.1.0": + version "0.1.0" + resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.1.0.tgz#71513d7322eea148430ebb17d45a292e498819e4" + integrity sha512-jm7U9oxXQ2N03q3+S9CzEAmMJaL2FqdAi4bOYdEBS0aAWAU29so35ZOs5i2uu4W29mK9oV9XS/4A5ggR1gOLEA== + dependencies: + "@langchain/core" ">=0.2.5 <0.3.0" + js-tiktoken "^1.0.12" + openai "^4.41.1" + zod "^3.22.4" + zod-to-json-schema "^3.22.3" + +"@langchain/textsplitters@~0.0.0": + version "0.0.2" + resolved "https://registry.yarnpkg.com/@langchain/textsplitters/-/textsplitters-0.0.2.tgz#500baa8341fb7fc86fca531a4192665a319504a3" + integrity sha512-6bQOuYHTGYlkgPY/8M5WPq4nnXZpEysGzRopQCYjg2WLcEoIPUMMrXsAaNNdvU3BOeMrhin8izvpDPD165hX6Q== + dependencies: + "@langchain/core" ">0.1.0 <0.3.0" + js-tiktoken "^1.0.12" + "@langtrase/trace-attributes@^3.0.8": version "3.0.8" resolved "https://registry.yarnpkg.com/@langtrase/trace-attributes/-/trace-attributes-3.0.8.tgz#ff6ae44cfc048a9da10a7949664b2060a71b6304" @@ -21142,10 +21176,10 @@ js-string-escape@^1.0.1: resolved "https://registry.yarnpkg.com/js-string-escape/-/js-string-escape-1.0.1.tgz#e2625badbc0d67c7533e9edc1068c587ae4137ef" integrity sha1-4mJbrbwNZ8dTPp7cEGjFh65BN+8= -js-tiktoken@^1.0.7, js-tiktoken@^1.0.8: - version "1.0.10" - resolved "https://registry.yarnpkg.com/js-tiktoken/-/js-tiktoken-1.0.10.tgz#2b343ec169399dcee8f9ef9807dbd4fafd3b30dc" - integrity sha512-ZoSxbGjvGyMT13x6ACo9ebhDha/0FHdKA+OsQcMOWcm1Zs7r90Rhk5lhERLzji+3rA7EKpXCgwXcM5fF3DMpdA== +js-tiktoken@^1.0.12: + version "1.0.12" + resolved "https://registry.yarnpkg.com/js-tiktoken/-/js-tiktoken-1.0.12.tgz#af0f5cf58e5e7318240d050c8413234019424211" + integrity sha512-L7wURW1fH9Qaext0VzaUDpFGVQgjkdE3Dgsy9/+yXyGEpBKnylTd0mU0bfbNkKDlXRb6TEsZkwuflu1B8uQbJQ== dependencies: base64-js "^1.5.1" @@ -21532,17 +21566,16 @@ kuler@^2.0.0: resolved "https://registry.yarnpkg.com/kuler/-/kuler-2.0.0.tgz#e2c570a3800388fb44407e851531c1d670b061b3" integrity sha512-Xq9nH7KlWZmXAtodXDDRE7vs6DU1gTU8zYDHDiWLSip45Egwq3plLHzPn27NgvzL2r1LMPC1vdqh98sQxtqj4A== -langchain@^0.1.30: - version "0.1.30" - resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.1.30.tgz#e1adb3f1849fcd5c596c668300afd5dc8cb37a97" - integrity sha512-5h/vNMmutQ98tbB0sPDlAileZVca6A2McFgGa3+D56Dm8mSSCzTQL2DngPA6h09DlKDpSr7+6PdFw5Hoj0ZDSw== +langchain@0.2.3: + version "0.2.3" + resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.3.tgz#c14bb05cf871b21bd63b84b3ab89580b1d62539f" + integrity sha512-T9xR7zd+Nj0oXy6WoYKmZLy0DlQiDLFPGYWdOXDxy+AvqlujoPdVQgDSpdqiOHvAjezrByAoKxoHCz5XMwTP/Q== dependencies: - "@anthropic-ai/sdk" "^0.9.1" - "@langchain/community" "~0.0.41" - "@langchain/core" "~0.1.44" - "@langchain/openai" "~0.0.19" + "@langchain/core" "~0.2.0" + "@langchain/openai" "~0.0.28" + "@langchain/textsplitters" "~0.0.0" binary-extensions "^2.2.0" - js-tiktoken "^1.0.7" + js-tiktoken "^1.0.12" js-yaml "^4.1.0" jsonpointer "^5.0.1" langchainhub "~0.0.8" @@ -21555,12 +21588,45 @@ langchain@^0.1.30: zod "^3.22.4" zod-to-json-schema "^3.22.3" +langchain@^0.2.3: + version "0.2.4" + resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.4.tgz#a85295e2510425cc5f4aca32bc9adeda88f6918d" + integrity sha512-zBsBuNREn/3IlWvIQqhQ2iqf6JJhyjjsB1Db/keDkcgThPI3EcblC1pqAXU2BIKHmpNUkHBR2bAUok5+xtgOcw== + dependencies: + "@langchain/core" "~0.2.0" + "@langchain/openai" "~0.1.0" + "@langchain/textsplitters" "~0.0.0" + binary-extensions "^2.2.0" + js-tiktoken "^1.0.12" + js-yaml "^4.1.0" + jsonpointer "^5.0.1" + langchainhub "~0.0.8" + langsmith "~0.1.30" + ml-distance "^4.0.0" + openapi-types "^12.1.3" + p-retry "4" + uuid "^9.0.0" + yaml "^2.2.1" + zod "^3.22.4" + zod-to-json-schema "^3.22.3" + langchainhub@~0.0.8: version "0.0.8" resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.14, langsmith@~0.1.1, langsmith@~0.1.7: +langsmith@^0.1.30, langsmith@~0.1.30: + version "0.1.30" + resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" + integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== + dependencies: + "@types/uuid" "^9.0.1" + commander "^10.0.1" + p-queue "^6.6.2" + p-retry "4" + uuid "^9.0.0" + +langsmith@~0.1.1, langsmith@~0.1.7: version "0.1.14" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.14.tgz#2b889dbcfb49547614df276a4a5a063092a1585d" integrity sha512-iEzQLLB7/0nRpAwNBAR7B7N64fyByg5UsNjSvLaCCkQ9AS68PSafjB8xQkyI8QXXrGjU1dEqDRoa8m4SUuRdUw== @@ -23340,6 +23406,11 @@ mustache@^2.3.2: resolved "https://registry.yarnpkg.com/mustache/-/mustache-2.3.2.tgz#a6d4d9c3f91d13359ab889a812954f9230a3d0c5" integrity sha512-KpMNwdQsYz3O/SBS1qJ/o3sqUJ5wSb8gb0pul8CO0S56b9Y2ALm8zCfsjPXsqGFfoNBkDwZuZIAjhsZI03gYVQ== +mustache@^4.2.0: + version "4.2.0" + resolved "https://registry.npmjs.org/mustache/-/mustache-4.2.0.tgz" + integrity sha512-71ippSywq5Yb7/tVYyGbkBggbU8H3u5Rz56fH60jGFgr8uHwxs+aSKeqmluIVzM0m0kB7xQjKS6qPfd0b2ZoqQ== + mutation-observer@^1.0.3: version "1.0.3" resolved "https://registry.yarnpkg.com/mutation-observer/-/mutation-observer-1.0.3.tgz#42e9222b101bca82e5ba9d5a7acf4a14c0f263d0" @@ -24136,7 +24207,7 @@ open@^8.0.9, open@^8.4.0, open@~8.4.0: is-docker "^2.1.1" is-wsl "^2.2.0" -openai@^4.24.1, openai@^4.26.0: +openai@^4.24.1: version "4.26.1" resolved "https://registry.yarnpkg.com/openai/-/openai-4.26.1.tgz#7b7c0225c09922445f68f3c4cdbd5775ed31108c" integrity sha512-DvWbjhWbappsFRatOWmu4Dp1/Q4RG9oOz6CfOSjy0/Drb8G+5iAiqWAO4PfpGIkhOOKtvvNfQri2SItl+U7LhQ== @@ -24151,6 +24222,20 @@ openai@^4.24.1, openai@^4.26.0: node-fetch "^2.6.7" web-streams-polyfill "^3.2.1" +openai@^4.41.1: + version "4.47.1" + resolved "https://registry.yarnpkg.com/openai/-/openai-4.47.1.tgz#1d23c7a8eb3d7bcdc69709cd905f4c9af0181dba" + integrity sha512-WWSxhC/69ZhYWxH/OBsLEirIjUcfpQ5+ihkXKp06hmeYXgBBIUCa9IptMzYx6NdkiOCsSGYCnTIsxaic3AjRCQ== + dependencies: + "@types/node" "^18.11.18" + "@types/node-fetch" "^2.6.4" + abort-controller "^3.0.0" + agentkeepalive "^4.2.1" + form-data-encoder "1.7.2" + formdata-node "^4.3.2" + node-fetch "^2.6.7" + web-streams-polyfill "^3.2.1" + openapi-sampler@^1.3.1: version "1.4.0" resolved "https://registry.yarnpkg.com/openapi-sampler/-/openapi-sampler-1.4.0.tgz#c133cad6250481f2ec7e48b16eb70062adb514c0" @@ -29011,7 +29096,7 @@ string-replace-loader@^2.2.0: loader-utils "^1.2.3" schema-utils "^1.0.0" -"string-width-cjs@npm:string-width@^4.2.0": +"string-width-cjs@npm:string-width@^4.2.0", "string-width@^1.0.2 || 2 || 3 || 4", string-width@^4.0.0, string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.2, string-width@^4.2.3: version "4.2.3" resolved "https://registry.yarnpkg.com/string-width/-/string-width-4.2.3.tgz#269c7117d27b05ad2e536830a8ec895ef9c6d010" integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g== @@ -29029,15 +29114,6 @@ string-width@^1.0.1: is-fullwidth-code-point "^1.0.0" strip-ansi "^3.0.0" -"string-width@^1.0.2 || 2 || 3 || 4", string-width@^4.0.0, string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.2, string-width@^4.2.3: - version "4.2.3" - resolved "https://registry.yarnpkg.com/string-width/-/string-width-4.2.3.tgz#269c7117d27b05ad2e536830a8ec895ef9c6d010" - integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g== - dependencies: - emoji-regex "^8.0.0" - is-fullwidth-code-point "^3.0.0" - strip-ansi "^6.0.1" - string-width@^5.0.1, string-width@^5.1.2: version "5.1.2" resolved "https://registry.yarnpkg.com/string-width/-/string-width-5.1.2.tgz#14f8daec6d81e7221d2a357e668cab73bdbca794" @@ -29147,7 +29223,7 @@ stringify-object@^3.2.1: is-obj "^1.0.1" is-regexp "^1.0.0" -"strip-ansi-cjs@npm:strip-ansi@^6.0.1": +"strip-ansi-cjs@npm:strip-ansi@^6.0.1", strip-ansi@^6.0.0, strip-ansi@^6.0.1: version "6.0.1" resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9" integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A== @@ -29161,13 +29237,6 @@ strip-ansi@^3.0.0, strip-ansi@^3.0.1: dependencies: ansi-regex "^2.0.0" -strip-ansi@^6.0.0, strip-ansi@^6.0.1: - version "6.0.1" - resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9" - integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A== - dependencies: - ansi-regex "^5.0.1" - strip-ansi@^7.0.1, strip-ansi@^7.1.0: version "7.1.0" resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-7.1.0.tgz#d5b6568ca689d8561370b0707685d22434faff45" @@ -30928,6 +30997,11 @@ uuid@^8.0.0, uuid@^8.3.0, uuid@^8.3.2: resolved "https://registry.yarnpkg.com/uuid/-/uuid-8.3.2.tgz#80d5b5ced271bb9af6c445f21a1a04c606cefbe2" integrity sha512-+NYs2QeMWy+GWFOEm9xnn6HCDp0l7QBD7ml8zLUmJ+93Q5NF0NocErnwkTkXVFNiX3/fpC6afS8Dhb/gz7R7eg== +uuid@^9.0.1: + version "9.0.1" + resolved "https://registry.yarnpkg.com/uuid/-/uuid-9.0.1.tgz#e188d4c8853cc722220392c424cd637f32293f30" + integrity sha512-b+1eJOlsR9K8HJpow9Ok3fiWOWSIcIzXodvv0rQjVoOVNpWMpxf1wZNpt4y9h10odCNrqnYp1OBzRktckBe3sA== + v8-compile-cache-lib@^3.0.1: version "3.0.1" resolved "https://registry.yarnpkg.com/v8-compile-cache-lib/-/v8-compile-cache-lib-3.0.1.tgz#6336e8d71965cb3d35a1bbb7868445a7c05264bf" @@ -32040,7 +32114,7 @@ workerpool@6.2.1: resolved "https://registry.yarnpkg.com/workerpool/-/workerpool-6.2.1.tgz#46fc150c17d826b86a008e5a4508656777e9c343" integrity sha512-ILEIE97kDZvF9Wb9f6h5aXK4swSlKGUcOEGiIYb2OOu/IrDU9iwj0fD//SsA6E5ibwJxpEvhullJY4Sl4GcpAw== -"wrap-ansi-cjs@npm:wrap-ansi@^7.0.0": +"wrap-ansi-cjs@npm:wrap-ansi@^7.0.0", wrap-ansi@^7.0.0: version "7.0.0" resolved "https://registry.yarnpkg.com/wrap-ansi/-/wrap-ansi-7.0.0.tgz#67e145cff510a6a6984bdf1152911d69d2eb9e43" integrity sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q== @@ -32066,15 +32140,6 @@ wrap-ansi@^6.2.0: string-width "^4.1.0" strip-ansi "^6.0.0" -wrap-ansi@^7.0.0: - version "7.0.0" - resolved "https://registry.yarnpkg.com/wrap-ansi/-/wrap-ansi-7.0.0.tgz#67e145cff510a6a6984bdf1152911d69d2eb9e43" - integrity sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q== - dependencies: - ansi-styles "^4.0.0" - string-width "^4.1.0" - strip-ansi "^6.0.0" - wrap-ansi@^8.1.0: version "8.1.0" resolved "https://registry.yarnpkg.com/wrap-ansi/-/wrap-ansi-8.1.0.tgz#56dc22368ee570face1b49819975d9b9a5ead214" From 14b0c656733061010f8c195ec0c6380543acadfa Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 4 Jun 2024 00:09:18 -0600 Subject: [PATCH 05/47] Primes initial context with required kb docs and updates kb retrieval to use kbDataClient instead of esStore --- package.json | 4 +- .../knowledge_base/helpers.ts | 86 ++++++++++++++ .../knowledge_base/index.ts | 72 +++++++++++- .../helpers/get_msearch_query_body.ts | 10 +- .../graphs/default_assistant_graph/graph.ts | 109 ++++++++---------- .../nodes/execute_tools.ts | 45 ++++++++ .../{ => nodes}/generate_chat_title.ts | 10 +- .../nodes/run_agent.ts | 58 ++++++++++ .../nodes/should_continue.ts | 28 +++++ .../server/routes/attack_discovery/helpers.ts | 4 + .../attack_discovery/post_attack_discovery.ts | 1 + .../server/lib/conversational_chain.test.ts | 3 +- .../alert_counts/alert_counts_tool.test.ts | 3 + .../attack_discovery_tool.test.ts | 4 + .../knowledge_base_retrieval_tool.ts | 30 ++--- .../knowledge_base_write_tool.ts | 9 +- yarn.lock | 55 +-------- 17 files changed, 382 insertions(+), 149 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts rename x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/{ => nodes}/generate_chat_title.ts (89%) create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/should_continue.ts diff --git a/package.json b/package.json index 07631972db714..ee8511723448d 100644 --- a/package.json +++ b/package.json @@ -920,7 +920,7 @@ "@kbn/xstate-utils": "link:packages/kbn-xstate-utils", "@kbn/zod-helpers": "link:packages/kbn-zod-helpers", "@langchain/community": "^0.2.4", - "@langchain/core": "^0.2.3", + "@langchain/core": "0.2.3", "@langchain/langgraph": "^0.0.21", "@langchain/openai": "^0.0.34", "@langtrase/trace-attributes": "^3.0.8", @@ -1057,7 +1057,7 @@ "jsonwebtoken": "^9.0.2", "jsts": "^1.6.2", "kea": "^2.6.0", - "langchain": "^0.2.3", + "langchain": "0.2.3", "langsmith": "^0.1.30", "launchdarkly-js-client-sdk": "^3.1.4", "launchdarkly-node-server-sdk": "^7.0.3", diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index dc7f64e1aeee1..839ac3e559ba1 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -6,6 +6,8 @@ */ import { errors } from '@elastic/elasticsearch'; +import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types'; +import { AuthenticatedUser } from '@kbn/core-security-common'; export const isModelAlreadyExistsError = (error: Error) => { return ( @@ -14,3 +16,87 @@ export const isModelAlreadyExistsError = (error: Error) => { error.body.error.type === 'status_exception') ); }; + +/** + * Returns an Elasticsearch query DSL that performs a vector search against the Knowledge Base for the given query/user/filter. + * + * @param filter - Optional filter to apply to the search + * @param kbResource - Specific resource tag to filter for, e.g. 'esql' or 'user' + * @param modelId - ID of the model to search with, e.g. `.elser_model_2` + * @param query - The search query provided by the user + * @param required - Whether to only include required entries + * @param user - The authenticated user + * @returns + */ +export const getKBVectorSearchQuery = ({ + filter, + kbResource, + modelId, + query, + required, + user, +}: { + filter?: QueryDslQueryContainer | undefined; + kbResource?: string | undefined; + modelId: string; + query: string; + required?: boolean | undefined; + user: AuthenticatedUser; +}): QueryDslQueryContainer => { + const resourceFilter = kbResource + ? [ + { + term: { + 'metadata.kbResource': kbResource, + }, + }, + ] + : []; + const requiredFilter = required + ? [ + { + term: { + 'metadata.required': required, + }, + }, + ] + : []; + + const userFilter = [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, + }, + ], + }, + }, + }, + }, + ]; + + return { + bool: { + must: [ + { + text_expansion: { + 'vector.tokens': { + model_id: modelId, + model_text: query, + }, + }, + }, + ...requiredFilter, + ...resourceFilter, + ...userFilter, + ], + filter, + }, + }; +}; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index d2dd1165f6894..1e867ac4ae468 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -11,13 +11,14 @@ import { } from '@elastic/elasticsearch/lib/api/types'; import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import type { KibanaRequest } from '@kbn/core-http-server'; -import type { Document } from 'langchain/document'; +import { Document } from 'langchain/document'; import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; import { KnowledgeBaseEntryCreateProps, KnowledgeBaseEntryResponse, } from '@kbn/elastic-assistant-common'; import pRetry from 'p-retry'; +import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/typesWithBodyKey'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader'; @@ -26,7 +27,7 @@ import { createKnowledgeBaseEntry, transformToCreateSchema } from './create_know import { EsKnowledgeBaseEntrySchema } from './types'; import { transformESSearchToKnowledgeBaseEntry } from './transforms'; import { ESQL_DOCS_LOADED_QUERY } from '../../routes/knowledge_base/constants'; -import { isModelAlreadyExistsError } from './helpers'; +import { getKBVectorSearchQuery, isModelAlreadyExistsError } from './helpers'; interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { ml: MlPluginSetup; @@ -220,8 +221,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { /** * Adds LangChain Documents to the knowledge base * - * @param documents - * @param authenticatedUser + * @param documents LangChain Documents to add to the knowledge base */ public addKnowledgeBaseDocuments = async ({ documents, @@ -265,6 +265,70 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { return created?.data ? transformESSearchToKnowledgeBaseEntry(created?.data) : []; }; + /** + * Performs similarity search to retrieve LangChain Documents from the knowledge base + */ + public getKnowledgeBaseDocuments = async ({ + filter, + kbResource, + query, + required, + }: { + filter?: QueryDslQueryContainer; + kbResource?: string; + query: string; + required?: boolean; + }): Promise => { + const user = this.options.currentUser; + if (user == null) { + throw new Error( + 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' + ); + } + + const esClient = await this.options.elasticsearchClientPromise; + const modelId = await this.options.getElserId(); + + const vectorSearchQuery = getKBVectorSearchQuery({ + filter, + kbResource, + modelId, + query, + required, + user, + }); + + try { + const result = await esClient.search({ + index: this.indexTemplateAndPattern.alias, + size: 10, + query: vectorSearchQuery, + }); + + const results = result.hits.hits.map( + (hit) => + new Document({ + pageContent: hit?._source?.text ?? '', + metadata: hit?._source?.metadata ?? {}, + }) + ); + + this.options.logger.debug( + `getKnowledgeBaseDocuments() - Similarity Search Query:\n ${JSON.stringify( + vectorSearchQuery + )}` + ); + this.options.logger.debug( + `getKnowledgeBaseDocuments() - Similarity Search Results:\n ${JSON.stringify(results)}` + ); + + return results; + } catch (e) { + this.options.logger.error(`Error performing KB Similarity Search: ${e.message}`); + return []; + } + }; + /** * Creates a new Knowledge Base Entry. * diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_msearch_query_body.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_msearch_query_body.ts index 7780eab2c2d1d..c93c3f2e30954 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_msearch_query_body.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_msearch_query_body.ts @@ -58,10 +58,10 @@ export const getMsearchQueryBody = ({ query: vectorSearchQuery, size: vectorSearchQuerySize, }, - // { index }, - // { - // query: termsSearchQuery, - // size: termsSearchQuerySize, - // }, + { index }, + { + query: termsSearchQuery, + size: termsSearchQuerySize, + }, ], }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts index 32ccf28e14da0..e3d84c74a9991 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts @@ -5,8 +5,7 @@ * 2.0. */ -import { ToolExecutor } from '@langchain/langgraph/prebuilt'; -import { RunnableConfig, RunnableLambda } from '@langchain/core/runnables'; +import { RunnableConfig } from '@langchain/core/runnables'; import { END, START, StateGraph, StateGraphArgs } from '@langchain/langgraph'; import { AgentAction, AgentFinish, AgentStep } from '@langchain/core/agents'; import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; @@ -17,8 +16,10 @@ import type { Logger } from '@kbn/logging'; import { BaseMessage } from '@langchain/core/messages'; import { BaseChatModel } from '@langchain/core/language_models/chat_models'; import { AgentState, NodeParamsBase } from './types'; -import { generateChatTitle } from './generate_chat_title'; import { AssistantDataClients } from '../../executors/types'; +import { shouldContinue } from './nodes/should_continue'; +import { AGENT_NODE, runAgent } from './nodes/run_agent'; +import { executeTools, TOOLS_NODE } from './nodes/execute_tools'; export const DEFAULT_ASSISTANT_GRAPH_ID = 'Default Security Assistant Graph'; @@ -68,75 +69,55 @@ export const getDefaultAssistantGraph = ({ }, }; + // Default node parameters const nodeParams: NodeParamsBase = { model: llm, logger, }; - // Create a tool executor - const toolExecutor = new ToolExecutor({ tools }); - - // Define logic that will be used to determine which conditional edge to go down - const shouldContinue = (state: AgentState) => { - logger.debug(`graph:shouldContinue:state\n${JSON.stringify(state, null, 2)}`); - if (state.agentOutcome && 'returnValues' in state.agentOutcome) { - return 'end'; - } - return 'continue'; - }; - - const runAgent = async (state: AgentState, config?: RunnableConfig) => { - logger.debug(`graph:runAgent:\nstate\n${JSON.stringify(state, null, 2)}`); - - const agentOutcome = await agentRunnable.invoke( - { - ...state, - chat_history: messages, - knowledge_history: 'The users favorite color is blue', // TODO: Plumb through initial retrieval - }, - config - ); - return { - agentOutcome, - }; - }; - - const executeTools = async (state: AgentState, config?: RunnableConfig) => { - logger.debug(`graph:executeTools:state\n${JSON.stringify(state, null, 2)}`); - const agentAction = state.agentOutcome; - if (!agentAction || 'returnValues' in agentAction) { - throw new Error('Agent has not been run yet'); - } - const out = await toolExecutor.invoke(agentAction, config); - return { - steps: [{ action: agentAction, observation: JSON.stringify(out, null, 2) }], - }; - }; - - // Create a new graph, with the default state from above - const workflow = new StateGraph({ channels: graphState }); - - // Define the nodes to cycle between - workflow.addNode('generateChatTitle', (state: AgentState) => - generateChatTitle({ + // Create nodes + // const generateChatTitleNode = (state: AgentState) => + // generateChatTitle({ + // ...nodeParams, + // conversationId, + // conversationsDataClient: dataClients?.conversationsDataClient, + // logger: logger.get(GENERATE_CHAT_TITLE_NODE), + // state, + // }); + const runAgentNode = (state: AgentState, config?: RunnableConfig) => + runAgent({ + ...nodeParams, + agentRunnable, + config, + dataClients, + logger: logger.get(AGENT_NODE), state, - conversationsDataClient: dataClients?.conversationsDataClient, - conversationId, + }); + const executeToolsNode = (state: AgentState, config?: RunnableConfig) => + executeTools({ ...nodeParams, - }) - ); - workflow.addNode('agent', new RunnableLambda({ func: runAgent })); - workflow.addNode('action', new RunnableLambda({ func: executeTools })); - - // Add conditional edge for determining if we shouldContinue - workflow.addConditionalEdges('agent', shouldContinue, { continue: 'action', end: END }); - - // Add edges for start, and between agent and action (action always followed by agent) - workflow.addEdge(START, 'generateChatTitle'); - workflow.addEdge('generateChatTitle', 'agent'); - workflow.addEdge('action', 'agent'); + config, + logger: logger.get(TOOLS_NODE), + state, + tools, + }); + const shouldContinueEdge = (state: AgentState) => shouldContinue({ ...nodeParams, state }); - return workflow.compile(); + // Put together a new graph using the nodes and default state from above + const graph = new StateGraph({ channels: graphState }); + // Define the nodes to cycle between + // TODO: Re-enable title generation and wire remainder of persistence within graph after https://github.com/elastic/kibana/pull/184485 + // graph.addNode('generateChatTitle', generateChatTitleNode); + graph.addNode(AGENT_NODE, runAgentNode); + graph.addNode(TOOLS_NODE, executeToolsNode); + // Add conditional edge for basic routing + graph.addConditionalEdges(AGENT_NODE, shouldContinueEdge, { continue: TOOLS_NODE, end: END }); + // Add edges, starting with chat title generation, then alternate between agent and action until finished + graph.addEdge(START, AGENT_NODE); + // graph.addEdge('generateChatTitle', AGENT_NODE); + graph.addEdge(TOOLS_NODE, AGENT_NODE); + // Compile the graph + return graph.compile(); } catch (e) { throw new Error(`Unable to compile DefaultAssistantGraph\n${e}`); } diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts new file mode 100644 index 0000000000000..b42455e14f6f1 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts @@ -0,0 +1,45 @@ +/* + * 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 { RunnableConfig } from '@langchain/core/runnables'; +import { StructuredTool } from '@langchain/core/tools'; +import { ToolExecutor } from '@langchain/langgraph/prebuilt'; +import { AgentState, NodeParamsBase } from '../types'; + +export interface ExecuteToolsParams extends NodeParamsBase { + state: AgentState; + config?: RunnableConfig; + tools: StructuredTool[]; +} + +export const TOOLS_NODE = 'tools'; + +/** + * Node to execute tools + * + * Note: Could maybe leverage `ToolNode` if tool selection state is pushed to `messages[]`. + * See: https://github.com/langchain-ai/langgraphjs/blob/0ef76d603b55c00a04f5793d1e6ab15af7c756cb/langgraph/src/prebuilt/tool_node.ts + * + * @param config - Any configuration that may've been supplied + * @param logger - The scoped logger + * @param state - The current state of the graph + * @param tools - The tools available to execute + */ +export const executeTools = async ({ config, logger, state, tools }: ExecuteToolsParams) => { + logger.debug(`Node state:\n${JSON.stringify(state, null, 2)}`); + + const toolExecutor = new ToolExecutor({ tools }); + const agentAction = state.agentOutcome; + + if (!agentAction || 'returnValues' in agentAction) { + throw new Error('Agent has not been run yet'); + } + const out = await toolExecutor.invoke(agentAction, config); + return { + steps: [{ action: agentAction, observation: JSON.stringify(out, null, 2) }], + }; +}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/generate_chat_title.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts similarity index 89% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/generate_chat_title.ts rename to x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts index 9cdf8fdb96943..bcba25eab0b0d 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/generate_chat_title.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts @@ -7,8 +7,8 @@ import { StringOutputParser } from '@langchain/core/output_parsers'; import { ChatPromptTemplate } from '@langchain/core/prompts'; -import { AgentState, NodeParamsBase } from './types'; -import { AIAssistantConversationsDataClient } from '../../../../ai_assistant_data_clients/conversations'; +import { AgentState, NodeParamsBase } from '../types'; +import { AIAssistantConversationsDataClient } from '../../../../../ai_assistant_data_clients/conversations'; export const GENERATE_CHAT_TITLE_PROMPT = ChatPromptTemplate.fromMessages([ [ @@ -27,14 +27,16 @@ export interface GenerateChatTitleParams extends NodeParamsBase { conversationId?: string; state: AgentState; } + +export const GENERATE_CHAT_TITLE_NODE = 'generateChatTitle'; + export const generateChatTitle = async ({ conversationsDataClient, logger, model, state, }: GenerateChatTitleParams) => { - logger.debug('node:generateChatTitle'); - logger.debug(`state:\n ${JSON.stringify(state, null, 2)}`); + logger.debug(`Node state:\n ${JSON.stringify(state, null, 2)}`); if (state.messages.length !== 0) { logger.debug('No need to generate chat title, messages already exist'); return; 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 new file mode 100644 index 0000000000000..b0353bb5d8ec7 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts @@ -0,0 +1,58 @@ +/* + * 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 { RunnableConfig } from '@langchain/core/runnables'; +import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; +import { AgentState, NodeParamsBase } from '../types'; +import { AssistantDataClients } from '../../../executors/types'; + +export interface RunAgentParams extends NodeParamsBase { + agentRunnable: AgentRunnableSequence; + dataClients?: AssistantDataClients; + state: AgentState; + config?: RunnableConfig; +} + +export const AGENT_NODE = 'agent'; + +const NO_HISTORY = '[No existing knowledge history]'; +/** + * Node to run the agent + * + * @param agentRunnable - The agent to run + * @param config - Any configuration that may've been supplied + * @param logger - The scoped logger + * @param dataClients - Data clients available for use + * @param state - The current state of the graph + */ +export const runAgent = async ({ + agentRunnable, + config, + dataClients, + logger, + state, +}: RunAgentParams) => { + logger.debug(`Node state:\n${JSON.stringify(state, null, 2)}`); + + const knowledgeHistory = await dataClients?.kbDataClient?.getKnowledgeBaseDocuments({ + kbResource: 'user', + required: true, + query: '', + }); + + const agentOutcome = await agentRunnable.invoke( + { + ...state, + chat_history: state.messages, // TODO: Message de-dupe with ...state spread + knowledge_history: knowledgeHistory?.length ? knowledgeHistory : NO_HISTORY, + }, + config + ); + return { + agentOutcome, + }; +}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/should_continue.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/should_continue.ts new file mode 100644 index 0000000000000..281963df363a8 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/should_continue.ts @@ -0,0 +1,28 @@ +/* + * 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 { AgentState, NodeParamsBase } from '../types'; + +export interface ShouldContinueParams extends NodeParamsBase { + state: AgentState; +} + +/** + * Node to determine which conditional edge to choose. Essentially the 'router' node. + * + * @param logger - The scoped logger + * @param state - The current state of the graph + */ +export const shouldContinue = ({ logger, state }: ShouldContinueParams) => { + logger.debug(`Node state:\n${JSON.stringify(state, null, 2)}`); + + if (state.agentOutcome && 'returnValues' in state.agentOutcome) { + return 'end'; + } + + return 'continue'; +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts index 65a3c96b27a0c..4f916be8105d4 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts @@ -6,6 +6,7 @@ */ import { KibanaRequest } from '@kbn/core/server'; +import { Logger } from '@kbn/logging'; import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { AttackDiscoveryPostRequestBody, @@ -40,6 +41,7 @@ export const getAssistantToolParams = ({ langChainTimeout, latestReplacements, llm, + logger, onNewReplacements, request, size, @@ -50,6 +52,7 @@ export const getAssistantToolParams = ({ langChainTimeout: number; latestReplacements: Replacements; llm: ActionsClientLlm; + logger: Logger; onNewReplacements: (newReplacements: Replacements) => void; request: KibanaRequest< unknown, @@ -65,6 +68,7 @@ export const getAssistantToolParams = ({ esClient, langChainTimeout, llm, + logger, modelExists: false, // not required for attack discovery onNewReplacements, replacements: latestReplacements, diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts index 5d9240dd1d97d..7859d635ccb30 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts @@ -127,6 +127,7 @@ export const postAttackDiscoveryRoute = ( latestReplacements, langChainTimeout: LANG_CHAIN_TIMEOUT, llm, + logger, onNewReplacements, request, size, diff --git a/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts b/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts index c865526115db1..ef01c20eeaf29 100644 --- a/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts +++ b/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts @@ -8,8 +8,7 @@ import type { Client } from '@elastic/elasticsearch'; import { createAssist as Assist } from '../utils/assist'; import { clipContext, ConversationalChain } from './conversational_chain'; -import { FakeListChatModel } from '@langchain/core/utils/testing'; -import { FakeListLLM } from 'langchain/llms/fake'; +import { FakeListChatModel, FakeListLLM } from '@langchain/core/utils/testing'; import { BaseChatModel } from '@langchain/core/language_models/chat_models'; import { experimental_StreamData, Message } from 'ai'; import { ChatPromptTemplate } from '@langchain/core/prompts'; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts index e2c3036477875..752f8e472a755 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts @@ -8,6 +8,7 @@ import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import type { KibanaRequest } from '@kbn/core-http-server'; import type { DynamicTool } from '@langchain/core/tools'; +import { loggerMock } from '@kbn/logging-mocks'; import { ALERT_COUNTS_TOOL } from './alert_counts_tool'; import type { RetrievalQAChain } from 'langchain/chains'; import type { ExecuteConnectorRequestBody } from '@kbn/elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen'; @@ -29,9 +30,11 @@ describe('AlertCountsTool', () => { const isEnabledKnowledgeBase = true; const chain = {} as unknown as RetrievalQAChain; const modelExists = true; + const logger = loggerMock.create(); const rest = { isEnabledKnowledgeBase, chain, + logger, modelExists, }; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.test.ts index a608673adf661..5d8fb0b51739a 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.test.ts @@ -11,6 +11,8 @@ import type { AttackDiscoveryPostRequestBody } from '@kbn/elastic-assistant-comm import type { ActionsClientLlm } from '@kbn/langchain/server'; import type { DynamicTool } from '@langchain/core/tools'; +import { loggerMock } from '@kbn/logging-mocks'; + import { ATTACK_DISCOVERY_TOOL } from './attack_discovery_tool'; import { mockAnonymizationFields } from '../mock/mock_anonymization_fields'; import { mockEmptyOpenAndAcknowledgedAlertsQueryResults } from '../mock/mock_empty_open_and_acknowledged_alerts_qery_results'; @@ -66,11 +68,13 @@ describe('AttackDiscoveryTool', () => { search: jest.fn(), } as unknown as ElasticsearchClient; const llm = jest.fn() as unknown as ActionsClientLlm; + const logger = loggerMock.create(); const rest = { anonymizationFields: mockAnonymizationFields, isEnabledKnowledgeBase: false, llm, + logger, modelExists: false, onNewReplacements: jest.fn(), size, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts index c06898b1ce3b9..47cb35e244d51 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts @@ -8,13 +8,16 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from 'zod'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import type { AIAssistantKnowledgeBaseDataClient } from '@kbn/elastic-assistant-plugin/server/ai_assistant_data_clients/knowledge_base'; import { APP_UI_ID } from '../../../../common'; -export type KnowledgeBaseRetrievalToolParams = AssistantToolParams; +export interface KnowledgeBaseRetrievalToolParams extends AssistantToolParams { + kbDataClient: AIAssistantKnowledgeBaseDataClient; +} const toolDetails = { description: - "Call this for fetching details from the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Call this function when the user asks for information about themself, like 'what is my favorite...' or 'using my saved....'. Input must always be the free-text query on a single line, with no other text. You are welcome to re-write the query to be a summary of items/things to search for in the knowledge base, as a vector search will be performed to return similar results when requested. If the results returned do not look relevant, disregard and tell the user you were unable to find the information they were looking for.", + "Call this for fetching details from the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Call this function when the user asks for information about themself, like 'what is my favorite...' or 'using my saved....'. Input must always be the free-text query on a single line, with no other text. You are welcome to re-write the query to be a summary of items/things to search for in the knowledge base, as a vector search will be performed to return similar results when requested. If the results returned do not look relevant, disregard and tell the user you were unable to find the information they were looking for. All requests include a `knowledge history` section which includes some existing knowledge of the user. DO NOT CALL THIS FUNCTION if the `knowledge history` sections appears to be able to answer the user's query.", id: 'knowledge-base-retrieval-tool', name: 'KnowledgeBaseRetrievalTool', }; @@ -22,14 +25,14 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is KnowledgeBaseRetrievalToolParams => { - const { chain, isEnabledKnowledgeBase, modelExists } = params; - return isEnabledKnowledgeBase && modelExists && chain != null; + const { kbDataClient, isEnabledKnowledgeBase, modelExists } = params; + return isEnabledKnowledgeBase && modelExists && kbDataClient != null; }, getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { chain, logger } = params as KnowledgeBaseRetrievalToolParams; - if (chain == null) return null; + const { kbDataClient, logger } = params as KnowledgeBaseRetrievalToolParams; + if (kbDataClient == null) return null; return new DynamicStructuredTool({ name: toolDetails.name, @@ -39,13 +42,14 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { }), func: async (input, _, cbManager) => { logger.debug(`KnowledgeBaseRetrievalToolParams:input\n ${JSON.stringify(input, null, 2)}`); - const result = await chain.invoke( - { - query: input.query, - }, - cbManager - ); - return result.text; + + const docs = await kbDataClient.getKnowledgeBaseDocuments({ + query: input.query, + kbResource: 'user', + required: false, + }); + + return JSON.stringify(docs); }, tags: ['knowledge-base'], }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts index 754955dd648ee..addb2a5580dfc 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts @@ -18,7 +18,7 @@ export interface KnowledgeBaseWriteToolParams extends AssistantToolParams { const toolDetails = { description: - "Call this for writing details to the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Input will be the summarized knowledge base entry to store, with no other text.", + "Call this for writing details to the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Input will be the summarized knowledge base entry to store, with no other text, and whether or not the entry is required.", id: 'knowledge-base-write-tool', name: 'KnowledgeBaseWriteTool', }; @@ -40,12 +40,17 @@ export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { description: toolDetails.description, schema: z.object({ query: z.string().describe(`Summary of items/things to save in the knowledge base`), + required: z + .boolean() + .describe( + `Whether or not the entry is required to always be included in conversations. Is only true if the user explicitly asks for it to be required or always included in conversations, otherwise this is always false.` + ), }), func: async (input, _, cbManager) => { logger.debug(`KnowledgeBaseWriteToolParams:input\n ${JSON.stringify(input, null, 2)}`); const knowledgeBaseEntry: KnowledgeBaseEntryCreateProps = { - metadata: { kbResource: 'user', source: 'conversation', required: false }, + metadata: { kbResource: 'user', source: 'conversation', required: input.required }, text: input.query, }; diff --git a/yarn.lock b/yarn.lock index ff76fc98269a9..271fa2d53a762 100644 --- a/yarn.lock +++ b/yarn.lock @@ -6882,7 +6882,7 @@ zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@0.2.3", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@>=0.2.5 <0.3.0", "@langchain/core@~0.2.0": +"@langchain/core@0.2.3", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@~0.2.0": version "0.2.3" resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.3.tgz#7faa82f92b0c7843506e827a38bfcbb60f009d13" integrity sha512-mVuFHSLpPQ4yOHNXeoSA3LnmIMuFmUiit5rvbYcPZqM6SrB2zCNN2nD4Ty5+3H5X4tYItDoSqsTuUNUQySXRQw== @@ -6900,24 +6900,6 @@ zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/core@^0.2.3": - version "0.2.5" - resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.5.tgz#4983469d04615f7f66ae573319939ed7301f6889" - integrity sha512-tMaKRFVewFn8crQwlbXGjT7hlMdX1yXHap1ebBx7Bb2C3C9AeZ+sXbX11m27yamypNlVVegwUcisw3YCaDkZJA== - dependencies: - ansi-styles "^5.0.0" - camelcase "6" - decamelize "1.2.0" - js-tiktoken "^1.0.12" - langsmith "~0.1.30" - ml-distance "^4.0.0" - mustache "^4.2.0" - p-queue "^6.6.2" - p-retry "4" - uuid "^9.0.0" - zod "^3.22.4" - zod-to-json-schema "^3.22.3" - "@langchain/langgraph@^0.0.21": version "0.0.21" resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.21.tgz#5037597a954abad9ed5f0a1742226f5fcf27e7d7" @@ -6937,17 +6919,6 @@ zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/openai@~0.1.0": - version "0.1.0" - resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.1.0.tgz#71513d7322eea148430ebb17d45a292e498819e4" - integrity sha512-jm7U9oxXQ2N03q3+S9CzEAmMJaL2FqdAi4bOYdEBS0aAWAU29so35ZOs5i2uu4W29mK9oV9XS/4A5ggR1gOLEA== - dependencies: - "@langchain/core" ">=0.2.5 <0.3.0" - js-tiktoken "^1.0.12" - openai "^4.41.1" - zod "^3.22.4" - zod-to-json-schema "^3.22.3" - "@langchain/textsplitters@~0.0.0": version "0.0.2" resolved "https://registry.yarnpkg.com/@langchain/textsplitters/-/textsplitters-0.0.2.tgz#500baa8341fb7fc86fca531a4192665a319504a3" @@ -21588,34 +21559,12 @@ langchain@0.2.3: zod "^3.22.4" zod-to-json-schema "^3.22.3" -langchain@^0.2.3: - version "0.2.4" - resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.4.tgz#a85295e2510425cc5f4aca32bc9adeda88f6918d" - integrity sha512-zBsBuNREn/3IlWvIQqhQ2iqf6JJhyjjsB1Db/keDkcgThPI3EcblC1pqAXU2BIKHmpNUkHBR2bAUok5+xtgOcw== - dependencies: - "@langchain/core" "~0.2.0" - "@langchain/openai" "~0.1.0" - "@langchain/textsplitters" "~0.0.0" - binary-extensions "^2.2.0" - js-tiktoken "^1.0.12" - js-yaml "^4.1.0" - jsonpointer "^5.0.1" - langchainhub "~0.0.8" - langsmith "~0.1.30" - ml-distance "^4.0.0" - openapi-types "^12.1.3" - p-retry "4" - uuid "^9.0.0" - yaml "^2.2.1" - zod "^3.22.4" - zod-to-json-schema "^3.22.3" - langchainhub@~0.0.8: version "0.0.8" resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.30, langsmith@~0.1.30: +langsmith@^0.1.30: version "0.1.30" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== From 83cdc37d66c411b2393af59343ebb3fc06e33c7e Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 4 Jun 2024 00:12:37 -0600 Subject: [PATCH 06/47] Remove chat title generation node --- .../graphs/default_assistant_graph/graph.ts | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts index e3d84c74a9991..12c72fe54ec3e 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts @@ -76,14 +76,6 @@ export const getDefaultAssistantGraph = ({ }; // Create nodes - // const generateChatTitleNode = (state: AgentState) => - // generateChatTitle({ - // ...nodeParams, - // conversationId, - // conversationsDataClient: dataClients?.conversationsDataClient, - // logger: logger.get(GENERATE_CHAT_TITLE_NODE), - // state, - // }); const runAgentNode = (state: AgentState, config?: RunnableConfig) => runAgent({ ...nodeParams, @@ -106,15 +98,12 @@ export const getDefaultAssistantGraph = ({ // Put together a new graph using the nodes and default state from above const graph = new StateGraph({ channels: graphState }); // Define the nodes to cycle between - // TODO: Re-enable title generation and wire remainder of persistence within graph after https://github.com/elastic/kibana/pull/184485 - // graph.addNode('generateChatTitle', generateChatTitleNode); graph.addNode(AGENT_NODE, runAgentNode); graph.addNode(TOOLS_NODE, executeToolsNode); // Add conditional edge for basic routing graph.addConditionalEdges(AGENT_NODE, shouldContinueEdge, { continue: TOOLS_NODE, end: END }); - // Add edges, starting with chat title generation, then alternate between agent and action until finished + // Add edges, alternating between agent and action until finished graph.addEdge(START, AGENT_NODE); - // graph.addEdge('generateChatTitle', AGENT_NODE); graph.addEdge(TOOLS_NODE, AGENT_NODE); // Compile the graph return graph.compile(); From b04da0e78e78d3c5241ff4cd76bb78238d61646f Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 4 Jun 2024 00:36:42 -0600 Subject: [PATCH 07/47] yarn.lock update --- yarn.lock | 34 ++++------------------------------ 1 file changed, 4 insertions(+), 30 deletions(-) diff --git a/yarn.lock b/yarn.lock index 8b4dca506de26..536c9848a077d 100644 --- a/yarn.lock +++ b/yarn.lock @@ -21564,7 +21564,7 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.30: +langsmith@^0.1.30, langsmith@~0.1.1, langsmith@~0.1.7: version "0.1.30" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== @@ -21575,17 +21575,6 @@ langsmith@^0.1.30: p-retry "4" uuid "^9.0.0" -langsmith@~0.1.1, langsmith@~0.1.7: - version "0.1.14" - resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.14.tgz#2b889dbcfb49547614df276a4a5a063092a1585d" - integrity sha512-iEzQLLB7/0nRpAwNBAR7B7N64fyByg5UsNjSvLaCCkQ9AS68PSafjB8xQkyI8QXXrGjU1dEqDRoa8m4SUuRdUw== - dependencies: - "@types/uuid" "^9.0.1" - commander "^10.0.1" - p-queue "^6.6.2" - p-retry "4" - uuid "^9.0.0" - language-subtag-registry@~0.3.2: version "0.3.21" resolved "https://registry.yarnpkg.com/language-subtag-registry/-/language-subtag-registry-0.3.21.tgz#04ac218bea46f04cb039084602c6da9e788dd45a" @@ -24156,22 +24145,7 @@ open@^8.0.9, open@^8.4.0, open@~8.4.0: is-docker "^2.1.1" is-wsl "^2.2.0" -openai@^4.24.1: - version "4.26.1" - resolved "https://registry.yarnpkg.com/openai/-/openai-4.26.1.tgz#7b7c0225c09922445f68f3c4cdbd5775ed31108c" - integrity sha512-DvWbjhWbappsFRatOWmu4Dp1/Q4RG9oOz6CfOSjy0/Drb8G+5iAiqWAO4PfpGIkhOOKtvvNfQri2SItl+U7LhQ== - dependencies: - "@types/node" "^18.11.18" - "@types/node-fetch" "^2.6.4" - abort-controller "^3.0.0" - agentkeepalive "^4.2.1" - digest-fetch "^1.3.0" - form-data-encoder "1.7.2" - formdata-node "^4.3.2" - node-fetch "^2.6.7" - web-streams-polyfill "^3.2.1" - -openai@^4.41.1: +openai@^4.24.1, openai@^4.41.1: version "4.47.1" resolved "https://registry.yarnpkg.com/openai/-/openai-4.47.1.tgz#1d23c7a8eb3d7bcdc69709cd905f4c9af0181dba" integrity sha512-WWSxhC/69ZhYWxH/OBsLEirIjUcfpQ5+ihkXKp06hmeYXgBBIUCa9IptMzYx6NdkiOCsSGYCnTIsxaic3AjRCQ== @@ -30931,7 +30905,7 @@ uuid-browser@^3.1.0: resolved "https://registry.yarnpkg.com/uuid-browser/-/uuid-browser-3.1.0.tgz#0f05a40aef74f9e5951e20efbf44b11871e56410" integrity sha1-DwWkCu90+eWVHiDvv0SxGHHlZBA= -uuid@9.0.0, uuid@^9, uuid@^9.0.0: +uuid@9.0.0: version "9.0.0" resolved "https://registry.yarnpkg.com/uuid/-/uuid-9.0.0.tgz#592f550650024a38ceb0c562f2f6aa435761efb5" integrity sha512-MXcSTerfPa4uqyzStbRoTgt5XIe3x5+42+q1sDuy3R5MDk66URdLMOZe5aPX/SQd+kuYAh0FdP/pO28IkQyTeg== @@ -30946,7 +30920,7 @@ uuid@^8.0.0, uuid@^8.3.0, uuid@^8.3.2: resolved "https://registry.yarnpkg.com/uuid/-/uuid-8.3.2.tgz#80d5b5ced271bb9af6c445f21a1a04c606cefbe2" integrity sha512-+NYs2QeMWy+GWFOEm9xnn6HCDp0l7QBD7ml8zLUmJ+93Q5NF0NocErnwkTkXVFNiX3/fpC6afS8Dhb/gz7R7eg== -uuid@^9.0.1: +uuid@^9, uuid@^9.0.0, uuid@^9.0.1: version "9.0.1" resolved "https://registry.yarnpkg.com/uuid/-/uuid-9.0.1.tgz#e188d4c8853cc722220392c424cd637f32293f30" integrity sha512-b+1eJOlsR9K8HJpow9Ok3fiWOWSIcIzXodvv0rQjVoOVNpWMpxf1wZNpt4y9h10odCNrqnYp1OBzRktckBe3sA== From 5bc866be98904f0a125238b740b09cfbc78f5e08 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 4 Jun 2024 01:12:49 -0600 Subject: [PATCH 08/47] yarn.lock update2 --- yarn.lock | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/yarn.lock b/yarn.lock index 536c9848a077d..6ea3e40fa8ec2 100644 --- a/yarn.lock +++ b/yarn.lock @@ -12832,11 +12832,6 @@ bare-path@^2.0.0, bare-path@^2.1.0: dependencies: bare-os "^2.1.0" -base-64@^0.1.0: - version "0.1.0" - resolved "https://registry.yarnpkg.com/base-64/-/base-64-0.1.0.tgz#780a99c84e7d600260361511c4877613bf24f6bb" - integrity sha512-Y5gU45svrR5tI2Vt/X9GPd3L0HNIKzGu202EjxrXMpuc2V2CiKgemAbUUsqYmZJvPtCXoUKjNZwBJzsNScUbXA== - base64-js@1.3.1: version "1.3.1" resolved "https://registry.yarnpkg.com/base64-js/-/base64-js-1.3.1.tgz#58ece8cb75dd07e71ed08c736abc5fac4dbf8df1" @@ -15851,14 +15846,6 @@ diffie-hellman@^5.0.0: miller-rabin "^4.0.0" randombytes "^2.0.0" -digest-fetch@^1.3.0: - version "1.3.0" - resolved "https://registry.yarnpkg.com/digest-fetch/-/digest-fetch-1.3.0.tgz#898e69264d00012a23cf26e8a3e40320143fc661" - integrity sha512-CGJuv6iKNM7QyZlM2T3sPAdZWd/p9zQiRNS9G+9COUCwzWFTs0Xp8NF5iePx7wtvhDykReiRRrSeNb4oMmB8lA== - dependencies: - base-64 "^0.1.0" - md5 "^2.3.0" - dir-glob@^2.2.2: version "2.2.2" resolved "https://registry.yarnpkg.com/dir-glob/-/dir-glob-2.2.2.tgz#fa09f0694153c8918b18ba0deafae94769fc50c4" From 8477e13754b524ce5429d320eb82334f421444f4 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 4 Jun 2024 16:21:37 +0200 Subject: [PATCH 09/47] fix types --- .../graphs/default_assistant_graph/graph.ts | 9 ++++++--- .../graphs/default_assistant_graph/helpers.ts | 12 ++++++------ .../server/lib/conversational_chain.test.ts | 4 ++-- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts index 12c72fe54ec3e..779bf20a61720 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts @@ -10,7 +10,6 @@ import { END, START, StateGraph, StateGraphArgs } from '@langchain/langgraph'; import { AgentAction, AgentFinish, AgentStep } from '@langchain/core/agents'; import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; import { StructuredTool } from '@langchain/core/tools'; -import type { CompiledStateGraph } from '@langchain/langgraph/dist/graph/state'; import type { Logger } from '@kbn/logging'; import { BaseMessage } from '@langchain/core/messages'; @@ -33,6 +32,8 @@ interface GetDefaultAssistantGraphParams { tools: StructuredTool[]; } +export type DefaultAssistantGraph = ReturnType; + /** * Returns a compiled default assistant graph */ @@ -44,7 +45,7 @@ export const getDefaultAssistantGraph = ({ logger, messages, tools, -}: GetDefaultAssistantGraphParams): CompiledStateGraph => { +}: GetDefaultAssistantGraphParams) => { try { // Default graph state const graphState: StateGraphArgs['channels'] = { @@ -96,7 +97,9 @@ export const getDefaultAssistantGraph = ({ const shouldContinueEdge = (state: AgentState) => shouldContinue({ ...nodeParams, state }); // Put together a new graph using the nodes and default state from above - const graph = new StateGraph({ channels: graphState }); + const graph = new StateGraph, '__start__' | 'agent' | 'tools'>({ + channels: graphState, + }); // Define the nodes to cycle between graph.addNode(AGENT_NODE, runAgentNode); graph.addNode(TOOLS_NODE, executeToolsNode); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts index f2c318b32a2ab..383b3e9f5cee8 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts @@ -11,15 +11,14 @@ import { streamFactory, StreamResponseWithHeaders } from '@kbn/ml-response-strea import { transformError } from '@kbn/securitysolution-es-utils'; import type { KibanaRequest } from '@kbn/core-http-server'; import type { ExecuteConnectorRequestBody, TraceData } from '@kbn/elastic-assistant-common'; -import { CompiledStateGraph } from '@langchain/langgraph/dist/graph/state'; -import { DEFAULT_ASSISTANT_GRAPH_ID } from './graph'; +import { DEFAULT_ASSISTANT_GRAPH_ID, DefaultAssistantGraph } from './graph'; import type { OnLlmResponse, TraceOptions } from '../../executors/types'; import type { APMTracer } from '../../tracers/apm_tracer'; import { withAssistantSpan } from '../../tracers/with_assistant_span'; interface StreamGraphParams { apmTracer: APMTracer; - assistantGraph: CompiledStateGraph; + assistantGraph: DefaultAssistantGraph; inputs: { input: string }; logger: Logger; onLlmResponse?: OnLlmResponse; @@ -100,6 +99,7 @@ export const streamGraph = async ({ const msg = chunk.message; if (msg.tool_call_chunks && msg.tool_call_chunks.length > 0) { + /* empty */ } else if (!didEnd) { if (msg.response_metadata.finish_reason === 'stop') { handleStreamEnd(finalMessage); @@ -111,7 +111,7 @@ export const streamGraph = async ({ } } - processEvent(); + await processEvent(); } catch (err) { // if I throw an error here, it crashes the server. Not sure how to get around that. // If I put await on this function the error works properly, but when there is not an error @@ -129,14 +129,14 @@ export const streamGraph = async ({ }; // Start processing events, do not await! Return `responseWithHeaders` immediately - processEvent(); + await processEvent(); return responseWithHeaders; }; interface InvokeGraphParams { apmTracer: APMTracer; - assistantGraph: CompiledStateGraph; + assistantGraph: DefaultAssistantGraph; inputs: { input: string }; onLlmResponse?: OnLlmResponse; traceOptions?: TraceOptions; diff --git a/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts b/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts index ef01c20eeaf29..641f3f334eeb6 100644 --- a/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts +++ b/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts @@ -8,7 +8,7 @@ import type { Client } from '@elastic/elasticsearch'; import { createAssist as Assist } from '../utils/assist'; import { clipContext, ConversationalChain } from './conversational_chain'; -import { FakeListChatModel, FakeListLLM } from '@langchain/core/utils/testing'; +import { FakeListChatModel, FakeStreamingLLM } from '@langchain/core/utils/testing'; import { BaseChatModel } from '@langchain/core/language_models/chat_models'; import { experimental_StreamData, Message } from 'ai'; import { ChatPromptTemplate } from '@langchain/core/prompts'; @@ -75,7 +75,7 @@ describe('conversational chain', () => { ? new FakeListChatModel({ responses, }) - : new FakeListLLM({ responses }); + : new FakeStreamingLLM({ responses }); const aiClient = Assist({ es_client: mockElasticsearchClient as unknown as Client, From 98752735f4d8648e101e7600570d23a06e3a40a3 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 6 Jun 2024 09:17:48 +0200 Subject: [PATCH 10/47] Bump --- package.json | 8 +-- yarn.lock | 138 ++++++++++++++++++++++----------------------------- 2 files changed, 62 insertions(+), 84 deletions(-) diff --git a/package.json b/package.json index f1fe4bdb88532..ab841c0c23f3b 100644 --- a/package.json +++ b/package.json @@ -935,10 +935,10 @@ "@paralleldrive/cuid2": "^2.2.2", "@reduxjs/toolkit": "1.9.7", "@slack/webhook": "^7.0.1", - "@smithy/eventstream-codec": "^2.0.12", - "@smithy/eventstream-serde-node": "^2.1.1", - "@smithy/types": "^2.9.1", - "@smithy/util-utf8": "^2.0.0", + "@smithy/eventstream-codec": "^3.0.0", + "@smithy/eventstream-serde-node": "3.0.0", + "@smithy/types": "^3.0.0", + "@smithy/util-utf8": "^3.0.0", "@tanstack/react-query": "^4.29.12", "@tanstack/react-query-devtools": "^4.29.12", "@turf/along": "6.0.1", diff --git a/yarn.lock b/yarn.lock index b8f012f0d144e..99f6c5b91ade7 100644 --- a/yarn.lock +++ b/yarn.lock @@ -96,12 +96,12 @@ tslib "^1.11.1" "@aws-sdk/types@^3.222.0": - version "3.433.0" - resolved "https://registry.yarnpkg.com/@aws-sdk/types/-/types-3.433.0.tgz#0f94eae2a4a3525ca872c9ab04e143c01806d755" - integrity sha512-0jEE2mSrNDd8VGFjTc1otYrwYPIkzZJEIK90ZxisKvQ/EURGBhNzWn7ejWB9XCMFT6XumYLBR0V9qq5UPisWtA== + version "3.577.0" + resolved "https://registry.yarnpkg.com/@aws-sdk/types/-/types-3.577.0.tgz#7700784d368ce386745f8c340d9d68cea4716f90" + integrity sha512-FT2JZES3wBKN/alfmhlo+3ZOq/XJ0C7QOZcDNrpKjB0kqYoKjhVKZ/Hx6ArR0czkKfHzBBEs6y40ebIHx2nSmA== dependencies: - "@smithy/types" "^2.4.0" - tslib "^2.5.0" + "@smithy/types" "^3.0.0" + tslib "^2.6.2" "@aws-sdk/util-utf8-browser@^3.0.0": version "3.259.0" @@ -6749,7 +6749,7 @@ zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@0.2.3", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@~0.2.0": +"@langchain/core@0.2.3", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@^0.2.3", "@langchain/core@~0.2.0": version "0.2.3" resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.3.tgz#7faa82f92b0c7843506e827a38bfcbb60f009d13" integrity sha512-mVuFHSLpPQ4yOHNXeoSA3LnmIMuFmUiit5rvbYcPZqM6SrB2zCNN2nD4Ty5+3H5X4tYItDoSqsTuUNUQySXRQw== @@ -6775,7 +6775,7 @@ "@langchain/core" ">0.1.61 <0.3.0" uuid "^9.0.1" -"@langchain/openai@^0.0.34": +"@langchain/openai@^0.0.34", "@langchain/openai@~0.0.28": version "0.0.34" resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.34.tgz#36c9bca0721ab9f7e5d40927e7c0429cacbd5b56" integrity sha512-M+CW4oXle5fdoz2T2SwdOef8pl3/1XmUx1vjn2mXUVM/128aO0l23FMF0SNBsAbRV6P+p/TuzjodchJbi0Ht/A== @@ -6786,17 +6786,6 @@ zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/openai@~0.0.28": - version "0.0.33" - resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.33.tgz#af88d815ff0095018c879d3a1a5a32b2795b5c69" - integrity sha512-hTBo9y9bHtFvMT5ySBW7TrmKhLSA91iNahigeqAFBVrLmBDz+6rzzLFc1mpq6JEAR3fZKdaUXqso3nB23jIpTw== - dependencies: - "@langchain/core" ">0.1.56 <0.3.0" - js-tiktoken "^1.0.12" - openai "^4.41.1" - zod "^3.22.4" - zod-to-json-schema "^3.22.3" - "@langchain/textsplitters@~0.0.0": version "0.0.2" resolved "https://registry.yarnpkg.com/@langchain/textsplitters/-/textsplitters-0.0.2.tgz#500baa8341fb7fc86fca531a4192665a319504a3" @@ -7975,70 +7964,70 @@ "@types/node" ">=18.0.0" axios "^1.6.0" -"@smithy/eventstream-codec@^2.0.12", "@smithy/eventstream-codec@^2.1.1": - version "2.1.1" - resolved "https://registry.yarnpkg.com/@smithy/eventstream-codec/-/eventstream-codec-2.1.1.tgz#4405ab0f9c77d439c575560c4886e59ee17d6d38" - integrity sha512-E8KYBxBIuU4c+zrpR22VsVrOPoEDzk35bQR3E+xm4k6Pa6JqzkDOdMyf9Atac5GPNKHJBdVaQ4JtjdWX2rl/nw== +"@smithy/eventstream-codec@^3.0.0": + version "3.0.0" + resolved "https://registry.yarnpkg.com/@smithy/eventstream-codec/-/eventstream-codec-3.0.0.tgz#81d30391220f73d41f432f65384b606d67673e46" + integrity sha512-PUtyEA0Oik50SaEFCZ0WPVtF9tz/teze2fDptW6WRXl+RrEenH8UbEjudOz8iakiMl3lE3lCVqYf2Y+znL8QFQ== dependencies: "@aws-crypto/crc32" "3.0.0" - "@smithy/types" "^2.9.1" - "@smithy/util-hex-encoding" "^2.1.1" - tslib "^2.5.0" + "@smithy/types" "^3.0.0" + "@smithy/util-hex-encoding" "^3.0.0" + tslib "^2.6.2" -"@smithy/eventstream-serde-node@^2.1.1": - version "2.1.1" - resolved "https://registry.yarnpkg.com/@smithy/eventstream-serde-node/-/eventstream-serde-node-2.1.1.tgz#2e1afa27f9c7eb524c1c53621049c5e4e3cea6a5" - integrity sha512-LF882q/aFidFNDX7uROAGxq3H0B7rjyPkV6QDn6/KDQ+CG7AFkRccjxRf1xqajq/Pe4bMGGr+VKAaoF6lELIQw== +"@smithy/eventstream-serde-node@3.0.0": + version "3.0.0" + resolved "https://registry.yarnpkg.com/@smithy/eventstream-serde-node/-/eventstream-serde-node-3.0.0.tgz#6519523fbb429307be29b151b8ba35bcca2b6e64" + integrity sha512-baRPdMBDMBExZXIUAoPGm/hntixjt/VFpU6+VmCyiYJYzRHRxoaI1MN+5XE+hIS8AJ2GCHLMFEIOLzq9xx1EgQ== dependencies: - "@smithy/eventstream-serde-universal" "^2.1.1" - "@smithy/types" "^2.9.1" - tslib "^2.5.0" + "@smithy/eventstream-serde-universal" "^3.0.0" + "@smithy/types" "^3.0.0" + tslib "^2.6.2" -"@smithy/eventstream-serde-universal@^2.1.1": - version "2.1.1" - resolved "https://registry.yarnpkg.com/@smithy/eventstream-serde-universal/-/eventstream-serde-universal-2.1.1.tgz#0f5eec9ad033017973a67bafb5549782499488d2" - integrity sha512-LR0mMT+XIYTxk4k2fIxEA1BPtW3685QlqufUEUAX1AJcfFfxNDKEvuCRZbO8ntJb10DrIFVJR9vb0MhDCi0sAQ== +"@smithy/eventstream-serde-universal@^3.0.0": + version "3.0.0" + resolved "https://registry.yarnpkg.com/@smithy/eventstream-serde-universal/-/eventstream-serde-universal-3.0.0.tgz#cb8441a73fbde4cbaa68e4a21236f658d914a073" + integrity sha512-HNFfShmotWGeAoW4ujP8meV9BZavcpmerDbPIjkJbxKbN8RsUcpRQ/2OyIxWNxXNH2GWCAxuSB7ynmIGJlQ3Dw== dependencies: - "@smithy/eventstream-codec" "^2.1.1" - "@smithy/types" "^2.9.1" - tslib "^2.5.0" + "@smithy/eventstream-codec" "^3.0.0" + "@smithy/types" "^3.0.0" + tslib "^2.6.2" -"@smithy/is-array-buffer@^2.0.0": - version "2.0.0" - resolved "https://registry.yarnpkg.com/@smithy/is-array-buffer/-/is-array-buffer-2.0.0.tgz#8fa9b8040651e7ba0b2f6106e636a91354ff7d34" - integrity sha512-z3PjFjMyZNI98JFRJi/U0nGoLWMSJlDjAW4QUX2WNZLas5C0CmVV6LJ01JI0k90l7FvpmixjWxPFmENSClQ7ug== +"@smithy/is-array-buffer@^3.0.0": + version "3.0.0" + resolved "https://registry.yarnpkg.com/@smithy/is-array-buffer/-/is-array-buffer-3.0.0.tgz#9a95c2d46b8768946a9eec7f935feaddcffa5e7a" + integrity sha512-+Fsu6Q6C4RSJiy81Y8eApjEB5gVtM+oFKTffg+jSuwtvomJJrhUJBu2zS8wjXSgH/g1MKEWrzyChTBe6clb5FQ== dependencies: - tslib "^2.5.0" + tslib "^2.6.2" -"@smithy/types@^2.4.0", "@smithy/types@^2.9.1": - version "2.9.1" - resolved "https://registry.yarnpkg.com/@smithy/types/-/types-2.9.1.tgz#ed04d4144eed3b8bd26d20fc85aae8d6e357ebb9" - integrity sha512-vjXlKNXyprDYDuJ7UW5iobdmyDm6g8dDG+BFUncAg/3XJaN45Gy5RWWWUVgrzIK7S4R1KWgIX5LeJcfvSI24bw== +"@smithy/types@^3.0.0": + version "3.0.0" + resolved "https://registry.yarnpkg.com/@smithy/types/-/types-3.0.0.tgz#00231052945159c64ffd8b91e8909d8d3006cb7e" + integrity sha512-VvWuQk2RKFuOr98gFhjca7fkBS+xLLURT8bUjk5XQoV0ZLm7WPwWPPY3/AwzTLuUBDeoKDCthfe1AsTUWaSEhw== dependencies: - tslib "^2.5.0" + tslib "^2.6.2" -"@smithy/util-buffer-from@^2.0.0": - version "2.0.0" - resolved "https://registry.yarnpkg.com/@smithy/util-buffer-from/-/util-buffer-from-2.0.0.tgz#7eb75d72288b6b3001bc5f75b48b711513091deb" - integrity sha512-/YNnLoHsR+4W4Vf2wL5lGv0ksg8Bmk3GEGxn2vEQt52AQaPSCuaO5PM5VM7lP1K9qHRKHwrPGktqVoAHKWHxzw== +"@smithy/util-buffer-from@^3.0.0": + version "3.0.0" + resolved "https://registry.yarnpkg.com/@smithy/util-buffer-from/-/util-buffer-from-3.0.0.tgz#559fc1c86138a89b2edaefc1e6677780c24594e3" + integrity sha512-aEOHCgq5RWFbP+UDPvPot26EJHjOC+bRgse5A8V3FSShqd5E5UN4qc7zkwsvJPPAVsf73QwYcHN1/gt/rtLwQA== dependencies: - "@smithy/is-array-buffer" "^2.0.0" - tslib "^2.5.0" + "@smithy/is-array-buffer" "^3.0.0" + tslib "^2.6.2" -"@smithy/util-hex-encoding@^2.1.1": - version "2.1.1" - resolved "https://registry.yarnpkg.com/@smithy/util-hex-encoding/-/util-hex-encoding-2.1.1.tgz#978252b9fb242e0a59bae4ead491210688e0d15f" - integrity sha512-3UNdP2pkYUUBGEXzQI9ODTDK+Tcu1BlCyDBaRHwyxhA+8xLP8agEKQq4MGmpjqb4VQAjq9TwlCQX0kP6XDKYLg== +"@smithy/util-hex-encoding@^3.0.0": + version "3.0.0" + resolved "https://registry.yarnpkg.com/@smithy/util-hex-encoding/-/util-hex-encoding-3.0.0.tgz#32938b33d5bf2a15796cd3f178a55b4155c535e6" + integrity sha512-eFndh1WEK5YMUYvy3lPlVmYY/fZcQE1D8oSf41Id2vCeIkKJXPcYDCZD+4+xViI6b1XSd7tE+s5AmXzz5ilabQ== dependencies: - tslib "^2.5.0" + tslib "^2.6.2" -"@smithy/util-utf8@^2.0.0": - version "2.0.0" - resolved "https://registry.yarnpkg.com/@smithy/util-utf8/-/util-utf8-2.0.0.tgz#b4da87566ea7757435e153799df9da717262ad42" - integrity sha512-rctU1VkziY84n5OXe3bPNpKR001ZCME2JCaBBFgtiM2hfKbHFudc/BkMuPab8hRbLd0j3vbnBTTZ1igBf0wgiQ== +"@smithy/util-utf8@^3.0.0": + version "3.0.0" + resolved "https://registry.yarnpkg.com/@smithy/util-utf8/-/util-utf8-3.0.0.tgz#1a6a823d47cbec1fd6933e5fc87df975286d9d6a" + integrity sha512-rUeT12bxFnplYDe815GXbq/oixEGHfRFFtcTF3YdDi/JaENIM6aSYYLJydG83UNzLXeRI5K8abYd/8Sp/QM0kA== dependencies: - "@smithy/util-buffer-from" "^2.0.0" - tslib "^2.5.0" + "@smithy/util-buffer-from" "^3.0.0" + tslib "^2.6.2" "@statoscope/extensions@5.28.1": version "5.28.1" @@ -21278,7 +21267,7 @@ kuler@^2.0.0: resolved "https://registry.yarnpkg.com/kuler/-/kuler-2.0.0.tgz#e2c570a3800388fb44407e851531c1d670b061b3" integrity sha512-Xq9nH7KlWZmXAtodXDDRE7vs6DU1gTU8zYDHDiWLSip45Egwq3plLHzPn27NgvzL2r1LMPC1vdqh98sQxtqj4A== -langchain@0.2.3: +langchain@0.2.3, langchain@^0.2.3: version "0.2.3" resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.3.tgz#c14bb05cf871b21bd63b84b3ab89580b1d62539f" integrity sha512-T9xR7zd+Nj0oXy6WoYKmZLy0DlQiDLFPGYWdOXDxy+AvqlujoPdVQgDSpdqiOHvAjezrByAoKxoHCz5XMwTP/Q== @@ -21305,7 +21294,7 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.30: +langsmith@^0.1.30, langsmith@~0.1.1, langsmith@~0.1.7: version "0.1.30" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== @@ -21316,17 +21305,6 @@ langsmith@^0.1.30: p-retry "4" uuid "^9.0.0" -langsmith@~0.1.1, langsmith@~0.1.7: - version "0.1.28" - resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.28.tgz#fbe01352d0b993fd11d4085dd337b1cec17ef28d" - integrity sha512-IQUbo7I7rEE6QYBhrcgwqvlkcUsHlia0yTQpDwWdITw/VJx1f7gLPjNdbwWE+jvOZ4HcD7gCf2HR6zFXputu5A== - dependencies: - "@types/uuid" "^9.0.1" - commander "^10.0.1" - p-queue "^6.6.2" - p-retry "4" - uuid "^9.0.0" - language-subtag-registry@~0.3.2: version "0.3.21" resolved "https://registry.yarnpkg.com/language-subtag-registry/-/language-subtag-registry-0.3.21.tgz#04ac218bea46f04cb039084602c6da9e788dd45a" @@ -29905,7 +29883,7 @@ tslib@^1.10.0, tslib@^1.11.1, tslib@^1.8.1, tslib@^1.9.0, tslib@^1.9.3: resolved "https://registry.yarnpkg.com/tslib/-/tslib-1.14.1.tgz#cf2d38bdc34a134bcaf1091c41f6619e2f672d00" integrity sha512-Xni35NKzjgMrwevysHTCArtLDpPvye8zV/0E4EyYn43P7/7qvQwPh9BGkHewbMulVntbigmcT7rdX3BNo9wRJg== -tslib@^2.0.0, tslib@^2.0.1, tslib@^2.0.3, tslib@^2.1.0, tslib@^2.3.1, tslib@^2.4.0, tslib@^2.5.0, tslib@^2.5.2: +tslib@^2.0.0, tslib@^2.0.1, tslib@^2.0.3, tslib@^2.1.0, tslib@^2.3.1, tslib@^2.4.0, tslib@^2.5.0, tslib@^2.5.2, tslib@^2.6.2: version "2.6.2" resolved "https://registry.yarnpkg.com/tslib/-/tslib-2.6.2.tgz#703ac29425e7b37cd6fd456e92404d46d1f3e4ae" integrity sha512-AEYxH93jGFPn/a2iVAwW87VuUIkR1FVUKB77NwMF7nBTDkDrrT/Hpt/IrCJ0QXhW27jTBDcf5ZY7w6RiqTMw2Q== From 5ba88d8e6b56ca4995f5af7a539f8d47dab20c22 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 10 Jun 2024 15:06:03 +0200 Subject: [PATCH 11/47] WIP --- .../impl/capabilities/index.ts | 6 +- .../conversation_sidepanel/index.tsx | 2 +- .../impl/assistant/index.tsx | 16 +- .../assistant/use_conversation/helpers.ts | 29 ++- .../application/connector/methods/get/get.ts | 2 + .../execute_custom_llm_chain/index.ts | 2 + .../server/lib/langchain/executors/types.ts | 2 + .../graphs/default_assistant_graph/_graph.ts | 245 ++++++++++++++++++ .../graphs/default_assistant_graph/index.ts | 2 + .../server/routes/attack_discovery/helpers.ts | 4 + .../attack_discovery/post_attack_discovery.ts | 1 + .../routes/post_actions_connector_execute.ts | 4 +- .../server/routes/request_context_factory.ts | 2 + .../plugins/elastic_assistant/server/types.ts | 4 + .../public/functions/index.ts | 4 +- .../public/functions/visualize_esql.tsx | 15 +- .../public/index.ts | 2 + .../common/experimental_features.ts | 2 +- x-pack/plugins/security_solution/kibana.jsonc | 3 +- .../custom_codeblock_markdown_plugin.tsx | 37 ++- .../custom_codeblock/esql_code_block.tsx | 160 ++++++++++++ .../public/assistant/get_comments/index.tsx | 5 + .../assistant/get_comments/stream/index.tsx | 10 +- .../get_comments/stream/message_text.tsx | 15 +- .../scripts/run_cypress/parallel.ts | 5 +- .../esql_language_knowledge_base_tool copy.ts | 125 +++++++++ .../esql_language_knowledge_base_tool.ts | 111 +++++++- .../test/security_solution_cypress/config.ts | 3 +- 28 files changed, 766 insertions(+), 52 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/_graph.ts create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool copy.ts diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index 9c734cc4b3c13..8449ba2a3b912 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -14,6 +14,8 @@ export type AssistantFeatures = { [K in keyof typeof defaultAssistantFeatures]: * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: false, - assistantModelEvaluation: false, + assistantKnowledgeBaseByDefault: true, + assistantModelEvaluation: true, }); + +// console.log('a') diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_sidepanel/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_sidepanel/index.tsx index 2af44aa21acb6..d22a1e13d468a 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_sidepanel/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_sidepanel/index.tsx @@ -212,7 +212,7 @@ export const ConversationSidePanel = React.memo( onClick: () => setDeleteConversationItem(conversation), iconType: 'trash', iconSize: 's', - disabled: conversation.isDefault, + // disabled: conversation.isDefault, 'aria-label': i18n.DELETE_CONVERSATION_ARIA_LABEL, 'data-test-subj': 'delete-option', }} diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx index 18f2c6be2a863..bd71d41f9b155 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -219,6 +219,7 @@ const AssistantComponent: React.FC = ({ [conversations, currentConversation?.id, getConversation] ); + const esqlInlineEditRef = useRef(); useEffect(() => { if (conversationsLoaded && Object.keys(conversations).length > 0) { setCurrentConversation((prev) => { @@ -586,6 +587,7 @@ const AssistantComponent: React.FC = ({ setIsStreaming, currentUserAvatar, isFlyoutMode, + esqlInlineEditRef, })} {...(!isFlyoutMode ? { @@ -634,12 +636,11 @@ const AssistantComponent: React.FC = ({ ), [ + getComments, abortStream, - refetchCurrentConversation, currentConversation, - editingSystemPromptId, - getComments, showAnonymizedValues, + refetchCurrentConversation, handleRegenerateResponse, isEnabledKnowledgeBase, isEnabledRAGAlerts, @@ -647,6 +648,7 @@ const AssistantComponent: React.FC = ({ currentUserAvatar, isFlyoutMode, selectedPromptContextsCount, + editingSystemPromptId, isNewConversation, isSettingsModalVisible, promptContexts, @@ -1089,6 +1091,14 @@ const AssistantComponent: React.FC = ({ + +
+ ); } diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts index de766085e1aee..fd5280f0f7046 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts @@ -20,6 +20,20 @@ export interface CodeBlockDetails { export type QueryType = 'eql' | 'esql' | 'kql' | 'dsl' | 'json' | 'no-type' | 'sql'; +// If your codeblocks aren't getting tagged with the right language, add keywords to the array. +export const MARKDOWN_TYPES = { + eql: ['Event Query Language', 'EQL sequence query', 'EQL'], + esql: ['Elasticsearch Query Language', 'ESQL', 'ES|QL', 'SQL'], + kql: ['Kibana Query Language', 'KQL Query', 'KQL'], + dsl: [ + 'Elasticsearch QueryDSL', + 'Elasticsearch Query DSL', + 'Elasticsearch DSL', + 'Query DSL', + 'DSL', + ], +}; + /** * `analyzeMarkdown` is a helper that enriches content returned from a query * with action buttons @@ -32,26 +46,13 @@ export type QueryType = 'eql' | 'esql' | 'kql' | 'dsl' | 'json' | 'no-type' | 's export const analyzeMarkdown = (markdown: string): CodeBlockDetails[] => { const codeBlockRegex = /```(\w+)?\s([\s\S]*?)```/g; const matches = [...markdown.matchAll(codeBlockRegex)]; - // If your codeblocks aren't getting tagged with the right language, add keywords to the array. - const types = { - eql: ['Event Query Language', 'EQL sequence query', 'EQL'], - esql: ['Elasticsearch Query Language', 'ESQL', 'ES|QL', 'SQL'], - kql: ['Kibana Query Language', 'KQL Query', 'KQL'], - dsl: [ - 'Elasticsearch QueryDSL', - 'Elasticsearch Query DSL', - 'Elasticsearch DSL', - 'Query DSL', - 'DSL', - ], - }; const result: CodeBlockDetails[] = matches.map((match) => { let type = match[1] || 'no-type'; if (type === 'no-type' || type === 'json') { const start = match.index || 0; const precedingText = markdown.slice(0, start); - for (const [typeKey, keywords] of Object.entries(types)) { + for (const [typeKey, keywords] of Object.entries(MARKDOWN_TYPES)) { if (keywords.some((kw) => precedingText.toLowerCase().includes(kw.toLowerCase()))) { type = typeKey; break; diff --git a/x-pack/plugins/actions/server/application/connector/methods/get/get.ts b/x-pack/plugins/actions/server/application/connector/methods/get/get.ts index 2d4a94f5615d7..46a2cb3821ff2 100644 --- a/x-pack/plugins/actions/server/application/connector/methods/get/get.ts +++ b/x-pack/plugins/actions/server/application/connector/methods/get/get.ts @@ -61,7 +61,9 @@ export async function get({ connector = { id, actionTypeId: foundInMemoryConnector.actionTypeId, + isMissingSecrets: foundInMemoryConnector.isMissingSecrets, name: foundInMemoryConnector.name, + config: foundInMemoryConnector.config, isPreconfigured: foundInMemoryConnector.isPreconfigured, isSystemAction: foundInMemoryConnector.isSystemAction, isDeprecated: isConnectorDeprecated(foundInMemoryConnector), diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts index 7ea1e5fb3c9b9..42f4604003b93 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts @@ -48,6 +48,7 @@ export const callAgentExecutor: AgentExecutor = async ({ request, size, traceOptions, + search, }) => { const isOpenAI = llmType === 'openai'; const llmClass = isOpenAI ? ActionsClientChatOpenAI : ActionsClientSimpleChatModel; @@ -101,6 +102,7 @@ export const callAgentExecutor: AgentExecutor = async ({ replacements, request, size, + search, }; const tools: ToolInterface[] = assistantTools.flatMap( diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index bd07099e312b3..543d02400a0e8 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -14,6 +14,7 @@ import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain'; import { ExecuteConnectorRequestBody, Message, Replacements } from '@kbn/elastic-assistant-common'; import { StreamResponseWithHeaders } from '@kbn/ml-response-stream/server'; import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; +import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; import { ResponseBody } from '../types'; import type { AssistantTool } from '../../../types'; import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store'; @@ -44,6 +45,7 @@ export interface AgentExecutorParams { conversationId?: string; dataClients?: AssistantDataClients; esClient: ElasticsearchClient; + search: ReturnType; esStore: ElasticsearchStore; langChainMessages: BaseMessage[]; llmType?: string; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/_graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/_graph.ts new file mode 100644 index 0000000000000..fa5b0302d4db0 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/_graph.ts @@ -0,0 +1,245 @@ +/* + * 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. + */ + +// /* +// * 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 { forOwn, get } from 'lodash'; +// import { ToolExecutor } from '@langchain/langgraph/prebuilt'; +// import { RunnableConfig, RunnableLambda } from '@langchain/core/runnables'; +// import { END, START, StateGraph, StateGraphArgs } from '@langchain/langgraph'; +// import { AgentAction, AgentFinish, AgentStep } from '@langchain/core/agents'; +// import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; +// import { StructuredTool } from '@langchain/core/tools'; +// import type { CompiledStateGraph } from '@langchain/langgraph/dist/graph/state'; +// import type { Logger } from '@kbn/logging'; + +// import { BaseMessage } from '@langchain/core/messages'; +// import { BaseChatModel } from '@langchain/core/language_models/chat_models'; +// import { ElasticsearchClient, IScopedClusterClient } from '@kbn/core-elasticsearch-server'; +// import { getESQLQueryColumns } from '@kbn/esql-utils'; +// import { AgentState, NodeParamsBase } from './types'; +// import { generateChatTitle } from './generate_chat_title'; +// import { AssistantDataClients } from '../../executors/types'; + +// export const DEFAULT_ASSISTANT_GRAPH_ID = 'Default Security Assistant Graph'; + +// interface GetDefaultAssistantGraphParams { +// agentRunnable: AgentRunnableSequence; +// dataClients?: AssistantDataClients; +// conversationId?: string; +// llm: BaseChatModel; +// logger: Logger; +// messages: BaseMessage[]; +// tools: StructuredTool[]; +// esClient: ElasticsearchClient; +// search: IScopedClusterClient; +// } + +// /** +// * Returns a compiled default assistant graph +// */ +// export const getDefaultAssistantGraph = ({ +// agentRunnable, +// conversationId, +// dataClients, +// llm, +// logger, +// messages, +// tools, +// esClient, +// search, +// }: GetDefaultAssistantGraphParams): CompiledStateGraph => { +// try { +// // Default graph state +// const graphState: StateGraphArgs['channels'] = { +// input: { +// value: (x: string, y?: string) => y ?? x, +// default: () => '', +// }, +// steps: { +// value: (x: AgentStep[], y: AgentStep[]) => x.concat(y), +// default: () => [], +// }, +// agentOutcome: { +// value: ( +// x: AgentAction | AgentFinish | undefined, +// y?: AgentAction | AgentFinish | undefined +// ) => y ?? x, +// default: () => undefined, +// }, +// messages: { +// value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y), +// default: () => messages, +// }, +// }; + +// const nodeParams: NodeParamsBase = { +// model: llm, +// logger, +// }; + +// // Create a tool executor +// const toolExecutor = new ToolExecutor({ tools }); + +// // Define logic that will be used to determine which conditional edge to go down +// const shouldContinue = (state: AgentState) => { +// logger.error(`graph:shouldContinue:state\n${JSON.stringify(state, null, 2)}`); +// if (state.agentOutcome && 'returnValues' in state.agentOutcome) { +// return 'end'; +// } +// return 'continue'; +// }; + +// const runAgent = async (state: AgentState, config?: RunnableConfig) => { +// logger.error(`graph:runAgent:\nstate\n${JSON.stringify(state, null, 2)}`); + +// const agentOutcome = await agentRunnable.invoke( +// { +// ...state, +// chat_history: messages, +// knowledge_history: 'The users favorite color is blue', // TODO: Plumb through initial retrieval +// }, +// config +// ); +// return { +// agentOutcome, +// }; +// }; + +// const executeTools = async (state: AgentState, config?: RunnableConfig) => { +// logger.error(`graph:executeTools:state\n${JSON.stringify(state, null, 2)}`); +// const agentAction = state.agentOutcome; +// if (!agentAction || 'returnValues' in agentAction) { +// throw new Error('Agent has not been run yet'); +// } +// const out = await toolExecutor.invoke(agentAction, config); +// return { +// steps: [{ action: agentAction, observation: JSON.stringify(out, null, 2) }], +// }; +// }; + +// const workflow = new StateGraph({ +// channels: { +// input: { +// value: (x: string, y?: string) => y ?? x, +// default: () => '', +// }, +// indices: { +// value: (x: string, y?: string) => y ?? x, +// default: () => [], +// }, +// }, +// }); + +// workflow.addNode('listIndicesEsDb', async (state: AgentState) => { +// console.error('listIndicesEsDb, state', state); +// const response = await esClient.indices.get({ index: '*' }); + +// console.error('response', JSON.stringify(response, null, 2)); + +// function transformInputToOutput(input: Record) { +// const output: Record = {}; + +// forOwn(input, (value, _key) => { +// const key = Object.keys(value.aliases)[0] || _key; +// const properties = get(value, 'mappings.properties', {}); +// output[key] = transformProperties(properties); +// }); + +// return output; +// } + +// function transformProperties(properties) { +// const result: Record = {}; + +// forOwn(properties, (value, key) => { +// if (value.type) { +// result[key] = value.type; +// } else if (value.properties) { +// result[key] = transformProperties(value.properties); +// } else if (value.fields) { +// result[key] = { fields: {} }; +// forOwn(value.fields, (fieldValue, fieldKey) => { +// result[key].fields[fieldKey] = { type: fieldValue.type }; +// if (fieldValue.ignore_above !== undefined) { +// result[key].fields[fieldKey].ignore_above = fieldValue.ignore_above; +// } +// }); +// } else { +// result[key] = value; +// } +// }); + +// return result; +// } + +// return { +// ...state, +// indices: transformInputToOutput(response), +// }; +// }); + +// workflow.addNode('generateESQLQuery', async (state: AgentState) => { +// const ecsDuplicatesPrompt = ECS_DUPLICATES_PROMPT; +// const outputParser = new JsonOutputParser(); +// const ecsDuplicatesGraph = ecsDuplicatesPrompt.pipe(model) + +// const currentMapping = await ecsDuplicatesGraph.invoke({ +// indices: state.indices, +// }); + +// return { +// ...state, +// }; +// }); + +// workflow.addNode('getColumns', async (state: AgentState) => { +// const columns = getESQLQueryColumns({ esqlQuery: state.query, search }); + +// console.error('columns', columns); + +// return { +// ...state, +// columns, +// }; +// }); +// // // Create a new graph, with the default state from above +// // const workflow = new StateGraph({ channels: graphState }); + +// // // Define the nodes to cycle between +// // workflow.addNode('generateChatTitle', (state: AgentState) => +// // generateChatTitle({ +// // state, +// // conversationsDataClient: dataClients?.conversationsDataClient, +// // conversationId, +// // ...nodeParams, +// // }) +// // ); +// // workflow.addNode('agent', new RunnableLambda({ func: runAgent })); +// // workflow.addNode('action', new RunnableLambda({ func: executeTools })); + +// // // Add conditional edge for determining if we shouldContinue +// // workflow.addConditionalEdges('agent', shouldContinue, { continue: 'action', end: END }); + +// // // Add edges for start, and between agent and action (action always followed by agent) +// workflow.addEdge(START, 'listIndicesEsDb'); +// workflow.addEdge('listIndicesEsDb', 'generateESQLQuery'); +// workflow.addEdge('generateESQLQuery', 'getColumns'); +// workflow.addEdge('getColumns', END); +// // workflow.addEdge('generateChatTitle', 'agent'); +// // workflow.addEdge('action', 'agent'); + +// return workflow.compile(); +// } catch (e) { +// throw new Error(`Unable to compile DefaultAssistantGraph\n${e}`); +// } +// }; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts index 1e40f6b2fe127..625ad960a27e0 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts @@ -45,6 +45,7 @@ export const callAssistantGraph: AgentExecutor = async ({ request, size, traceOptions, + search, }) => { const logger = parentLogger.get('defaultAssistantGraph'); const isOpenAI = llmType === 'openai'; @@ -93,6 +94,7 @@ export const callAssistantGraph: AgentExecutor = async ({ replacements, request, size, + search, }; const tools: StructuredTool[] = assistantTools.flatMap( diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts index 4f916be8105d4..6ce84574902e6 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts @@ -17,6 +17,7 @@ import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/s import { v4 as uuidv4 } from 'uuid'; import { ActionsClientLlm } from '@kbn/langchain/server'; +import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; import { AssistantToolParams } from '../../types'; export const REQUIRED_FOR_ATTACK_DISCOVERY: AnonymizationFieldResponse[] = [ @@ -45,6 +46,7 @@ export const getAssistantToolParams = ({ onNewReplacements, request, size, + search, }: { alertsIndexPattern: string; anonymizationFields?: AnonymizationFieldResponse[]; @@ -60,6 +62,7 @@ export const getAssistantToolParams = ({ ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody >; size: number; + search: DataPluginStart['search']['asScoped']; }): AssistantToolParams => ({ alertsIndexPattern, anonymizationFields: [...(anonymizationFields ?? []), ...REQUIRED_FOR_ATTACK_DISCOVERY], @@ -74,4 +77,5 @@ export const getAssistantToolParams = ({ replacements: latestReplacements, request, size, + search, // not required for attach discovery }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts index 7859d635ccb30..d8b27dd3d2767 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts @@ -131,6 +131,7 @@ export const postAttackDiscoveryRoute = ( onNewReplacements, request, size, + search: assistantContext.search, }); // invoke the attack discovery tool: diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 1ab0ef84c2574..80d9217fc622d 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -72,6 +72,7 @@ export const postActionsConnectorExecuteRoute = ( }, }, }, + // eslint-disable-next-line complexity async (context, request, response) => { const abortSignal = getRequestAbortedSignal(request.events.aborted$); @@ -186,7 +187,7 @@ export const postActionsConnectorExecuteRoute = ( model: request.body.model, messages: [ { - role: 'assistant', + role: 'system', content: i18n.translate( 'xpack.elasticAssistantPlugin.server.autoTitlePromptDescription', { @@ -370,6 +371,7 @@ export const postActionsConnectorExecuteRoute = ( conversationId, dataClients, esClient, + search: assistantContext.search, esStore, isStream: request.body.subAction !== 'invokeAI', llmType: getLlmType(actionTypeId), diff --git a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts index 0a0864882df16..1e5082a0b72dc 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -63,6 +63,8 @@ export class RequestContextFactory implements IRequestContextFactory { actions: startPlugins.actions, + search: startPlugins.data.search.asScoped(request), + logger: this.logger, getServerBasePath: () => core.http.basePath.serverBasePath, diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index f12bacde983df..3014f7e61e93d 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -39,6 +39,7 @@ import { ActionsClientSimpleChatModel, } from '@kbn/langchain/server'; +import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/conversations'; import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; import { AIAssistantDataClient } from './ai_assistant_data_clients'; @@ -97,6 +98,7 @@ export interface ElasticAssistantPluginSetupDependencies { } export interface ElasticAssistantPluginStartDependencies { actions: ActionsPluginStart; + data: DataPluginStart; spaces?: SpacesPluginStart; security: SecurityPluginStart; } @@ -104,6 +106,7 @@ export interface ElasticAssistantPluginStartDependencies { export interface ElasticAssistantApiRequestHandlerContext { core: CoreRequestHandlerContext; actions: ActionsPluginStart; + search: ReturnType; getRegisteredFeatures: GetRegisteredFeatures; getRegisteredTools: GetRegisteredTools; logger: Logger; @@ -225,4 +228,5 @@ export interface AssistantToolParams { ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody >; size?: number; + search: ReturnType; } diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/index.ts index 0dd8fc2eed38e..ba3e16ca79712 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/index.ts @@ -9,7 +9,9 @@ import type { RegisterRenderFunctionDefinition } from '@kbn/observability-ai-ass import type { ObservabilityAIAssistantAppPluginStartDependencies } from '../types'; import { registerChangesRenderFunction } from './changes'; import { registerLensRenderFunction } from './lens'; -import { registerVisualizeQueryRenderFunction } from './visualize_esql'; +import { registerVisualizeQueryRenderFunction, VisualizeESQL } from './visualize_esql'; + +export { VisualizeESQL }; export async function registerFunctions({ registerRenderFunction, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx index 334db28a8b89a..1cc93868cdf9e 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx @@ -112,6 +112,7 @@ export function VisualizeESQL({ preferredChartType, ObservabilityAIAssistantMultipaneFlyoutContext, errorMessages, + esqlInlineEditRef, }: VisualizeESQLProps) { // fetch the pattern from the query const indexPattern = getIndexPatternFromESQLQuery(query); @@ -230,10 +231,19 @@ export function VisualizeESQL({ ReactDOM.unmountComponentAtNode(chatFlyoutSecondSlotHandler.container); } }, - container: chatFlyoutSecondSlotHandler?.container, + container: esqlInlineEditRef ?? chatFlyoutSecondSlotHandler?.container, }; } - }, [chatFlyoutSecondSlotHandler, lensInput, lensLoadEvent, onActionClick, query]); + }, [ + chatFlyoutSecondSlotHandler, + esqlInlineEditRef, + lensInput, + lensLoadEvent, + onActionClick, + query, + ]); + + console.error('triggerOptions', triggerOptions); if (!lensHelpersAsync.value || !dataViewAsync.value || !lensInput) { return ; @@ -279,6 +289,7 @@ export function VisualizeESQL({ onClick={() => { chatFlyoutSecondSlotHandler?.setVisibility?.(true); if (triggerOptions) { + console.error('triggerOptions', triggerOptions); uiActions.getTrigger('IN_APP_EMBEDDABLE_EDIT_TRIGGER').exec(triggerOptions); } }} diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/index.ts index 5de1c30de7c4c..dcb7116b6458b 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/index.ts @@ -14,6 +14,8 @@ import type { ObservabilityAIAssistantAppPublicStart, } from './types'; +export { VisualizeESQL } from './functions'; + export const plugin: PluginInitializer< ObservabilityAIAssistantAppPublicSetup, ObservabilityAIAssistantAppPublicStart, diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index a3248bc4374ba..b8f7d7a1e3021 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -133,7 +133,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables the Assistant Knowledge Base by default, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, /** * Enables the new user details flyout displayed on the Alerts table. diff --git a/x-pack/plugins/security_solution/kibana.jsonc b/x-pack/plugins/security_solution/kibana.jsonc index 14741ec4c37e0..5b2ba81fc5e3e 100644 --- a/x-pack/plugins/security_solution/kibana.jsonc +++ b/x-pack/plugins/security_solution/kibana.jsonc @@ -53,7 +53,8 @@ "notifications", "savedSearch", "unifiedDocViewer", - "charts" + "charts", + "observabilityAIAssistantApp" ], "optionalPlugins": [ "cloudExperiments", diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_codeblock_markdown_plugin.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_codeblock_markdown_plugin.tsx index 19f566537a2b6..8825974f9c9c6 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_codeblock_markdown_plugin.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_codeblock_markdown_plugin.tsx @@ -6,7 +6,8 @@ */ import type { Node } from 'unist'; -import type { Parent } from 'mdast'; +import type { Parent, Content, PhrasingContent } from 'mdast'; +import { MARKDOWN_TYPES } from '@kbn/elastic-assistant/impl/assistant/use_conversation/helpers'; export const customCodeBlockLanguagePlugin = () => { const visitor = (node: Node, parent?: Parent) => { @@ -17,14 +18,32 @@ export const customCodeBlockLanguagePlugin = () => { }); } - if ( - node.type === 'code' && - (node.lang === 'eql' || - node.lang === 'esql' || - node.lang === 'kql' || - node.lang === 'dsl' || - node.lang === 'json') - ) { + if (node.type === 'code' && !node.lang) { + try { + const index = parent?.children.indexOf(node as Content); + + if (index) { + const previousText = (parent?.children[index - 1]?.children as PhrasingContent[]) + ?.map((child) => child.value) + .join(' '); + for (const [typeKey, keywords] of Object.entries(MARKDOWN_TYPES)) { + if (keywords.some((kw) => previousText.toLowerCase().includes(kw.toLowerCase()))) { + node.lang = typeKey; + break; + } + } + } + } catch (e) { + /* empty */ + } + } + + if (node.type === 'code' && node.lang === 'esql') { + node.type = 'esql'; + return; + } + + if (node.type === 'code' && ['eql', 'kql', 'dsl', 'json'].includes(node.lang as string)) { node.type = 'customCodeBlock'; } }; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx new file mode 100644 index 0000000000000..b5abf8154837f --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -0,0 +1,160 @@ +/* + * 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 { + EuiButtonEmpty, + EuiCodeBlock, + EuiFlexGroup, + EuiFlexItem, + EuiPanel, + useEuiTheme, +} from '@elastic/eui'; +import { css } from '@emotion/css'; +import type { UseChatSend } from '@kbn/elastic-assistant/impl/assistant/chat_send/use_chat_send'; +import { i18n } from '@kbn/i18n'; +import React, { useCallback, useEffect } from 'react'; +import { VisualizeESQL } from '@kbn/observability-ai-assistant-app-plugin/public'; +import { lastValueFrom } from 'rxjs'; +import { useQuery } from '@tanstack/react-query'; +import { useKibana } from '../../../common/lib/kibana'; + +export function EsqlCodeBlock({ + value, + actionsDisabled, + handleSendMessage, + ...rest +}: { + value: string; + actionsDisabled: boolean; + handleSendMessage: UseChatSend['handleSendMessage']; +}) { + console.error('value', value); + console.error('rest', rest); + const { lens, dataViews, uiActions, data } = useKibana().services; + console.error('sss', useKibana(), useKibana().services.expressions.getTypes()); + const theme = useEuiTheme(); + + const { data: columns } = useQuery({ + queryFn: async () => { + return lastValueFrom( + data.search.search( + { + params: { + query: value, + version: '2024.04.01', + dropNullColumns: true, + }, + }, + { + strategy: 'esql_async', + isSearchStored: false, + } + ) + ); + }, + select: (dataz) => { + console.error('dataz', dataz); + return dataz.rawResponse.columns.map((column) => ({ + id: column.name, + name: column.name, + meta: { + type: column.type === 'long' ? 'number' : column.type, + }, + })); + }, + keepPreviousData: true, + }); + + console.error('columns', columns, [ + { + id: 'count', + name: 'count', + meta: { + type: 'number', + }, + }, + { + id: 'minute', + name: 'minute', + meta: { + type: 'date', + }, + }, + ]); + + return ( + <> + + + + + {value} + + + + + + {} + // onActionClick({ type: ChatActionClickType.executeEsqlQuery, query: value }) + } + disabled={actionsDisabled} + > + {i18n.translate('xpack.observabilityAiAssistant.runThisQuery', { + defaultMessage: 'Display results', + })} + + + + + {i18n.translate('xpack.observabilityAiAssistant.visualizeThisQuery', { + defaultMessage: 'Visualize this query', + })} + + + + + + + {}} + ObservabilityAIAssistantMultipaneFlyoutContext={{}} + query={value.trim()} + esqlInlineEditRef={rest.esqlInlineEditRef.current} + /> + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx index 10d5a15c800ae..67df2e93ecbc9 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx @@ -63,6 +63,7 @@ export const getComments = ({ isFlyoutMode, currentUserAvatar, setIsStreaming, + esqlInlineEditRef, }: { abortStream: () => void; currentConversation?: Conversation; @@ -83,6 +84,8 @@ export const getComments = ({ // should only happen when no apiConfig is present const actionTypeId = currentConversation.apiConfig?.actionTypeId ?? ''; + console.error('esqlInlineEditRef', esqlInlineEditRef); + const extraLoadingComment = isFetchingResponse ? [ { @@ -177,6 +180,7 @@ export const getComments = ({ regenerateMessage={regenerateMessageOfConversation} setIsStreaming={setIsStreaming} transformMessage={transformMessage} + esqlInlineEditRef={esqlInlineEditRef} /> ), }; @@ -202,6 +206,7 @@ export const getComments = ({ refetchCurrentConversation={refetchCurrentConversation} setIsStreaming={setIsStreaming} transformMessage={transformMessage} + esqlInlineEditRef={esqlInlineEditRef} /> ), }; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx index 631d8b507ceed..181ee317c9e98 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx @@ -44,6 +44,7 @@ export const StreamComment = ({ regenerateMessage, setIsStreaming, transformMessage, + esqlInlineEditRef, }: Props) => { const { error, isLoading, isStreaming, pendingMessage, setComplete } = useStream({ refetchCurrentConversation, @@ -108,7 +109,14 @@ export const StreamComment = ({ return ( } + body={ + + } error={error ? new Error(error) : undefined} controls={controls} /> diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx index c40b0c04043ad..d1a048940daf8 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx @@ -25,6 +25,7 @@ import { euiThemeVars } from '@kbn/ui-theme'; import type { Node } from 'unist'; import { customCodeBlockLanguagePlugin } from '../custom_codeblock/custom_codeblock_markdown_plugin'; import { CustomCodeBlock } from '../custom_codeblock/custom_code_block'; +import { EsqlCodeBlock } from '../custom_codeblock/esql_code_block'; interface Props { content: string; @@ -98,7 +99,7 @@ const loadingCursorPlugin = () => { }; }; -const getPluginDependencies = () => { +const getPluginDependencies = ({ esqlInlineEditRef }) => { const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); const processingPlugins = getDefaultEuiMarkdownProcessingPlugins(); @@ -108,6 +109,14 @@ const getPluginDependencies = () => { processingPlugins[1][1].components = { ...components, cursor: Cursor, + esql: (props) => { + return ( + <> + + + + ); + }, customCodeBlock: (props) => { return ( <> @@ -143,12 +152,12 @@ const getPluginDependencies = () => { }; }; -export function MessageText({ loading, content, index }: Props) { +export function MessageText({ loading, content, index, esqlInlineEditRef }: Props) { const containerClassName = css` overflow-wrap: anywhere; `; - const { parsingPluginList, processingPluginList } = getPluginDependencies(); + const { parsingPluginList, processingPluginList } = getPluginDependencies({ esqlInlineEditRef }); return ( diff --git a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts index 4b0f586779270..aec694c14b09c 100644 --- a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts +++ b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts @@ -337,10 +337,7 @@ ${JSON.stringify( procs, config, installDir: options?.installDir, - extraKbnOpts: - options?.installDir || options?.ci || !isOpen - ? [] - : ['--dev', '--no-dev-config', '--no-dev-credentials'], + extraKbnOpts: options?.installDir || options?.ci || !isOpen ? [] : ['--dev'], onEarlyExit, inspect: argv.inspect, }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool copy.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool copy.ts new file mode 100644 index 0000000000000..3e8326cc54aae --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool copy.ts @@ -0,0 +1,125 @@ +/* + * 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 { DynamicStructuredTool } from '@langchain/core/tools'; +import { z } from 'zod'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { forOwn, get } from 'lodash'; +import { JsonOutputParser } from '@langchain/core/output_parsers'; +import { getESQLQueryColumns } from '@kbn/esql-utils'; +import { APP_UI_ID } from '../../../../common'; + +export type EsqlKnowledgeBaseToolParams = AssistantToolParams; + +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. Only output valid ES|QL queries as described above. Do not add any additional text to describe your output.', + id: 'esql-knowledge-base-tool', + name: 'ESQLKnowledgeBaseTool', +}; +export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { + ...toolDetails, + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is EsqlKnowledgeBaseToolParams => { + const { chain, isEnabledKnowledgeBase, modelExists } = params; + return isEnabledKnowledgeBase && modelExists && chain != null; + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + + const { chain, esClient, search } = params as EsqlKnowledgeBaseToolParams; + if (chain == null) return null; + + console.error('search', search); + + return new DynamicStructuredTool({ + name: toolDetails.name, + description: toolDetails.description, + schema: z.object({ + question: z.string().describe(`The user's exact question about ESQL`), + }), + func: async (input, _, cbManager) => { + console.error('input.questions', input); + const response = await esClient.indices.get({ index: '*' }); + + console.error('response', JSON.stringify(response, null, 2)); + + function transformInputToOutput(input2: Record) { + const output: Record = {}; + + forOwn(input2, (value, _key) => { + const key = Object.keys(value.aliases)[0] || _key; + const properties = get(value, 'mappings.properties', {}); + output[key] = transformProperties(properties); + }); + + return output; + } + + function transformProperties(properties) { + const result: Record = {}; + + forOwn(properties, (value, key) => { + if (value.type) { + result[key] = value.type; + } else if (value.properties) { + result[key] = transformProperties(value.properties); + } else if (value.fields) { + result[key] = { fields: {} }; + forOwn(value.fields, (fieldValue, fieldKey) => { + result[key].fields[fieldKey] = { type: fieldValue.type }; + if (fieldValue.ignore_above !== undefined) { + result[key].fields[fieldKey].ignore_above = fieldValue.ignore_above; + } + }); + } else { + result[key] = value; + } + }); + + return result; + } + + const indices = transformInputToOutput(response); + + const result = await chain.invoke( + { + query: ` + CONTEXT:\`\`\` + Available indices: ${JSON.stringify(Object.keys(indices), null, 2)} + Object where the key is the name of the index and the value is the mapping: ${JSON.stringify( + indices + )}. + \`\`\` + + ${input.question} + `, + }, + cbManager + ); + + console.error('result', JSON.stringify(result, null, 2)); + + const esqlQuery = result.text + .match(/(?<=""")[\s\S]*?(?=""")/g) + .join('') + .replace('\n', '') + .trim(); + + try { + const esqlQueryColumns = await getESQLQueryColumns({ esqlQuery, search }); + console.error('esqlQueryColumns', esqlQueryColumns); + } catch (e) { + console.error('e', e); + } + + return result.text; + }, + tags: ['esql', 'query-generation', 'knowledge-base'], + }); + }, +}; 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 692753a22dea0..7a5e18680a89d 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 @@ -8,8 +8,41 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from 'zod'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { forOwn, get } from 'lodash'; +import { JsonOutputParser, StringOutputParser } from '@langchain/core/output_parsers'; +import { getESQLQueryColumns } from '@kbn/esql-utils'; +import { ChatPromptTemplate } from '@langchain/core/prompts'; import { APP_UI_ID } from '../../../../common'; +export const ECS_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ + [ + 'system', + `As an expert user of Elastic Security, please generate an accurate and valid ESQL query to detect the use case below. Your response should be formatted to be able to use immediately in an Elastic Security timeline or detection rule. Take your time with the answer, check your knowledge really well on all the functions I am asking for. For ES|QL answers specifically, you should only ever answer with what's available in your private knowledge. I cannot afford for queries to be inaccurate. Assume I am using the Elastic Common Schema and Elastic Agent.. + +Here is some context for you to reference for your task, read it carefully as you will get questions about it later: + + +{indicesMapping} + +`, + ], + [ + 'human', + `{input}. + +Example response format: + +A: Please find the ESQL query below: +\`\`\`esql +FROM logs +| SORT @timestamp DESC +| LIMIT 5 +\`\`\` +"`, + ], + ['ai', 'Please find the ESQL query below:'], +]); + export type EsqlKnowledgeBaseToolParams = AssistantToolParams; const toolDetails = { @@ -28,9 +61,13 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { chain } = params as EsqlKnowledgeBaseToolParams; + const { chain, esClient, search, llm } = params as EsqlKnowledgeBaseToolParams; if (chain == null) return null; + console.error('params', Object.keys(params)); + + console.error('search', search); + return new DynamicStructuredTool({ name: toolDetails.name, description: toolDetails.description, @@ -38,13 +75,71 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { question: z.string().describe(`The user's exact question about ESQL`), }), func: async (input, _, cbManager) => { - const result = await chain.invoke( - { - query: input.question, - }, - cbManager - ); - return result.text; + const response = await esClient.indices.get({ index: '*' }); + + function transformInputToOutput(input2: Record) { + const output: Record = {}; + + forOwn(input2, (value, _key) => { + const key = Object.keys(value.aliases)[0] || _key; + const properties = get(value, 'mappings.properties', {}); + output[key] = transformProperties(properties); + }); + + return output; + } + + function transformProperties(properties) { + const result: Record = {}; + + forOwn(properties, (value, key) => { + if (value.type) { + result[key] = value.type; + } else if (value.properties) { + result[key] = transformProperties(value.properties); + } else if (value.fields) { + result[key] = { fields: {} }; + forOwn(value.fields, (fieldValue, fieldKey) => { + result[key].fields[fieldKey] = { type: fieldValue.type }; + if (fieldValue.ignore_above !== undefined) { + result[key].fields[fieldKey].ignore_above = fieldValue.ignore_above; + } + }); + } else { + result[key] = value; + } + }); + + return result; + } + + const indices = transformInputToOutput(response); + + const graph = ECS_MAIN_PROMPT.pipe(llm).pipe(new StringOutputParser()); + + const result = await graph.invoke({ + input: input.question, + indicesMapping: JSON.stringify(Object.keys(indices), null, 2), + }); + + console.error('result', JSON.stringify(result, null, 2)); + + const esqlQuery = result + ?.match(/(?<=```esql)[\s\S]*?(?=```)/g) + .join('') + .replace('\n', '') + .trim(); + + console.error('esqlQuery', esqlQuery); + + try { + const esqlQueryColumns = await getESQLQueryColumns({ esqlQuery, search: search.search }); + console.error('esqlQueryColumns', esqlQueryColumns); + } catch (e) { + console.error('e', e); + } + + return result; }, tags: ['esql', 'query-generation', 'knowledge-base'], }); diff --git a/x-pack/test/security_solution_cypress/config.ts b/x-pack/test/security_solution_cypress/config.ts index 606df49c4f90e..99817362af73d 100644 --- a/x-pack/test/security_solution_cypress/config.ts +++ b/x-pack/test/security_solution_cypress/config.ts @@ -47,9 +47,10 @@ export default async function ({ readConfigFile }: FtrConfigProviderContext) { `--xpack.securitySolution.enableExperimental=${JSON.stringify([ 'alertSuppressionForEsqlRuleEnabled', 'bulkCustomHighlightedFieldsEnabled', + 'assistantKnowledgeBaseByDefault', ])}`, // mock cloud to enable the guided onboarding tour in e2e tests - '--xpack.cloud.id=test', + // '--xpack.cloud.id=test', `--home.disableWelcomeScreen=true`, // Specify which version of the detection-rules package to install // `--xpack.securitySolution.prebuiltRulesPackageVersion=8.3.1`, From aa97cc07939e3c6eb73d911c5ec4c74418215472 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 14 Jun 2024 11:40:55 +0200 Subject: [PATCH 12/47] WIP --- .../assistant/assistant_overlay/index.tsx | 65 ++-- .../impl/assistant/index.tsx | 10 - .../graphs/default_assistant_graph/graph.ts | 21 +- .../public/functions/visualize_esql.tsx | 14 +- x-pack/plugins/security_solution/kibana.jsonc | 4 +- .../assistant/content/quick_prompts/index.tsx | 8 - .../custom_codeblock/esql_code_block.tsx | 338 +++++++++++++----- .../public/assistant/get_comments/index.tsx | 5 - .../assistant/get_comments/stream/index.tsx | 10 +- .../get_comments/stream/message_text.tsx | 8 +- .../esql_language_knowledge_base_tool copy.ts | 125 ------- .../esql_language_knowledge_base_tool.ts | 174 ++++++++- 12 files changed, 473 insertions(+), 309 deletions(-) delete mode 100644 x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool copy.ts diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_overlay/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_overlay/index.tsx index 44907d8b1fd00..bc0d484de49e7 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_overlay/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_overlay/index.tsx @@ -6,11 +6,10 @@ */ import React, { useCallback, useEffect, useRef, useState } from 'react'; -import { EuiModal, EuiFlyoutResizable, useEuiTheme } from '@elastic/eui'; - +import { EuiModal, EuiFlyoutResizable } from '@elastic/eui'; import useEvent from 'react-use/lib/useEvent'; // eslint-disable-next-line @kbn/eslint/module_migration -import styled from 'styled-components'; +import styled, { createGlobalStyle } from 'styled-components'; import { css } from '@emotion/react'; import { ShowAssistantOverlayProps, @@ -37,8 +36,13 @@ export interface Props { currentUserAvatar?: UserAvatar; } +export const UnifiedTimelineGlobalStyles = createGlobalStyle` + body:has(.timeline-portal-overlay-mask) .euiOverlayMask { + z-index: 1003 !important; + } +`; + export const AssistantOverlay = React.memo(({ isFlyoutMode, currentUserAvatar }) => { - const { euiTheme } = useEuiTheme(); const [isModalVisible, setIsModalVisible] = useState(false); const [conversationTitle, setConversationTitle] = useState( WELCOME_CONVERSATION_TITLE @@ -132,32 +136,33 @@ export const AssistantOverlay = React.memo(({ isFlyoutMode, currentUserAv if (isFlyoutMode) { return ( - div { - height: 100%; - } - `} - onClose={handleCloseModal} - data-test-subj="ai-assistant-flyout" - paddingSize="none" - hideCloseButton - // EUI TODO: This z-index override of EuiOverlayMask is a workaround, and ideally should be resolved with a cleaner UI/UX flow long-term - maskProps={{ style: `z-index: ${(euiTheme.levels.flyout as number) + 3}` }} // we need this flyout to be above the timeline flyout (which has a z-index of 1002) - > - - + <> + div { + height: 100%; + } + `} + onClose={handleCloseModal} + data-test-subj="ai-assistant-flyout" + paddingSize="none" + hideCloseButton + > + + + + ); } diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx index 279ebddbfece7..2af6d8396c0ca 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -219,7 +219,6 @@ const AssistantComponent: React.FC = ({ [conversations, currentConversation?.id, getConversation] ); - const esqlInlineEditRef = useRef(); useEffect(() => { if (conversationsLoaded && Object.keys(conversations).length > 0) { setCurrentConversation((prev) => { @@ -593,7 +592,6 @@ const AssistantComponent: React.FC = ({ setIsStreaming, currentUserAvatar, isFlyoutMode, - esqlInlineEditRef, })} {...(!isFlyoutMode ? { @@ -1092,14 +1090,6 @@ const AssistantComponent: React.FC = ({ - -
- ); } diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts index 779bf20a61720..7f657892afa7d 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts @@ -86,14 +86,19 @@ export const getDefaultAssistantGraph = ({ logger: logger.get(AGENT_NODE), state, }); - const executeToolsNode = (state: AgentState, config?: RunnableConfig) => - executeTools({ - ...nodeParams, - config, - logger: logger.get(TOOLS_NODE), - state, - tools, - }); + const executeToolsNode = (state: AgentState, config?: RunnableConfig) => { + try { + return executeTools({ + ...nodeParams, + config, + logger: logger.get(TOOLS_NODE), + state, + tools, + }); + } catch (e) { + console.error('dupa', e); + } + }; const shouldContinueEdge = (state: AgentState) => shouldContinue({ ...nodeParams, state }); // Put together a new graph using the nodes and default state from above diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx index a716631e27164..97c46a6a2d759 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx @@ -116,7 +116,6 @@ export function VisualizeESQL({ preferredChartType, ObservabilityAIAssistantMultipaneFlyoutContext, errorMessages, - esqlInlineEditRef, }: VisualizeESQLProps) { // fetch the pattern from the query const indexPattern = getIndexPatternFromESQLQuery(query); @@ -242,19 +241,10 @@ export function VisualizeESQL({ ReactDOM.unmountComponentAtNode(chatFlyoutSecondSlotHandler.container); } }, - container: esqlInlineEditRef ?? chatFlyoutSecondSlotHandler?.container, + container: chatFlyoutSecondSlotHandler?.container, }; } - }, [ - chatFlyoutSecondSlotHandler, - esqlInlineEditRef, - lensInput, - lensLoadEvent, - onActionClick, - query, - ]); - - console.error('triggerOptions', triggerOptions); + }, [chatFlyoutSecondSlotHandler, lensInput, lensLoadEvent, onActionClick, query]); if (!lensHelpersAsync.value || !dataViewAsync.value || !lensInput) { return ; diff --git a/x-pack/plugins/security_solution/kibana.jsonc b/x-pack/plugins/security_solution/kibana.jsonc index 5b2ba81fc5e3e..7ca152629adb3 100644 --- a/x-pack/plugins/security_solution/kibana.jsonc +++ b/x-pack/plugins/security_solution/kibana.jsonc @@ -25,6 +25,7 @@ "ecsDataQualityDashboard", "elasticAssistant", "embeddable", + "esqlDataGrid", "eventLog", "features", "fieldFormats", @@ -53,8 +54,7 @@ "notifications", "savedSearch", "unifiedDocViewer", - "charts", - "observabilityAIAssistantApp" + "charts" ], "optionalPlugins": [ "cloudExperiments", diff --git a/x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx b/x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx index 799087f202e98..c35598f8898ba 100644 --- a/x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx @@ -8,7 +8,6 @@ import type { QuickPrompt } from '@kbn/elastic-assistant'; import * as i18n from './translations'; import { - KNOWLEDGE_BASE_CATEGORY, PROMPT_CONTEXT_ALERT_CATEGORY, PROMPT_CONTEXT_DETECTION_RULES_CATEGORY, PROMPT_CONTEXT_EVENT_CATEGORY, @@ -27,13 +26,6 @@ export const BASE_SECURITY_QUICK_PROMPTS: QuickPrompt[] = [ categories: [PROMPT_CONTEXT_ALERT_CATEGORY], isDefault: true, }, - { - title: i18n.ESQL_QUERY_GENERATION_TITLE, - prompt: i18n.ESQL_QUERY_GENERATION_PROMPT, - color: '#9170B8', - categories: [KNOWLEDGE_BASE_CATEGORY], - isDefault: true, - }, { title: i18n.RULE_CREATION_TITLE, prompt: i18n.RULE_CREATION_PROMPT, diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index b5abf8154837f..c58d565e1907a 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -6,21 +6,43 @@ */ import { EuiButtonEmpty, + EuiButtonIcon, EuiCodeBlock, EuiFlexGroup, EuiFlexItem, EuiPanel, + EuiToolTip, useEuiTheme, } from '@elastic/eui'; import { css } from '@emotion/css'; import type { UseChatSend } from '@kbn/elastic-assistant/impl/assistant/chat_send/use_chat_send'; import { i18n } from '@kbn/i18n'; -import React, { useCallback, useEffect } from 'react'; -import { VisualizeESQL } from '@kbn/observability-ai-assistant-app-plugin/public'; -import { lastValueFrom } from 'rxjs'; +import React, { useEffect, useState } from 'react'; import { useQuery } from '@tanstack/react-query'; +import { v4 as uuidv4 } from 'uuid'; +import { + getIndexPatternFromESQLQuery, + getESQLQueryColumns, + getESQLAdHocDataview, + getESQLResults, +} from '@kbn/esql-utils'; +import useAsync from 'react-use/lib/useAsync'; +import { ESQLDataGrid } from '@kbn/esql-datagrid/public'; +import type { TypedLensByValueInput } from '@kbn/lens-plugin/public'; +import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils'; import { useKibana } from '../../../common/lib/kibana'; +function generateId() { + return uuidv4(); +} + +const saveVisualizationLabel = i18n.translate( + 'xpack.observabilityAiAssistant.lensESQLFunction.save', + { + defaultMessage: 'Save visualization', + } +); + export function EsqlCodeBlock({ value, actionsDisabled, @@ -31,60 +53,112 @@ export function EsqlCodeBlock({ actionsDisabled: boolean; handleSendMessage: UseChatSend['handleSendMessage']; }) { - console.error('value', value); - console.error('rest', rest); - const { lens, dataViews, uiActions, data } = useKibana().services; - console.error('sss', useKibana(), useKibana().services.expressions.getTypes()); + const { lens, dataViews: dataViewService, uiActions, data } = useKibana().services; const theme = useEuiTheme(); - const { data: columns } = useQuery({ + const lensHelpersAsync = useAsync(() => { + return lens.stateHelperApi(); + }, [lens]); + + const { data: queryResults } = useQuery({ + queryKey: ['test'], + enabled: true, queryFn: async () => { - return lastValueFrom( - data.search.search( - { - params: { - query: value, - version: '2024.04.01', - dropNullColumns: true, - }, - }, - { - strategy: 'esql_async', - isSearchStored: false, - } - ) - ); + return getESQLResults({ + esqlQuery: value, + // esqlQuery: `FROM logs-* | STATS avg_bytes = AVG(http.response.body.bytes) BY host.name`, + search: data.search.search, + }); }, select: (dataz) => { - console.error('dataz', dataz); - return dataz.rawResponse.columns.map((column) => ({ - id: column.name, - name: column.name, - meta: { - type: column.type === 'long' ? 'number' : column.type, - }, - })); + return { + params: dataz.params, + rows: dataz.response.values, + columns: dataz.response.columns, + }; }, + refetchOnWindowFocus: false, keepPreviousData: true, }); - console.error('columns', columns, [ - { - id: 'count', - name: 'count', - meta: { - type: 'number', - }, - }, - { - id: 'minute', - name: 'minute', - meta: { - type: 'date', - }, - }, + // const rows = getESQLResults({ esqlQuery: value, search: data.search.search }); + + const indexPattern = getIndexPatternFromESQLQuery(value); + const formattedColumns = useAsync( + () => + getESQLQueryColumns({ + esqlQuery: value, + search: data.search.search, + }), + [value] + ); + + const dataViewAsync = useAsync(() => { + return getESQLAdHocDataview(indexPattern, dataViewService).then((dataView) => { + if (dataView.fields.getByName('@timestamp')?.type === 'date') { + dataView.timeFieldName = '@timestamp'; + } + return dataView; + }); + }, [indexPattern, dataViewService]); + + const [isSaveModalOpen, setIsSaveModalOpen] = useState(false); + const [isTableVisible, setIsTableVisible] = useState(false); + const [lensInput, setLensInput] = useState(undefined); + const [showVisualization, setShowVisualization] = useState(false); + + const preferredChartType = undefined; // 'XY'; + + // initialization + useEffect(() => { + if (lensHelpersAsync.value && dataViewAsync.value && !lensInput && formattedColumns.value) { + const context = { + dataViewSpec: dataViewAsync.value?.toSpec(), + fieldName: '', + textBasedColumns: formattedColumns.value, + query: { + esql: value, + }, + }; + + const chartSuggestions = lensHelpersAsync.value.suggestions( + context, + dataViewAsync.value, + [], + preferredChartType + ); + + if (chartSuggestions?.length) { + const [suggestion] = chartSuggestions; + + const attrs = getLensAttributesFromSuggestion({ + filters: [], + query: { + esql: value, + }, + suggestion, + dataView: dataViewAsync.value, + }) as TypedLensByValueInput['attributes']; + + const lensEmbeddableInput = { + attributes: attrs, + id: generateId(), + }; + setLensInput(lensEmbeddableInput); + } + } + }, [ + dataViewAsync.value, + formattedColumns.value, + lensHelpersAsync.value, + lensInput, + preferredChartType, + value, ]); + // if the Lens suggestions api suggests a table then we want to render a Discover table instead + const isLensInputTable = lensInput?.attributes?.visualizationType === 'lnsDatatable'; + return ( <> - - + + {!showVisualization && ( + + + + setShowVisualization(true)} + disabled={actionsDisabled} + > + {i18n.translate('xpack.observabilityAiAssistant.visualizeThisQuery', { + defaultMessage: 'Generate Visualization', + })} + + + + + )} + + + + {showVisualization && queryResults && formattedColumns.value && ( + <> + {!isLensInputTable && ( + <> - {} - // onActionClick({ type: ChatActionClickType.executeEsqlQuery, query: value }) - } - disabled={actionsDisabled} - > - {i18n.translate('xpack.observabilityAiAssistant.runThisQuery', { - defaultMessage: 'Display results', - })} - + + + + setIsTableVisible(!isTableVisible)} + data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton" + aria-label={ + isTableVisible + ? i18n.translate( + 'xpack.observabilityAiAssistant.lensESQLFunction.displayChart', + { + defaultMessage: 'Display chart', + } + ) + : i18n.translate( + 'xpack.observabilityAiAssistant.lensESQLFunction.displayTable', + { + defaultMessage: 'Display table', + } + ) + } + /> + + + + + setIsSaveModalOpen(true)} + data-test-subj="observabilityAiAssistantLensESQLSaveButton" + aria-label={saveVisualizationLabel} + /> + + + - - - {i18n.translate('xpack.observabilityAiAssistant.visualizeThisQuery', { - defaultMessage: 'Visualize this query', - })} - + + {isTableVisible ? ( + + ) : ( + + )} - - - - - {}} - ObservabilityAIAssistantMultipaneFlyoutContext={{}} - query={value.trim()} - esqlInlineEditRef={rest.esqlInlineEditRef.current} - /> + + )} + {/* hide the grid in case of errors (as the user can't fix them) */} + {isLensInputTable && ( + + + + )} + + )} + + {isSaveModalOpen ? ( + { + setIsSaveModalOpen(() => false); + }} + // For now, we don't want to allow saving ESQL charts to the library + isSaveable={false} + /> + ) : null} ); } diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx index 67df2e93ecbc9..10d5a15c800ae 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx @@ -63,7 +63,6 @@ export const getComments = ({ isFlyoutMode, currentUserAvatar, setIsStreaming, - esqlInlineEditRef, }: { abortStream: () => void; currentConversation?: Conversation; @@ -84,8 +83,6 @@ export const getComments = ({ // should only happen when no apiConfig is present const actionTypeId = currentConversation.apiConfig?.actionTypeId ?? ''; - console.error('esqlInlineEditRef', esqlInlineEditRef); - const extraLoadingComment = isFetchingResponse ? [ { @@ -180,7 +177,6 @@ export const getComments = ({ regenerateMessage={regenerateMessageOfConversation} setIsStreaming={setIsStreaming} transformMessage={transformMessage} - esqlInlineEditRef={esqlInlineEditRef} /> ), }; @@ -206,7 +202,6 @@ export const getComments = ({ refetchCurrentConversation={refetchCurrentConversation} setIsStreaming={setIsStreaming} transformMessage={transformMessage} - esqlInlineEditRef={esqlInlineEditRef} /> ), }; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx index 181ee317c9e98..631d8b507ceed 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx @@ -44,7 +44,6 @@ export const StreamComment = ({ regenerateMessage, setIsStreaming, transformMessage, - esqlInlineEditRef, }: Props) => { const { error, isLoading, isStreaming, pendingMessage, setComplete } = useStream({ refetchCurrentConversation, @@ -109,14 +108,7 @@ export const StreamComment = ({ return ( - } + body={} error={error ? new Error(error) : undefined} controls={controls} /> diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx index d1a048940daf8..60c83982a893c 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx @@ -99,7 +99,7 @@ const loadingCursorPlugin = () => { }; }; -const getPluginDependencies = ({ esqlInlineEditRef }) => { +const getPluginDependencies = () => { const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); const processingPlugins = getDefaultEuiMarkdownProcessingPlugins(); @@ -112,7 +112,7 @@ const getPluginDependencies = ({ esqlInlineEditRef }) => { esql: (props) => { return ( <> - + ); @@ -152,12 +152,12 @@ const getPluginDependencies = ({ esqlInlineEditRef }) => { }; }; -export function MessageText({ loading, content, index, esqlInlineEditRef }: Props) { +export function MessageText({ loading, content, index }: Props) { const containerClassName = css` overflow-wrap: anywhere; `; - const { parsingPluginList, processingPluginList } = getPluginDependencies({ esqlInlineEditRef }); + const { parsingPluginList, processingPluginList } = getPluginDependencies(); return ( diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool copy.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool copy.ts deleted file mode 100644 index 3e8326cc54aae..0000000000000 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool copy.ts +++ /dev/null @@ -1,125 +0,0 @@ -/* - * 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 { DynamicStructuredTool } from '@langchain/core/tools'; -import { z } from 'zod'; -import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; -import { forOwn, get } from 'lodash'; -import { JsonOutputParser } from '@langchain/core/output_parsers'; -import { getESQLQueryColumns } from '@kbn/esql-utils'; -import { APP_UI_ID } from '../../../../common'; - -export type EsqlKnowledgeBaseToolParams = AssistantToolParams; - -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. Only output valid ES|QL queries as described above. Do not add any additional text to describe your output.', - id: 'esql-knowledge-base-tool', - name: 'ESQLKnowledgeBaseTool', -}; -export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { - ...toolDetails, - sourceRegister: APP_UI_ID, - isSupported: (params: AssistantToolParams): params is EsqlKnowledgeBaseToolParams => { - const { chain, isEnabledKnowledgeBase, modelExists } = params; - return isEnabledKnowledgeBase && modelExists && chain != null; - }, - getTool(params: AssistantToolParams) { - if (!this.isSupported(params)) return null; - - const { chain, esClient, search } = params as EsqlKnowledgeBaseToolParams; - if (chain == null) return null; - - console.error('search', search); - - return new DynamicStructuredTool({ - name: toolDetails.name, - description: toolDetails.description, - schema: z.object({ - question: z.string().describe(`The user's exact question about ESQL`), - }), - func: async (input, _, cbManager) => { - console.error('input.questions', input); - const response = await esClient.indices.get({ index: '*' }); - - console.error('response', JSON.stringify(response, null, 2)); - - function transformInputToOutput(input2: Record) { - const output: Record = {}; - - forOwn(input2, (value, _key) => { - const key = Object.keys(value.aliases)[0] || _key; - const properties = get(value, 'mappings.properties', {}); - output[key] = transformProperties(properties); - }); - - return output; - } - - function transformProperties(properties) { - const result: Record = {}; - - forOwn(properties, (value, key) => { - if (value.type) { - result[key] = value.type; - } else if (value.properties) { - result[key] = transformProperties(value.properties); - } else if (value.fields) { - result[key] = { fields: {} }; - forOwn(value.fields, (fieldValue, fieldKey) => { - result[key].fields[fieldKey] = { type: fieldValue.type }; - if (fieldValue.ignore_above !== undefined) { - result[key].fields[fieldKey].ignore_above = fieldValue.ignore_above; - } - }); - } else { - result[key] = value; - } - }); - - return result; - } - - const indices = transformInputToOutput(response); - - const result = await chain.invoke( - { - query: ` - CONTEXT:\`\`\` - Available indices: ${JSON.stringify(Object.keys(indices), null, 2)} - Object where the key is the name of the index and the value is the mapping: ${JSON.stringify( - indices - )}. - \`\`\` - - ${input.question} - `, - }, - cbManager - ); - - console.error('result', JSON.stringify(result, null, 2)); - - const esqlQuery = result.text - .match(/(?<=""")[\s\S]*?(?=""")/g) - .join('') - .replace('\n', '') - .trim(); - - try { - const esqlQueryColumns = await getESQLQueryColumns({ esqlQuery, search }); - console.error('esqlQueryColumns', esqlQueryColumns); - } catch (e) { - console.error('e', e); - } - - return result.text; - }, - tags: ['esql', 'query-generation', 'knowledge-base'], - }); - }, -}; 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 7a5e18680a89d..d6a583b680f14 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 @@ -12,18 +12,22 @@ import { forOwn, get } from 'lodash'; import { JsonOutputParser, StringOutputParser } from '@langchain/core/output_parsers'; import { getESQLQueryColumns } from '@kbn/esql-utils'; import { ChatPromptTemplate } from '@langchain/core/prompts'; +import type { StateGraphArgs } from '@langchain/langgraph'; +import { StateGraph } from '@langchain/langgraph'; import { APP_UI_ID } from '../../../../common'; export const ECS_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ [ 'system', - `As an expert user of Elastic Security, please generate an accurate and valid ESQL query to detect the use case below. Your response should be formatted to be able to use immediately in an Elastic Security timeline or detection rule. Take your time with the answer, check your knowledge really well on all the functions I am asking for. For ES|QL answers specifically, you should only ever answer with what's available in your private knowledge. I cannot afford for queries to be inaccurate. Assume I am using the Elastic Common Schema and Elastic Agent.. + `As an expert user of Elastic Security, please generate an accurate and valid ESQL query to detect the use case below. Your response should be formatted to be able to use immediately in an Elastic Security timeline or detection rule. Take your time with the answer, check your knowledge really well on all the functions I am asking for. For ES|QL answers specifically, you should only ever answer with what's available in your private knowledge. I cannot afford for queries to be inaccurate. Assume I am using the Elastic Common Schema and Elastic Agent. Under any circumstances wrap index in quotes. + + If multiple indices are matched please try to use wildcard to match all indices. If you are unsure about the index name, please refer to the context provided below. Here is some context for you to reference for your task, read it carefully as you will get questions about it later: - -{indicesMapping} - + +{availableIndices} + `, ], [ @@ -51,6 +55,146 @@ const toolDetails = { id: 'esql-knowledge-base-tool', name: 'ESQLKnowledgeBaseTool', }; + +export interface CategorizationState { + rawSamples: string[]; + samples: string[]; + formattedSamples: string; + ecsTypes: string; + ecsCategories: string; + exAnswer: string; + lastExecutedChain: string; + packageName: string; + dataStreamName: string; + errors: object; + pipelineResults: object[]; + finalized: boolean; + reviewed: boolean; + currentPipeline: object; + currentProcessors: object[]; + invalidCategorization: object; + initialPipeline: object; + result: object; +} + +const graphState: StateGraphArgs['channels'] = { + lastExecutedChain: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + rawSamples: { + value: (x: string[], y?: string[]) => y ?? x, + default: () => [], + }, + samples: { + value: (x: string[], y?: string[]) => y ?? x, + default: () => [], + }, + formattedSamples: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + ecsTypes: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + ecsCategories: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + exAnswer: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + packageName: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + dataStreamName: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + finalized: { + value: (x: boolean, y?: boolean) => y ?? x, + default: () => false, + }, + reviewed: { + value: (x: boolean, y?: boolean) => y ?? x, + default: () => false, + }, + errors: { + value: (x: object, y?: object) => y ?? x, + default: () => ({}), + }, + pipelineResults: { + value: (x: object[], y?: object[]) => y ?? x, + default: () => [{}], + }, + currentPipeline: { + value: (x: object, y?: object) => y ?? x, + default: () => ({}), + }, + currentProcessors: { + value: (x: object[], y?: object[]) => y ?? x, + default: () => [], + }, + invalidCategorization: { + value: (x: object, y?: object) => y ?? x, + default: () => ({}), + }, + initialPipeline: { + value: (x: object, y?: object) => y ?? x, + default: () => ({}), + }, + result: { + value: (x: object, y?: object) => y ?? x, + default: () => ({}), + }, +}; + +function modelInput(state: CategorizationState): Partial { + // const samples = modifySamples(state); + // const formattedSamples = formatSamples(samples); + // const initialPipeline = JSON.parse(JSON.stringify(state.currentPipeline)); + return { + // exAnswer: JSON.stringify(CATEGORIZATION_EXAMPLE_ANSWER, null, 2), + // ecsCategories: JSON.stringify(ECS_CATEGORIES, null, 2), + // ecsTypes: JSON.stringify(ECS_TYPES, null, 2), + // samples, + // formattedSamples, + // initialPipeline, + finalized: false, + reviewed: false, + lastExecutedChain: 'modelInput', + }; +} + +function modelOutput(state: CategorizationState): Partial { + return { + finalized: true, + lastExecutedChain: 'modelOutput', + result: { + query: state.query, + rows: state.pipelineResults, + columns: state.currentPipeline, + }, + }; +} + +const handleGenerateQuery = (state: CategorizationState, model) => {}; + +const getEsqlGraph = (client, model) => { + const workflow = new StateGraph({ + channels: graphState, + }) + .addNode('modelInput', modelInput) + .addNode('modelOutput', modelOutput) + .addNode('handleGenerateQuery', (state: CategorizationState) => + handleGenerateQuery(state, model) + ) + .addNode('handleClassifyEsql', (state: CategorizationState) => {}); +}; + export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, @@ -75,7 +219,14 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { question: z.string().describe(`The user's exact question about ESQL`), }), func: async (input, _, cbManager) => { - const response = await esClient.indices.get({ index: '*' }); + let response; + try { + response = await esClient.indices.getDataStream(); + } catch (e) { + console.error('e', e); + } + + console.error('response', JSON.stringify(response, null, 2)); function transformInputToOutput(input2: Record) { const output: Record = {}; @@ -113,13 +264,17 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { return result; } - const indices = transformInputToOutput(response); + // const indices = transformInputToOutput(response); const graph = ECS_MAIN_PROMPT.pipe(llm).pipe(new StringOutputParser()); const result = await graph.invoke({ input: input.question, - indicesMapping: JSON.stringify(Object.keys(indices), null, 2), + availableIndices: JSON.stringify( + response?.data_streams.map((item) => item.name), + null, + 2 + ), }); console.error('result', JSON.stringify(result, null, 2)); @@ -128,6 +283,7 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { ?.match(/(?<=```esql)[\s\S]*?(?=```)/g) .join('') .replace('\n', '') + .replaceAll('"', '') .trim(); console.error('esqlQuery', esqlQuery); @@ -139,7 +295,9 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { console.error('e', e); } - return result; + console.error('trimee', result.replaceAll('"', '').trim()); + + return result.replaceAll('"', '').trim(); }, tags: ['esql', 'query-generation', 'knowledge-base'], }); From 825afe3d0d250bbe4b45adc80e935886ae88c1ec Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 19 Jun 2024 23:01:50 +0200 Subject: [PATCH 13/47] add cases --- .../assistant/use_conversation/helpers.ts | 5 +- x-pack/plugins/cases/kibana.jsonc | 2 + .../components/markdown_editor/editor.tsx | 4 +- .../custom_codeblock/custom_code_block.tsx | 45 +++ .../custom_codeblock_markdown_plugin.tsx | 54 +++ .../custom_codeblock/esql_code_block.tsx | 338 ++++++++++++++++++ .../components/markdown_editor/renderer.tsx | 10 +- .../scrollable_markdown_renderer.tsx | 10 +- .../components/markdown_editor/use_plugins.ts | 18 +- .../user_actions/comment/actions.tsx | 2 +- .../components/user_actions/comment/user.tsx | 1 + .../components/user_actions/markdown_form.tsx | 5 +- x-pack/plugins/cases/public/types.ts | 2 + .../custom_codeblock/esql_code_block.tsx | 36 +- .../public/assistant/get_comments/index.tsx | 1 + .../assistant/get_comments/stream/index.tsx | 11 +- .../get_comments/stream/message_text.tsx | 9 +- .../public/assistant/helpers.tsx | 5 +- 18 files changed, 521 insertions(+), 37 deletions(-) create mode 100644 x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_code_block.tsx create mode 100644 x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_codeblock_markdown_plugin.tsx create mode 100644 x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts index fd5280f0f7046..fc7f8aa8fefc8 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts @@ -12,6 +12,7 @@ import { Conversation } from '../../assistant_context/types'; export interface CodeBlockDetails { type: QueryType; content: string; + timestamp: string; start: number; end: number; getControlContainer?: () => Element | undefined; @@ -43,7 +44,7 @@ export const MARKDOWN_TYPES = { * * @param markdown */ -export const analyzeMarkdown = (markdown: string): CodeBlockDetails[] => { +export const analyzeMarkdown = (markdown: string, timestamp: string): CodeBlockDetails[] => { const codeBlockRegex = /```(\w+)?\s([\s\S]*?)```/g; const matches = [...markdown.matchAll(codeBlockRegex)]; @@ -63,7 +64,7 @@ export const analyzeMarkdown = (markdown: string): CodeBlockDetails[] => { const content = match[2].trim(); const start = match.index || 0; const end = start + match[0].length; - return { type: type as QueryType, content, start, end }; + return { type: type as QueryType, content, start, end, timestamp }; }); return result; diff --git a/x-pack/plugins/cases/kibana.jsonc b/x-pack/plugins/cases/kibana.jsonc index 84c04da1fe0f6..e4cee6a6d9c9c 100644 --- a/x-pack/plugins/cases/kibana.jsonc +++ b/x-pack/plugins/cases/kibana.jsonc @@ -15,7 +15,9 @@ "alerting", "actions", "data", + "dataViews", "embeddable", + "esqlDataGrid", "lens", "licensing", "features", diff --git a/x-pack/plugins/cases/public/components/markdown_editor/editor.tsx b/x-pack/plugins/cases/public/components/markdown_editor/editor.tsx index 92580b0e1da20..a1d7bc1cc0652 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/editor.tsx +++ b/x-pack/plugins/cases/public/components/markdown_editor/editor.tsx @@ -40,7 +40,9 @@ const MarkdownEditorComponent = forwardRef(null); useLensButtonToggle({ diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_code_block.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_code_block.tsx new file mode 100644 index 0000000000000..9b009c3f32240 --- /dev/null +++ b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_code_block.tsx @@ -0,0 +1,45 @@ +/* + * 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 { EuiCodeBlock, EuiFlexGroup, EuiFlexItem, EuiPanel, useEuiTheme } from '@elastic/eui'; +import { css } from '@emotion/css'; +import React from 'react'; + +const CustomCodeBlockComponent = ({ value }: { value: string }) => { + const theme = useEuiTheme(); + + return ( + + + + + {value} + + + + + ); +}; + +CustomCodeBlockComponent.displayName = 'CustomCodeBlock'; + +export const CustomCodeBlock = React.memo(CustomCodeBlockComponent); diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_codeblock_markdown_plugin.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_codeblock_markdown_plugin.tsx new file mode 100644 index 0000000000000..8825974f9c9c6 --- /dev/null +++ b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_codeblock_markdown_plugin.tsx @@ -0,0 +1,54 @@ +/* + * 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 type { Node } from 'unist'; +import type { Parent, Content, PhrasingContent } from 'mdast'; +import { MARKDOWN_TYPES } from '@kbn/elastic-assistant/impl/assistant/use_conversation/helpers'; + +export const customCodeBlockLanguagePlugin = () => { + const visitor = (node: Node, parent?: Parent) => { + if ('children' in node) { + const nodeAsParent = node as Parent; + nodeAsParent.children.forEach((child) => { + visitor(child, nodeAsParent); + }); + } + + if (node.type === 'code' && !node.lang) { + try { + const index = parent?.children.indexOf(node as Content); + + if (index) { + const previousText = (parent?.children[index - 1]?.children as PhrasingContent[]) + ?.map((child) => child.value) + .join(' '); + for (const [typeKey, keywords] of Object.entries(MARKDOWN_TYPES)) { + if (keywords.some((kw) => previousText.toLowerCase().includes(kw.toLowerCase()))) { + node.lang = typeKey; + break; + } + } + } + } catch (e) { + /* empty */ + } + } + + if (node.type === 'code' && node.lang === 'esql') { + node.type = 'esql'; + return; + } + + if (node.type === 'code' && ['eql', 'kql', 'dsl', 'json'].includes(node.lang as string)) { + node.type = 'customCodeBlock'; + } + }; + + return (tree: Node) => { + visitor(tree); + }; +}; diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx new file mode 100644 index 0000000000000..d8d44910813b6 --- /dev/null +++ b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx @@ -0,0 +1,338 @@ +/* + * 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 { + EuiButtonEmpty, + EuiButtonIcon, + EuiCodeBlock, + EuiFlexGroup, + EuiFlexItem, + EuiPanel, + EuiSpacer, + EuiToolTip, + useEuiTheme, +} from '@elastic/eui'; +import { css } from '@emotion/css'; +import { i18n } from '@kbn/i18n'; +import React, { useEffect, useState } from 'react'; +import { useQuery } from '@tanstack/react-query'; +import { v4 as uuidv4 } from 'uuid'; +import { + getIndexPatternFromESQLQuery, + getESQLQueryColumns, + getESQLAdHocDataview, + getESQLResults, +} from '@kbn/esql-utils'; +import useAsync from 'react-use/lib/useAsync'; +import { ESQLDataGrid } from '@kbn/esql-datagrid/public'; +import type { TypedLensByValueInput } from '@kbn/lens-plugin/public'; +import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils'; +import { useKibana } from '../../../../common/lib/kibana/kibana_react'; + +function generateId() { + return uuidv4(); +} + +const saveVisualizationLabel = i18n.translate('xpack.securityAiAssistant.lensESQLFunction.save', { + defaultMessage: 'Save visualization', +}); + +interface EsqlCodeBlockProps { + value: string; + timestamp: string; + actionsDisabled: boolean; +} + +const EsqlCodeBlockComponent: React.FC = ({ + value, + actionsDisabled = false, + timestamp, +}) => { + const { lens, dataViews: dataViewService, data } = useKibana().services; + const theme = useEuiTheme(); + + const lensHelpersAsync = useAsync(() => { + return lens.stateHelperApi(); + }, [lens]); + + const { data: queryResults } = useQuery({ + queryKey: ['test'], + enabled: true, + queryFn: async () => { + return getESQLResults({ + esqlQuery: value, + search: data.search.search, + filter: { + range: { + '@timestamp': { + lte: timestamp, + format: 'strict_date_optional_time', + }, + }, + }, + }); + }, + select: (dataz) => { + return { + params: dataz.params, + rows: dataz.response.values, + columns: dataz.response.columns, + }; + }, + refetchOnWindowFocus: false, + keepPreviousData: true, + }); + + const indexPattern = getIndexPatternFromESQLQuery(value); + const formattedColumns = useAsync( + () => + getESQLQueryColumns({ + esqlQuery: value, + search: data.search.search, + }), + [value] + ); + + const dataViewAsync = useAsync(() => { + return getESQLAdHocDataview(indexPattern, dataViewService).then((dataView) => { + if (dataView.fields.getByName('@timestamp')?.type === 'date') { + dataView.timeFieldName = '@timestamp'; + } + return dataView; + }); + }, [indexPattern, dataViewService]); + + const [isSaveModalOpen, setIsSaveModalOpen] = useState(false); + const [isTableVisible, setIsTableVisible] = useState(false); + const [lensInput, setLensInput] = useState(undefined); + const [showVisualization, setShowVisualization] = useState(false); + + const preferredChartType = undefined; // 'XY'; + + // initialization + useEffect(() => { + if (lensHelpersAsync.value && dataViewAsync.value && !lensInput && formattedColumns.value) { + const context = { + dataViewSpec: dataViewAsync.value?.toSpec(), + fieldName: '', + textBasedColumns: formattedColumns.value, + query: { + esql: value, + }, + }; + + const chartSuggestions = lensHelpersAsync.value.suggestions( + context, + dataViewAsync.value, + [], + preferredChartType + ); + + if (chartSuggestions?.length) { + const [suggestion] = chartSuggestions; + + const attrs = getLensAttributesFromSuggestion({ + filters: [], + query: { + esql: value, + }, + suggestion, + dataView: dataViewAsync.value, + }) as TypedLensByValueInput['attributes']; + + const lensEmbeddableInput = { + attributes: attrs, + id: generateId(), + }; + setLensInput(lensEmbeddableInput); + } + } + }, [ + dataViewAsync.value, + formattedColumns.value, + lensHelpersAsync.value, + lensInput, + preferredChartType, + value, + ]); + + // if the Lens suggestions api suggests a table then we want to render a Discover table instead + const isLensInputTable = lensInput?.attributes?.visualizationType === 'lnsDatatable'; + + return ( + <> + + + + + {value} + + + + {!showVisualization && ( + + + + setShowVisualization(true)} + disabled={actionsDisabled} + > + {i18n.translate('xpack.securityAiAssistant.visualizeThisQuery', { + defaultMessage: 'Generate Visualization', + })} + + + + + )} + + + + {showVisualization && queryResults && formattedColumns.value && ( + <> + {!isLensInputTable && ( + <> + + + + + setIsTableVisible(!isTableVisible)} + data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton" + aria-label={ + isTableVisible + ? i18n.translate( + 'xpack.observabilityAiAssistant.lensESQLFunction.displayChart', + { + defaultMessage: 'Display chart', + } + ) + : i18n.translate( + 'xpack.observabilityAiAssistant.lensESQLFunction.displayTable', + { + defaultMessage: 'Display table', + } + ) + } + /> + + + + + setIsSaveModalOpen(true)} + data-test-subj="observabilityAiAssistantLensESQLSaveButton" + aria-label={saveVisualizationLabel} + /> + + + + + + {isTableVisible ? ( + + ) : lensInput ? ( + + ) : null} + + + )} + {/* hide the grid in case of errors (as the user can't fix them) */} + {isLensInputTable && ( + + + + )} + + )} + + {isSaveModalOpen ? ( + { + setIsSaveModalOpen(() => false); + }} + // For now, we don't want to allow saving ESQL charts to the library + isSaveable={false} + /> + ) : null} + + ); +}; + +EsqlCodeBlockComponent.displayName = 'EsqlCodeBlock'; + +export const EsqlCodeBlock = React.memo(EsqlCodeBlockComponent); + +export const getEsqlRenderer = (timestamp) => (props) => { + return ( + <> + + + + ); +}; diff --git a/x-pack/plugins/cases/public/components/markdown_editor/renderer.tsx b/x-pack/plugins/cases/public/components/markdown_editor/renderer.tsx index f896da64cf021..b6ac0295c9cc0 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/renderer.tsx +++ b/x-pack/plugins/cases/public/components/markdown_editor/renderer.tsx @@ -16,6 +16,7 @@ interface Props { children: string; disableLinks?: boolean; textSize?: EuiMarkdownFormatProps['textSize']; + timestamp?: string; } const withDisabledLinks = (disableLinks?: boolean): React.FC => { @@ -28,8 +29,13 @@ const withDisabledLinks = (disableLinks?: boolean): React.FC return MarkdownLinkProcessingComponent; }; -const MarkdownRendererComponent: React.FC = ({ children, disableLinks, textSize }) => { - const { processingPlugins, parsingPlugins } = usePlugins(); +const MarkdownRendererComponent: React.FC = ({ + children, + disableLinks, + textSize, + timestamp, +}) => { + const { processingPlugins, parsingPlugins } = usePlugins({ timestamp }); // Deep clone of the processing plugins to prevent affecting the markdown editor. const processingPluginList = cloneDeep(processingPlugins); // This line of code is TS-compatible and it will break if [1][1] change in the future. diff --git a/x-pack/plugins/cases/public/components/markdown_editor/scrollable_markdown_renderer.tsx b/x-pack/plugins/cases/public/components/markdown_editor/scrollable_markdown_renderer.tsx index e4ca3d175cdf0..7db5a05e3acb1 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/scrollable_markdown_renderer.tsx +++ b/x-pack/plugins/cases/public/components/markdown_editor/scrollable_markdown_renderer.tsx @@ -20,7 +20,13 @@ export const getContentWrapperCss = (euiTheme: EuiThemeComputed<{}>) => css` -webkit-box-orient: vertical; `; -const ScrollableMarkdownRenderer = ({ content }: { content: string }) => { +const ScrollableMarkdownRenderer = ({ + content, + timestamp, +}: { + content: string; + timestamp?: string; +}) => { const { euiTheme } = useEuiTheme(); return (
{ css={getContentWrapperCss(euiTheme)} data-test-subj="scrollable-markdown" > - {content} + {content}
); }; diff --git a/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts b/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts index 2ed36feebcb4e..5b186d86302cc 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts +++ b/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts @@ -16,15 +16,26 @@ import type { TemporaryProcessingPluginsType } from './types'; import { KibanaServices, useApplicationCapabilities } from '../../common/lib/kibana'; import * as lensMarkdownPlugin from './plugins/lens'; import { ID as LensPluginId } from './plugins/lens/constants'; +import { getEsqlRenderer } from './plugins/custom_codeblock/esql_code_block'; +import { customCodeBlockLanguagePlugin } from './plugins/custom_codeblock/custom_codeblock_markdown_plugin'; -export const usePlugins = (disabledPlugins?: string[]) => { +export const usePlugins = ({ + disabledPlugins, + timestamp, +}: { + disabledPlugins?: string[]; + timestamp?: string; +} = {}) => { const kibanaConfig = KibanaServices.getConfig(); const timelinePlugins = useTimelineContext()?.editor_plugins; const appCapabilities = useApplicationCapabilities(); return useMemo(() => { const uiPlugins = getDefaultEuiMarkdownUiPlugins(); - const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); + const parsingPlugins = [ + customCodeBlockLanguagePlugin, + ...getDefaultEuiMarkdownParsingPlugins(), + ]; const processingPlugins = getDefaultEuiMarkdownProcessingPlugins() as TemporaryProcessingPluginsType; @@ -37,6 +48,8 @@ export const usePlugins = (disabledPlugins?: string[]) => { processingPlugins[1][1].components.timeline = timelinePlugins.processingPluginRenderer; } + processingPlugins[1][1].components.esql = getEsqlRenderer(timestamp); + if ( kibanaConfig?.markdownPlugins?.lens && !disabledPlugins?.includes(LensPluginId) && @@ -59,5 +72,6 @@ export const usePlugins = (disabledPlugins?: string[]) => { disabledPlugins, kibanaConfig?.markdownPlugins?.lens, timelinePlugins, + timestamp, ]); }; diff --git a/x-pack/plugins/cases/public/components/user_actions/comment/actions.tsx b/x-pack/plugins/cases/public/components/user_actions/comment/actions.tsx index 84f0ed768edf6..2b1cef6fc8ab8 100644 --- a/x-pack/plugins/cases/public/components/user_actions/comment/actions.tsx +++ b/x-pack/plugins/cases/public/components/user_actions/comment/actions.tsx @@ -54,7 +54,7 @@ export const createActionAttachmentUserActionBuilder = ({ timelineAvatarAriaLabel: actionIconName, actions: , children: comment.comment.trim().length > 0 && ( - + ), }, ]; diff --git a/x-pack/plugins/cases/public/components/user_actions/comment/user.tsx b/x-pack/plugins/cases/public/components/user_actions/comment/user.tsx index 67366ee4a81d3..6bea462851c00 100644 --- a/x-pack/plugins/cases/public/components/user_actions/comment/user.tsx +++ b/x-pack/plugins/cases/public/components/user_actions/comment/user.tsx @@ -99,6 +99,7 @@ export const createUserAttachmentUserActionBuilder = ({ ref={(element) => (commentRefs.current[comment.id] = element)} id={comment.id} content={comment.comment} + timestamp={comment.createdAt} isEditable={isEdit} caseId={caseId} onChangeEditable={handleManageMarkdownEditId} diff --git a/x-pack/plugins/cases/public/components/user_actions/markdown_form.tsx b/x-pack/plugins/cases/public/components/user_actions/markdown_form.tsx index 3866fe774ec14..75df5db9751f1 100644 --- a/x-pack/plugins/cases/public/components/user_actions/markdown_form.tsx +++ b/x-pack/plugins/cases/public/components/user_actions/markdown_form.tsx @@ -12,6 +12,7 @@ import { ScrollableMarkdown, EditableMarkdown } from '../markdown_editor'; interface UserActionMarkdownProps { content: string; + timestamp: string; id: string; caseId: string; isEditable: boolean; @@ -26,7 +27,7 @@ export interface UserActionMarkdownRefObject { const UserActionMarkdownComponent = forwardRef< UserActionMarkdownRefObject, UserActionMarkdownProps ->(({ id, content, caseId, isEditable, onChangeEditable, onSaveContent }, ref) => { +>(({ id, content, timestamp, caseId, isEditable, onChangeEditable, onSaveContent }, ref) => { const editorRef = useRef(); const fieldName = 'content'; @@ -43,7 +44,7 @@ const UserActionMarkdownComponent = forwardRef< formSchema={schema} /> ) : ( - + ); }); diff --git a/x-pack/plugins/cases/public/types.ts b/x-pack/plugins/cases/public/types.ts index c857446ea042c..888ee3df42104 100644 --- a/x-pack/plugins/cases/public/types.ts +++ b/x-pack/plugins/cases/public/types.ts @@ -30,6 +30,7 @@ import type { ContentManagementPublicStart } from '@kbn/content-management-plugi import type { UiActionsStart } from '@kbn/ui-actions-plugin/public'; import type { ServerlessPluginSetup, ServerlessPluginStart } from '@kbn/serverless/public'; +import type { DataViewsPublicPluginStart } from '@kbn/data-views-plugin/public'; import type { UseCasesAddToExistingCaseModal } from './components/all_cases/selector_modal/use_cases_add_to_existing_case_modal'; import type { UseCasesAddToNewCaseFlyout } from './components/create/flyout/use_cases_add_to_new_case_flyout'; import type { UseIsAddToCaseOpen } from './components/cases_context/state/use_is_add_to_case_open'; @@ -74,6 +75,7 @@ export interface CasesPublicSetupDependencies { export interface CasesPublicStartDependencies { apm?: ApmBase; data: DataPublicPluginStart; + dataViews: DataViewsPublicPluginStart; embeddable: EmbeddableStart; features: FeaturesPluginStart; files: FilesStart; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index c58d565e1907a..eb8fc4493fc97 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -15,7 +15,6 @@ import { useEuiTheme, } from '@elastic/eui'; import { css } from '@emotion/css'; -import type { UseChatSend } from '@kbn/elastic-assistant/impl/assistant/chat_send/use_chat_send'; import { i18n } from '@kbn/i18n'; import React, { useEffect, useState } from 'react'; import { useQuery } from '@tanstack/react-query'; @@ -36,24 +35,20 @@ function generateId() { return uuidv4(); } -const saveVisualizationLabel = i18n.translate( - 'xpack.observabilityAiAssistant.lensESQLFunction.save', - { - defaultMessage: 'Save visualization', - } -); +const saveVisualizationLabel = i18n.translate('xpack.securityAiAssistant.lensESQLFunction.save', { + defaultMessage: 'Save visualization', +}); export function EsqlCodeBlock({ value, actionsDisabled, - handleSendMessage, - ...rest + timestamp, }: { value: string; + timestamp: string; actionsDisabled: boolean; - handleSendMessage: UseChatSend['handleSendMessage']; }) { - const { lens, dataViews: dataViewService, uiActions, data } = useKibana().services; + const { lens, dataViews: dataViewService, data } = useKibana().services; const theme = useEuiTheme(); const lensHelpersAsync = useAsync(() => { @@ -66,8 +61,15 @@ export function EsqlCodeBlock({ queryFn: async () => { return getESQLResults({ esqlQuery: value, - // esqlQuery: `FROM logs-* | STATS avg_bytes = AVG(http.response.body.bytes) BY host.name`, search: data.search.search, + filter: { + range: { + '@timestamp': { + lte: timestamp, + format: 'strict_date_optional_time', + }, + }, + }, }); }, select: (dataz) => { @@ -81,8 +83,6 @@ export function EsqlCodeBlock({ keepPreviousData: true, }); - // const rows = getESQLResults({ esqlQuery: value, search: data.search.search }); - const indexPattern = getIndexPatternFromESQLQuery(value); const formattedColumns = useAsync( () => @@ -190,13 +190,13 @@ export function EsqlCodeBlock({ setShowVisualization(true)} disabled={actionsDisabled} > - {i18n.translate('xpack.observabilityAiAssistant.visualizeThisQuery', { + {i18n.translate('xpack.securityAiAssistant.visualizeThisQuery', { defaultMessage: 'Generate Visualization', })} @@ -279,14 +279,14 @@ export function EsqlCodeBlock({ flyoutType="overlay" isTableView /> - ) : ( + ) : lensInput ? ( - )} + ) : null} )} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx index 10d5a15c800ae..46b23998928c2 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx @@ -193,6 +193,7 @@ export const getComments = ({ actionTypeId={actionTypeId} abortStream={abortStream} content={transformedMessage.content} + timestamp={transformedMessage.timestamp} index={index} isControlsEnabled={isControlsEnabled} isEnabledLangChain={isEnabledLangChain} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx index 631d8b507ceed..367119c92486c 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx @@ -17,6 +17,7 @@ import { MessageText } from './message_text'; interface Props { abortStream: () => void; content?: string; + timestamp?: string; isEnabledLangChain: boolean; isError?: boolean; isFetching?: boolean; @@ -33,6 +34,7 @@ interface Props { export const StreamComment = ({ abortStream, content, + timestamp, actionTypeId, index, isControlsEnabled = false, @@ -108,7 +110,14 @@ export const StreamComment = ({ return ( } + body={ + + } error={error ? new Error(error) : undefined} controls={controls} /> diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx index 60c83982a893c..155536ee1bb24 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx @@ -29,6 +29,7 @@ import { EsqlCodeBlock } from '../custom_codeblock/esql_code_block'; interface Props { content: string; + timestamp?: string; index: number; loading: boolean; } @@ -99,7 +100,7 @@ const loadingCursorPlugin = () => { }; }; -const getPluginDependencies = () => { +const getPluginDependencies = (timestamp?: string) => { const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); const processingPlugins = getDefaultEuiMarkdownProcessingPlugins(); @@ -112,7 +113,7 @@ const getPluginDependencies = () => { esql: (props) => { return ( <> - + ); @@ -152,12 +153,12 @@ const getPluginDependencies = () => { }; }; -export function MessageText({ loading, content, index }: Props) { +export function MessageText({ loading, content, timestamp, index }: Props) { const containerClassName = css` overflow-wrap: anywhere; `; - const { parsingPluginList, processingPluginList } = getPluginDependencies(); + const { parsingPluginList, processingPluginList } = getPluginDependencies(timestamp); return ( diff --git a/x-pack/plugins/security_solution/public/assistant/helpers.tsx b/x-pack/plugins/security_solution/public/assistant/helpers.tsx index 61856ea1b17f3..32199234f5d89 100644 --- a/x-pack/plugins/security_solution/public/assistant/helpers.tsx +++ b/x-pack/plugins/security_solution/public/assistant/helpers.tsx @@ -65,14 +65,15 @@ export const augmentMessageCodeBlocks = ( currentConversation: Conversation, showAnonymizedValues: boolean ): CodeBlockDetails[][] => { - const cbd = currentConversation.messages.map(({ content }) => + const cbd = currentConversation.messages.map(({ content, timestamp }) => analyzeMarkdown( showAnonymizedValues ? content ?? '' : replaceAnonymizedValuesWithOriginalValues({ messageContent: content ?? '', replacements: currentConversation.replacements, - }) + }), + timestamp ) ); From 6468286f6caa39b946f35715f6b295c8f11324a0 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 20 Jun 2024 08:24:38 +0200 Subject: [PATCH 14/47] cleanup --- package.json | 8 +- .../conversation_sidepanel/index.tsx | 2 +- .../impl/assistant/index.tsx | 6 +- .../graphs/default_assistant_graph/_graph.ts | 245 ------------------ .../public/functions/index.ts | 4 +- .../public/index.ts | 2 - .../assistant/content/quick_prompts/index.tsx | 8 + yarn.lock | 33 ++- 8 files changed, 36 insertions(+), 272 deletions(-) delete mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/_graph.ts diff --git a/package.json b/package.json index 840fff2ea046b..efd020e91df66 100644 --- a/package.json +++ b/package.json @@ -80,7 +80,7 @@ "resolutions": { "**/@bazel/typescript/protobufjs": "6.11.4", "**/@hello-pangea/dnd": "16.6.0", - "**/@langchain/core": "0.2.7", + "**/@langchain/core": "0.2.3", "**/@types/node": "20.10.5", "**/@typescript-eslint/utils": "5.62.0", "**/chokidar": "^3.5.3", @@ -928,8 +928,8 @@ "@kbn/watcher-plugin": "link:x-pack/plugins/watcher", "@kbn/xstate-utils": "link:packages/kbn-xstate-utils", "@kbn/zod-helpers": "link:packages/kbn-zod-helpers", - "@langchain/community": "0.2.4", - "@langchain/core": "^0.2.7", + "@langchain/community": "^0.2.4", + "@langchain/core": "0.2.3", "@langchain/langgraph": "^0.0.23", "@langchain/openai": "^0.0.34", "@langtrase/trace-attributes": "^3.0.8", @@ -1069,7 +1069,7 @@ "jsts": "^1.6.2", "kea": "^2.6.0", "langchain": "0.2.3", - "langsmith": "^0.1.32", + "langsmith": "^0.1.30", "launchdarkly-js-client-sdk": "^3.3.0", "launchdarkly-node-server-sdk": "^7.0.3", "load-json-file": "^6.2.0", diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_sidepanel/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_sidepanel/index.tsx index d22a1e13d468a..2af44aa21acb6 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_sidepanel/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_sidepanel/index.tsx @@ -212,7 +212,7 @@ export const ConversationSidePanel = React.memo( onClick: () => setDeleteConversationItem(conversation), iconType: 'trash', iconSize: 's', - // disabled: conversation.isDefault, + disabled: conversation.isDefault, 'aria-label': i18n.DELETE_CONVERSATION_ARIA_LABEL, 'data-test-subj': 'delete-option', }} diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx index 2af6d8396c0ca..c220b464634c0 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -640,11 +640,12 @@ const AssistantComponent: React.FC = ({ ), [ - getComments, + refetchCurrentConversation, abortStream, currentConversation, + editingSystemPromptId, + getComments, showAnonymizedValues, - refetchCurrentConversation, handleRegenerateResponse, isEnabledKnowledgeBase, isEnabledRAGAlerts, @@ -652,7 +653,6 @@ const AssistantComponent: React.FC = ({ currentUserAvatar, isFlyoutMode, selectedPromptContextsCount, - editingSystemPromptId, isNewConversation, isSettingsModalVisible, promptContexts, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/_graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/_graph.ts deleted file mode 100644 index fa5b0302d4db0..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/_graph.ts +++ /dev/null @@ -1,245 +0,0 @@ -/* - * 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. - */ - -// /* -// * 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 { forOwn, get } from 'lodash'; -// import { ToolExecutor } from '@langchain/langgraph/prebuilt'; -// import { RunnableConfig, RunnableLambda } from '@langchain/core/runnables'; -// import { END, START, StateGraph, StateGraphArgs } from '@langchain/langgraph'; -// import { AgentAction, AgentFinish, AgentStep } from '@langchain/core/agents'; -// import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; -// import { StructuredTool } from '@langchain/core/tools'; -// import type { CompiledStateGraph } from '@langchain/langgraph/dist/graph/state'; -// import type { Logger } from '@kbn/logging'; - -// import { BaseMessage } from '@langchain/core/messages'; -// import { BaseChatModel } from '@langchain/core/language_models/chat_models'; -// import { ElasticsearchClient, IScopedClusterClient } from '@kbn/core-elasticsearch-server'; -// import { getESQLQueryColumns } from '@kbn/esql-utils'; -// import { AgentState, NodeParamsBase } from './types'; -// import { generateChatTitle } from './generate_chat_title'; -// import { AssistantDataClients } from '../../executors/types'; - -// export const DEFAULT_ASSISTANT_GRAPH_ID = 'Default Security Assistant Graph'; - -// interface GetDefaultAssistantGraphParams { -// agentRunnable: AgentRunnableSequence; -// dataClients?: AssistantDataClients; -// conversationId?: string; -// llm: BaseChatModel; -// logger: Logger; -// messages: BaseMessage[]; -// tools: StructuredTool[]; -// esClient: ElasticsearchClient; -// search: IScopedClusterClient; -// } - -// /** -// * Returns a compiled default assistant graph -// */ -// export const getDefaultAssistantGraph = ({ -// agentRunnable, -// conversationId, -// dataClients, -// llm, -// logger, -// messages, -// tools, -// esClient, -// search, -// }: GetDefaultAssistantGraphParams): CompiledStateGraph => { -// try { -// // Default graph state -// const graphState: StateGraphArgs['channels'] = { -// input: { -// value: (x: string, y?: string) => y ?? x, -// default: () => '', -// }, -// steps: { -// value: (x: AgentStep[], y: AgentStep[]) => x.concat(y), -// default: () => [], -// }, -// agentOutcome: { -// value: ( -// x: AgentAction | AgentFinish | undefined, -// y?: AgentAction | AgentFinish | undefined -// ) => y ?? x, -// default: () => undefined, -// }, -// messages: { -// value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y), -// default: () => messages, -// }, -// }; - -// const nodeParams: NodeParamsBase = { -// model: llm, -// logger, -// }; - -// // Create a tool executor -// const toolExecutor = new ToolExecutor({ tools }); - -// // Define logic that will be used to determine which conditional edge to go down -// const shouldContinue = (state: AgentState) => { -// logger.error(`graph:shouldContinue:state\n${JSON.stringify(state, null, 2)}`); -// if (state.agentOutcome && 'returnValues' in state.agentOutcome) { -// return 'end'; -// } -// return 'continue'; -// }; - -// const runAgent = async (state: AgentState, config?: RunnableConfig) => { -// logger.error(`graph:runAgent:\nstate\n${JSON.stringify(state, null, 2)}`); - -// const agentOutcome = await agentRunnable.invoke( -// { -// ...state, -// chat_history: messages, -// knowledge_history: 'The users favorite color is blue', // TODO: Plumb through initial retrieval -// }, -// config -// ); -// return { -// agentOutcome, -// }; -// }; - -// const executeTools = async (state: AgentState, config?: RunnableConfig) => { -// logger.error(`graph:executeTools:state\n${JSON.stringify(state, null, 2)}`); -// const agentAction = state.agentOutcome; -// if (!agentAction || 'returnValues' in agentAction) { -// throw new Error('Agent has not been run yet'); -// } -// const out = await toolExecutor.invoke(agentAction, config); -// return { -// steps: [{ action: agentAction, observation: JSON.stringify(out, null, 2) }], -// }; -// }; - -// const workflow = new StateGraph({ -// channels: { -// input: { -// value: (x: string, y?: string) => y ?? x, -// default: () => '', -// }, -// indices: { -// value: (x: string, y?: string) => y ?? x, -// default: () => [], -// }, -// }, -// }); - -// workflow.addNode('listIndicesEsDb', async (state: AgentState) => { -// console.error('listIndicesEsDb, state', state); -// const response = await esClient.indices.get({ index: '*' }); - -// console.error('response', JSON.stringify(response, null, 2)); - -// function transformInputToOutput(input: Record) { -// const output: Record = {}; - -// forOwn(input, (value, _key) => { -// const key = Object.keys(value.aliases)[0] || _key; -// const properties = get(value, 'mappings.properties', {}); -// output[key] = transformProperties(properties); -// }); - -// return output; -// } - -// function transformProperties(properties) { -// const result: Record = {}; - -// forOwn(properties, (value, key) => { -// if (value.type) { -// result[key] = value.type; -// } else if (value.properties) { -// result[key] = transformProperties(value.properties); -// } else if (value.fields) { -// result[key] = { fields: {} }; -// forOwn(value.fields, (fieldValue, fieldKey) => { -// result[key].fields[fieldKey] = { type: fieldValue.type }; -// if (fieldValue.ignore_above !== undefined) { -// result[key].fields[fieldKey].ignore_above = fieldValue.ignore_above; -// } -// }); -// } else { -// result[key] = value; -// } -// }); - -// return result; -// } - -// return { -// ...state, -// indices: transformInputToOutput(response), -// }; -// }); - -// workflow.addNode('generateESQLQuery', async (state: AgentState) => { -// const ecsDuplicatesPrompt = ECS_DUPLICATES_PROMPT; -// const outputParser = new JsonOutputParser(); -// const ecsDuplicatesGraph = ecsDuplicatesPrompt.pipe(model) - -// const currentMapping = await ecsDuplicatesGraph.invoke({ -// indices: state.indices, -// }); - -// return { -// ...state, -// }; -// }); - -// workflow.addNode('getColumns', async (state: AgentState) => { -// const columns = getESQLQueryColumns({ esqlQuery: state.query, search }); - -// console.error('columns', columns); - -// return { -// ...state, -// columns, -// }; -// }); -// // // Create a new graph, with the default state from above -// // const workflow = new StateGraph({ channels: graphState }); - -// // // Define the nodes to cycle between -// // workflow.addNode('generateChatTitle', (state: AgentState) => -// // generateChatTitle({ -// // state, -// // conversationsDataClient: dataClients?.conversationsDataClient, -// // conversationId, -// // ...nodeParams, -// // }) -// // ); -// // workflow.addNode('agent', new RunnableLambda({ func: runAgent })); -// // workflow.addNode('action', new RunnableLambda({ func: executeTools })); - -// // // Add conditional edge for determining if we shouldContinue -// // workflow.addConditionalEdges('agent', shouldContinue, { continue: 'action', end: END }); - -// // // Add edges for start, and between agent and action (action always followed by agent) -// workflow.addEdge(START, 'listIndicesEsDb'); -// workflow.addEdge('listIndicesEsDb', 'generateESQLQuery'); -// workflow.addEdge('generateESQLQuery', 'getColumns'); -// workflow.addEdge('getColumns', END); -// // workflow.addEdge('generateChatTitle', 'agent'); -// // workflow.addEdge('action', 'agent'); - -// return workflow.compile(); -// } catch (e) { -// throw new Error(`Unable to compile DefaultAssistantGraph\n${e}`); -// } -// }; diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/index.ts index ba3e16ca79712..0dd8fc2eed38e 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/index.ts @@ -9,9 +9,7 @@ import type { RegisterRenderFunctionDefinition } from '@kbn/observability-ai-ass import type { ObservabilityAIAssistantAppPluginStartDependencies } from '../types'; import { registerChangesRenderFunction } from './changes'; import { registerLensRenderFunction } from './lens'; -import { registerVisualizeQueryRenderFunction, VisualizeESQL } from './visualize_esql'; - -export { VisualizeESQL }; +import { registerVisualizeQueryRenderFunction } from './visualize_esql'; export async function registerFunctions({ registerRenderFunction, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/index.ts index dcb7116b6458b..5de1c30de7c4c 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/index.ts @@ -14,8 +14,6 @@ import type { ObservabilityAIAssistantAppPublicStart, } from './types'; -export { VisualizeESQL } from './functions'; - export const plugin: PluginInitializer< ObservabilityAIAssistantAppPublicSetup, ObservabilityAIAssistantAppPublicStart, diff --git a/x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx b/x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx index c35598f8898ba..799087f202e98 100644 --- a/x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx @@ -8,6 +8,7 @@ import type { QuickPrompt } from '@kbn/elastic-assistant'; import * as i18n from './translations'; import { + KNOWLEDGE_BASE_CATEGORY, PROMPT_CONTEXT_ALERT_CATEGORY, PROMPT_CONTEXT_DETECTION_RULES_CATEGORY, PROMPT_CONTEXT_EVENT_CATEGORY, @@ -26,6 +27,13 @@ export const BASE_SECURITY_QUICK_PROMPTS: QuickPrompt[] = [ categories: [PROMPT_CONTEXT_ALERT_CATEGORY], isDefault: true, }, + { + title: i18n.ESQL_QUERY_GENERATION_TITLE, + prompt: i18n.ESQL_QUERY_GENERATION_PROMPT, + color: '#9170B8', + categories: [KNOWLEDGE_BASE_CATEGORY], + isDefault: true, + }, { title: i18n.RULE_CREATION_TITLE, prompt: i18n.RULE_CREATION_PROMPT, diff --git a/yarn.lock b/yarn.lock index 65ddf06bbd2e8..1ed8696db1c03 100644 --- a/yarn.lock +++ b/yarn.lock @@ -6948,7 +6948,7 @@ resolved "https://registry.yarnpkg.com/@kwsites/promise-deferred/-/promise-deferred-1.1.1.tgz#8ace5259254426ccef57f3175bc64ed7095ed919" integrity sha512-GaHYm+c0O9MjZRu0ongGBRbinu8gVAMd2UZjji6jVmqKtZluZnptXGWhz1E8j8D2HJ3f/yMxKAUC0b+57wncIw== -"@langchain/community@0.2.4": +"@langchain/community@^0.2.4": version "0.2.4" resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.4.tgz#fb5feb4f4a01a1b33adfd28ce7126d0dedb3e6d1" integrity sha512-rwrPNQLyIe84TPqPYbYOfDA4G/ba1rdj7OtZg63dQmxIvNDOmUCh4xIQac2iuRUnM3o4Ben0Faa9qz+V5oPgIA== @@ -6965,16 +6965,16 @@ zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@0.2.7", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@^0.2.7", "@langchain/core@~0.2.0": - version "0.2.7" - resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.7.tgz#6a7c4a3a8f4cf884a4012c530e8d68c919186b77" - integrity sha512-FdFiNWhszFuUyAhYdY+l5DtPnAnWCAjXMnkLmUJ1J54NeUiUm7gy26Hnd4bkvaOQJ8ddHH/EX03ZwdoYfLv1jw== +"@langchain/core@0.2.3", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@~0.2.0": + version "0.2.3" + resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.3.tgz#7faa82f92b0c7843506e827a38bfcbb60f009d13" + integrity sha512-mVuFHSLpPQ4yOHNXeoSA3LnmIMuFmUiit5rvbYcPZqM6SrB2zCNN2nD4Ty5+3H5X4tYItDoSqsTuUNUQySXRQw== dependencies: ansi-styles "^5.0.0" camelcase "6" decamelize "1.2.0" js-tiktoken "^1.0.12" - langsmith "~0.1.30" + langsmith "~0.1.7" ml-distance "^4.0.0" mustache "^4.2.0" p-queue "^6.6.2" @@ -21720,10 +21720,10 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.32, langsmith@~0.1.1, langsmith@~0.1.30, langsmith@~0.1.7: - version "0.1.32" - resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.32.tgz#38938b0e8685522087b697b8200c488c6490c137" - integrity sha512-EUWHIH6fiOCGRYdzgwGoXwJxCMyUrL+bmUcxoVmkXoXoAGDOVinz8bqJLKbxotsQWqM64NKKsW85OTIutgNaMQ== +langsmith@^0.1.30, langsmith@~0.1.1, langsmith@~0.1.30, langsmith@~0.1.7: + version "0.1.30" + resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" + integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== dependencies: "@types/uuid" "^9.0.1" commander "^10.0.1" @@ -24337,9 +24337,9 @@ open@^8.0.9, open@^8.4.0, open@~8.4.0: is-wsl "^2.2.0" openai@^4.24.1, openai@^4.41.1: - version "4.51.0" - resolved "https://registry.yarnpkg.com/openai/-/openai-4.51.0.tgz#8ab08bba2441375e8e4ce6161f9ac987d2b2c157" - integrity sha512-UKuWc3/qQyklqhHM8CbdXCv0Z0obap6T0ECdcO5oATQxAbKE5Ky3YCXFQY207z+eGG6ez4U9wvAcuMygxhmStg== + version "4.47.1" + resolved "https://registry.yarnpkg.com/openai/-/openai-4.47.1.tgz#1d23c7a8eb3d7bcdc69709cd905f4c9af0181dba" + integrity sha512-WWSxhC/69ZhYWxH/OBsLEirIjUcfpQ5+ihkXKp06hmeYXgBBIUCa9IptMzYx6NdkiOCsSGYCnTIsxaic3AjRCQ== dependencies: "@types/node" "^18.11.18" "@types/node-fetch" "^2.6.4" @@ -28125,7 +28125,7 @@ semver@7.5.4: dependencies: lru-cache "^6.0.0" -semver@7.6.0, semver@^7.3.2, semver@^7.3.4, semver@^7.3.5, semver@^7.3.7, semver@^7.5.0, semver@^7.5.2, semver@^7.5.3, semver@^7.5.4: +semver@7.6.0: version "7.6.0" resolved "https://registry.yarnpkg.com/semver/-/semver-7.6.0.tgz#1a46a4db4bffcccd97b743b5005c8325f23d4e2d" integrity sha512-EnwXhrlwXMk9gKu5/flx5sv/an57AkRplG3hTK68W7FRDN+k+OWBj65M7719OkA82XLBxrcX0KSHj+X5COhOVg== @@ -28137,6 +28137,11 @@ semver@^6.0.0, semver@^6.1.0, semver@^6.1.2, semver@^6.3.0, semver@^6.3.1: resolved "https://registry.yarnpkg.com/semver/-/semver-6.3.1.tgz#556d2ef8689146e46dcea4bfdd095f3434dffcb4" integrity sha512-BR7VvDCVHO+q2xBEWskxS6DJE1qRnb7DxzUrogb71CWoSficBxYsiAGd+Kl0mmq/MprG9yArRkyrQxTO6XjMzA== +semver@^7.3.2, semver@^7.3.4, semver@^7.3.5, semver@^7.3.7, semver@^7.5.0, semver@^7.5.2, semver@^7.5.3, semver@^7.5.4: + version "7.6.2" + resolved "https://registry.yarnpkg.com/semver/-/semver-7.6.2.tgz#1e3b34759f896e8f14d6134732ce798aeb0c6e13" + integrity sha512-FNAIBWCx9qcRhoHcgcJ0gvU7SN1lYU2ZXuSfl04bSC5OpvDHFyJCjdNHomPXxjQlCBU67YW64PzY7/VIEH7F2w== + send@0.18.0: version "0.18.0" resolved "https://registry.yarnpkg.com/send/-/send-0.18.0.tgz#670167cc654b05f5aa4a767f9113bb371bc706be" From 5421b42063c449238713cb9ba3f62b8a90806fc7 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Sun, 23 Jun 2024 21:31:12 +0200 Subject: [PATCH 15/47] WIP --- package.json | 16 +- .../impl/capabilities/index.ts | 2 - .../custom_codeblock/esql_code_block.tsx | 224 +++---- .../get_comments/stream/message_text.tsx | 22 +- .../correct_common_esql_mistakes.test.ts | 182 ++++++ .../correct_common_esql_mistakes.ts | 287 +++++++++ .../esql_language_knowledge_base_tool.ts | 561 +++++++++++------- yarn.lock | 66 +-- 8 files changed, 988 insertions(+), 372 deletions(-) create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.test.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts diff --git a/package.json b/package.json index 55eb4bfd3daec..e966403807dfa 100644 --- a/package.json +++ b/package.json @@ -80,7 +80,8 @@ "resolutions": { "**/@bazel/typescript/protobufjs": "6.11.4", "**/@hello-pangea/dnd": "16.6.0", - "**/@langchain/core": "0.2.3", + "**/@langchain/core": "0.2.9", + "**/@langchain/openai": "0.2.0", "**/@types/node": "20.10.5", "**/@typescript-eslint/utils": "5.62.0", "**/chokidar": "^3.5.3", @@ -88,6 +89,7 @@ "**/globule/minimatch": "^3.1.2", "**/hoist-non-react-statics": "^3.3.2", "**/isomorphic-fetch/node-fetch": "^2.6.7", + "**/langchain": "0.2.6", "**/react-intl/**/@types/react": "^17.0.45", "**/remark-parse/trim": "1.0.1", "**/sharp": "0.32.6", @@ -930,10 +932,10 @@ "@kbn/watcher-plugin": "link:x-pack/plugins/watcher", "@kbn/xstate-utils": "link:packages/kbn-xstate-utils", "@kbn/zod-helpers": "link:packages/kbn-zod-helpers", - "@langchain/community": "^0.2.4", - "@langchain/core": "0.2.3", - "@langchain/langgraph": "^0.0.23", - "@langchain/openai": "^0.0.34", + "@langchain/community": "^0.2.13", + "@langchain/core": "^0.2.9", + "@langchain/langgraph": "^0.0.24", + "@langchain/openai": "^0.2.0", "@langtrase/trace-attributes": "^3.0.8", "@launchdarkly/node-server-sdk": "^9.4.5", "@loaders.gl/core": "^3.4.7", @@ -1070,8 +1072,8 @@ "jsonwebtoken": "^9.0.2", "jsts": "^1.6.2", "kea": "^2.6.0", - "langchain": "0.2.3", - "langsmith": "^0.1.30", + "langchain": "^0.2.6", + "langsmith": "^0.1.32", "launchdarkly-js-client-sdk": "^3.3.0", "launchdarkly-node-server-sdk": "^7.0.3", "load-json-file": "^6.2.0", diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index 3011ae16f879b..5f71f9d4fa3f3 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -22,5 +22,3 @@ export const defaultAssistantFeatures = Object.freeze({ assistantKnowledgeBaseByDefault: true, assistantModelEvaluation: true, }); - -// console.log('a') diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index eb8fc4493fc97..ecf5f39ed6307 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -16,7 +16,7 @@ import { } from '@elastic/eui'; import { css } from '@emotion/css'; import { i18n } from '@kbn/i18n'; -import React, { useEffect, useState } from 'react'; +import React, { useCallback, useEffect, useMemo, useState } from 'react'; import { useQuery } from '@tanstack/react-query'; import { v4 as uuidv4 } from 'uuid'; import { @@ -39,7 +39,7 @@ const saveVisualizationLabel = i18n.translate('xpack.securityAiAssistant.lensESQ defaultMessage: 'Save visualization', }); -export function EsqlCodeBlock({ +const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp, @@ -47,16 +47,16 @@ export function EsqlCodeBlock({ value: string; timestamp: string; actionsDisabled: boolean; -}) { - const { lens, dataViews: dataViewService, data } = useKibana().services; +}) => { const theme = useEuiTheme(); + const { lens, dataViews: dataViewService, data } = useKibana().services; - const lensHelpersAsync = useAsync(() => { + const { value: lensHelpersAsync } = useAsync(() => { return lens.stateHelperApi(); }, [lens]); const { data: queryResults } = useQuery({ - queryKey: ['test'], + queryKey: ['getESQLResults', value, timestamp], enabled: true, queryFn: async () => { return getESQLResults({ @@ -83,8 +83,8 @@ export function EsqlCodeBlock({ keepPreviousData: true, }); - const indexPattern = getIndexPatternFromESQLQuery(value); - const formattedColumns = useAsync( + const indexPattern = useMemo(() => getIndexPatternFromESQLQuery(value), [value]); + const { value: formattedColumns } = useAsync( () => getESQLQueryColumns({ esqlQuery: value, @@ -93,7 +93,7 @@ export function EsqlCodeBlock({ [value] ); - const dataViewAsync = useAsync(() => { + const { value: dataViewAsync } = useAsync(() => { return getESQLAdHocDataview(indexPattern, dataViewService).then((dataView) => { if (dataView.fields.getByName('@timestamp')?.type === 'date') { dataView.timeFieldName = '@timestamp'; @@ -111,19 +111,19 @@ export function EsqlCodeBlock({ // initialization useEffect(() => { - if (lensHelpersAsync.value && dataViewAsync.value && !lensInput && formattedColumns.value) { + if (lensHelpersAsync && dataViewAsync && !lensInput && formattedColumns) { const context = { - dataViewSpec: dataViewAsync.value?.toSpec(), + dataViewSpec: dataViewAsync?.toSpec(), fieldName: '', - textBasedColumns: formattedColumns.value, + textBasedColumns: formattedColumns, query: { esql: value, }, }; - const chartSuggestions = lensHelpersAsync.value.suggestions( + const chartSuggestions = lensHelpersAsync.suggestions( context, - dataViewAsync.value, + dataViewAsync, [], preferredChartType ); @@ -137,7 +137,7 @@ export function EsqlCodeBlock({ esql: value, }, suggestion, - dataView: dataViewAsync.value, + dataView: dataViewAsync, }) as TypedLensByValueInput['attributes']; const lensEmbeddableInput = { @@ -147,18 +147,13 @@ export function EsqlCodeBlock({ setLensInput(lensEmbeddableInput); } } - }, [ - dataViewAsync.value, - formattedColumns.value, - lensHelpersAsync.value, - lensInput, - preferredChartType, - value, - ]); + }, [dataViewAsync, formattedColumns, lensHelpersAsync, lensInput, preferredChartType, value]); // if the Lens suggestions api suggests a table then we want to render a Discover table instead const isLensInputTable = lensInput?.attributes?.visualizationType === 'lnsDatatable'; + const handleShowVisualization = useCallback(() => setShowVisualization(true), []); + return ( <> setShowVisualization(true)} + onClick={handleShowVisualization} disabled={actionsDisabled} > {i18n.translate('xpack.securityAiAssistant.visualizeThisQuery', { @@ -206,105 +201,108 @@ export function EsqlCodeBlock({ )} - - {showVisualization && queryResults && formattedColumns.value && ( - <> - {!isLensInputTable && ( - <> - - - - - setIsTableVisible(!isTableVisible)} - data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton" - aria-label={ + {showVisualization && queryResults && formattedColumns && ( + + + {!isLensInputTable && ( + <> + + + + - - - - - setIsSaveModalOpen(true)} - data-test-subj="observabilityAiAssistantLensESQLSaveButton" - aria-label={saveVisualizationLabel} - /> - - - - + > + setIsTableVisible(!isTableVisible)} + data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton" + aria-label={ + isTableVisible + ? i18n.translate( + 'xpack.observabilityAiAssistant.lensESQLFunction.displayChart', + { + defaultMessage: 'Display chart', + } + ) + : i18n.translate( + 'xpack.observabilityAiAssistant.lensESQLFunction.displayTable', + { + defaultMessage: 'Display table', + } + ) + } + /> + + + + + setIsSaveModalOpen(true)} + data-test-subj="observabilityAiAssistantLensESQLSaveButton" + aria-label={saveVisualizationLabel} + /> + + + + + + {isTableVisible ? ( + + ) : lensInput ? ( + + ) : null} + + + )} + {/* hide the grid in case of errors (as the user can't fix them) */} + {isLensInputTable && ( - {isTableVisible ? ( - - ) : lensInput ? ( - - ) : null} + - - )} - {/* hide the grid in case of errors (as the user can't fix them) */} - {isLensInputTable && ( - - - - )} - + )} + + )} {isSaveModalOpen ? ( @@ -319,4 +317,6 @@ export function EsqlCodeBlock({ ) : null} ); -} +}; + +export const EsqlCodeBlock = React.memo(EsqlCodeBlockComponent); diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx index 155536ee1bb24..3c88967a6db94 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx @@ -100,6 +100,14 @@ const loadingCursorPlugin = () => { }; }; +const getEsql = (timestamp?: string) => (props) => + ( + <> + + + + ); + const getPluginDependencies = (timestamp?: string) => { const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); @@ -110,14 +118,7 @@ const getPluginDependencies = (timestamp?: string) => { processingPlugins[1][1].components = { ...components, cursor: Cursor, - esql: (props) => { - return ( - <> - - - - ); - }, + esql: getEsql(timestamp), customCodeBlock: (props) => { return ( <> @@ -153,7 +154,7 @@ const getPluginDependencies = (timestamp?: string) => { }; }; -export function MessageText({ loading, content, timestamp, index }: Props) { +export const MessageText = React.memo(({ loading, content, timestamp, index }: Props) => { const containerClassName = css` overflow-wrap: anywhere; `; @@ -174,4 +175,5 @@ export function MessageText({ loading, content, timestamp, index }: Props) { ); -} +}); +MessageText.displayName = 'MessageText'; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.test.ts new file mode 100644 index 0000000000000..ad8e0f6cfd664 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.test.ts @@ -0,0 +1,182 @@ +/* + * 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 { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; + +describe('correctCommonEsqlMistakes', () => { + function normalize(input: string) { + return input.replaceAll(/[\t|\s]*\n[\t|\s]*/gms, '\n'); + } + + function expectQuery(input: string, expectedOutput: string) { + expect(normalize(correctCommonEsqlMistakes(input).output)).toEqual(normalize(expectedOutput)); + } + + it('replaces aliasing via the AS keyword with the = operator', () => { + expectQuery(`FROM logs-* | STATS COUNT() AS count`, 'FROM logs-*\n| STATS count = COUNT()'); + + expectQuery(`FROM logs-* | STATS COUNT() as count`, 'FROM logs-*\n| STATS count = COUNT()'); + + expectQuery( + `FROM logs-* | STATS AVG(transaction.duration.histogram) AS avg_request_latency, PERCENTILE(transaction.duration.histogram, 95) AS p95`, + `FROM logs-* + | STATS avg_request_latency = AVG(transaction.duration.histogram), p95 = PERCENTILE(transaction.duration.histogram, 95)` + ); + + expectQuery( + `FROM traces-apm* + | WHERE @timestamp >= NOW() - 24 hours + | STATS AVG(transaction.duration.us) AS avg_duration, SUM(success) AS total_successes, COUNT(*) AS total_requests BY service.name`, + `FROM traces-apm* + | WHERE @timestamp >= NOW() - 24 hours + | STATS avg_duration = AVG(transaction.duration.us), total_successes = SUM(success), total_requests = COUNT(*) BY service.name` + ); + }); + + it(`replaces " or ' escaping in FROM statements with backticks`, () => { + expectQuery(`FROM "logs-*" | LIMIT 10`, 'FROM logs-*\n| LIMIT 10'); + expectQuery(`FROM 'logs-*' | LIMIT 10`, 'FROM logs-*\n| LIMIT 10'); + expectQuery(`FROM logs-* | LIMIT 10`, 'FROM logs-*\n| LIMIT 10'); + }); + + it('replaces = as equal operator with ==', () => { + expectQuery( + `FROM logs-*\n| WHERE service.name = "foo"`, + `FROM logs-*\n| WHERE service.name == "foo"` + ); + + expectQuery( + `FROM logs-*\n| WHERE service.name = "foo" AND service.environment = "bar"`, + `FROM logs-*\n| WHERE service.name == "foo" AND service.environment == "bar"` + ); + + expectQuery( + `FROM logs-*\n| WHERE (service.name = "foo" AND service.environment = "bar") OR agent.name = "baz"`, + `FROM logs-*\n| WHERE (service.name == "foo" AND service.environment == "bar") OR agent.name == "baz"` + ); + + expectQuery( + `FROM logs-*\n| WHERE \`what=ever\` = "foo=bar"`, + `FROM logs-*\n| WHERE \`what=ever\` == "foo=bar"` + ); + }); + + it('replaces single-quote escaped strings with double-quote escaped strings', () => { + expectQuery( + `FROM nyc_taxis + | WHERE DATE_EXTRACT('hour', dropoff_datetime) >= 6 AND DATE_EXTRACT('hour', dropoff_datetime) < 10 + | LIMIT 10`, + `FROM nyc_taxis + | WHERE DATE_EXTRACT("hour", dropoff_datetime) >= 6 AND DATE_EXTRACT("hour", dropoff_datetime) < 10 + | LIMIT 10` + ); + expectQuery( + `FROM nyc_taxis + | WHERE DATE_EXTRACT('hour', "hh:mm a, 'of' d MMMM yyyy") >= 6 AND DATE_EXTRACT('hour', dropoff_datetime) < 10 + | LIMIT 10`, + `FROM nyc_taxis + | WHERE DATE_EXTRACT("hour", "hh:mm a, 'of' d MMMM yyyy") >= 6 AND DATE_EXTRACT("hour", dropoff_datetime) < 10 + | LIMIT 10` + ); + }); + + it(`verifies if the SORT key is in KEEP, and if it's not, it will include it`, () => { + expectQuery( + 'FROM logs-* \n| KEEP date \n| SORT @timestamp DESC', + 'FROM logs-*\n| KEEP date, @timestamp\n| SORT @timestamp DESC' + ); + + expectQuery( + `FROM logs-* | KEEP date, whatever | EVAL my_truncated_date_field = DATE_TRUNC(1 year, date) | SORT @timestamp, my_truncated_date_field DESC`, + 'FROM logs-*\n| KEEP date, whatever, @timestamp\n| EVAL my_truncated_date_field = DATE_TRUNC(1 year, date)\n| SORT @timestamp, my_truncated_date_field DESC' + ); + + expectQuery( + `FROM logs-* | KEEP date, whatever | RENAME whatever AS forever | SORT forever DESC`, + `FROM logs-*\n| KEEP date, whatever\n| RENAME whatever AS forever\n| SORT forever DESC` + ); + + expectQuery( + 'FROM employees\n| KEEP first_name, last_name\n| RENAME first_name AS fn, last_name AS ln', + 'FROM employees\n| KEEP first_name, last_name\n| RENAME first_name AS fn, last_name AS ln' + ); + }); + + it(`escapes the column name if SORT uses an expression`, () => { + expectQuery( + 'FROM logs-* \n| STATS COUNT(*) by service.name\n| SORT COUNT(*) DESC', + 'FROM logs-*\n| STATS COUNT(*) BY service.name\n| SORT `COUNT(*)` DESC' + ); + + expectQuery( + 'FROM logs-* \n| STATS COUNT(*) by service.name\n| SORT COUNT(*) DESC, @timestamp ASC', + 'FROM logs-*\n| STATS COUNT(*) BY service.name\n| SORT `COUNT(*)` DESC, @timestamp ASC' + ); + + expectQuery( + `FROM employees\n| KEEP first_name, last_name, height\n| SORT first_name ASC NULLS FIRST`, + `FROM employees\n| KEEP first_name, last_name, height\n| SORT first_name ASC NULLS FIRST` + ); + + expectQuery( + `FROM employees + | STATS my_count = COUNT() BY LEFT(last_name, 1) + | SORT \`LEFT(last_name, 1)\``, + `FROM employees + | STATS my_count = COUNT() BY LEFT(last_name, 1) + | SORT \`LEFT(last_name, 1)\`` + ); + }); + + it(`handles complicated queries correctly`, () => { + expectQuery( + `FROM "postgres-logs*" + | GROK message "%{TIMESTAMP_ISO8601:timestamp} %{TZ} \[%{NUMBER:process_id}\]: \[%{NUMBER:log_line}\] user=%{USER:user},db=%{USER:database},app=\[%{DATA:application}\],client=%{IP:client_ip} LOG: duration: %{NUMBER:duration:float} ms statement: %{GREEDYDATA:statement}" + | EVAL "@timestamp" = TO_DATETIME(timestamp) + | WHERE statement LIKE 'SELECT%' + | STATS avg_duration = AVG(duration)`, + `FROM postgres-logs* + | GROK message "%{TIMESTAMP_ISO8601:timestamp} %{TZ} \[%{NUMBER:process_id}\]: \[%{NUMBER:log_line}\] user=%{USER:user},db=%{USER:database},app=\[%{DATA:application}\],client=%{IP:client_ip} LOG: duration: %{NUMBER:duration:float} ms statement: %{GREEDYDATA:statement}" + | EVAL @timestamp = TO_DATETIME(timestamp) + | WHERE statement LIKE "SELECT%" + | STATS avg_duration = AVG(duration)` + ); + + expectQuery( + `FROM metrics-apm* + | WHERE metricset.name == "service_destination" AND @timestamp > NOW() - 24 hours + | EVAL total_events = span.destination.service.response_time.count + | EVAL total_latency = span.destination.service.response_time.sum.us + | EVAL is_failure = CASE(event.outcome == "failure", 1, 0) + | STATS + avg_throughput = AVG(total_events), + avg_latency_per_request = AVG(total_latency / total_events), + failure_rate = AVG(is_failure) + BY span.destination.service.resource`, + `FROM metrics-apm* + | WHERE metricset.name == "service_destination" AND @timestamp > NOW() - 24 hours + | EVAL total_events = span.destination.service.response_time.count + | EVAL total_latency = span.destination.service.response_time.sum.us + | EVAL is_failure = CASE(event.outcome == "failure", 1, 0) + | STATS avg_throughput = AVG(total_events), avg_latency_per_request = AVG(total_latency / total_events), failure_rate = AVG(is_failure) BY span.destination.service.resource` + ); + + expectQuery( + `FROM sample_data + | EVAL successful = CASE( + STARTS_WITH(message, "Connected to"), 1, + message == "Connection error", 0 + ) + | STATS success_rate = AVG(successful)`, + `FROM sample_data + | EVAL successful = CASE( + STARTS_WITH(message, "Connected to"), 1, + message == "Connection error", 0 + ) + | STATS success_rate = AVG(successful)` + ); + }); +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts new file mode 100644 index 0000000000000..09b6a59cac357 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts @@ -0,0 +1,287 @@ +/* + * 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. + */ + +const STRING_DELIMITER_TOKENS = ['`', "'", '"']; +const ESCAPE_TOKEN = '\\\\'; + +// this function splits statements by a certain token, +// and takes into account escaping, or function calls + +function split(value: string, splitToken: string) { + const statements: string[] = []; + + let delimiterToken: string | undefined; + + let groupingCount: number = 0; + + let currentStatement: string = ''; + + const trimmed = value.trim().split(''); + + for (let index = 0; index < trimmed.length; index++) { + const char = trimmed[index]; + + if ( + !delimiterToken && + groupingCount === 0 && + trimmed + .slice(index, index + splitToken.length) + .join('') + .toLowerCase() === splitToken.toLowerCase() + ) { + index += splitToken.length - 1; + statements.push(currentStatement.trim()); + currentStatement = ''; + continue; + } + + currentStatement += char; + + if (delimiterToken === char) { + // end identifier + delimiterToken = undefined; + } else if (!delimiterToken && trimmed[index - 1] !== ESCAPE_TOKEN) { + const applicableToken = STRING_DELIMITER_TOKENS.includes(char) ? char : undefined; + + if (applicableToken) { + // start identifier + delimiterToken = applicableToken; + continue; + } else if (char === '(') { + groupingCount++; + } else if (char === ')') { + groupingCount--; + } + } + } + + if (currentStatement) { + statements.push(currentStatement.trim()); + } + + return statements; +} + +export function splitIntoCommands(query: string) { + const commands: string[] = split(query, '|'); + + return commands.map((command) => { + const commandName = command.match(/^([A-Za-z]+)/)?.[1]; + + return { + name: commandName, + command, + }; + }); +} + +function replaceSingleQuotesWithDoubleQuotes(command: string) { + const regex = /'(?=(?:[^"]*"[^"]*")*[^"]*$)/g; + return command.replace(regex, '"'); +} + +function removeColumnQuotesAndEscape(column: string) { + const plainColumnIdentifier = column.replaceAll(/^"(.*)"$/g, `$1`).replaceAll(/^'(.*)'$/g, `$1`); + + if (isValidColumnName(plainColumnIdentifier)) { + return plainColumnIdentifier; + } + + return `\`${plainColumnIdentifier}\``; +} + +function replaceAsKeywordWithAssignments(command: string) { + return command.replaceAll(/^STATS\s*(.*)/g, (__, statsOperations: string) => { + return `STATS ${statsOperations.replaceAll( + /(,\s*)?(.*?)\s(AS|as)\s([`a-zA-Z0-9.\-_]+)/g, + '$1$4 = $2' + )}`; + }); +} + +function isValidColumnName(column: string) { + return Boolean(column.match(/^`.*`$/) || column.match(/^[@A-Za-z\._\-\d]+$/)); +} + +function escapeColumns(line: string) { + const [, command, body] = line.match(/^([A-Za-z_]+)(.*)$/ms) ?? ['', '', '']; + + const escapedBody = split(body.trim(), ',') + .map((statement) => { + const [lhs, rhs] = split(statement, '='); + if (!rhs) { + return lhs; + } + return `${removeColumnQuotesAndEscape(lhs)} = ${rhs}`; + }) + .join(', '); + + return `${command} ${escapedBody}`; +} + +function verifyKeepColumns( + keepCommand: string, + nextCommands: Array<{ name?: string; command: string }> +) { + const columnsInKeep = split(keepCommand.replace(/^KEEP\s*/, ''), ',').map((statement) => + split(statement, '=')?.[0].trim() + ); + + const availableColumns = columnsInKeep.concat(); + + for (const { name, command } of nextCommands) { + if (['STATS', 'KEEP', 'DROP', 'DISSECT', 'GROK', 'ENRICH', 'RENAME'].includes(name || '')) { + // these operations alter columns in a way that is hard to analyze, so we abort + break; + } + + const commandBody = command.replace(/^[A-Za-z]+\s*/, ''); + + if (name === 'EVAL') { + // EVAL creates new columns, make them available + const columnsInEval = split(commandBody, ',').map((column) => + split(column.trim(), '=')[0].trim() + ); + + columnsInEval.forEach((column) => { + availableColumns.push(column); + }); + } + + if (name === 'RENAME') { + // RENAME creates and removes columns + split(commandBody, ',').forEach((statement) => { + const [prevName, newName] = split(statement, 'AS').map((side) => side.trim()); + availableColumns.push(newName); + if (!availableColumns.includes(prevName)) { + columnsInKeep.push(prevName); + } + }); + } + + if (name === 'SORT') { + const columnsInSort = split(commandBody, ',').map((column) => + split(column.trim(), ' ')[0].trim() + ); + + columnsInSort.forEach((column) => { + if (isValidColumnName(column) && !availableColumns.includes(column)) { + columnsInKeep.push(column); + } + }); + } + } + + return `KEEP ${columnsInKeep.join(', ')}`; +} + +function escapeExpressionsInSort(sortCommand: string) { + const columnsInSort = split(sortCommand.replace(/^SORT\s*/, ''), ',') + .map((statement) => split(statement, '=')?.[0].trim()) + .map((columnAndSortOrder) => { + let [, column, sortOrder = ''] = + columnAndSortOrder.match(/^(.*?)\s+(ASC|DESC\s*([A-Z\s]+)?)?$/i) || []; + if (!column) { + return columnAndSortOrder; + } + + if (sortOrder) sortOrder = ` ${sortOrder}`; + + if (!column.match(/^`?[a-zA-Z0-9_\.@]+`?$/)) { + column = `\`${column}\``; + } + + return `${column}${sortOrder}`; + }); + + return `SORT ${columnsInSort.join(', ')}`; +} + +function ensureEqualityOperators(whereCommand: string) { + const body = whereCommand.split(/^WHERE /)[1]; + + const byChar = body.split(''); + + let next = ''; + let isColumnName = false; + byChar.forEach((char, index) => { + next += char; + + if (!isColumnName && char === '=' && byChar[index - 1] === ' ' && byChar[index + 1] === ' ') { + next += '='; + } + + if (!isColumnName && (char === '`' || char.match(/[a-z@]/i))) { + isColumnName = true; + } else if (isColumnName && (char === '`' || !char.match(/[a-z@0-9]/i))) { + isColumnName = false; + } + }); + + return `WHERE ${next}`; +} + +export function correctCommonEsqlMistakes(query: string): { + isCorrection: boolean; + input: string; + output: string; +} { + const commands = splitIntoCommands(query.trim()); + + const formattedCommands: string[] = commands.map(({ name, command }, index) => { + let formattedCommand = command; + + switch (name) { + case 'FROM': + formattedCommand = formattedCommand + .replaceAll(/FROM "(.*)"/g, 'FROM $1') + .replaceAll(/FROM '(.*)'/g, 'FROM $1') + .replaceAll(/FROM `(.*)`/g, 'FROM $1'); + break; + + case 'WHERE': + formattedCommand = replaceSingleQuotesWithDoubleQuotes(formattedCommand); + formattedCommand = ensureEqualityOperators(formattedCommand); + break; + + case 'EVAL': + formattedCommand = replaceSingleQuotesWithDoubleQuotes(formattedCommand); + formattedCommand = escapeColumns(formattedCommand); + break; + + case 'STATS': + formattedCommand = replaceAsKeywordWithAssignments(formattedCommand); + const [before, after] = split(formattedCommand, ' BY '); + formattedCommand = escapeColumns(before); + if (after) { + formattedCommand += ` BY ${after}`; + } + break; + + case 'KEEP': + formattedCommand = verifyKeepColumns(formattedCommand, commands.slice(index + 1)); + break; + + case 'SORT': + formattedCommand = escapeExpressionsInSort(formattedCommand); + break; + } + return formattedCommand; + }); + + const output = formattedCommands.join('\n| '); + + const originalFormattedQuery = commands.map((cmd) => cmd.command).join('\n| '); + + const isCorrection = output !== originalFormattedQuery; + + return { + input: query, + output, + isCorrection, + }; +} 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 d6a583b680f14..4078e958a539f 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 @@ -5,16 +5,28 @@ * 2.0. */ +import Fs from 'fs'; +import { keyBy, mapValues, once, pick, forOwn, get, isEmpty, map } from 'lodash'; import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from 'zod'; +import pLimit from 'p-limit'; +import Path from 'path'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; -import { forOwn, get } from 'lodash'; import { JsonOutputParser, StringOutputParser } from '@langchain/core/output_parsers'; +import type { StateGraphArgs } from '@langchain/langgraph'; +import { END, START, StateGraph } from '@langchain/langgraph'; +import type { BaseMessage } from '@langchain/core/messages'; +import type { RunnableConfig } from '@langchain/core/runnables'; +import type { ValidationResult } from '@kbn/esql-validation-autocomplete/src/validation/types'; import { getESQLQueryColumns } from '@kbn/esql-utils'; +import { promisify } from 'util'; import { ChatPromptTemplate } from '@langchain/core/prompts'; -import type { StateGraphArgs } from '@langchain/langgraph'; -import { StateGraph } from '@langchain/langgraph'; +import { validateQuery } from '@kbn/esql-validation-autocomplete'; +import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; import { APP_UI_ID } from '../../../../common'; +import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; + +export const INLINE_ESQL_QUERY_REGEX = /```esql\s*(.*?)\s*```/gms; export const ECS_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ [ @@ -33,12 +45,11 @@ Here is some context for you to reference for your task, read it carefully as yo [ 'human', `{input}. - Example response format: A: Please find the ESQL query below: \`\`\`esql -FROM logs +FROM logs-* | SORT @timestamp DESC | LIMIT 5 \`\`\` @@ -49,150 +60,349 @@ FROM logs export type EsqlKnowledgeBaseToolParams = AssistantToolParams; +const TOOL_NAME = 'ESQLKnowledgeBaseTool'; + 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. Only output valid ES|QL queries as described above. Do not add any additional text to describe your output.', id: 'esql-knowledge-base-tool', - name: 'ESQLKnowledgeBaseTool', + name: TOOL_NAME, + description: `You MUST use the "${TOOL_NAME}" function when the user wants to: + - visualize data + - 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. + + If the user asks for a query, and one of the dataset info functions was called and returned no results, you should still call the query function to generate an example query. + + 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.`, }; -export interface CategorizationState { - rawSamples: string[]; - samples: string[]; - formattedSamples: string; - ecsTypes: string; - ecsCategories: string; - exAnswer: string; - lastExecutedChain: string; - packageName: string; - dataStreamName: string; - errors: object; - pipelineResults: object[]; - finalized: boolean; - reviewed: boolean; - currentPipeline: object; - currentProcessors: object[]; - invalidCategorization: object; - initialPipeline: object; - result: object; +interface IState { + messages: BaseMessage[]; + esqlQuery: string; + documentation: { + functions: string[]; + commands: string[]; + intention: string; + }; + errors: ValidationResult['errors']; + availableIndices: Record; + invalidQueries: string[]; } -const graphState: StateGraphArgs['channels'] = { - lastExecutedChain: { - value: (x: string, y?: string) => y ?? x, - default: () => '', - }, - rawSamples: { - value: (x: string[], y?: string[]) => y ?? x, - default: () => [], - }, - samples: { - value: (x: string[], y?: string[]) => y ?? x, +// This defines the agent state +const graphState: StateGraphArgs['channels'] = { + messages: { + value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y), default: () => [], }, - formattedSamples: { - value: (x: string, y?: string) => y ?? x, - default: () => '', - }, - ecsTypes: { - value: (x: string, y?: string) => y ?? x, - default: () => '', - }, - ecsCategories: { - value: (x: string, y?: string) => y ?? x, - default: () => '', - }, - exAnswer: { - value: (x: string, y?: string) => y ?? x, - default: () => '', - }, - packageName: { + esqlQuery: { value: (x: string, y?: string) => y ?? x, default: () => '', }, - dataStreamName: { - value: (x: string, y?: string) => y ?? x, - default: () => '', - }, - finalized: { - value: (x: boolean, y?: boolean) => y ?? x, - default: () => false, - }, - reviewed: { - value: (x: boolean, y?: boolean) => y ?? x, - default: () => false, + documentation: { + value: ( + x: { + functions: string[]; + commands: string[]; + intention: string; + }, + y?: { + functions: string[]; + commands: string[]; + intention: string; + } + ) => y ?? x, + default: () => ({ + functions: [], + commands: [], + intention: '', + }), }, errors: { - value: (x: object, y?: object) => y ?? x, - default: () => ({}), - }, - pipelineResults: { - value: (x: object[], y?: object[]) => y ?? x, - default: () => [{}], - }, - currentPipeline: { - value: (x: object, y?: object) => y ?? x, - default: () => ({}), - }, - currentProcessors: { - value: (x: object[], y?: object[]) => y ?? x, + value: (x, y) => (y && !y.length ? y : x.concat(y)), default: () => [], }, - invalidCategorization: { - value: (x: object, y?: object) => y ?? x, - default: () => ({}), - }, - initialPipeline: { - value: (x: object, y?: object) => y ?? x, - default: () => ({}), + invalidQueries: { + value: (x, y) => x.concat(y), + default: () => [], }, - result: { - value: (x: object, y?: object) => y ?? x, + availableIndices: { + value: (x, y) => y ?? x, default: () => ({}), }, }; -function modelInput(state: CategorizationState): Partial { - // const samples = modifySamples(state); - // const formattedSamples = formatSamples(samples); - // const initialPipeline = JSON.parse(JSON.stringify(state.currentPipeline)); - return { - // exAnswer: JSON.stringify(CATEGORIZATION_EXAMPLE_ANSWER, null, 2), - // ecsCategories: JSON.stringify(ECS_CATEGORIES, null, 2), - // ecsTypes: JSON.stringify(ECS_TYPES, null, 2), - // samples, - // formattedSamples, - // initialPipeline, - finalized: false, - reviewed: false, - lastExecutedChain: 'modelInput', +const readFile = promisify(Fs.readFile); +const readdir = promisify(Fs.readdir); + +const loadSystemMessage = once(async () => { + const data = await readFile( + Path.join( + __dirname, + '../../../../../observability_solution/observability_ai_assistant_app/server/functions/query/system_message.txt' + ) + ); + return data.toString('utf-8'); +}); + +const loadEsqlDocs = once(async () => { + const dir = Path.join( + __dirname, + '../../../../../observability_solution/observability_ai_assistant_app/server/functions/query/esql_docs' + ); + const files = (await readdir(dir)).filter((file) => Path.extname(file) === '.txt'); + + if (!files.length) { + return {}; + } + + const limiter = pLimit(10); + return keyBy( + await Promise.all( + files.map((file) => + limiter(async () => { + const data = (await readFile(Path.join(dir, file))).toString('utf-8'); + const filename = Path.basename(file, '.txt'); + + const keyword = filename + .replace('esql-', '') + .replace('agg-', '') + .replaceAll('-', '_') + .toUpperCase(); + + return { + keyword: keyword === 'STATS_BY' ? 'STATS' : keyword, + data, + }; + }) + ) + ), + 'keyword' + ); +}); + +// function transformInputToOutput(input2: Record) { +// const output: Record = {}; + +// forOwn(input2, (value, _key) => { +// const key = Object.keys(value.aliases)[0] || _key; +// const properties = get(value, 'mappings.properties', {}); +// output[key] = transformProperties(properties); +// }); + +// return output; +// } + +// function transformProperties(properties) { +// const result: Record = {}; + +// forOwn(properties, (value, _key) => { +// const key = value.data_stream || _key; +// if (value.type) { +// result[key] = value.type; +// } else if (value.properties) { +// result[key] = transformProperties(value.properties); +// } else if (value.fields) { +// result[key] = { fields: {} }; +// forOwn(value.fields, (fieldValue, fieldKey) => { +// result[key].fields[fieldKey] = { type: fieldValue.type }; +// if (fieldValue.ignore_above !== undefined) { +// result[key].fields[fieldKey].ignore_above = fieldValue.ignore_above; +// } +// }); +// } else { +// result[key] = value; +// } +// }); + +// return result; +// } + +const getClassifyEsql = + ({ + userQuery, + llm, + esClient, + }: { + userQuery: string; + llm: NonNullable; + esClient: AssistantToolParams['esClient']; + }) => + async (state: IState, config?: RunnableConfig) => { + const [systemMessage, esqlDocs] = await Promise.all([loadSystemMessage(), loadEsqlDocs()]); + + const formatInstructions = `Respond only in valid JSON. The JSON object you return should match the following schema: + {{ commands: [], functions: [], intention: string }} + + Where commands is a list of processing or source commands that are referenced in the list of commands in this conversation. + Where functions is a list of functions that are referenced in the list of functions in this conversation. + Where intention is the user\'s intention. + `; + + const prompt = await ChatPromptTemplate.fromMessages([ + [ + 'system', + `${systemMessage} Answer the user query. Wrap the output in \`json\` tags\n{format_instructions}`, + ], + [ + 'user', + `Use this function to determine: + - what ES|QL functions and commands are candidates for answering the user's question + - whether the user has requested a query, and if so, it they want it to be executed, or just shown. + + All parameters are required. Make sure the functions and commands you request are available in the + system message. + + {query} + `, + ], + ]).partial({ + format_instructions: formatInstructions, + }); + + // Set up a parser + const parser = new JsonOutputParser(); + + const chainss = prompt.pipe(llm).pipe(parser); + + const result = await chainss.invoke({ + query: userQuery, + date: 'date', + msg: 'msg', + ip: 'ip', + }); + + const keywords = [ + ...(result.commands ?? []), + ...(result.functions ?? []), + 'SYNTAX', + 'OVERVIEW', + 'OPERATORS', + ].map((keyword) => keyword.toUpperCase()); + + const messagesToInclude = mapValues(pick(esqlDocs, keywords), ({ data }) => data); + + const allDataStreams = await esClient.indices.getDataStream(); + + return { + documentation: messagesToInclude, + availableIndices: map(allDataStreams.data_streams, 'name'), + }; }; -} -function modelOutput(state: CategorizationState): Partial { - return { - finalized: true, - lastExecutedChain: 'modelOutput', - result: { - query: state.query, - rows: state.pipelineResults, - columns: state.currentPipeline, - }, +const getGenerateQuery = + ({ userQuery, llm }: { userQuery: string; llm: NonNullable }) => + async (state: IState) => { + const [systemMessage] = await Promise.all([loadSystemMessage()]); + + const answerPrompt = await ChatPromptTemplate.fromMessages([ + [ + 'system', + `${systemMessage}\nDocumentation: {documentation}\nAvailable indices: {availableIndices}`, + ], + [ + 'user', + `Answer the user's question that was previously asked ("{query}...") using the attached documentation. Take into account any previous errors {errors} and invalid ES|QL queries {invalidQueries}. + + Format any ES|QL query as follows: + \`\`\`esql + + \`\`\` + + Respond in plain text. Do not attempt to use a function. + + You must use commands and functions for which you have requested documentation. + + DO NOT UNDER ANY CIRCUMSTANCES generate more than a single query. + If multiple queries are needed, do it as a follow-up step. Make this clear to the user. For example: + + Human: plot both yesterday's and today's data. + + Assistant: Here's how you can plot yesterday's data: + \`\`\`esql + + \`\`\` + + Let's see that first. We'll look at today's data next. + + Human: + + Assistant: Let's look at today's data: + + \`\`\`esql + + \`\`\` + + DO NOT UNDER ANY CIRCUMSTANCES use commands or functions that are not a capability of ES|QL + as mentioned in the system message and documentation. When converting queries from one language + to ES|QL, make sure that the functions are available and documented in ES|QL. + E.g., for SPL's LEN, use LENGTH. For IF, use CASE. + `, + ], + ]).partial({ + documentation: JSON.stringify(state.documentation), + availableIndices: JSON.stringify(state.availableIndices), + errors: state.errors.join('\n'), + invalidQueries: state.invalidQueries.join('\n'), + }); + + const finalChain = answerPrompt.pipe(llm).pipe(new StringOutputParser()); + + const finalResult = await finalChain.invoke({ + query: userQuery, + date: 'date', + msg: 'msg', + ip: 'ip', + }); + + const correctedResult = finalResult.replaceAll(INLINE_ESQL_QUERY_REGEX, (_match, query) => { + const correction = correctCommonEsqlMistakes(query); + // if (correction.isCorrection) { + // logger.error(`Corrected query, from: \n${correction.input}\nto:\n${correction.output}`); + // } + return `\`\`\`esql\n${correction.output}\n\`\`\``; + }); + + const esqlQuery = correctedResult.match(new RegExp(INLINE_ESQL_QUERY_REGEX, 'ms'))?.[1]; + + return { answer: correctedResult, esqlQuery }; }; -} -const handleGenerateQuery = (state: CategorizationState, model) => {}; +const getValidateQuery = + ({ search }: { search: AssistantToolParams['search'] }) => + async (state: IState, config?: RunnableConfig) => { + const { errors } = await validateQuery(state.esqlQuery, getAstAndSyntaxErrors, { + // setting this to true, we don't want to validate the index / fields existence + ignoreOnMissingCallbacks: true, + }); -const getEsqlGraph = (client, model) => { - const workflow = new StateGraph({ - channels: graphState, - }) - .addNode('modelInput', modelInput) - .addNode('modelOutput', modelOutput) - .addNode('handleGenerateQuery', (state: CategorizationState) => - handleGenerateQuery(state, model) - ) - .addNode('handleClassifyEsql', (state: CategorizationState) => {}); + if (!isEmpty(errors)) { + return { errors, invalidQueries: [state.esqlQuery] }; + } + + let esqlQueryColumns; + try { + esqlQueryColumns = await getESQLQueryColumns({ + esqlQuery: state.esqlQuery, + search: search.search, + }); + return { errors: [] }; + } catch (e) { + return { errors: [e?.message], invalidQueries: [state.esqlQuery] }; + } + }; + +const shouldRegenerate = (state: IState) => { + if (state.errors?.length) { + return 'generateQuery'; + } + + return END; }; export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { @@ -208,96 +418,31 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { const { chain, esClient, search, llm } = params as EsqlKnowledgeBaseToolParams; if (chain == null) return null; - console.error('params', Object.keys(params)); - - console.error('search', search); - return new DynamicStructuredTool({ name: toolDetails.name, description: toolDetails.description, schema: z.object({ - question: z.string().describe(`The user's exact question about ESQL`), + query: z.string().describe(`The user's exact question about ESQL`), }), func: async (input, _, cbManager) => { - let response; - try { - response = await esClient.indices.getDataStream(); - } catch (e) { - console.error('e', e); - } - - console.error('response', JSON.stringify(response, null, 2)); - - function transformInputToOutput(input2: Record) { - const output: Record = {}; - - forOwn(input2, (value, _key) => { - const key = Object.keys(value.aliases)[0] || _key; - const properties = get(value, 'mappings.properties', {}); - output[key] = transformProperties(properties); - }); - - return output; - } - - function transformProperties(properties) { - const result: Record = {}; - - forOwn(properties, (value, key) => { - if (value.type) { - result[key] = value.type; - } else if (value.properties) { - result[key] = transformProperties(value.properties); - } else if (value.fields) { - result[key] = { fields: {} }; - forOwn(value.fields, (fieldValue, fieldKey) => { - result[key].fields[fieldKey] = { type: fieldValue.type }; - if (fieldValue.ignore_above !== undefined) { - result[key].fields[fieldKey].ignore_above = fieldValue.ignore_above; - } - }); - } else { - result[key] = value; - } - }); - - return result; - } - - // const indices = transformInputToOutput(response); - - const graph = ECS_MAIN_PROMPT.pipe(llm).pipe(new StringOutputParser()); - - const result = await graph.invoke({ - input: input.question, - availableIndices: JSON.stringify( - response?.data_streams.map((item) => item.name), - null, - 2 - ), - }); - - console.error('result', JSON.stringify(result, null, 2)); - - const esqlQuery = result - ?.match(/(?<=```esql)[\s\S]*?(?=```)/g) - .join('') - .replace('\n', '') - .replaceAll('"', '') - .trim(); - - console.error('esqlQuery', esqlQuery); - - try { - const esqlQueryColumns = await getESQLQueryColumns({ esqlQuery, search: search.search }); - console.error('esqlQueryColumns', esqlQueryColumns); - } catch (e) { - console.error('e', e); - } - - console.error('trimee', result.replaceAll('"', '').trim()); - - return result.replaceAll('"', '').trim(); + const workflow = new StateGraph({ + channels: graphState, + }) + .addNode('classifyEsql', getClassifyEsql({ userQuery: input.query, llm, esClient })) + .addNode('generateQuery', getGenerateQuery({ userQuery: input.query, llm })) + .addNode('validateQuery', getValidateQuery({ search })) + .addEdge(START, 'classifyEsql') + .addEdge('classifyEsql', 'generateQuery') + .addEdge('generateQuery', 'validateQuery') + .addConditionalEdges('validateQuery', shouldRegenerate); + + const app = workflow.compile(); + + const query = await app.invoke({}, { recursionLimit: 20 }); + + console.error('graph result', query); + + return query.esqlQuery; }, tags: ['esql', 'query-generation', 'knowledge-base'], }); diff --git a/yarn.lock b/yarn.lock index 5d67cf25cdf77..fcd2d72be4939 100644 --- a/yarn.lock +++ b/yarn.lock @@ -6965,33 +6965,33 @@ resolved "https://registry.yarnpkg.com/@kwsites/promise-deferred/-/promise-deferred-1.1.1.tgz#8ace5259254426ccef57f3175bc64ed7095ed919" integrity sha512-GaHYm+c0O9MjZRu0ongGBRbinu8gVAMd2UZjji6jVmqKtZluZnptXGWhz1E8j8D2HJ3f/yMxKAUC0b+57wncIw== -"@langchain/community@^0.2.4": - version "0.2.4" - resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.4.tgz#fb5feb4f4a01a1b33adfd28ce7126d0dedb3e6d1" - integrity sha512-rwrPNQLyIe84TPqPYbYOfDA4G/ba1rdj7OtZg63dQmxIvNDOmUCh4xIQac2iuRUnM3o4Ben0Faa9qz+V5oPgIA== +"@langchain/community@^0.2.13": + version "0.2.13" + resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.13.tgz#7a327663d2a6006456ff136d2b16cb2b1a76b541" + integrity sha512-f0GZCGM5XP0r+H643GpUU4YelKHsUdhUY1Kb8rKpCoy8zgs1nUkiYDVylAf0ezwUOT4NYCEuwpw0jj8hQSLn1Q== dependencies: - "@langchain/core" "~0.2.0" - "@langchain/openai" "~0.0.28" + "@langchain/core" "~0.2.9" + "@langchain/openai" "~0.1.0" binary-extensions "^2.2.0" expr-eval "^2.0.2" flat "^5.0.2" js-yaml "^4.1.0" langchain "0.2.3" - langsmith "~0.1.1" + langsmith "~0.1.30" uuid "^9.0.0" zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@0.2.3", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.56 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@~0.2.0": - version "0.2.3" - resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.3.tgz#7faa82f92b0c7843506e827a38bfcbb60f009d13" - integrity sha512-mVuFHSLpPQ4yOHNXeoSA3LnmIMuFmUiit5rvbYcPZqM6SrB2zCNN2nD4Ty5+3H5X4tYItDoSqsTuUNUQySXRQw== +"@langchain/core@0.2.9", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@>=0.2.8 <0.3.0", "@langchain/core@^0.2.9", "@langchain/core@~0.2.0", "@langchain/core@~0.2.9": + version "0.2.9" + resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.9.tgz#8a22585ef11d2ca8742a8bbfe77dd25baedc7779" + integrity sha512-pJshopBZqMNF020q0OrrO+vfApWTZUlZecRYMM7TWA5M8/zvEyU/mgA9DlzeRjjDmG6pwF6dIKVjpl6fIGVXlQ== dependencies: ansi-styles "^5.0.0" camelcase "6" decamelize "1.2.0" js-tiktoken "^1.0.12" - langsmith "~0.1.7" + langsmith "~0.1.30" ml-distance "^4.0.0" mustache "^4.2.0" p-queue "^6.6.2" @@ -7000,22 +7000,22 @@ zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/langgraph@^0.0.23": - version "0.0.23" - resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.23.tgz#34b5ad5dc9fe644ee96bcfcf11197ec1d7f9e0e2" - integrity sha512-pXlcsBOseT5xdf9enUqbLQ/59LaZxgMI2dL2vFJ+EpcoK7bQnlzzhRtRPp+vubMyMeEKRoAXlaA9ObwpVi93CA== +"@langchain/langgraph@^0.0.24": + version "0.0.24" + resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.24.tgz#814e686dcef700d15a6f90e27e9c9f79d75faef4" + integrity sha512-fW9cnz62oKZFAlyO/4oEjXxthrqZPQtqyX4f7ttKEi0rJZKeuoohvOtnC8faq6nrtMtX9JpLixHjK0SgN7XN3g== dependencies: "@langchain/core" ">0.1.61 <0.3.0" uuid "^9.0.1" -"@langchain/openai@^0.0.34", "@langchain/openai@~0.0.28": - version "0.0.34" - resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.0.34.tgz#36c9bca0721ab9f7e5d40927e7c0429cacbd5b56" - integrity sha512-M+CW4oXle5fdoz2T2SwdOef8pl3/1XmUx1vjn2mXUVM/128aO0l23FMF0SNBsAbRV6P+p/TuzjodchJbi0Ht/A== +"@langchain/openai@0.2.0", "@langchain/openai@>=0.1.0 <0.3.0", "@langchain/openai@^0.2.0", "@langchain/openai@~0.1.0": + version "0.2.0" + resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.2.0.tgz#342e49d15b946fa01128d1bb81357e688e7cf567" + integrity sha512-gZd+0IOxpiKuh1m6KTT5vtUoOO72GEYyoU4+c6qAUucOEqQS0Vvz3lMGyNWLjK4x4Xpd+r8GAF5mj/jvghwP1A== dependencies: - "@langchain/core" ">0.1.56 <0.3.0" + "@langchain/core" ">=0.2.8 <0.3.0" js-tiktoken "^1.0.12" - openai "^4.41.1" + openai "^4.49.1" zod "^3.22.4" zod-to-json-schema "^3.22.3" @@ -21710,20 +21710,20 @@ kuler@^2.0.0: resolved "https://registry.yarnpkg.com/kuler/-/kuler-2.0.0.tgz#e2c570a3800388fb44407e851531c1d670b061b3" integrity sha512-Xq9nH7KlWZmXAtodXDDRE7vs6DU1gTU8zYDHDiWLSip45Egwq3plLHzPn27NgvzL2r1LMPC1vdqh98sQxtqj4A== -langchain@0.2.3: - version "0.2.3" - resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.3.tgz#c14bb05cf871b21bd63b84b3ab89580b1d62539f" - integrity sha512-T9xR7zd+Nj0oXy6WoYKmZLy0DlQiDLFPGYWdOXDxy+AvqlujoPdVQgDSpdqiOHvAjezrByAoKxoHCz5XMwTP/Q== +langchain@0.2.3, langchain@0.2.6, langchain@^0.2.6: + version "0.2.6" + resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.6.tgz#22249707ba800c38ec9a6cca36e383a881227393" + integrity sha512-vDAJHGu/lA4pn3hkyzSC6RiaZhtj0ozfRyG8L6J2vCnXyJV/lgk9uGMP2x645EBrSozBMHJBng1UYeaUR/1fQQ== dependencies: "@langchain/core" "~0.2.0" - "@langchain/openai" "~0.0.28" + "@langchain/openai" ">=0.1.0 <0.3.0" "@langchain/textsplitters" "~0.0.0" binary-extensions "^2.2.0" js-tiktoken "^1.0.12" js-yaml "^4.1.0" jsonpointer "^5.0.1" langchainhub "~0.0.8" - langsmith "~0.1.7" + langsmith "~0.1.30" ml-distance "^4.0.0" openapi-types "^12.1.3" p-retry "4" @@ -21737,7 +21737,7 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.30, langsmith@~0.1.1, langsmith@~0.1.30, langsmith@~0.1.7: +langsmith@^0.1.30, langsmith@~0.1.30: version "0.1.30" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== @@ -24353,10 +24353,10 @@ open@^8.0.9, open@^8.4.0, open@~8.4.0: is-docker "^2.1.1" is-wsl "^2.2.0" -openai@^4.24.1, openai@^4.41.1: - version "4.47.1" - resolved "https://registry.yarnpkg.com/openai/-/openai-4.47.1.tgz#1d23c7a8eb3d7bcdc69709cd905f4c9af0181dba" - integrity sha512-WWSxhC/69ZhYWxH/OBsLEirIjUcfpQ5+ihkXKp06hmeYXgBBIUCa9IptMzYx6NdkiOCsSGYCnTIsxaic3AjRCQ== +openai@^4.24.1, openai@^4.49.1: + version "4.51.0" + resolved "https://registry.yarnpkg.com/openai/-/openai-4.51.0.tgz#8ab08bba2441375e8e4ce6161f9ac987d2b2c157" + integrity sha512-UKuWc3/qQyklqhHM8CbdXCv0Z0obap6T0ECdcO5oATQxAbKE5Ky3YCXFQY207z+eGG6ez4U9wvAcuMygxhmStg== dependencies: "@types/node" "^18.11.18" "@types/node-fetch" "^2.6.4" From 2e88bfe98df1e2e2638434efc19ea17a623115cf Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 24 Jun 2024 13:18:04 +0200 Subject: [PATCH 16/47] cleanup --- x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx | 2 +- .../actions/server/application/connector/methods/get/get.ts | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx index c220b464634c0..dd96b4883c969 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -640,8 +640,8 @@ const AssistantComponent: React.FC = ({ ), [ - refetchCurrentConversation, abortStream, + refetchCurrentConversation, currentConversation, editingSystemPromptId, getComments, diff --git a/x-pack/plugins/actions/server/application/connector/methods/get/get.ts b/x-pack/plugins/actions/server/application/connector/methods/get/get.ts index 46a2cb3821ff2..2d4a94f5615d7 100644 --- a/x-pack/plugins/actions/server/application/connector/methods/get/get.ts +++ b/x-pack/plugins/actions/server/application/connector/methods/get/get.ts @@ -61,9 +61,7 @@ export async function get({ connector = { id, actionTypeId: foundInMemoryConnector.actionTypeId, - isMissingSecrets: foundInMemoryConnector.isMissingSecrets, name: foundInMemoryConnector.name, - config: foundInMemoryConnector.config, isPreconfigured: foundInMemoryConnector.isPreconfigured, isSystemAction: foundInMemoryConnector.isSystemAction, isDeprecated: isConnectorDeprecated(foundInMemoryConnector), From 48259c90da006efc67172f7c74e68a13a915f458 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 24 Jun 2024 15:23:30 +0200 Subject: [PATCH 17/47] fix --- package.json | 2 +- yarn.lock | 13 ++++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/package.json b/package.json index e966403807dfa..35b6dcf702c64 100644 --- a/package.json +++ b/package.json @@ -1073,7 +1073,7 @@ "jsts": "^1.6.2", "kea": "^2.6.0", "langchain": "^0.2.6", - "langsmith": "^0.1.32", + "langsmith": "^0.1.30", "launchdarkly-js-client-sdk": "^3.3.0", "launchdarkly-node-server-sdk": "^7.0.3", "load-json-file": "^6.2.0", diff --git a/yarn.lock b/yarn.lock index fcd2d72be4939..01b730ca8909c 100644 --- a/yarn.lock +++ b/yarn.lock @@ -21737,7 +21737,18 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.30, langsmith@~0.1.30: +langsmith@^0.1.30: + version "0.1.32" + resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.32.tgz#38938b0e8685522087b697b8200c488c6490c137" + integrity sha512-EUWHIH6fiOCGRYdzgwGoXwJxCMyUrL+bmUcxoVmkXoXoAGDOVinz8bqJLKbxotsQWqM64NKKsW85OTIutgNaMQ== + dependencies: + "@types/uuid" "^9.0.1" + commander "^10.0.1" + p-queue "^6.6.2" + p-retry "4" + uuid "^9.0.0" + +langsmith@~0.1.30: version "0.1.30" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== From a46732b0958ca4ab43276d75fb006f65ab64ad9a Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Mon, 24 Jun 2024 13:43:54 +0000 Subject: [PATCH 18/47] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/plugins/cases/tsconfig.json | 4 ++++ x-pack/plugins/security_solution/tsconfig.json | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/x-pack/plugins/cases/tsconfig.json b/x-pack/plugins/cases/tsconfig.json index 535f4e5e106dc..f1a1f4fef9194 100644 --- a/x-pack/plugins/cases/tsconfig.json +++ b/x-pack/plugins/cases/tsconfig.json @@ -75,6 +75,10 @@ "@kbn/core-logging-browser-mocks", "@kbn/data-views-plugin", "@kbn/core-http-router-server-internal", + "@kbn/elastic-assistant", + "@kbn/esql-utils", + "@kbn/esql-datagrid", + "@kbn/visualization-utils", ], "exclude": [ "target/**/*", diff --git a/x-pack/plugins/security_solution/tsconfig.json b/x-pack/plugins/security_solution/tsconfig.json index f1320dc3205be..94a668ea384e4 100644 --- a/x-pack/plugins/security_solution/tsconfig.json +++ b/x-pack/plugins/security_solution/tsconfig.json @@ -209,5 +209,9 @@ "@kbn/core-analytics-browser", "@kbn/core-i18n-browser", "@kbn/core-theme-browser", + "@kbn/esql-datagrid", + "@kbn/visualization-utils", + "@kbn/esql-validation-autocomplete", + "@kbn/esql-ast", ] } From 04140f9842b466d74b0b19bbe3663ca76471870a Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 24 Jun 2024 16:17:47 +0200 Subject: [PATCH 19/47] cleanup --- .../impl/capabilities/index.ts | 4 +- .../routes/post_actions_connector_execute.ts | 1 - .../common/experimental_features.ts | 7 +- .../scripts/run_cypress/parallel.ts | 5 +- .../esql_language_knowledge_base_tool.ts | 422 +----------------- .../graph_esql_language_tool.ts | 410 +++++++++++++++++ .../server/assistant/tools/index.ts | 4 +- .../security_solution/server/plugin.ts | 5 +- 8 files changed, 441 insertions(+), 417 deletions(-) create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index 5f71f9d4fa3f3..c1c101fd74cd8 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: true, - assistantModelEvaluation: true, + assistantKnowledgeBaseByDefault: false, + assistantModelEvaluation: false, }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 80d9217fc622d..ed1024a6c5395 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -72,7 +72,6 @@ export const postActionsConnectorExecuteRoute = ( }, }, }, - // eslint-disable-next-line complexity async (context, request, response) => { const abortSignal = getRequestAbortedSignal(request.events.aborted$); diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 4fb429635bfd5..809a87b2f53bc 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -127,7 +127,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables the Assistant Knowledge Base by default, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: true, + assistantKnowledgeBaseByDefault: false, /** * Enables the Managed User section inside the new user details flyout. @@ -267,6 +267,11 @@ export const allowedExperimentalValues = Object.freeze({ * Adds a new option to filter descendants of a process for Management / Event Filters */ filterProcessDescendantsForEventFiltersEnabled: false, + + /** + * Adds a new langgraph-based ESQL generation tool + */ + graphEsqlTool: true, }); type ExperimentalConfigKeys = Array; diff --git a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts index aec694c14b09c..4b0f586779270 100644 --- a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts +++ b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts @@ -337,7 +337,10 @@ ${JSON.stringify( procs, config, installDir: options?.installDir, - extraKbnOpts: options?.installDir || options?.ci || !isOpen ? [] : ['--dev'], + extraKbnOpts: + options?.installDir || options?.ci || !isOpen + ? [] + : ['--dev', '--no-dev-config', '--no-dev-credentials'], onEarlyExit, inspect: argv.inspect, }); 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 4078e958a539f..692753a22dea0 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 @@ -5,406 +5,19 @@ * 2.0. */ -import Fs from 'fs'; -import { keyBy, mapValues, once, pick, forOwn, get, isEmpty, map } from 'lodash'; import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from 'zod'; -import pLimit from 'p-limit'; -import Path from 'path'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; -import { JsonOutputParser, StringOutputParser } from '@langchain/core/output_parsers'; -import type { StateGraphArgs } from '@langchain/langgraph'; -import { END, START, StateGraph } from '@langchain/langgraph'; -import type { BaseMessage } from '@langchain/core/messages'; -import type { RunnableConfig } from '@langchain/core/runnables'; -import type { ValidationResult } from '@kbn/esql-validation-autocomplete/src/validation/types'; -import { getESQLQueryColumns } from '@kbn/esql-utils'; -import { promisify } from 'util'; -import { ChatPromptTemplate } from '@langchain/core/prompts'; -import { validateQuery } from '@kbn/esql-validation-autocomplete'; -import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; import { APP_UI_ID } from '../../../../common'; -import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; - -export const INLINE_ESQL_QUERY_REGEX = /```esql\s*(.*?)\s*```/gms; - -export const ECS_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ - [ - 'system', - `As an expert user of Elastic Security, please generate an accurate and valid ESQL query to detect the use case below. Your response should be formatted to be able to use immediately in an Elastic Security timeline or detection rule. Take your time with the answer, check your knowledge really well on all the functions I am asking for. For ES|QL answers specifically, you should only ever answer with what's available in your private knowledge. I cannot afford for queries to be inaccurate. Assume I am using the Elastic Common Schema and Elastic Agent. Under any circumstances wrap index in quotes. - - If multiple indices are matched please try to use wildcard to match all indices. If you are unsure about the index name, please refer to the context provided below. - -Here is some context for you to reference for your task, read it carefully as you will get questions about it later: - - -{availableIndices} - -`, - ], - [ - 'human', - `{input}. -Example response format: - -A: Please find the ESQL query below: -\`\`\`esql -FROM logs-* -| SORT @timestamp DESC -| LIMIT 5 -\`\`\` -"`, - ], - ['ai', 'Please find the ESQL query below:'], -]); export type EsqlKnowledgeBaseToolParams = AssistantToolParams; -const TOOL_NAME = 'ESQLKnowledgeBaseTool'; - 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. Only output valid ES|QL queries as described above. Do not add any additional text to describe your output.', id: 'esql-knowledge-base-tool', - name: TOOL_NAME, - description: `You MUST use the "${TOOL_NAME}" function when the user wants to: - - visualize data - - 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. - - If the user asks for a query, and one of the dataset info functions was called and returned no results, you should still call the query function to generate an example query. - - 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.`, -}; - -interface IState { - messages: BaseMessage[]; - esqlQuery: string; - documentation: { - functions: string[]; - commands: string[]; - intention: string; - }; - errors: ValidationResult['errors']; - availableIndices: Record; - invalidQueries: string[]; -} - -// This defines the agent state -const graphState: StateGraphArgs['channels'] = { - messages: { - value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y), - default: () => [], - }, - esqlQuery: { - value: (x: string, y?: string) => y ?? x, - default: () => '', - }, - documentation: { - value: ( - x: { - functions: string[]; - commands: string[]; - intention: string; - }, - y?: { - functions: string[]; - commands: string[]; - intention: string; - } - ) => y ?? x, - default: () => ({ - functions: [], - commands: [], - intention: '', - }), - }, - errors: { - value: (x, y) => (y && !y.length ? y : x.concat(y)), - default: () => [], - }, - invalidQueries: { - value: (x, y) => x.concat(y), - default: () => [], - }, - availableIndices: { - value: (x, y) => y ?? x, - default: () => ({}), - }, + name: 'ESQLKnowledgeBaseTool', }; - -const readFile = promisify(Fs.readFile); -const readdir = promisify(Fs.readdir); - -const loadSystemMessage = once(async () => { - const data = await readFile( - Path.join( - __dirname, - '../../../../../observability_solution/observability_ai_assistant_app/server/functions/query/system_message.txt' - ) - ); - return data.toString('utf-8'); -}); - -const loadEsqlDocs = once(async () => { - const dir = Path.join( - __dirname, - '../../../../../observability_solution/observability_ai_assistant_app/server/functions/query/esql_docs' - ); - const files = (await readdir(dir)).filter((file) => Path.extname(file) === '.txt'); - - if (!files.length) { - return {}; - } - - const limiter = pLimit(10); - return keyBy( - await Promise.all( - files.map((file) => - limiter(async () => { - const data = (await readFile(Path.join(dir, file))).toString('utf-8'); - const filename = Path.basename(file, '.txt'); - - const keyword = filename - .replace('esql-', '') - .replace('agg-', '') - .replaceAll('-', '_') - .toUpperCase(); - - return { - keyword: keyword === 'STATS_BY' ? 'STATS' : keyword, - data, - }; - }) - ) - ), - 'keyword' - ); -}); - -// function transformInputToOutput(input2: Record) { -// const output: Record = {}; - -// forOwn(input2, (value, _key) => { -// const key = Object.keys(value.aliases)[0] || _key; -// const properties = get(value, 'mappings.properties', {}); -// output[key] = transformProperties(properties); -// }); - -// return output; -// } - -// function transformProperties(properties) { -// const result: Record = {}; - -// forOwn(properties, (value, _key) => { -// const key = value.data_stream || _key; -// if (value.type) { -// result[key] = value.type; -// } else if (value.properties) { -// result[key] = transformProperties(value.properties); -// } else if (value.fields) { -// result[key] = { fields: {} }; -// forOwn(value.fields, (fieldValue, fieldKey) => { -// result[key].fields[fieldKey] = { type: fieldValue.type }; -// if (fieldValue.ignore_above !== undefined) { -// result[key].fields[fieldKey].ignore_above = fieldValue.ignore_above; -// } -// }); -// } else { -// result[key] = value; -// } -// }); - -// return result; -// } - -const getClassifyEsql = - ({ - userQuery, - llm, - esClient, - }: { - userQuery: string; - llm: NonNullable; - esClient: AssistantToolParams['esClient']; - }) => - async (state: IState, config?: RunnableConfig) => { - const [systemMessage, esqlDocs] = await Promise.all([loadSystemMessage(), loadEsqlDocs()]); - - const formatInstructions = `Respond only in valid JSON. The JSON object you return should match the following schema: - {{ commands: [], functions: [], intention: string }} - - Where commands is a list of processing or source commands that are referenced in the list of commands in this conversation. - Where functions is a list of functions that are referenced in the list of functions in this conversation. - Where intention is the user\'s intention. - `; - - const prompt = await ChatPromptTemplate.fromMessages([ - [ - 'system', - `${systemMessage} Answer the user query. Wrap the output in \`json\` tags\n{format_instructions}`, - ], - [ - 'user', - `Use this function to determine: - - what ES|QL functions and commands are candidates for answering the user's question - - whether the user has requested a query, and if so, it they want it to be executed, or just shown. - - All parameters are required. Make sure the functions and commands you request are available in the - system message. - - {query} - `, - ], - ]).partial({ - format_instructions: formatInstructions, - }); - - // Set up a parser - const parser = new JsonOutputParser(); - - const chainss = prompt.pipe(llm).pipe(parser); - - const result = await chainss.invoke({ - query: userQuery, - date: 'date', - msg: 'msg', - ip: 'ip', - }); - - const keywords = [ - ...(result.commands ?? []), - ...(result.functions ?? []), - 'SYNTAX', - 'OVERVIEW', - 'OPERATORS', - ].map((keyword) => keyword.toUpperCase()); - - const messagesToInclude = mapValues(pick(esqlDocs, keywords), ({ data }) => data); - - const allDataStreams = await esClient.indices.getDataStream(); - - return { - documentation: messagesToInclude, - availableIndices: map(allDataStreams.data_streams, 'name'), - }; - }; - -const getGenerateQuery = - ({ userQuery, llm }: { userQuery: string; llm: NonNullable }) => - async (state: IState) => { - const [systemMessage] = await Promise.all([loadSystemMessage()]); - - const answerPrompt = await ChatPromptTemplate.fromMessages([ - [ - 'system', - `${systemMessage}\nDocumentation: {documentation}\nAvailable indices: {availableIndices}`, - ], - [ - 'user', - `Answer the user's question that was previously asked ("{query}...") using the attached documentation. Take into account any previous errors {errors} and invalid ES|QL queries {invalidQueries}. - - Format any ES|QL query as follows: - \`\`\`esql - - \`\`\` - - Respond in plain text. Do not attempt to use a function. - - You must use commands and functions for which you have requested documentation. - - DO NOT UNDER ANY CIRCUMSTANCES generate more than a single query. - If multiple queries are needed, do it as a follow-up step. Make this clear to the user. For example: - - Human: plot both yesterday's and today's data. - - Assistant: Here's how you can plot yesterday's data: - \`\`\`esql - - \`\`\` - - Let's see that first. We'll look at today's data next. - - Human: - - Assistant: Let's look at today's data: - - \`\`\`esql - - \`\`\` - - DO NOT UNDER ANY CIRCUMSTANCES use commands or functions that are not a capability of ES|QL - as mentioned in the system message and documentation. When converting queries from one language - to ES|QL, make sure that the functions are available and documented in ES|QL. - E.g., for SPL's LEN, use LENGTH. For IF, use CASE. - `, - ], - ]).partial({ - documentation: JSON.stringify(state.documentation), - availableIndices: JSON.stringify(state.availableIndices), - errors: state.errors.join('\n'), - invalidQueries: state.invalidQueries.join('\n'), - }); - - const finalChain = answerPrompt.pipe(llm).pipe(new StringOutputParser()); - - const finalResult = await finalChain.invoke({ - query: userQuery, - date: 'date', - msg: 'msg', - ip: 'ip', - }); - - const correctedResult = finalResult.replaceAll(INLINE_ESQL_QUERY_REGEX, (_match, query) => { - const correction = correctCommonEsqlMistakes(query); - // if (correction.isCorrection) { - // logger.error(`Corrected query, from: \n${correction.input}\nto:\n${correction.output}`); - // } - return `\`\`\`esql\n${correction.output}\n\`\`\``; - }); - - const esqlQuery = correctedResult.match(new RegExp(INLINE_ESQL_QUERY_REGEX, 'ms'))?.[1]; - - return { answer: correctedResult, esqlQuery }; - }; - -const getValidateQuery = - ({ search }: { search: AssistantToolParams['search'] }) => - async (state: IState, config?: RunnableConfig) => { - const { errors } = await validateQuery(state.esqlQuery, getAstAndSyntaxErrors, { - // setting this to true, we don't want to validate the index / fields existence - ignoreOnMissingCallbacks: true, - }); - - if (!isEmpty(errors)) { - return { errors, invalidQueries: [state.esqlQuery] }; - } - - let esqlQueryColumns; - try { - esqlQueryColumns = await getESQLQueryColumns({ - esqlQuery: state.esqlQuery, - search: search.search, - }); - return { errors: [] }; - } catch (e) { - return { errors: [e?.message], invalidQueries: [state.esqlQuery] }; - } - }; - -const shouldRegenerate = (state: IState) => { - if (state.errors?.length) { - return 'generateQuery'; - } - - return END; -}; - export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, @@ -415,34 +28,23 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { chain, esClient, search, llm } = params as EsqlKnowledgeBaseToolParams; + const { chain } = params as EsqlKnowledgeBaseToolParams; if (chain == null) return null; return new DynamicStructuredTool({ name: toolDetails.name, description: toolDetails.description, schema: z.object({ - query: z.string().describe(`The user's exact question about ESQL`), + question: z.string().describe(`The user's exact question about ESQL`), }), func: async (input, _, cbManager) => { - const workflow = new StateGraph({ - channels: graphState, - }) - .addNode('classifyEsql', getClassifyEsql({ userQuery: input.query, llm, esClient })) - .addNode('generateQuery', getGenerateQuery({ userQuery: input.query, llm })) - .addNode('validateQuery', getValidateQuery({ search })) - .addEdge(START, 'classifyEsql') - .addEdge('classifyEsql', 'generateQuery') - .addEdge('generateQuery', 'validateQuery') - .addConditionalEdges('validateQuery', shouldRegenerate); - - const app = workflow.compile(); - - const query = await app.invoke({}, { recursionLimit: 20 }); - - console.error('graph result', query); - - return query.esqlQuery; + const result = await chain.invoke( + { + query: input.question, + }, + cbManager + ); + return result.text; }, tags: ['esql', 'query-generation', 'knowledge-base'], }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts new file mode 100644 index 0000000000000..39335479cc719 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -0,0 +1,410 @@ +/* + * 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 Fs from 'fs'; +import { keyBy, mapValues, once, pick, isEmpty, map } from 'lodash'; +import { DynamicStructuredTool } from '@langchain/core/tools'; +import { z } from 'zod'; +import pLimit from 'p-limit'; +import Path from 'path'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { JsonOutputParser, StringOutputParser } from '@langchain/core/output_parsers'; +import type { StateGraphArgs } from '@langchain/langgraph'; +import { END, START, StateGraph } from '@langchain/langgraph'; +import type { BaseMessage } from '@langchain/core/messages'; +import type { RunnableConfig } from '@langchain/core/runnables'; +import type { ValidationResult } from '@kbn/esql-validation-autocomplete/src/validation/types'; +import { getESQLQueryColumns } from '@kbn/esql-utils'; +import { promisify } from 'util'; +import { ChatPromptTemplate } from '@langchain/core/prompts'; +import { validateQuery } from '@kbn/esql-validation-autocomplete'; +import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; +import { APP_UI_ID } from '../../../../common'; +import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; + +export const INLINE_ESQL_QUERY_REGEX = /```esql\s*(.*?)\s*```/gms; + +export const ECS_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ + [ + 'system', + `As an expert user of Elastic Security, please generate an accurate and valid ESQL query to detect the use case below. Your response should be formatted to be able to use immediately in an Elastic Security timeline or detection rule. Take your time with the answer, check your knowledge really well on all the functions I am asking for. For ES|QL answers specifically, you should only ever answer with what's available in your private knowledge. I cannot afford for queries to be inaccurate. Assume I am using the Elastic Common Schema and Elastic Agent. Under any circumstances wrap index in quotes. + + If multiple indices are matched please try to use wildcard to match all indices. If you are unsure about the index name, please refer to the context provided below. + +Here is some context for you to reference for your task, read it carefully as you will get questions about it later: + + +{availableIndices} + +`, + ], + [ + 'human', + `{input}. +Example response format: + +A: Please find the ESQL query below: +\`\`\`esql +FROM logs-* +| SORT @timestamp DESC +| LIMIT 5 +\`\`\` +"`, + ], + ['ai', 'Please find the ESQL query below:'], +]); + +export type GraphESQLToolParams = AssistantToolParams; + +const TOOL_NAME = 'GraphESQLTool'; + +const toolDetails = { + id: 'esql-knowledge-base-tool', + name: TOOL_NAME, + description: `You MUST use the "${TOOL_NAME}" function when the user wants to: + - visualize data + - 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. + + If the user asks for a query, and one of the dataset info functions was called and returned no results, you should still call the query function to generate an example query. + + 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.`, +}; + +interface IState { + messages: BaseMessage[]; + esqlQuery: string; + documentation: { + functions: string[]; + commands: string[]; + intention: string; + }; + errors: ValidationResult['errors']; + availableIndices: Record; + invalidQueries: string[]; +} + +// This defines the agent state +const graphState: StateGraphArgs['channels'] = { + messages: { + value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y), + default: () => [], + }, + esqlQuery: { + value: (x: string, y?: string) => y ?? x, + default: () => '', + }, + documentation: { + value: ( + x: { + functions: string[]; + commands: string[]; + intention: string; + }, + y?: { + functions: string[]; + commands: string[]; + intention: string; + } + ) => y ?? x, + default: () => ({ + functions: [], + commands: [], + intention: '', + }), + }, + errors: { + value: (x, y) => (y && !y.length ? y : x.concat(y)), + default: () => [], + }, + invalidQueries: { + value: (x, y) => x.concat(y), + default: () => [], + }, + availableIndices: { + value: (x, y) => y ?? x, + default: () => ({}), + }, +}; + +const readFile = promisify(Fs.readFile); +const readdir = promisify(Fs.readdir); + +const loadSystemMessage = once(async () => { + const data = await readFile( + Path.join( + __dirname, + '../../../../../observability_solution/observability_ai_assistant_app/server/functions/query/system_message.txt' + ) + ); + return data.toString('utf-8'); +}); + +const loadEsqlDocs = once(async () => { + const dir = Path.join( + __dirname, + '../../../../../observability_solution/observability_ai_assistant_app/server/functions/query/esql_docs' + ); + const files = (await readdir(dir)).filter((file) => Path.extname(file) === '.txt'); + + if (!files.length) { + return {}; + } + + const limiter = pLimit(10); + return keyBy( + await Promise.all( + files.map((file) => + limiter(async () => { + const data = (await readFile(Path.join(dir, file))).toString('utf-8'); + const filename = Path.basename(file, '.txt'); + + const keyword = filename + .replace('esql-', '') + .replace('agg-', '') + .replaceAll('-', '_') + .toUpperCase(); + + return { + keyword: keyword === 'STATS_BY' ? 'STATS' : keyword, + data, + }; + }) + ) + ), + 'keyword' + ); +}); + +const getClassifyEsql = + ({ + userQuery, + llm, + esClient, + }: { + userQuery: string; + llm: NonNullable; + esClient: AssistantToolParams['esClient']; + }) => + async (state: IState, config?: RunnableConfig) => { + const [systemMessage, esqlDocs] = await Promise.all([loadSystemMessage(), loadEsqlDocs()]); + + const formatInstructions = `Respond only in valid JSON. The JSON object you return should match the following schema: + {{ commands: [], functions: [], intention: string }} + + Where commands is a list of processing or source commands that are referenced in the list of commands in this conversation. + Where functions is a list of functions that are referenced in the list of functions in this conversation. + Where intention is the user\'s intention. + `; + + const prompt = await ChatPromptTemplate.fromMessages([ + [ + 'system', + `${systemMessage} Answer the user query. Wrap the output in \`json\` tags\n{format_instructions}`, + ], + [ + 'user', + `Use this function to determine: + - what ES|QL functions and commands are candidates for answering the user's question + - whether the user has requested a query, and if so, it they want it to be executed, or just shown. + + All parameters are required. Make sure the functions and commands you request are available in the + system message. + + {query} + `, + ], + ]).partial({ + format_instructions: formatInstructions, + }); + + // Set up a parser + const parser = new JsonOutputParser(); + + const chainss = prompt.pipe(llm).pipe(parser); + + const result = await chainss.invoke({ + query: userQuery, + date: 'date', + msg: 'msg', + ip: 'ip', + }); + + const keywords = [ + ...(result.commands ?? []), + ...(result.functions ?? []), + 'SYNTAX', + 'OVERVIEW', + 'OPERATORS', + ].map((keyword) => keyword.toUpperCase()); + + const messagesToInclude = mapValues(pick(esqlDocs, keywords), ({ data }) => data); + + const allDataStreams = await esClient.indices.getDataStream(); + + return { + documentation: messagesToInclude, + availableIndices: map(allDataStreams.data_streams, 'name'), + }; + }; + +const getGenerateQuery = + ({ userQuery, llm }: { userQuery: string; llm: NonNullable }) => + async (state: IState) => { + const [systemMessage] = await Promise.all([loadSystemMessage()]); + + const answerPrompt = await ChatPromptTemplate.fromMessages([ + [ + 'system', + `${systemMessage}\nDocumentation: {documentation}\nAvailable indices: {availableIndices}`, + ], + [ + 'user', + `Answer the user's question that was previously asked ("{query}...") using the attached documentation. Take into account any previous errors {errors} and invalid ES|QL queries {invalidQueries}. + + Format any ES|QL query as follows: + \`\`\`esql + + \`\`\` + + Respond in plain text. Do not attempt to use a function. + + You must use commands and functions for which you have requested documentation. + + DO NOT UNDER ANY CIRCUMSTANCES generate more than a single query. + If multiple queries are needed, do it as a follow-up step. Make this clear to the user. For example: + + Human: plot both yesterday's and today's data. + + Assistant: Here's how you can plot yesterday's data: + \`\`\`esql + + \`\`\` + + Let's see that first. We'll look at today's data next. + + Human: + + Assistant: Let's look at today's data: + + \`\`\`esql + + \`\`\` + + DO NOT UNDER ANY CIRCUMSTANCES use commands or functions that are not a capability of ES|QL + as mentioned in the system message and documentation. When converting queries from one language + to ES|QL, make sure that the functions are available and documented in ES|QL. + E.g., for SPL's LEN, use LENGTH. For IF, use CASE. + `, + ], + ]).partial({ + documentation: JSON.stringify(state.documentation), + availableIndices: JSON.stringify(state.availableIndices), + errors: state.errors.join('\n'), + invalidQueries: state.invalidQueries.join('\n'), + }); + + const finalChain = answerPrompt.pipe(llm).pipe(new StringOutputParser()); + + const finalResult = await finalChain.invoke({ + query: userQuery, + date: 'date', + msg: 'msg', + ip: 'ip', + }); + + const correctedResult = finalResult.replaceAll(INLINE_ESQL_QUERY_REGEX, (_match, query) => { + const correction = correctCommonEsqlMistakes(query); + // if (correction.isCorrection) { + // logger.error(`Corrected query, from: \n${correction.input}\nto:\n${correction.output}`); + // } + return `\`\`\`esql\n${correction.output}\n\`\`\``; + }); + + const esqlQuery = correctedResult.match(new RegExp(INLINE_ESQL_QUERY_REGEX, 'ms'))?.[1]; + + return { answer: correctedResult, esqlQuery }; + }; + +const getValidateQuery = + ({ search }: { search: AssistantToolParams['search'] }) => + async (state: IState, config?: RunnableConfig) => { + const { errors } = await validateQuery(state.esqlQuery, getAstAndSyntaxErrors, { + // setting this to true, we don't want to validate the index / fields existence + ignoreOnMissingCallbacks: true, + }); + + if (!isEmpty(errors)) { + return { errors, invalidQueries: [state.esqlQuery] }; + } + + try { + await getESQLQueryColumns({ + esqlQuery: state.esqlQuery, + search: search.search, + }); + return { errors: [] }; + } catch (e) { + return { errors: [e?.message], invalidQueries: [state.esqlQuery] }; + } + }; + +const shouldRegenerate = (state: IState) => { + if (state.errors?.length) { + return 'generateQuery'; + } + + return END; +}; + +export const GRAPH_ESQL_TOOL: AssistantTool = { + ...toolDetails, + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is GraphESQLToolParams => { + const { chain, isEnabledKnowledgeBase, modelExists } = params; + return isEnabledKnowledgeBase && modelExists && chain != null; + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + + const { chain, esClient, search, llm } = params as GraphESQLToolParams; + if (chain == null) return null; + + return new DynamicStructuredTool({ + name: toolDetails.name, + description: toolDetails.description, + schema: z.object({ + query: z.string().describe(`The user's exact question about ESQL`), + }), + func: async (input, _, cbManager) => { + const workflow = new StateGraph({ + channels: graphState, + }) + .addNode('classifyEsql', getClassifyEsql({ userQuery: input.query, llm, esClient })) + .addNode('generateQuery', getGenerateQuery({ userQuery: input.query, llm })) + .addNode('validateQuery', getValidateQuery({ search })) + .addEdge(START, 'classifyEsql') + .addEdge('classifyEsql', 'generateQuery') + .addEdge('generateQuery', 'validateQuery') + .addConditionalEdges('validateQuery', shouldRegenerate); + + const app = workflow.compile(); + + const query = await app.invoke({}, { recursionLimit: 20 }); + + return query.esqlQuery; + }, + tags: ['esql', 'query-generation', 'knowledge-base'], + }); + }, +}; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index 0e5ea3a8f69d1..807e7e543be28 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -8,15 +8,17 @@ import type { AssistantTool } from '@kbn/elastic-assistant-plugin/server'; import { ALERT_COUNTS_TOOL } from './alert_counts/alert_counts_tool'; +import { GRAPH_ESQL_TOOL } from './esql_language_knowledge_base/graph_esql_language_tool'; import { ESQL_KNOWLEDGE_BASE_TOOL } from './esql_language_knowledge_base/esql_language_knowledge_base_tool'; import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool'; import { ATTACK_DISCOVERY_TOOL } from './attack_discovery/attack_discovery_tool'; import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from './knowledge_base/knowledge_base_retrieval_tool'; import { KNOWLEDGE_BASE_WRITE_TOOL } from './knowledge_base/knowledge_base_write_tool'; -export const getAssistantTools = (): AssistantTool[] => [ +export const getAssistantTools = (graphEsqlToolEnabled?: boolean): AssistantTool[] => [ ALERT_COUNTS_TOOL, ATTACK_DISCOVERY_TOOL, + ...(graphEsqlToolEnabled ? [GRAPH_ESQL_TOOL] : []), ESQL_KNOWLEDGE_BASE_TOOL, KNOWLEDGE_BASE_RETRIEVAL_TOOL, KNOWLEDGE_BASE_WRITE_TOOL, diff --git a/x-pack/plugins/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index 5b5b833dd2d4b..3720954163085 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -565,7 +565,10 @@ export class Plugin implements ISecuritySolutionPlugin { this.licensing$ = plugins.licensing.license$; // Assistant Tool and Feature Registration - plugins.elasticAssistant.registerTools(APP_UI_ID, getAssistantTools()); + plugins.elasticAssistant.registerTools( + APP_UI_ID, + getAssistantTools(config.experimentalFeatures.graphEsqlTool) + ); plugins.elasticAssistant.registerFeatures(APP_UI_ID, { assistantKnowledgeBaseByDefault: config.experimentalFeatures.assistantKnowledgeBaseByDefault, assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, From 55bd7c555ebbdc34400ed394bc4ff5641d67a13b Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 24 Jun 2024 16:20:25 +0200 Subject: [PATCH 20/47] fix --- .../custom_codeblock/esql_code_block.tsx | 38 +++++++------------ .../custom_codeblock/esql_code_block.tsx | 15 +++----- 2 files changed, 19 insertions(+), 34 deletions(-) diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx index d8d44910813b6..29f23d96466c4 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx @@ -36,7 +36,7 @@ function generateId() { return uuidv4(); } -const saveVisualizationLabel = i18n.translate('xpack.securityAiAssistant.lensESQLFunction.save', { +const saveVisualizationLabel = i18n.translate('xpack.cases.lensESQLFunction.save', { defaultMessage: 'Save visualization', }); @@ -220,18 +220,12 @@ const EsqlCodeBlockComponent: React.FC = ({ = ({ data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton" aria-label={ isTableVisible - ? i18n.translate( - 'xpack.observabilityAiAssistant.lensESQLFunction.displayChart', - { - defaultMessage: 'Display chart', - } - ) - : i18n.translate( - 'xpack.observabilityAiAssistant.lensESQLFunction.displayTable', - { - defaultMessage: 'Display table', - } - ) + ? i18n.translate('xpack.cases.lensESQLFunction.displayChart', { + defaultMessage: 'Display chart', + }) + : i18n.translate('xpack.cases.lensESQLFunction.displayTable', { + defaultMessage: 'Display table', + }) } /> diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index ecf5f39ed6307..d69f8b7491ef0 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -213,17 +213,14 @@ const EsqlCodeBlockComponent = ({ content={ isTableVisible ? i18n.translate( - 'xpack.observabilityAiAssistant.lensESQLFunction.visualization', + 'xpack.securitySolution.lensESQLFunction.visualization', { defaultMessage: 'Visualization', } ) - : i18n.translate( - 'xpack.observabilityAiAssistant.lensESQLFunction.table', - { - defaultMessage: 'Table of results', - } - ) + : i18n.translate('xpack.securitySolution.lensESQLFunction.table', { + defaultMessage: 'Table of results', + }) } > Date: Mon, 24 Jun 2024 16:59:45 +0200 Subject: [PATCH 21/47] fix --- .../plugins/custom_codeblock/esql_code_block.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx index 29f23d96466c4..032fdc0d1ffe9 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx @@ -316,11 +316,11 @@ EsqlCodeBlockComponent.displayName = 'EsqlCodeBlock'; export const EsqlCodeBlock = React.memo(EsqlCodeBlockComponent); -export const getEsqlRenderer = (timestamp) => (props) => { - return ( +// eslint-disable-next-line react/display-name +export const getEsqlRenderer = (timestamp: string | undefined) => (props) => + ( <> ); -}; From 652891a41507b2a1641e62d68fac348f2c7b8889 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 26 Jun 2024 13:56:47 +0200 Subject: [PATCH 22/47] fix --- .../graph_esql_language_tool.ts | 6 ++---- .../security_solution/server/assistant/tools/index.ts | 3 +-- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index 39335479cc719..00601df5b8afc 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -63,7 +63,7 @@ export type GraphESQLToolParams = AssistantToolParams; const TOOL_NAME = 'GraphESQLTool'; const toolDetails = { - id: 'esql-knowledge-base-tool', + id: 'graph-esql-tool', name: TOOL_NAME, description: `You MUST use the "${TOOL_NAME}" function when the user wants to: - visualize data @@ -325,9 +325,7 @@ const getGenerateQuery = const correctedResult = finalResult.replaceAll(INLINE_ESQL_QUERY_REGEX, (_match, query) => { const correction = correctCommonEsqlMistakes(query); - // if (correction.isCorrection) { - // logger.error(`Corrected query, from: \n${correction.input}\nto:\n${correction.output}`); - // } + return `\`\`\`esql\n${correction.output}\n\`\`\``; }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index 807e7e543be28..f4e9e4f31ff72 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -18,8 +18,7 @@ import { KNOWLEDGE_BASE_WRITE_TOOL } from './knowledge_base/knowledge_base_write export const getAssistantTools = (graphEsqlToolEnabled?: boolean): AssistantTool[] => [ ALERT_COUNTS_TOOL, ATTACK_DISCOVERY_TOOL, - ...(graphEsqlToolEnabled ? [GRAPH_ESQL_TOOL] : []), - ESQL_KNOWLEDGE_BASE_TOOL, + ...(graphEsqlToolEnabled ? [GRAPH_ESQL_TOOL] : [ESQL_KNOWLEDGE_BASE_TOOL]), KNOWLEDGE_BASE_RETRIEVAL_TOOL, KNOWLEDGE_BASE_WRITE_TOOL, OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL, From f70dbf3cb116c869c3c023837275b7818761a257 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 26 Jun 2024 21:14:43 +0200 Subject: [PATCH 23/47] fix --- .../plugins/elastic_assistant/server/types.ts | 6 +- .../graph_esql_language_tool.ts | 59 +++++++++++-------- 2 files changed, 39 insertions(+), 26 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index cfeac08289140..0635286484673 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -206,13 +206,13 @@ export interface AIAssistantPrompts { * Interfaces for registering tools to be used by the elastic assistant */ -export interface AssistantTool { +export interface AssistantTool { id: string; name: string; description: string; sourceRegister: string; isSupported: (params: AssistantToolParams) => boolean; - getTool: (params: AssistantToolParams) => Tool | DynamicStructuredTool | null; + getTool: (params: AssistantToolParams) => Tool | DynamicStructuredTool | null; } export interface AssistantToolParams { @@ -223,7 +223,7 @@ export interface AssistantToolParams { esClient: ElasticsearchClient; kbDataClient?: AIAssistantKnowledgeBaseDataClient; langChainTimeout?: number; - llm?: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; + llm?: ActionsClientChatOpenAI | ActionsClientSimpleChatModel; logger: Logger; modelExists: boolean; onNewReplacements?: (newReplacements: Replacements) => void; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index 00601df5b8afc..1cf4089cf4c9c 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -90,7 +90,7 @@ interface IState { intention: string; }; errors: ValidationResult['errors']; - availableIndices: Record; + availableIndices: string[]; invalidQueries: string[]; } @@ -133,7 +133,7 @@ const graphState: StateGraphArgs['channels'] = { }, availableIndices: { value: (x, y) => y ?? x, - default: () => ({}), + default: () => [], }, }; @@ -248,7 +248,10 @@ const getClassifyEsql = 'OPERATORS', ].map((keyword) => keyword.toUpperCase()); - const messagesToInclude = mapValues(pick(esqlDocs, keywords), ({ data }) => data); + const messagesToInclude = mapValues( + pick(esqlDocs, keywords), + ({ data }) => data + ) as unknown as IState['documentation']; const allDataStreams = await esClient.indices.getDataStream(); @@ -365,7 +368,11 @@ const shouldRegenerate = (state: IState) => { return END; }; -export const GRAPH_ESQL_TOOL: AssistantTool = { +const schema = z.object({ + query: z.string().describe(`The user's exact question about ESQL`), +}); + +export const GRAPH_ESQL_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is GraphESQLToolParams => { @@ -381,26 +388,32 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { return new DynamicStructuredTool({ name: toolDetails.name, description: toolDetails.description, - schema: z.object({ - query: z.string().describe(`The user's exact question about ESQL`), - }), + schema, func: async (input, _, cbManager) => { - const workflow = new StateGraph({ - channels: graphState, - }) - .addNode('classifyEsql', getClassifyEsql({ userQuery: input.query, llm, esClient })) - .addNode('generateQuery', getGenerateQuery({ userQuery: input.query, llm })) - .addNode('validateQuery', getValidateQuery({ search })) - .addEdge(START, 'classifyEsql') - .addEdge('classifyEsql', 'generateQuery') - .addEdge('generateQuery', 'validateQuery') - .addConditionalEdges('validateQuery', shouldRegenerate); - - const app = workflow.compile(); - - const query = await app.invoke({}, { recursionLimit: 20 }); - - return query.esqlQuery; + if (llm) { + const workflow = new StateGraph({ + channels: graphState, + }) + .addNode('classifyEsql', getClassifyEsql({ userQuery: input.query, llm, esClient })) + .addNode('generateQuery', getGenerateQuery({ userQuery: input.query, llm })) + .addNode('validateQuery', getValidateQuery({ search })) + .addEdge(START, 'classifyEsql') + .addEdge('classifyEsql', 'generateQuery') + .addEdge('generateQuery', 'validateQuery') + .addConditionalEdges('validateQuery', shouldRegenerate); + + const app = workflow.compile(); + + let query; + + try { + query = await app.invoke({}, { recursionLimit: 20 }); + } catch (e) { + return 'error'; + } + + return query.esqlQuery; + } }, tags: ['esql', 'query-generation', 'knowledge-base'], }); From c4a58045ae0cbe25a3e50ca9848197c8d4622756 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 26 Jun 2024 22:21:33 +0200 Subject: [PATCH 24/47] fix --- x-pack/plugins/elastic_assistant/server/types.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 0635286484673..baf5e93c92b00 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -33,11 +33,7 @@ import { } from '@kbn/elastic-assistant-common'; import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; -import { - ActionsClientChatOpenAI, - ActionsClientLlm, - ActionsClientSimpleChatModel, -} from '@kbn/langchain/server'; +import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; import { AttackDiscoveryDataClient } from './ai_assistant_data_clients/attack_discovery'; From af26d67cd99337ae24a2879509b261443ba4270e Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 26 Jun 2024 22:52:26 +0200 Subject: [PATCH 25/47] fix --- .../plugins/custom_codeblock/esql_code_block.tsx | 4 ++-- x-pack/plugins/elastic_assistant/server/types.ts | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx index 032fdc0d1ffe9..bf0124fbab925 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx @@ -193,13 +193,13 @@ const EsqlCodeBlockComponent: React.FC = ({ setShowVisualization(true)} disabled={actionsDisabled} > - {i18n.translate('xpack.securityAiAssistant.visualizeThisQuery', { + {i18n.translate('xpack.cases.lensESQLFunction.visualizeThisQuery', { defaultMessage: 'Generate Visualization', })} diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index baf5e93c92b00..e87c7f8da372b 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -33,7 +33,7 @@ import { } from '@kbn/elastic-assistant-common'; import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; -import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; +import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; import { AttackDiscoveryDataClient } from './ai_assistant_data_clients/attack_discovery'; From 2ff191f7667037d1ecd545dd157e604947a275bf Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 26 Jun 2024 23:20:57 +0200 Subject: [PATCH 26/47] fix --- .../plugins/elastic_assistant/server/types.ts | 5 +- .../esql_language_knowledge_base_tool.ts | 11 ++-- .../graph_esql_language_tool.ts | 55 +++++++++---------- .../knowledge_base_retrieval_tool.ts | 7 ++- .../knowledge_base_write_tool.ts | 12 +++- 5 files changed, 52 insertions(+), 38 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index e87c7f8da372b..6fd0a908602f7 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -33,9 +33,10 @@ import { } from '@kbn/elastic-assistant-common'; import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; -import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; +import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; +import { z } from 'zod'; import { AttackDiscoveryDataClient } from './ai_assistant_data_clients/attack_discovery'; import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/conversations'; import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; @@ -202,7 +203,7 @@ export interface AIAssistantPrompts { * Interfaces for registering tools to be used by the elastic assistant */ -export interface AssistantTool { +export interface AssistantTool { id: string; name: string; description: string; 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 692753a22dea0..6c6e5e5bf66e7 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 @@ -18,7 +18,12 @@ const toolDetails = { id: 'esql-knowledge-base-tool', name: 'ESQLKnowledgeBaseTool', }; -export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { + +const schema = z.object({ + question: z.string().describe(`The user's exact question about ESQL`), +}); + +export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is EsqlKnowledgeBaseToolParams => { @@ -34,9 +39,7 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { return new DynamicStructuredTool({ name: toolDetails.name, description: toolDetails.description, - schema: z.object({ - question: z.string().describe(`The user's exact question about ESQL`), - }), + schema, func: async (input, _, cbManager) => { const result = await chain.invoke( { diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index 1cf4089cf4c9c..b1a992599d8ca 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -369,51 +369,46 @@ const shouldRegenerate = (state: IState) => { }; const schema = z.object({ - query: z.string().describe(`The user's exact question about ESQL`), + question: z.string().describe(`The user's exact question about ESQL`), }); export const GRAPH_ESQL_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, - isSupported: (params: AssistantToolParams): params is GraphESQLToolParams => { - const { chain, isEnabledKnowledgeBase, modelExists } = params; - return isEnabledKnowledgeBase && modelExists && chain != null; - }, + isSupported: () => true, getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { chain, esClient, search, llm } = params as GraphESQLToolParams; - if (chain == null) return null; + const { esClient, search, llm } = params; + if (!llm) return null; return new DynamicStructuredTool({ name: toolDetails.name, description: toolDetails.description, schema, func: async (input, _, cbManager) => { - if (llm) { - const workflow = new StateGraph({ - channels: graphState, - }) - .addNode('classifyEsql', getClassifyEsql({ userQuery: input.query, llm, esClient })) - .addNode('generateQuery', getGenerateQuery({ userQuery: input.query, llm })) - .addNode('validateQuery', getValidateQuery({ search })) - .addEdge(START, 'classifyEsql') - .addEdge('classifyEsql', 'generateQuery') - .addEdge('generateQuery', 'validateQuery') - .addConditionalEdges('validateQuery', shouldRegenerate); - - const app = workflow.compile(); - - let query; - - try { - query = await app.invoke({}, { recursionLimit: 20 }); - } catch (e) { - return 'error'; - } - - return query.esqlQuery; + const workflow = new StateGraph({ + channels: graphState, + }) + .addNode('classifyEsql', getClassifyEsql({ userQuery: input.question, llm, esClient })) + .addNode('generateQuery', getGenerateQuery({ userQuery: input.question, llm })) + .addNode('validateQuery', getValidateQuery({ search })) + .addEdge(START, 'classifyEsql') + .addEdge('classifyEsql', 'generateQuery') + .addEdge('generateQuery', 'validateQuery') + .addConditionalEdges('validateQuery', shouldRegenerate); + + const app = workflow.compile(); + + let query; + + try { + query = await app.invoke({ question: input.question }, { recursionLimit: 20 }); + } catch (e) { + return 'error'; } + + return query.esqlQuery; }, tags: ['esql', 'query-generation', 'knowledge-base'], }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts index 47cb35e244d51..4414fad57489e 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts @@ -21,7 +21,12 @@ const toolDetails = { id: 'knowledge-base-retrieval-tool', name: 'KnowledgeBaseRetrievalTool', }; -export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { + +const schema = z.object({ + query: z.string().describe(`Summary of items/things to search for in the knowledge base`), +}); + +export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is KnowledgeBaseRetrievalToolParams => { diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts index addb2a5580dfc..1cf9fed3d802b 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts @@ -22,7 +22,17 @@ const toolDetails = { id: 'knowledge-base-write-tool', name: 'KnowledgeBaseWriteTool', }; -export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { + +const schema = z.object({ + query: z.string().describe(`Summary of items/things to save in the knowledge base`), + required: z + .boolean() + .describe( + `Whether or not the entry is required to always be included in conversations. Is only true if the user explicitly asks for it to be required or always included in conversations, otherwise this is always false.` + ), +}); + +export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is KnowledgeBaseWriteToolParams => { From d0b1ac7bc5343ea4af400c4371b8465f2f03d501 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 26 Jun 2024 23:46:02 +0200 Subject: [PATCH 27/47] fix --- .../get_comments/custom_codeblock/esql_code_block.tsx | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index d69f8b7491ef0..352d21a1e93c4 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -35,7 +35,7 @@ function generateId() { return uuidv4(); } -const saveVisualizationLabel = i18n.translate('xpack.securityAiAssistant.lensESQLFunction.save', { +const saveVisualizationLabel = i18n.translate('xpack.securitySolution.lensESQLFunction.save', { defaultMessage: 'Save visualization', }); @@ -191,7 +191,7 @@ const EsqlCodeBlockComponent = ({ onClick={handleShowVisualization} disabled={actionsDisabled} > - {i18n.translate('xpack.securityAiAssistant.visualizeThisQuery', { + {i18n.translate('xpack.securitySolution.lensESQLFunction.visualizeThisQuery', { defaultMessage: 'Generate Visualization', })} @@ -229,7 +229,7 @@ const EsqlCodeBlockComponent = ({ isTableVisible ? 'visBarVerticalStacked' : 'tableDensityExpanded' } onClick={() => setIsTableVisible(!isTableVisible)} - data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton" + data-test-subj="securityAiAssistantLensESQLDisplayTableButton" aria-label={ isTableVisible ? i18n.translate( @@ -254,7 +254,7 @@ const EsqlCodeBlockComponent = ({ size="xs" iconType="save" onClick={() => setIsSaveModalOpen(true)} - data-test-subj="observabilityAiAssistantLensESQLSaveButton" + data-test-subj="securityAiAssistantLensESQLSaveButton" aria-label={saveVisualizationLabel} /> From c075b4abd3118a52f78e9685a28c12f033694066 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 27 Jun 2024 00:08:04 +0200 Subject: [PATCH 28/47] fix --- yarn.lock | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/yarn.lock b/yarn.lock index d42031adb97f7..caf6d70922e46 100644 --- a/yarn.lock +++ b/yarn.lock @@ -21745,7 +21745,7 @@ langchainhub@~0.0.8: resolved "https://registry.yarnpkg.com/langchainhub/-/langchainhub-0.0.8.tgz#fd4b96dc795e22e36c1a20bad31b61b0c33d3110" integrity sha512-Woyb8YDHgqqTOZvWIbm2CaFDGfZ4NTSyXV687AG4vXEfoNo7cGQp7nhl7wL3ehenKWmNEmcxCLgOZzW8jE6lOQ== -langsmith@^0.1.30: +langsmith@^0.1.30, langsmith@~0.1.30: version "0.1.32" resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.32.tgz#38938b0e8685522087b697b8200c488c6490c137" integrity sha512-EUWHIH6fiOCGRYdzgwGoXwJxCMyUrL+bmUcxoVmkXoXoAGDOVinz8bqJLKbxotsQWqM64NKKsW85OTIutgNaMQ== @@ -21756,17 +21756,6 @@ langsmith@^0.1.30: p-retry "4" uuid "^9.0.0" -langsmith@~0.1.30: - version "0.1.30" - resolved "https://registry.yarnpkg.com/langsmith/-/langsmith-0.1.30.tgz#3000e441605b26e15a87fb991a3929c944edbc0a" - integrity sha512-g8f10H1iiRjCweXJjgM3Y9xl6ApCa1OThDvc0BlSDLVrGVPy1on9wT39vAzYkeadC7oG48p7gfpGlYH3kLkJ9Q== - dependencies: - "@types/uuid" "^9.0.1" - commander "^10.0.1" - p-queue "^6.6.2" - p-retry "4" - uuid "^9.0.0" - language-subtag-registry@~0.3.2: version "0.3.21" resolved "https://registry.yarnpkg.com/language-subtag-registry/-/language-subtag-registry-0.3.21.tgz#04ac218bea46f04cb039084602c6da9e788dd45a" From 1e65fbf63d4115145c8fd37d68071c96c2477d51 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 27 Jun 2024 01:58:18 +0200 Subject: [PATCH 29/47] cleanup --- .../use_conversation/helpers.test.ts | 4 +- x-pack/plugins/cases/kibana.jsonc | 2 - .../public/components/__mock__/timeline.tsx | 8 +- .../custom_codeblock/custom_code_block.tsx | 45 --- .../custom_codeblock_markdown_plugin.tsx | 54 --- .../custom_codeblock/esql_code_block.tsx | 326 ------------------ .../components/markdown_editor/types.ts | 8 +- .../components/markdown_editor/use_plugins.ts | 20 +- .../components/timeline_context/index.tsx | 20 +- .../user_actions/markdown_form.test.tsx | 1 + x-pack/plugins/cases/public/types.ts | 2 - x-pack/plugins/cases/tsconfig.json | 6 +- .../custom_codeblock/custom_code_block.tsx | 5 +- .../custom_codeblock/esql_code_block.tsx | 10 +- .../get_comments/stream/message_text.tsx | 33 +- .../public/cases/pages/index.tsx | 10 +- .../correct_common_esql_mistakes.ts | 2 + .../test/security_solution_cypress/config.ts | 3 +- 18 files changed, 79 insertions(+), 480 deletions(-) delete mode 100644 x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_code_block.tsx delete mode 100644 x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_codeblock_markdown_plugin.tsx delete mode 100644 x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.test.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.test.ts index f348049eec8b6..86e4c18407d46 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.test.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.test.ts @@ -56,11 +56,11 @@ This query will filter the events based on the condition that the ${tilde}user.n describe('useConversation helpers', () => { describe('analyzeMarkdown', () => { it('should identify dsl Query successfully.', () => { - const result = analyzeMarkdown(markDownWithDSLQuery); + const result = analyzeMarkdown(markDownWithDSLQuery, new Date().toISOString()); expect(result[0].type).toBe('dsl'); }); it('should identify kql Query successfully.', () => { - const result = analyzeMarkdown(markDownWithKQLQuery); + const result = analyzeMarkdown(markDownWithKQLQuery, new Date().toISOString()); expect(result[0].type).toBe('kql'); }); }); diff --git a/x-pack/plugins/cases/kibana.jsonc b/x-pack/plugins/cases/kibana.jsonc index e4cee6a6d9c9c..84c04da1fe0f6 100644 --- a/x-pack/plugins/cases/kibana.jsonc +++ b/x-pack/plugins/cases/kibana.jsonc @@ -15,9 +15,7 @@ "alerting", "actions", "data", - "dataViews", "embeddable", - "esqlDataGrid", "lens", "licensing", "features", diff --git a/x-pack/plugins/cases/public/components/__mock__/timeline.tsx b/x-pack/plugins/cases/public/components/__mock__/timeline.tsx index d576b0ef1732c..778b82a221887 100644 --- a/x-pack/plugins/cases/public/components/__mock__/timeline.tsx +++ b/x-pack/plugins/cases/public/components/__mock__/timeline.tsx @@ -12,8 +12,12 @@ const mockTimelineComponent = (name: string) => {nam export const timelineIntegrationMock = { editor_plugins: { - parsingPlugin: jest.fn(), - processingPluginRenderer: () => mockTimelineComponent('plugin-renderer'), + parsingPlugins: [], + processingPluginRenderer: { + timeline: () => mockTimelineComponent('plugin-renderer'), + esql: () => mockTimelineComponent('esql-plugin-renderer'), + customCodeBlock: () => mockTimelineComponent('custom-code-block-plugin-renderer'), + }, uiPlugin: { name: 'mock-timeline', button: { label: 'mock-timeline-button', iconType: 'mock-timeline-icon' }, diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_code_block.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_code_block.tsx deleted file mode 100644 index 9b009c3f32240..0000000000000 --- a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_code_block.tsx +++ /dev/null @@ -1,45 +0,0 @@ -/* - * 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 { EuiCodeBlock, EuiFlexGroup, EuiFlexItem, EuiPanel, useEuiTheme } from '@elastic/eui'; -import { css } from '@emotion/css'; -import React from 'react'; - -const CustomCodeBlockComponent = ({ value }: { value: string }) => { - const theme = useEuiTheme(); - - return ( - - - - - {value} - - - - - ); -}; - -CustomCodeBlockComponent.displayName = 'CustomCodeBlock'; - -export const CustomCodeBlock = React.memo(CustomCodeBlockComponent); diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_codeblock_markdown_plugin.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_codeblock_markdown_plugin.tsx deleted file mode 100644 index 8825974f9c9c6..0000000000000 --- a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/custom_codeblock_markdown_plugin.tsx +++ /dev/null @@ -1,54 +0,0 @@ -/* - * 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 type { Node } from 'unist'; -import type { Parent, Content, PhrasingContent } from 'mdast'; -import { MARKDOWN_TYPES } from '@kbn/elastic-assistant/impl/assistant/use_conversation/helpers'; - -export const customCodeBlockLanguagePlugin = () => { - const visitor = (node: Node, parent?: Parent) => { - if ('children' in node) { - const nodeAsParent = node as Parent; - nodeAsParent.children.forEach((child) => { - visitor(child, nodeAsParent); - }); - } - - if (node.type === 'code' && !node.lang) { - try { - const index = parent?.children.indexOf(node as Content); - - if (index) { - const previousText = (parent?.children[index - 1]?.children as PhrasingContent[]) - ?.map((child) => child.value) - .join(' '); - for (const [typeKey, keywords] of Object.entries(MARKDOWN_TYPES)) { - if (keywords.some((kw) => previousText.toLowerCase().includes(kw.toLowerCase()))) { - node.lang = typeKey; - break; - } - } - } - } catch (e) { - /* empty */ - } - } - - if (node.type === 'code' && node.lang === 'esql') { - node.type = 'esql'; - return; - } - - if (node.type === 'code' && ['eql', 'kql', 'dsl', 'json'].includes(node.lang as string)) { - node.type = 'customCodeBlock'; - } - }; - - return (tree: Node) => { - visitor(tree); - }; -}; diff --git a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx deleted file mode 100644 index bf0124fbab925..0000000000000 --- a/x-pack/plugins/cases/public/components/markdown_editor/plugins/custom_codeblock/esql_code_block.tsx +++ /dev/null @@ -1,326 +0,0 @@ -/* - * 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 { - EuiButtonEmpty, - EuiButtonIcon, - EuiCodeBlock, - EuiFlexGroup, - EuiFlexItem, - EuiPanel, - EuiSpacer, - EuiToolTip, - useEuiTheme, -} from '@elastic/eui'; -import { css } from '@emotion/css'; -import { i18n } from '@kbn/i18n'; -import React, { useEffect, useState } from 'react'; -import { useQuery } from '@tanstack/react-query'; -import { v4 as uuidv4 } from 'uuid'; -import { - getIndexPatternFromESQLQuery, - getESQLQueryColumns, - getESQLAdHocDataview, - getESQLResults, -} from '@kbn/esql-utils'; -import useAsync from 'react-use/lib/useAsync'; -import { ESQLDataGrid } from '@kbn/esql-datagrid/public'; -import type { TypedLensByValueInput } from '@kbn/lens-plugin/public'; -import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils'; -import { useKibana } from '../../../../common/lib/kibana/kibana_react'; - -function generateId() { - return uuidv4(); -} - -const saveVisualizationLabel = i18n.translate('xpack.cases.lensESQLFunction.save', { - defaultMessage: 'Save visualization', -}); - -interface EsqlCodeBlockProps { - value: string; - timestamp: string; - actionsDisabled: boolean; -} - -const EsqlCodeBlockComponent: React.FC = ({ - value, - actionsDisabled = false, - timestamp, -}) => { - const { lens, dataViews: dataViewService, data } = useKibana().services; - const theme = useEuiTheme(); - - const lensHelpersAsync = useAsync(() => { - return lens.stateHelperApi(); - }, [lens]); - - const { data: queryResults } = useQuery({ - queryKey: ['test'], - enabled: true, - queryFn: async () => { - return getESQLResults({ - esqlQuery: value, - search: data.search.search, - filter: { - range: { - '@timestamp': { - lte: timestamp, - format: 'strict_date_optional_time', - }, - }, - }, - }); - }, - select: (dataz) => { - return { - params: dataz.params, - rows: dataz.response.values, - columns: dataz.response.columns, - }; - }, - refetchOnWindowFocus: false, - keepPreviousData: true, - }); - - const indexPattern = getIndexPatternFromESQLQuery(value); - const formattedColumns = useAsync( - () => - getESQLQueryColumns({ - esqlQuery: value, - search: data.search.search, - }), - [value] - ); - - const dataViewAsync = useAsync(() => { - return getESQLAdHocDataview(indexPattern, dataViewService).then((dataView) => { - if (dataView.fields.getByName('@timestamp')?.type === 'date') { - dataView.timeFieldName = '@timestamp'; - } - return dataView; - }); - }, [indexPattern, dataViewService]); - - const [isSaveModalOpen, setIsSaveModalOpen] = useState(false); - const [isTableVisible, setIsTableVisible] = useState(false); - const [lensInput, setLensInput] = useState(undefined); - const [showVisualization, setShowVisualization] = useState(false); - - const preferredChartType = undefined; // 'XY'; - - // initialization - useEffect(() => { - if (lensHelpersAsync.value && dataViewAsync.value && !lensInput && formattedColumns.value) { - const context = { - dataViewSpec: dataViewAsync.value?.toSpec(), - fieldName: '', - textBasedColumns: formattedColumns.value, - query: { - esql: value, - }, - }; - - const chartSuggestions = lensHelpersAsync.value.suggestions( - context, - dataViewAsync.value, - [], - preferredChartType - ); - - if (chartSuggestions?.length) { - const [suggestion] = chartSuggestions; - - const attrs = getLensAttributesFromSuggestion({ - filters: [], - query: { - esql: value, - }, - suggestion, - dataView: dataViewAsync.value, - }) as TypedLensByValueInput['attributes']; - - const lensEmbeddableInput = { - attributes: attrs, - id: generateId(), - }; - setLensInput(lensEmbeddableInput); - } - } - }, [ - dataViewAsync.value, - formattedColumns.value, - lensHelpersAsync.value, - lensInput, - preferredChartType, - value, - ]); - - // if the Lens suggestions api suggests a table then we want to render a Discover table instead - const isLensInputTable = lensInput?.attributes?.visualizationType === 'lnsDatatable'; - - return ( - <> - - - - - {value} - - - - {!showVisualization && ( - - - - setShowVisualization(true)} - disabled={actionsDisabled} - > - {i18n.translate('xpack.cases.lensESQLFunction.visualizeThisQuery', { - defaultMessage: 'Generate Visualization', - })} - - - - - )} - - - - {showVisualization && queryResults && formattedColumns.value && ( - <> - {!isLensInputTable && ( - <> - - - - - setIsTableVisible(!isTableVisible)} - data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton" - aria-label={ - isTableVisible - ? i18n.translate('xpack.cases.lensESQLFunction.displayChart', { - defaultMessage: 'Display chart', - }) - : i18n.translate('xpack.cases.lensESQLFunction.displayTable', { - defaultMessage: 'Display table', - }) - } - /> - - - - - setIsSaveModalOpen(true)} - data-test-subj="observabilityAiAssistantLensESQLSaveButton" - aria-label={saveVisualizationLabel} - /> - - - - - - {isTableVisible ? ( - - ) : lensInput ? ( - - ) : null} - - - )} - {/* hide the grid in case of errors (as the user can't fix them) */} - {isLensInputTable && ( - - - - )} - - )} - - {isSaveModalOpen ? ( - { - setIsSaveModalOpen(() => false); - }} - // For now, we don't want to allow saving ESQL charts to the library - isSaveable={false} - /> - ) : null} - - ); -}; - -EsqlCodeBlockComponent.displayName = 'EsqlCodeBlock'; - -export const EsqlCodeBlock = React.memo(EsqlCodeBlockComponent); - -// eslint-disable-next-line react/display-name -export const getEsqlRenderer = (timestamp: string | undefined) => (props) => - ( - <> - - - - ); diff --git a/x-pack/plugins/cases/public/components/markdown_editor/types.ts b/x-pack/plugins/cases/public/components/markdown_editor/types.ts index 82ff738a178bf..846bded0ed9f8 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/types.ts +++ b/x-pack/plugins/cases/public/components/markdown_editor/types.ts @@ -21,7 +21,13 @@ export type TemporaryProcessingPluginsType = [ [ typeof rehype2react, Parameters[0] & { - components: { a: FunctionComponent; lens: unknown; timeline: unknown }; + components: { + a: FunctionComponent; + lens: unknown; + timeline: unknown; + customCodeBlock: unknown; + esql: unknown; + }; } ], ...PluggableList diff --git a/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts b/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts index 5b186d86302cc..4685df90f5480 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts +++ b/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts @@ -16,8 +16,6 @@ import type { TemporaryProcessingPluginsType } from './types'; import { KibanaServices, useApplicationCapabilities } from '../../common/lib/kibana'; import * as lensMarkdownPlugin from './plugins/lens'; import { ID as LensPluginId } from './plugins/lens/constants'; -import { getEsqlRenderer } from './plugins/custom_codeblock/esql_code_block'; -import { customCodeBlockLanguagePlugin } from './plugins/custom_codeblock/custom_codeblock_markdown_plugin'; export const usePlugins = ({ disabledPlugins, @@ -32,24 +30,26 @@ export const usePlugins = ({ return useMemo(() => { const uiPlugins = getDefaultEuiMarkdownUiPlugins(); - const parsingPlugins = [ - customCodeBlockLanguagePlugin, - ...getDefaultEuiMarkdownParsingPlugins(), - ]; + const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); const processingPlugins = getDefaultEuiMarkdownProcessingPlugins() as TemporaryProcessingPluginsType; if (timelinePlugins) { uiPlugins.push(timelinePlugins.uiPlugin); - parsingPlugins.push(timelinePlugins.parsingPlugin); + parsingPlugins.concat(timelinePlugins.parsingPlugins); // This line of code is TS-compatible and it will break if [1][1] change in the future. - processingPlugins[1][1].components.timeline = timelinePlugins.processingPluginRenderer; + processingPlugins[1][1].components.timeline = + timelinePlugins.processingPluginRenderer.timeline; + processingPlugins[1][1].components.esql = (props: { + value: string; + actionsDisabled: boolean; + }) => timelinePlugins.processingPluginRenderer.esql({ timestamp: timestamp ?? '', ...props }); + processingPlugins[1][1].components.customCodeBlock = + timelinePlugins.processingPluginRenderer.customCodeBlock; } - processingPlugins[1][1].components.esql = getEsqlRenderer(timestamp); - if ( kibanaConfig?.markdownPlugins?.lens && !disabledPlugins?.includes(LensPluginId) && diff --git a/x-pack/plugins/cases/public/components/timeline_context/index.tsx b/x-pack/plugins/cases/public/components/timeline_context/index.tsx index 3b5e4949150c0..0a1f680cf1135 100644 --- a/x-pack/plugins/cases/public/components/timeline_context/index.tsx +++ b/x-pack/plugins/cases/public/components/timeline_context/index.tsx @@ -30,10 +30,22 @@ interface TimelineProcessingPluginRendererProps { export interface CasesTimelineIntegration { editor_plugins: { - parsingPlugin: Plugin; - processingPluginRenderer: React.FC< - TimelineProcessingPluginRendererProps & { position: EuiMarkdownAstNodePosition } - >; + parsingPlugins: Plugin[]; + processingPluginRenderer: { + timeline: React.FC< + TimelineProcessingPluginRendererProps & { position: EuiMarkdownAstNodePosition } + >; + esql: ({ + timestamp, + value, + actionsDisabled, + }: { + value: string; + timestamp: string; + actionsDisabled: boolean; + }) => JSX.Element; + customCodeBlock: ({ value }: { value: string }) => JSX.Element; + }; uiPlugin: EuiMarkdownEditorUiPlugin; }; hooks: { diff --git a/x-pack/plugins/cases/public/components/user_actions/markdown_form.test.tsx b/x-pack/plugins/cases/public/components/user_actions/markdown_form.test.tsx index 87bae9a4624a0..d8758363c29dd 100644 --- a/x-pack/plugins/cases/public/components/user_actions/markdown_form.test.tsx +++ b/x-pack/plugins/cases/public/components/user_actions/markdown_form.test.tsx @@ -30,6 +30,7 @@ const defaultProps = { draftStorageKey, onChangeEditable, onSaveContent, + timestamp: new Date().toISOString(), }; describe('UserActionMarkdown ', () => { diff --git a/x-pack/plugins/cases/public/types.ts b/x-pack/plugins/cases/public/types.ts index 888ee3df42104..c857446ea042c 100644 --- a/x-pack/plugins/cases/public/types.ts +++ b/x-pack/plugins/cases/public/types.ts @@ -30,7 +30,6 @@ import type { ContentManagementPublicStart } from '@kbn/content-management-plugi import type { UiActionsStart } from '@kbn/ui-actions-plugin/public'; import type { ServerlessPluginSetup, ServerlessPluginStart } from '@kbn/serverless/public'; -import type { DataViewsPublicPluginStart } from '@kbn/data-views-plugin/public'; import type { UseCasesAddToExistingCaseModal } from './components/all_cases/selector_modal/use_cases_add_to_existing_case_modal'; import type { UseCasesAddToNewCaseFlyout } from './components/create/flyout/use_cases_add_to_new_case_flyout'; import type { UseIsAddToCaseOpen } from './components/cases_context/state/use_is_add_to_case_open'; @@ -75,7 +74,6 @@ export interface CasesPublicSetupDependencies { export interface CasesPublicStartDependencies { apm?: ApmBase; data: DataPublicPluginStart; - dataViews: DataViewsPublicPluginStart; embeddable: EmbeddableStart; features: FeaturesPluginStart; files: FilesStart; diff --git a/x-pack/plugins/cases/tsconfig.json b/x-pack/plugins/cases/tsconfig.json index f1a1f4fef9194..b4a116484ca30 100644 --- a/x-pack/plugins/cases/tsconfig.json +++ b/x-pack/plugins/cases/tsconfig.json @@ -74,11 +74,7 @@ "@kbn/core-logging-server-mocks", "@kbn/core-logging-browser-mocks", "@kbn/data-views-plugin", - "@kbn/core-http-router-server-internal", - "@kbn/elastic-assistant", - "@kbn/esql-utils", - "@kbn/esql-datagrid", - "@kbn/visualization-utils", + "@kbn/core-http-router-server-internal" ], "exclude": [ "target/**/*", diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx index 90e406f500ba8..db22c84825f05 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx @@ -8,7 +8,10 @@ import { EuiCodeBlock, EuiFlexGroup, EuiFlexItem, EuiPanel, useEuiTheme } from ' import { css } from '@emotion/css'; import React from 'react'; -export const CustomCodeBlock = ({ value }: { value: string }) => { +export interface CustomCodeBlockProps { + value: string; +} +export const CustomCodeBlock = ({ value }: CustomCodeBlockProps) => { const theme = useEuiTheme(); return ( diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index 352d21a1e93c4..c56d054877710 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -39,15 +39,13 @@ const saveVisualizationLabel = i18n.translate('xpack.securitySolution.lensESQLFu defaultMessage: 'Save visualization', }); -const EsqlCodeBlockComponent = ({ - value, - actionsDisabled, - timestamp, -}: { +export interface EsqlCodeBlockProps { value: string; timestamp: string; actionsDisabled: boolean; -}) => { +} + +const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeBlockProps) => { const theme = useEuiTheme(); const { lens, dataViews: dataViewService, data } = useKibana().services; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx index 3c88967a6db94..1420e3598f7ba 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx @@ -24,7 +24,9 @@ import { euiThemeVars } from '@kbn/ui-theme'; import type { Node } from 'unist'; import { customCodeBlockLanguagePlugin } from '../custom_codeblock/custom_codeblock_markdown_plugin'; +import type { CustomCodeBlockProps } from '../custom_codeblock/custom_code_block'; import { CustomCodeBlock } from '../custom_codeblock/custom_code_block'; +import type { EsqlCodeBlockProps } from '../custom_codeblock/esql_code_block'; import { EsqlCodeBlock } from '../custom_codeblock/esql_code_block'; interface Props { @@ -100,13 +102,19 @@ const loadingCursorPlugin = () => { }; }; -const getEsql = (timestamp?: string) => (props) => - ( - <> - - - - ); +export const Esql = ({ timestamp, ...props }: EsqlCodeBlockProps) => ( + <> + + + +); + +export const CodeBlock = (props: CustomCodeBlockProps) => ( + <> + + + +); const getPluginDependencies = (timestamp?: string) => { const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); @@ -118,15 +126,8 @@ const getPluginDependencies = (timestamp?: string) => { processingPlugins[1][1].components = { ...components, cursor: Cursor, - esql: getEsql(timestamp), - customCodeBlock: (props) => { - return ( - <> - - - - ); - }, + esql: Esql, + customCodeBlock: CodeBlock, table: (props) => ( <> diff --git a/x-pack/plugins/security_solution/public/cases/pages/index.tsx b/x-pack/plugins/security_solution/public/cases/pages/index.tsx index 54ece0123738f..d963a1c7c5d02 100644 --- a/x-pack/plugins/security_solution/public/cases/pages/index.tsx +++ b/x-pack/plugins/security_solution/public/cases/pages/index.tsx @@ -35,6 +35,8 @@ import * as timelineMarkdownPlugin from '../../common/components/markdown_editor import { DetailsPanel } from '../../timelines/components/side_panel'; import { useFetchAlertData } from './use_fetch_alert_data'; import { useUpsellingMessage } from '../../common/hooks/use_upselling'; +import { Esql, CodeBlock } from '../../assistant/get_comments/stream/message_text'; +import { customCodeBlockLanguagePlugin } from '../../assistant/get_comments/custom_codeblock/custom_codeblock_markdown_plugin'; const TimelineDetailsPanel = () => { const { browserFields, runtimeMappings } = useSourcererDataView(SourcererScopeName.detections); @@ -183,8 +185,12 @@ const CaseContainerComponent: React.FC = () => { showAlertDetails, timelineIntegration: { editor_plugins: { - parsingPlugin: timelineMarkdownPlugin.parser, - processingPluginRenderer: timelineMarkdownPlugin.renderer, + parsingPlugins: [timelineMarkdownPlugin.parser, customCodeBlockLanguagePlugin], + processingPluginRenderer: { + timeline: timelineMarkdownPlugin.renderer, + esql: Esql, + customCodeBlock: CodeBlock, + }, uiPlugin: timelineMarkdownPlugin.plugin({ interactionsUpsellingMessage }), }, hooks: { diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts index 09b6a59cac357..d0ecf0671a234 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts @@ -36,6 +36,7 @@ function split(value: string, splitToken: string) { index += splitToken.length - 1; statements.push(currentStatement.trim()); currentStatement = ''; + // eslint-disable-next-line no-continue continue; } @@ -50,6 +51,7 @@ function split(value: string, splitToken: string) { if (applicableToken) { // start identifier delimiterToken = applicableToken; + // eslint-disable-next-line no-continue continue; } else if (char === '(') { groupingCount++; diff --git a/x-pack/test/security_solution_cypress/config.ts b/x-pack/test/security_solution_cypress/config.ts index 0fce26844c1bf..6e65ab15324a6 100644 --- a/x-pack/test/security_solution_cypress/config.ts +++ b/x-pack/test/security_solution_cypress/config.ts @@ -47,11 +47,10 @@ export default async function ({ readConfigFile }: FtrConfigProviderContext) { `--xpack.securitySolution.enableExperimental=${JSON.stringify([ 'alertSuppressionForEsqlRuleEnabled', 'bulkCustomHighlightedFieldsEnabled', - 'assistantKnowledgeBaseByDefault', 'manualRuleRunEnabled', ])}`, // mock cloud to enable the guided onboarding tour in e2e tests - // '--xpack.cloud.id=test', + '--xpack.cloud.id=test', `--home.disableWelcomeScreen=true`, // Specify which version of the detection-rules package to install // `--xpack.securitySolution.prebuiltRulesPackageVersion=8.3.1`, From 3ae0385ea11a65d0039dc9ec85495cb7af393c4c Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 27 Jun 2024 02:01:37 +0200 Subject: [PATCH 30/47] fix --- x-pack/plugins/cases/tsconfig.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/cases/tsconfig.json b/x-pack/plugins/cases/tsconfig.json index b4a116484ca30..535f4e5e106dc 100644 --- a/x-pack/plugins/cases/tsconfig.json +++ b/x-pack/plugins/cases/tsconfig.json @@ -74,7 +74,7 @@ "@kbn/core-logging-server-mocks", "@kbn/core-logging-browser-mocks", "@kbn/data-views-plugin", - "@kbn/core-http-router-server-internal" + "@kbn/core-http-router-server-internal", ], "exclude": [ "target/**/*", From 77073b26004619ff27e900bdedcd0cdd6b8a1cdd Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 27 Jun 2024 08:23:51 +0200 Subject: [PATCH 31/47] fix --- .../impl/capabilities/index.ts | 2 +- .../components/markdown_editor/use_plugins.ts | 3 +- .../components/timeline_context/index.tsx | 5 ++- .../common/experimental_features.ts | 2 +- .../custom_codeblock/custom_code_block.tsx | 2 +- .../custom_codeblock/esql_code_block.tsx | 42 +++++++++++-------- .../get_comments/stream/message_text.tsx | 2 +- .../public/cases/pages/index.tsx | 5 ++- .../security_solution/server/plugin.ts | 2 +- 9 files changed, 40 insertions(+), 25 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index c1c101fd74cd8..54c24f6ce7b8f 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, assistantModelEvaluation: false, }); diff --git a/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts b/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts index 4685df90f5480..e2a74b3ad13ef 100644 --- a/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts +++ b/x-pack/plugins/cases/public/components/markdown_editor/use_plugins.ts @@ -37,7 +37,8 @@ export const usePlugins = ({ if (timelinePlugins) { uiPlugins.push(timelinePlugins.uiPlugin); - parsingPlugins.concat(timelinePlugins.parsingPlugins); + parsingPlugins.unshift(timelinePlugins.parsingPlugins.customCodeBlock); + parsingPlugins.push(timelinePlugins.parsingPlugins.timeline); // This line of code is TS-compatible and it will break if [1][1] change in the future. processingPlugins[1][1].components.timeline = diff --git a/x-pack/plugins/cases/public/components/timeline_context/index.tsx b/x-pack/plugins/cases/public/components/timeline_context/index.tsx index 0a1f680cf1135..5a48d6f879824 100644 --- a/x-pack/plugins/cases/public/components/timeline_context/index.tsx +++ b/x-pack/plugins/cases/public/components/timeline_context/index.tsx @@ -30,7 +30,10 @@ interface TimelineProcessingPluginRendererProps { export interface CasesTimelineIntegration { editor_plugins: { - parsingPlugins: Plugin[]; + parsingPlugins: { + timeline: Plugin; + customCodeBlock: Plugin; + }; processingPluginRenderer: { timeline: React.FC< TimelineProcessingPluginRendererProps & { position: EuiMarkdownAstNodePosition } diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 13c477d285d3c..444c8409ba30e 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -127,7 +127,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables the Assistant Knowledge Base by default, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, /** * Enables the Managed User section inside the new user details flyout. diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx index db22c84825f05..eac01423340f0 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx @@ -34,7 +34,7 @@ export const CustomCodeBlock = ({ value }: CustomCodeBlockProps) => { > - + {value} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index c56d054877710..f03e9111ef2ca 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -7,6 +7,7 @@ import { EuiButtonEmpty, EuiButtonIcon, + EuiCallOut, EuiCodeBlock, EuiFlexGroup, EuiFlexItem, @@ -53,7 +54,7 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB return lens.stateHelperApi(); }, [lens]); - const { data: queryResults } = useQuery({ + const { data: queryResults, error: queryResultsError } = useQuery({ queryKey: ['getESQLResults', value, timestamp], enabled: true, queryFn: async () => { @@ -61,28 +62,28 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB esqlQuery: value, search: data.search.search, filter: { - range: { - '@timestamp': { - lte: timestamp, - format: 'strict_date_optional_time', - }, - }, + range: timestamp + ? { + '@timestamp': { + lte: timestamp, + format: 'strict_date_optional_time', + }, + } + : {}, }, }); }, - select: (dataz) => { - return { - params: dataz.params, - rows: dataz.response.values, - columns: dataz.response.columns, - }; - }, + select: (queryResultsData) => ({ + params: queryResultsData.params, + rows: queryResultsData.response.values, + columns: queryResultsData.response.columns, + }), refetchOnWindowFocus: false, keepPreviousData: true, }); const indexPattern = useMemo(() => getIndexPatternFromESQLQuery(value), [value]); - const { value: formattedColumns } = useAsync( + const { value: formattedColumns, error: formattedColumnsError } = useAsync( () => getESQLQueryColumns({ esqlQuery: value, @@ -91,7 +92,7 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB [value] ); - const { value: dataViewAsync } = useAsync(() => { + const { value: dataViewAsync, error: dataViewError } = useAsync(() => { return getESQLAdHocDataview(indexPattern, dataViewService).then((dataView) => { if (dataView.fields.getByName('@timestamp')?.type === 'date') { dataView.timeFieldName = '@timestamp'; @@ -173,7 +174,7 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB > - + {value} @@ -199,6 +200,13 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB )} + {showVisualization && (queryResultsError || formattedColumnsError || dataViewError) && ( + + +

{`${queryResultsError || formattedColumnsError || dataViewError}`}

+ + + )} {showVisualization && queryResults && formattedColumns && ( diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx index 1420e3598f7ba..2e93a32cca9e5 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx @@ -126,7 +126,7 @@ const getPluginDependencies = (timestamp?: string) => { processingPlugins[1][1].components = { ...components, cursor: Cursor, - esql: Esql, + esql: (props) => , customCodeBlock: CodeBlock, table: (props) => ( <> diff --git a/x-pack/plugins/security_solution/public/cases/pages/index.tsx b/x-pack/plugins/security_solution/public/cases/pages/index.tsx index d963a1c7c5d02..91b260e8940d9 100644 --- a/x-pack/plugins/security_solution/public/cases/pages/index.tsx +++ b/x-pack/plugins/security_solution/public/cases/pages/index.tsx @@ -185,7 +185,10 @@ const CaseContainerComponent: React.FC = () => { showAlertDetails, timelineIntegration: { editor_plugins: { - parsingPlugins: [timelineMarkdownPlugin.parser, customCodeBlockLanguagePlugin], + parsingPlugins: { + timeline: timelineMarkdownPlugin.parser, + customCodeBlock: customCodeBlockLanguagePlugin, + }, processingPluginRenderer: { timeline: timelineMarkdownPlugin.renderer, esql: Esql, diff --git a/x-pack/plugins/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index 3720954163085..7da418e4c5869 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -567,7 +567,7 @@ export class Plugin implements ISecuritySolutionPlugin { // Assistant Tool and Feature Registration plugins.elasticAssistant.registerTools( APP_UI_ID, - getAssistantTools(config.experimentalFeatures.graphEsqlTool) + getAssistantTools(config.experimentalFeatures.aiAssistantGraphEsqlTool) ); plugins.elasticAssistant.registerFeatures(APP_UI_ID, { assistantKnowledgeBaseByDefault: config.experimentalFeatures.assistantKnowledgeBaseByDefault, From 8f1e3628dced7a3a7b5d4b7be2a6c08fd91dab59 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 27 Jun 2024 08:40:14 +0200 Subject: [PATCH 32/47] fix --- x-pack/plugins/elastic_assistant/server/types.ts | 4 ++-- .../custom_codeblock/esql_code_block.tsx | 4 ++-- .../esql_language_knowledge_base_tool.test.ts | 3 +++ .../esql_language_knowledge_base_tool.ts | 11 +++++------ .../graph_esql_language_tool.ts | 11 +++++------ .../server/assistant/tools/index.ts | 4 ++++ .../knowledge_base_retrieval_tool.ts | 9 +++------ .../knowledge_base/knowledge_base_write_tool.ts | 14 +++----------- 8 files changed, 27 insertions(+), 33 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 6fd0a908602f7..d097c7927cf4a 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -203,13 +203,13 @@ export interface AIAssistantPrompts { * Interfaces for registering tools to be used by the elastic assistant */ -export interface AssistantTool { +export interface AssistantTool { id: string; name: string; description: string; sourceRegister: string; isSupported: (params: AssistantToolParams) => boolean; - getTool: (params: AssistantToolParams) => Tool | DynamicStructuredTool | null; + getTool: (params: AssistantToolParams) => Tool | DynamicStructuredTool | null; } export interface AssistantToolParams { diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index f03e9111ef2ca..62985f5eae137 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -49,6 +49,7 @@ export interface EsqlCodeBlockProps { const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeBlockProps) => { const theme = useEuiTheme(); const { lens, dataViews: dataViewService, data } = useKibana().services; + const [showVisualization, setShowVisualization] = useState(false); const { value: lensHelpersAsync } = useAsync(() => { return lens.stateHelperApi(); @@ -56,7 +57,7 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB const { data: queryResults, error: queryResultsError } = useQuery({ queryKey: ['getESQLResults', value, timestamp], - enabled: true, + enabled: showVisualization, queryFn: async () => { return getESQLResults({ esqlQuery: value, @@ -104,7 +105,6 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB const [isSaveModalOpen, setIsSaveModalOpen] = useState(false); const [isTableVisible, setIsTableVisible] = useState(false); const [lensInput, setLensInput] = useState(undefined); - const [showVisualization, setShowVisualization] = useState(false); const preferredChartType = undefined; // 'XY'; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts index 29b10e9fb0275..7907d834ce2f3 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts @@ -12,6 +12,7 @@ import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import type { KibanaRequest } from '@kbn/core-http-server'; import type { ExecuteConnectorRequestBody } from '@kbn/elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen'; import { loggerMock } from '@kbn/logging-mocks'; +import { createSearchRequestHandlerContext } from '@kbn/data-plugin/server/search/mocks'; describe('EsqlLanguageKnowledgeBaseTool', () => { const chain = {} as RetrievalQAChain; @@ -28,12 +29,14 @@ describe('EsqlLanguageKnowledgeBaseTool', () => { size: 20, }, } as unknown as KibanaRequest; + const search = createSearchRequestHandlerContext(); const logger = loggerMock.create(); const rest = { chain, esClient, logger, request, + search, }; describe('isSupported', () => { 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 6c6e5e5bf66e7..c1cce234309fa 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 @@ -9,6 +9,7 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from 'zod'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; import { APP_UI_ID } from '../../../../common'; +import type { LangchainZodAny } from '..'; export type EsqlKnowledgeBaseToolParams = AssistantToolParams; @@ -19,11 +20,7 @@ const toolDetails = { name: 'ESQLKnowledgeBaseTool', }; -const schema = z.object({ - question: z.string().describe(`The user's exact question about ESQL`), -}); - -export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { +export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is EsqlKnowledgeBaseToolParams => { @@ -39,7 +36,9 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { return new DynamicStructuredTool({ name: toolDetails.name, description: toolDetails.description, - schema, + schema: z.object({ + question: z.string().describe(`The user's exact question about ESQL`), + }) as unknown as LangchainZodAny, func: async (input, _, cbManager) => { const result = await chain.invoke( { diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index b1a992599d8ca..574ccfec20534 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -25,6 +25,7 @@ import { validateQuery } from '@kbn/esql-validation-autocomplete'; import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; import { APP_UI_ID } from '../../../../common'; import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; +import type { LangchainZodAny } from '..'; export const INLINE_ESQL_QUERY_REGEX = /```esql\s*(.*?)\s*```/gms; @@ -368,11 +369,7 @@ const shouldRegenerate = (state: IState) => { return END; }; -const schema = z.object({ - question: z.string().describe(`The user's exact question about ESQL`), -}); - -export const GRAPH_ESQL_TOOL: AssistantTool = { +export const GRAPH_ESQL_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: () => true, @@ -385,7 +382,9 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { return new DynamicStructuredTool({ name: toolDetails.name, description: toolDetails.description, - schema, + schema: z.object({ + question: z.string().describe(`The user's exact question about ESQL`), + }) as unknown as LangchainZodAny, func: async (input, _, cbManager) => { const workflow = new StateGraph({ channels: graphState, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index f4e9e4f31ff72..b843999f6eaf1 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -7,6 +7,7 @@ import type { AssistantTool } from '@kbn/elastic-assistant-plugin/server'; +import type { z } from 'zod'; import { ALERT_COUNTS_TOOL } from './alert_counts/alert_counts_tool'; import { GRAPH_ESQL_TOOL } from './esql_language_knowledge_base/graph_esql_language_tool'; import { ESQL_KNOWLEDGE_BASE_TOOL } from './esql_language_knowledge_base/esql_language_knowledge_base_tool'; @@ -15,6 +16,9 @@ import { ATTACK_DISCOVERY_TOOL } from './attack_discovery/attack_discovery_tool' import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from './knowledge_base/knowledge_base_retrieval_tool'; import { KNOWLEDGE_BASE_WRITE_TOOL } from './knowledge_base/knowledge_base_write_tool'; +// eslint-disable-next-line @typescript-eslint/no-explicit-any +export type LangchainZodAny = z.ZodObject; + export const getAssistantTools = (graphEsqlToolEnabled?: boolean): AssistantTool[] => [ ALERT_COUNTS_TOOL, ATTACK_DISCOVERY_TOOL, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts index 4414fad57489e..154f477070ee7 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts @@ -10,6 +10,7 @@ import { z } from 'zod'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; import type { AIAssistantKnowledgeBaseDataClient } from '@kbn/elastic-assistant-plugin/server/ai_assistant_data_clients/knowledge_base'; import { APP_UI_ID } from '../../../../common'; +import type { LangchainZodAny } from '..'; export interface KnowledgeBaseRetrievalToolParams extends AssistantToolParams { kbDataClient: AIAssistantKnowledgeBaseDataClient; @@ -22,11 +23,7 @@ const toolDetails = { name: 'KnowledgeBaseRetrievalTool', }; -const schema = z.object({ - query: z.string().describe(`Summary of items/things to search for in the knowledge base`), -}); - -export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { +export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is KnowledgeBaseRetrievalToolParams => { @@ -44,7 +41,7 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { description: toolDetails.description, schema: z.object({ query: z.string().describe(`Summary of items/things to search for in the knowledge base`), - }), + }) as unknown as LangchainZodAny, func: async (input, _, cbManager) => { logger.debug(`KnowledgeBaseRetrievalToolParams:input\n ${JSON.stringify(input, null, 2)}`); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts index 1cf9fed3d802b..5b23ddab5129b 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts @@ -11,6 +11,7 @@ import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant- import type { AIAssistantKnowledgeBaseDataClient } from '@kbn/elastic-assistant-plugin/server/ai_assistant_data_clients/knowledge_base'; import type { KnowledgeBaseEntryCreateProps } from '@kbn/elastic-assistant-common'; import { APP_UI_ID } from '../../../../common'; +import type { LangchainZodAny } from '..'; export interface KnowledgeBaseWriteToolParams extends AssistantToolParams { kbDataClient: AIAssistantKnowledgeBaseDataClient; @@ -23,16 +24,7 @@ const toolDetails = { name: 'KnowledgeBaseWriteTool', }; -const schema = z.object({ - query: z.string().describe(`Summary of items/things to save in the knowledge base`), - required: z - .boolean() - .describe( - `Whether or not the entry is required to always be included in conversations. Is only true if the user explicitly asks for it to be required or always included in conversations, otherwise this is always false.` - ), -}); - -export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { +export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is KnowledgeBaseWriteToolParams => { @@ -55,7 +47,7 @@ export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { .describe( `Whether or not the entry is required to always be included in conversations. Is only true if the user explicitly asks for it to be required or always included in conversations, otherwise this is always false.` ), - }), + }) as unknown as LangchainZodAny, func: async (input, _, cbManager) => { logger.debug(`KnowledgeBaseWriteToolParams:input\n ${JSON.stringify(input, null, 2)}`); From d70622fc392f5f0adb9b6a90213e1673fd2899fa Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 27 Jun 2024 08:52:02 +0200 Subject: [PATCH 33/47] fix --- .../langchain/graphs/default_assistant_graph/helpers.ts | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts index 383b3e9f5cee8..ae9006e030f55 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts @@ -109,6 +109,11 @@ export const streamGraph = async ({ } } } + } else if (event.event === 'on_llm_end') { + const generations = event.data.output?.generations[0]; + if (generations && generations[0]?.generationInfo.finish_reason === 'stop') { + handleStreamEnd(finalMessage); + } } await processEvent(); @@ -129,7 +134,7 @@ export const streamGraph = async ({ }; // Start processing events, do not await! Return `responseWithHeaders` immediately - await processEvent(); + processEvent(); return responseWithHeaders; }; From d4d70638fee8064980a23f497d63df2cabf0f3dc Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Thu, 27 Jun 2024 11:41:01 +0000 Subject: [PATCH 34/47] [CI] Auto-commit changed files from 'node scripts/eslint --no-cache --fix' --- x-pack/plugins/elastic_assistant/server/types.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index d097c7927cf4a..34a003321061a 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -36,7 +36,6 @@ import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server' import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; -import { z } from 'zod'; import { AttackDiscoveryDataClient } from './ai_assistant_data_clients/attack_discovery'; import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/conversations'; import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; From 1b12104bd624939363170d6918aed867eeefc4c8 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 28 Jun 2024 21:22:12 +0200 Subject: [PATCH 35/47] fix --- .../langchain/graphs/default_assistant_graph/helpers.ts | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts index ae9006e030f55..383b3e9f5cee8 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts @@ -109,11 +109,6 @@ export const streamGraph = async ({ } } } - } else if (event.event === 'on_llm_end') { - const generations = event.data.output?.generations[0]; - if (generations && generations[0]?.generationInfo.finish_reason === 'stop') { - handleStreamEnd(finalMessage); - } } await processEvent(); @@ -134,7 +129,7 @@ export const streamGraph = async ({ }; // Start processing events, do not await! Return `responseWithHeaders` immediately - processEvent(); + await processEvent(); return responseWithHeaders; }; From 3dd68128814909f8e73abd4f5cdc181a80e04551 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 28 Jun 2024 22:43:38 +0200 Subject: [PATCH 36/47] bump elser memory --- .buildkite/scripts/steps/cloud/deploy.json | 2 +- .../kbn-elastic-assistant/impl/assistant/index.tsx | 7 ------- .../graph_esql_language_tool.ts | 14 ++++++++++---- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/.buildkite/scripts/steps/cloud/deploy.json b/.buildkite/scripts/steps/cloud/deploy.json index da1918a2a6953..336fed24afbe3 100644 --- a/.buildkite/scripts/steps/cloud/deploy.json +++ b/.buildkite/scripts/steps/cloud/deploy.json @@ -141,7 +141,7 @@ ], "id": "ml", "size": { - "value": 1024, + "value": 4096, "resource": "memory" } } diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx index dd96b4883c969..acc1861910121 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -519,13 +519,6 @@ const AssistantComponent: React.FC = ({ isFetchedAnonymizationFields, ]); - useEffect(() => {}, [ - areConnectorsFetched, - connectors, - conversationsLoaded, - currentConversation, - isLoading, - ]); const createCodeBlockPortals = useCallback( () => diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index 574ccfec20534..b378461113bfd 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -376,8 +376,8 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { esClient, search, llm } = params; - if (!llm) return null; + const { chain, esClient, search, llm } = params; + if (!llm || !chain) return null; return new DynamicStructuredTool({ name: toolDetails.name, @@ -400,11 +400,17 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { const app = workflow.compile(); let query; - try { query = await app.invoke({ question: input.question }, { recursionLimit: 20 }); } catch (e) { - return 'error'; + // Fallback to KnowledgeBase tool + const result = await chain.invoke( + { + query: input.question, + }, + cbManager + ); + return result.text; } return query.esqlQuery; From a1e8dcc6ebfb36be69c2a41ebf94471dec1dc8e7 Mon Sep 17 00:00:00 2001 From: Pedro Jaramillo Date: Fri, 28 Jun 2024 15:01:08 -0700 Subject: [PATCH 37/47] Fix lint --- x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx index acc1861910121..2760de2e63279 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -519,7 +519,6 @@ const AssistantComponent: React.FC = ({ isFetchedAnonymizationFields, ]); - const createCodeBlockPortals = useCallback( () => messageCodeBlocks?.map((codeBlocks: CodeBlockDetails[], i: number) => { From 4a8e974a31f5f76178ca02493233df822a26682b Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Sat, 29 Jun 2024 01:57:27 +0200 Subject: [PATCH 38/47] rely on knowledgebase --- .../graphs/default_assistant_graph/helpers.ts | 7 +- .../nodes/execute_tools.ts | 10 +- .../graph_esql_language_tool.ts | 196 +++--------------- 3 files changed, 49 insertions(+), 164 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts index 383b3e9f5cee8..ca96ab04616e8 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts @@ -109,9 +109,14 @@ export const streamGraph = async ({ } } } + } else if (event.event === 'on_llm_end') { + const generations = event.data.output?.generations[0]; + if (generations?.[0]?.generationInfo?.finish_reason === 'stop') { + handleStreamEnd(finalMessage); + } } - await processEvent(); + processEvent(); } catch (err) { // if I throw an error here, it crashes the server. Not sure how to get around that. // If I put await on this function the error works properly, but when there is not an error diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts index b42455e14f6f1..eab12b3e0ab24 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts @@ -38,7 +38,15 @@ export const executeTools = async ({ config, logger, state, tools }: ExecuteTool if (!agentAction || 'returnValues' in agentAction) { throw new Error('Agent has not been run yet'); } - const out = await toolExecutor.invoke(agentAction, config); + + let out; + try { + out = await toolExecutor.invoke(agentAction, config); + } catch (err) { + return { + steps: [{ action: agentAction, error: JSON.stringify(err, null, 2) }], + }; + } return { steps: [{ action: agentAction, observation: JSON.stringify(out, null, 2) }], }; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index b378461113bfd..22286a6725085 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -6,13 +6,12 @@ */ import Fs from 'fs'; -import { keyBy, mapValues, once, pick, isEmpty, map } from 'lodash'; +import { once, isEmpty, map } from 'lodash'; import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from 'zod'; -import pLimit from 'p-limit'; import Path from 'path'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; -import { JsonOutputParser, StringOutputParser } from '@langchain/core/output_parsers'; +import { StringOutputParser } from '@langchain/core/output_parsers'; import type { StateGraphArgs } from '@langchain/langgraph'; import { END, START, StateGraph } from '@langchain/langgraph'; import type { BaseMessage } from '@langchain/core/messages'; @@ -38,9 +37,9 @@ export const ECS_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ Here is some context for you to reference for your task, read it carefully as you will get questions about it later: - -{availableIndices} - + +{availableDataStreams} + `, ], [ @@ -85,13 +84,8 @@ const toolDetails = { interface IState { messages: BaseMessage[]; esqlQuery: string; - documentation: { - functions: string[]; - commands: string[]; - intention: string; - }; errors: ValidationResult['errors']; - availableIndices: string[]; + availableDataStreams: string[]; invalidQueries: string[]; } @@ -105,25 +99,6 @@ const graphState: StateGraphArgs['channels'] = { value: (x: string, y?: string) => y ?? x, default: () => '', }, - documentation: { - value: ( - x: { - functions: string[]; - commands: string[]; - intention: string; - }, - y?: { - functions: string[]; - commands: string[]; - intention: string; - } - ) => y ?? x, - default: () => ({ - functions: [], - commands: [], - intention: '', - }), - }, errors: { value: (x, y) => (y && !y.length ? y : x.concat(y)), default: () => [], @@ -132,14 +107,13 @@ const graphState: StateGraphArgs['channels'] = { value: (x, y) => x.concat(y), default: () => [], }, - availableIndices: { + availableDataStreams: { value: (x, y) => y ?? x, default: () => [], }, }; const readFile = promisify(Fs.readFile); -const readdir = promisify(Fs.readdir); const loadSystemMessage = once(async () => { const data = await readFile( @@ -151,114 +125,13 @@ const loadSystemMessage = once(async () => { return data.toString('utf-8'); }); -const loadEsqlDocs = once(async () => { - const dir = Path.join( - __dirname, - '../../../../../observability_solution/observability_ai_assistant_app/server/functions/query/esql_docs' - ); - const files = (await readdir(dir)).filter((file) => Path.extname(file) === '.txt'); - - if (!files.length) { - return {}; - } - - const limiter = pLimit(10); - return keyBy( - await Promise.all( - files.map((file) => - limiter(async () => { - const data = (await readFile(Path.join(dir, file))).toString('utf-8'); - const filename = Path.basename(file, '.txt'); - - const keyword = filename - .replace('esql-', '') - .replace('agg-', '') - .replaceAll('-', '_') - .toUpperCase(); - - return { - keyword: keyword === 'STATS_BY' ? 'STATS' : keyword, - data, - }; - }) - ) - ), - 'keyword' - ); -}); - -const getClassifyEsql = - ({ - userQuery, - llm, - esClient, - }: { - userQuery: string; - llm: NonNullable; - esClient: AssistantToolParams['esClient']; - }) => - async (state: IState, config?: RunnableConfig) => { - const [systemMessage, esqlDocs] = await Promise.all([loadSystemMessage(), loadEsqlDocs()]); - - const formatInstructions = `Respond only in valid JSON. The JSON object you return should match the following schema: - {{ commands: [], functions: [], intention: string }} - - Where commands is a list of processing or source commands that are referenced in the list of commands in this conversation. - Where functions is a list of functions that are referenced in the list of functions in this conversation. - Where intention is the user\'s intention. - `; - - const prompt = await ChatPromptTemplate.fromMessages([ - [ - 'system', - `${systemMessage} Answer the user query. Wrap the output in \`json\` tags\n{format_instructions}`, - ], - [ - 'user', - `Use this function to determine: - - what ES|QL functions and commands are candidates for answering the user's question - - whether the user has requested a query, and if so, it they want it to be executed, or just shown. - - All parameters are required. Make sure the functions and commands you request are available in the - system message. - - {query} - `, - ], - ]).partial({ - format_instructions: formatInstructions, - }); - - // Set up a parser - const parser = new JsonOutputParser(); - - const chainss = prompt.pipe(llm).pipe(parser); - - const result = await chainss.invoke({ - query: userQuery, - date: 'date', - msg: 'msg', - ip: 'ip', - }); - - const keywords = [ - ...(result.commands ?? []), - ...(result.functions ?? []), - 'SYNTAX', - 'OVERVIEW', - 'OPERATORS', - ].map((keyword) => keyword.toUpperCase()); - - const messagesToInclude = mapValues( - pick(esqlDocs, keywords), - ({ data }) => data - ) as unknown as IState['documentation']; - +const getDataStreams = + ({ esClient }: { esClient: AssistantToolParams['esClient'] }) => + async () => { const allDataStreams = await esClient.indices.getDataStream(); return { - documentation: messagesToInclude, - availableIndices: map(allDataStreams.data_streams, 'name'), + availableDataStreams: map(allDataStreams.data_streams, 'name'), }; }; @@ -268,23 +141,20 @@ const getGenerateQuery = const [systemMessage] = await Promise.all([loadSystemMessage()]); const answerPrompt = await ChatPromptTemplate.fromMessages([ - [ - 'system', - `${systemMessage}\nDocumentation: {documentation}\nAvailable indices: {availableIndices}`, - ], + ['system', `${systemMessage}\nAvailable data streams: {availableDataStreams}`], [ 'user', - `Answer the user's question that was previously asked ("{query}...") using the attached documentation. Take into account any previous errors {errors} and invalid ES|QL queries {invalidQueries}. + `Answer the user's question that was previously asked ("{query}..."). Take into account any previous errors {errors} and invalid ES|QL queries {invalidQueries}. Format any ES|QL query as follows: + \`\`\`esql \`\`\` + Respond in plain text. Do not attempt to use a function. - You must use commands and functions for which you have requested documentation. - DO NOT UNDER ANY CIRCUMSTANCES generate more than a single query. If multiple queries are needed, do it as a follow-up step. Make this clear to the user. For example: @@ -306,14 +176,13 @@ const getGenerateQuery = \`\`\` DO NOT UNDER ANY CIRCUMSTANCES use commands or functions that are not a capability of ES|QL - as mentioned in the system message and documentation. When converting queries from one language + as mentioned in the system message. When converting queries from one language to ES|QL, make sure that the functions are available and documented in ES|QL. E.g., for SPL's LEN, use LENGTH. For IF, use CASE. `, ], ]).partial({ - documentation: JSON.stringify(state.documentation), - availableIndices: JSON.stringify(state.availableIndices), + availableDataStreams: JSON.stringify(state.availableDataStreams), errors: state.errors.join('\n'), invalidQueries: state.invalidQueries.join('\n'), }); @@ -357,7 +226,10 @@ const getValidateQuery = }); return { errors: [] }; } catch (e) { - return { errors: [e?.message], invalidQueries: [state.esqlQuery] }; + return { + errors: e?.message.match(new RegExp(/Unknown column.*/)), + invalidQueries: [state.esqlQuery], + }; } }; @@ -372,7 +244,10 @@ const shouldRegenerate = (state: IState) => { export const GRAPH_ESQL_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, - isSupported: () => true, + isSupported: (params: AssistantToolParams): params is AssistantToolParams => { + const { kbDataClient, isEnabledKnowledgeBase, modelExists } = params; + return isEnabledKnowledgeBase && modelExists && kbDataClient != null; + }, getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; @@ -389,11 +264,11 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { const workflow = new StateGraph({ channels: graphState, }) - .addNode('classifyEsql', getClassifyEsql({ userQuery: input.question, llm, esClient })) + .addNode('getDataStreams', getDataStreams({ esClient })) .addNode('generateQuery', getGenerateQuery({ userQuery: input.question, llm })) .addNode('validateQuery', getValidateQuery({ search })) - .addEdge(START, 'classifyEsql') - .addEdge('classifyEsql', 'generateQuery') + .addEdge(START, 'getDataStreams') + .addEdge('getDataStreams', 'generateQuery') .addEdge('generateQuery', 'validateQuery') .addConditionalEdges('validateQuery', shouldRegenerate); @@ -403,17 +278,14 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { try { query = await app.invoke({ question: input.question }, { recursionLimit: 20 }); } catch (e) { - // Fallback to KnowledgeBase tool - const result = await chain.invoke( - { - query: input.question, - }, - cbManager - ); - return result.text; + return e; } - return query.esqlQuery; + return ` + \`\`\`esql + ${query.esqlQuery} + \`\`\` + `; }, tags: ['esql', 'query-generation', 'knowledge-base'], }); From 60c07df0595710016fc4a40fcfcba85777e48122 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Sat, 29 Jun 2024 14:12:00 +0200 Subject: [PATCH 39/47] fix --- x-pack/plugins/elastic_assistant/kibana.jsonc | 1 + .../execute_custom_llm_chain/index.ts | 2 + .../server/lib/langchain/executors/types.ts | 2 + .../graphs/default_assistant_graph/index.ts | 2 + .../nodes/execute_tools.ts | 2 +- .../server/routes/attack_discovery/helpers.ts | 7 - .../attack_discovery/post_attack_discovery.ts | 1 - .../routes/post_actions_connector_execute.ts | 1 + .../server/routes/request_context_factory.ts | 6 + .../plugins/elastic_assistant/server/types.ts | 14 +- .../scripts/run_cypress/parallel.ts | 2 +- .../graph_esql_language_tool.ts | 186 ++++++++++++------ 12 files changed, 151 insertions(+), 75 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/kibana.jsonc b/x-pack/plugins/elastic_assistant/kibana.jsonc index 9879ba274d209..6db16861bf934 100644 --- a/x-pack/plugins/elastic_assistant/kibana.jsonc +++ b/x-pack/plugins/elastic_assistant/kibana.jsonc @@ -10,6 +10,7 @@ "requiredPlugins": [ "actions", "data", + "dataViews", "ml", "taskManager", "licensing", diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts index fda963c8ffdac..9d4e43be96744 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts @@ -50,6 +50,7 @@ export const callAgentExecutor: AgentExecutor = async ({ size, traceOptions, search, + dataViews, }) => { const isOpenAI = llmType === 'openai'; const llmClass = isOpenAI ? ActionsClientChatOpenAI : ActionsClientSimpleChatModel; @@ -104,6 +105,7 @@ export const callAgentExecutor: AgentExecutor = async ({ request, size, search, + dataViews, }; const tools: ToolInterface[] = assistantTools.flatMap( diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index 543d02400a0e8..7c21cf8579964 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -15,6 +15,7 @@ import { ExecuteConnectorRequestBody, Message, Replacements } from '@kbn/elastic import { StreamResponseWithHeaders } from '@kbn/ml-response-stream/server'; import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; +import { DataViewsService } from '@kbn/data-views-plugin/server'; import { ResponseBody } from '../types'; import type { AssistantTool } from '../../../types'; import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store'; @@ -46,6 +47,7 @@ export interface AgentExecutorParams { dataClients?: AssistantDataClients; esClient: ElasticsearchClient; search: ReturnType; + dataViews: DataViewsService; esStore: ElasticsearchStore; langChainMessages: BaseMessage[]; llmType?: string; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts index 625ad960a27e0..8fd11a69b6911 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts @@ -46,6 +46,7 @@ export const callAssistantGraph: AgentExecutor = async ({ size, traceOptions, search, + dataViews, }) => { const logger = parentLogger.get('defaultAssistantGraph'); const isOpenAI = llmType === 'openai'; @@ -95,6 +96,7 @@ export const callAssistantGraph: AgentExecutor = async ({ request, size, search, + dataViews, }; const tools: StructuredTool[] = assistantTools.flatMap( diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts index eab12b3e0ab24..47e201650a54d 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts @@ -44,7 +44,7 @@ export const executeTools = async ({ config, logger, state, tools }: ExecuteTool out = await toolExecutor.invoke(agentAction, config); } catch (err) { return { - steps: [{ action: agentAction, error: JSON.stringify(err, null, 2) }], + steps: [{ action: agentAction, observation: JSON.stringify(`Error: ${err}`, null, 2) }], }; } return { diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts index cfff0fe9a62f6..9dca7ee46cbda 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts @@ -26,7 +26,6 @@ import { transformError } from '@kbn/securitysolution-es-utils'; import type { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server'; import moment from 'moment/moment'; import { uniq } from 'lodash/fp'; -import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; import { getLangSmithTracer } from '../evaluate/utils'; import { getLlmType } from '../utils'; import type { GetRegisteredTools } from '../../services/app_context'; @@ -66,7 +65,6 @@ export const getAssistantToolParams = ({ latestReplacements, onNewReplacements, request, - search, size, }: { actions: ActionsPluginStart; @@ -86,7 +84,6 @@ export const getAssistantToolParams = ({ unknown, ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody >; - search: ReturnType; size: number; }) => { const traceOptions = { @@ -121,7 +118,6 @@ export const getAssistantToolParams = ({ logger, onNewReplacements, request, - search, size, }); }; @@ -136,7 +132,6 @@ const formatAssistantToolParams = ({ logger, onNewReplacements, request, - search, size, }: { alertsIndexPattern: string; @@ -152,7 +147,6 @@ const formatAssistantToolParams = ({ unknown, ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody >; - search: ReturnType; size: number; }): AssistantToolParams => ({ alertsIndexPattern, @@ -167,7 +161,6 @@ const formatAssistantToolParams = ({ onNewReplacements, replacements: latestReplacements, request, - search, size, }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts index 9c0f6887d863a..8ff2cd72ee36c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts @@ -134,7 +134,6 @@ export const postAttackDiscoveryRoute = ( onNewReplacements, request, size, - search: assistantContext.search, }); // invoke the attack discovery tool: diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index ed1024a6c5395..1b64b8278401f 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -369,6 +369,7 @@ export const postActionsConnectorExecuteRoute = ( connectorId, conversationId, dataClients, + dataViews: assistantContext.dataViews, esClient, search: assistantContext.search, esStore, diff --git a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts index 3e8d5827cf68e..9b648ae7403b8 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -65,6 +65,12 @@ export class RequestContextFactory implements IRequestContextFactory { search: startPlugins.data.search.asScoped(request), + dataViews: await startPlugins.dataViews.dataViewsServiceFactory( + coreContext.savedObjects.client, + coreContext.elasticsearch.client.asCurrentUser, + request + ), + logger: this.logger, getServerBasePath: () => core.http.basePath.serverBasePath, diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index b3d8708fd9e7b..409d17859f172 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -34,9 +34,14 @@ import { } from '@kbn/elastic-assistant-common'; import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; -import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; +import { + ActionsClientChatOpenAI, + ActionsClientLlm, + ActionsClientSimpleChatModel, +} from '@kbn/langchain/server'; import { DataPluginStart } from '@kbn/data-plugin/server/plugin'; +import { DataViewsServerPluginStart, DataViewsService } from '@kbn/data-views-plugin/server'; import { AttackDiscoveryDataClient } from './ai_assistant_data_clients/attack_discovery'; import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/conversations'; import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; @@ -97,6 +102,7 @@ export interface ElasticAssistantPluginSetupDependencies { export interface ElasticAssistantPluginStartDependencies { actions: ActionsPluginStart; data: DataPluginStart; + dataViews: DataViewsServerPluginStart; spaces?: SpacesPluginStart; security: SecurityServiceStart; } @@ -105,6 +111,7 @@ export interface ElasticAssistantApiRequestHandlerContext { core: CoreRequestHandlerContext; actions: ActionsPluginStart; search: ReturnType; + dataViews: DataViewsService; getRegisteredFeatures: GetRegisteredFeatures; getRegisteredTools: GetRegisteredTools; logger: Logger; @@ -220,7 +227,7 @@ export interface AssistantToolParams { esClient: ElasticsearchClient; kbDataClient?: AIAssistantKnowledgeBaseDataClient; langChainTimeout?: number; - llm?: ActionsClientChatOpenAI | ActionsClientSimpleChatModel; + llm?: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; logger: Logger; modelExists: boolean; onNewReplacements?: (newReplacements: Replacements) => void; @@ -231,5 +238,6 @@ export interface AssistantToolParams { ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody >; size?: number; - search: ReturnType; + dataViews?: DataViewsService; + search?: ReturnType; } diff --git a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts index 4b0f586779270..a92608b070afb 100644 --- a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts +++ b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts @@ -340,7 +340,7 @@ ${JSON.stringify( extraKbnOpts: options?.installDir || options?.ci || !isOpen ? [] - : ['--dev', '--no-dev-config', '--no-dev-credentials'], + : ['--dev', '--no-dev-credentials'], onEarlyExit, inspect: argv.inspect, }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index 22286a6725085..d7d78dbdb9289 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -6,7 +6,7 @@ */ import Fs from 'fs'; -import { once, isEmpty, map } from 'lodash'; +import { once, isEmpty, map, without } from 'lodash'; import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from 'zod'; import Path from 'path'; @@ -16,15 +16,16 @@ import type { StateGraphArgs } from '@langchain/langgraph'; import { END, START, StateGraph } from '@langchain/langgraph'; import type { BaseMessage } from '@langchain/core/messages'; import type { RunnableConfig } from '@langchain/core/runnables'; -import type { ValidationResult } from '@kbn/esql-validation-autocomplete/src/validation/types'; -import { getESQLQueryColumns } from '@kbn/esql-utils'; +import { getESQLQueryColumns, getIndexPatternFromESQLQuery } from '@kbn/esql-utils'; import { promisify } from 'util'; import { ChatPromptTemplate } from '@langchain/core/prompts'; import { validateQuery } from '@kbn/esql-validation-autocomplete'; +import type { EditorError, ESQLMessage } from '@kbn/esql-ast'; import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; import { APP_UI_ID } from '../../../../common'; import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; import type { LangchainZodAny } from '..'; +import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; export const INLINE_ESQL_QUERY_REGEX = /```esql\s*(.*?)\s*```/gms; @@ -37,16 +38,16 @@ export const ECS_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ Here is some context for you to reference for your task, read it carefully as you will get questions about it later: - -{availableDataStreams} - + +{availableDataViews} + `, ], [ 'human', `{input}. Example response format: - + A: Please find the ESQL query below: \`\`\`esql FROM logs-* @@ -84,8 +85,9 @@ const toolDetails = { interface IState { messages: BaseMessage[]; esqlQuery: string; - errors: ValidationResult['errors']; - availableDataStreams: string[]; + errors: string[]; + availableDataViews: string[]; + dataViewFields: string[]; invalidQueries: string[]; } @@ -107,7 +109,11 @@ const graphState: StateGraphArgs['channels'] = { value: (x, y) => x.concat(y), default: () => [], }, - availableDataStreams: { + availableDataViews: { + value: (x, y) => y ?? x, + default: () => [], + }, + dataViewFields: { value: (x, y) => y ?? x, default: () => [], }, @@ -126,64 +132,88 @@ const loadSystemMessage = once(async () => { }); const getDataStreams = - ({ esClient }: { esClient: AssistantToolParams['esClient'] }) => + ({ dataViews }: { dataViews: AssistantToolParams['dataViews'] }) => async () => { - const allDataStreams = await esClient.indices.getDataStream(); + const allDataStreams = await dataViews?.getTitles(); return { - availableDataStreams: map(allDataStreams.data_streams, 'name'), + availableDataViews: allDataStreams, }; }; const getGenerateQuery = - ({ userQuery, llm }: { userQuery: string; llm: NonNullable }) => + ({ + userQuery, + llm, + }: { + userQuery: string; + llm: ActionsClientChatOpenAI | ActionsClientSimpleChatModel; + }) => async (state: IState) => { const [systemMessage] = await Promise.all([loadSystemMessage()]); const answerPrompt = await ChatPromptTemplate.fromMessages([ - ['system', `${systemMessage}\nAvailable data streams: {availableDataStreams}`], + [ + 'system', + `${systemMessage} + + Available indices: + {availableDataViews} + `, + ], [ 'user', - `Answer the user's question that was previously asked ("{query}..."). Take into account any previous errors {errors} and invalid ES|QL queries {invalidQueries}. + `Answer the user's question that was previously asked ("{query}..."). Take into account any previous errors: +{errors} - Format any ES|QL query as follows: - - \`\`\`esql - - \`\`\` - +and invalid ES|QL queries: +{invalidQueries} - Respond in plain text. Do not attempt to use a function. +If errors were related to "Unknown column" make sure to check the Available index fields: +{dataViewFields} - DO NOT UNDER ANY CIRCUMSTANCES generate more than a single query. - If multiple queries are needed, do it as a follow-up step. Make this clear to the user. For example: +Format any ES|QL query as follows: + +\`\`\`esql + +\`\`\` + + +Respond in plain text. Do not attempt to use a function. - Human: plot both yesterday's and today's data. +DO NOT UNDER ANY CIRCUMSTANCES generate more than a single query. +If multiple queries are needed, do it as a follow-up step. Make this clear to the user. For example: - Assistant: Here's how you can plot yesterday's data: - \`\`\`esql - - \`\`\` +Human: plot both yesterday's and today's data. - Let's see that first. We'll look at today's data next. +Assistant: Here's how you can plot yesterday's data: + +\`\`\`esql + +\`\`\` + - Human: +Let's see that first. We'll look at today's data next. - Assistant: Let's look at today's data: +Human: - \`\`\`esql - - \`\`\` +Assistant: Let's look at today's data: - DO NOT UNDER ANY CIRCUMSTANCES use commands or functions that are not a capability of ES|QL - as mentioned in the system message. When converting queries from one language - to ES|QL, make sure that the functions are available and documented in ES|QL. - E.g., for SPL's LEN, use LENGTH. For IF, use CASE. - `, + +\`\`\`esql + +\`\`\` + + +DO NOT UNDER ANY CIRCUMSTANCES use commands or functions that are not a capability of ES|QL +as mentioned in the system message. When converting queries from one language +to ES|QL, make sure that the functions are available and documented in ES|QL. +E.g., for SPL's LEN, use LENGTH. For IF, use CASE.`, ], ]).partial({ - availableDataStreams: JSON.stringify(state.availableDataStreams), + availableDataViews: state.availableDataViews.join('\n'), errors: state.errors.join('\n'), + dataViewFields: state.dataViewFields.join('\n'), invalidQueries: state.invalidQueries.join('\n'), }); @@ -208,7 +238,13 @@ const getGenerateQuery = }; const getValidateQuery = - ({ search }: { search: AssistantToolParams['search'] }) => + ({ + dataViews, + search, + }: { + dataViews: AssistantToolParams['dataViews']; + search: AssistantToolParams['search']; + }) => async (state: IState, config?: RunnableConfig) => { const { errors } = await validateQuery(state.esqlQuery, getAstAndSyntaxErrors, { // setting this to true, we don't want to validate the index / fields existence @@ -216,21 +252,49 @@ const getValidateQuery = }); if (!isEmpty(errors)) { - return { errors, invalidQueries: [state.esqlQuery] }; - } - - try { - await getESQLQueryColumns({ - esqlQuery: state.esqlQuery, - search: search.search, - }); - return { errors: [] }; - } catch (e) { return { - errors: e?.message.match(new RegExp(/Unknown column.*/)), + errors: map( + errors, + (error) => (error as ESQLMessage)?.text || (error as EditorError)?.message + ), invalidQueries: [state.esqlQuery], }; } + + if (search?.search) { + try { + await getESQLQueryColumns({ + esqlQuery: state.esqlQuery, + search: search.search, + }); + return { errors: [] }; + } catch (e) { + let dataViewFields; + let availableDataViews = state.availableDataViews; + if (dataViews) { + const indexPattern = getIndexPatternFromESQLQuery(state.esqlQuery); + let indexFields; + try { + indexFields = await dataViews.getFieldsForWildcard({ pattern: indexPattern }); + } catch (err) { + availableDataViews = without(availableDataViews, indexPattern); + } + + dataViewFields = map(indexFields, (field) => `${field.name} (${field.type})`); + } + + // console.error('eeeeee', e); + + return { + errors: e?.message.match(new RegExp(/Unknown column.*/)) ?? e?.message, + dataViewFields, + availableDataViews, + invalidQueries: [state.esqlQuery], + }; + } + } + + return {}; }; const shouldRegenerate = (state: IState) => { @@ -251,7 +315,7 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { chain, esClient, search, llm } = params; + const { chain, dataViews, search, llm } = params; if (!llm || !chain) return null; return new DynamicStructuredTool({ @@ -264,9 +328,9 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { const workflow = new StateGraph({ channels: graphState, }) - .addNode('getDataStreams', getDataStreams({ esClient })) + .addNode('getDataStreams', getDataStreams({ dataViews })) .addNode('generateQuery', getGenerateQuery({ userQuery: input.question, llm })) - .addNode('validateQuery', getValidateQuery({ search })) + .addNode('validateQuery', getValidateQuery({ dataViews, search })) .addEdge(START, 'getDataStreams') .addEdge('getDataStreams', 'generateQuery') .addEdge('generateQuery', 'validateQuery') @@ -281,11 +345,9 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { return e; } - return ` - \`\`\`esql - ${query.esqlQuery} - \`\`\` - `; + // console.error('query', query.esqlQuery); + + return query.esqlQuery; }, tags: ['esql', 'query-generation', 'knowledge-base'], }); From afff8e7b67664223ae3370620cae979af83898a4 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Sat, 29 Jun 2024 18:09:11 +0200 Subject: [PATCH 40/47] test --- .../graph_esql_language_tool.ts | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index d7d78dbdb9289..1c281eff4ba10 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -22,10 +22,14 @@ import { ChatPromptTemplate } from '@langchain/core/prompts'; import { validateQuery } from '@kbn/esql-validation-autocomplete'; import type { EditorError, ESQLMessage } from '@kbn/esql-ast'; import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; +import type { + ActionsClientChatOpenAI, + ActionsClientLlm, + ActionsClientSimpleChatModel, +} from '@kbn/langchain/server'; import { APP_UI_ID } from '../../../../common'; import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; import type { LangchainZodAny } from '..'; -import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server'; export const INLINE_ESQL_QUERY_REGEX = /```esql\s*(.*?)\s*```/gms; @@ -147,7 +151,7 @@ const getGenerateQuery = llm, }: { userQuery: string; - llm: ActionsClientChatOpenAI | ActionsClientSimpleChatModel; + llm: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; }) => async (state: IState) => { const [systemMessage] = await Promise.all([loadSystemMessage()]); @@ -340,7 +344,7 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { let query; try { - query = await app.invoke({ question: input.question }, { recursionLimit: 20 }); + query = await app.invoke({ question: input.question }); } catch (e) { return e; } From 36646db5103b1a6939aef0675194fa5dfe2613c8 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Sat, 29 Jun 2024 16:21:11 +0000 Subject: [PATCH 41/47] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/plugins/elastic_assistant/tsconfig.json | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index f63a8da530196..2414e0f286185 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -45,6 +45,7 @@ "@kbn/core-saved-objects-api-server", "@kbn/langchain", "@kbn/stack-connectors-plugin", + "@kbn/data-views-plugin", ], "exclude": [ "target/**/*", From 90f1991afd9ba4f718cb812a04824218c762067e Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Sun, 30 Jun 2024 23:24:39 +0200 Subject: [PATCH 42/47] bump --- package.json | 16 +++---- .../routes/post_actions_connector_execute.ts | 2 + yarn.lock | 42 +++++++++---------- 3 files changed, 31 insertions(+), 29 deletions(-) diff --git a/package.json b/package.json index a1a81073b7afb..5d2fff5e88cfd 100644 --- a/package.json +++ b/package.json @@ -80,8 +80,8 @@ "resolutions": { "**/@bazel/typescript/protobufjs": "6.11.4", "**/@hello-pangea/dnd": "16.6.0", - "**/@langchain/core": "0.2.9", - "**/@langchain/openai": "0.2.0", + "**/@langchain/core": "0.2.11", + "**/@langchain/openai": "0.2.1", "**/@types/node": "20.10.5", "**/@typescript-eslint/utils": "5.62.0", "**/chokidar": "^3.5.3", @@ -89,7 +89,7 @@ "**/globule/minimatch": "^3.1.2", "**/hoist-non-react-statics": "^3.3.2", "**/isomorphic-fetch/node-fetch": "^2.6.7", - "**/langchain": "0.2.6", + "**/langchain": "0.2.8", "**/react-intl/**/@types/react": "^17.0.45", "**/remark-parse/trim": "1.0.1", "**/sharp": "0.32.6", @@ -934,10 +934,10 @@ "@kbn/watcher-plugin": "link:x-pack/plugins/watcher", "@kbn/xstate-utils": "link:packages/kbn-xstate-utils", "@kbn/zod-helpers": "link:packages/kbn-zod-helpers", - "@langchain/community": "^0.2.13", - "@langchain/core": "^0.2.9", - "@langchain/langgraph": "^0.0.24", - "@langchain/openai": "^0.2.0", + "@langchain/community": "^0.2.15", + "@langchain/core": "^0.2.11", + "@langchain/langgraph": "^0.0.25", + "@langchain/openai": "^0.2.1", "@langtrase/trace-attributes": "^3.0.8", "@launchdarkly/node-server-sdk": "^9.4.6", "@loaders.gl/core": "^3.4.7", @@ -1074,7 +1074,7 @@ "jsonwebtoken": "^9.0.2", "jsts": "^1.6.2", "kea": "^2.6.0", - "langchain": "^0.2.6", + "langchain": "^0.2.8", "langsmith": "^0.1.30", "launchdarkly-js-client-sdk": "^3.4.0", "launchdarkly-node-server-sdk": "^7.0.3", diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 1b64b8278401f..ae45c9bb3642c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -5,6 +5,8 @@ * 2.0. */ +/* eslint-disable complexity */ + import { IRouter, Logger } from '@kbn/core/server'; import { transformError } from '@kbn/securitysolution-es-utils'; import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; diff --git a/yarn.lock b/yarn.lock index d2882fc425d09..e4de6481455c8 100644 --- a/yarn.lock +++ b/yarn.lock @@ -6973,10 +6973,10 @@ resolved "https://registry.yarnpkg.com/@kwsites/promise-deferred/-/promise-deferred-1.1.1.tgz#8ace5259254426ccef57f3175bc64ed7095ed919" integrity sha512-GaHYm+c0O9MjZRu0ongGBRbinu8gVAMd2UZjji6jVmqKtZluZnptXGWhz1E8j8D2HJ3f/yMxKAUC0b+57wncIw== -"@langchain/community@^0.2.13": - version "0.2.13" - resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.13.tgz#7a327663d2a6006456ff136d2b16cb2b1a76b541" - integrity sha512-f0GZCGM5XP0r+H643GpUU4YelKHsUdhUY1Kb8rKpCoy8zgs1nUkiYDVylAf0ezwUOT4NYCEuwpw0jj8hQSLn1Q== +"@langchain/community@^0.2.15": + version "0.2.15" + resolved "https://registry.yarnpkg.com/@langchain/community/-/community-0.2.15.tgz#91ad5af44bfc72e83f6b6bd58cf5b29e53effed6" + integrity sha512-WOsNQGhriXh5tqRWfX3nthWO6RoVtM5kceX2GbJhqk09KV4R+1QmrOyph3OrpmjRA/YuSf0a+94LHY+c/QGolw== dependencies: "@langchain/core" "~0.2.9" "@langchain/openai" "~0.1.0" @@ -6990,10 +6990,10 @@ zod "^3.22.3" zod-to-json-schema "^3.22.5" -"@langchain/core@0.2.9", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@>=0.2.8 <0.3.0", "@langchain/core@^0.2.9", "@langchain/core@~0.2.0", "@langchain/core@~0.2.9": - version "0.2.9" - resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.9.tgz#8a22585ef11d2ca8742a8bbfe77dd25baedc7779" - integrity sha512-pJshopBZqMNF020q0OrrO+vfApWTZUlZecRYMM7TWA5M8/zvEyU/mgA9DlzeRjjDmG6pwF6dIKVjpl6fIGVXlQ== +"@langchain/core@0.2.11", "@langchain/core@>0.1.0 <0.3.0", "@langchain/core@>0.1.61 <0.3.0", "@langchain/core@>=0.2.8 <0.3.0", "@langchain/core@>=0.2.9 <0.3.0", "@langchain/core@^0.2.11", "@langchain/core@~0.2.9": + version "0.2.11" + resolved "https://registry.yarnpkg.com/@langchain/core/-/core-0.2.11.tgz#5f47467e20e56b250831baef20083657c6facb4c" + integrity sha512-d4SNL7WI0c3oHrV4WxCRH1/TNqdePXEzYjYwIb4aEH6lW1aM0utGhLbNthX+aYkOL4Ynx2FoG4h91ECIipiKWQ== dependencies: ansi-styles "^5.0.0" camelcase "6" @@ -7008,18 +7008,18 @@ zod "^3.22.4" zod-to-json-schema "^3.22.3" -"@langchain/langgraph@^0.0.24": - version "0.0.24" - resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.24.tgz#814e686dcef700d15a6f90e27e9c9f79d75faef4" - integrity sha512-fW9cnz62oKZFAlyO/4oEjXxthrqZPQtqyX4f7ttKEi0rJZKeuoohvOtnC8faq6nrtMtX9JpLixHjK0SgN7XN3g== +"@langchain/langgraph@^0.0.25": + version "0.0.25" + resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.0.25.tgz#2582d8652e2dda722f0c5043c1d0254a778e2486" + integrity sha512-DiTnB5Psm0y7TSgHdK4r/r+xzLohbN4zMQL+5Wk3EmOGX45ioBp98AqL8hYdyxKgHM6cjoIFHavHF7EhMg+ugQ== dependencies: "@langchain/core" ">0.1.61 <0.3.0" uuid "^9.0.1" -"@langchain/openai@0.2.0", "@langchain/openai@>=0.1.0 <0.3.0", "@langchain/openai@^0.2.0", "@langchain/openai@~0.1.0": - version "0.2.0" - resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.2.0.tgz#342e49d15b946fa01128d1bb81357e688e7cf567" - integrity sha512-gZd+0IOxpiKuh1m6KTT5vtUoOO72GEYyoU4+c6qAUucOEqQS0Vvz3lMGyNWLjK4x4Xpd+r8GAF5mj/jvghwP1A== +"@langchain/openai@0.2.1", "@langchain/openai@>=0.1.0 <0.3.0", "@langchain/openai@^0.2.1", "@langchain/openai@~0.1.0": + version "0.2.1" + resolved "https://registry.yarnpkg.com/@langchain/openai/-/openai-0.2.1.tgz#2c0c2cb6bd7839d8ce342c97099c6e35f2dde40d" + integrity sha512-Ti3C6ZIUPaueIPAfMljMnLu3GSGNq5KmrlHeWkIbrLShOBlzj4xj7mRfR73oWgAC0qivfxdkfbB0e+WCY+oRJw== dependencies: "@langchain/core" ">=0.2.8 <0.3.0" js-tiktoken "^1.0.12" @@ -21774,12 +21774,12 @@ kuler@^2.0.0: resolved "https://registry.yarnpkg.com/kuler/-/kuler-2.0.0.tgz#e2c570a3800388fb44407e851531c1d670b061b3" integrity sha512-Xq9nH7KlWZmXAtodXDDRE7vs6DU1gTU8zYDHDiWLSip45Egwq3plLHzPn27NgvzL2r1LMPC1vdqh98sQxtqj4A== -langchain@0.2.3, langchain@0.2.6, langchain@^0.2.6: - version "0.2.6" - resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.6.tgz#22249707ba800c38ec9a6cca36e383a881227393" - integrity sha512-vDAJHGu/lA4pn3hkyzSC6RiaZhtj0ozfRyG8L6J2vCnXyJV/lgk9uGMP2x645EBrSozBMHJBng1UYeaUR/1fQQ== +langchain@0.2.3, langchain@0.2.8, langchain@^0.2.8: + version "0.2.8" + resolved "https://registry.yarnpkg.com/langchain/-/langchain-0.2.8.tgz#9bd77f5c12071d0ccb637c04fc33415e5369e5aa" + integrity sha512-kb2IOMA71xH8e6EXFg0l4S+QSMC/c796pj1+7mPBkR91HHwoyHZhFRrBaZv4tV+Td+Ba91J2uEDBmySklZLpNQ== dependencies: - "@langchain/core" "~0.2.0" + "@langchain/core" ">=0.2.9 <0.3.0" "@langchain/openai" ">=0.1.0 <0.3.0" "@langchain/textsplitters" "~0.0.0" binary-extensions "^2.2.0" From aa7ff2b3c7bb62dc2218b9d5a09ce57e5efc6569 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 1 Jul 2024 18:40:34 +0200 Subject: [PATCH 43/47] fix system message --- .../graphs/default_assistant_graph/index.ts | 11 +++- .../post_actions_connector_execute.test.ts | 4 ++ .../plugins/elastic_assistant/server/types.ts | 2 + .../server/lib/conversational_chain.test.ts | 9 ++- .../scripts/run_cypress/parallel.ts | 2 +- .../graph_esql_language_tool.ts | 55 ++++++------------- 6 files changed, 40 insertions(+), 43 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts index 8fd11a69b6911..cd9b2ee1dd393 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts @@ -86,6 +86,7 @@ export const callAssistantGraph: AgentExecutor = async ({ anonymizationFields, chain, esClient, + esStore, isEnabledKnowledgeBase, kbDataClient: dataClients?.kbDataClient, llm: model, @@ -131,7 +132,15 @@ export const callAssistantGraph: AgentExecutor = async ({ const inputs = { input: latestMessage[0].content as string }; if (isStream) { - return streamGraph({ apmTracer, assistantGraph, inputs, logger, onLlmResponse, request }); + return streamGraph({ + apmTracer, + assistantGraph, + inputs, + logger, + onLlmResponse, + request, + traceOptions, + }); } const graphResponse = await invokeGraph({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts index 5ee8d8e83c846..e7feaeb20cf93 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts @@ -109,6 +109,10 @@ const mockContext = { elastic_cloud_user: false, metadata: { _reserved: false }, }), + getAIAssistantKnowledgeBaseDataClient: jest.fn().mockResolvedValue({ + indexTemplateAndPattern: {}, + getKnowledgeBaseDocuments: jest.fn(), + }), getAIAssistantConversationsDataClient: jest.fn().mockResolvedValue({ getConversation: jest.fn().mockResolvedValue(existingConversation), updateConversation: jest.fn().mockResolvedValue(existingConversation), diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 409d17859f172..52971e6791da0 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -47,6 +47,7 @@ import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/ import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; import { AIAssistantDataClient } from './ai_assistant_data_clients'; import { AIAssistantKnowledgeBaseDataClient } from './ai_assistant_data_clients/knowledge_base'; +import { ElasticsearchStore } from './lib/langchain/elasticsearch_store/elasticsearch_store'; export const PLUGIN_ID = 'elasticAssistant' as const; @@ -225,6 +226,7 @@ export interface AssistantToolParams { isEnabledKnowledgeBase: boolean; chain?: RetrievalQAChain; esClient: ElasticsearchClient; + esStore: ElasticsearchStore; kbDataClient?: AIAssistantKnowledgeBaseDataClient; langChainTimeout?: number; llm?: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; diff --git a/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts b/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts index fa6ee8e6e1abd..d1533247048f2 100644 --- a/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts +++ b/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts @@ -433,7 +433,10 @@ describe('conversational chain', () => { expectedDocs: [ { documents: [ - { metadata: { _id: '1', _index: 'index' } }, + { + metadata: { _id: '1', _index: 'index' }, + pageContent: '', + }, { metadata: { _id: '1', _index: 'website' }, pageContent: expect.any(String), @@ -444,8 +447,8 @@ describe('conversational chain', () => { ], // Even with body_content of 1000, the token count should be below the model limit of 100 expectedTokens: [ - { type: 'context_token_count', count: 70 }, - { type: 'prompt_token_count', count: 97 }, + { type: 'context_token_count', count: 73 }, + { type: 'prompt_token_count', count: 100 }, ], expectedHasClipped: true, expectedSearchRequest: [ diff --git a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts index a92608b070afb..4b0f586779270 100644 --- a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts +++ b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts @@ -340,7 +340,7 @@ ${JSON.stringify( extraKbnOpts: options?.installDir || options?.ci || !isOpen ? [] - : ['--dev', '--no-dev-credentials'], + : ['--dev', '--no-dev-config', '--no-dev-credentials'], onEarlyExit, inspect: argv.inspect, }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index 6c87f3b362350..09c98730a4289 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -5,11 +5,9 @@ * 2.0. */ -import Fs from 'fs'; -import { once, isEmpty, map, without } from 'lodash'; +import { isEmpty, map, without } from 'lodash'; import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from 'zod'; -import Path from 'path'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; import { StringOutputParser } from '@langchain/core/output_parsers'; import type { StateGraphArgs } from '@langchain/langgraph'; @@ -17,7 +15,6 @@ import { END, START, StateGraph } from '@langchain/langgraph'; import type { BaseMessage } from '@langchain/core/messages'; import type { RunnableConfig } from '@langchain/core/runnables'; import { getESQLQueryColumns, getIndexPatternFromESQLQuery } from '@kbn/esql-utils'; -import { promisify } from 'util'; import { ChatPromptTemplate } from '@langchain/core/prompts'; import { validateQuery } from '@kbn/esql-validation-autocomplete'; import type { EditorError, ESQLMessage } from '@kbn/esql-ast'; @@ -27,10 +24,10 @@ import type { ActionsClientLlm, ActionsClientSimpleChatModel, } from '@kbn/langchain/server'; +import type { ElasticsearchStore } from '@kbn/elastic-assistant-plugin/server/lib/langchain/elasticsearch_store/elasticsearch_store'; import { APP_UI_ID } from '../../../../common'; import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; import type { LangchainZodAny } from '..'; -import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from '../knowledge_base/knowledge_base_retrieval_tool'; export const INLINE_ESQL_QUERY_REGEX = /```esql\s*(.*?)\s*```/gms; @@ -94,18 +91,6 @@ const graphState: StateGraphArgs['channels'] = { }, }; -const readFile = promisify(Fs.readFile); - -const loadSystemMessage = once(async () => { - const data = await readFile( - Path.join( - __dirname, - '../../../../../observability_solution/observability_ai_assistant_app/server/functions/query/system_message.txt' - ) - ); - return data.toString('utf-8'); -}); - const getDataStreams = ({ dataViews }: { dataViews: AssistantToolParams['dataViews'] }) => async () => { @@ -120,21 +105,26 @@ const getGenerateQuery = ({ userQuery, llm, + esStore, }: { userQuery: string; llm: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; + esStore: ElasticsearchStore; }) => async (state: IState) => { - const [systemMessage] = await Promise.all([loadSystemMessage()]); + const knowledgeBaseDocs = await esStore.similaritySearch(userQuery); + const documentation = map(knowledgeBaseDocs, 'pageContent').join('\n'); const answerPrompt = await ChatPromptTemplate.fromMessages([ [ 'system', - `${systemMessage} - As an expert user of Elastic Security, please generate an accurate and valid ESQL query to detect the use case below. Your response should be formatted to be able to use immediately in an Elastic Security timeline or detection rule. Take your time with the answer, check your knowledge really well on all the functions I am asking for. For ES|QL answers specifically, you should only ever answer with what's available in your private knowledge. I cannot afford for queries to be inaccurate. Assume I am using the Elastic Common Schema and Elastic Agent. Under any circumstances wrap index in quotes. + `As an expert user of Elastic Security, please generate an accurate and valid ESQL query to detect the use case below. Your response should be formatted to be able to use immediately in an Elastic Security timeline or detection rule. Take your time with the answer, check your knowledge really well on all the functions I am asking for. For ES|QL answers specifically, you should only ever answer with what's available in your private knowledge. I cannot afford for queries to be inaccurate. Assume I am using the Elastic Common Schema and Elastic Agent. Under any circumstances wrap index in quotes. If multiple indices are matched please try to use wildcard to match all indices. If you are unsure about the index name, please refer to the context provided below. + ES|QL documentation: + {documentation} + Available indices: {availableDataViews} `, @@ -190,19 +180,19 @@ E.g., for SPL's LEN, use LENGTH. For IF, use CASE.`, ], ['ai', 'Please find the ESQL query below:'], ]).partial({ + documentation, availableDataViews: state.availableDataViews.join('\n'), errors: state.errors.join('\n'), dataViewFields: state.dataViewFields.join('\n'), invalidQueries: state.invalidQueries.join('\n'), }); - const finalChain = answerPrompt.pipe(llm).pipe(new StringOutputParser()); + const finalChain = answerPrompt + .pipe(llm as ActionsClientChatOpenAI | ActionsClientSimpleChatModel) + .pipe(new StringOutputParser()); const finalResult = await finalChain.invoke({ query: userQuery, - date: 'date', - msg: 'msg', - ip: 'ip', }); const correctedResult = finalResult.replaceAll(INLINE_ESQL_QUERY_REGEX, (_match, query) => { @@ -292,16 +282,8 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { chain, dataViews, search, llm } = params; - if (!llm || !chain) return null; - - const retrievalTool = KNOWLEDGE_BASE_RETRIEVAL_TOOL.getTool(params); - - const boundModel = llm as ActionsClientChatOpenAI | ActionsClientSimpleChatModel; - - if (retrievalTool && boundModel?.bindTools) { - boundModel.bindTools([retrievalTool]); - } + const { dataViews, search, llm, esStore } = params; + if (!llm) return null; return new DynamicStructuredTool({ name: toolDetails.name, @@ -314,10 +296,7 @@ export const GRAPH_ESQL_TOOL: AssistantTool = { channels: graphState, }) .addNode('getDataStreams', getDataStreams({ dataViews })) - .addNode( - 'generateQuery', - getGenerateQuery({ userQuery: input.question, llm: boundModel }) - ) + .addNode('generateQuery', getGenerateQuery({ userQuery: input.question, llm, esStore })) .addNode('validateQuery', getValidateQuery({ dataViews, search })) .addEdge(START, 'getDataStreams') .addEdge('getDataStreams', 'generateQuery') From c4d85dc266c2b001548dcbeed3c51c33ae3cc3c7 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 1 Jul 2024 23:26:58 +0200 Subject: [PATCH 44/47] cleanup --- .buildkite/scripts/steps/cloud/deploy.json | 2 +- .../kbn-elastic-assistant-common/impl/capabilities/index.ts | 2 +- .../plugins/security_solution/common/experimental_features.ts | 4 ++-- .../correct_common_esql_mistakes.test.ts | 3 +++ .../correct_common_esql_mistakes.ts | 2 ++ .../esql_language_knowledge_base_tool.ts | 1 - .../tools/knowledge_base/knowledge_base_retrieval_tool.ts | 1 - .../tools/knowledge_base/knowledge_base_write_tool.ts | 1 - 8 files changed, 9 insertions(+), 7 deletions(-) diff --git a/.buildkite/scripts/steps/cloud/deploy.json b/.buildkite/scripts/steps/cloud/deploy.json index 336fed24afbe3..da1918a2a6953 100644 --- a/.buildkite/scripts/steps/cloud/deploy.json +++ b/.buildkite/scripts/steps/cloud/deploy.json @@ -141,7 +141,7 @@ ], "id": "ml", "size": { - "value": 4096, + "value": 1024, "resource": "memory" } } diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index 54c24f6ce7b8f..c1c101fd74cd8 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: true, + assistantKnowledgeBaseByDefault: false, assistantModelEvaluation: false, }); diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index bdcdaca7b9f84..bfeb02db69102 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -125,7 +125,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables the Assistant Knowledge Base by default, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: true, + assistantKnowledgeBaseByDefault: false, /** * Enables the Managed User section inside the new user details flyout. @@ -257,7 +257,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Adds a new langgraph-based ESQL generation tool */ - aiAssistantGraphEsqlTool: true, + aiAssistantGraphEsqlTool: false, }); type ExperimentalConfigKeys = Array; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.test.ts index ad8e0f6cfd664..40dc79851c052 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.test.ts @@ -4,6 +4,9 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + +// Origin: x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_common_esql_mistakes.test.ts + import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; describe('correctCommonEsqlMistakes', () => { diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts index d0ecf0671a234..957338203479b 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/correct_common_esql_mistakes.ts @@ -5,6 +5,8 @@ * 2.0. */ +// Origin: x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_common_esql_mistakes.ts + const STRING_DELIMITER_TOKENS = ['`', "'", '"']; const ESCAPE_TOKEN = '\\\\'; 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 c1cce234309fa..e5448280c5117 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 @@ -19,7 +19,6 @@ const toolDetails = { id: 'esql-knowledge-base-tool', name: 'ESQLKnowledgeBaseTool', }; - export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts index 154f477070ee7..6c9d644c02685 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts @@ -22,7 +22,6 @@ const toolDetails = { id: 'knowledge-base-retrieval-tool', name: 'KnowledgeBaseRetrievalTool', }; - export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts index 5b23ddab5129b..7d5f7e968ba75 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts @@ -23,7 +23,6 @@ const toolDetails = { id: 'knowledge-base-write-tool', name: 'KnowledgeBaseWriteTool', }; - export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { ...toolDetails, sourceRegister: APP_UI_ID, From b4a1168140e291c2fd97fbfaa3d9aeb7158902e1 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 1 Jul 2024 23:31:16 +0200 Subject: [PATCH 45/47] i18n --- .../custom_codeblock/esql_code_block.tsx | 21 ++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx index 62985f5eae137..f8623cdba3af0 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/esql_code_block.tsx @@ -153,6 +153,10 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB const handleShowVisualization = useCallback(() => setShowVisualization(true), []); + const handleToogleTableView = useCallback(() => setIsTableVisible((prev) => !prev), []); + + const handleCloseSaveModal = useCallback(() => setIsSaveModalOpen(false), []); + return ( <> {showVisualization && (queryResultsError || formattedColumnsError || dataViewError) && ( - +

{`${queryResultsError || formattedColumnsError || dataViewError}`}

 @@ -234,7 +247,7 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB iconType={ isTableVisible ? 'visBarVerticalStacked' : 'tableDensityExpanded' } - onClick={() => setIsTableVisible(!isTableVisible)} + onClick={handleToogleTableView} data-test-subj="securityAiAssistantLensESQLDisplayTableButton" aria-label={ isTableVisible @@ -311,9 +324,7 @@ const EsqlCodeBlockComponent = ({ value, actionsDisabled, timestamp }: EsqlCodeB {isSaveModalOpen ? ( { - setIsSaveModalOpen(() => false); - }} + onClose={handleCloseSaveModal} // For now, we don't want to allow saving ESQL charts to the library isSaveable={false} /> From 53b998a4147f85085acfa8ebbb29e0ad9aa913cd Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 2 Jul 2024 00:59:55 +0200 Subject: [PATCH 46/47] fix types --- .../elastic_assistant/server/__mocks__/request_context.ts | 4 ++++ .../lib/langchain/execute_custom_llm_chain/index.test.ts | 4 ++++ .../server/lib/langchain/execute_custom_llm_chain/index.ts | 1 + .../lib/langchain/graphs/default_assistant_graph/helpers.ts | 4 ++-- .../elastic_assistant/server/routes/evaluate/post_evaluate.ts | 2 ++ x-pack/plugins/elastic_assistant/server/types.ts | 2 +- .../esql_language_knowledge_base/graph_esql_language_tool.ts | 4 ++-- 7 files changed, 16 insertions(+), 5 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts index c6a9951e89e3e..78b2a35e53205 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -9,6 +9,8 @@ import { licensingMock } from '@kbn/licensing-plugin/server/mocks'; import { actionsClientMock } from '@kbn/actions-plugin/server/actions_client/actions_client.mock'; import { MockedKeys } from '@kbn/utility-types-jest'; import { AwaitedProperties } from '@kbn/utility-types'; +import { dataViewsService as dataViewsServiceMock } from '@kbn/data-views-plugin/server/mocks'; +import { dataPluginMock } from '@kbn/data-plugin/server/mocks'; import { ElasticAssistantApiRequestHandlerContext, ElasticAssistantRequestHandlerContext, @@ -132,6 +134,8 @@ const createElasticAssistantRequestContextMock = ( getSpaceId: jest.fn(), core: clients.core, telemetry: clients.elasticAssistant.telemetry, + search: dataPluginMock.createRequestHandlerContext().search, + dataViews: dataViewsServiceMock, }; }; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts index e1e8cdc50eee0..c6b50d4d6998f 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts @@ -11,6 +11,8 @@ import { coreMock } from '@kbn/core/server/mocks'; import { KibanaRequest } from '@kbn/core/server'; import { loggerMock } from '@kbn/logging-mocks'; import { initializeAgentExecutorWithOptions } from 'langchain/agents'; +import { dataViewsService as dataViewsServiceMock } from '@kbn/data-views-plugin/server/mocks'; +import { dataPluginMock } from '@kbn/data-plugin/server/mocks'; import { mockActionResponse } from '../../../__mocks__/action_result_data'; import { langChainMessages } from '../../../__mocks__/lang_chain_messages'; @@ -104,6 +106,8 @@ const defaultProps: AgentExecutorParams = { langChainMessages, logger: mockLogger, onNewReplacements: jest.fn(), + search: dataPluginMock.createRequestHandlerContext().search, + dataViews: dataViewsServiceMock, request: mockRequest, replacements: {}, }; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts index 9d4e43be96744..624f2ab9cd52a 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts @@ -96,6 +96,7 @@ export const callAgentExecutor: AgentExecutor = async ({ anonymizationFields, chain, esClient, + esStore, isEnabledKnowledgeBase, llm, logger, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts index 189a6b485784c..a53bcd8f23fad 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts @@ -116,7 +116,7 @@ export const streamGraph = async ({ } } - return processEvent(); + void processEvent(); } catch (err) { // if I throw an error here, it crashes the server. Not sure how to get around that. // If I put await on this function the error works properly, but when there is not an error @@ -134,7 +134,7 @@ export const streamGraph = async ({ }; // Start processing events, do not await! Return `responseWithHeaders` immediately - await processEvent(); + void processEvent(); return responseWithHeaders; }; diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index ef1950b5e90ad..be12b592fcab1 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -203,6 +203,8 @@ export const postEvaluateRoute = ( llmType: 'openai', logger, request: skeletonRequest, + search: assistantContext.search, + dataViews: assistantContext.dataViews, traceOptions: { exampleId, projectName, diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 52971e6791da0..bf9c33d63d6a7 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -226,7 +226,7 @@ export interface AssistantToolParams { isEnabledKnowledgeBase: boolean; chain?: RetrievalQAChain; esClient: ElasticsearchClient; - esStore: ElasticsearchStore; + esStore?: ElasticsearchStore; kbDataClient?: AIAssistantKnowledgeBaseDataClient; langChainTimeout?: number; llm?: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts index 09c98730a4289..bceddef76e8a1 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/graph_esql_language_tool.ts @@ -109,10 +109,10 @@ const getGenerateQuery = }: { userQuery: string; llm: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; - esStore: ElasticsearchStore; + esStore?: ElasticsearchStore; }) => async (state: IState) => { - const knowledgeBaseDocs = await esStore.similaritySearch(userQuery); + const knowledgeBaseDocs = (await esStore?.similaritySearch(userQuery)) ?? []; const documentation = map(knowledgeBaseDocs, 'pageContent').join('\n'); const answerPrompt = await ChatPromptTemplate.fromMessages([ From 45e586bf1590f7150195403d2b6e8f5185ecddf8 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 11 Jul 2024 13:43:09 -0600 Subject: [PATCH 47/47] Up ML node memory for ci-cloud-deploy, and enable assistantKnowledgeBaseByDefault and assistantModelEvaluation FF --- .buildkite/scripts/steps/cloud/deploy.json | 2 +- .../plugins/security_solution/common/experimental_features.ts | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.buildkite/scripts/steps/cloud/deploy.json b/.buildkite/scripts/steps/cloud/deploy.json index da1918a2a6953..336fed24afbe3 100644 --- a/.buildkite/scripts/steps/cloud/deploy.json +++ b/.buildkite/scripts/steps/cloud/deploy.json @@ -141,7 +141,7 @@ ], "id": "ml", "size": { - "value": 1024, + "value": 4096, "resource": "memory" } } diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 1476e29f71505..27b19dc4e91a6 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -120,12 +120,12 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables the Assistant Model Evaluation advanced setting and API endpoint, introduced in `8.11.0`. */ - assistantModelEvaluation: false, + assistantModelEvaluation: true, /** * Enables the Assistant Knowledge Base by default, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, /** * Enables the Managed User section inside the new user details flyout.