From c2ce299b87a7fd134f9c3f83dfe6430a039ad65d Mon Sep 17 00:00:00 2001 From: Jason Botzas-Coluni <44372106+jaybcee@users.noreply.github.com> Date: Fri, 19 Dec 2025 10:39:35 -0500 Subject: [PATCH 1/5] Add refusal field to assistant conversations (#243423) (cherry picked from commit fe9c9fd75355732baf10cda2b70fe4bb9f36652b) # Conflicts: # x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/events.ts # x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/messages.ts # x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts # x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts # x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts # x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts # x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts --- oas_docs/output/kibana.serverless.yaml | 3 ++ oas_docs/output/kibana.yaml | 3 ++ .../inference-common/src/chat_complete/api.ts | 4 ++ .../src/chat_complete/events.ts | 8 ++++ .../src/chat_complete/messages.ts | 4 ++ .../src/chat_model/from_inference/chunks.ts | 4 +- .../src/chat_model/from_inference/messages.ts | 2 + ...sistant_api_2023_10_31.bundled.schema.yaml | 3 ++ ...sistant_api_2023_10_31.bundled.schema.yaml | 3 ++ .../conversations/common_attributes.gen.ts | 4 ++ .../common_attributes.schema.yaml | 3 ++ .../adapters/openai/from_openai.ts | 1 + .../utils/chunks_into_message.ts | 3 +- .../chat_complete/utils/merge_chunks.ts | 4 ++ .../chat_complete/utils/stream_to_response.ts | 1 + .../connector_types/inference/helpers.ts | 4 ++ .../append_conversation_messages.test.ts | 45 +++++++++++++++++++ .../append_conversation_messages.ts | 1 + .../conversations/create_conversation.ts | 1 + .../conversations/field_maps_configuration.ts | 5 +++ .../conversations/transforms.ts | 1 + .../conversations/types.ts | 2 + .../conversations/update_conversation.ts | 2 + .../server/lib/langchain/executors/types.ts | 3 +- .../graphs/default_assistant_graph/helpers.ts | 17 +++++-- .../server/routes/chat/chat_complete_route.ts | 4 +- .../server/routes/helpers.ts | 6 +++ .../routes/post_actions_connector_execute.ts | 4 +- 28 files changed, 136 insertions(+), 9 deletions(-) diff --git a/oas_docs/output/kibana.serverless.yaml b/oas_docs/output/kibana.serverless.yaml index 9ffca4433824c..3af7fe70c75a4 100644 --- a/oas_docs/output/kibana.serverless.yaml +++ b/oas_docs/output/kibana.serverless.yaml @@ -58784,6 +58784,9 @@ components: reader: $ref: '#/components/schemas/Security_AI_Assistant_API_Reader' description: Message content. + refusal: + description: Refusal reason returned by the model when content is filtered. + type: string role: $ref: '#/components/schemas/Security_AI_Assistant_API_MessageRole' description: Message role. diff --git a/oas_docs/output/kibana.yaml b/oas_docs/output/kibana.yaml index 7bcb6c6dc5a89..8bf51a45d0af3 100644 --- a/oas_docs/output/kibana.yaml +++ b/oas_docs/output/kibana.yaml @@ -68212,6 +68212,9 @@ components: reader: $ref: '#/components/schemas/Security_AI_Assistant_API_Reader' description: Message content. + refusal: + description: Refusal reason returned by the model when content is filtered. + type: string role: $ref: '#/components/schemas/Security_AI_Assistant_API_MessageRole' description: Message role. diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts index d6d6b308f3cb2..31f2291b35106 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts @@ -196,6 +196,10 @@ export interface ChatCompleteResponse< * The text content of the LLM response. */ content: string; + /** + * Optional refusal reason returned by the model when content is filtered. + */ + refusal?: string; /** * The eventual tool calls performed by the LLM. */ diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/events.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/events.ts index e5308bb566f7f..bbaa9dde9a7d6 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/events.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/events.ts @@ -35,6 +35,10 @@ export type ChatCompletionMessageEvent['toolCalls']; + /** + * Optional refusal reason returned by the model when content is filtered. + */ + refusal?: string; /** * Optional deanonymized input messages metadata */ @@ -83,6 +87,10 @@ export type ChatCompletionChunkEvent = InferenceTaskEventBase< * The content chunk */ content: string; + /** + * Optional refusal reason chunk. + */ + refusal?: string; /** * The tool call chunks */ diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/messages.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/messages.ts index 7fdd78670fa1b..95f6e5aa554cc 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/messages.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/messages.ts @@ -59,6 +59,10 @@ export type AssistantMessage = MessageBase & { * Note that LLM with parallel tool invocation can potentially call multiple tools at the same time. */ toolCalls?: ToolCall[]; + /** + * Optional refusal reason returned by the model when content is filtered. + */ + refusal?: string | null; }; /** diff --git a/x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/chunks.ts b/x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/chunks.ts index a64e59c99f876..5cc6e6a54e92a 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/chunks.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/chunks.ts @@ -23,10 +23,12 @@ export const completionChunkToLangchain = (chunk: ChatCompletionChunkEvent): AIM }; }); + const additionalKwargs = chunk.refusal ? { refusal: chunk.refusal } : {}; + return new AIMessageChunk({ content: chunk.content, tool_call_chunks: toolCallChunks, - additional_kwargs: {}, + additional_kwargs: additionalKwargs, response_metadata: {}, }); }; diff --git a/x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/messages.ts b/x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/messages.ts index 6e4388e6d9b4a..ac0da8f210a9e 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/messages.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/messages.ts @@ -9,8 +9,10 @@ import type { ChatCompleteResponse } from '@kbn/inference-common'; import { AIMessage } from '@langchain/core/messages'; export const responseToLangchainMessage = (response: ChatCompleteResponse): AIMessage => { + const additionalKwargs = response.refusal ? { refusal: response.refusal } : undefined; return new AIMessage({ content: response.content, + ...(additionalKwargs ? { additional_kwargs: additionalKwargs } : {}), tool_calls: response.toolCalls.map((toolCall) => { return { id: toolCall.toolCallId, diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml index 7aaaf4b064a33..27b886c1dabd8 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml @@ -2739,6 +2739,9 @@ components: reader: $ref: '#/components/schemas/Reader' description: Message content. + refusal: + description: Refusal reason returned by the model when content is filtered. + type: string role: $ref: '#/components/schemas/MessageRole' description: Message role. diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml index d0ba85456497c..d7f45cd52c987 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml @@ -2739,6 +2739,9 @@ components: reader: $ref: '#/components/schemas/Reader' description: Message content. + refusal: + description: Refusal reason returned by the model when content is filtered. + type: string role: $ref: '#/components/schemas/MessageRole' description: Message role. diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts index 4d1c14e6d45f4..512078acc5e97 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts @@ -248,6 +248,10 @@ export const Message = z.object({ * Message content. */ content: z.string(), + /** + * Refusal reason returned by the model when content is filtered. + */ + refusal: z.string().optional(), /** * Message content. */ diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml index 48d8cdbe892d2..3e30a3e7236bf 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml @@ -257,6 +257,9 @@ components: type: string description: Message content. example: 'Hello, how can I assist you today?' + refusal: + type: string + description: Refusal reason returned by the model when content is filtered. reader: $ref: '#/components/schemas/Reader' description: Message content. diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/from_openai.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/from_openai.ts index a40bd40a9b616..a0a2c29419f45 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/from_openai.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/from_openai.ts @@ -18,6 +18,7 @@ export function chunkFromOpenAI(chunk: OpenAI.ChatCompletionChunk): ChatCompleti return { type: ChatCompletionEventType.ChatCompletionChunk, content: delta.content ?? '', + refusal: delta.refusal ?? undefined, tool_calls: delta.tool_calls?.map((toolCall) => { return { diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/chunks_into_message.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/chunks_into_message.ts index 41c6e60bef31a..2be6154d81459 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/chunks_into_message.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/chunks_into_message.ts @@ -45,7 +45,7 @@ export function chunksIntoMessage({ logger.debug(() => `Received completed message: ${JSON.stringify(concatenatedChunk)}`); - const { content, tool_calls: toolCalls } = concatenatedChunk; + const { content, refusal, tool_calls: toolCalls } = concatenatedChunk; const activeSpan = trace.getActiveSpan(); if (activeSpan) { setChoice(activeSpan, { content, toolCalls }); @@ -56,6 +56,7 @@ export function chunksIntoMessage({ return { type: ChatCompletionEventType.ChatCompletionMessage, content, + ...(refusal ? { refusal } : {}), toolCalls: validatedToolCalls, }; }) diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/merge_chunks.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/merge_chunks.ts index 31d4a0ad277ca..c8c32e7028168 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/merge_chunks.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/merge_chunks.ts @@ -9,6 +9,7 @@ import { ChatCompletionChunkEvent, UnvalidatedToolCall } from '@kbn/inference-co interface UnvalidatedMessage { content: string; + refusal?: string; tool_calls: UnvalidatedToolCall[]; } @@ -19,6 +20,9 @@ export const mergeChunks = (chunks: ChatCompletionChunkEvent[]): UnvalidatedMess const message = chunks.reduce( (prev, chunk) => { prev.content += chunk.content ?? ''; + if (chunk.refusal) { + prev.refusal = chunk.refusal; + } chunk.tool_calls?.forEach((toolCall) => { let prevToolCall = prev.tool_calls[toolCall.index]; diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/stream_to_response.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/stream_to_response.ts index b831a5da6b304..b018a26207136 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/stream_to_response.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/utils/stream_to_response.ts @@ -33,6 +33,7 @@ export const streamToResponse = return { content: messageEvent.content, + refusal: messageEvent.refusal, toolCalls: messageEvent.toolCalls, tokens: tokenEvent?.tokens, deanonymized_input: messageEvent.deanonymized_input, diff --git a/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/helpers.ts b/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/helpers.ts index 05faf03a6ff3f..74728cddf0dd1 100644 --- a/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/helpers.ts +++ b/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/helpers.ts @@ -47,6 +47,9 @@ export function chunksIntoMessage(obs$: Observable) (prev, chunk) => { if (chunk.choices.length > 0 && !chunk.usage) { prev.choices[0].message.content += chunk.choices[0].message.content ?? ''; + if (chunk.choices[0].message.refusal) { + prev.choices[0].message.refusal = chunk.choices[0].message.refusal; + } chunk.choices[0].message.tool_calls?.forEach((toolCall) => { if (toolCall.index !== undefined) { @@ -89,6 +92,7 @@ export function chunksIntoMessage(obs$: Observable) { message: { content: '', + refusal: null, role: 'assistant', }, }, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts index ec7d8e6a91e20..98ef04f2b7339 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts @@ -323,6 +323,51 @@ describe('appendConversationMessages', () => { }) ); }); + + it('preserves refusal reason when present on messages', async () => { + const messageWithRefusal = createMockMessage({ + refusal: 'Detected harmful input content: INSULTS', + }); + setupSuccessfulTest(); + + await callAppendConversationMessages([messageWithRefusal]); + + expect(dataWriter.bulk).toHaveBeenCalledWith( + expect.objectContaining({ + documentsToUpdate: expect.arrayContaining([ + expect.objectContaining({ + messages: expect.arrayContaining([ + expect.objectContaining({ + refusal: 'Detected harmful input content: INSULTS', + }), + ]), + }), + ]), + }) + ); + }); + + it('generates UUID for messages without id', async () => { + const messageWithoutId = createMockMessage({ id: undefined }); + setupSuccessfulTest(); + + await callAppendConversationMessages([messageWithoutId]); + + expect(dataWriter.bulk).toHaveBeenCalledWith( + expect.objectContaining({ + documentsToUpdate: expect.arrayContaining([ + expect.objectContaining({ + messages: expect.arrayContaining([ + expect.objectContaining({ + id: expect.any(String), + content: 'test content', + }), + ]), + }), + ]), + }) + ); + }); }); describe('transformToUpdateScheme', () => { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.ts index f5c096665b322..4695e0fb5f782 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.ts @@ -104,6 +104,7 @@ export const transformToUpdateScheme = ( messages: messages?.map((message) => ({ '@timestamp': message.timestamp, content: message.content, + ...(message.refusal ? { refusal: message.refusal } : {}), is_error: message.isError, reader: message.reader, role: message.role, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.ts index ee706cf61fea4..0d8f62eb8719c 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.ts @@ -94,6 +94,7 @@ export const transformToCreateScheme = ( messages: messages?.map((message) => ({ '@timestamp': message.timestamp, content: message.content, + ...(message.refusal ? { refusal: message.refusal } : {}), is_error: message.isError, reader: message.reader, role: message.role, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/field_maps_configuration.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/field_maps_configuration.ts index acc8bde1a2339..6fc9e085ba3b6 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/field_maps_configuration.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/field_maps_configuration.ts @@ -72,6 +72,11 @@ export const conversationsFieldMap: FieldMap = { array: false, required: false, }, + 'messages.refusal': { + type: 'text', + array: false, + required: false, + }, 'messages.reader': { type: 'object', array: false, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts index 3fba0b0e07357..8f8a9fa71ecc9 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts @@ -50,6 +50,7 @@ export const transformESToConversation = ( messageContent: message.content, replacements, }), + ...(message.refusal ? { refusal: message.refusal } : {}), ...(message.is_error ? { isError: message.is_error } : {}), ...(message.reader ? { reader: message.reader } : {}), ...(message.user ? { user: message.user } : {}), diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts index 5bee0fcb2a5f4..5d57fdf897e4f 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts @@ -33,6 +33,7 @@ export interface EsConversationSchema { messages?: Array<{ '@timestamp': string; content: string; + refusal?: string; reader?: Reader; role: MessageRole; is_error?: boolean; @@ -70,6 +71,7 @@ export interface CreateMessageSchema { messages?: Array<{ '@timestamp': string; content: string; + refusal?: string; reader?: Reader; role: MessageRole; is_error?: boolean; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts index 2dd40b8429b3a..483ab838e2858 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts @@ -28,6 +28,7 @@ export interface UpdateConversationSchema { messages?: Array<{ '@timestamp': string; content: string; + refusal?: string; reader?: Reader; role: MessageRole; is_error?: boolean; @@ -133,6 +134,7 @@ export const transformToUpdateScheme = ( messages: messages.map((message) => ({ '@timestamp': message.timestamp, content: message.content, + ...(message.refusal ? { refusal: message.refusal } : {}), is_error: message.isError, reader: message.reader, role: message.role, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index 02adb6753f15c..1e9ca9a19fa58 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -35,7 +35,8 @@ import { AIAssistantDataClient } from '../../../ai_assistant_data_clients'; export type OnLlmResponse = ( content: string, traceData?: Message['traceData'], - isError?: boolean + isError?: boolean, + refusal?: string ) => Promise; export interface AssistantDataClients { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts index 08ebde6746e63..4e244f0503556 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts @@ -71,7 +71,7 @@ export const streamGraph = async ({ } = streamFactory<{ type: string; payload: string }>(request.headers, logger, false, false); let didEnd = false; - const handleStreamEnd = (finalResponse: string, isError = false) => { + const handleStreamEnd = (finalResponse: string, isError = false, refusal?: string) => { if (didEnd) { return; } @@ -92,7 +92,8 @@ export const streamGraph = async ({ transactionId: streamingSpan?.transaction?.ids?.['transaction.id'], traceId: streamingSpan?.ids?.['trace.id'], }, - isError + isError, + refusal ).catch(() => {}); } streamEnd(); @@ -129,7 +130,11 @@ export const streamGraph = async ({ !data.output.lc_kwargs?.tool_calls?.length && !didEnd ) { - handleStreamEnd(data.output.content); + const refusal = + typeof data.output?.additional_kwargs?.refusal === 'string' + ? (data.output.additional_kwargs.refusal as string) + : undefined; + handleStreamEnd(data.output.content, false, refusal); } else if ( // This is the end of one model invocation but more message will follow as there are tool calls. If this chunk contains text content, add a newline separator to the stream to visually separate the chunks. event === 'on_chat_model_end' && @@ -206,8 +211,12 @@ export const invokeGraph = async ({ const lastMessage = result.messages[result.messages.length - 1]; const output = lastMessage.text; const conversationId = result.conversationId; + const refusal = + typeof lastMessage?.additional_kwargs?.refusal === 'string' + ? (lastMessage.additional_kwargs.refusal as string) + : undefined; if (onLlmResponse) { - await onLlmResponse(output, traceData); + await onLlmResponse(output, traceData, false, refusal); } return { output, traceData, conversationId }; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts index 09dd6bccebb3d..0b28879388898 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts @@ -212,7 +212,8 @@ export const chatCompleteRoute = ( const onLlmResponse = async ( content: string, traceData: Message['traceData'] = {}, - isError = false + isError = false, + refusal ): Promise => { if (conversationId && conversationsDataClient) { const { prunedContent, prunedContentReferencesStore } = pruneContentReferences( @@ -224,6 +225,7 @@ export const chatCompleteRoute = ( conversationId, conversationsDataClient, messageContent: prunedContent, + messageRefusal: refusal, replacements: latestReplacements, isError, traceData, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts index 99e2f3d8bb3b6..0395196273360 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts @@ -105,11 +105,13 @@ export const getPluginNameFromRequest = ({ export const getMessageFromRawResponse = ({ rawContent, metadata, + refusal, isError, traceData, }: { rawContent?: string; metadata?: MessageMetadata; + refusal?: string; traceData?: TraceData; isError?: boolean; }): Message => { @@ -118,6 +120,7 @@ export const getMessageFromRawResponse = ({ return { role: 'assistant', content: rawContent, + ...(refusal ? { refusal } : {}), timestamp: dateTimeString, metadata, isError, @@ -175,6 +178,7 @@ export const getSystemPromptFromUserConversation = async ({ export interface AppendAssistantMessageToConversationParams { conversationsDataClient: AIAssistantConversationsDataClient; messageContent: string; + messageRefusal?: string; replacements: Replacements; conversationId: string; contentReferences: ContentReferences; @@ -184,6 +188,7 @@ export interface AppendAssistantMessageToConversationParams { export const appendAssistantMessageToConversation = async ({ conversationsDataClient, messageContent, + messageRefusal, replacements, conversationId, contentReferences, @@ -207,6 +212,7 @@ export const appendAssistantMessageToConversation = async ({ messageContent, replacements, }), + refusal: messageRefusal, metadata: !isEmpty(metadata) ? metadata : undefined, traceData, isError, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index e1c872ad1587a..89b0afc1b3b95 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -142,7 +142,8 @@ export const postActionsConnectorExecuteRoute = ( onLlmResponse = async ( content: string, traceData: Message['traceData'] = {}, - isError = false + isError = false, + refusal ): Promise => { if (conversationsDataClient && conversationId) { const { prunedContent, prunedContentReferencesStore } = pruneContentReferences( @@ -154,6 +155,7 @@ export const postActionsConnectorExecuteRoute = ( conversationId, conversationsDataClient, messageContent: prunedContent, + messageRefusal: refusal, replacements: latestReplacements, isError, traceData, From bcdb76d258a9677738ad7098688f81cc3a28aaaf Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 19 Dec 2025 10:11:50 -0700 Subject: [PATCH 2/5] rm mistakenly committed test --- .../append_conversation_messages.test.ts | 22 ------------------- 1 file changed, 22 deletions(-) diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts index 98ef04f2b7339..0a24b09b68130 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts @@ -346,28 +346,6 @@ describe('appendConversationMessages', () => { }) ); }); - - it('generates UUID for messages without id', async () => { - const messageWithoutId = createMockMessage({ id: undefined }); - setupSuccessfulTest(); - - await callAppendConversationMessages([messageWithoutId]); - - expect(dataWriter.bulk).toHaveBeenCalledWith( - expect.objectContaining({ - documentsToUpdate: expect.arrayContaining([ - expect.objectContaining({ - messages: expect.arrayContaining([ - expect.objectContaining({ - id: expect.any(String), - content: 'test content', - }), - ]), - }), - ]), - }) - ); - }); }); describe('transformToUpdateScheme', () => { From 299a99cf34c37f5d5546f65eb981c502236d9a67 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 19 Dec 2025 10:15:04 -0700 Subject: [PATCH 3/5] type fixing --- .../server/routes/chat/chat_complete_route.ts | 2 +- .../plugins/elastic_assistant/server/routes/helpers.ts | 3 ++- .../server/routes/post_actions_connector_execute.ts | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts index 0b28879388898..db7843ef7f6e2 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts @@ -213,7 +213,7 @@ export const chatCompleteRoute = ( content: string, traceData: Message['traceData'] = {}, isError = false, - refusal + refusal?: string ): Promise => { if (conversationId && conversationsDataClient) { const { prunedContent, prunedContentReferencesStore } = pruneContentReferences( diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts index 0395196273360..fa8c999d1282c 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts @@ -254,7 +254,8 @@ export interface LangChainExecuteParams { onLlmResponse?: ( content: string, traceData?: Message['traceData'], - isError?: boolean + isError?: boolean, + refusal?: string, ) => Promise; response: KibanaResponseFactory; responseLanguage?: string; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 89b0afc1b3b95..97f7f85c39ed0 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -143,7 +143,7 @@ export const postActionsConnectorExecuteRoute = ( content: string, traceData: Message['traceData'] = {}, isError = false, - refusal + refusal?: string ): Promise => { if (conversationsDataClient && conversationId) { const { prunedContent, prunedContentReferencesStore } = pruneContentReferences( From cab2830c9df147e93beadee44a8af0cc8b55faa3 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 19 Dec 2025 10:21:43 -0700 Subject: [PATCH 4/5] small fix --- .../security/plugins/elastic_assistant/server/routes/helpers.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts index fa8c999d1282c..ec34a5872bc9b 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts @@ -255,7 +255,7 @@ export interface LangChainExecuteParams { content: string, traceData?: Message['traceData'], isError?: boolean, - refusal?: string, + refusal?: string ) => Promise; response: KibanaResponseFactory; responseLanguage?: string; From d24f8859a1f04ccef931bd8e4b3276c4e5b5c1e9 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Mon, 22 Dec 2025 11:29:40 -0700 Subject: [PATCH 5/5] fix tests --- .../graphs/default_assistant_graph/helpers.test.ts | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.test.ts index 559de20d3c02d..603ede3b24361 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.test.ts @@ -108,7 +108,8 @@ describe('streamGraph', () => { expect(mockOnLlmResponse).toHaveBeenCalledWith( 'final message', { transactionId: 'transactionId', traceId: 'traceId' }, - false + false, + undefined ); }); }); @@ -177,7 +178,8 @@ describe('streamGraph', () => { expect(mockOnLlmResponse).toHaveBeenCalledWith( 'content', { transactionId: 'transactionId', traceId: 'traceId' }, - false + false, + undefined ); }); }); @@ -239,7 +241,8 @@ describe('streamGraph', () => { expect(mockOnLlmResponse).toHaveBeenCalledWith( 'Look at these rare IP addresses.', { transactionId: 'transactionId', traceId: 'traceId' }, - false + false, + undefined ); }); };