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..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 @@ -323,6 +323,29 @@ 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', + }), + ]), + }), + ]), + }) + ); + }); }); 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.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 ); }); }; 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..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 @@ -212,7 +212,8 @@ export const chatCompleteRoute = ( const onLlmResponse = async ( content: string, traceData: Message['traceData'] = {}, - isError = false + isError = false, + refusal?: string ): 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..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 @@ -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, @@ -248,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 e1c872ad1587a..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 @@ -142,7 +142,8 @@ export const postActionsConnectorExecuteRoute = ( onLlmResponse = async ( content: string, traceData: Message['traceData'] = {}, - isError = false + isError = false, + refusal?: string ): 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,