From e88a785335207fd837d7d4f417bfac5e497dd595 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 13 Oct 2023 08:23:53 -0600 Subject: [PATCH 01/52] wip --- .../impl/assistant/api.tsx | 33 ++++++++++++++++--- .../stack_connectors/common/openai/schema.ts | 1 + .../stack_connectors/common/openai/types.ts | 1 + .../server/connector_types/openai/openai.ts | 10 ++++-- 4 files changed, 37 insertions(+), 8 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index c7c1254656d61..3641db2217484 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -56,8 +56,8 @@ export const fetchConnectorExecuteAction = async ({ const requestBody = { params: { - subActionParams: body, - subAction: 'invokeAI', + subActionParams: { body: JSON.stringify(body), stream: true }, + subAction: 'stream', }, }; @@ -66,6 +66,28 @@ export const fetchConnectorExecuteAction = async ({ ? `/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute` : `/api/actions/connector/${apiConfig?.connectorId}/_execute`; + // http + // .post<{ + // connector_id: string; + // status: string; + // data: string; + // service_message?: string; + // }>(path, { + // // headers: { + // // 'Content-Type': 'dont-compress-this', + // // }, + // body: JSON.stringify(requestBody), + // signal, + // asResponse: true, + // rawResponse: true, + // }) + // .then((rez) => { + // console.log('it worked', rez); + // }) + // .catch((err) => { + // console.log('it broke', err); + // }); + const response = await http.fetch<{ connector_id: string; status: string; @@ -73,13 +95,14 @@ export const fetchConnectorExecuteAction = async ({ service_message?: string; }>(path, { method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, body: JSON.stringify(requestBody), signal, + // asResponse: true, + // rawResponse: true, }); + console.log('response', response); + if (response.status !== 'ok' || !response.data) { if (response.service_message) { return { diff --git a/x-pack/plugins/stack_connectors/common/openai/schema.ts b/x-pack/plugins/stack_connectors/common/openai/schema.ts index fa14aa61fa5b3..b37427191b48b 100644 --- a/x-pack/plugins/stack_connectors/common/openai/schema.ts +++ b/x-pack/plugins/stack_connectors/common/openai/schema.ts @@ -42,6 +42,7 @@ export const InvokeAIActionParamsSchema = schema.object({ schema.nullable(schema.oneOf([schema.string(), schema.arrayOf(schema.string())])) ), temperature: schema.maybe(schema.number()), + stream: schema.boolean({ defaultValue: false }), }); export const InvokeAIActionResponseSchema = schema.string(); diff --git a/x-pack/plugins/stack_connectors/common/openai/types.ts b/x-pack/plugins/stack_connectors/common/openai/types.ts index 86e2c172846dc..4c78b86a4a99a 100644 --- a/x-pack/plugins/stack_connectors/common/openai/types.ts +++ b/x-pack/plugins/stack_connectors/common/openai/types.ts @@ -23,6 +23,7 @@ export type Secrets = TypeOf; export type RunActionParams = TypeOf; export type InvokeAIActionParams = TypeOf; export type InvokeAIActionResponse = TypeOf; +export type StreamingResponse = TypeOf; export type RunActionResponse = TypeOf; export type DashboardActionParams = TypeOf; export type DashboardActionResponse = TypeOf; diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts index 21c7bc4abdcc0..51036f1514035 100644 --- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts +++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts @@ -27,7 +27,6 @@ import { DashboardActionParams, DashboardActionResponse, InvokeAIActionParams, - InvokeAIActionResponse, } from '../../../common/openai/types'; import { initDashboard } from './create_dashboard'; import { @@ -128,6 +127,7 @@ export class OpenAIConnector extends SubActionConnector { * @param stream flag indicating whether it is a streaming request or not */ public async streamApi({ body, stream }: StreamActionParams): Promise { + console.log('START STREAM'); const executeBody = getRequestWithStreamOption( this.provider, this.url, @@ -144,6 +144,7 @@ export class OpenAIConnector extends SubActionConnector { data: executeBody, ...axiosOptions, }); + console.log('END STREAM'); return stream ? pipeStreamingResponse(response) : response.data; } @@ -187,13 +188,16 @@ export class OpenAIConnector extends SubActionConnector { * Sends the stringified input to the runApi method. Returns the trimmed completion from the response. * @param body An object containing array of message objects, and possible other OpenAI properties */ - public async invokeAI(body: InvokeAIActionParams): Promise { - const res = await this.runApi({ body: JSON.stringify(body) }); + public async invokeAI(params: InvokeAIActionParams): Promise { + const { stream, ...body } = params; + const res = await this.streamApi({ body: JSON.stringify(body), stream }); if (res.choices && res.choices.length > 0 && res.choices[0].message?.content) { const result = res.choices[0].message.content.trim(); return result; } + console.log('returning now', res); + return res; return 'An error occurred sending your message. \n\nAPI Error: The response from OpenAI was in an unrecognized format.'; } From 556e4bdf99fab3e3f3cd3d81f437ec52eea3c1a6 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 13 Oct 2023 10:48:32 -0600 Subject: [PATCH 02/52] more wip --- .../impl/assistant/api.tsx | 33 +++--------------- .../server/lib/langchain/execute.ts | 34 +++++++++++++++++++ .../lib/langchain/llm/actions_client_llm.ts | 4 +++ .../routes/post_actions_connector_execute.ts | 9 +++++ .../schemas/post_actions_connector_execute.ts | 1 + 5 files changed, 52 insertions(+), 29 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index 3641db2217484..6cd74e75e5b6c 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -56,44 +56,19 @@ export const fetchConnectorExecuteAction = async ({ const requestBody = { params: { - subActionParams: { body: JSON.stringify(body), stream: true }, - subAction: 'stream', + subActionParams: { ...body, stream: true }, + subAction: 'invokeAI', }, + assistantLangChain, }; try { - const path = assistantLangChain - ? `/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute` - : `/api/actions/connector/${apiConfig?.connectorId}/_execute`; - - // http - // .post<{ - // connector_id: string; - // status: string; - // data: string; - // service_message?: string; - // }>(path, { - // // headers: { - // // 'Content-Type': 'dont-compress-this', - // // }, - // body: JSON.stringify(requestBody), - // signal, - // asResponse: true, - // rawResponse: true, - // }) - // .then((rez) => { - // console.log('it worked', rez); - // }) - // .catch((err) => { - // console.log('it broke', err); - // }); - const response = await http.fetch<{ connector_id: string; status: string; data: string; service_message?: string; - }>(path, { + }>(`/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`, { method: 'POST', body: JSON.stringify(requestBody), signal, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts new file mode 100644 index 0000000000000..8537f971199cc --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts @@ -0,0 +1,34 @@ +/* + * 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 { get } from 'lodash/fp'; +import { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server'; +import { KibanaRequest } from '@kbn/core-http-server'; +import { RequestBody } from './types'; + +interface Props { + actions: ActionsPluginStart; + connectorId: string; + request: KibanaRequest; +} + +export const executeAction = async ({ actions, request, connectorId }: Props) => { + const actionsClient = await actions.getActionsClientWithRequest(request); + const actionResult = await actionsClient.execute({ + actionId: connectorId, + params: request.body.params, + }); + const content = get('data', actionResult); + // if (typeof content !== 'string') { + // throw new Error(`content should be a string, but it had an unexpected type: ${typeof content}`); + // } + return { + connector_id: connectorId, + data: content, // the response from the actions framework + status: 'ok', + }; +}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.ts index e4403b64d6e0d..ebc3965ab4fdb 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.ts @@ -65,6 +65,7 @@ export class ActionsClientLlm extends LLM { } async _call(prompt: string): Promise { + console.log('CALLL'); // convert the Langchain prompt to an assistant message: const assistantMessage = getMessageContentAndRole(prompt); this.#logger.debug( @@ -85,8 +86,11 @@ export class ActionsClientLlm extends LLM { // create an actions client from the authenticated request context: const actionsClient = await this.#actions.getActionsClientWithRequest(this.#request); + console.log('NOW WHAT?????', requestBody); const actionResult = await actionsClient.execute(requestBody); + console.log('NOW WHAT?????', actionResult); + if (actionResult.status === 'error') { throw new Error( `${LLM_TYPE}: action result status is error: ${actionResult?.message} - ${actionResult?.serviceMessage}` 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 8f620dac06faa..905388b276ed4 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 @@ -7,6 +7,7 @@ import { IRouter, Logger } from '@kbn/core/server'; import { transformError } from '@kbn/securitysolution-es-utils'; +import { executeAction } from '../lib/langchain/execute'; import { POST_ACTIONS_CONNECTOR_EXECUTE } from '../../common/constants'; import { getLangChainMessages } from '../lib/langchain/helpers'; import { buildResponse } from '../lib/build_response'; @@ -42,6 +43,14 @@ export const postActionsConnectorExecuteRoute = ( // get a scoped esClient for assistant memory const esClient = (await context.core).elasticsearch.client.asCurrentUser; + if (!request.body.assistantLangChain) { + const result = await executeAction({ actions, request, connectorId }); + console.log('this is the reader', result.data); + + return response.ok({ + body: result, + }); + } // convert the assistant messages to LangChain messages: const langChainMessages = getLangChainMessages( diff --git a/x-pack/plugins/elastic_assistant/server/schemas/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/schemas/post_actions_connector_execute.ts index b30ccd94e105b..7a8d52e725722 100644 --- a/x-pack/plugins/elastic_assistant/server/schemas/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/schemas/post_actions_connector_execute.ts @@ -34,6 +34,7 @@ export const PostActionsConnectorExecuteBody = t.type({ ]), subAction: t.string, }), + assistantLangChain: t.boolean, }); export type PostActionsConnectorExecuteBodyInputs = t.TypeOf< From bed104be8a1cc39a350dbbff9aa4539e4cf24a8b Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 16 Oct 2023 09:43:14 -0600 Subject: [PATCH 03/52] wip --- .../impl/assistant/api.tsx | 38 ++-- .../impl/assistant/chat_send/chat_stream.tsx | 168 ++++++++++++++++++ .../assistant/chat_send/use_chat_send.tsx | 11 +- .../impl/assistant/use_conversation/index.tsx | 49 ++++- .../server/lib/langchain/execute.ts | 32 +++- .../routes/post_actions_connector_execute.ts | 8 +- 6 files changed, 280 insertions(+), 26 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/chat_stream.tsx diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index 6cd74e75e5b6c..c3a1154f2280a 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -9,6 +9,7 @@ import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/public/common'; import { HttpSetup, IHttpFetchError } from '@kbn/core-http-browser'; +import { IncomingMessage } from 'http'; import type { Conversation, Message } from '../assistant_context/types'; import { API_ERROR } from './translations'; import { MODEL_GPT_3_5_TURBO } from '../connectorland/models/model_selector/model_selector'; @@ -53,30 +54,45 @@ export const fetchConnectorExecuteAction = async ({ // Azure OpenAI and Bedrock invokeAI both expect this body format messages: outboundMessages, }; - + const isStream = true; const requestBody = { params: { - subActionParams: { ...body, stream: true }, + subActionParams: { ...body, stream: isStream }, subAction: 'invokeAI', }, assistantLangChain, }; try { - const response = await http.fetch<{ - connector_id: string; - status: string; - data: string; - service_message?: string; - }>(`/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`, { + const response = await http.fetch< + | { + connector_id: string; + status: string; + data: string; + service_message?: string; + } + | IncomingMessage + >(`/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`, { method: 'POST', body: JSON.stringify(requestBody), signal, - // asResponse: true, - // rawResponse: true, + asResponse: isStream, + rawResponse: isStream, }); - console.log('response', response); + if (isStream) { + const reader = response?.response?.body?.getReader(); + + console.log('is typeof incoming message'); + if (!reader) { + throw new Error('Could not get reader from response'); + } + return { + response: reader, + isStream: true, + isError: false, + }; + } if (response.status !== 'ok' || !response.data) { if (response.service_message) { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/chat_stream.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/chat_stream.tsx new file mode 100644 index 0000000000000..03603b3b43a2f --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/chat_stream.tsx @@ -0,0 +1,168 @@ +/* + * 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 { + BehaviorSubject, + catchError, + concatMap, + delay, + filter as rxJsFilter, + finalize, + map, + Observable, + of, + scan, + share, + shareReplay, + tap, + timestamp, +} from 'rxjs'; +import { cloneDeep } from 'lodash'; +import { AbortError } from '@kbn/kibana-utils-plugin/common'; +import { ConversationRole, Message } from '../../assistant_context/types'; + +export function readableStreamReaderIntoObservable( + readableStreamReader: ReadableStreamDefaultReader +): Observable { + return new Observable((subscriber) => { + let lineBuffer: string = ''; + + async function read(): Promise { + const { done, value } = await readableStreamReader.read(); + if (done) { + if (lineBuffer) { + subscriber.next(lineBuffer); + } + subscriber.complete(); + + return; + } + + const textChunk = new TextDecoder().decode(value); + + const lines = textChunk.split('\n'); + lines[0] = lineBuffer + lines[0]; + + lineBuffer = lines.pop() || ''; + + for (const line of lines) { + subscriber.next(line); + } + + return read(); + } + + read().catch((err) => subscriber.error(err)); + + return () => { + readableStreamReader.cancel().catch(() => {}); + }; + }).pipe(share()); +} + +const role: ConversationRole = 'assistant'; +export const chatStream = (reader: ReadableStreamDefaultReader) => { + const subject = new BehaviorSubject<{ message: Message }>({ + message: { role }, + }); + readableStreamReaderIntoObservable(reader).pipe( + // lines start with 'data: ' + map((line) => line.substring(6)), + // a message completes with the line '[DONE]' + rxJsFilter((line) => !!line && line !== '[DONE]'), + // parse the JSON, add the type + map((line) => JSON.parse(line) as {} | { error: { message: string } }), + // validate the message. in some cases OpenAI + // will throw halfway through the message + tap((line) => { + if ('error' in line) { + throw new Error(line.error.message); + } + }), + // there also might be some metadata that we need + // to exclude + rxJsFilter( + (line): line is { object: string } => + 'object' in line && line.object === 'chat.completion.chunk' + ), + // this is how OpenAI signals that the context window + // limit has been exceeded + tap((line) => { + if (line.choices[0].finish_reason === 'length') { + throw new Error(`Token limit reached`); + } + }), + // merge the messages + scan( + (acc, { choices }) => { + acc.message.content += choices[0].delta.content ?? ''; + acc.message.function_call.name += choices[0].delta.function_call?.name ?? ''; + acc.message.function_call.arguments += choices[0].delta.function_call?.arguments ?? ''; + return cloneDeep(acc); + }, + { + message: { + content: '', + function_call: { + name: '', + arguments: '', + trigger: 'assistant' as const, + }, + role: 'assistant', + }, + } + ), + // convert an error into state + catchError((error) => + of({ + ...subject.value, + error, + aborted: error instanceof AbortError, + }) + ) + ); + const MIN_DELAY = 35; + + const pendingMessages$ = subject.pipe( + // make sure the request is only triggered once, + // even with multiple subscribers + shareReplay(1), + // if the Observable is no longer subscribed, + // abort the running request + finalize(() => { + // controller.abort(); + }), + // append a timestamp of when each value was emitted + timestamp(), + // use the previous timestamp to calculate a target + // timestamp for emitting the next value + scan((acc, value) => { + const lastTimestamp = acc.timestamp || 0; + const emitAt = Math.max(lastTimestamp + MIN_DELAY, value.timestamp); + return { + timestamp: emitAt, + value: value.value, + }; + }), + // add the delay based on the elapsed time + // using concatMap(of(value).pipe(delay(50)) + // leads to browser issues because timers + // are throttled when the tab is not active + concatMap((value) => { + const now = Date.now(); + const delayFor = value.timestamp - now; + + if (delayFor <= 0) { + return of(value.value); + } + + return of(value.value).pipe(delay(delayFor)); + }) + ); + + return pendingMessages$; +}; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index 3e1c194097888..637c4f899f82d 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -54,7 +54,8 @@ export const useChatSend = ({ setUserPrompt, }: UseChatSendProps): UseChatSend => { const { isLoading, sendMessages } = useSendMessages(); - const { appendMessage, appendReplacements, clearConversation } = useConversation(); + const { appendMessage, appendReplacements, appendStreamMessage, clearConversation } = + useConversation(); const handlePromptChange = (prompt: string) => { setPromptTextPreview(prompt); @@ -95,6 +96,13 @@ export const useChatSend = ({ apiConfig: currentConversation.apiConfig, messages: updatedMessages, }); + if (rawResponse.isStream) { + console.log('appendStreamMessage bout to call', rawResponse); + return appendStreamMessage({ + conversationId: currentConversation.id, + reader: rawResponse.response, + }); + } const responseMessage: Message = getMessageFromRawResponse(rawResponse); appendMessage({ conversationId: currentConversation.id, message: responseMessage }); }, @@ -109,6 +117,7 @@ export const useChatSend = ({ http, appendReplacements, editingSystemPromptId, + appendStreamMessage, ] ); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx index 0cf4a4bdc9439..9536d157fb795 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx @@ -5,8 +5,9 @@ * 2.0. */ -import { useCallback } from 'react'; +import { useCallback, useEffect, useState } from 'react'; +import { chatStream } from '../chat_send/chat_stream'; import { useAssistantContext } from '../../assistant_context'; import { Conversation, Message } from '../../assistant_context/types'; import * as i18n from './translations'; @@ -56,6 +57,13 @@ interface SetConversationProps { } interface UseConversation { + appendStreamMessage: ({ + conversationId, + reader, + }: { + conversationId: string; + reader: ReadableStream; + }) => void; appendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[]; appendReplacements: ({ conversationId, @@ -71,6 +79,44 @@ interface UseConversation { export const useConversation = (): UseConversation => { const { allSystemPrompts, assistantTelemetry, setConversations } = useAssistantContext(); + const [pendingMessage, setPendingMessage] = useState(); + useEffect(() => { + console.log('pendingMessage', pendingMessage); + }, [pendingMessage]); + const appendStreamMessage = useCallback(({ conversationId, reader }) => { + console.log('appending appendStreamMessage'); + let lastPendingMessage; + chatStream(reader).subscribe({ + next: (msg) => { + lastPendingMessage = msg; + setPendingMessage(() => msg); + }, + complete: () => { + setPendingMessage(lastPendingMessage); + }, + }); + console.log('lastPendingMessage', lastPendingMessage); + // assistantTelemetry?.reportAssistantMessageSent({ conversationId, role: message.role }); + // let messages: Message[] = []; + // setConversations((prev: Record) => { + // const prevConversation: Conversation | undefined = prev[conversationId]; + // + // if (prevConversation != null) { + // messages = [...prevConversation.messages, message]; + // const newConversation = { + // ...prevConversation, + // messages, + // }; + // return { + // ...prev, + // [conversationId]: newConversation, + // }; + // } else { + // return prev; + // } + // }); + // return messages; + }, []); /** * Append a message to the conversation[] for a given conversationId */ @@ -264,6 +310,7 @@ export const useConversation = (): UseConversation => { return { appendMessage, appendReplacements, + appendStreamMessage, clearConversation, createConversation, deleteConversation, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts index 8537f971199cc..8467458d3cf58 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts @@ -8,6 +8,7 @@ import { get } from 'lodash/fp'; import { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server'; import { KibanaRequest } from '@kbn/core-http-server'; +import { IncomingMessage } from 'http'; import { RequestBody } from './types'; interface Props { @@ -15,20 +16,33 @@ interface Props { connectorId: string; request: KibanaRequest; } +interface StaticResponse { + connector_id: string; + data: string; + status: string; +} -export const executeAction = async ({ actions, request, connectorId }: Props) => { +export const executeAction = async ({ + actions, + request, + connectorId, +}: Props): Promise => { const actionsClient = await actions.getActionsClientWithRequest(request); const actionResult = await actionsClient.execute({ actionId: connectorId, params: request.body.params, }); const content = get('data', actionResult); - // if (typeof content !== 'string') { - // throw new Error(`content should be a string, but it had an unexpected type: ${typeof content}`); - // } - return { - connector_id: connectorId, - data: content, // the response from the actions framework - status: 'ok', - }; + if (typeof content === 'string') { + return { + connector_id: connectorId, + data: content, // the response from the actions framework + status: 'ok', + }; + } + // is data stream + if (content instanceof IncomingMessage) { + return content; + } + throw new Error('Unexpected action result'); }; 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 905388b276ed4..2fea4770b0ccf 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 @@ -40,18 +40,18 @@ export const postActionsConnectorExecuteRoute = ( // get the actions plugin start contract from the request context: const actions = (await context.elasticAssistant).actions; - - // get a scoped esClient for assistant memory - const esClient = (await context.core).elasticsearch.client.asCurrentUser; + // if not langchain, call execute action directly and return the response: if (!request.body.assistantLangChain) { const result = await executeAction({ actions, request, connectorId }); - console.log('this is the reader', result.data); return response.ok({ body: result, }); } + // get a scoped esClient for assistant memory + const esClient = (await context.core).elasticsearch.client.asCurrentUser; + // convert the assistant messages to LangChain messages: const langChainMessages = getLangChainMessages( request.body.params.subActionParams.messages From 939afd72a20288255a4b74f0651324913699122a Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 16 Oct 2023 11:26:12 -0600 Subject: [PATCH 04/52] rm logs --- .../server/lib/langchain/llm/actions_client_llm.ts | 4 ---- .../server/connector_types/openai/openai.ts | 9 +++++---- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.ts index ebc3965ab4fdb..e4403b64d6e0d 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.ts @@ -65,7 +65,6 @@ export class ActionsClientLlm extends LLM { } async _call(prompt: string): Promise { - console.log('CALLL'); // convert the Langchain prompt to an assistant message: const assistantMessage = getMessageContentAndRole(prompt); this.#logger.debug( @@ -86,11 +85,8 @@ export class ActionsClientLlm extends LLM { // create an actions client from the authenticated request context: const actionsClient = await this.#actions.getActionsClientWithRequest(this.#request); - console.log('NOW WHAT?????', requestBody); const actionResult = await actionsClient.execute(requestBody); - console.log('NOW WHAT?????', actionResult); - if (actionResult.status === 'error') { throw new Error( `${LLM_TYPE}: action result status is error: ${actionResult?.message} - ${actionResult?.serviceMessage}` diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts index 51036f1514035..cb2cc5883439e 100644 --- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts +++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts @@ -27,6 +27,8 @@ import { DashboardActionParams, DashboardActionResponse, InvokeAIActionParams, + InvokeAIActionResponse, + StreamingResponse, } from '../../../common/openai/types'; import { initDashboard } from './create_dashboard'; import { @@ -127,7 +129,6 @@ export class OpenAIConnector extends SubActionConnector { * @param stream flag indicating whether it is a streaming request or not */ public async streamApi({ body, stream }: StreamActionParams): Promise { - console.log('START STREAM'); const executeBody = getRequestWithStreamOption( this.provider, this.url, @@ -144,7 +145,6 @@ export class OpenAIConnector extends SubActionConnector { data: executeBody, ...axiosOptions, }); - console.log('END STREAM'); return stream ? pipeStreamingResponse(response) : response.data; } @@ -188,7 +188,9 @@ export class OpenAIConnector extends SubActionConnector { * Sends the stringified input to the runApi method. Returns the trimmed completion from the response. * @param body An object containing array of message objects, and possible other OpenAI properties */ - public async invokeAI(params: InvokeAIActionParams): Promise { + public async invokeAI( + params: InvokeAIActionParams + ): Promise { const { stream, ...body } = params; const res = await this.streamApi({ body: JSON.stringify(body), stream }); @@ -196,7 +198,6 @@ export class OpenAIConnector extends SubActionConnector { const result = res.choices[0].message.content.trim(); return result; } - console.log('returning now', res); return res; return 'An error occurred sending your message. \n\nAPI Error: The response from OpenAI was in an unrecognized format.'; From 693fcae1382f1099e38badcc6897b9833e9f58ed Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 16 Oct 2023 13:43:04 -0600 Subject: [PATCH 05/52] wip --- .../impl/assistant/api.tsx | 52 +++++---- .../assistant/chat_send/use_chat_send.tsx | 8 -- .../impl/assistant/helpers.ts | 6 +- .../impl/assistant/use_conversation/index.tsx | 60 ++++------- .../impl/assistant_context/types.tsx | 3 +- .../public/assistant/get_comments/index.tsx | 100 +++++++++++------- 6 files changed, 120 insertions(+), 109 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index c3a1154f2280a..44dde91e2a602 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -8,7 +8,8 @@ import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/public/common'; import { HttpSetup, IHttpFetchError } from '@kbn/core-http-browser'; - +// TODO: Why do i get this error here? its imported in other places without issue +// eslint-disable-next-line import/no-nodejs-modules import { IncomingMessage } from 'http'; import type { Conversation, Message } from '../assistant_context/types'; import { API_ERROR } from './translations'; @@ -25,8 +26,9 @@ export interface FetchConnectorExecuteAction { } export interface FetchConnectorExecuteResponse { - response: string; + response: string | ReadableStreamDefaultReader; isError: boolean; + isStream: boolean; } export const fetchConnectorExecuteAction = async ({ @@ -64,28 +66,27 @@ export const fetchConnectorExecuteAction = async ({ }; try { - const response = await http.fetch< - | { - connector_id: string; - status: string; - data: string; - service_message?: string; + if (isStream) { + const response = await http.fetch( + `/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`, + { + method: 'POST', + body: JSON.stringify(requestBody), + signal, + asResponse: isStream, + rawResponse: isStream, } - | IncomingMessage - >(`/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`, { - method: 'POST', - body: JSON.stringify(requestBody), - signal, - asResponse: isStream, - rawResponse: isStream, - }); + ); - if (isStream) { const reader = response?.response?.body?.getReader(); console.log('is typeof incoming message'); if (!reader) { - throw new Error('Could not get reader from response'); + return { + response: `${API_ERROR}\n\nCould not get reader from response`, + isError: true, + isStream: false, + }; } return { response: reader, @@ -94,26 +95,41 @@ export const fetchConnectorExecuteAction = async ({ }; } + const response = await http.fetch<{ + connector_id: string; + status: string; + data: string; + service_message?: string; + }>(`/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`, { + method: 'POST', + body: JSON.stringify(requestBody), + signal, + }); + if (response.status !== 'ok' || !response.data) { if (response.service_message) { return { response: `${API_ERROR}\n\n${response.service_message}`, isError: true, + isStream: false, }; } return { response: API_ERROR, isError: true, + isStream: false, }; } return { response: assistantLangChain ? getFormattedMessageContent(response.data) : response.data, isError: false, + isStream: false, }; } catch (error) { return { response: API_ERROR, isError: true, + isStream: false, }; } }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index 637c4f899f82d..53087e577dded 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -96,13 +96,6 @@ export const useChatSend = ({ apiConfig: currentConversation.apiConfig, messages: updatedMessages, }); - if (rawResponse.isStream) { - console.log('appendStreamMessage bout to call', rawResponse); - return appendStreamMessage({ - conversationId: currentConversation.id, - reader: rawResponse.response, - }); - } const responseMessage: Message = getMessageFromRawResponse(rawResponse); appendMessage({ conversationId: currentConversation.id, message: responseMessage }); }, @@ -117,7 +110,6 @@ export const useChatSend = ({ http, appendReplacements, editingSystemPromptId, - appendStreamMessage, ] ); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.ts index b4eb89a092600..61001d95e8a3e 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.ts @@ -12,12 +12,14 @@ import type { Message } from '../assistant_context/types'; import { enterpriseMessaging, WELCOME_CONVERSATION } from './use_conversation/sample_conversations'; export const getMessageFromRawResponse = (rawResponse: FetchConnectorExecuteResponse): Message => { - const { response, isError } = rawResponse; + const { response, isStream, isError } = rawResponse; const dateTimeString = new Date().toLocaleString(); // TODO: Pull from response if (rawResponse) { return { role: 'assistant', - content: response, + ...(isStream + ? { reader: response as ReadableStreamDefaultReader } + : { content: response as string }), timestamp: dateTimeString, isError, }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx index 9536d157fb795..659b6b5c7de70 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx @@ -7,7 +7,6 @@ import { useCallback, useEffect, useState } from 'react'; -import { chatStream } from '../chat_send/chat_stream'; import { useAssistantContext } from '../../assistant_context'; import { Conversation, Message } from '../../assistant_context/types'; import * as i18n from './translations'; @@ -57,13 +56,7 @@ interface SetConversationProps { } interface UseConversation { - appendStreamMessage: ({ - conversationId, - reader, - }: { - conversationId: string; - reader: ReadableStream; - }) => void; + appendStreamMessage: ({ conversationId, message }: AppendMessageProps) => void; appendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[]; appendReplacements: ({ conversationId, @@ -83,39 +76,28 @@ export const useConversation = (): UseConversation => { useEffect(() => { console.log('pendingMessage', pendingMessage); }, [pendingMessage]); + const appendStreamMessage = useCallback(({ conversationId, reader }) => { - console.log('appending appendStreamMessage'); - let lastPendingMessage; - chatStream(reader).subscribe({ - next: (msg) => { - lastPendingMessage = msg; - setPendingMessage(() => msg); - }, - complete: () => { - setPendingMessage(lastPendingMessage); - }, + // assistantTelemetry?.reportAssistantMessageSent({ conversationId, role }); + let messages: Message[] = []; + setConversations((prev: Record) => { + const prevConversation: Conversation | undefined = prev[conversationId]; + + if (prevConversation != null) { + messages = [...prevConversation.messages, { reader, role: 'assistant' }]; + const newConversation = { + ...prevConversation, + messages, + }; + return { + ...prev, + [conversationId]: newConversation, + }; + } else { + return prev; + } }); - console.log('lastPendingMessage', lastPendingMessage); - // assistantTelemetry?.reportAssistantMessageSent({ conversationId, role: message.role }); - // let messages: Message[] = []; - // setConversations((prev: Record) => { - // const prevConversation: Conversation | undefined = prev[conversationId]; - // - // if (prevConversation != null) { - // messages = [...prevConversation.messages, message]; - // const newConversation = { - // ...prevConversation, - // messages, - // }; - // return { - // ...prev, - // [conversationId]: newConversation, - // }; - // } else { - // return prev; - // } - // }); - // return messages; + return messages; }, []); /** * Append a message to the conversation[] for a given conversationId diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx index 651eeee17f21e..f9ade4b9abc1e 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx @@ -15,7 +15,8 @@ export interface MessagePresentation { } export interface Message { role: ConversationRole; - content: string; + reader?: ReadableStreamDefaultReader; + content?: string; timestamp: string; isError?: boolean; presentation?: MessagePresentation; 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 d91caae855d22..cc6a626a5496f 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 @@ -13,6 +13,7 @@ import React from 'react'; import { AssistantAvatar } from '@kbn/elastic-assistant'; import { css } from '@emotion/react/dist/emotion-react.cjs'; import { euiThemeVars } from '@kbn/ui-theme'; +import type { EuiPanelProps } from '@elastic/eui/src/components/panel'; import { CommentActions } from '../comment_actions'; import * as i18n from './translations'; @@ -28,35 +29,20 @@ export const getComments = ({ currentConversation.messages.map((message, index) => { const isUser = message.role === 'user'; const replacements = currentConversation.replacements; - const messageContentWithReplacements = - replacements != null - ? Object.keys(replacements).reduce( - (acc, replacement) => acc.replaceAll(replacement, replacements[replacement]), - message.content - ) - : message.content; - const transformedMessage = { - ...message, - content: messageContentWithReplacements, + const errorStyles = { + eventColor: 'danger' as EuiPanelProps['color'], + css: css` + .euiCommentEvent { + border: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; + } + .euiCommentEvent__header { + padding: 0 !important; + border-block-end: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; + } + `, }; - return { - actions: , - children: - index !== currentConversation.messages.length - 1 ? ( - - - {showAnonymizedValues ? message.content : transformedMessage.content} - - - ) : ( - - - {showAnonymizedValues ? message.content : transformedMessage.content} - - - - ), + const messageProps = { timelineAvatar: isUser ? ( ) : ( @@ -66,19 +52,51 @@ export const getComments = ({ message.timestamp.length === 0 ? new Date().toLocaleString() : message.timestamp ), username: isUser ? i18n.YOU : i18n.ASSISTANT, - ...(message.isError - ? { - eventColor: 'danger', - css: css` - .euiCommentEvent { - border: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; - } - .euiCommentEvent__header { - padding: 0 !important; - border-block-end: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; - } - `, - } - : {}), + ...(message.isError ? errorStyles : {}), + }; + + if (message.content && message.content.length) { + const messageContentWithReplacements = + replacements != null + ? Object.keys(replacements).reduce( + (acc, replacement) => acc.replaceAll(replacement, replacements[replacement]), + message.content + ) + : message.content; + const transformedMessage = { + ...message, + content: messageContentWithReplacements, + }; + + return { + ...messageProps, + actions: , + children: + index !== currentConversation.messages.length - 1 ? ( + + + {showAnonymizedValues ? message.content : transformedMessage.content} + + + ) : ( + + + {showAnonymizedValues ? message.content : transformedMessage.content} + + + + ), + }; + } + if (message.reader) { + return { + ...messageProps, + children: {'hello world i need to make this a stream here dawg'}, + }; + } + return { + ...messageProps, + children: {'oh no an error happened'}, + ...errorStyles, }; }); From f58a5e17770fffc2684bb5d6b97dab825ff47678 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 16 Oct 2023 16:10:30 -0600 Subject: [PATCH 06/52] its streaming baby! --- .../public/assistant/get_comments/index.tsx | 10 +- .../public/assistant/get_comments/stream.tsx | 45 ++++ .../assistant/get_comments/stream_obs.tsx | 207 ++++++++++++++++++ .../public/assistant/helpers.tsx | 2 +- 4 files changed, 262 insertions(+), 2 deletions(-) create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream_obs.tsx 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 cc6a626a5496f..739804792c9f1 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 @@ -14,6 +14,7 @@ import { AssistantAvatar } from '@kbn/elastic-assistant'; import { css } from '@emotion/react/dist/emotion-react.cjs'; import { euiThemeVars } from '@kbn/ui-theme'; import type { EuiPanelProps } from '@elastic/eui/src/components/panel'; +import { StreamComment } from './stream_obs'; import { CommentActions } from '../comment_actions'; import * as i18n from './translations'; @@ -91,7 +92,14 @@ export const getComments = ({ if (message.reader) { return { ...messageProps, - children: {'hello world i need to make this a stream here dawg'}, + children: ( + + ), }; } return { diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx new file mode 100644 index 0000000000000..53fdcafa0a2d5 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream.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 React, { useEffect, useState } from 'react'; + +export const StreamComment = ({ reader }: { reader: ReadableStreamDefaultReader }) => { + const [data, setData] = useState(''); + const [loading, setLoading] = useState(false); + + useEffect(() => { + const fetchData = async () => { + try { + setLoading(true); + const decoder = new TextDecoder(); + + while (true) { + const { value, done } = await reader.read(); + if (done) { + setLoading(false); + break; + } + + const decodedChunk = decoder.decode(value, { stream: true }); + setData((prevValue: string) => `${prevValue}${decodedChunk}`); + } + } catch (error) { + setLoading(false); + // Handle other errors + } + }; + + fetchData(); + }, [reader]); + + return ( + + {loading && {'Fetching response...'}} + {data} + + ); +}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream_obs.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream_obs.tsx new file mode 100644 index 0000000000000..5b520c596a844 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream_obs.tsx @@ -0,0 +1,207 @@ +/* + * 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 React, { useEffect, useMemo } from 'react'; +import { + EuiFlexGroup, + EuiFlexItem, + EuiText, + EuiLoadingSpinner, + EuiIcon, + EuiMarkdownFormat, +} from '@elastic/eui'; +import { i18n } from '@kbn/i18n'; +import { concatMap, delay, Observable, of } from 'rxjs'; +import useObservable from 'react-use/lib/useObservable'; + +export interface PromptObservableState { + chunks: Chunk[]; + message?: string; + error?: string; + loading: boolean; +} + +interface ChunkChoice { + index: 0; + delta: { role: string; content: string }; + finish_reason: null | string; +} + +interface Chunk { + id: string; + object: string; + created: number; + model: string; + choices: ChunkChoice[]; +} + +const cursorCss = ` + @keyframes blink { + 0% { + opacity: 1; + } + 50% { + opacity: 0; + } + 100% { + opacity: 1; + } + } + + animation: blink 1s infinite; + width: 10px; + height: 16px; + vertical-align: middle; + display: inline-block; + background: rgba(0, 0, 0, 0.25); +`; + +function getMessageFromChunks(chunks: Chunk[]) { + let message = ''; + chunks.forEach((chunk) => { + message += chunk.choices[0]?.delta.content ?? ''; + }); + return message; +} + +interface Props { + index: number; + isLastComment: boolean; + lastCommentRef: React.MutableRefObject; + reader: ReadableStreamDefaultReader; +} + +export const StreamComment = ({ index, isLastComment, lastCommentRef, reader }: Props) => { + const response$ = useMemo( + () => + new Observable((observer) => { + observer.next({ chunks: [], loading: true }); + + const decoder = new TextDecoder(); + + const chunks: Chunk[] = []; + + let prev: string = ''; + + function read() { + reader.read().then(({ done, value }: { done: boolean; value?: Uint8Array }) => { + try { + if (done) { + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: false, + }); + observer.complete(); + return; + } + + let lines: string[] = (prev + decoder.decode(value)).split('\n'); + + const lastLine: string = lines[lines.length - 1]; + + const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; + + if (isPartialChunk) { + prev = lastLine; + lines.pop(); + } else { + prev = ''; + } + + lines = lines.map((str) => str.substr(6)).filter((str) => !!str && str !== '[DONE]'); + + const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); + + nextChunks.forEach((chunk) => { + chunks.push(chunk); + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: true, + }); + }); + } catch (err) { + observer.error(err); + return; + } + read(); + }); + } + + read(); + + return () => { + reader.cancel(); + }; + }).pipe(concatMap((value) => of(value).pipe(delay(50)))), + [reader] + ); + + const response = useObservable(response$); + + useEffect(() => {}, [response$]); + + let content = response?.message ?? ''; + + let state: 'init' | 'loading' | 'streaming' | 'error' | 'complete' = 'init'; + + if (response?.loading) { + state = content ? 'streaming' : 'loading'; + } else if (response && 'error' in response && response.error) { + state = 'error'; + content = response.error; + } else if (content) { + state = 'complete'; + } + + let inner: React.ReactNode; + + if (state === 'complete' || state === 'streaming') { + inner = ( + <> + + {content} + + {state === 'streaming' ? : <>} + + ); + } else if (state === 'init' || state === 'loading') { + inner = ( + + + + + + + {i18n.translate('xpack.observability.coPilotPrompt.chatLoading', { + defaultMessage: 'Waiting for a response...', + })} + + + + ); + } else { + inner = ( + + + + + + {content} + + + ); + } + + return ( + <> + {inner} + {isLastComment && } + + ); +}; diff --git a/x-pack/plugins/security_solution/public/assistant/helpers.tsx b/x-pack/plugins/security_solution/public/assistant/helpers.tsx index 6915796ca89ca..32d577171bd29 100644 --- a/x-pack/plugins/security_solution/public/assistant/helpers.tsx +++ b/x-pack/plugins/security_solution/public/assistant/helpers.tsx @@ -84,7 +84,7 @@ export const augmentMessageCodeBlocks = ( const cbd = currentConversation.messages.map(({ content }) => analyzeMarkdown( getMessageContentWithReplacements({ - messageContent: content, + messageContent: content ?? '', // TODO: streaming?? replacements: currentConversation.replacements, }) ) From 1ebf9b67c313f420dbcbabbe1f7353fb1bd9b562 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Wed, 18 Oct 2023 12:43:36 -0600 Subject: [PATCH 07/52] update w main --- .../server/lib/langchain/execute.ts | 48 ------------------- 1 file changed, 48 deletions(-) delete mode 100644 x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts deleted file mode 100644 index 8467458d3cf58..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute.ts +++ /dev/null @@ -1,48 +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 { get } from 'lodash/fp'; -import { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server'; -import { KibanaRequest } from '@kbn/core-http-server'; -import { IncomingMessage } from 'http'; -import { RequestBody } from './types'; - -interface Props { - actions: ActionsPluginStart; - connectorId: string; - request: KibanaRequest; -} -interface StaticResponse { - connector_id: string; - data: string; - status: string; -} - -export const executeAction = async ({ - actions, - request, - connectorId, -}: Props): Promise => { - const actionsClient = await actions.getActionsClientWithRequest(request); - const actionResult = await actionsClient.execute({ - actionId: connectorId, - params: request.body.params, - }); - const content = get('data', actionResult); - if (typeof content === 'string') { - return { - connector_id: connectorId, - data: content, // the response from the actions framework - status: 'ok', - }; - } - // is data stream - if (content instanceof IncomingMessage) { - return content; - } - throw new Error('Unexpected action result'); -}; From c37a1d7ec40c4dfa358ddb4d573936b88f9c5fb1 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 19 Oct 2023 11:03:56 -0600 Subject: [PATCH 08/52] an attempt --- .../impl/assistant/api.tsx | 2 +- .../actions/server/lib/action_executor.ts | 12 +- .../lib/get_token_count_from_openai_stream.ts | 1 + .../elastic_assistant/server/lib/executor.ts | 7 +- .../routes/post_actions_connector_execute.ts | 9 - .../buttons/ask_assistant_button.stories.tsx | 51 ++++ .../buttons/ask_assistant_button.tsx | 111 +++++++++ ...xpand_conversation_list_button.stories.tsx | 37 +++ .../hide_expand_conversation_list_button.tsx | 32 +++ .../buttons/new_chat_button.stories.tsx | 19 ++ .../get_comments/buttons/new_chat_button.tsx | 19 ++ .../regenerate_response_button.stories.tsx | 19 ++ .../buttons/regenerate_response_button.tsx | 26 ++ .../buttons/start_chat_button.stories.tsx | 19 ++ .../buttons/start_chat_button.tsx | 25 ++ .../stop_generating_button.stories.tsx | 19 ++ .../buttons/stop_generating_button.tsx | 27 +++ .../get_comments/esql_code_block.tsx | 75 ++++++ .../get_comments/failed_to_load_response.tsx | 27 +++ .../public/assistant/get_comments/index.tsx | 2 +- .../assistant/get_comments/message_panel.tsx | 37 +++ .../assistant/get_comments/message_text.tsx | 196 +++++++++++++++ .../public/assistant/get_comments/stream.tsx | 224 +++++++++++++++--- 23 files changed, 955 insertions(+), 41 deletions(-) create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.stories.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.stories.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.stories.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.stories.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.stories.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.stories.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/esql_code_block.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/failed_to_load_response.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/message_panel.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index 44dde91e2a602..08bc75d0cda70 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -80,7 +80,7 @@ export const fetchConnectorExecuteAction = async ({ const reader = response?.response?.body?.getReader(); - console.log('is typeof incoming message'); + console.log('is typeof incoming message', reader); if (!reader) { return { response: `${API_ERROR}\n\nCould not get reader from response`, diff --git a/x-pack/plugins/actions/server/lib/action_executor.ts b/x-pack/plugins/actions/server/lib/action_executor.ts index 9cd70d4c7bf91..1ca2ce865cd28 100644 --- a/x-pack/plugins/actions/server/lib/action_executor.ts +++ b/x-pack/plugins/actions/server/lib/action_executor.ts @@ -354,9 +354,19 @@ export class ActionExecutor { }; if (result.data instanceof Readable) { + let body: string; + if (!(validatedParams as { subActionParams: { body: string } }).subActionParams?.body) { + const { stream: _, ...rest } = ( + validatedParams as { subActionParams: { [a: string]: string } } + ).subActionParams; + body = JSON.stringify(rest); + } else { + body = (validatedParams as { subActionParams: { body: string } }).subActionParams + .body; + } getTokenCountFromOpenAIStream({ responseStream: result.data.pipe(new PassThrough()), - body: (validatedParams as { subActionParams: { body: string } }).subActionParams.body, + body, }) .then(({ total, prompt, completion }) => { event.kibana!.action!.execution!.gen_ai!.usage = { diff --git a/x-pack/plugins/actions/server/lib/get_token_count_from_openai_stream.ts b/x-pack/plugins/actions/server/lib/get_token_count_from_openai_stream.ts index 74c89f716171e..2de568a4d66d0 100644 --- a/x-pack/plugins/actions/server/lib/get_token_count_from_openai_stream.ts +++ b/x-pack/plugins/actions/server/lib/get_token_count_from_openai_stream.ts @@ -64,6 +64,7 @@ export async function getTokenCountFromOpenAIStream({ }); try { + console.log('finished', responseStream); await finished(responseStream); } catch { // no need to handle this explicitly diff --git a/x-pack/plugins/elastic_assistant/server/lib/executor.ts b/x-pack/plugins/elastic_assistant/server/lib/executor.ts index 936e3781731d8..247f3a164706f 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/executor.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/executor.ts @@ -8,6 +8,7 @@ import { get } from 'lodash/fp'; import { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server'; import { KibanaRequest } from '@kbn/core-http-server'; +import { PassThrough, Readable } from 'stream'; import { RequestBody } from './langchain/types'; interface Props { @@ -25,7 +26,7 @@ export const executeAction = async ({ actions, request, connectorId, -}: Props): Promise => { +}: Props): Promise => { const actionsClient = await actions.getActionsClientWithRequest(request); const actionResult = await actionsClient.execute({ actionId: connectorId, @@ -39,5 +40,7 @@ export const executeAction = async ({ status: 'ok', }; } - throw new Error('Unexpected action result'); + const readable = get('data', actionResult); + + return (readable as Readable).pipe(new PassThrough()) as Readable; }; 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 30fda687d1f85..4c2299c5599a3 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 @@ -41,15 +41,6 @@ export const postActionsConnectorExecuteRoute = ( // get the actions plugin start contract from the request context: const actions = (await context.elasticAssistant).actions; - // if not langchain, call execute action directly and return the response: - if (!request.body.assistantLangChain) { - const result = await executeAction({ actions, request, connectorId }); - - return response.ok({ - body: result, - }); - } - // if not langchain, call execute action directly and return the response: if (!request.body.assistantLangChain) { const result = await executeAction({ actions, request, connectorId }); diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.stories.tsx new file mode 100644 index 0000000000000..9c0d19e3c4b75 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.stories.tsx @@ -0,0 +1,51 @@ +/* + * 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 React from 'react'; +import { ComponentStory } from '@storybook/react'; +import { EuiButtonSize } from '@elastic/eui'; + +import { AskAssistantButton as Component, AskAssistantButtonProps } from './ask_assistant_button'; + +export default { + component: Component, + title: 'app/Atoms/AskAiAssistantButton', + argTypes: { + size: { + options: ['xs', 's', 'm'] as EuiButtonSize[], + control: { type: 'radio' }, + }, + fill: { + control: { + type: 'boolean', + }, + }, + flush: { + control: { + type: 'boolean', + if: { arg: 'variant', eq: 'empty' }, + }, + }, + variant: { + options: ['basic', 'empty', 'iconOnly'], + control: { type: 'radio' }, + }, + }, +}; + +const Template: ComponentStory = (props: AskAssistantButtonProps) => ( + +); + +const defaultProps = { + fill: true, + size: 'm' as EuiButtonSize, + variant: 'basic' as const, +}; + +export const AskAiAssistantButton = Template.bind({}); +AskAiAssistantButton.args = defaultProps; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.tsx new file mode 100644 index 0000000000000..13f1443bdc34d --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.tsx @@ -0,0 +1,111 @@ +/* + * 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 React from 'react'; +import type { EuiButtonSize, EuiButtonEmptySizes } from '@elastic/eui'; +import { EuiButton, EuiButtonEmpty, EuiToolTip, EuiButtonIcon } from '@elastic/eui'; +import { i18n } from '@kbn/i18n'; + +export type AskAssistantButtonProps = ( + | { + variant: 'basic'; + size: EuiButtonSize; + fill?: boolean; + flush?: false; + } + | { + variant: 'empty'; + size: EuiButtonEmptySizes; + fill?: false; + flush?: 'both'; + } + | { + variant: 'iconOnly'; + size: EuiButtonSize; + fill?: boolean; + flush?: false; + } +) & { + onClick: () => void; +}; + +export function AskAssistantButton({ + fill, + flush, + size, + variant, + onClick, +}: AskAssistantButtonProps) { + const buttonLabel = i18n.translate( + 'xpack.securitySolution.aiAssistant.askAssistantButton.buttonLabel', + { + defaultMessage: 'Ask Assistant', + } + ); + + switch (variant) { + case 'basic': + return ( + + {buttonLabel} + + ); + + case 'empty': + return ( + + {buttonLabel} + + ); + + case 'iconOnly': + return ( + + + + ); + } +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.stories.tsx new file mode 100644 index 0000000000000..d49ccdfb3823d --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.stories.tsx @@ -0,0 +1,37 @@ +/* + * 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 React from 'react'; +import { ComponentStory } from '@storybook/react'; + +import { + HideExpandConversationListButton as Component, + HideExpandConversationListButtonProps, +} from './hide_expand_conversation_list_button'; + +export default { + component: Component, + title: 'app/Atoms/HideExpandConversationListButton', + argTypes: { + isExpanded: { + control: { + type: 'boolean', + }, + }, + }, +}; + +const Template: ComponentStory = ( + props: HideExpandConversationListButtonProps +) => ; + +const defaultProps = { + isExpanded: true, +}; + +export const HideExpandConversationListButton = Template.bind({}); +HideExpandConversationListButton.args = defaultProps; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.tsx new file mode 100644 index 0000000000000..3e2b89aa193de --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.tsx @@ -0,0 +1,32 @@ +/* + * 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 React from 'react'; +import { EuiButtonEmpty } from '@elastic/eui'; +import { i18n } from '@kbn/i18n'; + +export type HideExpandConversationListButtonProps = React.ComponentProps & { + isExpanded: boolean; +}; + +export function HideExpandConversationListButton(props: HideExpandConversationListButtonProps) { + return ( + + {props.isExpanded + ? i18n.translate('xpack.securitySolution.aiAssistant.hideExpandConversationButton.hide', { + defaultMessage: 'Hide chats', + }) + : i18n.translate('xpack.securitySolution.aiAssistant.hideExpandConversationButton.show', { + defaultMessage: 'Show chats', + })} + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.stories.tsx new file mode 100644 index 0000000000000..f4e0cae677ef0 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.stories.tsx @@ -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 { ComponentMeta, ComponentStoryObj } from '@storybook/react'; +import { NewChatButton as Component } from './new_chat_button'; + +const meta: ComponentMeta = { + component: Component, + title: 'app/Atoms/NewChatButton', +}; + +export default meta; + +export const NewChatButton: ComponentStoryObj = { + args: {}, +}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.tsx new file mode 100644 index 0000000000000..9bfe281625159 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.tsx @@ -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 React from 'react'; +import { EuiButton } from '@elastic/eui'; +import { i18n } from '@kbn/i18n'; + +export function NewChatButton(props: React.ComponentProps) { + return ( + + {i18n.translate('xpack.securitySolution.aiAssistant.newChatButton', { + defaultMessage: 'New chat', + })} + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.stories.tsx new file mode 100644 index 0000000000000..bfbb506a9e90a --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.stories.tsx @@ -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 { ComponentMeta, ComponentStoryObj } from '@storybook/react'; +import { RegenerateResponseButton as Component } from './regenerate_response_button'; + +const meta: ComponentMeta = { + component: Component, + title: 'app/Atoms/RegenerateResponseButton', +}; + +export default meta; + +export const RegenerateResponseButton: ComponentStoryObj = { + args: {}, +}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.tsx new file mode 100644 index 0000000000000..3e7ac90dcacda --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.tsx @@ -0,0 +1,26 @@ +/* + * 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 { EuiButtonEmptyProps } from '@elastic/eui'; +import { EuiButtonEmpty } from '@elastic/eui'; +import React from 'react'; +import { i18n } from '@kbn/i18n'; + +export function RegenerateResponseButton(props: Partial) { + return ( + + {i18n.translate('xpack.securitySolution.aiAssistant.regenerateResponseButtonLabel', { + defaultMessage: 'Regenerate', + })} + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.stories.tsx new file mode 100644 index 0000000000000..de9116900b61b --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.stories.tsx @@ -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 { ComponentMeta, ComponentStoryObj } from '@storybook/react'; +import { StartChatButton as Component } from './start_chat_button'; + +const meta: ComponentMeta = { + component: Component, + title: 'app/Atoms/StartChatButton', +}; + +export default meta; + +export const StartChatButton: ComponentStoryObj = { + args: {}, +}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.tsx new file mode 100644 index 0000000000000..a0354c9533d7b --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.tsx @@ -0,0 +1,25 @@ +/* + * 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 React from 'react'; +import { EuiButton } from '@elastic/eui'; +import { i18n } from '@kbn/i18n'; + +export function StartChatButton(props: React.ComponentProps) { + return ( + + {i18n.translate('xpack.securitySolution.aiAssistant.button.startChat', { + defaultMessage: 'Start chat', + })} + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.stories.tsx new file mode 100644 index 0000000000000..acf27b4a01274 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.stories.tsx @@ -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 { ComponentMeta, ComponentStoryObj } from '@storybook/react'; +import { StopGeneratingButton as Component } from './stop_generating_button'; + +const meta: ComponentMeta = { + component: Component, + title: 'app/Atoms/StopGeneratingButton', +}; + +export default meta; + +export const StopGeneratingButton: ComponentStoryObj = { + args: {}, +}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.tsx new file mode 100644 index 0000000000000..9a8f1e8ae3ae7 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.tsx @@ -0,0 +1,27 @@ +/* + * 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 { EuiButtonEmptyProps } from '@elastic/eui'; +import { EuiButtonEmpty } from '@elastic/eui'; +import React from 'react'; +import { i18n } from '@kbn/i18n'; + +export function StopGeneratingButton(props: Partial) { + return ( + + {i18n.translate('xpack.securitySolution.aiAssistant.stopGeneratingButtonLabel', { + defaultMessage: 'Stop generating', + })} + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/esql_code_block.tsx new file mode 100644 index 0000000000000..c9eda951c7fe8 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/esql_code_block.tsx @@ -0,0 +1,75 @@ +/* + * 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 { i18n } from '@kbn/i18n'; +import React from 'react'; +import type { ChatActionClickHandler } from './message_text'; + +export function EsqlCodeBlock({ + value, + actionsDisabled, + onActionClick, +}: { + value: string; + actionsDisabled: boolean; + onActionClick: ChatActionClickHandler; +}) { + const theme = useEuiTheme(); + + return ( + + + + + {value} + + + + + + onActionClick({ type: 'executeEsqlQuery', query: value })} + disabled={actionsDisabled} + > + {i18n.translate('xpack.securitySolution.aiAssistant.runThisQuery', { + defaultMessage: 'Run this query', + })} + + + + + + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/failed_to_load_response.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/failed_to_load_response.tsx new file mode 100644 index 0000000000000..5161f5a2b0298 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/failed_to_load_response.tsx @@ -0,0 +1,27 @@ +/* + * 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 { EuiFlexGroup, EuiFlexItem, EuiIcon, EuiText } from '@elastic/eui'; +import React from 'react'; +import { i18n } from '@kbn/i18n'; + +export function FailedToLoadResponse() { + return ( + + + + + + + {i18n.translate('xpack.securitySolution.aiAssistant.failedLoadingResponseText', { + defaultMessage: 'Failed to load response', + })} + + + + ); +} 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 e9334232be8c7..50fb1472ec960 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 @@ -14,7 +14,7 @@ import { AssistantAvatar } from '@kbn/elastic-assistant'; import { css } from '@emotion/react'; import { euiThemeVars } from '@kbn/ui-theme'; import type { EuiPanelProps } from '@elastic/eui/src/components/panel'; -import { StreamComment } from './stream_obs'; +import { StreamComment } from './stream'; import { CommentActions } from '../comment_actions'; import * as i18n from './translations'; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/message_panel.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/message_panel.tsx new file mode 100644 index 0000000000000..78a1e8fae4778 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/message_panel.tsx @@ -0,0 +1,37 @@ +/* + * 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 { EuiHorizontalRule, EuiSpacer } from '@elastic/eui'; +import React from 'react'; +import { FailedToLoadResponse } from './failed_to_load_response'; + +interface Props { + error?: Error; + body?: React.ReactNode; + controls?: React.ReactNode; +} + +export function MessagePanel(props: Props) { + return ( + <> + {props.body} + {props.error ? ( + <> + {props.body ? : null} + + + ) : null} + {props.controls ? ( + <> + + + + {props.controls} + + ) : null} + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx new file mode 100644 index 0000000000000..d2a7f3b90b5f2 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx @@ -0,0 +1,196 @@ +/* + * 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 { + EuiMarkdownFormat, + EuiSpacer, + EuiText, + getDefaultEuiMarkdownParsingPlugins, + getDefaultEuiMarkdownProcessingPlugins, +} from '@elastic/eui'; +import { css } from '@emotion/css'; +import classNames from 'classnames'; +import type { Code, InlineCode, Parent, Text } from 'mdast'; +import React, { useMemo, useRef } from 'react'; +import type { Node } from 'unist'; +import { EsqlCodeBlock } from './esql_code_block'; + +export type ChatActionClickHandler = (payload: { type: string; query: string }) => Promise; +interface Props { + content: string; + loading: boolean; + onActionClick: ChatActionClickHandler; +} + +const ANIMATION_TIME = 1; + +const cursorCss = css` + @keyframes blink { + 0% { + opacity: 0; + } + 50% { + opacity: 1; + } + 100% { + opacity: 0; + } + } + + animation: blink ${ANIMATION_TIME}s infinite; + width: 10px; + height: 16px; + vertical-align: middle; + display: inline-block; + background: rgba(0, 0, 0, 0.25); +`; + +const Cursor = () => ; + +// a weird combination of different whitespace chars to make sure it stays +// invisible even when we cannot properly parse the text while still being +// unique +const CURSOR = ` ᠎  `; + +const loadingCursorPlugin = () => { + 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 !== 'text' && node.type !== 'inlineCode' && node.type !== 'code') { + return; + } + + const textNode = node as Text | InlineCode | Code; + + const indexOfCursor = textNode.value.indexOf(CURSOR); + if (indexOfCursor === -1) { + return; + } + + textNode.value = textNode.value.replace(CURSOR, ''); + + const indexOfNode = parent?.children.indexOf(textNode); + parent?.children.splice(indexOfNode + 1, 0, { + type: 'cursor' as Text['type'], + value: CURSOR, + }); + }; + + return (tree: Node) => { + visitor(tree); + }; +}; + +const esqlLanguagePlugin = () => { + 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 === 'esql') { + node.type = 'esql'; + } + }; + + return (tree: Node) => { + visitor(tree); + }; +}; + +export function MessageText({ loading, content, onActionClick }: Props) { + console.log('content??', content); + const containerClassName = css` + overflow-wrap: break-word; + `; + + const onActionClickRef = useRef(onActionClick); + + onActionClickRef.current = onActionClick; + + const { parsingPluginList, processingPluginList } = useMemo(() => { + const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); + + const processingPlugins = getDefaultEuiMarkdownProcessingPlugins(); + + const { components } = processingPlugins[1][1]; + + processingPlugins[1][1].components = { + ...components, + cursor: Cursor, + esql: (props) => { + return ( + <> + + + + ); + }, + table: (props) => ( + <> +
+ {' '} + + + + + ), + th: (props) => { + const { children, ...rest } = props; + return ( + + ); + }, + tr: (props) => , + td: (props) => { + const { children, ...rest } = props; + return ( + + ); + }, + }; + + return { + parsingPluginList: [loadingCursorPlugin, esqlLanguagePlugin, ...parsingPlugins], + processingPluginList: processingPlugins, + }; + }, [loading]); + + return ( + + + {`${content}${loading ? CURSOR : ''}`} + + + ); +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx index 53fdcafa0a2d5..7806bb40a39d6 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx @@ -5,41 +5,211 @@ * 2.0. */ -import React, { useEffect, useState } from 'react'; +import React, { useEffect, useMemo } from 'react'; +import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui'; +import { concatMap, delay, Observable, of } from 'rxjs'; +import useObservable from 'react-use/lib/useObservable'; +import { StopGeneratingButton } from './buttons/stop_generating_button'; +import { RegenerateResponseButton } from './buttons/regenerate_response_button'; +import { MessagePanel } from './message_panel'; +import { MessageText } from './message_text'; -export const StreamComment = ({ reader }: { reader: ReadableStreamDefaultReader }) => { - const [data, setData] = useState(''); - const [loading, setLoading] = useState(false); +export interface PromptObservableState { + chunks: Chunk[]; + message?: string; + error?: string; + loading: boolean; +} + +interface ChunkChoice { + index: 0; + delta: { role: string; content: string }; + finish_reason: null | string; +} + +interface Chunk { + id: string; + object: string; + created: number; + model: string; + choices: ChunkChoice[]; +} + +const cursorCss = ` + @keyframes blink { + 0% { + opacity: 1; + } + 50% { + opacity: 0; + } + 100% { + opacity: 1; + } + } + + animation: blink 1s infinite; + width: 10px; + height: 16px; + vertical-align: middle; + display: inline-block; + background: rgba(0, 0, 0, 0.25); +`; + +function getMessageFromChunks(chunks: Chunk[]) { + let message = ''; + chunks.forEach((chunk) => { + message += chunk.choices[0]?.delta.content ?? ''; + }); + return message; +} + +interface Props { + index: number; + isLastComment: boolean; + lastCommentRef: React.MutableRefObject; + reader: ReadableStreamDefaultReader; +} + +export const StreamComment = ({ index, isLastComment, lastCommentRef, reader }: Props) => { + const response$ = useMemo( + () => + new Observable((observer) => { + observer.next({ chunks: [], loading: true }); - useEffect(() => { - const fetchData = async () => { - try { - setLoading(true); const decoder = new TextDecoder(); - while (true) { - const { value, done } = await reader.read(); - if (done) { - setLoading(false); - break; - } + const chunks: Chunk[] = []; + + let prev: string = ''; + + function read() { + reader.read().then(({ done, value }: { done: boolean; value?: Uint8Array }) => { + try { + if (done) { + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: false, + }); + observer.complete(); + return; + } + + let lines: string[] = (prev + decoder.decode(value)).split('\n'); + + const lastLine: string = lines[lines.length - 1]; + + const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; + + if (isPartialChunk) { + prev = lastLine; + lines.pop(); + } else { + prev = ''; + } - const decodedChunk = decoder.decode(value, { stream: true }); - setData((prevValue: string) => `${prevValue}${decodedChunk}`); + lines = lines.map((str) => str.substr(6)).filter((str) => !!str && str !== '[DONE]'); + + const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); + + nextChunks.forEach((chunk) => { + chunks.push(chunk); + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: true, + }); + }); + } catch (err) { + observer.error(err); + return; + } + read(); + }); } - } catch (error) { - setLoading(false); - // Handle other errors - } - }; - fetchData(); - }, [reader]); + read(); + + return () => { + reader.cancel(); + }; + }).pipe(concatMap((value) => of(value).pipe(delay(50)))), + [reader] + ); + + const response = useObservable(response$); + + useEffect(() => {}, [response$]); + + let content = response?.message ?? ''; + + let state: 'init' | 'loading' | 'streaming' | 'error' | 'complete' = 'init'; + + if (response?.loading) { + state = content ? 'streaming' : 'loading'; + console.log('state', state); + } else if (response && 'error' in response && response.error) { + state = 'error'; + content = response.error; + } else if (content) { + state = 'complete'; + } + + const isLoading = state === 'init' || state === 'loading'; + const isStreaming = state === 'streaming'; return ( - - {loading && {'Fetching response...'}} - {data} - + <> + {}} + /> + } + error={response?.error ? new Error(response?.error) : undefined} + controls={ + isLoading || isStreaming ? ( + { + console.log('stop generating'); + // subscription?.unsubscribe(); + // setLoading(false); + // setDisplayedMessages((prevMessages) => + // prevMessages.concat({ + // '@timestamp': new Date().toISOString(), + // message: { + // ...pendingMessage!.message, + // }, + // }) + // ); + // setPendingMessage((prev) => ({ + // message: { + // role: MessageRole.Assistant, + // ...prev?.message, + // }, + // aborted: true, + // error: new AbortError(), + // })); + }} + /> + ) : ( + + + { + console.log('RegenerateResponseButton'); + // reloadRecalledMessages(); + }} + /> + + + ) + } + /> + {isLastComment && } + ); }; From fa00236b3021a3351309b9321dbb0cf40480723d Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 19 Oct 2023 11:10:12 -0600 Subject: [PATCH 09/52] rm extra files --- .../buttons/ask_assistant_button.stories.tsx | 51 ----- .../buttons/ask_assistant_button.tsx | 111 ---------- ...xpand_conversation_list_button.stories.tsx | 37 ---- .../hide_expand_conversation_list_button.tsx | 32 --- .../buttons/new_chat_button.stories.tsx | 19 -- .../get_comments/buttons/new_chat_button.tsx | 19 -- .../buttons/start_chat_button.stories.tsx | 19 -- .../buttons/start_chat_button.tsx | 25 --- .../public/assistant/get_comments/stream.tsx | 1 - .../assistant/get_comments/stream_obs.tsx | 207 ------------------ 10 files changed, 521 deletions(-) delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.stories.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.stories.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.stories.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.stories.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream_obs.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.stories.tsx deleted file mode 100644 index 9c0d19e3c4b75..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.stories.tsx +++ /dev/null @@ -1,51 +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 React from 'react'; -import { ComponentStory } from '@storybook/react'; -import { EuiButtonSize } from '@elastic/eui'; - -import { AskAssistantButton as Component, AskAssistantButtonProps } from './ask_assistant_button'; - -export default { - component: Component, - title: 'app/Atoms/AskAiAssistantButton', - argTypes: { - size: { - options: ['xs', 's', 'm'] as EuiButtonSize[], - control: { type: 'radio' }, - }, - fill: { - control: { - type: 'boolean', - }, - }, - flush: { - control: { - type: 'boolean', - if: { arg: 'variant', eq: 'empty' }, - }, - }, - variant: { - options: ['basic', 'empty', 'iconOnly'], - control: { type: 'radio' }, - }, - }, -}; - -const Template: ComponentStory = (props: AskAssistantButtonProps) => ( - -); - -const defaultProps = { - fill: true, - size: 'm' as EuiButtonSize, - variant: 'basic' as const, -}; - -export const AskAiAssistantButton = Template.bind({}); -AskAiAssistantButton.args = defaultProps; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.tsx deleted file mode 100644 index 13f1443bdc34d..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/ask_assistant_button.tsx +++ /dev/null @@ -1,111 +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 React from 'react'; -import type { EuiButtonSize, EuiButtonEmptySizes } from '@elastic/eui'; -import { EuiButton, EuiButtonEmpty, EuiToolTip, EuiButtonIcon } from '@elastic/eui'; -import { i18n } from '@kbn/i18n'; - -export type AskAssistantButtonProps = ( - | { - variant: 'basic'; - size: EuiButtonSize; - fill?: boolean; - flush?: false; - } - | { - variant: 'empty'; - size: EuiButtonEmptySizes; - fill?: false; - flush?: 'both'; - } - | { - variant: 'iconOnly'; - size: EuiButtonSize; - fill?: boolean; - flush?: false; - } -) & { - onClick: () => void; -}; - -export function AskAssistantButton({ - fill, - flush, - size, - variant, - onClick, -}: AskAssistantButtonProps) { - const buttonLabel = i18n.translate( - 'xpack.securitySolution.aiAssistant.askAssistantButton.buttonLabel', - { - defaultMessage: 'Ask Assistant', - } - ); - - switch (variant) { - case 'basic': - return ( - - {buttonLabel} - - ); - - case 'empty': - return ( - - {buttonLabel} - - ); - - case 'iconOnly': - return ( - - - - ); - } -} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.stories.tsx deleted file mode 100644 index d49ccdfb3823d..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.stories.tsx +++ /dev/null @@ -1,37 +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 React from 'react'; -import { ComponentStory } from '@storybook/react'; - -import { - HideExpandConversationListButton as Component, - HideExpandConversationListButtonProps, -} from './hide_expand_conversation_list_button'; - -export default { - component: Component, - title: 'app/Atoms/HideExpandConversationListButton', - argTypes: { - isExpanded: { - control: { - type: 'boolean', - }, - }, - }, -}; - -const Template: ComponentStory = ( - props: HideExpandConversationListButtonProps -) => ; - -const defaultProps = { - isExpanded: true, -}; - -export const HideExpandConversationListButton = Template.bind({}); -HideExpandConversationListButton.args = defaultProps; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.tsx deleted file mode 100644 index 3e2b89aa193de..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/hide_expand_conversation_list_button.tsx +++ /dev/null @@ -1,32 +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 React from 'react'; -import { EuiButtonEmpty } from '@elastic/eui'; -import { i18n } from '@kbn/i18n'; - -export type HideExpandConversationListButtonProps = React.ComponentProps & { - isExpanded: boolean; -}; - -export function HideExpandConversationListButton(props: HideExpandConversationListButtonProps) { - return ( - - {props.isExpanded - ? i18n.translate('xpack.securitySolution.aiAssistant.hideExpandConversationButton.hide', { - defaultMessage: 'Hide chats', - }) - : i18n.translate('xpack.securitySolution.aiAssistant.hideExpandConversationButton.show', { - defaultMessage: 'Show chats', - })} - - ); -} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.stories.tsx deleted file mode 100644 index f4e0cae677ef0..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.stories.tsx +++ /dev/null @@ -1,19 +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 { ComponentMeta, ComponentStoryObj } from '@storybook/react'; -import { NewChatButton as Component } from './new_chat_button'; - -const meta: ComponentMeta = { - component: Component, - title: 'app/Atoms/NewChatButton', -}; - -export default meta; - -export const NewChatButton: ComponentStoryObj = { - args: {}, -}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.tsx deleted file mode 100644 index 9bfe281625159..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/new_chat_button.tsx +++ /dev/null @@ -1,19 +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 React from 'react'; -import { EuiButton } from '@elastic/eui'; -import { i18n } from '@kbn/i18n'; - -export function NewChatButton(props: React.ComponentProps) { - return ( - - {i18n.translate('xpack.securitySolution.aiAssistant.newChatButton', { - defaultMessage: 'New chat', - })} - - ); -} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.stories.tsx deleted file mode 100644 index de9116900b61b..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.stories.tsx +++ /dev/null @@ -1,19 +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 { ComponentMeta, ComponentStoryObj } from '@storybook/react'; -import { StartChatButton as Component } from './start_chat_button'; - -const meta: ComponentMeta = { - component: Component, - title: 'app/Atoms/StartChatButton', -}; - -export default meta; - -export const StartChatButton: ComponentStoryObj = { - args: {}, -}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.tsx deleted file mode 100644 index a0354c9533d7b..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/start_chat_button.tsx +++ /dev/null @@ -1,25 +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 React from 'react'; -import { EuiButton } from '@elastic/eui'; -import { i18n } from '@kbn/i18n'; - -export function StartChatButton(props: React.ComponentProps) { - return ( - - {i18n.translate('xpack.securitySolution.aiAssistant.button.startChat', { - defaultMessage: 'Start chat', - })} - - ); -} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx index 7806bb40a39d6..9cec6a146eca1 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx @@ -148,7 +148,6 @@ export const StreamComment = ({ index, isLastComment, lastCommentRef, reader }: if (response?.loading) { state = content ? 'streaming' : 'loading'; - console.log('state', state); } else if (response && 'error' in response && response.error) { state = 'error'; content = response.error; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream_obs.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream_obs.tsx deleted file mode 100644 index 5b520c596a844..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream_obs.tsx +++ /dev/null @@ -1,207 +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 React, { useEffect, useMemo } from 'react'; -import { - EuiFlexGroup, - EuiFlexItem, - EuiText, - EuiLoadingSpinner, - EuiIcon, - EuiMarkdownFormat, -} from '@elastic/eui'; -import { i18n } from '@kbn/i18n'; -import { concatMap, delay, Observable, of } from 'rxjs'; -import useObservable from 'react-use/lib/useObservable'; - -export interface PromptObservableState { - chunks: Chunk[]; - message?: string; - error?: string; - loading: boolean; -} - -interface ChunkChoice { - index: 0; - delta: { role: string; content: string }; - finish_reason: null | string; -} - -interface Chunk { - id: string; - object: string; - created: number; - model: string; - choices: ChunkChoice[]; -} - -const cursorCss = ` - @keyframes blink { - 0% { - opacity: 1; - } - 50% { - opacity: 0; - } - 100% { - opacity: 1; - } - } - - animation: blink 1s infinite; - width: 10px; - height: 16px; - vertical-align: middle; - display: inline-block; - background: rgba(0, 0, 0, 0.25); -`; - -function getMessageFromChunks(chunks: Chunk[]) { - let message = ''; - chunks.forEach((chunk) => { - message += chunk.choices[0]?.delta.content ?? ''; - }); - return message; -} - -interface Props { - index: number; - isLastComment: boolean; - lastCommentRef: React.MutableRefObject; - reader: ReadableStreamDefaultReader; -} - -export const StreamComment = ({ index, isLastComment, lastCommentRef, reader }: Props) => { - const response$ = useMemo( - () => - new Observable((observer) => { - observer.next({ chunks: [], loading: true }); - - const decoder = new TextDecoder(); - - const chunks: Chunk[] = []; - - let prev: string = ''; - - function read() { - reader.read().then(({ done, value }: { done: boolean; value?: Uint8Array }) => { - try { - if (done) { - observer.next({ - chunks, - message: getMessageFromChunks(chunks), - loading: false, - }); - observer.complete(); - return; - } - - let lines: string[] = (prev + decoder.decode(value)).split('\n'); - - const lastLine: string = lines[lines.length - 1]; - - const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; - - if (isPartialChunk) { - prev = lastLine; - lines.pop(); - } else { - prev = ''; - } - - lines = lines.map((str) => str.substr(6)).filter((str) => !!str && str !== '[DONE]'); - - const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); - - nextChunks.forEach((chunk) => { - chunks.push(chunk); - observer.next({ - chunks, - message: getMessageFromChunks(chunks), - loading: true, - }); - }); - } catch (err) { - observer.error(err); - return; - } - read(); - }); - } - - read(); - - return () => { - reader.cancel(); - }; - }).pipe(concatMap((value) => of(value).pipe(delay(50)))), - [reader] - ); - - const response = useObservable(response$); - - useEffect(() => {}, [response$]); - - let content = response?.message ?? ''; - - let state: 'init' | 'loading' | 'streaming' | 'error' | 'complete' = 'init'; - - if (response?.loading) { - state = content ? 'streaming' : 'loading'; - } else if (response && 'error' in response && response.error) { - state = 'error'; - content = response.error; - } else if (content) { - state = 'complete'; - } - - let inner: React.ReactNode; - - if (state === 'complete' || state === 'streaming') { - inner = ( - <> - - {content} - - {state === 'streaming' ? : <>} - - ); - } else if (state === 'init' || state === 'loading') { - inner = ( - - - - - - - {i18n.translate('xpack.observability.coPilotPrompt.chatLoading', { - defaultMessage: 'Waiting for a response...', - })} - - - - ); - } else { - inner = ( - - - - - - {content} - - - ); - } - - return ( - <> - {inner} - {isLastComment && } - - ); -}; From f2b354cd676ae061eac57ba348067e109dcbefa1 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 19 Oct 2023 11:11:37 -0600 Subject: [PATCH 10/52] rm more --- .../impl/assistant/chat_send/chat_stream.tsx | 168 ------------------ 1 file changed, 168 deletions(-) delete mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/chat_stream.tsx diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/chat_stream.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/chat_stream.tsx deleted file mode 100644 index 03603b3b43a2f..0000000000000 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/chat_stream.tsx +++ /dev/null @@ -1,168 +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 { - BehaviorSubject, - catchError, - concatMap, - delay, - filter as rxJsFilter, - finalize, - map, - Observable, - of, - scan, - share, - shareReplay, - tap, - timestamp, -} from 'rxjs'; -import { cloneDeep } from 'lodash'; -import { AbortError } from '@kbn/kibana-utils-plugin/common'; -import { ConversationRole, Message } from '../../assistant_context/types'; - -export function readableStreamReaderIntoObservable( - readableStreamReader: ReadableStreamDefaultReader -): Observable { - return new Observable((subscriber) => { - let lineBuffer: string = ''; - - async function read(): Promise { - const { done, value } = await readableStreamReader.read(); - if (done) { - if (lineBuffer) { - subscriber.next(lineBuffer); - } - subscriber.complete(); - - return; - } - - const textChunk = new TextDecoder().decode(value); - - const lines = textChunk.split('\n'); - lines[0] = lineBuffer + lines[0]; - - lineBuffer = lines.pop() || ''; - - for (const line of lines) { - subscriber.next(line); - } - - return read(); - } - - read().catch((err) => subscriber.error(err)); - - return () => { - readableStreamReader.cancel().catch(() => {}); - }; - }).pipe(share()); -} - -const role: ConversationRole = 'assistant'; -export const chatStream = (reader: ReadableStreamDefaultReader) => { - const subject = new BehaviorSubject<{ message: Message }>({ - message: { role }, - }); - readableStreamReaderIntoObservable(reader).pipe( - // lines start with 'data: ' - map((line) => line.substring(6)), - // a message completes with the line '[DONE]' - rxJsFilter((line) => !!line && line !== '[DONE]'), - // parse the JSON, add the type - map((line) => JSON.parse(line) as {} | { error: { message: string } }), - // validate the message. in some cases OpenAI - // will throw halfway through the message - tap((line) => { - if ('error' in line) { - throw new Error(line.error.message); - } - }), - // there also might be some metadata that we need - // to exclude - rxJsFilter( - (line): line is { object: string } => - 'object' in line && line.object === 'chat.completion.chunk' - ), - // this is how OpenAI signals that the context window - // limit has been exceeded - tap((line) => { - if (line.choices[0].finish_reason === 'length') { - throw new Error(`Token limit reached`); - } - }), - // merge the messages - scan( - (acc, { choices }) => { - acc.message.content += choices[0].delta.content ?? ''; - acc.message.function_call.name += choices[0].delta.function_call?.name ?? ''; - acc.message.function_call.arguments += choices[0].delta.function_call?.arguments ?? ''; - return cloneDeep(acc); - }, - { - message: { - content: '', - function_call: { - name: '', - arguments: '', - trigger: 'assistant' as const, - }, - role: 'assistant', - }, - } - ), - // convert an error into state - catchError((error) => - of({ - ...subject.value, - error, - aborted: error instanceof AbortError, - }) - ) - ); - const MIN_DELAY = 35; - - const pendingMessages$ = subject.pipe( - // make sure the request is only triggered once, - // even with multiple subscribers - shareReplay(1), - // if the Observable is no longer subscribed, - // abort the running request - finalize(() => { - // controller.abort(); - }), - // append a timestamp of when each value was emitted - timestamp(), - // use the previous timestamp to calculate a target - // timestamp for emitting the next value - scan((acc, value) => { - const lastTimestamp = acc.timestamp || 0; - const emitAt = Math.max(lastTimestamp + MIN_DELAY, value.timestamp); - return { - timestamp: emitAt, - value: value.value, - }; - }), - // add the delay based on the elapsed time - // using concatMap(of(value).pipe(delay(50)) - // leads to browser issues because timers - // are throttled when the tab is not active - concatMap((value) => { - const now = Date.now(); - const delayFor = value.timestamp - now; - - if (delayFor <= 0) { - return of(value.value); - } - - return of(value.value).pipe(delay(delayFor)); - }) - ); - - return pendingMessages$; -}; From eb856ac5ca29dd8fc7736b617471df83ea246452 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 19 Oct 2023 11:27:40 -0600 Subject: [PATCH 11/52] console logs --- .../public/assistant/get_comments/message_text.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx index d2a7f3b90b5f2..71db33ee28c47 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx @@ -57,6 +57,7 @@ const CURSOR = ` ᠎  `; const loadingCursorPlugin = () => { const visitor = (node: Node, parent?: Parent) => { + console.log('node loadingCursorPlugin?', node); if ('children' in node) { const nodeAsParent = node as Parent; nodeAsParent.children.forEach((child) => { @@ -91,6 +92,7 @@ const loadingCursorPlugin = () => { const esqlLanguagePlugin = () => { const visitor = (node: Node, parent?: Parent) => { + console.log('node esqlLanguagePlugin?', node); if ('children' in node) { const nodeAsParent = node as Parent; nodeAsParent.children.forEach((child) => { @@ -109,7 +111,6 @@ const esqlLanguagePlugin = () => { }; export function MessageText({ loading, content, onActionClick }: Props) { - console.log('content??', content); const containerClassName = css` overflow-wrap: break-word; `; @@ -183,9 +184,8 @@ export function MessageText({ loading, content, onActionClick }: Props) { }, [loading]); return ( - + From 8a71d4ec45397dbe9fb1967cbe7088cfc3429a05 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 19 Oct 2023 15:35:13 -0600 Subject: [PATCH 12/52] saved to memory and stop button works --- .../assistant/chat_send/use_chat_send.tsx | 3 +- .../impl/assistant/index.tsx | 4 +- .../impl/assistant/use_conversation/index.tsx | 59 ++++++------ .../impl/assistant_context/index.tsx | 6 +- .../public/assistant/get_comments/index.tsx | 29 +++++- .../assistant/get_comments/message_text.tsx | 9 +- .../public/assistant/get_comments/stream.tsx | 96 ++++++++++++------- 7 files changed, 131 insertions(+), 75 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index 53087e577dded..3e1c194097888 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -54,8 +54,7 @@ export const useChatSend = ({ setUserPrompt, }: UseChatSendProps): UseChatSend => { const { isLoading, sendMessages } = useSendMessages(); - const { appendMessage, appendReplacements, appendStreamMessage, clearConversation } = - useConversation(); + const { appendMessage, appendReplacements, clearConversation } = useConversation(); const handlePromptChange = (prompt: string) => { setPromptTextPreview(prompt); 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 cdfd8187e7a2f..dd7e4462f6f98 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -93,7 +93,7 @@ const AssistantComponent: React.FC = ({ [selectedPromptContexts] ); - const { createConversation } = useConversation(); + const { amendMessage, createConversation } = useConversation(); // Connector details const { data: connectors, isSuccess: areConnectorsFetched } = useLoadConnectors({ http }); @@ -339,6 +339,7 @@ const AssistantComponent: React.FC = ({ currentConversation, lastCommentRef, showAnonymizedValues, + amendMessage, })} css={css` margin-right: 20px; @@ -368,6 +369,7 @@ const AssistantComponent: React.FC = ({ ), [ + amendMessage, currentConversation, editingSystemPromptId, getComments, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx index 659b6b5c7de70..745cfe0a8168e 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx @@ -5,7 +5,7 @@ * 2.0. */ -import { useCallback, useEffect, useState } from 'react'; +import { useCallback } from 'react'; import { useAssistantContext } from '../../assistant_context'; import { Conversation, Message } from '../../assistant_context/types'; @@ -31,7 +31,7 @@ export const DEFAULT_CONVERSATION_STATE: Conversation = { }, }; -interface AppendMessageProps { +export interface AppendMessageProps { conversationId: string; message: Message; } @@ -58,6 +58,7 @@ interface SetConversationProps { interface UseConversation { appendStreamMessage: ({ conversationId, message }: AppendMessageProps) => void; appendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[]; + amendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[]; appendReplacements: ({ conversationId, replacements, @@ -72,33 +73,35 @@ interface UseConversation { export const useConversation = (): UseConversation => { const { allSystemPrompts, assistantTelemetry, setConversations } = useAssistantContext(); - const [pendingMessage, setPendingMessage] = useState(); - useEffect(() => { - console.log('pendingMessage', pendingMessage); - }, [pendingMessage]); + /** + * Replaces the last message of conversation[] for a given conversationId + */ + const amendMessage = useCallback( + ({ conversationId, message }) => { + let messages: Message[] = []; + setConversations((prev: Record) => { + const prevConversation: Conversation | undefined = prev[conversationId]; - const appendStreamMessage = useCallback(({ conversationId, reader }) => { - // assistantTelemetry?.reportAssistantMessageSent({ conversationId, role }); - let messages: Message[] = []; - setConversations((prev: Record) => { - const prevConversation: Conversation | undefined = prev[conversationId]; + if (prevConversation != null) { + prevConversation.messages.pop(); + messages = [...prevConversation.messages, message]; + const newConversation = { + ...prevConversation, + messages, + }; + return { + ...prev, + [conversationId]: newConversation, + }; + } else { + return prev; + } + }); + return messages; + }, + [setConversations] + ); - if (prevConversation != null) { - messages = [...prevConversation.messages, { reader, role: 'assistant' }]; - const newConversation = { - ...prevConversation, - messages, - }; - return { - ...prev, - [conversationId]: newConversation, - }; - } else { - return prev; - } - }); - return messages; - }, []); /** * Append a message to the conversation[] for a given conversationId */ @@ -290,9 +293,9 @@ export const useConversation = (): UseConversation => { ); return { + amendMessage, appendMessage, appendReplacements, - appendStreamMessage, clearConversation, createConversation, deleteConversation, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index f296246736a78..6c99970d01a97 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -13,6 +13,7 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react'; import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public'; import { useLocalStorage } from 'react-use'; import type { DocLinksStart } from '@kbn/core-doc-links-browser'; +import { AppendMessageProps } from '../assistant/use_conversation'; import { updatePromptContexts } from './helpers'; import type { PromptContext, @@ -35,7 +36,7 @@ import { SYSTEM_PROMPT_LOCAL_STORAGE_KEY, } from './constants'; import { CONVERSATIONS_TAB, SettingsTabs } from '../assistant/settings/assistant_settings'; -import { AssistantAvailability, AssistantTelemetry } from './types'; +import { AssistantAvailability, AssistantTelemetry, Message } from './types'; export interface ShowAssistantOverlayProps { showOverlay: boolean; @@ -104,10 +105,11 @@ export interface UseAssistantContext { currentConversation, lastCommentRef, showAnonymizedValues, + amendMessage, }: { currentConversation: Conversation; lastCommentRef: React.MutableRefObject; - + amendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[]; showAnonymizedValues: boolean; }) => EuiCommentProps[]; http: HttpSetup; 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 50fb1472ec960..6252c64bd875f 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 @@ -6,7 +6,7 @@ */ import type { EuiCommentProps } from '@elastic/eui'; -import type { Conversation } from '@kbn/elastic-assistant'; +import type { Conversation, Message } from '@kbn/elastic-assistant'; import { EuiAvatar, EuiMarkdownFormat, EuiText, tint } from '@elastic/eui'; import React from 'react'; @@ -19,15 +19,36 @@ import { CommentActions } from '../comment_actions'; import * as i18n from './translations'; export const getComments = ({ + amendMessage, currentConversation, lastCommentRef, showAnonymizedValues, }: { + amendMessage: ({ + conversationId, + message, + }: { + conversationId: string; + message: Message; + }) => Message[]; currentConversation: Conversation; lastCommentRef: React.MutableRefObject; showAnonymizedValues: boolean; -}): EuiCommentProps[] => - currentConversation.messages.map((message, index) => { +}): EuiCommentProps[] => { + const amendMessageOfConversation = (content: string, index: number) => { + const replacementMessage = currentConversation.messages[index]; + console.log('replacementMessage', { replacementMessage, content }); + amendMessage({ + conversationId: currentConversation.id, + message: { + content, + role: replacementMessage.role, + timestamp: replacementMessage.timestamp, + }, + }); + }; + + return currentConversation.messages.map((message, index) => { const isUser = message.role === 'user'; const replacements = currentConversation.replacements; const errorStyles = { @@ -94,6 +115,7 @@ export const getComments = ({ ...messageProps, children: ( { const visitor = (node: Node, parent?: Parent) => { - console.log('node loadingCursorPlugin?', node); if ('children' in node) { const nodeAsParent = node as Parent; nodeAsParent.children.forEach((child) => { @@ -92,7 +91,6 @@ const loadingCursorPlugin = () => { const esqlLanguagePlugin = () => { const visitor = (node: Node, parent?: Parent) => { - console.log('node esqlLanguagePlugin?', node); if ('children' in node) { const nodeAsParent = node as Parent; nodeAsParent.children.forEach((child) => { @@ -103,6 +101,13 @@ const esqlLanguagePlugin = () => { if (node.type === 'code' && node.lang === 'esql') { node.type = 'esql'; } + // TODO: make these renderers + if ( + (node.type === 'code' && node.lang === 'kql') || + (node.type === 'code' && node.lang === 'eql') + ) { + node.type = 'esql'; + } }; return (tree: Node) => { diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx index 9cec6a146eca1..895510ffef8ea 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx @@ -5,10 +5,10 @@ * 2.0. */ -import React, { useEffect, useMemo } from 'react'; +import React, { useCallback, useEffect, useMemo, useState } from 'react'; import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui'; -import { concatMap, delay, Observable, of } from 'rxjs'; -import useObservable from 'react-use/lib/useObservable'; +import type { Subscription } from 'rxjs'; +import { concatMap, delay, Observable, of, share } from 'rxjs'; import { StopGeneratingButton } from './buttons/stop_generating_button'; import { RegenerateResponseButton } from './buttons/regenerate_response_button'; import { MessagePanel } from './message_panel'; @@ -35,27 +35,6 @@ interface Chunk { choices: ChunkChoice[]; } -const cursorCss = ` - @keyframes blink { - 0% { - opacity: 1; - } - 50% { - opacity: 0; - } - 100% { - opacity: 1; - } - } - - animation: blink 1s infinite; - width: 10px; - height: 16px; - vertical-align: middle; - display: inline-block; - background: rgba(0, 0, 0, 0.25); -`; - function getMessageFromChunks(chunks: Chunk[]) { let message = ''; chunks.forEach((chunk) => { @@ -65,14 +44,21 @@ function getMessageFromChunks(chunks: Chunk[]) { } interface Props { + amendMessage: (message: string, index: number) => void; index: number; isLastComment: boolean; lastCommentRef: React.MutableRefObject; reader: ReadableStreamDefaultReader; } -export const StreamComment = ({ index, isLastComment, lastCommentRef, reader }: Props) => { - const response$ = useMemo( +export const StreamComment = ({ + amendMessage, + index, + isLastComment, + lastCommentRef, + reader, +}: Props) => { + const observer$ = useMemo( () => new Observable((observer) => { observer.next({ chunks: [], loading: true }); @@ -135,23 +121,58 @@ export const StreamComment = ({ index, isLastComment, lastCommentRef, reader }: reader.cancel(); }; }).pipe(concatMap((value) => of(value).pipe(delay(50)))), + [reader] ); + // const response = useObservable(observer$); + + const [pendingMessage, setPendingMessage] = useState(); + + // const [recalledMessages, setRecalledMessages] = useState(undefined); - const response = useObservable(response$); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(); + const [subscription, setSubscription] = useState(); - useEffect(() => {}, [response$]); + const onCompleteStream = useCallback(() => { + amendMessage(pendingMessage ?? '', index); + }, [amendMessage, index, pendingMessage]); - let content = response?.message ?? ''; + const [complete, setComplete] = useState(false); + + useEffect(() => { + if (complete) { + onCompleteStream(); + } + }, [complete, onCompleteStream]); + + useEffect(() => { + console.log('observer$', { observer$ }); + const newSubscription = observer$.pipe(share()).subscribe({ + next: ({ message, loading: isLoading }) => { + setLoading(isLoading); + + setPendingMessage(message); + }, + complete: () => { + console.log('on complete'); + setComplete(true); + setLoading(false); + }, + error: (err) => { + setError(err.message); + }, + }); + setSubscription(newSubscription); + }, [observer$]); let state: 'init' | 'loading' | 'streaming' | 'error' | 'complete' = 'init'; - if (response?.loading) { - state = content ? 'streaming' : 'loading'; - } else if (response && 'error' in response && response.error) { + if (loading) { + state = pendingMessage ? 'streaming' : 'loading'; + } else if (error) { state = 'error'; - content = response.error; - } else if (content) { + } else if (pendingMessage) { state = 'complete'; } @@ -163,18 +184,19 @@ export const StreamComment = ({ index, isLastComment, lastCommentRef, reader }: {}} /> } - error={response?.error ? new Error(response?.error) : undefined} + error={error ? new Error(error) : undefined} controls={ isLoading || isStreaming ? ( { + subscription?.unsubscribe(); + setComplete(true); console.log('stop generating'); - // subscription?.unsubscribe(); // setLoading(false); // setDisplayedMessages((prevMessages) => // prevMessages.concat({ From 6d14e5e1d2186ccae758357bc4fd7e2dcdf3283b Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 19 Oct 2023 15:38:53 -0600 Subject: [PATCH 13/52] rm logs --- x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx | 1 - .../actions/server/lib/get_token_count_from_openai_stream.ts | 1 - 2 files changed, 2 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index 08bc75d0cda70..42e5669b1d1d6 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -80,7 +80,6 @@ export const fetchConnectorExecuteAction = async ({ const reader = response?.response?.body?.getReader(); - console.log('is typeof incoming message', reader); if (!reader) { return { response: `${API_ERROR}\n\nCould not get reader from response`, diff --git a/x-pack/plugins/actions/server/lib/get_token_count_from_openai_stream.ts b/x-pack/plugins/actions/server/lib/get_token_count_from_openai_stream.ts index 2de568a4d66d0..74c89f716171e 100644 --- a/x-pack/plugins/actions/server/lib/get_token_count_from_openai_stream.ts +++ b/x-pack/plugins/actions/server/lib/get_token_count_from_openai_stream.ts @@ -64,7 +64,6 @@ export async function getTokenCountFromOpenAIStream({ }); try { - console.log('finished', responseStream); await finished(responseStream); } catch { // no need to handle this explicitly From b821018c457f4c4256cea38eedc981b90dd2ff80 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 23 Oct 2023 10:09:31 -0600 Subject: [PATCH 14/52] wip --- .../impl/assistant/index.tsx | 20 ++- .../regenerate_response_button.stories.tsx | 0 .../buttons/regenerate_response_button.tsx | 0 .../stop_generating_button.stories.tsx | 0 .../buttons/stop_generating_button.tsx | 0 .../{ => stream}/esql_code_block.tsx | 0 .../{ => stream}/failed_to_load_response.tsx | 0 .../assistant/get_comments/stream/index.tsx | 90 +++++++++++++ .../{ => stream}/message_panel.tsx | 0 .../{ => stream}/message_text.tsx | 0 .../{stream.tsx => stream/use_stream.tsx} | 122 +++++------------- 11 files changed, 136 insertions(+), 96 deletions(-) rename x-pack/plugins/security_solution/public/assistant/get_comments/{ => stream}/buttons/regenerate_response_button.stories.tsx (100%) rename x-pack/plugins/security_solution/public/assistant/get_comments/{ => stream}/buttons/regenerate_response_button.tsx (100%) rename x-pack/plugins/security_solution/public/assistant/get_comments/{ => stream}/buttons/stop_generating_button.stories.tsx (100%) rename x-pack/plugins/security_solution/public/assistant/get_comments/{ => stream}/buttons/stop_generating_button.tsx (100%) rename x-pack/plugins/security_solution/public/assistant/get_comments/{ => stream}/esql_code_block.tsx (100%) rename x-pack/plugins/security_solution/public/assistant/get_comments/{ => stream}/failed_to_load_response.tsx (100%) create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx rename x-pack/plugins/security_solution/public/assistant/get_comments/{ => stream}/message_panel.tsx (100%) rename x-pack/plugins/security_solution/public/assistant/get_comments/{ => stream}/message_text.tsx (100%) rename x-pack/plugins/security_solution/public/assistant/get_comments/{stream.tsx => stream/use_stream.tsx} (62%) 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 dd7e4462f6f98..bd422adec563d 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -321,12 +321,20 @@ const AssistantComponent: React.FC = ({ const createCodeBlockPortals = useCallback( () => - messageCodeBlocks?.map((codeBlocks: CodeBlockDetails[]) => { - return codeBlocks.map((codeBlock: CodeBlockDetails) => { - const getElement = codeBlock.getControlContainer; - const element = getElement?.(); - return element ? createPortal(codeBlock.button, element) : <>; - }); + messageCodeBlocks?.map((codeBlocks: CodeBlockDetails[], i: number) => { + return ( + + {codeBlocks.map((codeBlock: CodeBlockDetails, j: number) => { + const getElement = codeBlock.getControlContainer; + const element = getElement?.(); + return ( + + {element ? createPortal(codeBlock.button, element) : <>} + + ); + })} + + ); }), [messageCodeBlocks] ); diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.stories.tsx similarity index 100% rename from x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.stories.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.stories.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.tsx similarity index 100% rename from x-pack/plugins/security_solution/public/assistant/get_comments/buttons/regenerate_response_button.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.stories.tsx similarity index 100% rename from x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.stories.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.stories.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.tsx similarity index 100% rename from x-pack/plugins/security_solution/public/assistant/get_comments/buttons/stop_generating_button.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/esql_code_block.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/esql_code_block.tsx similarity index 100% rename from x-pack/plugins/security_solution/public/assistant/get_comments/esql_code_block.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/esql_code_block.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/failed_to_load_response.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/failed_to_load_response.tsx similarity index 100% rename from x-pack/plugins/security_solution/public/assistant/get_comments/failed_to_load_response.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/failed_to_load_response.tsx 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 new file mode 100644 index 0000000000000..ed8ba5ecaf5ce --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.tsx @@ -0,0 +1,90 @@ +/* + * 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 React from 'react'; +import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui'; +import { useStream } from './use_stream'; +import { StopGeneratingButton } from './buttons/stop_generating_button'; +import { RegenerateResponseButton } from './buttons/regenerate_response_button'; +import { MessagePanel } from './message_panel'; +import { MessageText } from './message_text'; + +interface Props { + amendMessage: (message: string, index: number) => void; + index: number; + isLastComment: boolean; + lastCommentRef: React.MutableRefObject; + reader: ReadableStreamDefaultReader; +} + +export const StreamComment = ({ + amendMessage, + index, + isLastComment, + lastCommentRef, + reader, +}: Props) => { + const { subscription, setComplete, isLoading, isStreaming, pendingMessage, error } = useStream({ + amendMessage, + index, + reader, + }); + return ( + <> + {}} + /> + } + error={error ? new Error(error) : undefined} + controls={ + isLoading || isStreaming ? ( + { + subscription?.unsubscribe(); + setComplete(true); + console.log('stop generating'); + // setLoading(false); + // setDisplayedMessages((prevMessages) => + // prevMessages.concat({ + // '@timestamp': new Date().toISOString(), + // message: { + // ...pendingMessage!.message, + // }, + // }) + // ); + // setPendingMessage((prev) => ({ + // message: { + // role: MessageRole.Assistant, + // ...prev?.message, + // }, + // aborted: true, + // error: new AbortError(), + // })); + }} + /> + ) : ( + + + { + console.log('RegenerateResponseButton'); + // reloadRecalledMessages(); + }} + /> + + + ) + } + /> + {isLastComment && } + + ); +}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/message_panel.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_panel.tsx similarity index 100% rename from x-pack/plugins/security_solution/public/assistant/get_comments/message_panel.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_panel.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx similarity index 100% rename from x-pack/plugins/security_solution/public/assistant/get_comments/message_text.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx similarity index 62% rename from x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx rename to x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx index 895510ffef8ea..700fb8a97c6da 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx @@ -5,14 +5,23 @@ * 2.0. */ -import React, { useCallback, useEffect, useMemo, useState } from 'react'; -import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui'; +import { useCallback, useEffect, useMemo, useState } from 'react'; import type { Subscription } from 'rxjs'; import { concatMap, delay, Observable, of, share } from 'rxjs'; -import { StopGeneratingButton } from './buttons/stop_generating_button'; -import { RegenerateResponseButton } from './buttons/regenerate_response_button'; -import { MessagePanel } from './message_panel'; -import { MessageText } from './message_text'; +interface UseStreamProps { + amendMessage: (message: string, index: number) => void; + index: number; + reader: ReadableStreamDefaultReader; +} + +interface UseStream { + error: string | undefined; + isLoading: boolean; + isStreaming: boolean; + pendingMessage: string; + setComplete: (complete: boolean) => void; + subscription: Subscription | undefined; +} export interface PromptObservableState { chunks: Chunk[]; @@ -35,29 +44,7 @@ interface Chunk { choices: ChunkChoice[]; } -function getMessageFromChunks(chunks: Chunk[]) { - let message = ''; - chunks.forEach((chunk) => { - message += chunk.choices[0]?.delta.content ?? ''; - }); - return message; -} - -interface Props { - amendMessage: (message: string, index: number) => void; - index: number; - isLastComment: boolean; - lastCommentRef: React.MutableRefObject; - reader: ReadableStreamDefaultReader; -} - -export const StreamComment = ({ - amendMessage, - index, - isLastComment, - lastCommentRef, - reader, -}: Props) => { +export const useStream = ({ amendMessage, index, reader }: UseStreamProps): UseStream => { const observer$ = useMemo( () => new Observable((observer) => { @@ -124,12 +111,8 @@ export const StreamComment = ({ [reader] ); - // const response = useObservable(observer$); const [pendingMessage, setPendingMessage] = useState(); - - // const [recalledMessages, setRecalledMessages] = useState(undefined); - const [loading, setLoading] = useState(false); const [error, setError] = useState(); const [subscription, setSubscription] = useState(); @@ -147,7 +130,6 @@ export const StreamComment = ({ }, [complete, onCompleteStream]); useEffect(() => { - console.log('observer$', { observer$ }); const newSubscription = observer$.pipe(share()).subscribe({ next: ({ message, loading: isLoading }) => { setLoading(isLoading); @@ -155,7 +137,6 @@ export const StreamComment = ({ setPendingMessage(message); }, complete: () => { - console.log('on complete'); setComplete(true); setLoading(false); }, @@ -178,59 +159,20 @@ export const StreamComment = ({ const isLoading = state === 'init' || state === 'loading'; const isStreaming = state === 'streaming'; - - return ( - <> - {}} - /> - } - error={error ? new Error(error) : undefined} - controls={ - isLoading || isStreaming ? ( - { - subscription?.unsubscribe(); - setComplete(true); - console.log('stop generating'); - // setLoading(false); - // setDisplayedMessages((prevMessages) => - // prevMessages.concat({ - // '@timestamp': new Date().toISOString(), - // message: { - // ...pendingMessage!.message, - // }, - // }) - // ); - // setPendingMessage((prev) => ({ - // message: { - // role: MessageRole.Assistant, - // ...prev?.message, - // }, - // aborted: true, - // error: new AbortError(), - // })); - }} - /> - ) : ( - - - { - console.log('RegenerateResponseButton'); - // reloadRecalledMessages(); - }} - /> - - - ) - } - /> - {isLastComment && } - - ); + return { + error, + isLoading, + isStreaming, + pendingMessage: pendingMessage ?? '', + setComplete, + subscription, + }; }; + +function getMessageFromChunks(chunks: Chunk[]) { + let message = ''; + chunks.forEach((chunk) => { + message += chunk.choices[0]?.delta.content ?? ''; + }); + return message; +} From 682d69a1644608349c372ef0ac41a6fd7ae16400 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 23 Oct 2023 11:27:22 -0600 Subject: [PATCH 15/52] use stream for all comments --- .../impl/assistant/use_conversation/index.tsx | 14 +- .../public/assistant/get_comments/index.tsx | 85 ++++----- .../assistant/get_comments/stream/index.tsx | 100 ++++++----- .../get_comments/stream/use_stream.tsx | 169 +++++++++--------- 4 files changed, 183 insertions(+), 185 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx index 745cfe0a8168e..3baaa67b58b26 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx @@ -35,6 +35,10 @@ export interface AppendMessageProps { conversationId: string; message: Message; } +interface AmendMessageProps { + conversationId: string; + content: string; +} interface AppendReplacementsProps { conversationId: string; @@ -57,8 +61,8 @@ interface SetConversationProps { interface UseConversation { appendStreamMessage: ({ conversationId, message }: AppendMessageProps) => void; - appendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[]; - amendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[]; + appendMessage: ({ conversationId, message }: AppendMessageProps) => Message[]; + amendMessage: ({ conversationId, content }: AmendMessageProps) => Message[]; appendReplacements: ({ conversationId, replacements, @@ -77,14 +81,14 @@ export const useConversation = (): UseConversation => { * Replaces the last message of conversation[] for a given conversationId */ const amendMessage = useCallback( - ({ conversationId, message }) => { + ({ conversationId, content }) => { let messages: Message[] = []; setConversations((prev: Record) => { const prevConversation: Conversation | undefined = prev[conversationId]; if (prevConversation != null) { - prevConversation.messages.pop(); - messages = [...prevConversation.messages, message]; + const message = prevConversation.messages.pop(); + messages = [...prevConversation.messages, { ...message, content }]; const newConversation = { ...prevConversation, messages, 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 6252c64bd875f..92b46ff23dfb0 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 @@ -7,7 +7,7 @@ import type { EuiCommentProps } from '@elastic/eui'; import type { Conversation, Message } from '@kbn/elastic-assistant'; -import { EuiAvatar, EuiMarkdownFormat, EuiText, tint } from '@elastic/eui'; +import { EuiAvatar, tint } from '@elastic/eui'; import React from 'react'; import { AssistantAvatar } from '@kbn/elastic-assistant'; @@ -26,25 +26,19 @@ export const getComments = ({ }: { amendMessage: ({ conversationId, - message, + content, }: { conversationId: string; - message: Message; + content: string; }) => Message[]; currentConversation: Conversation; lastCommentRef: React.MutableRefObject; showAnonymizedValues: boolean; }): EuiCommentProps[] => { - const amendMessageOfConversation = (content: string, index: number) => { - const replacementMessage = currentConversation.messages[index]; - console.log('replacementMessage', { replacementMessage, content }); + const amendMessageOfConversation = (content: string) => { amendMessage({ conversationId: currentConversation.id, - message: { - content, - role: replacementMessage.role, - timestamp: replacementMessage.timestamp, - }, + content, }); }; @@ -77,40 +71,7 @@ export const getComments = ({ ...(message.isError ? errorStyles : {}), }; - if (message.content && message.content.length) { - const messageContentWithReplacements = - replacements != null - ? Object.keys(replacements).reduce( - (acc, replacement) => acc.replaceAll(replacement, replacements[replacement]), - message.content - ) - : message.content; - const transformedMessage = { - ...message, - content: messageContentWithReplacements, - }; - - return { - ...messageProps, - actions: , - children: - index !== currentConversation.messages.length - 1 ? ( - - - {showAnonymizedValues ? message.content : transformedMessage.content} - - - ) : ( - - - {showAnonymizedValues ? message.content : transformedMessage.content} - - - - ), - }; - } - if (message.reader) { + if (!(message.content && message.content.length)) { return { ...messageProps, children: ( @@ -124,10 +85,40 @@ export const getComments = ({ ), }; } + + const messageContentWithReplacements = + replacements != null + ? Object.keys(replacements).reduce( + (acc, replacement) => acc.replaceAll(replacement, replacements[replacement]), + message.content + ) + : message.content; + const transformedMessage = { + ...message, + content: messageContentWithReplacements, + }; + return { ...messageProps, - children: {'oh no an error happened'}, - ...errorStyles, + actions: , + children: ( + <> + + {index !== currentConversation.messages.length - 1 ? null : } + + ), }; + + // return { + // ...messageProps, + // children: {'oh no an error happened'}, + // ...errorStyles, + // }; }); }; 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 ed8ba5ecaf5ce..42b3a59b66bbe 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 @@ -5,7 +5,7 @@ * 2.0. */ -import React from 'react'; +import React, { useMemo } from 'react'; import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui'; import { useStream } from './use_stream'; import { StopGeneratingButton } from './buttons/stop_generating_button'; @@ -14,75 +14,81 @@ import { MessagePanel } from './message_panel'; import { MessageText } from './message_text'; interface Props { - amendMessage: (message: string, index: number) => void; - index: number; + amendMessage: (message: string) => void; + content?: string; isLastComment: boolean; lastCommentRef: React.MutableRefObject; - reader: ReadableStreamDefaultReader; + reader?: ReadableStreamDefaultReader; } export const StreamComment = ({ amendMessage, - index, + content, isLastComment, lastCommentRef, reader, }: Props) => { - const { subscription, setComplete, isLoading, isStreaming, pendingMessage, error } = useStream({ + const { error, isLoading, isStreaming, pendingMessage, setComplete, subscription } = useStream({ amendMessage, - index, + content, reader, }); + const message = content ?? pendingMessage; + const controls = useMemo( + () => + reader != null ? ( + isLoading || isStreaming ? ( + { + subscription?.unsubscribe(); + setComplete(true); + console.log('stop generating'); + // setLoading(false); + // setDisplayedMessages((prevMessages) => + // prevMessages.concat({ + // '@timestamp': new Date().toISOString(), + // message: { + // ...pendingMessage!.message, + // }, + // }) + // ); + // setPendingMessage((prev) => ({ + // message: { + // role: MessageRole.Assistant, + // ...prev?.message, + // }, + // aborted: true, + // error: new AbortError(), + // })); + }} + /> + ) : isLastComment ? ( + + + { + console.log('RegenerateResponseButton'); + // reloadRecalledMessages(); + }} + /> + + + ) : null + ) : null, + [isLastComment, isLoading, isStreaming, reader, setComplete, subscription] + ); return ( <> {}} /> } error={error ? new Error(error) : undefined} - controls={ - isLoading || isStreaming ? ( - { - subscription?.unsubscribe(); - setComplete(true); - console.log('stop generating'); - // setLoading(false); - // setDisplayedMessages((prevMessages) => - // prevMessages.concat({ - // '@timestamp': new Date().toISOString(), - // message: { - // ...pendingMessage!.message, - // }, - // }) - // ); - // setPendingMessage((prev) => ({ - // message: { - // role: MessageRole.Assistant, - // ...prev?.message, - // }, - // aborted: true, - // error: new AbortError(), - // })); - }} - /> - ) : ( - - - { - console.log('RegenerateResponseButton'); - // reloadRecalledMessages(); - }} - /> - - - ) - } + controls={controls} /> {isLastComment && } diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx index 700fb8a97c6da..9ec0565f16a2e 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx @@ -8,20 +8,6 @@ import { useCallback, useEffect, useMemo, useState } from 'react'; import type { Subscription } from 'rxjs'; import { concatMap, delay, Observable, of, share } from 'rxjs'; -interface UseStreamProps { - amendMessage: (message: string, index: number) => void; - index: number; - reader: ReadableStreamDefaultReader; -} - -interface UseStream { - error: string | undefined; - isLoading: boolean; - isStreaming: boolean; - pendingMessage: string; - setComplete: (complete: boolean) => void; - subscription: Subscription | undefined; -} export interface PromptObservableState { chunks: Chunk[]; @@ -44,72 +30,91 @@ interface Chunk { choices: ChunkChoice[]; } -export const useStream = ({ amendMessage, index, reader }: UseStreamProps): UseStream => { - const observer$ = useMemo( - () => - new Observable((observer) => { - observer.next({ chunks: [], loading: true }); - - const decoder = new TextDecoder(); - - const chunks: Chunk[] = []; - - let prev: string = ''; - - function read() { - reader.read().then(({ done, value }: { done: boolean; value?: Uint8Array }) => { - try { - if (done) { - observer.next({ - chunks, - message: getMessageFromChunks(chunks), - loading: false, - }); - observer.complete(); - return; - } - - let lines: string[] = (prev + decoder.decode(value)).split('\n'); - - const lastLine: string = lines[lines.length - 1]; - - const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; - - if (isPartialChunk) { - prev = lastLine; - lines.pop(); - } else { - prev = ''; - } - - lines = lines.map((str) => str.substr(6)).filter((str) => !!str && str !== '[DONE]'); +interface UseStreamProps { + amendMessage: (message: string) => void; + content?: string; + reader?: ReadableStreamDefaultReader; +} - const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); +interface UseStream { + error: string | undefined; + isLoading: boolean; + isStreaming: boolean; + pendingMessage: string; + setComplete: (complete: boolean) => void; + subscription: Subscription | undefined; +} - nextChunks.forEach((chunk) => { - chunks.push(chunk); - observer.next({ - chunks, - message: getMessageFromChunks(chunks), - loading: true, - }); +export const useStream = ({ amendMessage, content, reader }: UseStreamProps): UseStream => { + const observer$ = useMemo( + () => + content == null && reader != null + ? new Observable((observer) => { + observer.next({ chunks: [], loading: true }); + + const decoder = new TextDecoder(); + + const chunks: Chunk[] = []; + + let prev: string = ''; + + function read() { + reader?.read().then(({ done, value }: { done: boolean; value?: Uint8Array }) => { + try { + if (done) { + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: false, + }); + observer.complete(); + return; + } + + let lines: string[] = (prev + decoder.decode(value)).split('\n'); + + const lastLine: string = lines[lines.length - 1]; + + const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; + + if (isPartialChunk) { + prev = lastLine; + lines.pop(); + } else { + prev = ''; + } + + lines = lines + .map((str) => str.substr(6)) + .filter((str) => !!str && str !== '[DONE]'); + + const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); + + nextChunks.forEach((chunk) => { + chunks.push(chunk); + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: true, + }); + }); + } catch (err) { + observer.error(err); + return; + } + read(); }); - } catch (err) { - observer.error(err); - return; } - read(); - }); - } - read(); + read(); - return () => { - reader.cancel(); - }; - }).pipe(concatMap((value) => of(value).pipe(delay(50)))), + return () => { + reader.cancel(); + }; + }).pipe(concatMap((value) => of(value).pipe(delay(50)))) + : new Observable(), - [reader] + [content, reader] ); const [pendingMessage, setPendingMessage] = useState(); @@ -118,8 +123,8 @@ export const useStream = ({ amendMessage, index, reader }: UseStreamProps): UseS const [subscription, setSubscription] = useState(); const onCompleteStream = useCallback(() => { - amendMessage(pendingMessage ?? '', index); - }, [amendMessage, index, pendingMessage]); + amendMessage(pendingMessage ?? ''); + }, [amendMessage, pendingMessage]); const [complete, setComplete] = useState(false); @@ -147,18 +152,10 @@ export const useStream = ({ amendMessage, index, reader }: UseStreamProps): UseS setSubscription(newSubscription); }, [observer$]); - let state: 'init' | 'loading' | 'streaming' | 'error' | 'complete' = 'init'; - - if (loading) { - state = pendingMessage ? 'streaming' : 'loading'; - } else if (error) { - state = 'error'; - } else if (pendingMessage) { - state = 'complete'; - } + const { isLoading, isStreaming } = useMemo(() => { + return { isLoading: loading, isStreaming: loading && pendingMessage != null }; + }, [loading, pendingMessage]); - const isLoading = state === 'init' || state === 'loading'; - const isStreaming = state === 'streaming'; return { error, isLoading, From 4fe7b815192386a4209cbeb8529ec3941c2dbb9a Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 23 Oct 2023 14:44:08 -0600 Subject: [PATCH 16/52] regenerate button done --- .../assistant/chat_send/use_chat_send.tsx | 23 ++++- .../impl/assistant/index.tsx | 15 +++ .../impl/assistant/use_conversation/index.tsx | 40 +++++++- .../impl/assistant_context/index.tsx | 10 +- .../public/assistant/get_comments/index.tsx | 36 +++---- .../assistant/get_comments/stream/index.tsx | 95 +++++++------------ .../get_comments/stream/use_stream.tsx | 17 ++-- 7 files changed, 142 insertions(+), 94 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index 3e1c194097888..c269ed76cfd45 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -34,6 +34,7 @@ interface UseChatSend { handleOnChatCleared: () => void; handlePromptChange: (prompt: string) => void; handleSendMessage: (promptText: string) => void; + handleRegenerateResponse: () => void; isLoading: boolean; } @@ -54,7 +55,8 @@ export const useChatSend = ({ setUserPrompt, }: UseChatSendProps): UseChatSend => { const { isLoading, sendMessages } = useSendMessages(); - const { appendMessage, appendReplacements, clearConversation } = useConversation(); + const { appendMessage, appendReplacements, clearConversation, removeLastMessage } = + useConversation(); const handlePromptChange = (prompt: string) => { setPromptTextPreview(prompt); @@ -112,6 +114,24 @@ export const useChatSend = ({ ] ); + const handleRegenerateResponse = useCallback(async () => { + const updatedMessages = removeLastMessage(currentConversation.id); + const rawResponse = await sendMessages({ + http, + apiConfig: currentConversation.apiConfig, + messages: updatedMessages, + }); + const responseMessage: Message = getMessageFromRawResponse(rawResponse); + appendMessage({ conversationId: currentConversation.id, message: responseMessage }); + }, [ + appendMessage, + currentConversation.apiConfig, + currentConversation.id, + http, + removeLastMessage, + sendMessages, + ]); + const handleButtonSendMessage = useCallback( (message: string) => { handleSendMessage(message); @@ -146,6 +166,7 @@ export const useChatSend = ({ handleOnChatCleared, handlePromptChange, handleSendMessage, + handleRegenerateResponse, isLoading, }; }; 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 bd422adec563d..ff0e0536ea309 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -31,6 +31,7 @@ import { css } from '@emotion/react'; import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/common/openai/constants'; import { ActionConnectorProps } from '@kbn/triggers-actions-ui-plugin/public/types'; +import { useChatSend } from './chat_send/use_chat_send'; import { ChatSend } from './chat_send'; import { BlockBotCallToAction } from './block_bot/cta'; import { AssistantHeader } from './assistant_header'; @@ -339,6 +340,18 @@ const AssistantComponent: React.FC = ({ [messageCodeBlocks] ); + const { handleRegenerateResponse } = useChatSend({ + allSystemPrompts, + currentConversation, + setPromptTextPreview, + setUserPrompt, + editingSystemPromptId, + http, + setEditingSystemPromptId, + selectedPromptContexts, + setSelectedPromptContexts, + }); + const chatbotComments = useMemo( () => ( <> @@ -348,6 +361,7 @@ const AssistantComponent: React.FC = ({ lastCommentRef, showAnonymizedValues, amendMessage, + regenerateMessage: handleRegenerateResponse, })} css={css` margin-right: 20px; @@ -382,6 +396,7 @@ const AssistantComponent: React.FC = ({ editingSystemPromptId, getComments, handleOnSystemPromptSelectionChange, + handleRegenerateResponse, isSettingsModalVisible, promptContexts, promptTextPreview, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx index 3baaa67b58b26..a7ce6ef4ded7a 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx @@ -60,7 +60,6 @@ interface SetConversationProps { } interface UseConversation { - appendStreamMessage: ({ conversationId, message }: AppendMessageProps) => void; appendMessage: ({ conversationId, message }: AppendMessageProps) => Message[]; amendMessage: ({ conversationId, content }: AmendMessageProps) => Message[]; appendReplacements: ({ @@ -70,6 +69,7 @@ interface UseConversation { clearConversation: (conversationId: string) => void; createConversation: ({ conversationId, messages }: CreateConversationProps) => Conversation; deleteConversation: (conversationId: string) => void; + removeLastMessage: (conversationId: string) => Message[]; setApiConfig: ({ conversationId, apiConfig }: SetApiConfigProps) => void; setConversation: ({ conversation }: SetConversationProps) => void; } @@ -78,21 +78,52 @@ export const useConversation = (): UseConversation => { const { allSystemPrompts, assistantTelemetry, setConversations } = useAssistantContext(); /** - * Replaces the last message of conversation[] for a given conversationId + * Removes the last message of conversation[] for a given conversationId + */ + const removeLastMessage = useCallback( + (conversationId: string) => { + let messages: Message[] = []; + setConversations((prev: Record) => { + const prevConversation: Conversation | undefined = prev[conversationId]; + + if (prevConversation != null) { + prevConversation.messages.pop(); + messages = prevConversation.messages; + const newConversation = { + ...prevConversation, + messages, + }; + console.log('removeLastMessage newConversation', newConversation); + return { + ...prev, + [conversationId]: newConversation, + }; + } else { + return prev; + } + }); + return messages; + }, + [setConversations] + ); + + /** + * Updates the last message of conversation[] for a given conversationId with provided content */ const amendMessage = useCallback( - ({ conversationId, content }) => { + ({ conversationId, content }: AmendMessageProps) => { let messages: Message[] = []; setConversations((prev: Record) => { const prevConversation: Conversation | undefined = prev[conversationId]; if (prevConversation != null) { - const message = prevConversation.messages.pop(); + const message = prevConversation.messages.pop() as unknown as Message; messages = [...prevConversation.messages, { ...message, content }]; const newConversation = { ...prevConversation, messages, }; + console.log('amendMessage newConversation', newConversation); return { ...prev, [conversationId]: newConversation, @@ -303,6 +334,7 @@ export const useConversation = (): UseConversation => { clearConversation, createConversation, deleteConversation, + removeLastMessage, setApiConfig, setConversation, }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index 6c99970d01a97..553c42278abed 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -13,7 +13,6 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react'; import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public'; import { useLocalStorage } from 'react-use'; import type { DocLinksStart } from '@kbn/core-doc-links-browser'; -import { AppendMessageProps } from '../assistant/use_conversation'; import { updatePromptContexts } from './helpers'; import type { PromptContext, @@ -109,7 +108,14 @@ export interface UseAssistantContext { }: { currentConversation: Conversation; lastCommentRef: React.MutableRefObject; - amendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[]; + amendMessage: ({ + conversationId, + content, + }: { + conversationId: string; + content: string; + }) => Message[]; + regenerateMessage: () => void; showAnonymizedValues: boolean; }) => EuiCommentProps[]; http: HttpSetup; 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 92b46ff23dfb0..458f02799ade8 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 @@ -22,6 +22,7 @@ export const getComments = ({ amendMessage, currentConversation, lastCommentRef, + regenerateMessage, showAnonymizedValues, }: { amendMessage: ({ @@ -33,6 +34,7 @@ export const getComments = ({ }) => Message[]; currentConversation: Conversation; lastCommentRef: React.MutableRefObject; + regenerateMessage: (conversationId: string) => void; showAnonymizedValues: boolean; }): EuiCommentProps[] => { const amendMessageOfConversation = (content: string) => { @@ -42,6 +44,10 @@ export const getComments = ({ }); }; + const regenerateMessageOfConversation = () => { + regenerateMessage(currentConversation.id); + }; + return currentConversation.messages.map((message, index) => { const isUser = message.role === 'user'; const replacements = currentConversation.replacements; @@ -70,18 +76,22 @@ export const getComments = ({ username: isUser ? i18n.YOU : i18n.ASSISTANT, ...(message.isError ? errorStyles : {}), }; + const isLastComment = index === currentConversation.messages.length - 1; + // message still needs to stream, no response manipulation if (!(message.content && message.content.length)) { return { ...messageProps, children: ( - + <> + + {isLastComment ? : null} + ), }; } @@ -107,18 +117,12 @@ export const getComments = ({ amendMessage={amendMessageOfConversation} content={showAnonymizedValues ? message.content : transformedMessage.content} reader={message.reader} - lastCommentRef={lastCommentRef} - isLastComment={index === currentConversation.messages.length - 1} + regenerateMessage={regenerateMessageOfConversation} + isLastComment={isLastComment} /> - {index !== currentConversation.messages.length - 1 ? null : } + {isLastComment ? : null} ), }; - - // return { - // ...messageProps, - // children: {'oh no an error happened'}, - // ...errorStyles, - // }; }); }; 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 42b3a59b66bbe..fed130334355f 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,7 +17,7 @@ interface Props { amendMessage: (message: string) => void; content?: string; isLastComment: boolean; - lastCommentRef: React.MutableRefObject; + regenerateMessage: () => void; reader?: ReadableStreamDefaultReader; } @@ -25,72 +25,47 @@ export const StreamComment = ({ amendMessage, content, isLastComment, - lastCommentRef, reader, + regenerateMessage, }: Props) => { - const { error, isLoading, isStreaming, pendingMessage, setComplete, subscription } = useStream({ + const { error, isLoading, isStreaming, pendingMessage, setComplete } = useStream({ amendMessage, content, reader, }); const message = content ?? pendingMessage; - const controls = useMemo( - () => - reader != null ? ( - isLoading || isStreaming ? ( - { - subscription?.unsubscribe(); - setComplete(true); - console.log('stop generating'); - // setLoading(false); - // setDisplayedMessages((prevMessages) => - // prevMessages.concat({ - // '@timestamp': new Date().toISOString(), - // message: { - // ...pendingMessage!.message, - // }, - // }) - // ); - // setPendingMessage((prev) => ({ - // message: { - // role: MessageRole.Assistant, - // ...prev?.message, - // }, - // aborted: true, - // error: new AbortError(), - // })); - }} - /> - ) : isLastComment ? ( - - - { - console.log('RegenerateResponseButton'); - // reloadRecalledMessages(); - }} - /> - - - ) : null - ) : null, - [isLastComment, isLoading, isStreaming, reader, setComplete, subscription] - ); + const controls = useMemo(() => { + if (reader == null || !isLastComment) { + return; + } + if (isLoading || isStreaming) { + return ( + { + setComplete(true); + }} + /> + ); + } + return ( + + + + + + ); + }, [isLastComment, isLoading, isStreaming, reader, regenerateMessage, setComplete]); return ( - <> - {}} - /> - } - error={error ? new Error(error) : undefined} - controls={controls} - /> - {isLastComment && } - + {}} + /> + } + error={error ? new Error(error) : undefined} + controls={controls} + /> ); }; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx index 9ec0565f16a2e..96e6c3121e73d 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx @@ -42,7 +42,6 @@ interface UseStream { isStreaming: boolean; pendingMessage: string; setComplete: (complete: boolean) => void; - subscription: Subscription | undefined; } export const useStream = ({ amendMessage, content, reader }: UseStreamProps): UseStream => { @@ -113,7 +112,6 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us }; }).pipe(concatMap((value) => of(value).pipe(delay(50)))) : new Observable(), - [content, reader] ); @@ -123,13 +121,16 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us const [subscription, setSubscription] = useState(); const onCompleteStream = useCallback(() => { + subscription?.unsubscribe(); + setLoading(false); amendMessage(pendingMessage ?? ''); - }, [amendMessage, pendingMessage]); + }, [amendMessage, pendingMessage, subscription]); const [complete, setComplete] = useState(false); useEffect(() => { if (complete) { + setComplete(false); onCompleteStream(); } }, [complete, onCompleteStream]); @@ -138,7 +139,6 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us const newSubscription = observer$.pipe(share()).subscribe({ next: ({ message, loading: isLoading }) => { setLoading(isLoading); - setPendingMessage(message); }, complete: () => { @@ -152,17 +152,12 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us setSubscription(newSubscription); }, [observer$]); - const { isLoading, isStreaming } = useMemo(() => { - return { isLoading: loading, isStreaming: loading && pendingMessage != null }; - }, [loading, pendingMessage]); - return { error, - isLoading, - isStreaming, + isLoading: loading, + isStreaming: loading && pendingMessage != null, pendingMessage: pendingMessage ?? '', setComplete, - subscription, }; }; From e8fc9f069c02380dc59e80c1db6c3b4118fa8417 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 24 Oct 2023 10:46:07 -0600 Subject: [PATCH 17/52] fixing --- .../impl/assistant/use_conversation/index.tsx | 2 -- .../assistant/comment_actions/index.tsx | 12 +++++++----- .../regenerate_response_button.stories.tsx | 19 ------------------- .../stop_generating_button.stories.tsx | 19 ------------------- .../get_comments/stream/message_text.tsx | 2 +- 5 files changed, 8 insertions(+), 46 deletions(-) delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.stories.tsx delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.stories.tsx diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx index a7ce6ef4ded7a..0f815db18409c 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx @@ -93,7 +93,6 @@ export const useConversation = (): UseConversation => { ...prevConversation, messages, }; - console.log('removeLastMessage newConversation', newConversation); return { ...prev, [conversationId]: newConversation, @@ -123,7 +122,6 @@ export const useConversation = (): UseConversation => { ...prevConversation, messages, }; - console.log('amendMessage newConversation', newConversation); return { ...prev, [conversationId]: newConversation, diff --git a/x-pack/plugins/security_solution/public/assistant/comment_actions/index.tsx b/x-pack/plugins/security_solution/public/assistant/comment_actions/index.tsx index b7cac488319dc..5a3e4a3d1c210 100644 --- a/x-pack/plugins/security_solution/public/assistant/comment_actions/index.tsx +++ b/x-pack/plugins/security_solution/public/assistant/comment_actions/index.tsx @@ -41,17 +41,19 @@ const CommentActionsComponent: React.FC = ({ message }) => { [dispatch] ); + const content = message.content ?? ''; + const onAddNoteToTimeline = useCallback(() => { updateAndAssociateNode({ associateNote, - newNote: message.content, + newNote: content, updateNewNote: () => {}, updateNote, user: '', // TODO: attribute assistant messages }); toasts.addSuccess(i18n.ADDED_NOTE_TO_TIMELINE); - }, [associateNote, message.content, toasts, updateNote]); + }, [associateNote, content, toasts, updateNote]); // Attach to case support const selectCaseModal = cases.hooks.useCasesAddToExistingCaseModal({ @@ -65,13 +67,13 @@ const CommentActionsComponent: React.FC = ({ message }) => { selectCaseModal.open({ getAttachments: () => [ { - comment: message.content, + comment: content, type: AttachmentType.user, owner: i18n.ELASTIC_AI_ASSISTANT, }, ], }); - }, [message.content, selectCaseModal, showAssistantOverlay]); + }, [content, selectCaseModal, showAssistantOverlay]); return ( @@ -99,7 +101,7 @@ const CommentActionsComponent: React.FC = ({ message }) => { - + {(copy) => ( = { - component: Component, - title: 'app/Atoms/RegenerateResponseButton', -}; - -export default meta; - -export const RegenerateResponseButton: ComponentStoryObj = { - args: {}, -}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.stories.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.stories.tsx deleted file mode 100644 index acf27b4a01274..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.stories.tsx +++ /dev/null @@ -1,19 +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 { ComponentMeta, ComponentStoryObj } from '@storybook/react'; -import { StopGeneratingButton as Component } from './stop_generating_button'; - -const meta: ComponentMeta = { - component: Component, - title: 'app/Atoms/StopGeneratingButton', -}; - -export default meta; - -export const StopGeneratingButton: ComponentStoryObj = { - args: {}, -}; 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 b76d24ce39381..9ef70caf8f524 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 @@ -77,7 +77,7 @@ const loadingCursorPlugin = () => { textNode.value = textNode.value.replace(CURSOR, ''); - const indexOfNode = parent?.children.indexOf(textNode); + const indexOfNode = parent?.children.indexOf(textNode) ?? 0; parent?.children.splice(indexOfNode + 1, 0, { type: 'cursor' as Text['type'], value: CURSOR, From a6c749e66b59f74bba4fdd225c2bf35a208b334f Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 24 Oct 2023 16:14:53 -0600 Subject: [PATCH 18/52] test --- .../assistant/get_comments/stream/mock.ts | 637 ++++++++++++++++++ .../get_comments/stream/use_stream.test.tsx | 139 ++++ .../get_comments/stream/use_stream.tsx | 95 ++- 3 files changed, 820 insertions(+), 51 deletions(-) create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/mock.ts create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/mock.ts b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/mock.ts new file mode 100644 index 0000000000000..ebba2d432d51c --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/mock.ts @@ -0,0 +1,637 @@ +/* + * 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. + */ + +// Uint8Array returned by OpenAI +const one = { + '0': 100, + '1': 97, + '2': 116, + '3': 97, + '4': 58, + '5': 32, + '6': 123, + '7': 34, + '8': 105, + '9': 100, + '10': 34, + '11': 58, + '12': 34, + '13': 99, + '14': 104, + '15': 97, + '16': 116, + '17': 99, + '18': 109, + '19': 112, + '20': 108, + '21': 45, + '22': 56, + '23': 68, + '24': 74, + '25': 71, + '26': 72, + '27': 117, + '28': 73, + '29': 67, + '30': 74, + '31': 68, + '32': 118, + '33': 103, + '34': 116, + '35': 71, + '36': 49, + '37': 98, + '38': 97, + '39': 65, + '40': 106, + '41': 74, + '42': 112, + '43': 103, + '44': 88, + '45': 120, + '46': 120, + '47': 90, + '48': 52, + '49': 73, + '50': 121, + '51': 34, + '52': 44, + '53': 34, + '54': 111, + '55': 98, + '56': 106, + '57': 101, + '58': 99, + '59': 116, + '60': 34, + '61': 58, + '62': 34, + '63': 99, + '64': 104, + '65': 97, + '66': 116, + '67': 46, + '68': 99, + '69': 111, + '70': 109, + '71': 112, + '72': 108, + '73': 101, + '74': 116, + '75': 105, + '76': 111, + '77': 110, + '78': 46, + '79': 99, + '80': 104, + '81': 117, + '82': 110, + '83': 107, + '84': 34, + '85': 44, + '86': 34, + '87': 99, + '88': 114, + '89': 101, + '90': 97, + '91': 116, + '92': 101, + '93': 100, + '94': 34, + '95': 58, + '96': 49, + '97': 54, + '98': 57, + '99': 56, + '100': 49, + '101': 56, + '102': 50, + '103': 57, + '104': 57, + '105': 55, + '106': 44, + '107': 34, + '108': 109, + '109': 111, + '110': 100, + '111': 101, + '112': 108, + '113': 34, + '114': 58, + '115': 34, + '116': 103, + '117': 112, + '118': 116, + '119': 45, + '120': 52, + '121': 45, + '122': 48, + '123': 54, + '124': 49, + '125': 51, + '126': 34, + '127': 44, + '128': 34, + '129': 99, + '130': 104, + '131': 111, + '132': 105, + '133': 99, + '134': 101, + '135': 115, + '136': 34, + '137': 58, + '138': 91, + '139': 123, + '140': 34, + '141': 105, + '142': 110, + '143': 100, + '144': 101, + '145': 120, + '146': 34, + '147': 58, + '148': 48, + '149': 44, + '150': 34, + '151': 100, + '152': 101, + '153': 108, + '154': 116, + '155': 97, + '156': 34, + '157': 58, + '158': 123, + '159': 34, + '160': 114, + '161': 111, + '162': 108, + '163': 101, + '164': 34, + '165': 58, + '166': 34, + '167': 97, + '168': 115, + '169': 115, + '170': 105, + '171': 115, + '172': 116, + '173': 97, + '174': 110, + '175': 116, + '176': 34, + '177': 44, + '178': 34, + '179': 99, + '180': 111, + '181': 110, + '182': 116, + '183': 101, + '184': 110, + '185': 116, + '186': 34, + '187': 58, + '188': 34, + '189': 34, + '190': 125, + '191': 44, + '192': 34, + '193': 102, + '194': 105, + '195': 110, + '196': 105, + '197': 115, + '198': 104, + '199': 95, + '200': 114, + '201': 101, + '202': 97, + '203': 115, + '204': 111, + '205': 110, + '206': 34, + '207': 58, + '208': 110, + '209': 117, + '210': 108, + '211': 108, + '212': 125, + '213': 93, + '214': 125, + '215': 10, + '216': 10, + '217': 100, + '218': 97, + '219': 116, + '220': 97, + '221': 58, + '222': 32, + '223': 123, + '224': 34, + '225': 105, + '226': 100, + '227': 34, + '228': 58, + '229': 34, + '230': 99, + '231': 104, + '232': 97, + '233': 116, + '234': 99, + '235': 109, + '236': 112, + '237': 108, + '238': 45, + '239': 56, + '240': 68, + '241': 74, + '242': 71, + '243': 72, + '244': 117, + '245': 73, + '246': 67, + '247': 74, + '248': 68, + '249': 118, + '250': 103, + '251': 116, + '252': 71, + '253': 49, + '254': 98, + '255': 97, + '256': 65, + '257': 106, + '258': 74, + '259': 112, + '260': 103, + '261': 88, + '262': 120, + '263': 120, + '264': 90, + '265': 52, + '266': 73, + '267': 121, + '268': 34, + '269': 44, + '270': 34, + '271': 111, + '272': 98, + '273': 106, + '274': 101, + '275': 99, + '276': 116, + '277': 34, + '278': 58, + '279': 34, + '280': 99, + '281': 104, + '282': 97, + '283': 116, + '284': 46, + '285': 99, + '286': 111, + '287': 109, + '288': 112, + '289': 108, + '290': 101, + '291': 116, + '292': 105, + '293': 111, + '294': 110, + '295': 46, + '296': 99, + '297': 104, + '298': 117, + '299': 110, + '300': 107, + '301': 34, + '302': 44, + '303': 34, + '304': 99, + '305': 114, + '306': 101, + '307': 97, + '308': 116, + '309': 101, + '310': 100, + '311': 34, + '312': 58, + '313': 49, + '314': 54, + '315': 57, + '316': 56, + '317': 49, + '318': 56, + '319': 50, + '320': 57, + '321': 57, + '322': 55, + '323': 44, + '324': 34, + '325': 109, + '326': 111, + '327': 100, + '328': 101, + '329': 108, + '330': 34, + '331': 58, + '332': 34, + '333': 103, + '334': 112, + '335': 116, + '336': 45, + '337': 52, + '338': 45, + '339': 48, + '340': 54, + '341': 49, + '342': 51, + '343': 34, + '344': 44, + '345': 34, + '346': 99, + '347': 104, + '348': 111, + '349': 105, + '350': 99, + '351': 101, + '352': 115, + '353': 34, + '354': 58, + '355': 91, + '356': 123, + '357': 34, + '358': 105, + '359': 110, + '360': 100, + '361': 101, + '362': 120, + '363': 34, + '364': 58, + '365': 48, + '366': 44, + '367': 34, + '368': 100, + '369': 101, + '370': 108, + '371': 116, + '372': 97, + '373': 34, + '374': 58, + '375': 123, + '376': 34, + '377': 99, + '378': 111, + '379': 110, + '380': 116, + '381': 101, + '382': 110, + '383': 116, + '384': 34, + '385': 58, + '386': 34, + '387': 67, + '388': 104, + '389': 34, + '390': 125, + '391': 44, + '392': 34, + '393': 102, + '394': 105, + '395': 110, + '396': 105, + '397': 115, + '398': 104, + '399': 95, + '400': 114, + '401': 101, + '402': 97, + '403': 115, + '404': 111, + '405': 110, + '406': 34, + '407': 58, + '408': 110, + '409': 117, + '410': 108, + '411': 108, + '412': 125, + '413': 93, + '414': 125, + '415': 10, + '416': 10, +}; +const two = { + '0': 100, + '1': 97, + '2': 116, + '3': 97, + '4': 58, + '5': 32, + '6': 123, + '7': 34, + '8': 105, + '9': 100, + '10': 34, + '11': 58, + '12': 34, + '13': 99, + '14': 104, + '15': 97, + '16': 116, + '17': 99, + '18': 109, + '19': 112, + '20': 108, + '21': 45, + '22': 56, + '23': 68, + '24': 74, + '25': 71, + '26': 72, + '27': 117, + '28': 73, + '29': 67, + '30': 74, + '31': 68, + '32': 118, + '33': 103, + '34': 116, + '35': 71, + '36': 49, + '37': 98, + '38': 97, + '39': 65, + '40': 106, + '41': 74, + '42': 112, + '43': 103, + '44': 88, + '45': 120, + '46': 120, + '47': 90, + '48': 52, + '49': 73, + '50': 121, + '51': 34, + '52': 44, + '53': 34, + '54': 111, + '55': 98, + '56': 106, + '57': 101, + '58': 99, + '59': 116, + '60': 34, + '61': 58, + '62': 34, + '63': 99, + '64': 104, + '65': 97, + '66': 116, + '67': 46, + '68': 99, + '69': 111, + '70': 109, + '71': 112, + '72': 108, + '73': 101, + '74': 116, + '75': 105, + '76': 111, + '77': 110, + '78': 46, + '79': 99, + '80': 104, + '81': 117, + '82': 110, + '83': 107, + '84': 34, + '85': 44, + '86': 34, + '87': 99, + '88': 114, + '89': 101, + '90': 97, + '91': 116, + '92': 101, + '93': 100, + '94': 34, + '95': 58, + '96': 49, + '97': 54, + '98': 57, + '99': 56, + '100': 49, + '101': 56, + '102': 50, + '103': 57, + '104': 57, + '105': 55, + '106': 44, + '107': 34, + '108': 109, + '109': 111, + '110': 100, + '111': 101, + '112': 108, + '113': 34, + '114': 58, + '115': 34, + '116': 103, + '117': 112, + '118': 116, + '119': 45, + '120': 52, + '121': 45, + '122': 48, + '123': 54, + '124': 49, + '125': 51, + '126': 34, + '127': 44, + '128': 34, + '129': 99, + '130': 104, + '131': 111, + '132': 105, + '133': 99, + '134': 101, + '135': 115, + '136': 34, + '137': 58, + '138': 91, + '139': 123, + '140': 34, + '141': 105, + '142': 110, + '143': 100, + '144': 101, + '145': 120, + '146': 34, + '147': 58, + '148': 48, + '149': 44, + '150': 34, + '151': 100, + '152': 101, + '153': 108, + '154': 116, + '155': 97, + '156': 34, + '157': 58, + '158': 123, + '159': 34, + '160': 99, + '161': 111, + '162': 110, + '163': 116, + '164': 101, + '165': 110, + '166': 116, + '167': 34, + '168': 58, + '169': 34, + '170': 101, + '171': 100, + '172': 100, + '173': 97, + '174': 114, + '175': 34, + '176': 125, + '177': 44, + '178': 34, + '179': 102, + '180': 105, + '181': 110, + '182': 105, + '183': 115, + '184': 104, + '185': 95, + '186': 114, + '187': 101, + '188': 97, + '189': 115, + '190': 111, + '191': 110, + '192': 34, + '193': 58, + '194': 110, + '195': 117, + '196': 108, + '197': 108, + '198': 125, + '199': 93, + '200': 125, + '201': 10, + '202': 10, +}; +export const mockUint8Arrays = [one, two]; +export const getReaderValue = (obj: { [key: string]: number }): Buffer => { + const parsedArray = Object.values(obj); + return Buffer.from(new Uint8Array(parsedArray)); +}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx new file mode 100644 index 0000000000000..e53bdb73cf4d0 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx @@ -0,0 +1,139 @@ +/* + * 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 { renderHook } from '@testing-library/react-hooks'; +import { useStream } from './use_stream'; +import { getReaderValue, mockUint8Arrays } from './mock'; + +const amendMessage = jest.fn(); +const reader = jest.fn(); +const cancel = jest.fn(); +const readerComplete = { + read: reader + .mockResolvedValueOnce({ + done: false, + value: getReaderValue(mockUint8Arrays[0]), + }) + .mockResolvedValueOnce({ + done: false, + value: getReaderValue(mockUint8Arrays[1]), + }) + .mockResolvedValue({ + done: true, + }), + cancel, + releaseLock: jest.fn(), + closed: jest.fn().mockResolvedValue(true), +} as unknown as ReadableStreamDefaultReader; + +const defaultProps = { amendMessage, reader: readerComplete }; +describe('useStream', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('Should stream response. isLoading/isStreaming are true while streaming, isLoading/isStreaming are false when streaming completes', async () => { + const { result, waitFor } = renderHook(() => useStream(defaultProps)); + expect(reader).toHaveBeenCalledTimes(1); + await waitFor(() => { + expect(result.current).toEqual({ + error: undefined, + isLoading: false, + isStreaming: false, + pendingMessage: '', + setComplete: expect.any(Function), + }); + }); + await waitFor(() => { + expect(result.current).toEqual({ + error: undefined, + isLoading: true, + isStreaming: false, + pendingMessage: '', + setComplete: expect.any(Function), + }); + }); + await waitFor(() => { + expect(result.current).toEqual({ + error: undefined, + isLoading: true, + isStreaming: true, + pendingMessage: 'Ch', + setComplete: expect.any(Function), + }); + }); + + await waitFor(() => { + expect(result.current).toEqual({ + error: undefined, + isLoading: false, + isStreaming: false, + pendingMessage: 'Cheddar', + setComplete: expect.any(Function), + }); + }); + expect(reader).toHaveBeenCalledTimes(3); + }); + + it('should not call observable when content is provided', () => { + renderHook(() => + useStream({ + ...defaultProps, + content: 'test content', + }) + ); + expect(reader).not.toHaveBeenCalled(); + }); + + it('should handle a stream error and update UseStream object accordingly', async () => { + const errorMessage = 'Test error message'; + const errorReader = { + read: jest + .fn() + .mockResolvedValueOnce({ + done: false, + value: getReaderValue(mockUint8Arrays[0]), + }) + .mockRejectedValue(new Error(errorMessage)), + cancel, + releaseLock: jest.fn(), + closed: jest.fn().mockResolvedValue(true), + } as unknown as ReadableStreamDefaultReader; + const { result, waitForNextUpdate } = renderHook(() => + useStream({ + amendMessage, + reader: errorReader, + }) + ); + expect(result.current.error).toBeUndefined(); + + await waitForNextUpdate(); + + expect(result.current.error).toBe(errorMessage); + expect(result.current.isLoading).toBe(false); + expect(result.current.pendingMessage).toBe(''); + expect(cancel).toHaveBeenCalled(); + }); + + it('should handle an empty content and reader object and return an empty observable', () => { + const { result } = renderHook(() => + useStream({ + ...defaultProps, + content: '', + reader: undefined, + }) + ); + + expect(result.current).toEqual({ + error: undefined, + isLoading: false, + isStreaming: false, + pendingMessage: '', + setComplete: expect.any(Function), + }); + }); +}); diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx index 96e6c3121e73d..07563c5ccf46f 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx @@ -7,7 +7,7 @@ import { useCallback, useEffect, useMemo, useState } from 'react'; import type { Subscription } from 'rxjs'; -import { concatMap, delay, Observable, of, share } from 'rxjs'; +import { delay, Observable, share } from 'rxjs'; export interface PromptObservableState { chunks: Chunk[]; @@ -52,65 +52,62 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us observer.next({ chunks: [], loading: true }); const decoder = new TextDecoder(); - + let prev = ''; const chunks: Chunk[] = []; - let prev: string = ''; - function read() { - reader?.read().then(({ done, value }: { done: boolean; value?: Uint8Array }) => { - try { - if (done) { + reader + ?.read() + .then(({ done, value }: { done: boolean; value?: Uint8Array }) => { + try { + if (done) { + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: false, + }); + observer.complete(); + return; + } + + const lines: string[] = (prev + decoder.decode(value)).split('\n'); + const lastLine: string = lines.pop() || ''; + const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; + + if (isPartialChunk) { + prev = lastLine; + } else { + prev = ''; + } + + const nextChunks: Chunk[] = lines + .map((str) => str.substr(6)) + .filter((str) => !!str && str !== '[DONE]') + .map((line) => JSON.parse(line)); + + chunks.push(...nextChunks); observer.next({ chunks, message: getMessageFromChunks(chunks), - loading: false, + loading: true, }); - observer.complete(); + } catch (err) { + observer.error(err); return; } - - let lines: string[] = (prev + decoder.decode(value)).split('\n'); - - const lastLine: string = lines[lines.length - 1]; - - const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; - - if (isPartialChunk) { - prev = lastLine; - lines.pop(); - } else { - prev = ''; - } - - lines = lines - .map((str) => str.substr(6)) - .filter((str) => !!str && str !== '[DONE]'); - - const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); - - nextChunks.forEach((chunk) => { - chunks.push(chunk); - observer.next({ - chunks, - message: getMessageFromChunks(chunks), - loading: true, - }); - }); - } catch (err) { + read(); + }) + .catch((err) => { observer.error(err); - return; - } - read(); - }); + }); } read(); return () => { - reader.cancel(); + reader?.cancel(); }; - }).pipe(concatMap((value) => of(value).pipe(delay(50)))) + }).pipe(delay(50)) : new Observable(), [content, reader] ); @@ -137,13 +134,13 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us useEffect(() => { const newSubscription = observer$.pipe(share()).subscribe({ - next: ({ message, loading: isLoading }) => { + next: (all) => { + const { message, loading: isLoading } = all; setLoading(isLoading); setPendingMessage(message); }, complete: () => { setComplete(true); - setLoading(false); }, error: (err) => { setError(err.message); @@ -162,9 +159,5 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us }; function getMessageFromChunks(chunks: Chunk[]) { - let message = ''; - chunks.forEach((chunk) => { - message += chunk.choices[0]?.delta.content ?? ''; - }); - return message; + return chunks.map((chunk) => chunk.choices[0]?.delta.content ?? '').join(''); } From b6c976a0abe6595c2efdba0a47327d2e5cccf8e7 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 24 Oct 2023 16:20:07 -0600 Subject: [PATCH 19/52] align regenerate button right --- .../public/assistant/get_comments/stream/index.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 fed130334355f..c0b237348b983 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 @@ -48,7 +48,7 @@ export const StreamComment = ({ ); } return ( - + From ebaab5997844f192fdca5ba338c01b7122d2313b Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 24 Oct 2023 16:31:47 -0600 Subject: [PATCH 20/52] fix --- .../get_comments/stream/message_text.tsx | 41 +++---------------- 1 file changed, 6 insertions(+), 35 deletions(-) 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 9ef70caf8f524..a1a2a5d150376 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 @@ -16,7 +16,8 @@ import classNames from 'classnames'; import type { Code, InlineCode, Parent, Text } from 'mdast'; import React, { useMemo, useRef } from 'react'; import type { Node } from 'unist'; -import { EsqlCodeBlock } from './esql_code_block'; +import { customCodeBlockLanguagePlugin } from '../custom_codeblock/custom_codeblock_markdown_plugin'; +import { CustomCodeBlock } from '../custom_codeblock/custom_code_block'; export type ChatActionClickHandler = (payload: { type: string; query: string }) => Promise; interface Props { @@ -89,32 +90,6 @@ const loadingCursorPlugin = () => { }; }; -const esqlLanguagePlugin = () => { - 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 === 'esql') { - node.type = 'esql'; - } - // TODO: make these renderers - if ( - (node.type === 'code' && node.lang === 'kql') || - (node.type === 'code' && node.lang === 'eql') - ) { - node.type = 'esql'; - } - }; - - return (tree: Node) => { - visitor(tree); - }; -}; - export function MessageText({ loading, content, onActionClick }: Props) { const containerClassName = css` overflow-wrap: break-word; @@ -133,19 +108,15 @@ export function MessageText({ loading, content, onActionClick }: Props) { processingPlugins[1][1].components = { ...components, - cursor: Cursor, - esql: (props) => { + customCodeBlock: (props) => { return ( <> - + ); }, + cursor: Cursor, table: (props) => ( <>
@@ -183,7 +154,7 @@ export function MessageText({ loading, content, onActionClick }: Props) { }; return { - parsingPluginList: [loadingCursorPlugin, esqlLanguagePlugin, ...parsingPlugins], + parsingPluginList: [loadingCursorPlugin, customCodeBlockLanguagePlugin, ...parsingPlugins], processingPluginList: processingPlugins, }; }, [loading]); From d52e9cf9a107815a5458f89800041f72714743db Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Tue, 24 Oct 2023 23:09:52 +0000 Subject: [PATCH 21/52] [CI] Auto-commit changed files from 'node scripts/eslint --no-cache --fix' --- .../security_solution/public/assistant/get_comments/index.tsx | 2 -- 1 file changed, 2 deletions(-) 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 e8a45dbc5e885..458f02799ade8 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 @@ -17,8 +17,6 @@ import type { EuiPanelProps } from '@elastic/eui/src/components/panel'; import { StreamComment } from './stream'; import { CommentActions } from '../comment_actions'; import * as i18n from './translations'; -import { customCodeBlockLanguagePlugin } from './custom_codeblock/custom_codeblock_markdown_plugin'; -import { CustomCodeBlock } from './custom_codeblock/custom_code_block'; export const getComments = ({ amendMessage, From c4f8a8e119ad007b8510a705496e7daadea059e2 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Wed, 25 Oct 2023 13:41:41 -0600 Subject: [PATCH 22/52] fixing --- .../get_comments/stream/message_text.tsx | 124 +++++++++--------- .../stack_connectors/common/openai/schema.ts | 5 +- .../connector_types/openai/openai.test.ts | 5 +- .../server/connector_types/openai/openai.ts | 17 +-- 4 files changed, 74 insertions(+), 77 deletions(-) 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 a1a2a5d150376..ca2760dddab13 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 @@ -14,16 +14,14 @@ import { import { css } from '@emotion/css'; import classNames from 'classnames'; import type { Code, InlineCode, Parent, Text } from 'mdast'; -import React, { useMemo, useRef } from 'react'; +import React from 'react'; import type { Node } from 'unist'; import { customCodeBlockLanguagePlugin } from '../custom_codeblock/custom_codeblock_markdown_plugin'; import { CustomCodeBlock } from '../custom_codeblock/custom_code_block'; -export type ChatActionClickHandler = (payload: { type: string; query: string }) => Promise; interface Props { content: string; loading: boolean; - onActionClick: ChatActionClickHandler; } const ANIMATION_TIME = 1; @@ -90,74 +88,72 @@ const loadingCursorPlugin = () => { }; }; -export function MessageText({ loading, content, onActionClick }: Props) { - const containerClassName = css` - overflow-wrap: break-word; - `; - - const onActionClickRef = useRef(onActionClick); - - onActionClickRef.current = onActionClick; +const getPluginDependencies = () => { + const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); - const { parsingPluginList, processingPluginList } = useMemo(() => { - const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); + const processingPlugins = getDefaultEuiMarkdownProcessingPlugins(); - const processingPlugins = getDefaultEuiMarkdownProcessingPlugins(); + const { components } = processingPlugins[1][1]; - const { components } = processingPlugins[1][1]; - - processingPlugins[1][1].components = { - ...components, - customCodeBlock: (props) => { - return ( - <> - - - - ); - }, - cursor: Cursor, - table: (props) => ( + processingPlugins[1][1].components = { + ...components, + customCodeBlock: (props) => { + return ( <> -
- {' '} -
+ + + {children} + + +
+
+ + {children} + +
+
- + - ), - th: (props) => { - const { children, ...rest } = props; - return ( -
- - - {children} - + ); + }, + cursor: Cursor, + table: (props) => ( + <> +
+ {' '} + + + + + ), + th: (props) => { + const { children, ...rest } = props; + return ( + - ); - }, - tr: (props) => , - td: (props) => { - const { children, ...rest } = props; - return ( - - ); - }, - }; - - return { - parsingPluginList: [loadingCursorPlugin, customCodeBlockLanguagePlugin, ...parsingPlugins], - processingPluginList: processingPlugins, - }; - }, [loading]); + + + ); + }, + tr: (props) => , + td: (props) => { + const { children, ...rest } = props; + return ( + + ); + }, + }; + + return { + parsingPluginList: [loadingCursorPlugin, customCodeBlockLanguagePlugin, ...parsingPlugins], + processingPluginList: processingPlugins, + }; +}; + +export function MessageText({ loading, content }: Props) { + const containerClassName = css` + overflow-wrap: break-word; + `; + + const { parsingPluginList, processingPluginList } = getPluginDependencies(); return ( diff --git a/x-pack/plugins/stack_connectors/common/openai/schema.ts b/x-pack/plugins/stack_connectors/common/openai/schema.ts index aa89622c7fb6b..d8bc52de6cedf 100644 --- a/x-pack/plugins/stack_connectors/common/openai/schema.ts +++ b/x-pack/plugins/stack_connectors/common/openai/schema.ts @@ -42,7 +42,10 @@ export const InvokeAIActionParamsSchema = schema.object({ schema.nullable(schema.oneOf([schema.string(), schema.arrayOf(schema.string())])) ), temperature: schema.maybe(schema.number()), - stream: schema.boolean({ defaultValue: false }), + // TODO: Remove in part 2 of streaming work for security solution + // tracked here: https://github.com/elastic/security-team/issues/7363 + // `stream` is a temporary parameter while the feature is developed behind a feature flag + stream: schema.maybe(schema.boolean({ defaultValue: false })), }); export const InvokeAIActionResponseSchema = schema.object({ diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts index 0a4a6a2931d8d..a311ef2e965c2 100644 --- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts +++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts @@ -17,6 +17,7 @@ import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; import { actionsMock } from '@kbn/actions-plugin/server/mocks'; import { RunActionResponseSchema, StreamingResponseSchema } from '../../../common/openai/schema'; import { initDashboard } from './create_dashboard'; +import { InvokeAIActionResponse } from '../../../common/openai/types'; jest.mock('./create_dashboard'); describe('OpenAIConnector', () => { @@ -266,7 +267,9 @@ describe('OpenAIConnector', () => { describe('invokeAI', () => { it('the API call is successful with correct parameters', async () => { - const response = await connector.invokeAI(sampleOpenAiBody); + const response: InvokeAIActionResponse = (await connector.invokeAI( + sampleOpenAiBody + )) as unknown as InvokeAIActionResponse; expect(mockRequest).toBeCalledTimes(1); expect(mockRequest).toHaveBeenCalledWith({ url: 'https://api.openai.com/v1/chat/completions', diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts index 0c473a8aad197..acaf120ce4c46 100644 --- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts +++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts @@ -28,7 +28,6 @@ import { DashboardActionResponse, InvokeAIActionParams, InvokeAIActionResponse, - StreamingResponse, } from '../../../common/openai/types'; import { initDashboard } from './create_dashboard'; import { @@ -190,22 +189,18 @@ export class OpenAIConnector extends SubActionConnector { */ public async invokeAI( params: InvokeAIActionParams - ): Promise { - const { stream, ...body } = params; + ): Promise { + // TODO: Remove in part 2 of streaming work for security solution + // tracked here: https://github.com/elastic/security-team/issues/7363 + // `stream` is a temporary parameter while the feature is developed behind a feature flag + const { stream = false, ...body } = params; const res = await this.streamApi({ body: JSON.stringify(body), stream }); + // TODO: Remove in part 2 of streaming work for security solution if (res.choices && res.choices.length > 0 && res.choices[0].message?.content) { const result = res.choices[0].message.content.trim(); return { message: result, usage: res.usage }; } return res; - - return { - message: - 'An error occurred sending your message. \n\nAPI Error: The response from OpenAI was in an unrecognized format.', - ...(res.usage - ? { usage: res.usage } - : { usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 } }), - }; } } From cd8a3b548322ed7663041266a2f748cc66c158bd Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Wed, 25 Oct 2023 16:43:42 -0600 Subject: [PATCH 23/52] trying to make it better but made it worse --- .../impl/assistant/api.tsx | 25 +++++++++++++------ .../actions/server/lib/action_executor.ts | 3 +++ .../elastic_assistant/server/lib/executor.ts | 17 +++++++++++-- .../routes/post_actions_connector_execute.ts | 1 + .../server/connector_types/openai/openai.ts | 1 + 5 files changed, 38 insertions(+), 9 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index 41d2103e83eaf..acc81270846bd 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -57,13 +57,21 @@ export const fetchConnectorExecuteAction = async ({ messages: outboundMessages, }; const isStream = true; - const requestBody = { - params: { - subActionParams: { ...body, stream: isStream }, - subAction: 'invokeAI', - }, - assistantLangChain, - }; + const requestBody = isStream + ? { + params: { + subActionParams: body, + subAction: 'stream', + }, + assistantLangChain, + } + : { + params: { + subActionParams: body, + subAction: 'invokeAI', + }, + assistantLangChain, + }; try { if (isStream) { @@ -94,6 +102,9 @@ export const fetchConnectorExecuteAction = async ({ }; } + // TODO: Remove in part 2 of streaming work for security solution + // tracked here: https://github.com/elastic/security-team/issues/7363 + // This is a temporary code to support the non-streaming API const response = await http.fetch<{ connector_id: string; status: string; diff --git a/x-pack/plugins/actions/server/lib/action_executor.ts b/x-pack/plugins/actions/server/lib/action_executor.ts index 1ca2ce865cd28..6cdcc75d582dd 100644 --- a/x-pack/plugins/actions/server/lib/action_executor.ts +++ b/x-pack/plugins/actions/server/lib/action_executor.ts @@ -356,6 +356,9 @@ export class ActionExecutor { if (result.data instanceof Readable) { let body: string; if (!(validatedParams as { subActionParams: { body: string } }).subActionParams?.body) { + // TODO: Remove in part 2 of streaming work for security solution + // tracked here: https://github.com/elastic/security-team/issues/7363 + // `stream` is a temporary parameter while the feature is developed behind a feature flagr const { stream: _, ...rest } = ( validatedParams as { subActionParams: { [a: string]: string } } ).subActionParams; diff --git a/x-pack/plugins/elastic_assistant/server/lib/executor.ts b/x-pack/plugins/elastic_assistant/server/lib/executor.ts index 247f3a164706f..595287690cf38 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/executor.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/executor.ts @@ -28,10 +28,21 @@ export const executeAction = async ({ connectorId, }: Props): Promise => { const actionsClient = await actions.getActionsClientWithRequest(request); + + console.log('one'); const actionResult = await actionsClient.execute({ actionId: connectorId, - params: request.body.params, + params: { + ...request.body.params, + subActionParams: + // TODO: Remove in part 2 of streaming work for security solution + // tracked here: https://github.com/elastic/security-team/issues/7363 + request.body.params.subAction === 'invokeAI' + ? request.body.params.subActionParams + : { body: JSON.stringify(request.body.params.subActionParams), stream: true }, + }, }); + console.log('two', actionResult); const content = get('data.message', actionResult); if (typeof content === 'string') { return { @@ -42,5 +53,7 @@ export const executeAction = async ({ } const readable = get('data', actionResult); - return (readable as Readable).pipe(new PassThrough()) as Readable; + console.log('typeof', typeof readable); + console.log('three', readable); + return (readable as Readable).pipe(new PassThrough()); }; 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 ed6fa78c71db5..b886532f9ff88 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 @@ -44,6 +44,7 @@ export const postActionsConnectorExecuteRoute = ( // if not langchain, call execute action directly and return the response: if (!request.body.assistantLangChain) { logger.debug('Executing via actions framework directly, assistantLangChain: false'); + const result = await executeAction({ actions, request, connectorId }); return response.ok({ body: result, diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts index acaf120ce4c46..88e73d3ff9f6a 100644 --- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts +++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts @@ -193,6 +193,7 @@ export class OpenAIConnector extends SubActionConnector { // TODO: Remove in part 2 of streaming work for security solution // tracked here: https://github.com/elastic/security-team/issues/7363 // `stream` is a temporary parameter while the feature is developed behind a feature flag + // why not call stream directly from the securit solution api? because the body has no validation const { stream = false, ...body } = params; const res = await this.streamApi({ body: JSON.stringify(body), stream }); From 2dcb57da6439b6287d68d12c48053a01d1b84495 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Wed, 25 Oct 2023 17:05:38 -0600 Subject: [PATCH 24/52] stash dumb changes --- .../actions/server/lib/action_executor.ts | 15 +----------- .../stack_connectors/common/openai/schema.ts | 4 ---- .../stack_connectors/common/openai/types.ts | 1 - .../connector_types/openai/openai.test.ts | 5 +--- .../server/connector_types/openai/openai.ts | 23 +++++++++---------- 5 files changed, 13 insertions(+), 35 deletions(-) diff --git a/x-pack/plugins/actions/server/lib/action_executor.ts b/x-pack/plugins/actions/server/lib/action_executor.ts index 6cdcc75d582dd..9cd70d4c7bf91 100644 --- a/x-pack/plugins/actions/server/lib/action_executor.ts +++ b/x-pack/plugins/actions/server/lib/action_executor.ts @@ -354,22 +354,9 @@ export class ActionExecutor { }; if (result.data instanceof Readable) { - let body: string; - if (!(validatedParams as { subActionParams: { body: string } }).subActionParams?.body) { - // TODO: Remove in part 2 of streaming work for security solution - // tracked here: https://github.com/elastic/security-team/issues/7363 - // `stream` is a temporary parameter while the feature is developed behind a feature flagr - const { stream: _, ...rest } = ( - validatedParams as { subActionParams: { [a: string]: string } } - ).subActionParams; - body = JSON.stringify(rest); - } else { - body = (validatedParams as { subActionParams: { body: string } }).subActionParams - .body; - } getTokenCountFromOpenAIStream({ responseStream: result.data.pipe(new PassThrough()), - body, + body: (validatedParams as { subActionParams: { body: string } }).subActionParams.body, }) .then(({ total, prompt, completion }) => { event.kibana!.action!.execution!.gen_ai!.usage = { diff --git a/x-pack/plugins/stack_connectors/common/openai/schema.ts b/x-pack/plugins/stack_connectors/common/openai/schema.ts index d8bc52de6cedf..1154c808154b1 100644 --- a/x-pack/plugins/stack_connectors/common/openai/schema.ts +++ b/x-pack/plugins/stack_connectors/common/openai/schema.ts @@ -42,10 +42,6 @@ export const InvokeAIActionParamsSchema = schema.object({ schema.nullable(schema.oneOf([schema.string(), schema.arrayOf(schema.string())])) ), temperature: schema.maybe(schema.number()), - // TODO: Remove in part 2 of streaming work for security solution - // tracked here: https://github.com/elastic/security-team/issues/7363 - // `stream` is a temporary parameter while the feature is developed behind a feature flag - stream: schema.maybe(schema.boolean({ defaultValue: false })), }); export const InvokeAIActionResponseSchema = schema.object({ diff --git a/x-pack/plugins/stack_connectors/common/openai/types.ts b/x-pack/plugins/stack_connectors/common/openai/types.ts index 4c78b86a4a99a..86e2c172846dc 100644 --- a/x-pack/plugins/stack_connectors/common/openai/types.ts +++ b/x-pack/plugins/stack_connectors/common/openai/types.ts @@ -23,7 +23,6 @@ export type Secrets = TypeOf; export type RunActionParams = TypeOf; export type InvokeAIActionParams = TypeOf; export type InvokeAIActionResponse = TypeOf; -export type StreamingResponse = TypeOf; export type RunActionResponse = TypeOf; export type DashboardActionParams = TypeOf; export type DashboardActionResponse = TypeOf; diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts index a311ef2e965c2..0a4a6a2931d8d 100644 --- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts +++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts @@ -17,7 +17,6 @@ import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; import { actionsMock } from '@kbn/actions-plugin/server/mocks'; import { RunActionResponseSchema, StreamingResponseSchema } from '../../../common/openai/schema'; import { initDashboard } from './create_dashboard'; -import { InvokeAIActionResponse } from '../../../common/openai/types'; jest.mock('./create_dashboard'); describe('OpenAIConnector', () => { @@ -267,9 +266,7 @@ describe('OpenAIConnector', () => { describe('invokeAI', () => { it('the API call is successful with correct parameters', async () => { - const response: InvokeAIActionResponse = (await connector.invokeAI( - sampleOpenAiBody - )) as unknown as InvokeAIActionResponse; + const response = await connector.invokeAI(sampleOpenAiBody); expect(mockRequest).toBeCalledTimes(1); expect(mockRequest).toHaveBeenCalledWith({ url: 'https://api.openai.com/v1/chat/completions', diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts index 88e73d3ff9f6a..7413ba56090a1 100644 --- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts +++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts @@ -187,21 +187,20 @@ export class OpenAIConnector extends SubActionConnector { * Sends the stringified input to the runApi method. Returns the trimmed completion from the response. * @param body An object containing array of message objects, and possible other OpenAI properties */ - public async invokeAI( - params: InvokeAIActionParams - ): Promise { - // TODO: Remove in part 2 of streaming work for security solution - // tracked here: https://github.com/elastic/security-team/issues/7363 - // `stream` is a temporary parameter while the feature is developed behind a feature flag - // why not call stream directly from the securit solution api? because the body has no validation - const { stream = false, ...body } = params; - const res = await this.streamApi({ body: JSON.stringify(body), stream }); - - // TODO: Remove in part 2 of streaming work for security solution + public async invokeAI(body: InvokeAIActionParams): Promise { + const res = await this.runApi({ body: JSON.stringify(body) }); + if (res.choices && res.choices.length > 0 && res.choices[0].message?.content) { const result = res.choices[0].message.content.trim(); return { message: result, usage: res.usage }; } - return res; + + return { + message: + 'An error occurred sending your message. \n\nAPI Error: The response from OpenAI was in an unrecognized format.', + ...(res.usage + ? { usage: res.usage } + : { usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 } }), + }; } } From 358881ade1bd4102c9f1975cbcf5bbb1d9ef1dc6 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 26 Oct 2023 09:21:51 -0500 Subject: [PATCH 25/52] add comment --- .../packages/kbn-elastic-assistant/impl/assistant/api.tsx | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index acc81270846bd..0f7eade9cc31e 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -56,6 +56,12 @@ export const fetchConnectorExecuteAction = async ({ // Azure OpenAI and Bedrock invokeAI both expect this body format messages: outboundMessages, }; + + // TODO: Remove in part 2 of streaming work for security solution + // tracked here: https://github.com/elastic/security-team/issues/7363 + // My "Feature Flag", turn to false before merging + // In part 2 I will make enhancements to invokeAI to make it work with both openA, but to keep it to a Security Soltuion only review on this PR, + // I'm calling the stream action directly const isStream = true; const requestBody = isStream ? { From 5154e32e2ad63201c3e19d5dfe5f67873d269b3f Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 26 Oct 2023 11:41:29 -0500 Subject: [PATCH 26/52] fix stream to be pretty --- .../assistant/get_comments/stream/index.tsx | 8 +- .../get_comments/stream/message_text.tsx | 2 +- .../get_comments/stream/use_stream.test.tsx | 9 -- .../get_comments/stream/use_stream.tsx | 95 +++++++++++-------- 4 files changed, 59 insertions(+), 55 deletions(-) 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 c0b237348b983..6395ddaff8550 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 @@ -57,13 +57,7 @@ export const StreamComment = ({ }, [isLastComment, isLoading, isStreaming, reader, regenerateMessage, setComplete]); 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 ca2760dddab13..554d97a8e4f34 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 @@ -97,6 +97,7 @@ const getPluginDependencies = () => { processingPlugins[1][1].components = { ...components, + cursor: Cursor, customCodeBlock: (props) => { return ( <> @@ -105,7 +106,6 @@ const getPluginDependencies = () => { ); }, - cursor: Cursor, table: (props) => ( <>
diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx index e53bdb73cf4d0..368f7b90fee81 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx @@ -39,15 +39,6 @@ describe('useStream', () => { it('Should stream response. isLoading/isStreaming are true while streaming, isLoading/isStreaming are false when streaming completes', async () => { const { result, waitFor } = renderHook(() => useStream(defaultProps)); expect(reader).toHaveBeenCalledTimes(1); - await waitFor(() => { - expect(result.current).toEqual({ - error: undefined, - isLoading: false, - isStreaming: false, - pendingMessage: '', - setComplete: expect.any(Function), - }); - }); await waitFor(() => { expect(result.current).toEqual({ error: undefined, diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx index 07563c5ccf46f..b949246b56f7a 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx @@ -7,21 +7,18 @@ import { useCallback, useEffect, useMemo, useState } from 'react'; import type { Subscription } from 'rxjs'; -import { delay, Observable, share } from 'rxjs'; - +import { concatMap, delay, Observable, of, scan, share, shareReplay, timestamp } from 'rxjs'; export interface PromptObservableState { chunks: Chunk[]; message?: string; error?: string; loading: boolean; } - interface ChunkChoice { index: 0; delta: { role: string; content: string }; finish_reason: null | string; } - interface Chunk { id: string; object: string; @@ -29,13 +26,11 @@ interface Chunk { model: string; choices: ChunkChoice[]; } - interface UseStreamProps { amendMessage: (message: string) => void; content?: string; reader?: ReadableStreamDefaultReader; } - interface UseStream { error: string | undefined; isLoading: boolean; @@ -43,18 +38,16 @@ interface UseStream { pendingMessage: string; setComplete: (complete: boolean) => void; } - +const MIN_DELAY = 35; export const useStream = ({ amendMessage, content, reader }: UseStreamProps): UseStream => { const observer$ = useMemo( () => content == null && reader != null ? new Observable((observer) => { observer.next({ chunks: [], loading: true }); - const decoder = new TextDecoder(); - let prev = ''; const chunks: Chunk[] = []; - + let prev: string = ''; function read() { reader ?.read() @@ -69,27 +62,26 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us observer.complete(); return; } - - const lines: string[] = (prev + decoder.decode(value)).split('\n'); - const lastLine: string = lines.pop() || ''; + let lines: string[] = (prev + decoder.decode(value)).split('\n'); + const lastLine: string = lines[lines.length - 1]; const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; - if (isPartialChunk) { prev = lastLine; + lines.pop(); } else { prev = ''; } - - const nextChunks: Chunk[] = lines + lines = lines .map((str) => str.substr(6)) - .filter((str) => !!str && str !== '[DONE]') - .map((line) => JSON.parse(line)); - - chunks.push(...nextChunks); - observer.next({ - chunks, - message: getMessageFromChunks(chunks), - loading: true, + .filter((str) => !!str && str !== '[DONE]'); + const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); + nextChunks.forEach((chunk) => { + chunks.push(chunk); + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: true, + }); }); } catch (err) { observer.error(err); @@ -101,54 +93,78 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us observer.error(err); }); } - read(); - return () => { - reader?.cancel(); + reader.cancel(); }; - }).pipe(delay(50)) + }).pipe( + // make sure the request is only triggered once, + // even with multiple subscribers + shareReplay(1), + // append a timestamp of when each value was emitted + timestamp(), + // use the previous timestamp to calculate a target + // timestamp for emitting the next value + scan((acc, value) => { + const lastTimestamp = acc.timestamp || 0; + const emitAt = Math.max(lastTimestamp + MIN_DELAY, value.timestamp); + return { + timestamp: emitAt, + value: value.value, + }; + }), + // add the delay based on the elapsed time + // using concatMap(of(value).pipe(delay(50)) + // leads to browser issues because timers + // are throttled when the tab is not active + concatMap((value) => { + const now = Date.now(); + console.log('VALUE?', value); + const delayFor = value.timestamp - now; + + if (delayFor <= 0) { + return of(value.value); + } + + return of(value.value).pipe(delay(delayFor)); + }) + ) : new Observable(), [content, reader] ); - const [pendingMessage, setPendingMessage] = useState(); const [loading, setLoading] = useState(false); const [error, setError] = useState(); const [subscription, setSubscription] = useState(); - const onCompleteStream = useCallback(() => { subscription?.unsubscribe(); setLoading(false); amendMessage(pendingMessage ?? ''); }, [amendMessage, pendingMessage, subscription]); - const [complete, setComplete] = useState(false); - useEffect(() => { if (complete) { setComplete(false); onCompleteStream(); } }, [complete, onCompleteStream]); - useEffect(() => { const newSubscription = observer$.pipe(share()).subscribe({ - next: (all) => { - const { message, loading: isLoading } = all; + next: ({ message, loading: isLoading }) => { setLoading(isLoading); setPendingMessage(message); }, complete: () => { setComplete(true); + setLoading(false); }, error: (err) => { setError(err.message); + setLoading(false); }, }); setSubscription(newSubscription); }, [observer$]); - return { error, isLoading: loading, @@ -157,7 +173,10 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us setComplete, }; }; - function getMessageFromChunks(chunks: Chunk[]) { - return chunks.map((chunk) => chunk.choices[0]?.delta.content ?? '').join(''); + let message = ''; + chunks.forEach((chunk) => { + message += chunk.choices[0]?.delta.content ?? ''; + }); + return message; } From db30dd8ea8caf608677aebd0ea39388fb2cd1992 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 26 Oct 2023 11:52:21 -0500 Subject: [PATCH 27/52] add comments --- .../get_comments/stream/use_stream.tsx | 31 ++++++++++++++++--- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx index b949246b56f7a..6ec474fb5f784 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx @@ -7,7 +7,18 @@ import { useCallback, useEffect, useMemo, useState } from 'react'; import type { Subscription } from 'rxjs'; -import { concatMap, delay, Observable, of, scan, share, shareReplay, timestamp } from 'rxjs'; +import { + concatMap, + delay, + finalize, + Observable, + of, + scan, + share, + shareReplay, + timestamp, +} from 'rxjs'; + export interface PromptObservableState { chunks: Chunk[]; message?: string; @@ -32,13 +43,25 @@ interface UseStreamProps { reader?: ReadableStreamDefaultReader; } interface UseStream { + // The error message, if an error occurs during streaming. error: string | undefined; + // Indicates whether the streaming is in progress or not isLoading: boolean; + // Indicates whether the streaming is in progress and there is a pending message. isStreaming: boolean; + // The pending message from the streaming data. pendingMessage: string; + // A function to mark the streaming as complete setComplete: (complete: boolean) => void; } const MIN_DELAY = 35; +/** + * A hook that takes a ReadableStreamDefaultReader and returns an object with properties and functions + * that can be used to handle streaming data from a readable stream + * @param amendMessage - handles the amended message + * @param content - the content of the message. If provided, the function will not use the reader to stream data. + * @param reader - The readable stream reader used to stream data. If provided, the function will use this reader to stream data. + */ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): UseStream => { const observer$ = useMemo( () => @@ -127,7 +150,9 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us } return of(value.value).pipe(delay(delayFor)); - }) + }), + // set loading to false when the observable completes or errors out + finalize(() => setLoading(false)) ) : new Observable(), [content, reader] @@ -156,11 +181,9 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us }, complete: () => { setComplete(true); - setLoading(false); }, error: (err) => { setError(err.message); - setLoading(false); }, }); setSubscription(newSubscription); From 6b79554ac4d02eddcd7268fd3b9e239ce7eb78df Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 26 Oct 2023 13:17:43 -0500 Subject: [PATCH 28/52] cleanup --- .../get_comments/stream/stream_observable.ts | 118 +++++++++++++++ .../assistant/get_comments/stream/types.ts | 25 ++++ .../get_comments/stream/use_stream.tsx | 138 +----------------- 3 files changed, 151 insertions(+), 130 deletions(-) create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/types.ts diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts new file mode 100644 index 0000000000000..1e5fe552835dc --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts @@ -0,0 +1,118 @@ +/* + * 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 { concatMap, delay, finalize, Observable, of, scan, shareReplay, timestamp } from 'rxjs'; +import type { Dispatch, SetStateAction } from 'react'; +import type { PromptObservableState, Chunk } from './types'; + +const MIN_DELAY = 35; +/** + * Returns an Observable that reads data from a ReadableStream and emits values representing the state of the data processing. + * + * @param reader - The ReadableStreamDefaultReader used to read data from the stream. + * @param setLoading - A function to update the loading state. + * @returns {Observable} An Observable that emits PromptObservableState + */ +export const getStreamObservable = ( + reader: ReadableStreamDefaultReader, + setLoading: Dispatch> +): Observable => + new Observable((observer) => { + observer.next({ chunks: [], loading: true }); + const decoder = new TextDecoder(); + const chunks: Chunk[] = []; + let prev: string = ''; + function read() { + reader + ?.read() + .then(({ done, value }: { done: boolean; value?: Uint8Array }) => { + try { + if (done) { + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: false, + }); + observer.complete(); + return; + } + let lines: string[] = (prev + decoder.decode(value)).split('\n'); + const lastLine: string = lines[lines.length - 1]; + const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; + if (isPartialChunk) { + prev = lastLine; + lines.pop(); + } else { + prev = ''; + } + lines = lines.map((str) => str.substring(6)).filter((str) => !!str && str !== '[DONE]'); + const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); + nextChunks.forEach((chunk) => { + chunks.push(chunk); + observer.next({ + chunks, + message: getMessageFromChunks(chunks), + loading: true, + }); + }); + } catch (err) { + observer.error(err); + return; + } + read(); + }) + .catch((err) => { + observer.error(err); + }); + } + read(); + return () => { + reader.cancel(); + }; + }).pipe( + // make sure the request is only triggered once, + // even with multiple subscribers + shareReplay(1), + // append a timestamp of when each value was emitted + timestamp(), + // use the previous timestamp to calculate a target + // timestamp for emitting the next value + scan((acc, value) => { + const lastTimestamp = acc.timestamp || 0; + const emitAt = Math.max(lastTimestamp + MIN_DELAY, value.timestamp); + return { + timestamp: emitAt, + value: value.value, + }; + }), + // add the delay based on the elapsed time + // using concatMap(of(value).pipe(delay(50)) + // leads to browser issues because timers + // are throttled when the tab is not active + concatMap((value) => { + const now = Date.now(); + const delayFor = value.timestamp - now; + + if (delayFor <= 0) { + return of(value.value); + } + + return of(value.value).pipe(delay(delayFor)); + }), + // set loading to false when the observable completes or errors out + finalize(() => setLoading(false)) + ); + +function getMessageFromChunks(chunks: Chunk[]) { + let message = ''; + chunks.forEach((chunk) => { + message += chunk.choices[0]?.delta.content ?? ''; + }); + return message; +} + +export const getDumbObservable = () => new Observable(); diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/types.ts b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/types.ts new file mode 100644 index 0000000000000..3cf45852ddb11 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/types.ts @@ -0,0 +1,25 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +export interface PromptObservableState { + chunks: Chunk[]; + message?: string; + error?: string; + loading: boolean; +} +export interface ChunkChoice { + index: 0; + delta: { role: string; content: string }; + finish_reason: null | string; +} +export interface Chunk { + id: string; + object: string; + created: number; + model: string; + choices: ChunkChoice[]; +} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx index 6ec474fb5f784..e756ba87c8f35 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx @@ -7,36 +7,9 @@ import { useCallback, useEffect, useMemo, useState } from 'react'; import type { Subscription } from 'rxjs'; -import { - concatMap, - delay, - finalize, - Observable, - of, - scan, - share, - shareReplay, - timestamp, -} from 'rxjs'; +import { share } from 'rxjs'; +import { getDumbObservable, getStreamObservable } from './stream_observable'; -export interface PromptObservableState { - chunks: Chunk[]; - message?: string; - error?: string; - loading: boolean; -} -interface ChunkChoice { - index: 0; - delta: { role: string; content: string }; - finish_reason: null | string; -} -interface Chunk { - id: string; - object: string; - created: number; - model: string; - choices: ChunkChoice[]; -} interface UseStreamProps { amendMessage: (message: string) => void; content?: string; @@ -54,7 +27,6 @@ interface UseStream { // A function to mark the streaming as complete setComplete: (complete: boolean) => void; } -const MIN_DELAY = 35; /** * A hook that takes a ReadableStreamDefaultReader and returns an object with properties and functions * that can be used to handle streaming data from a readable stream @@ -63,104 +35,17 @@ const MIN_DELAY = 35; * @param reader - The readable stream reader used to stream data. If provided, the function will use this reader to stream data. */ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): UseStream => { + const [pendingMessage, setPendingMessage] = useState(); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(); + const [subscription, setSubscription] = useState(); const observer$ = useMemo( () => content == null && reader != null - ? new Observable((observer) => { - observer.next({ chunks: [], loading: true }); - const decoder = new TextDecoder(); - const chunks: Chunk[] = []; - let prev: string = ''; - function read() { - reader - ?.read() - .then(({ done, value }: { done: boolean; value?: Uint8Array }) => { - try { - if (done) { - observer.next({ - chunks, - message: getMessageFromChunks(chunks), - loading: false, - }); - observer.complete(); - return; - } - let lines: string[] = (prev + decoder.decode(value)).split('\n'); - const lastLine: string = lines[lines.length - 1]; - const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; - if (isPartialChunk) { - prev = lastLine; - lines.pop(); - } else { - prev = ''; - } - lines = lines - .map((str) => str.substr(6)) - .filter((str) => !!str && str !== '[DONE]'); - const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); - nextChunks.forEach((chunk) => { - chunks.push(chunk); - observer.next({ - chunks, - message: getMessageFromChunks(chunks), - loading: true, - }); - }); - } catch (err) { - observer.error(err); - return; - } - read(); - }) - .catch((err) => { - observer.error(err); - }); - } - read(); - return () => { - reader.cancel(); - }; - }).pipe( - // make sure the request is only triggered once, - // even with multiple subscribers - shareReplay(1), - // append a timestamp of when each value was emitted - timestamp(), - // use the previous timestamp to calculate a target - // timestamp for emitting the next value - scan((acc, value) => { - const lastTimestamp = acc.timestamp || 0; - const emitAt = Math.max(lastTimestamp + MIN_DELAY, value.timestamp); - return { - timestamp: emitAt, - value: value.value, - }; - }), - // add the delay based on the elapsed time - // using concatMap(of(value).pipe(delay(50)) - // leads to browser issues because timers - // are throttled when the tab is not active - concatMap((value) => { - const now = Date.now(); - console.log('VALUE?', value); - const delayFor = value.timestamp - now; - - if (delayFor <= 0) { - return of(value.value); - } - - return of(value.value).pipe(delay(delayFor)); - }), - // set loading to false when the observable completes or errors out - finalize(() => setLoading(false)) - ) - : new Observable(), + ? getStreamObservable(reader, setLoading) + : getDumbObservable(), [content, reader] ); - const [pendingMessage, setPendingMessage] = useState(); - const [loading, setLoading] = useState(false); - const [error, setError] = useState(); - const [subscription, setSubscription] = useState(); const onCompleteStream = useCallback(() => { subscription?.unsubscribe(); setLoading(false); @@ -196,10 +81,3 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us setComplete, }; }; -function getMessageFromChunks(chunks: Chunk[]) { - let message = ''; - chunks.forEach((chunk) => { - message += chunk.choices[0]?.delta.content ?? ''; - }); - return message; -} From b7013231e3a42ad7c52dcb438e8b9894fc57556e Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 26 Oct 2023 14:17:07 -0500 Subject: [PATCH 29/52] better UX --- .../impl/assistant/chat_send/index.tsx | 17 +- .../assistant/chat_send/use_chat_send.tsx | 3 +- .../impl/assistant/index.tsx | 26 +-- .../assistant/use_send_messages/index.tsx | 9 +- .../impl/assistant_context/index.tsx | 2 + .../public/assistant/get_comments/index.tsx | 155 +++++++++++------- .../assistant/get_comments/stream/index.tsx | 14 +- 7 files changed, 136 insertions(+), 90 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.tsx index 88db2e124ceab..2cabf27839236 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.tsx @@ -8,11 +8,11 @@ import React, { useCallback, useEffect, useMemo, useRef } from 'react'; import { css } from '@emotion/react'; import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui'; -import { useChatSend, UseChatSendProps } from './use_chat_send'; +import { UseChatSend } from './use_chat_send'; import { ChatActions } from '../chat_actions'; import { PromptTextArea } from '../prompt_textarea'; -export interface Props extends UseChatSendProps { +export interface Props extends UseChatSend { isDisabled: boolean; shouldRefocusPrompt: boolean; userPrompt: string | null; @@ -26,15 +26,12 @@ export const ChatSend: React.FC = ({ isDisabled, userPrompt, shouldRefocusPrompt, - ...rest + handleButtonSendMessage, + handleOnChatCleared, + handlePromptChange, + handleSendMessage, + isLoading, }) => { - const { - handleButtonSendMessage, - handleOnChatCleared, - handlePromptChange, - handleSendMessage, - isLoading, - } = useChatSend(rest); // For auto-focusing prompt within timeline const promptTextAreaRef = useRef(null); useEffect(() => { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index c269ed76cfd45..83a2ad33e55fb 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -29,7 +29,7 @@ export interface UseChatSendProps { setUserPrompt: React.Dispatch>; } -interface UseChatSend { +export interface UseChatSend { handleButtonSendMessage: (m: string) => void; handleOnChatCleared: () => void; handlePromptChange: (prompt: string) => void; @@ -55,6 +55,7 @@ export const useChatSend = ({ setUserPrompt, }: UseChatSendProps): UseChatSend => { const { isLoading, sendMessages } = useSendMessages(); + const { appendMessage, appendReplacements, clearConversation, removeLastMessage } = useConversation(); 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 ff0e0536ea309..ac83710477aa4 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -340,7 +340,14 @@ const AssistantComponent: React.FC = ({ [messageCodeBlocks] ); - const { handleRegenerateResponse } = useChatSend({ + const { + handleButtonSendMessage, + handleOnChatCleared, + handlePromptChange, + handleSendMessage, + handleRegenerateResponse, + isLoading: isLoadingChatSend, + } = useChatSend({ allSystemPrompts, currentConversation, setPromptTextPreview, @@ -362,6 +369,7 @@ const AssistantComponent: React.FC = ({ showAnonymizedValues, amendMessage, regenerateMessage: handleRegenerateResponse, + isFetchingResponse: isLoadingChatSend, })} css={css` margin-right: 20px; @@ -397,6 +405,7 @@ const AssistantComponent: React.FC = ({ getComments, handleOnSystemPromptSelectionChange, handleRegenerateResponse, + isLoadingChatSend, isSettingsModalVisible, promptContexts, promptTextPreview, @@ -507,18 +516,15 @@ const AssistantComponent: React.FC = ({ isWelcomeSetup={isWelcomeSetup} /> {!isDisabled && ( { }, [knowledgeBase.assistantLangChain] ); - - return { isLoading, sendMessages }; + const retThis = useMemo(() => { + console.log('isLoading, at root?', isLoading); + return { isLoading, sendMessages }; + }, [isLoading, sendMessages]); + return retThis; }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index 553c42278abed..efc08ca0d0183 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -105,8 +105,10 @@ export interface UseAssistantContext { lastCommentRef, showAnonymizedValues, amendMessage, + isFetchingResponse, }: { currentConversation: Conversation; + isFetchingResponse: boolean; lastCommentRef: React.MutableRefObject; amendMessage: ({ conversationId, 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 458f02799ade8..47bb3728f9f5d 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 @@ -7,7 +7,7 @@ import type { EuiCommentProps } from '@elastic/eui'; import type { Conversation, Message } from '@kbn/elastic-assistant'; -import { EuiAvatar, tint } from '@elastic/eui'; +import { EuiAvatar, EuiLoadingSpinner, tint } from '@elastic/eui'; import React from 'react'; import { AssistantAvatar } from '@kbn/elastic-assistant'; @@ -21,6 +21,7 @@ import * as i18n from './translations'; export const getComments = ({ amendMessage, currentConversation, + isFetchingResponse, lastCommentRef, regenerateMessage, showAnonymizedValues, @@ -33,6 +34,7 @@ export const getComments = ({ content: string; }) => Message[]; currentConversation: Conversation; + isFetchingResponse: boolean; lastCommentRef: React.MutableRefObject; regenerateMessage: (conversationId: string) => void; showAnonymizedValues: boolean; @@ -48,44 +50,101 @@ export const getComments = ({ regenerateMessage(currentConversation.id); }; - return currentConversation.messages.map((message, index) => { - const isUser = message.role === 'user'; - const replacements = currentConversation.replacements; - const errorStyles = { - eventColor: 'danger' as EuiPanelProps['color'], - css: css` - .euiCommentEvent { - border: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; - } - .euiCommentEvent__header { - padding: 0 !important; - border-block-end: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; - } - `, - }; + console.log('get comments return', { + isFetchingResponse, + currentConversation, + }); + const extraLoadingComment = isFetchingResponse + ? [ + { + username: i18n.ASSISTANT, + timelineAvatar: , + timestamp: '...', + children: ( + <> + + + + ), + }, + ] + : []; + + return [ + ...currentConversation.messages.map((message, index) => { + const isLastComment = index === currentConversation.messages.length - 1; + const isUser = message.role === 'user'; + const replacements = currentConversation.replacements; + const errorStyles = { + eventColor: 'danger' as EuiPanelProps['color'], + css: css` + .euiCommentEvent { + border: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; + } + .euiCommentEvent__header { + padding: 0 !important; + border-block-end: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; + } + `, + }; + + const messageProps = { + timelineAvatar: isUser ? ( + + ) : ( + + ), + timestamp: i18n.AT( + message.timestamp.length === 0 ? new Date().toLocaleString() : message.timestamp + ), + username: isUser ? i18n.YOU : i18n.ASSISTANT, + ...(message.isError ? errorStyles : {}), + }; + + // message still needs to stream, no response manipulation + if (!(message.content && message.content.length)) { + return { + ...messageProps, + children: ( + <> + + {isLastComment ? : null} + + ), + }; + } - const messageProps = { - timelineAvatar: isUser ? ( - - ) : ( - - ), - timestamp: i18n.AT( - message.timestamp.length === 0 ? new Date().toLocaleString() : message.timestamp - ), - username: isUser ? i18n.YOU : i18n.ASSISTANT, - ...(message.isError ? errorStyles : {}), - }; - const isLastComment = index === currentConversation.messages.length - 1; + const messageContentWithReplacements = + replacements != null + ? Object.keys(replacements).reduce( + (acc, replacement) => acc.replaceAll(replacement, replacements[replacement]), + message.content + ) + : message.content; + const transformedMessage = { + ...message, + content: messageContentWithReplacements, + }; - // message still needs to stream, no response manipulation - if (!(message.content && message.content.length)) { return { ...messageProps, + actions: , children: ( <> ), }; - } - - const messageContentWithReplacements = - replacements != null - ? Object.keys(replacements).reduce( - (acc, replacement) => acc.replaceAll(replacement, replacements[replacement]), - message.content - ) - : message.content; - const transformedMessage = { - ...message, - content: messageContentWithReplacements, - }; - - return { - ...messageProps, - actions: , - children: ( - <> - - {isLastComment ? : null} - - ), - }; - }); + }), + ...extraLoadingComment, + ]; }; 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 6395ddaff8550..4d8607dcb0fdc 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 @@ interface Props { amendMessage: (message: string) => void; content?: string; isLastComment: boolean; + isFetching?: boolean; regenerateMessage: () => void; reader?: ReadableStreamDefaultReader; } @@ -27,18 +28,23 @@ export const StreamComment = ({ isLastComment, reader, regenerateMessage, + isFetching = false, }: Props) => { const { error, isLoading, isStreaming, pendingMessage, setComplete } = useStream({ amendMessage, content, reader, }); - const message = content ?? pendingMessage; + const message = useMemo(() => content ?? pendingMessage, [content, pendingMessage]); + const isAnythingLoading = useMemo( + () => isFetching || isLoading || isStreaming, + [isFetching, isLoading, isStreaming] + ); const controls = useMemo(() => { if (reader == null || !isLastComment) { return; } - if (isLoading || isStreaming) { + if (isAnythingLoading) { return ( { @@ -54,10 +60,10 @@ export const StreamComment = ({ ); - }, [isLastComment, isLoading, isStreaming, reader, regenerateMessage, setComplete]); + }, [isAnythingLoading, isLastComment, reader, regenerateMessage, setComplete]); return ( } + body={} error={error ? new Error(error) : undefined} controls={controls} /> From 8ff7c3b19ad713f977e215f55816e100ddeac6a7 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 26 Oct 2023 14:18:35 -0500 Subject: [PATCH 30/52] cleanup --- .../impl/assistant/use_send_messages/index.tsx | 9 +++------ x-pack/plugins/elastic_assistant/server/lib/executor.ts | 5 +---- .../public/assistant/get_comments/index.tsx | 4 ---- 3 files changed, 4 insertions(+), 14 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_send_messages/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_send_messages/index.tsx index e963b61e4ec87..f9f63aa8ef8ac 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_send_messages/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_send_messages/index.tsx @@ -5,7 +5,7 @@ * 2.0. */ -import { useCallback, useMemo, useState } from 'react'; +import { useCallback, useState } from 'react'; import { HttpSetup } from '@kbn/core-http-browser'; @@ -48,9 +48,6 @@ export const useSendMessages = (): UseSendMessages => { }, [knowledgeBase.assistantLangChain] ); - const retThis = useMemo(() => { - console.log('isLoading, at root?', isLoading); - return { isLoading, sendMessages }; - }, [isLoading, sendMessages]); - return retThis; + + return { isLoading, sendMessages }; }; diff --git a/x-pack/plugins/elastic_assistant/server/lib/executor.ts b/x-pack/plugins/elastic_assistant/server/lib/executor.ts index 595287690cf38..3f269e206e4e6 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/executor.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/executor.ts @@ -29,7 +29,6 @@ export const executeAction = async ({ }: Props): Promise => { const actionsClient = await actions.getActionsClientWithRequest(request); - console.log('one'); const actionResult = await actionsClient.execute({ actionId: connectorId, params: { @@ -42,7 +41,7 @@ export const executeAction = async ({ : { body: JSON.stringify(request.body.params.subActionParams), stream: true }, }, }); - console.log('two', actionResult); + const content = get('data.message', actionResult); if (typeof content === 'string') { return { @@ -53,7 +52,5 @@ export const executeAction = async ({ } const readable = get('data', actionResult); - console.log('typeof', typeof readable); - console.log('three', readable); return (readable as Readable).pipe(new PassThrough()); }; 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 47bb3728f9f5d..42654ad4f5db1 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 @@ -50,10 +50,6 @@ export const getComments = ({ regenerateMessage(currentConversation.id); }; - console.log('get comments return', { - isFetchingResponse, - currentConversation, - }); const extraLoadingComment = isFetchingResponse ? [ { From f3258a40ef293d193f10c8fb67b5bc2fc2d6d68a Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 30 Oct 2023 08:51:47 -0500 Subject: [PATCH 31/52] update amendMessage --- .../impl/assistant/use_conversation/index.tsx | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx index 0f815db18409c..f152e4ce44189 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx @@ -61,7 +61,7 @@ interface SetConversationProps { interface UseConversation { appendMessage: ({ conversationId, message }: AppendMessageProps) => Message[]; - amendMessage: ({ conversationId, content }: AmendMessageProps) => Message[]; + amendMessage: ({ conversationId, content }: AmendMessageProps) => void; appendReplacements: ({ conversationId, replacements, @@ -87,8 +87,7 @@ export const useConversation = (): UseConversation => { const prevConversation: Conversation | undefined = prev[conversationId]; if (prevConversation != null) { - prevConversation.messages.pop(); - messages = prevConversation.messages; + messages = prevConversation.messages.slice(0, prevConversation.messages.length - 1); const newConversation = { ...prevConversation, messages, @@ -111,13 +110,12 @@ export const useConversation = (): UseConversation => { */ const amendMessage = useCallback( ({ conversationId, content }: AmendMessageProps) => { - let messages: Message[] = []; setConversations((prev: Record) => { const prevConversation: Conversation | undefined = prev[conversationId]; if (prevConversation != null) { const message = prevConversation.messages.pop() as unknown as Message; - messages = [...prevConversation.messages, { ...message, content }]; + const messages = [...prevConversation.messages, { ...message, content }]; const newConversation = { ...prevConversation, messages, @@ -130,7 +128,6 @@ export const useConversation = (): UseConversation => { return prev; } }); - return messages; }, [setConversations] ); From 30f2624f9731cf104ac1f734acfaf113511e0478 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 30 Oct 2023 10:31:59 -0500 Subject: [PATCH 32/52] tests wip --- .../stream/stream_observable.test.ts | 239 ++++++++++++++++++ 1 file changed, 239 insertions(+) create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.test.ts diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.test.ts b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.test.ts new file mode 100644 index 0000000000000..70265e20cfe80 --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.test.ts @@ -0,0 +1,239 @@ +/* + * 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 { getStreamObservable } from './stream_observable'; +// import { getReaderValue, mockUint8Arrays } from './mock'; +import type { PromptObservableState } from './types'; +// export const getReadableStreamMock = (): ReadableStreamDefaultReader => +// ({ +// read: jest +// .fn() +// .mockResolvedValueOnce({ +// done: false, +// value: getReaderValue(mockUint8Arrays[0]), +// }) +// .mockResolvedValueOnce({ +// done: false, +// value: getReaderValue(mockUint8Arrays[1]), +// }) +// .mockResolvedValue({ +// done: true, +// }), +// cancel: jest.fn(), +// releaseLock: jest.fn(), +// closed: jest.fn().mockResolvedValue(true), +// } as unknown as ReadableStreamDefaultReader); +// const mockedStream = getReadableStreamMock(); + +describe('getStreamObservable', () => { + const mockReader = { + read: jest.fn(), + cancel: jest.fn(), + }; + + const setLoading = jest.fn(); + + beforeEach(() => { + jest.clearAllMocks(); + }); + + it.only('should emit loading state and chunks', (done) => { + const expectedStates: PromptObservableState[] = [ + { chunks: [], loading: true }, + { + chunks: [ + { + id: '1', + object: 'chunk', + created: 1635633600000, + model: 'model-1', + choices: [ + { + index: 0, + delta: { role: 'role-1', content: 'content-1' }, + finish_reason: null, + }, + ], + }, + ], + message: 'content-1', + loading: true, + }, + { + chunks: [ + { + id: '1', + object: 'chunk', + created: 1635633600000, + model: 'model-1', + choices: [ + { + index: 0, + delta: { role: 'role-1', content: 'content-1' }, + finish_reason: null, + }, + ], + }, + ], + message: 'content-1', + loading: false, + }, + ]; + + const source = getStreamObservable(mockReader, setLoading); + const emittedStates: PromptObservableState[] = []; + + source.subscribe({ + next: (state) => emittedStates.push(state), + complete: () => { + expect(emittedStates).toEqual(expectedStates); + expect(setLoading).toHaveBeenCalledWith(false); + done(); + }, + error: (err) => done(err), + }); + + // Mock the read function to return a resolved promise with expected data + mockReader.read.mockResolvedValueOnce({ + done: false, + value: new Uint8Array( + new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[1].chunks[0])}`) + ), + }); + mockReader.read.mockResolvedValueOnce({ + done: false, + value: new Uint8Array(new TextEncoder().encode('data: [DONE]\n')), + }); + mockReader.read.mockResolvedValueOnce({ + done: false, + value: new Uint8Array( + new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[1].chunks[0])}`) + ), + }); + mockReader.read.mockResolvedValueOnce({ + done: false, + value: new Uint8Array(new TextEncoder().encode('data: [DONE]\n')), + }); + mockReader.read.mockResolvedValueOnce({ + done: true, + value: undefined, + }); + }); + + it('should handle errors', (done) => { + const source = getStreamObservable(mockReader, setLoading); + const error = new Error('Test Error'); + + source.subscribe({ + next: (state) => {}, + complete: () => done(new Error('Should not complete')), + error: (err) => { + expect(err).toEqual(error); + expect(setLoading).toHaveBeenCalledWith(false); + done(); + }, + }); + + // Simulate an error + mockReader.read.mockRejectedValue(error); + }); + + it('should complete the observable when the reader is done', (done) => { + const source = getStreamObservable(mockReader, setLoading); + + source.subscribe({ + next: (state) => {}, + complete: () => { + expect(setLoading).toHaveBeenCalledWith(false); + expect(mockReader.cancel).toHaveBeenCalled(); + done(); + }, + error: (err) => done(err), + }); + + // Simulate completion + mockReader.read.mockResolvedValue({ done: true, value: undefined }); + }); + + it('should emit values with the correct delay', (done) => { + const expectedStates: PromptObservableState[] = [ + { chunks: [], loading: true }, + { + chunks: [ + { + id: '1', + object: 'chunk', + created: 1635633600000, + model: 'model-1', + choices: [ + { + index: 0, + delta: { role: 'role-1', content: 'content-1' }, + finish_reason: null, + }, + ], + }, + ], + message: 'content-1', + loading: true, + }, + { + chunks: [ + { + id: '1', + object: 'chunk', + created: 1635633600000, + model: 'model-1', + choices: [ + { + index: 0, + delta: { role: 'role-1', content: 'content-1' }, + finish_reason: null, + }, + ], + }, + ], + message: 'content-1', + loading: false, + }, + ]; + + const source = getStreamObservable(mockReader, setLoading); + const emittedStates: PromptObservableState[] = []; + + source.subscribe({ + next: (state) => emittedStates.push(state), + complete: () => { + expect(emittedStates).toEqual(expectedStates); + expect(setLoading).toHaveBeenCalledWith(false); + done(); + }, + error: (err) => done(err), + }); + + // Simulate reading data with a delay + mockReader.read.mockResolvedValueOnce({ + done: false, + value: new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[1].chunks[0])}`), + }); + mockReader.read.mockResolvedValueOnce({ + done: false, + value: new TextEncoder().encode('data: [DONE]\n'), + }); + mockReader.read.mockResolvedValueOnce({ + done: false, + value: new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[2].chunks[0])}`), + }); + mockReader.read.mockResolvedValueOnce({ + done: false, + value: new TextEncoder().encode('data: [DONE]\n'), + }); + mockReader.read.mockResolvedValueOnce({ + done: true, + value: undefined, + }); + }); +}); From 87814845f02582580cc0926a3af0ceb207c8dc39 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 30 Oct 2023 16:05:57 -0500 Subject: [PATCH 33/52] stream observable tests --- .../assistant/get_comments/stream/mock.ts | 637 ------------------ .../stream/stream_observable.test.ts | 193 ++---- .../get_comments/stream/stream_observable.ts | 30 +- .../get_comments/stream/use_stream.test.tsx | 31 +- 4 files changed, 79 insertions(+), 812 deletions(-) delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/mock.ts diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/mock.ts b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/mock.ts deleted file mode 100644 index ebba2d432d51c..0000000000000 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/mock.ts +++ /dev/null @@ -1,637 +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. - */ - -// Uint8Array returned by OpenAI -const one = { - '0': 100, - '1': 97, - '2': 116, - '3': 97, - '4': 58, - '5': 32, - '6': 123, - '7': 34, - '8': 105, - '9': 100, - '10': 34, - '11': 58, - '12': 34, - '13': 99, - '14': 104, - '15': 97, - '16': 116, - '17': 99, - '18': 109, - '19': 112, - '20': 108, - '21': 45, - '22': 56, - '23': 68, - '24': 74, - '25': 71, - '26': 72, - '27': 117, - '28': 73, - '29': 67, - '30': 74, - '31': 68, - '32': 118, - '33': 103, - '34': 116, - '35': 71, - '36': 49, - '37': 98, - '38': 97, - '39': 65, - '40': 106, - '41': 74, - '42': 112, - '43': 103, - '44': 88, - '45': 120, - '46': 120, - '47': 90, - '48': 52, - '49': 73, - '50': 121, - '51': 34, - '52': 44, - '53': 34, - '54': 111, - '55': 98, - '56': 106, - '57': 101, - '58': 99, - '59': 116, - '60': 34, - '61': 58, - '62': 34, - '63': 99, - '64': 104, - '65': 97, - '66': 116, - '67': 46, - '68': 99, - '69': 111, - '70': 109, - '71': 112, - '72': 108, - '73': 101, - '74': 116, - '75': 105, - '76': 111, - '77': 110, - '78': 46, - '79': 99, - '80': 104, - '81': 117, - '82': 110, - '83': 107, - '84': 34, - '85': 44, - '86': 34, - '87': 99, - '88': 114, - '89': 101, - '90': 97, - '91': 116, - '92': 101, - '93': 100, - '94': 34, - '95': 58, - '96': 49, - '97': 54, - '98': 57, - '99': 56, - '100': 49, - '101': 56, - '102': 50, - '103': 57, - '104': 57, - '105': 55, - '106': 44, - '107': 34, - '108': 109, - '109': 111, - '110': 100, - '111': 101, - '112': 108, - '113': 34, - '114': 58, - '115': 34, - '116': 103, - '117': 112, - '118': 116, - '119': 45, - '120': 52, - '121': 45, - '122': 48, - '123': 54, - '124': 49, - '125': 51, - '126': 34, - '127': 44, - '128': 34, - '129': 99, - '130': 104, - '131': 111, - '132': 105, - '133': 99, - '134': 101, - '135': 115, - '136': 34, - '137': 58, - '138': 91, - '139': 123, - '140': 34, - '141': 105, - '142': 110, - '143': 100, - '144': 101, - '145': 120, - '146': 34, - '147': 58, - '148': 48, - '149': 44, - '150': 34, - '151': 100, - '152': 101, - '153': 108, - '154': 116, - '155': 97, - '156': 34, - '157': 58, - '158': 123, - '159': 34, - '160': 114, - '161': 111, - '162': 108, - '163': 101, - '164': 34, - '165': 58, - '166': 34, - '167': 97, - '168': 115, - '169': 115, - '170': 105, - '171': 115, - '172': 116, - '173': 97, - '174': 110, - '175': 116, - '176': 34, - '177': 44, - '178': 34, - '179': 99, - '180': 111, - '181': 110, - '182': 116, - '183': 101, - '184': 110, - '185': 116, - '186': 34, - '187': 58, - '188': 34, - '189': 34, - '190': 125, - '191': 44, - '192': 34, - '193': 102, - '194': 105, - '195': 110, - '196': 105, - '197': 115, - '198': 104, - '199': 95, - '200': 114, - '201': 101, - '202': 97, - '203': 115, - '204': 111, - '205': 110, - '206': 34, - '207': 58, - '208': 110, - '209': 117, - '210': 108, - '211': 108, - '212': 125, - '213': 93, - '214': 125, - '215': 10, - '216': 10, - '217': 100, - '218': 97, - '219': 116, - '220': 97, - '221': 58, - '222': 32, - '223': 123, - '224': 34, - '225': 105, - '226': 100, - '227': 34, - '228': 58, - '229': 34, - '230': 99, - '231': 104, - '232': 97, - '233': 116, - '234': 99, - '235': 109, - '236': 112, - '237': 108, - '238': 45, - '239': 56, - '240': 68, - '241': 74, - '242': 71, - '243': 72, - '244': 117, - '245': 73, - '246': 67, - '247': 74, - '248': 68, - '249': 118, - '250': 103, - '251': 116, - '252': 71, - '253': 49, - '254': 98, - '255': 97, - '256': 65, - '257': 106, - '258': 74, - '259': 112, - '260': 103, - '261': 88, - '262': 120, - '263': 120, - '264': 90, - '265': 52, - '266': 73, - '267': 121, - '268': 34, - '269': 44, - '270': 34, - '271': 111, - '272': 98, - '273': 106, - '274': 101, - '275': 99, - '276': 116, - '277': 34, - '278': 58, - '279': 34, - '280': 99, - '281': 104, - '282': 97, - '283': 116, - '284': 46, - '285': 99, - '286': 111, - '287': 109, - '288': 112, - '289': 108, - '290': 101, - '291': 116, - '292': 105, - '293': 111, - '294': 110, - '295': 46, - '296': 99, - '297': 104, - '298': 117, - '299': 110, - '300': 107, - '301': 34, - '302': 44, - '303': 34, - '304': 99, - '305': 114, - '306': 101, - '307': 97, - '308': 116, - '309': 101, - '310': 100, - '311': 34, - '312': 58, - '313': 49, - '314': 54, - '315': 57, - '316': 56, - '317': 49, - '318': 56, - '319': 50, - '320': 57, - '321': 57, - '322': 55, - '323': 44, - '324': 34, - '325': 109, - '326': 111, - '327': 100, - '328': 101, - '329': 108, - '330': 34, - '331': 58, - '332': 34, - '333': 103, - '334': 112, - '335': 116, - '336': 45, - '337': 52, - '338': 45, - '339': 48, - '340': 54, - '341': 49, - '342': 51, - '343': 34, - '344': 44, - '345': 34, - '346': 99, - '347': 104, - '348': 111, - '349': 105, - '350': 99, - '351': 101, - '352': 115, - '353': 34, - '354': 58, - '355': 91, - '356': 123, - '357': 34, - '358': 105, - '359': 110, - '360': 100, - '361': 101, - '362': 120, - '363': 34, - '364': 58, - '365': 48, - '366': 44, - '367': 34, - '368': 100, - '369': 101, - '370': 108, - '371': 116, - '372': 97, - '373': 34, - '374': 58, - '375': 123, - '376': 34, - '377': 99, - '378': 111, - '379': 110, - '380': 116, - '381': 101, - '382': 110, - '383': 116, - '384': 34, - '385': 58, - '386': 34, - '387': 67, - '388': 104, - '389': 34, - '390': 125, - '391': 44, - '392': 34, - '393': 102, - '394': 105, - '395': 110, - '396': 105, - '397': 115, - '398': 104, - '399': 95, - '400': 114, - '401': 101, - '402': 97, - '403': 115, - '404': 111, - '405': 110, - '406': 34, - '407': 58, - '408': 110, - '409': 117, - '410': 108, - '411': 108, - '412': 125, - '413': 93, - '414': 125, - '415': 10, - '416': 10, -}; -const two = { - '0': 100, - '1': 97, - '2': 116, - '3': 97, - '4': 58, - '5': 32, - '6': 123, - '7': 34, - '8': 105, - '9': 100, - '10': 34, - '11': 58, - '12': 34, - '13': 99, - '14': 104, - '15': 97, - '16': 116, - '17': 99, - '18': 109, - '19': 112, - '20': 108, - '21': 45, - '22': 56, - '23': 68, - '24': 74, - '25': 71, - '26': 72, - '27': 117, - '28': 73, - '29': 67, - '30': 74, - '31': 68, - '32': 118, - '33': 103, - '34': 116, - '35': 71, - '36': 49, - '37': 98, - '38': 97, - '39': 65, - '40': 106, - '41': 74, - '42': 112, - '43': 103, - '44': 88, - '45': 120, - '46': 120, - '47': 90, - '48': 52, - '49': 73, - '50': 121, - '51': 34, - '52': 44, - '53': 34, - '54': 111, - '55': 98, - '56': 106, - '57': 101, - '58': 99, - '59': 116, - '60': 34, - '61': 58, - '62': 34, - '63': 99, - '64': 104, - '65': 97, - '66': 116, - '67': 46, - '68': 99, - '69': 111, - '70': 109, - '71': 112, - '72': 108, - '73': 101, - '74': 116, - '75': 105, - '76': 111, - '77': 110, - '78': 46, - '79': 99, - '80': 104, - '81': 117, - '82': 110, - '83': 107, - '84': 34, - '85': 44, - '86': 34, - '87': 99, - '88': 114, - '89': 101, - '90': 97, - '91': 116, - '92': 101, - '93': 100, - '94': 34, - '95': 58, - '96': 49, - '97': 54, - '98': 57, - '99': 56, - '100': 49, - '101': 56, - '102': 50, - '103': 57, - '104': 57, - '105': 55, - '106': 44, - '107': 34, - '108': 109, - '109': 111, - '110': 100, - '111': 101, - '112': 108, - '113': 34, - '114': 58, - '115': 34, - '116': 103, - '117': 112, - '118': 116, - '119': 45, - '120': 52, - '121': 45, - '122': 48, - '123': 54, - '124': 49, - '125': 51, - '126': 34, - '127': 44, - '128': 34, - '129': 99, - '130': 104, - '131': 111, - '132': 105, - '133': 99, - '134': 101, - '135': 115, - '136': 34, - '137': 58, - '138': 91, - '139': 123, - '140': 34, - '141': 105, - '142': 110, - '143': 100, - '144': 101, - '145': 120, - '146': 34, - '147': 58, - '148': 48, - '149': 44, - '150': 34, - '151': 100, - '152': 101, - '153': 108, - '154': 116, - '155': 97, - '156': 34, - '157': 58, - '158': 123, - '159': 34, - '160': 99, - '161': 111, - '162': 110, - '163': 116, - '164': 101, - '165': 110, - '166': 116, - '167': 34, - '168': 58, - '169': 34, - '170': 101, - '171': 100, - '172': 100, - '173': 97, - '174': 114, - '175': 34, - '176': 125, - '177': 44, - '178': 34, - '179': 102, - '180': 105, - '181': 110, - '182': 105, - '183': 115, - '184': 104, - '185': 95, - '186': 114, - '187': 101, - '188': 97, - '189': 115, - '190': 111, - '191': 110, - '192': 34, - '193': 58, - '194': 110, - '195': 117, - '196': 108, - '197': 108, - '198': 125, - '199': 93, - '200': 125, - '201': 10, - '202': 10, -}; -export const mockUint8Arrays = [one, two]; -export const getReaderValue = (obj: { [key: string]: number }): Buffer => { - const parsedArray = Object.values(obj); - return Buffer.from(new Uint8Array(parsedArray)); -}; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.test.ts b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.test.ts index 70265e20cfe80..9a63621021cc3 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.test.ts +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.test.ts @@ -7,40 +7,23 @@ import { getStreamObservable } from './stream_observable'; // import { getReaderValue, mockUint8Arrays } from './mock'; import type { PromptObservableState } from './types'; -// export const getReadableStreamMock = (): ReadableStreamDefaultReader => -// ({ -// read: jest -// .fn() -// .mockResolvedValueOnce({ -// done: false, -// value: getReaderValue(mockUint8Arrays[0]), -// }) -// .mockResolvedValueOnce({ -// done: false, -// value: getReaderValue(mockUint8Arrays[1]), -// }) -// .mockResolvedValue({ -// done: true, -// }), -// cancel: jest.fn(), -// releaseLock: jest.fn(), -// closed: jest.fn().mockResolvedValue(true), -// } as unknown as ReadableStreamDefaultReader); -// const mockedStream = getReadableStreamMock(); - +import { Subject } from 'rxjs'; describe('getStreamObservable', () => { const mockReader = { read: jest.fn(), cancel: jest.fn(), }; + const typedReader = mockReader as unknown as ReadableStreamDefaultReader; + const setLoading = jest.fn(); beforeEach(() => { jest.clearAllMocks(); }); - it.only('should emit loading state and chunks', (done) => { + it('should emit loading state and chunks', (done) => { + const completeSubject = new Subject(); const expectedStates: PromptObservableState[] = [ { chunks: [], loading: true }, { @@ -83,157 +66,67 @@ describe('getStreamObservable', () => { }, ]; - const source = getStreamObservable(mockReader, setLoading); + mockReader.read + .mockResolvedValueOnce({ + done: false, + value: new Uint8Array( + new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[1].chunks[0])}`) + ), + }) + .mockResolvedValueOnce({ + done: false, + value: new Uint8Array(new TextEncoder().encode(``)), + }) + .mockResolvedValueOnce({ + done: false, + value: new Uint8Array(new TextEncoder().encode('data: [DONE]\n')), + }) + .mockResolvedValue({ + done: true, + }); + + const source = getStreamObservable(typedReader, setLoading); const emittedStates: PromptObservableState[] = []; source.subscribe({ next: (state) => emittedStates.push(state), complete: () => { expect(emittedStates).toEqual(expectedStates); - expect(setLoading).toHaveBeenCalledWith(false); done(); + + completeSubject.subscribe({ + next: () => { + expect(setLoading).toHaveBeenCalledWith(false); + expect(typedReader.cancel).toHaveBeenCalled(); + done(); + }, + }); }, error: (err) => done(err), }); - - // Mock the read function to return a resolved promise with expected data - mockReader.read.mockResolvedValueOnce({ - done: false, - value: new Uint8Array( - new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[1].chunks[0])}`) - ), - }); - mockReader.read.mockResolvedValueOnce({ - done: false, - value: new Uint8Array(new TextEncoder().encode('data: [DONE]\n')), - }); - mockReader.read.mockResolvedValueOnce({ - done: false, - value: new Uint8Array( - new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[1].chunks[0])}`) - ), - }); - mockReader.read.mockResolvedValueOnce({ - done: false, - value: new Uint8Array(new TextEncoder().encode('data: [DONE]\n')), - }); - mockReader.read.mockResolvedValueOnce({ - done: true, - value: undefined, - }); }); it('should handle errors', (done) => { - const source = getStreamObservable(mockReader, setLoading); + const completeSubject = new Subject(); const error = new Error('Test Error'); + // Simulate an error + mockReader.read.mockRejectedValue(error); + const source = getStreamObservable(typedReader, setLoading); source.subscribe({ next: (state) => {}, complete: () => done(new Error('Should not complete')), error: (err) => { expect(err).toEqual(error); - expect(setLoading).toHaveBeenCalledWith(false); done(); - }, - }); - - // Simulate an error - mockReader.read.mockRejectedValue(error); - }); - - it('should complete the observable when the reader is done', (done) => { - const source = getStreamObservable(mockReader, setLoading); - - source.subscribe({ - next: (state) => {}, - complete: () => { - expect(setLoading).toHaveBeenCalledWith(false); - expect(mockReader.cancel).toHaveBeenCalled(); - done(); - }, - error: (err) => done(err), - }); - - // Simulate completion - mockReader.read.mockResolvedValue({ done: true, value: undefined }); - }); - - it('should emit values with the correct delay', (done) => { - const expectedStates: PromptObservableState[] = [ - { chunks: [], loading: true }, - { - chunks: [ - { - id: '1', - object: 'chunk', - created: 1635633600000, - model: 'model-1', - choices: [ - { - index: 0, - delta: { role: 'role-1', content: 'content-1' }, - finish_reason: null, - }, - ], - }, - ], - message: 'content-1', - loading: true, - }, - { - chunks: [ - { - id: '1', - object: 'chunk', - created: 1635633600000, - model: 'model-1', - choices: [ - { - index: 0, - delta: { role: 'role-1', content: 'content-1' }, - finish_reason: null, - }, - ], + completeSubject.subscribe({ + next: () => { + expect(setLoading).toHaveBeenCalledWith(false); + expect(typedReader.cancel).toHaveBeenCalled(); + done(); }, - ], - message: 'content-1', - loading: false, + }); }, - ]; - - const source = getStreamObservable(mockReader, setLoading); - const emittedStates: PromptObservableState[] = []; - - source.subscribe({ - next: (state) => emittedStates.push(state), - complete: () => { - expect(emittedStates).toEqual(expectedStates); - expect(setLoading).toHaveBeenCalledWith(false); - done(); - }, - error: (err) => done(err), - }); - - // Simulate reading data with a delay - mockReader.read.mockResolvedValueOnce({ - done: false, - value: new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[1].chunks[0])}`), - }); - mockReader.read.mockResolvedValueOnce({ - done: false, - value: new TextEncoder().encode('data: [DONE]\n'), - }); - mockReader.read.mockResolvedValueOnce({ - done: false, - value: new TextEncoder().encode(`data: ${JSON.stringify(expectedStates[2].chunks[0])}`), - }); - mockReader.read.mockResolvedValueOnce({ - done: false, - value: new TextEncoder().encode('data: [DONE]\n'), - }); - mockReader.read.mockResolvedValueOnce({ - done: true, - value: undefined, }); }); }); diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts index 1e5fe552835dc..126cb81814d11 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts @@ -25,10 +25,9 @@ export const getStreamObservable = ( observer.next({ chunks: [], loading: true }); const decoder = new TextDecoder(); const chunks: Chunk[] = []; - let prev: string = ''; function read() { reader - ?.read() + .read() .then(({ done, value }: { done: boolean; value?: Uint8Array }) => { try { if (done) { @@ -40,17 +39,16 @@ export const getStreamObservable = ( observer.complete(); return; } - let lines: string[] = (prev + decoder.decode(value)).split('\n'); - const lastLine: string = lines[lines.length - 1]; - const isPartialChunk: boolean = !!lastLine && lastLine !== 'data: [DONE]'; - if (isPartialChunk) { - prev = lastLine; - lines.pop(); - } else { - prev = ''; - } - lines = lines.map((str) => str.substring(6)).filter((str) => !!str && str !== '[DONE]'); - const nextChunks: Chunk[] = lines.map((line) => JSON.parse(line)); + + const nextChunks: Chunk[] = decoder + .decode(value) + .split('\n') + // every line starts with "data: ", we remove it and are left with stringified JSON or the string "[DONE]" + .map((str) => str.substring(6)) + // filter out empty lines and the "[DONE]" string + .filter((str) => !!str && str !== '[DONE]') + .map((line) => JSON.parse(line)); + nextChunks.forEach((chunk) => { chunks.push(chunk); observer.next({ @@ -108,11 +106,7 @@ export const getStreamObservable = ( ); function getMessageFromChunks(chunks: Chunk[]) { - let message = ''; - chunks.forEach((chunk) => { - message += chunk.choices[0]?.delta.content ?? ''; - }); - return message; + return chunks.map((chunk) => chunk.choices[0]?.delta.content ?? '').join(''); } export const getDumbObservable = () => new Observable(); diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx index 368f7b90fee81..4fbecfac870e1 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.test.tsx @@ -7,20 +7,36 @@ import { renderHook } from '@testing-library/react-hooks'; import { useStream } from './use_stream'; -import { getReaderValue, mockUint8Arrays } from './mock'; const amendMessage = jest.fn(); const reader = jest.fn(); const cancel = jest.fn(); +const exampleChunk = { + id: '1', + object: 'chunk', + created: 1635633600000, + model: 'model-1', + choices: [ + { + index: 0, + delta: { role: 'role-1', content: 'content-1' }, + finish_reason: null, + }, + ], +}; const readerComplete = { read: reader .mockResolvedValueOnce({ done: false, - value: getReaderValue(mockUint8Arrays[0]), + value: new Uint8Array(new TextEncoder().encode(`data: ${JSON.stringify(exampleChunk)}`)), }) .mockResolvedValueOnce({ done: false, - value: getReaderValue(mockUint8Arrays[1]), + value: new Uint8Array(new TextEncoder().encode(``)), + }) + .mockResolvedValueOnce({ + done: false, + value: new Uint8Array(new TextEncoder().encode('data: [DONE]\n')), }) .mockResolvedValue({ done: true, @@ -53,7 +69,7 @@ describe('useStream', () => { error: undefined, isLoading: true, isStreaming: true, - pendingMessage: 'Ch', + pendingMessage: 'content-1', setComplete: expect.any(Function), }); }); @@ -63,11 +79,12 @@ describe('useStream', () => { error: undefined, isLoading: false, isStreaming: false, - pendingMessage: 'Cheddar', + pendingMessage: 'content-1', setComplete: expect.any(Function), }); }); - expect(reader).toHaveBeenCalledTimes(3); + + expect(reader).toHaveBeenCalledTimes(4); }); it('should not call observable when content is provided', () => { @@ -87,7 +104,7 @@ describe('useStream', () => { .fn() .mockResolvedValueOnce({ done: false, - value: getReaderValue(mockUint8Arrays[0]), + value: new Uint8Array(new TextEncoder().encode(`data: ${JSON.stringify(exampleChunk)}`)), }) .mockRejectedValue(new Error(errorMessage)), cancel, From a34a98054c103488ee6914755bf04b01a86181d2 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 09:04:40 -0500 Subject: [PATCH 34/52] fix tests --- .../impl/assistant/api.test.tsx | 14 +++++--- .../impl/assistant/api.tsx | 3 +- .../impl/assistant/chat_send/index.test.tsx | 32 +++---------------- .../impl/assistant/chat_send/index.tsx | 6 ++-- 4 files changed, 18 insertions(+), 37 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.test.tsx index e8feefbfd2533..b1d9145a9e612 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.test.tsx @@ -95,7 +95,7 @@ describe('API tests', () => { const result = await fetchConnectorExecuteAction(testProps); - expect(result).toEqual({ response: API_ERROR, isError: true }); + expect(result).toEqual({ response: API_ERROR, isStream: false, isError: true }); }); it('returns API_ERROR when there are no choices', async () => { @@ -109,7 +109,7 @@ describe('API tests', () => { const result = await fetchConnectorExecuteAction(testProps); - expect(result).toEqual({ response: API_ERROR, isError: true }); + expect(result).toEqual({ response: API_ERROR, isStream: false, isError: true }); }); it('returns the value of the action_input property when assistantLangChain is true, and `content` has properly prefixed and suffixed JSON with the action_input property', async () => { @@ -129,7 +129,11 @@ describe('API tests', () => { const result = await fetchConnectorExecuteAction(testProps); - expect(result).toEqual({ response: 'value from action_input', isError: false }); + expect(result).toEqual({ + response: 'value from action_input', + isStream: false, + isError: false, + }); }); it('returns the original content when assistantLangChain is true, and `content` has properly formatted JSON WITHOUT the action_input property', async () => { @@ -149,7 +153,7 @@ describe('API tests', () => { const result = await fetchConnectorExecuteAction(testProps); - expect(result).toEqual({ response, isError: false }); + expect(result).toEqual({ response, isStream: false, isError: false }); }); it('returns the original when assistantLangChain is true, and `content` is not JSON', async () => { @@ -169,7 +173,7 @@ describe('API tests', () => { const result = await fetchConnectorExecuteAction(testProps); - expect(result).toEqual({ response, isError: false }); + expect(result).toEqual({ response, isStream: false, isError: false }); }); }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index 0f7eade9cc31e..68d419d5dc789 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -62,7 +62,7 @@ export const fetchConnectorExecuteAction = async ({ // My "Feature Flag", turn to false before merging // In part 2 I will make enhancements to invokeAI to make it work with both openA, but to keep it to a Security Soltuion only review on this PR, // I'm calling the stream action directly - const isStream = true; + const isStream = false; const requestBody = isStream ? { params: { @@ -119,6 +119,7 @@ export const fetchConnectorExecuteAction = async ({ }>(`/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`, { method: 'POST', body: JSON.stringify(requestBody), + headers: { 'Content-Type': 'application/json' }, signal, }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.test.tsx index bd54136fae4f1..a8e8e064524a3 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.test.tsx @@ -9,50 +9,26 @@ import React from 'react'; import { fireEvent, render, waitFor } from '@testing-library/react'; import { ChatSend, Props } from '.'; import { TestProviders } from '../../mock/test_providers/test_providers'; -import { useChatSend } from './use_chat_send'; -import { defaultSystemPrompt, mockSystemPrompt } from '../../mock/system_prompt'; -import { emptyWelcomeConvo } from '../../mock/conversation'; -import { HttpSetup } from '@kbn/core-http-browser'; jest.mock('./use_chat_send'); -const testProps: Props = { - selectedPromptContexts: {}, - allSystemPrompts: [defaultSystemPrompt, mockSystemPrompt], - currentConversation: emptyWelcomeConvo, - http: { - basePath: { - basePath: '/mfg', - serverBasePath: '/mfg', - }, - anonymousPaths: {}, - externalUrl: {}, - } as unknown as HttpSetup, - editingSystemPromptId: defaultSystemPrompt.id, - setEditingSystemPromptId: () => {}, - setPromptTextPreview: () => {}, - setSelectedPromptContexts: () => {}, - setUserPrompt: () => {}, - isDisabled: false, - shouldRefocusPrompt: false, - userPrompt: '', -}; const handleButtonSendMessage = jest.fn(); const handleOnChatCleared = jest.fn(); const handlePromptChange = jest.fn(); const handleSendMessage = jest.fn(); -const chatSend = { +const testProps: Props = { handleButtonSendMessage, handleOnChatCleared, handlePromptChange, handleSendMessage, isLoading: false, + isDisabled: false, + shouldRefocusPrompt: false, + userPrompt: '', }; - describe('ChatSend', () => { beforeEach(() => { jest.clearAllMocks(); - (useChatSend as jest.Mock).mockReturnValue(chatSend); }); it('the prompt updates when the text area changes', async () => { const { getByTestId } = render(, { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.tsx index 2cabf27839236..fe8a2be756047 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.tsx @@ -23,14 +23,14 @@ export interface Props extends UseChatSend { * Allows the user to clear the chat and switch between different system prompts. */ export const ChatSend: React.FC = ({ - isDisabled, - userPrompt, - shouldRefocusPrompt, handleButtonSendMessage, handleOnChatCleared, handlePromptChange, handleSendMessage, + isDisabled, isLoading, + shouldRefocusPrompt, + userPrompt, }) => { // For auto-focusing prompt within timeline const promptTextAreaRef = useRef(null); From 14d7ce1038cc0497df25ba7a94ebe29cc4f15b03 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 09:40:15 -0500 Subject: [PATCH 35/52] more tests --- .../chat_send/use_chat_send.test.tsx | 15 ++ .../assistant/use_conversation/index.test.tsx | 82 ++++++++- .../server/lib/executor.test.ts | 156 ++++++++++++++++++ .../elastic_assistant/server/lib/executor.ts | 10 +- .../routes/post_actions_connector_execute.ts | 2 +- 5 files changed, 260 insertions(+), 5 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/lib/executor.test.ts diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx index 0b139177d02b0..8dfa8699048ac 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx @@ -23,6 +23,7 @@ const setSelectedPromptContexts = jest.fn(); const setUserPrompt = jest.fn(); const sendMessages = jest.fn(); const appendMessage = jest.fn(); +const removeLastMessage = jest.fn(); const appendReplacements = jest.fn(); const clearConversation = jest.fn(); @@ -55,6 +56,7 @@ describe('use chat send', () => { (useConversation as jest.Mock).mockReturnValue({ appendMessage, appendReplacements, + removeLastMessage, clearConversation, }); }); @@ -106,4 +108,17 @@ describe('use chat send', () => { expect(appendMessage.mock.calls[0][0].message.content).toEqual(`\n\n${promptText}`); }); }); + it('handleRegenerateResponse removes the last message of the conversation, resends the convo to GenAI, and appends the message received', async () => { + const { result } = renderHook(() => + useChatSend({ ...testProps, currentConversation: welcomeConvo }) + ); + + result.current.handleRegenerateResponse(); + expect(removeLastMessage).toHaveBeenCalledWith('Welcome'); + + await waitFor(() => { + expect(sendMessages).toHaveBeenCalled(); + expect(appendMessage.mock.calls[0][0].message.content).toEqual(robotMessage.response); + }); + }); }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx index d52460bd95a0f..de83c8c107ad2 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx @@ -16,10 +16,15 @@ const message = { role: 'user' as ConversationRole, timestamp: '10/04/2023, 1:00:36 PM', }; +const anotherMessage = { + content: 'I am a robot', + role: 'assistant' as ConversationRole, + timestamp: '10/04/2023, 1:00:46 PM', +}; const mockConvo = { id: 'new-convo', - messages: [message], + messages: [message, anotherMessage], apiConfig: { defaultSystemPromptId: 'default-system-prompt' }, theme: { title: 'Elastic AI Assistant', @@ -253,4 +258,79 @@ describe('useConversation', () => { }); }); }); + + it('should remove the last message from a conversation when called with valid conversationId', async () => { + await act(async () => { + const setConversations = jest.fn(); + const { result, waitForNextUpdate } = renderHook(() => useConversation(), { + wrapper: ({ children }) => ( + ({ + [alertConvo.id]: alertConvo, + [welcomeConvo.id]: welcomeConvo, + [mockConvo.id]: mockConvo, + }), + setConversations, + }} + > + {children} + + ), + }); + await waitForNextUpdate(); + + const removeResult = result.current.removeLastMessage('new-convo'); + + expect(removeResult).toEqual([message]); + expect(setConversations).toHaveBeenCalledWith({ + [alertConvo.id]: alertConvo, + [welcomeConvo.id]: welcomeConvo, + [mockConvo.id]: { ...mockConvo, messages: [message] }, + }); + }); + }); + + it.only('amendMessage updates the last message of conversation[] for a given conversationId with provided content', async () => { + await act(async () => { + const setConversations = jest.fn(); + const { result, waitForNextUpdate } = renderHook(() => useConversation(), { + wrapper: ({ children }) => ( + ({ + [alertConvo.id]: alertConvo, + [welcomeConvo.id]: welcomeConvo, + [mockConvo.id]: mockConvo, + }), + setConversations, + }} + > + {children} + + ), + }); + await waitForNextUpdate(); + + result.current.amendMessage({ + conversationId: 'new-convo', + content: 'hello world', + }); + + expect(setConversations).toHaveBeenCalledWith({ + [alertConvo.id]: alertConvo, + [welcomeConvo.id]: welcomeConvo, + [mockConvo.id]: { + ...mockConvo, + messages: [ + message, + { + ...anotherMessage, + content: 'hello world', + }, + ], + }, + }); + }); + }); }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/executor.test.ts b/x-pack/plugins/elastic_assistant/server/lib/executor.test.ts new file mode 100644 index 0000000000000..c43e99e6caf12 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/executor.test.ts @@ -0,0 +1,156 @@ +/* + * 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 { executeAction, Props } from './executor'; +import { PassThrough } from 'stream'; + +describe('executeAction', () => { + it('should execute an action and return a StaticResponse when the response from the actions framework is a string', async () => { + const actions = { + getActionsClientWithRequest: jest.fn().mockResolvedValue({ + execute: jest.fn().mockResolvedValue({ + data: { + message: 'Test message', + }, + }), + }), + }; + const request = { + body: { + params: { + subActionParams: {}, + }, + }, + }; + const connectorId = 'testConnectorId'; + + const result = await executeAction({ actions, request, connectorId } as unknown as Props); + + expect(result).toEqual({ + connector_id: connectorId, + data: 'Test message', + status: 'ok', + }); + }); + + it('should execute an action and return a Readable object when the response from the actions framework is a stream', async () => { + const readableStream = new PassThrough(); + const actions = { + getActionsClientWithRequest: jest.fn().mockResolvedValue({ + execute: jest.fn().mockResolvedValue({ + data: readableStream, + }), + }), + }; + const request = { + body: { + params: { + subActionParams: {}, + }, + }, + }; + const connectorId = 'testConnectorId'; + + const result = await executeAction({ actions, request, connectorId } as unknown as Props); + + expect(JSON.stringify(result)).toStrictEqual( + JSON.stringify(readableStream.pipe(new PassThrough())) + ); + }); + + it('should throw an error if the actions plugin fails to retrieve the actions client', async () => { + const actions = { + getActionsClientWithRequest: jest + .fn() + .mockRejectedValue(new Error('Failed to retrieve actions client')), + }; + const request = { + body: { + params: { + subActionParams: {}, + }, + }, + }; + const connectorId = 'testConnectorId'; + + await expect( + executeAction({ actions, request, connectorId } as unknown as Props) + ).rejects.toThrowError('Failed to retrieve actions client'); + }); + + it('should throw an error if the actions client fails to execute the action', async () => { + const actions = { + getActionsClientWithRequest: jest.fn().mockResolvedValue({ + execute: jest.fn().mockRejectedValue(new Error('Failed to execute action')), + }), + }; + const request = { + body: { + params: { + subActionParams: {}, + }, + }, + }; + const connectorId = 'testConnectorId'; + + await expect( + executeAction({ actions, request, connectorId } as unknown as Props) + ).rejects.toThrowError('Failed to execute action'); + }); + + it('should throw an error when the response from the actions framework is null or undefined', async () => { + const actions = { + getActionsClientWithRequest: jest.fn().mockResolvedValue({ + execute: jest.fn().mockResolvedValue({ + data: null, + }), + }), + }; + const request = { + body: { + params: { + subActionParams: {}, + }, + }, + }; + const connectorId = 'testConnectorId'; + + try { + await executeAction({ actions, request, connectorId } as unknown as Props); + } catch (e) { + expect(e.message).toBe('Unexpected action result'); + } + }); + + it('should return a StaticResponse object with connector_id, data, and status fields when the response from the actions framework is a string', async () => { + const actions = { + getActionsClientWithRequest: jest.fn().mockResolvedValue({ + execute: jest.fn().mockResolvedValue({ + data: { + message: 'Test message', + }, + }), + }), + }; + const request = { + body: { + params: { + subActionParams: {}, + }, + }, + }; + const connectorId = 'testConnectorId'; + + const result = await executeAction({ actions, request, connectorId } as unknown as Props); + + expect(result).toEqual({ + connector_id: connectorId, + data: 'Test message', + status: 'ok', + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/executor.ts b/x-pack/plugins/elastic_assistant/server/lib/executor.ts index 3f269e206e4e6..6e8a44f479bac 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/executor.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/executor.ts @@ -11,7 +11,7 @@ import { KibanaRequest } from '@kbn/core-http-server'; import { PassThrough, Readable } from 'stream'; import { RequestBody } from './langchain/types'; -interface Props { +export interface Props { actions: ActionsPluginStart; connectorId: string; request: KibanaRequest; @@ -50,7 +50,11 @@ export const executeAction = async ({ status: 'ok', }; } - const readable = get('data', actionResult); + const readable = get('data', actionResult) as Readable; - return (readable as Readable).pipe(new PassThrough()); + if (typeof readable?.read !== 'function') { + throw new Error('Unexpected action result'); + } + + return readable.pipe(new PassThrough()); }; 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 b886532f9ff88..299d8ade24a3f 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 @@ -41,10 +41,10 @@ export const postActionsConnectorExecuteRoute = ( // get the actions plugin start contract from the request context: const actions = (await context.elasticAssistant).actions; + // if not langchain, call execute action directly and return the response: if (!request.body.assistantLangChain) { logger.debug('Executing via actions framework directly, assistantLangChain: false'); - const result = await executeAction({ actions, request, connectorId }); return response.ok({ body: result, From cc389fec8c21ac394dcd11a6d245f14876017ed1 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 09:44:28 -0500 Subject: [PATCH 36/52] fix types --- .../impl/assistant/chat_send/index.test.tsx | 2 ++ .../kbn-elastic-assistant/impl/assistant_context/index.tsx | 4 ++-- .../impl/connectorland/connector_setup/index.tsx | 2 +- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.test.tsx index a8e8e064524a3..a7eac5c362ca0 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/index.test.tsx @@ -16,11 +16,13 @@ const handleButtonSendMessage = jest.fn(); const handleOnChatCleared = jest.fn(); const handlePromptChange = jest.fn(); const handleSendMessage = jest.fn(); +const handleRegenerateResponse = jest.fn(); const testProps: Props = { handleButtonSendMessage, handleOnChatCleared, handlePromptChange, handleSendMessage, + handleRegenerateResponse, isLoading: false, isDisabled: false, shouldRefocusPrompt: false, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index efc08ca0d0183..310c975b097bf 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -35,7 +35,7 @@ import { SYSTEM_PROMPT_LOCAL_STORAGE_KEY, } from './constants'; import { CONVERSATIONS_TAB, SettingsTabs } from '../assistant/settings/assistant_settings'; -import { AssistantAvailability, AssistantTelemetry, Message } from './types'; +import { AssistantAvailability, AssistantTelemetry } from './types'; export interface ShowAssistantOverlayProps { showOverlay: boolean; @@ -116,7 +116,7 @@ export interface UseAssistantContext { }: { conversationId: string; content: string; - }) => Message[]; + }) => void; regenerateMessage: () => void; showAnonymizedValues: boolean; }) => EuiCommentProps[]; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/connectorland/connector_setup/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/connectorland/connector_setup/index.tsx index e51d356679bc3..c80cfe0ed1007 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/connectorland/connector_setup/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/connectorland/connector_setup/index.tsx @@ -131,7 +131,7 @@ export const useConnectorSetup = ({ (message?.presentation?.stream ?? false) && currentMessageIndex !== length - 1; return ( Date: Tue, 31 Oct 2023 10:43:19 -0500 Subject: [PATCH 37/52] more tests --- .../buttons/regenerate_response_button.tsx | 2 +- .../stream/buttons/stop_generating_button.tsx | 2 +- .../get_comments/stream/esql_code_block.tsx | 75 ------------------ .../get_comments/stream/index.test.tsx | 79 +++++++++++++++++++ .../assistant/get_comments/stream/index.tsx | 2 + .../get_comments/stream/message_panel.tsx | 4 +- .../get_comments/stream/message_text.tsx | 5 +- 7 files changed, 89 insertions(+), 80 deletions(-) delete mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/esql_code_block.tsx create mode 100644 x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.tsx index 3e7ac90dcacda..e9121a9ed9f20 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/regenerate_response_button.tsx @@ -14,7 +14,7 @@ export function RegenerateResponseButton(props: Partial) { return ( diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.tsx index 9a8f1e8ae3ae7..5144e82b1125f 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/buttons/stop_generating_button.tsx @@ -13,7 +13,7 @@ import { i18n } from '@kbn/i18n'; export function StopGeneratingButton(props: Partial) { return ( - - - - {value} - - - - - - onActionClick({ type: 'executeEsqlQuery', query: value })} - disabled={actionsDisabled} - > - {i18n.translate('xpack.securitySolution.aiAssistant.runThisQuery', { - defaultMessage: 'Run this query', - })} - - - - - - - ); -} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx new file mode 100644 index 0000000000000..1d27c13e56a5f --- /dev/null +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx @@ -0,0 +1,79 @@ +/* + * 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 React from 'react'; +import { render, screen, fireEvent } from '@testing-library/react'; +import { StreamComment } from '.'; +import { useStream } from './use_stream'; +const mockSetComplete = jest.fn(); + +jest.mock('./use_stream'); + +const content = 'Test Content'; +const testProps = { + amendMessage: jest.fn(), + content, + isLastComment: true, + regenerateMessage: jest.fn(), +}; + +describe('StreamComment', () => { + beforeEach(() => { + jest.clearAllMocks(); + (useStream as jest.Mock).mockReturnValue({ + error: null, + isLoading: false, + isStreaming: false, + pendingMessage: 'Test Message', + setComplete: mockSetComplete, + }); + }); + it('renders content correctly', () => { + render(); + + expect(screen.getByText(content)).toBeInTheDocument(); + }); + + it('renders cursor when content is loading', () => { + render(); + expect(screen.getByTestId('cursor')).toBeInTheDocument(); + expect(screen.queryByTestId('stopGeneratingButton')).not.toBeInTheDocument(); + }); + + it('renders cursor and stopGeneratingButton when reader is loading', () => { + render(); + expect(screen.getByTestId('stopGeneratingButton')).toBeInTheDocument(); + expect(screen.getByTestId('cursor')).toBeInTheDocument(); + }); + + it('renders controls correctly when not loading', () => { + render(); + + expect(screen.getByTestId('regenerateResponseButton')).toBeInTheDocument(); + }); + + it('calls setComplete when StopGeneratingButton is clicked', () => { + render(); + + fireEvent.click(screen.getByTestId('stopGeneratingButton')); + + expect(mockSetComplete).toHaveBeenCalled(); + }); + + it('displays an error message correctly', () => { + (useStream as jest.Mock).mockReturnValue({ + error: 'Test Error Message', + isLoading: false, + isStreaming: false, + pendingMessage: 'Test Message', + setComplete: mockSetComplete, + }); + render(); + + expect(screen.getByTestId('messsage-error')).toBeInTheDocument(); + }); +}); 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 4d8607dcb0fdc..510f81c3f8fb4 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 @@ -61,6 +61,8 @@ export const StreamComment = ({ ); }, [isAnythingLoading, isLastComment, reader, regenerateMessage, setComplete]); + + console.log('error???', error); return ( } diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_panel.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_panel.tsx index 78a1e8fae4778..91c7f6ea4247d 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_panel.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/message_panel.tsx @@ -19,10 +19,10 @@ export function MessagePanel(props: Props) { <> {props.body} {props.error ? ( - <> + {props.body ? : null} - + ) : null} {props.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 554d97a8e4f34..c0eeb9d53f191 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 @@ -47,7 +47,9 @@ const cursorCss = css` background: rgba(0, 0, 0, 0.25); `; -const Cursor = () => ; +const Cursor = () => ( + +); // a weird combination of different whitespace chars to make sure it stays // invisible even when we cannot properly parse the text while still being @@ -158,6 +160,7 @@ export function MessageText({ loading, content }: Props) { return ( From db40b5a93be48caefca91d19b2a34624de1356bf Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 10:47:21 -0500 Subject: [PATCH 38/52] fix types in test --- .../public/assistant/get_comments/stream/index.test.tsx | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx index 1d27c13e56a5f..ef7bb7140543f 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx @@ -21,6 +21,8 @@ const testProps = { regenerateMessage: jest.fn(), }; +const mockReader = jest.fn() as unknown as ReadableStreamDefaultReader; + describe('StreamComment', () => { beforeEach(() => { jest.clearAllMocks(); @@ -45,19 +47,19 @@ describe('StreamComment', () => { }); it('renders cursor and stopGeneratingButton when reader is loading', () => { - render(); + render(); expect(screen.getByTestId('stopGeneratingButton')).toBeInTheDocument(); expect(screen.getByTestId('cursor')).toBeInTheDocument(); }); it('renders controls correctly when not loading', () => { - render(); + render(); expect(screen.getByTestId('regenerateResponseButton')).toBeInTheDocument(); }); it('calls setComplete when StopGeneratingButton is clicked', () => { - render(); + render(); fireEvent.click(screen.getByTestId('stopGeneratingButton')); From 2e92ed86e2bb0b252fcfe7f5feae211891bc4aeb Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 10:49:36 -0500 Subject: [PATCH 39/52] type and lint fixes --- .../impl/assistant/use_conversation/index.test.tsx | 2 +- .../impl/assistant_context/index.tsx | 12 ++++++++++++ .../public/assistant/get_comments/index.tsx | 10 ++-------- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx index de83c8c107ad2..562a252bf8111 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx @@ -291,7 +291,7 @@ describe('useConversation', () => { }); }); - it.only('amendMessage updates the last message of conversation[] for a given conversationId with provided content', async () => { + it('amendMessage updates the last message of conversation[] for a given conversationId with provided content', async () => { await act(async () => { const setConversations = jest.fn(); const { result, waitForNextUpdate } = renderHook(() => useConversation(), { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index 310c975b097bf..654c5d5fd65d4 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -64,12 +64,24 @@ export interface AssistantProviderProps { docLinks: Omit; children: React.ReactNode; getComments: ({ + amendMessage, currentConversation, + isFetchingResponse, lastCommentRef, + regenerateMessage, showAnonymizedValues, }: { + amendMessage: ({ + conversationId, + content, + }: { + conversationId: string; + content: string; + }) => void; currentConversation: Conversation; + isFetchingResponse: boolean; lastCommentRef: React.MutableRefObject; + regenerateMessage: (conversationId: string) => void; showAnonymizedValues: boolean; }) => EuiCommentProps[]; http: HttpSetup; 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 42654ad4f5db1..749d17479cedd 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 @@ -6,7 +6,7 @@ */ import type { EuiCommentProps } from '@elastic/eui'; -import type { Conversation, Message } from '@kbn/elastic-assistant'; +import type { Conversation } from '@kbn/elastic-assistant'; import { EuiAvatar, EuiLoadingSpinner, tint } from '@elastic/eui'; import React from 'react'; @@ -26,13 +26,7 @@ export const getComments = ({ regenerateMessage, showAnonymizedValues, }: { - amendMessage: ({ - conversationId, - content, - }: { - conversationId: string; - content: string; - }) => Message[]; + amendMessage: ({ conversationId, content }: { conversationId: string; content: string }) => void; currentConversation: Conversation; isFetchingResponse: boolean; lastCommentRef: React.MutableRefObject; From 5cf4f063f440fd1d0b464ae73fcc7d60687a2f65 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 10:53:39 -0500 Subject: [PATCH 40/52] more fixing --- .../server/lib/langchain/helpers.ts | 8 +-- .../assistant/get_comments/index.test.tsx | 70 ++++++++++--------- 2 files changed, 42 insertions(+), 36 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/helpers.ts index c13977ddb1e7d..5f21ef9707d44 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/helpers.ts @@ -13,13 +13,13 @@ export const getLangChainMessage = ( ): BaseMessage => { switch (assistantMessage.role) { case 'system': - return new SystemMessage(assistantMessage.content); + return new SystemMessage(assistantMessage.content ?? ''); case 'user': - return new HumanMessage(assistantMessage.content); + return new HumanMessage(assistantMessage.content ?? ''); case 'assistant': - return new AIMessage(assistantMessage.content); + return new AIMessage(assistantMessage.content ?? ''); default: - return new HumanMessage(assistantMessage.content); + return new HumanMessage(assistantMessage.content ?? ''); } }; diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/index.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/index.test.tsx index 920da4272ec69..6c6713c33494f 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/index.test.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/index.test.tsx @@ -9,43 +9,49 @@ import { getComments } from '.'; import type { ConversationRole } from '@kbn/elastic-assistant/impl/assistant_context/types'; const user: ConversationRole = 'user'; +const currentConversation = { + apiConfig: {}, + id: '1', + messages: [ + { + role: user, + content: 'Hello {name}', + timestamp: '2022-01-01', + isError: false, + }, + ], +}; +const lastCommentRef = { current: null }; +const showAnonymizedValues = false; +const testProps = { + amendMessage: jest.fn(), + regenerateMessage: jest.fn(), + isFetchingResponse: false, + currentConversation, + lastCommentRef, + showAnonymizedValues, +}; describe('getComments', () => { it('Does not add error state message has no error', () => { - const currentConversation = { - apiConfig: {}, - id: '1', - messages: [ - { - role: user, - content: 'Hello {name}', - timestamp: '2022-01-01', - isError: false, - }, - ], - }; - const lastCommentRef = { current: null }; - const showAnonymizedValues = false; - - const result = getComments({ currentConversation, lastCommentRef, showAnonymizedValues }); + const result = getComments(testProps); expect(result[0].eventColor).toEqual(undefined); }); it('Adds error state when message has error', () => { - const currentConversation = { - apiConfig: {}, - id: '1', - messages: [ - { - role: user, - content: 'Hello {name}', - timestamp: '2022-01-01', - isError: true, - }, - ], - }; - const lastCommentRef = { current: null }; - const showAnonymizedValues = false; - - const result = getComments({ currentConversation, lastCommentRef, showAnonymizedValues }); + const result = getComments({ + ...testProps, + currentConversation: { + apiConfig: {}, + id: '1', + messages: [ + { + role: user, + content: 'Hello {name}', + timestamp: '2022-01-01', + isError: true, + }, + ], + }, + }); expect(result[0].eventColor).toEqual('danger'); }); }); From fcc4dc7156dee9442cc3ed656cb129d081b89306 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 12:58:08 -0500 Subject: [PATCH 41/52] transform message added to stream --- .../public/assistant/get_comments/index.tsx | 79 +++++++++++++++---- .../assistant/get_comments/stream/index.tsx | 11 ++- 2 files changed, 71 insertions(+), 19 deletions(-) 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 749d17479cedd..b047de65d8e2c 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 @@ -6,7 +6,7 @@ */ import type { EuiCommentProps } from '@elastic/eui'; -import type { Conversation } from '@kbn/elastic-assistant'; +import type { Conversation, Message } from '@kbn/elastic-assistant'; import { EuiAvatar, EuiLoadingSpinner, tint } from '@elastic/eui'; import React from 'react'; @@ -14,10 +14,42 @@ import { AssistantAvatar } from '@kbn/elastic-assistant'; import { css } from '@emotion/react'; import { euiThemeVars } from '@kbn/ui-theme'; import type { EuiPanelProps } from '@elastic/eui/src/components/panel'; +import { getMessageContentWithReplacements } from '../helpers'; import { StreamComment } from './stream'; import { CommentActions } from '../comment_actions'; import * as i18n from './translations'; +export interface ContentMessage extends Message { + content: string; +} +const transformMessageWithReplacements = ({ + message, + content, + showAnonymizedValues, + replacements, + times = 0, +}: { + message: Message; + content: string; + showAnonymizedValues: boolean; + replacements?: Record; + times?: number; +}): ContentMessage => { + console.log(`called transformMessageWithReplacements ${times}`, replacements); + if (showAnonymizedValues || !replacements) { + console.log('return without transform', { showAnonymizedValues, replacements: !replacements }); + return { ...message, content }; + } + + return { + ...message, + content: getMessageContentWithReplacements({ + messageContent: content, + replacements, + }), + }; +}; + export const getComments = ({ amendMessage, currentConversation, @@ -57,6 +89,7 @@ export const getComments = ({ content="" regenerateMessage={regenerateMessageOfConversation} isLastComment + transformMessage={() => ({ content: '' })} isFetching /> @@ -97,8 +130,25 @@ export const getComments = ({ ...(message.isError ? errorStyles : {}), }; - // message still needs to stream, no response manipulation + const transformMessage = (content: string, times?: number) => + transformMessageWithReplacements({ + message, + content, + showAnonymizedValues, + replacements, + times, + }); + + console.log('about to return', { + conditionForStreaming: !(message.content && message.content.length), + messageContent: message.content, + }); + // message still needs to stream, no actions returned and replacements handled by streamer if (!(message.content && message.content.length)) { + console.log( + 'returns this StreamComment and typeof transformMessage is ', + typeof transformMessage + ); return { ...messageProps, children: ( @@ -107,6 +157,7 @@ export const getComments = ({ amendMessage={amendMessageOfConversation} reader={message.reader} regenerateMessage={regenerateMessageOfConversation} + transformMessage={transformMessage} isLastComment={isLastComment} /> {isLastComment ? : null} @@ -114,19 +165,11 @@ export const getComments = ({ ), }; } - - const messageContentWithReplacements = - replacements != null - ? Object.keys(replacements).reduce( - (acc, replacement) => acc.replaceAll(replacement, replacements[replacement]), - message.content - ) - : message.content; - const transformedMessage = { - ...message, - content: messageContentWithReplacements, - }; - + const transformedMessage = transformMessage(message.content ?? ''); + console.log('transformMessage', { + message: message.content, + transformedMessage: transformedMessage.content, + }); return { ...messageProps, actions: , @@ -134,10 +177,12 @@ export const getComments = ({ <> {isLastComment ? : null} 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 510f81c3f8fb4..97d1e58669e33 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 @@ -7,6 +7,7 @@ import React, { useMemo } from 'react'; import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui'; +import type { ContentMessage } from '..'; import { useStream } from './use_stream'; import { StopGeneratingButton } from './buttons/stop_generating_button'; import { RegenerateResponseButton } from './buttons/regenerate_response_button'; @@ -19,6 +20,7 @@ interface Props { isLastComment: boolean; isFetching?: boolean; regenerateMessage: () => void; + transformMessage: (message: string) => ContentMessage; reader?: ReadableStreamDefaultReader; } @@ -28,6 +30,7 @@ export const StreamComment = ({ isLastComment, reader, regenerateMessage, + transformMessage, isFetching = false, }: Props) => { const { error, isLoading, isStreaming, pendingMessage, setComplete } = useStream({ @@ -35,7 +38,12 @@ export const StreamComment = ({ content, reader, }); - const message = useMemo(() => content ?? pendingMessage, [content, pendingMessage]); + + const message = useMemo( + // only transform streaming message, transform happens upstream for content message + () => content ?? transformMessage(pendingMessage).content, + [content, transformMessage, pendingMessage] + ); const isAnythingLoading = useMemo( () => isFetching || isLoading || isStreaming, [isFetching, isLoading, isStreaming] @@ -62,7 +70,6 @@ export const StreamComment = ({ ); }, [isAnythingLoading, isLastComment, reader, regenerateMessage, setComplete]); - console.log('error???', error); return ( } From 8a6c54701f6737e950bdee64a36ab5e9f00ba394 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 13:02:27 -0500 Subject: [PATCH 42/52] cleanup --- .../public/assistant/get_comments/index.tsx | 36 ++++++------------- 1 file changed, 11 insertions(+), 25 deletions(-) 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 b047de65d8e2c..b0579df971e7d 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 @@ -27,26 +27,21 @@ const transformMessageWithReplacements = ({ content, showAnonymizedValues, replacements, - times = 0, }: { message: Message; content: string; showAnonymizedValues: boolean; replacements?: Record; - times?: number; }): ContentMessage => { - console.log(`called transformMessageWithReplacements ${times}`, replacements); - if (showAnonymizedValues || !replacements) { - console.log('return without transform', { showAnonymizedValues, replacements: !replacements }); - return { ...message, content }; - } - return { ...message, - content: getMessageContentWithReplacements({ - messageContent: content, - replacements, - }), + content: + showAnonymizedValues || !replacements + ? content + : getMessageContentWithReplacements({ + messageContent: content, + replacements, + }), }; }; @@ -89,7 +84,7 @@ export const getComments = ({ content="" regenerateMessage={regenerateMessageOfConversation} isLastComment - transformMessage={() => ({ content: '' })} + transformMessage={() => ({ content: '' } as unknown as ContentMessage)} isFetching /> @@ -139,16 +134,8 @@ export const getComments = ({ times, }); - console.log('about to return', { - conditionForStreaming: !(message.content && message.content.length), - messageContent: message.content, - }); // message still needs to stream, no actions returned and replacements handled by streamer if (!(message.content && message.content.length)) { - console.log( - 'returns this StreamComment and typeof transformMessage is ', - typeof transformMessage - ); return { ...messageProps, children: ( @@ -165,11 +152,10 @@ export const getComments = ({ ), }; } + + // transform message here so we can send correct message to CommentActions const transformedMessage = transformMessage(message.content ?? ''); - console.log('transformMessage', { - message: message.content, - transformedMessage: transformedMessage.content, - }); + return { ...messageProps, actions: , From 480ef875b88bc03c997206381a1896f9c17e28bb Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 13:05:16 -0500 Subject: [PATCH 43/52] fix --- .../security_solution/public/assistant/get_comments/index.tsx | 3 +-- x-pack/plugins/security_solution/public/assistant/helpers.tsx | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) 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 b0579df971e7d..133b2e2e5c0cb 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 @@ -125,13 +125,12 @@ export const getComments = ({ ...(message.isError ? errorStyles : {}), }; - const transformMessage = (content: string, times?: number) => + const transformMessage = (content: string) => transformMessageWithReplacements({ message, content, showAnonymizedValues, replacements, - times, }); // message still needs to stream, no actions returned and replacements handled by streamer diff --git a/x-pack/plugins/security_solution/public/assistant/helpers.tsx b/x-pack/plugins/security_solution/public/assistant/helpers.tsx index 11bf5bba8b7d6..beca67f55eeb4 100644 --- a/x-pack/plugins/security_solution/public/assistant/helpers.tsx +++ b/x-pack/plugins/security_solution/public/assistant/helpers.tsx @@ -90,7 +90,7 @@ export const augmentMessageCodeBlocks = ( const cbd = currentConversation.messages.map(({ content }) => analyzeMarkdown( getMessageContentWithReplacements({ - messageContent: content ?? '', // TODO: streaming?? + messageContent: content ?? '', replacements: currentConversation.replacements, }) ) From d92c8045cda29e8ae47e16f75939bc9d9f8b2f43 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 31 Oct 2023 14:18:11 -0500 Subject: [PATCH 44/52] fix type --- .../public/assistant/get_comments/stream/index.test.tsx | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx index ef7bb7140543f..7f741ddfa9675 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx @@ -19,6 +19,7 @@ const testProps = { content, isLastComment: true, regenerateMessage: jest.fn(), + transformMessage: jest.fn(), }; const mockReader = jest.fn() as unknown as ReadableStreamDefaultReader; From 1c6d412df9b039544d51d49a4c3cae366286a4fc Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 3 Nov 2023 10:43:49 -0600 Subject: [PATCH 45/52] PR fixing --- .../kbn-elastic-assistant/impl/assistant/api.tsx | 7 ++----- .../impl/assistant/use_conversation/index.tsx | 13 ++++++++----- .../assistant/get_comments/stream/message_text.tsx | 5 ++++- .../get_comments/stream/stream_observable.ts | 2 +- .../assistant/get_comments/stream/use_stream.tsx | 4 ++-- 5 files changed, 17 insertions(+), 14 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx index 68d419d5dc789..69e6d39d85e11 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx @@ -8,9 +8,6 @@ import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/public/common'; import { HttpSetup, IHttpFetchError } from '@kbn/core-http-browser'; -// TODO: Why do i get this error here? its imported in other places without issue -// eslint-disable-next-line import/no-nodejs-modules -import { IncomingMessage } from 'http'; import type { Conversation, Message } from '../assistant_context/types'; import { API_ERROR } from './translations'; import { MODEL_GPT_3_5_TURBO } from '../connectorland/models/model_selector/model_selector'; @@ -62,7 +59,7 @@ export const fetchConnectorExecuteAction = async ({ // My "Feature Flag", turn to false before merging // In part 2 I will make enhancements to invokeAI to make it work with both openA, but to keep it to a Security Soltuion only review on this PR, // I'm calling the stream action directly - const isStream = false; + const isStream = !assistantLangChain && false; const requestBody = isStream ? { params: { @@ -81,7 +78,7 @@ export const fetchConnectorExecuteAction = async ({ try { if (isStream) { - const response = await http.fetch( + const response = await http.fetch( `/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`, { method: 'POST', diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx index f152e4ce44189..3bd9f3fcbff71 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.tsx @@ -31,7 +31,7 @@ export const DEFAULT_CONVERSATION_STATE: Conversation = { }, }; -export interface AppendMessageProps { +interface AppendMessageProps { conversationId: string; message: Message; } @@ -114,11 +114,14 @@ export const useConversation = (): UseConversation => { const prevConversation: Conversation | undefined = prev[conversationId]; if (prevConversation != null) { - const message = prevConversation.messages.pop() as unknown as Message; - const messages = [...prevConversation.messages, { ...message, content }]; + const { messages, ...rest } = prevConversation; + const message = messages[messages.length - 1]; + const updatedMessages = message + ? [...messages.slice(0, -1), { ...message, content }] + : [...messages]; const newConversation = { - ...prevConversation, - messages, + ...rest, + messages: updatedMessages, }; return { ...prev, 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 c0eeb9d53f191..c38d94a70a28c 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 @@ -15,6 +15,9 @@ import { css } from '@emotion/css'; import classNames from 'classnames'; import type { Code, InlineCode, Parent, Text } from 'mdast'; import React from 'react'; +import { transparentize } from '@elastic/eui'; +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'; @@ -44,7 +47,7 @@ const cursorCss = css` height: 16px; vertical-align: middle; display: inline-block; - background: rgba(0, 0, 0, 0.25); + background: ${transparentize(euiThemeVars.euiColorDarkShade, 0.25)}; `; const Cursor = () => ( diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts index 126cb81814d11..83f9b4cf8ead3 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts @@ -109,4 +109,4 @@ function getMessageFromChunks(chunks: Chunk[]) { return chunks.map((chunk) => chunk.choices[0]?.delta.content ?? '').join(''); } -export const getDumbObservable = () => new Observable(); +export const getPlaceholderObservable = () => new Observable(); diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx index e756ba87c8f35..148338f2afafa 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/use_stream.tsx @@ -8,7 +8,7 @@ import { useCallback, useEffect, useMemo, useState } from 'react'; import type { Subscription } from 'rxjs'; import { share } from 'rxjs'; -import { getDumbObservable, getStreamObservable } from './stream_observable'; +import { getPlaceholderObservable, getStreamObservable } from './stream_observable'; interface UseStreamProps { amendMessage: (message: string) => void; @@ -43,7 +43,7 @@ export const useStream = ({ amendMessage, content, reader }: UseStreamProps): Us () => content == null && reader != null ? getStreamObservable(reader, setLoading) - : getDumbObservable(), + : getPlaceholderObservable(), [content, reader] ); const onCompleteStream = useCallback(() => { From cae753a6a658a43da83a1f1279a36bc8b8503da6 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 3 Nov 2023 11:41:35 -0600 Subject: [PATCH 46/52] solves weird streaming issue --- .../assistant/get_comments/stream/index.tsx | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) 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 97d1e58669e33..0fe3b4b1a3ecd 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 @@ -5,7 +5,7 @@ * 2.0. */ -import React, { useMemo } from 'react'; +import React, { useEffect, useMemo, useRef } from 'react'; import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui'; import type { ContentMessage } from '..'; import { useStream } from './use_stream'; @@ -39,6 +39,22 @@ export const StreamComment = ({ reader, }); + const currentState = useRef({ isStreaming, pendingMessage, setComplete }); + + useEffect(() => { + currentState.current = { isStreaming, pendingMessage, amendMessage }; + }, [amendMessage, isStreaming, pendingMessage]); + + useEffect( + () => () => { + if (currentState.current.isStreaming && currentState.current.pendingMessage.length > 0) { + currentState.current.amendMessage(currentState.current.pendingMessage ?? ''); + } + }, + // store values in current state to detect true unmount + [] + ); + const message = useMemo( // only transform streaming message, transform happens upstream for content message () => content ?? transformMessage(pendingMessage).content, From d72bd006adfe44b1bd9b81828c1649e627f8c39e Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 3 Nov 2023 11:42:16 -0600 Subject: [PATCH 47/52] lint fix --- .../public/assistant/get_comments/stream/message_text.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 c38d94a70a28c..c1e6cd133d095 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 @@ -10,12 +10,12 @@ import { EuiText, getDefaultEuiMarkdownParsingPlugins, getDefaultEuiMarkdownProcessingPlugins, + transparentize, } from '@elastic/eui'; import { css } from '@emotion/css'; import classNames from 'classnames'; import type { Code, InlineCode, Parent, Text } from 'mdast'; import React from 'react'; -import { transparentize } from '@elastic/eui'; import { euiThemeVars } from '@kbn/ui-theme'; import type { Node } from 'unist'; From 61c6ed4af745d44c4de2b4879ac70f7805d3c893 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 3 Nov 2023 11:44:22 -0600 Subject: [PATCH 48/52] comment code --- .../public/assistant/get_comments/stream/index.tsx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 0fe3b4b1a3ecd..d978912d1599a 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 @@ -47,11 +47,12 @@ export const StreamComment = ({ useEffect( () => () => { + // if the component is unmounted while streaming, amend the message with the pending message if (currentState.current.isStreaming && currentState.current.pendingMessage.length > 0) { currentState.current.amendMessage(currentState.current.pendingMessage ?? ''); } }, - // store values in current state to detect true unmount + // store values in currentState to detect true unmount [] ); From 301aa5151af9c625252f083117dda26199815dc6 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 3 Nov 2023 13:15:24 -0600 Subject: [PATCH 49/52] fix type --- .../public/assistant/get_comments/stream/index.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 d978912d1599a..cc630882eba86 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 @@ -39,7 +39,7 @@ export const StreamComment = ({ reader, }); - const currentState = useRef({ isStreaming, pendingMessage, setComplete }); + const currentState = useRef({ isStreaming, pendingMessage, amendMessage }); useEffect(() => { currentState.current = { isStreaming, pendingMessage, amendMessage }; From a9202586644404ebbda30fb76febac6cfec0323c Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 3 Nov 2023 13:18:59 -0600 Subject: [PATCH 50/52] update w main --- .../public/assistant/get_comments/index.tsx | 19 ++----------------- 1 file changed, 2 insertions(+), 17 deletions(-) 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 133b2e2e5c0cb..3b80e64f0f6f6 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 @@ -7,13 +7,10 @@ import type { EuiCommentProps } from '@elastic/eui'; import type { Conversation, Message } from '@kbn/elastic-assistant'; -import { EuiAvatar, EuiLoadingSpinner, tint } from '@elastic/eui'; +import { EuiAvatar, EuiLoadingSpinner } from '@elastic/eui'; import React from 'react'; import { AssistantAvatar } from '@kbn/elastic-assistant'; -import { css } from '@emotion/react'; -import { euiThemeVars } from '@kbn/ui-theme'; -import type { EuiPanelProps } from '@elastic/eui/src/components/panel'; import { getMessageContentWithReplacements } from '../helpers'; import { StreamComment } from './stream'; import { CommentActions } from '../comment_actions'; @@ -99,18 +96,6 @@ export const getComments = ({ const isLastComment = index === currentConversation.messages.length - 1; const isUser = message.role === 'user'; const replacements = currentConversation.replacements; - const errorStyles = { - eventColor: 'danger' as EuiPanelProps['color'], - css: css` - .euiCommentEvent { - border: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; - } - .euiCommentEvent__header { - padding: 0 !important; - border-block-end: 1px solid ${tint(euiThemeVars.euiColorDanger, 0.75)}; - } - `, - }; const messageProps = { timelineAvatar: isUser ? ( @@ -122,7 +107,7 @@ export const getComments = ({ message.timestamp.length === 0 ? new Date().toLocaleString() : message.timestamp ), username: isUser ? i18n.YOU : i18n.ASSISTANT, - ...(message.isError ? errorStyles : {}), + eventColor: message.isError ? 'danger' : undefined, }; const transformMessage = (content: string) => From 40e2b79c5313cb861d1dc161dcc4eee47510724d Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 6 Nov 2023 12:41:25 -0700 Subject: [PATCH 51/52] fixed --- .../kbn-elastic-assistant/impl/assistant/index.tsx | 5 ++++- .../public/assistant/get_comments/index.tsx | 8 ++++++-- .../assistant/get_comments/stream/index.test.tsx | 1 + .../public/assistant/get_comments/stream/index.tsx | 10 ++++++---- .../assistant/get_comments/stream/message_text.tsx | 5 ++++- 5 files changed, 21 insertions(+), 8 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 ac83710477aa4..b5a38ef96036f 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -189,7 +189,10 @@ const AssistantComponent: React.FC = ({ const [messageCodeBlocks, setMessageCodeBlocks] = useState(); const [_, setCodeBlockControlsVisible] = useState(false); useLayoutEffect(() => { - setMessageCodeBlocks(augmentMessageCodeBlocks(currentConversation)); + // need in order for code block controls to be added to the DOM + setTimeout(() => { + setMessageCodeBlocks(augmentMessageCodeBlocks(currentConversation)); + }, 0); }, [augmentMessageCodeBlocks, currentConversation]); const isSendingDisabled = useMemo(() => { 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 3b80e64f0f6f6..1ed6512309a4b 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 @@ -83,6 +83,8 @@ export const getComments = ({ isLastComment transformMessage={() => ({ content: '' } as unknown as ContentMessage)} isFetching + // we never need to append to a code block in the loading comment, which is what this index is used for + index={999} /> @@ -126,10 +128,11 @@ export const getComments = ({ <> {isLastComment ? : null} @@ -148,10 +151,11 @@ export const getComments = ({ {isLastComment ? : null} diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx index 7f741ddfa9675..7813e45829d1c 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx @@ -17,6 +17,7 @@ const content = 'Test Content'; const testProps = { amendMessage: jest.fn(), content, + index: 1, isLastComment: true, regenerateMessage: jest.fn(), transformMessage: jest.fn(), 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 cc630882eba86..db394f39bfa32 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,21 +17,23 @@ import { MessageText } from './message_text'; interface Props { amendMessage: (message: string) => void; content?: string; - isLastComment: boolean; isFetching?: boolean; + isLastComment: boolean; + index: number; + reader?: ReadableStreamDefaultReader; regenerateMessage: () => void; transformMessage: (message: string) => ContentMessage; - reader?: ReadableStreamDefaultReader; } export const StreamComment = ({ amendMessage, content, + index, + isFetching = false, isLastComment, reader, regenerateMessage, transformMessage, - isFetching = false, }: Props) => { const { error, isLoading, isStreaming, pendingMessage, setComplete } = useStream({ amendMessage, @@ -89,7 +91,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 c1e6cd133d095..415800a04609d 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,6 +24,7 @@ import { CustomCodeBlock } from '../custom_codeblock/custom_code_block'; interface Props { content: string; + index: number; loading: boolean; } @@ -153,7 +154,7 @@ const getPluginDependencies = () => { }; }; -export function MessageText({ loading, content }: Props) { +export function MessageText({ loading, content, index }: Props) { const containerClassName = css` overflow-wrap: break-word; `; @@ -163,6 +164,8 @@ export function MessageText({ loading, content }: Props) { return ( Date: Mon, 6 Nov 2023 15:10:03 -0700 Subject: [PATCH 52/52] fix scrolling --- .../impl/assistant/index.tsx | 89 ++++++++++--------- .../impl/assistant_context/index.tsx | 4 - .../assistant/get_comments/index.test.tsx | 2 - .../public/assistant/get_comments/index.tsx | 67 ++++++-------- .../timeline/tabs_content/index.tsx | 1 + 5 files changed, 74 insertions(+), 89 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 b5a38ef96036f..86e0f3a460055 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -52,6 +52,7 @@ import { ConnectorMissingCallout } from '../connectorland/connector_missing_call export interface Props { conversationId?: string; + embeddedLayout?: boolean; promptContextId?: string; shouldRefocusPrompt?: boolean; showTitle?: boolean; @@ -64,6 +65,7 @@ export interface Props { */ const AssistantComponent: React.FC = ({ conversationId, + embeddedLayout = false, promptContextId = '', shouldRefocusPrompt = false, showTitle = true, @@ -167,17 +169,11 @@ const AssistantComponent: React.FC = ({ const { comments: connectorComments, prompt: connectorPrompt } = useConnectorSetup({ conversation: blockBotConversation, - onSetupComplete: () => { - bottomRef.current?.scrollIntoView({ behavior: 'auto' }); - }, }); const currentTitle: string | JSX.Element = isWelcomeSetup && blockBotConversation.theme?.title ? blockBotConversation.theme?.title : title; - const bottomRef = useRef(null); - const lastCommentRef = useRef(null); - const [promptTextPreview, setPromptTextPreview] = useState(''); const [autoPopulatedOnce, setAutoPopulatedOnce] = useState(false); const [userPrompt, setUserPrompt] = useState(null); @@ -217,17 +213,22 @@ const AssistantComponent: React.FC = ({ }, []); // End drill in `Add To Timeline` action - // Scroll to bottom on conversation change - useEffect(() => { - bottomRef.current?.scrollIntoView({ behavior: 'auto' }); - }, []); + // Start Scrolling + const commentsContainerRef = useRef(null); + useEffect(() => { - setTimeout(() => { - bottomRef.current?.scrollIntoView({ behavior: 'auto' }); - }, 0); - }, [currentConversation.messages.length, selectedPromptContextsCount]); - //// - // + const parent = commentsContainerRef.current?.parentElement; + if (!parent) { + return; + } + // when scrollHeight changes, parent is scrolled to bottom + parent.scrollTop = parent.scrollHeight; + }); + + const getWrapper = (children: React.ReactNode, isCommentContainer: boolean) => + isCommentContainer ? {children} : <>{children}; + + // End Scrolling const selectedSystemPrompt = useMemo( () => getDefaultSystemPrompt({ allSystemPrompts, conversation: currentConversation }), @@ -368,7 +369,6 @@ const AssistantComponent: React.FC = ({ = ({ setSelectedPromptContexts={setSelectedPromptContexts} /> )} - -
), [ @@ -421,15 +419,12 @@ const AssistantComponent: React.FC = ({ const comments = useMemo(() => { if (isDisabled) { return ( - <> - - - + ); } @@ -446,7 +441,7 @@ const AssistantComponent: React.FC = ({ [assistantTelemetry, selectedConversationId] ); - return ( + return getWrapper( <> = ({ )} - {comments} - - {!isDisabled && showMissingConnectorCallout && areConnectorsFetched && ( + {getWrapper( <> - - - - 0} - isSettingsModalVisible={isSettingsModalVisible} - setIsSettingsModalVisible={setIsSettingsModalVisible} - /> - - - + {comments} + + {!isDisabled && showMissingConnectorCallout && areConnectorsFetched && ( + <> + + + + 0} + isSettingsModalVisible={isSettingsModalVisible} + setIsSettingsModalVisible={setIsSettingsModalVisible} + /> + + + + )} + , + !embeddedLayout )} = ({ /> )} - + , + embeddedLayout ); }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index 654c5d5fd65d4..dd508b54ce906 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -67,7 +67,6 @@ export interface AssistantProviderProps { amendMessage, currentConversation, isFetchingResponse, - lastCommentRef, regenerateMessage, showAnonymizedValues, }: { @@ -80,7 +79,6 @@ export interface AssistantProviderProps { }) => void; currentConversation: Conversation; isFetchingResponse: boolean; - lastCommentRef: React.MutableRefObject; regenerateMessage: (conversationId: string) => void; showAnonymizedValues: boolean; }) => EuiCommentProps[]; @@ -114,14 +112,12 @@ export interface UseAssistantContext { conversations: Record; getComments: ({ currentConversation, - lastCommentRef, showAnonymizedValues, amendMessage, isFetchingResponse, }: { currentConversation: Conversation; isFetchingResponse: boolean; - lastCommentRef: React.MutableRefObject; amendMessage: ({ conversationId, content, diff --git a/x-pack/plugins/security_solution/public/assistant/get_comments/index.test.tsx b/x-pack/plugins/security_solution/public/assistant/get_comments/index.test.tsx index 6c6713c33494f..8e6c9f20f7395 100644 --- a/x-pack/plugins/security_solution/public/assistant/get_comments/index.test.tsx +++ b/x-pack/plugins/security_solution/public/assistant/get_comments/index.test.tsx @@ -21,14 +21,12 @@ const currentConversation = { }, ], }; -const lastCommentRef = { current: null }; const showAnonymizedValues = false; const testProps = { amendMessage: jest.fn(), regenerateMessage: jest.fn(), isFetchingResponse: false, currentConversation, - lastCommentRef, showAnonymizedValues, }; describe('getComments', () => { 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 1ed6512309a4b..9c547b3033112 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 @@ -46,14 +46,12 @@ export const getComments = ({ amendMessage, currentConversation, isFetchingResponse, - lastCommentRef, regenerateMessage, showAnonymizedValues, }: { amendMessage: ({ conversationId, content }: { conversationId: string; content: string }) => void; currentConversation: Conversation; isFetchingResponse: boolean; - lastCommentRef: React.MutableRefObject; regenerateMessage: (conversationId: string) => void; showAnonymizedValues: boolean; }): EuiCommentProps[] => { @@ -75,19 +73,16 @@ export const getComments = ({ timelineAvatar: , timestamp: '...', children: ( - <> - ({ content: '' } as unknown as ContentMessage)} - isFetching - // we never need to append to a code block in the loading comment, which is what this index is used for - index={999} - /> - - + ({ content: '' } as unknown as ContentMessage)} + isFetching + // we never need to append to a code block in the loading comment, which is what this index is used for + index={999} + /> ), }, ] @@ -125,17 +120,14 @@ export const getComments = ({ return { ...messageProps, children: ( - <> - - {isLastComment ? : null} - + ), }; } @@ -147,19 +139,16 @@ export const getComments = ({ ...messageProps, actions: , children: ( - <> - - {isLastComment ? : null} - + ), }; }), diff --git a/x-pack/plugins/security_solution/public/timelines/components/timeline/tabs_content/index.tsx b/x-pack/plugins/security_solution/public/timelines/components/timeline/tabs_content/index.tsx index 04580b1546b4c..b62e36e20b938 100644 --- a/x-pack/plugins/security_solution/public/timelines/components/timeline/tabs_content/index.tsx +++ b/x-pack/plugins/security_solution/public/timelines/components/timeline/tabs_content/index.tsx @@ -108,6 +108,7 @@ const AssistantTab: React.FC<{
+ + + {children} -
-
- - {children} - -
-
+
+ + {children} + +
+