Add refusal field to assistant conversations

oas_docs/output/kibana.serverless.yaml

@@ -83927,6 +83927,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.

oas_docs/output/kibana.yaml

@@ -94445,6 +94445,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.

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts

@@ -197,6 +197,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.
    */

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/events.ts

@@ -32,6 +32,10 @@ export type ChatCompletionMessageEvent<TToolOptions extends ToolOptions = ToolOp
        * The text content of the LLM response.
        */
       content: string;
+      /**
+       * Optional refusal reason returned by the model when content is filtered.
+       */
+      refusal?: string;
       /**
        * Optional deanonymized input messages metadata
        */
@@ -84,6 +88,10 @@ export type ChatCompletionChunkEvent = InferenceTaskEventBase<
      * The content chunk
      */
     content: string;
+    /**
+     * Optional refusal reason chunk.
+     */
+    refusal?: string;
     /**
      * The tool call chunks
      */

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/messages.ts

@@ -57,6 +57,10 @@ export type AssistantMessage<TToolCalls extends ToolCall[] | undefined = ToolCal
      * Can be null if the LLM called a tool.
      */
     content: string | null;
+    /**
+     * Optional refusal reason returned by the model when content is filtered.
+     */
+    refusal?: string | null;
     // make sure `toolCalls` inherits the optionality from `TToolCalls`
   } & (TToolCalls extends ToolCall[]
       ? {

x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/from_inference/chunks.ts

@@ -26,10 +26,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: {},
   });
 };

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,

x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml

@@ -2925,6 +2925,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.

x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml

@@ -2925,6 +2925,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.

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts

@@ -421,6 +421,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.
    */

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml

@@ -443,6 +443,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.

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 {

x-pack/platform/plugins/shared/inference/server/chat_complete/utils/chunks_into_message.ts

@@ -45,7 +45,7 @@ export function chunksIntoMessage<TToolOptions extends ToolOptions>({
 
           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 });
@@ -59,6 +59,7 @@ export function chunksIntoMessage<TToolOptions extends ToolOptions>({
           return {
             type: ChatCompletionEventType.ChatCompletionMessage,
             content,
+            ...(refusal ? { refusal } : {}),
             toolCalls: validatedToolCalls,
           };
         })

x-pack/platform/plugins/shared/inference/server/chat_complete/utils/merge_chunks.ts

@@ -9,6 +9,7 @@ import type { ChatCompletionChunkEvent, UnvalidatedToolCall } from '@kbn/inferen
 
 interface UnvalidatedMessage {
   content: string;
+  refusal?: string;
   tool_calls: UnvalidatedToolCall[];
 }
 
@@ -19,6 +20,9 @@ export const mergeChunks = (chunks: ChatCompletionChunkEvent[]): UnvalidatedMess
   const message = chunks.reduce<UnvalidatedMessage>(
     (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];

x-pack/platform/plugins/shared/inference/server/chat_complete/utils/stream_to_response.ts

@@ -35,6 +35,7 @@ export const streamToResponse = <TToolOptions extends ToolOptions = ToolOptions>
 
         return {
           content: messageEvent.content,
+          refusal: messageEvent.refusal,
           toolCalls: messageEvent.toolCalls,
           tokens: tokenEvent?.tokens,
           deanonymized_input: messageEvent.deanonymized_input,

x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/helpers.ts

@@ -47,6 +47,9 @@ export function chunksIntoMessage(obs$: Observable<UnifiedChatCompleteResponse>)
           (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<UnifiedChatCompleteResponse>)
               {
                 message: {
                   content: '',
+                  refusal: null,
                   role: 'assistant',
                 },
               },

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.test.ts

@@ -326,6 +326,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',
+              }),
+            ]),
+          }),
+        ]),
+      })
+    );
+  });
+
   it('generates UUID for messages without id', async () => {
     const messageWithoutId = createMockMessage({ id: undefined });
     setupSuccessfulTest();

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.ts

@@ -106,6 +106,7 @@ export const transformToUpdateScheme = (
       '@timestamp': message.timestamp,
       id: message.id ?? uuidv4(),
       content: message.content,
+      ...(message.refusal ? { refusal: message.refusal } : {}),
       is_error: message.isError,
       reader: message.reader,
       role: message.role,

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.ts

@@ -96,6 +96,7 @@ export const transformToCreateScheme = (
       '@timestamp': message.timestamp,
       id: message.id ?? uuidv4(),
       content: message.content,
+      ...(message.refusal ? { refusal: message.refusal } : {}),
       is_error: message.isError,
       reader: message.reader,
       role: message.role,

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/field_maps_configuration.ts

@@ -92,6 +92,11 @@ export const conversationsFieldMap: FieldMap = {
     array: false,
     required: false,
   },
+  'messages.refusal': {
+    type: 'text',
+    array: false,
+    required: false,
+  },
   'messages.reader': {
     type: 'object',
     array: false,

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts

@@ -51,6 +51,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 } : {}),

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts

@@ -29,6 +29,7 @@ export interface EsConversationSchema {
     '@timestamp': string;
     id?: string;
     content: string;
+    refusal?: string;
     reader?: Reader;
     role: MessageRole;
     is_error?: boolean;
@@ -77,6 +78,7 @@ export interface CreateMessageSchema {
     '@timestamp': string;
     id: string;
     content: string;
+    refusal?: string;
     reader?: Reader;
     role: MessageRole;
     is_error?: boolean;

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts

@@ -29,6 +29,7 @@ export interface UpdateConversationSchema {
     '@timestamp': string;
     id: string;
     content: string;
+    refusal?: string;
     reader?: Reader;
     role: MessageRole;
     is_error?: boolean;
@@ -136,6 +137,7 @@ export const transformToUpdateScheme = (
             '@timestamp': message.timestamp,
             id: message.id ?? uuidv4(),
             content: message.content,
+            ...(message.refusal ? { refusal: message.refusal } : {}),
             is_error: message.isError,
             reader: message.reader,
             role: message.role,

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts

@@ -35,6 +35,7 @@ import type { AIAssistantDataClient } from '../../../ai_assistant_data_clients';
 
 export type OnLlmResponse = (args: {
   content: string;
+  refusal?: string;
   interruptValue?: InterruptValue;
   traceData?: Message['traceData'];
   isError?: boolean;

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts

@@ -106,12 +106,14 @@ export const streamGraph = async ({
 
   const handleFinalContent = (args: {
     finalResponse: string;
+    refusal?: string;
     isError: boolean;
     interruptValue?: InterruptValue;
   }) => {
     if (onLlmResponse) {
       onLlmResponse({
         content: args.finalResponse,
+        refusal: args.refusal,
         interruptValue: args.interruptValue,
         traceData: {
           transactionId: streamingSpan?.transaction?.ids?.['transaction.id'],
@@ -151,7 +153,11 @@ export const streamGraph = async ({
           !data.output.lc_kwargs?.tool_calls?.length &&
           !didEnd
         ) {
-          handleFinalContent({ finalResponse: data.output.content, isError: false });
+          const refusal =
+            typeof data.output?.additional_kwargs?.refusal === 'string'
+              ? (data.output.additional_kwargs.refusal as string)
+              : undefined;
+          handleFinalContent({ finalResponse: data.output.content, refusal, isError: false });
         } 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' &&
@@ -234,10 +240,15 @@ 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({
         content: output,
         traceData,
+        ...(refusal ? { refusal } : {}),
       });
     }
 

x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts

@@ -212,6 +212,7 @@ export const chatCompleteRoute = (
 
           const onLlmResponse: OnLlmResponse = async ({
             content,
+            refusal,
             traceData = {},
             isError = false,
           }): Promise<void> => {
@@ -225,6 +226,7 @@ export const chatCompleteRoute = (
                 conversationId,
                 conversationsDataClient,
                 messageContent: prunedContent,
+                messageRefusal: refusal,
                 replacements: latestReplacements,
                 isError,
                 traceData,