From 8c33a8b757f21be111ad6b04fc86ddcc58099e2f Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Wed, 11 Jun 2025 11:00:40 +0200 Subject: [PATCH 1/3] [Obs AI Assistant] Anonymization support (#223351) Re-submit of https://github.com/elastic/kibana/pull/216352 as it has merge conflicts and we don't have write permissions for Sandra's remote. To test, add the following to your kibana.yml: ``` uiSettings: overrides: "observability:aiAssistantAnonymizationRules": - id: "ner" type: "ner" enabled: true - id: "beach" type: "regex" enabled: true pattern: "sandy" ``` --------- Co-authored-by: Sandra Gonzales Co-authored-by: Sandra G (cherry picked from commit 71ec37a2a688ca5b0791658e7cc7d0634e7139a9) # Conflicts: # x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_timeline.tsx # x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts --- .../src/chat_complete/events.ts | 11 +- .../kbn-ai-assistant/src/chat/chat_item.tsx | 24 +- ...chat_item_content_inline_prompt_editor.tsx | 9 +- .../src/chat/chat_timeline.tsx | 56 +++- .../kbn-ai-assistant/src/utils/builders.ts | 1 + .../common/index.ts | 9 +- .../common/types.ts | 36 +++ .../common/ui_settings/settings_keys.ts | 1 + .../common/utils/anonymization/redaction.ts | 35 +++ .../public/analytics/schemas/common.ts | 87 ------ .../components/message_panel/message_text.tsx | 24 +- .../public/index.ts | 2 + .../public/plugin.tsx | 3 +- .../server/routes/runtime_types.ts | 9 + .../service/anonymization/chunk_text.ts | 29 ++ .../service/anonymization/deanonymize_text.ts | 59 ++++ .../detect_regex_entities.test.ts | 100 ++++++ .../anonymization/detect_regex_entities.ts | 70 +++++ .../service/anonymization/get_entity_hash.ts | 15 + .../get_redactable_message_parts.ts | 33 ++ .../server/service/anonymization/index.ts | 288 ++++++++++++++++++ .../server/service/client/index.test.ts | 17 +- .../server/service/client/index.ts | 70 +++-- .../server/service/index.ts | 21 +- .../conversation_component_template.ts | 4 + .../common/ui_settings.ts | 38 +++ .../routes/components/settings_page.test.tsx | 6 +- .../server/config.ts | 6 +- .../server/plugin.ts | 4 +- .../anonymization/anonymization.spec.ts | 158 ++++++++++ .../apis/observability/ai_assistant/index.ts | 1 + .../ai_assistant/utils/advanced_settings.ts | 2 +- 32 files changed, 1089 insertions(+), 139 deletions(-) create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/common/utils/anonymization/redaction.ts delete mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/public/analytics/schemas/common.ts create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/chunk_text.ts create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/deanonymize_text.ts create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/detect_regex_entities.test.ts create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/detect_regex_entities.ts create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/get_entity_hash.ts create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/get_redactable_message_parts.ts create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/index.ts create mode 100644 x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/anonymization/anonymization.spec.ts 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 77bd535d0fda1..2ae609333a17c 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,7 +35,16 @@ export type ChatCompletionMessageEvent['toolCalls']; } >; - +// with unredactions +export interface ChatCompletionUnredactedMessageEvent< + TToolOptions extends ToolOptions = ToolOptions +> extends ChatCompletionMessageEvent { + unredactions: Array<{ + entity: string; + class_name: string; + hash: string; + }>; +} /** * Represent a partial tool call present in a chunk event. * diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_item.tsx b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_item.tsx index 706a72f3a5d23..1241b989ee4d9 100644 --- a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_item.tsx +++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_item.tsx @@ -27,7 +27,19 @@ import { ChatItemActions } from './chat_item_actions'; import { ChatItemAvatar } from './chat_item_avatar'; import { ChatItemContentInlinePromptEditor } from './chat_item_content_inline_prompt_editor'; import { ChatTimelineItem } from './chat_timeline'; - +// Helper function to extract plain text from a React node. +function extractTextFromReactNode(node: React.ReactNode): string { + if (typeof node === 'string' || typeof node === 'number') { + return node.toString(); + } + if (Array.isArray(node)) { + return node.map(extractTextFromReactNode).join(''); + } + if (React.isValidElement(node)) { + return extractTextFromReactNode(node.props.children); + } + return ''; +} export interface ChatItemProps extends Omit { onActionClick: ChatActionClickHandler; onEditSubmit: (message: Message) => void; @@ -36,6 +48,7 @@ export interface ChatItemProps extends Omit { onSendTelemetry: (eventWithPayload: TelemetryEventTypeWithPayload) => void; onStopGeneratingClick: () => void; isConversationOwnedByCurrentUser: boolean; + displayContent?: React.ReactNode; } const moreCompactHeaderClassName = css` @@ -95,6 +108,7 @@ export function ChatItem({ onRegenerateClick, onSendTelemetry, onStopGeneratingClick, + anonymizedHighlightedContent, }: ChatItemProps) { const accordionId = useGeneratedHtmlId({ prefix: 'chat' }); @@ -135,17 +149,21 @@ export function ChatItem({ return onEditSubmit(newMessage); }; + // extract text if content is not a string. const handleCopyToClipboard = () => { - navigator.clipboard.writeText(content || ''); + const copyText = typeof content === 'string' ? content : extractTextFromReactNode(content); + navigator.clipboard.writeText(copyText || ''); }; - let contentElement: React.ReactNode = + let contentElement: React.ReactNode; + contentElement = content || loading || error ? ( void; onSubmit: (message: Message) => void; + anonymizedHighlightedContent?: React.ReactNode; } const textContainerClassName = css` @@ -48,6 +49,7 @@ export function ChatItemContentInlinePromptEditor({ onActionClick, onSendTelemetry, onSubmit, + anonymizedHighlightedContent, }: Props) { return !editing ? ( - + ) : ( void; } +// helper using detected entity positions to transform user messages into react node to add text highlighting +function highlightContent( + content: string, + detectedEntities: Array<{ start_pos: number; end_pos: number; entity: string }> +): React.ReactNode { + // Sort the entities by start position + const sortedEntities = [...detectedEntities].sort((a, b) => a.start_pos - b.start_pos); + const parts: Array = []; + let lastIndex = 0; + sortedEntities.forEach((entity, index) => { + // Add the text before the entity + if (entity.start_pos > lastIndex) { + parts.push(content.substring(lastIndex, entity.start_pos)); + } + // Wrap the sensitive text in a span with highlight styles + parts.push( + + {content.substring(entity.start_pos, entity.end_pos)} + + ); + lastIndex = entity.end_pos; + }); + // Add any remaining text after the last entity + if (lastIndex < content.length) { + parts.push(content.substring(lastIndex)); + } + return parts; +} const euiCommentListClassName = css` padding-bottom: 32px; `; @@ -103,6 +135,20 @@ export function ChatTimeline({ display: none; } `; + const { + services: { uiSettings }, + } = useKibana(); + + const { anonymizationEnabled } = useMemo(() => { + try { + const rules = uiSettings?.get(aiAssistantAnonymizationRules); + return { + anonymizationEnabled: Array.isArray(rules) && rules.some((rule) => rule.enabled), + }; + } catch (e) { + return { anonymizationEnabled: false }; + } + }, [uiSettings]); const items = useMemo(() => { const timelineItems = getTimelineItemsfromConversation({ @@ -121,8 +167,13 @@ export function ChatTimeline({ let currentGroup: ChatTimelineItem[] | null = null; for (const item of timelineItems) { + const { role, content, unredactions } = item.message.message; if (item.display.hide || !item) continue; + if (anonymizationEnabled && role === 'user' && content && unredactions) { + item.anonymizedHighlightedContent = highlightContent(content, unredactions); + } + if (item.display.collapsed) { if (currentGroup) { currentGroup.push(item); @@ -147,6 +198,7 @@ export function ChatTimeline({ isConversationOwnedByCurrentUser, isArchived, onActionClick, + anonymizationEnabled, ]); return ( diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/utils/builders.ts b/x-pack/platform/packages/shared/kbn-ai-assistant/src/utils/builders.ts index 219d243b98993..d3dbe60837ec2 100644 --- a/x-pack/platform/packages/shared/kbn-ai-assistant/src/utils/builders.ts +++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/utils/builders.ts @@ -20,6 +20,7 @@ type BuildMessageProps = DeepPartial & { name: string; trigger: MessageRole.Assistant | MessageRole.User | MessageRole.Elastic; }; + unredactions?: Message['message']['unredactions']; }; }; diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/common/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/common/index.ts index 59b77b2b12385..f60db3258a756 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/common/index.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/common/index.ts @@ -5,7 +5,13 @@ * 2.0. */ -export type { Message, Conversation, KnowledgeBaseEntry, ConversationCreateRequest } from './types'; +export type { + Message, + Conversation, + KnowledgeBaseEntry, + ConversationCreateRequest, + AnonymizationRule, +} from './types'; export { KnowledgeBaseEntryRole, MessageRole, @@ -47,6 +53,7 @@ export { aiAssistantLogsIndexPattern, aiAssistantSimulatedFunctionCalling, aiAssistantSearchConnectorIndexPattern, + aiAssistantAnonymizationRules, } from './ui_settings/settings_keys'; export { concatenateChatCompletionChunks } from './utils/concatenate_chat_completion_chunks'; diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/common/types.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/common/types.ts index 030c0705b3086..d6b6e7b787572 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/common/types.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/common/types.ts @@ -29,11 +29,31 @@ export interface PendingMessage { aborted?: boolean; error?: any; } +export interface DetectedEntity { + entity: string; + class_name: string; + start_pos: number; + end_pos: number; + hash: string; + type: 'ner' | 'regex'; +} + +export type DetectedEntityType = DetectedEntity['type']; +export interface Unredaction { + entity: string; + class_name: string; + start_pos: number; + end_pos: number; + type: 'ner' | 'regex'; +} + +export type UnredactionType = Unredaction['type']; export interface Message { '@timestamp': string; message: { content?: string; + unredactions?: Unredaction[]; name?: string; role: MessageRole; function_call?: { @@ -160,3 +180,19 @@ export enum ConversationAccess { SHARED = 'shared', PRIVATE = 'private', } + +export interface InferenceChunk { + chunkText: string; + charStartOffset: number; +} + +export interface AnonymizationRule { + id: string; + entityClass: string; + type: 'regex' | 'ner'; + pattern?: string; + enabled: boolean; + builtIn: boolean; + description?: string; + normalize?: boolean; +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/common/ui_settings/settings_keys.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/common/ui_settings/settings_keys.ts index fdcb34f22ecde..ae5947f0e707a 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/common/ui_settings/settings_keys.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/common/ui_settings/settings_keys.ts @@ -14,3 +14,4 @@ export const aiAssistantSimulatedFunctionCalling = export const aiAssistantSearchConnectorIndexPattern = 'observability:aiAssistantSearchConnectorIndexPattern'; export const aiAssistantPreferredAIAssistantType = 'aiAssistant:preferredAIAssistantType'; +export const aiAssistantAnonymizationRules = 'observability:aiAssistantAnonymizationRules'; diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/common/utils/anonymization/redaction.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/common/utils/anonymization/redaction.ts new file mode 100644 index 0000000000000..f4a37d73fefe0 --- /dev/null +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/common/utils/anonymization/redaction.ts @@ -0,0 +1,35 @@ +/* + * 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 { DetectedEntity } from '../../types'; +/** Regex matching object‑hash placeholders (40 hex chars) */ +export const HASH_REGEX = /[0-9a-f]{40}/g; +/** + * Replace each entity span in the original with its hash. + */ + +export function redactEntities(original: string, entities: DetectedEntity[]): string { + let redacted = original; + entities + .slice() + .sort((a, b) => b.start_pos - a.start_pos) + .forEach((e) => { + redacted = redacted.slice(0, e.start_pos) + e.hash + redacted.slice(e.end_pos); + }); + + return redacted; +} +/** + * Replace every placeholder in a string with its real value + * (taken from `hashMap`). + */ +export function unhashString( + contentWithHashes: string, + hashMap: Map +): string { + return contentWithHashes.replace(HASH_REGEX, (h) => hashMap.get(h)?.value ?? h); +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/public/analytics/schemas/common.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/public/analytics/schemas/common.ts deleted file mode 100644 index 4a2739ef82c35..0000000000000 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/public/analytics/schemas/common.ts +++ /dev/null @@ -1,87 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -import type { RootSchema } from '@kbn/core/public'; -import { AssistantScope } from '@kbn/ai-assistant-common'; -import type { Message } from '../../../common'; - -export const messageSchema: RootSchema = { - '@timestamp': { - type: 'text', - _meta: { - description: 'The timestamp of the message.', - }, - }, - message: { - properties: { - content: { - type: 'text', - _meta: { - description: 'The response generated by the LLM.', - optional: true, - }, - }, - name: { - type: 'text', - _meta: { - description: 'The name of the function that was executed.', - optional: true, - }, - }, - role: { - type: 'text', - _meta: { - description: 'The actor that generated the response.', - }, - }, - data: { - type: 'text', - _meta: { - description: '', - optional: true, - }, - }, - function_call: { - _meta: { - description: 'The function call that was executed.', - optional: true, - }, - properties: { - name: { - type: 'text', - _meta: { - description: 'The name of the function that was executed.', - optional: false, - }, - }, - arguments: { - type: 'text', - _meta: { - description: 'The arguments that were used when executing the function.', - optional: true, - }, - }, - trigger: { - type: 'text', - _meta: { - description: 'The actor which triggered the execution of this function.', - }, - }, - }, - }, - }, - }, - scopes: { - type: 'array', - items: { - type: 'text', - _meta: { - description: 'The scopes that were used when generating the message.', - }, - }, - }, -}; diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx b/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx index a4e652eda80cd..74c699b841e74 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx @@ -27,6 +27,7 @@ interface Props { content: string; loading: boolean; onActionClick: ChatActionClickHandler; + anonymizedHighlightedContent?: React.ReactNode; } const ANIMATION_TIME = 1; @@ -115,7 +116,12 @@ const esqlLanguagePlugin = () => { }; }; -export function MessageText({ loading, content, onActionClick }: Props) { +export function MessageText({ + loading, + content, + onActionClick, + anonymizedHighlightedContent, +}: Props) { const containerClassName = css` overflow-wrap: anywhere; `; @@ -187,13 +193,15 @@ export function MessageText({ loading, content, onActionClick }: Props) { return ( - - {`${content}${loading ? CURSOR : ''}`} - + {anonymizedHighlightedContent || ( + + {`${content}${loading ? CURSOR : ''}`} + + )} ); } diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/public/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/public/index.ts index 20684793de075..d405cfe8d62b7 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/public/index.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/public/index.ts @@ -93,6 +93,7 @@ export { ObservabilityAIAssistantTelemetryEventType } from './analytics/telemetr export { createFunctionRequestMessage } from '../common/utils/create_function_request_message'; export { createFunctionResponseMessage } from '../common/utils/create_function_response_message'; +export { redactEntities, unhashString } from '../common/utils/anonymization/redaction'; export type { ObservabilityAIAssistantAPIClientRequestParamsOf, @@ -106,6 +107,7 @@ export { useKibana } from './hooks/use_kibana'; export { aiAssistantLogsIndexPattern, aiAssistantSimulatedFunctionCalling, + aiAssistantAnonymizationRules, aiAssistantSearchConnectorIndexPattern, aiAssistantPreferredAIAssistantType, } from '../common/ui_settings/settings_keys'; diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/public/plugin.tsx b/x-pack/platform/plugins/shared/observability_ai_assistant/public/plugin.tsx index 2753b750dc288..285be306e9522 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/public/plugin.tsx +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/public/plugin.tsx @@ -46,7 +46,8 @@ export class ObservabilityAIAssistantPlugin constructor(context: PluginInitializerContext) { this.logger = context.logger.get(); - this.scopeFromConfig = context.config.get().scope; + const config = context.config.get(); + this.scopeFromConfig = config.scope; } setup( coreSetup: CoreSetup, diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/runtime_types.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/runtime_types.ts index da1b5ba19d79f..1970efc8b205d 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/runtime_types.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/runtime_types.ts @@ -15,6 +15,14 @@ import { type StarterPrompt, } from '../../common/types'; +export const unredactionRt = t.type({ + entity: t.string, + class_name: t.string, + start_pos: t.number, + end_pos: t.number, + type: t.union([t.literal('ner'), t.literal('regex')]), +}); + export const messageRt: t.Type = t.type({ '@timestamp': t.string, message: t.intersection([ @@ -45,6 +53,7 @@ export const messageRt: t.Type = t.type({ arguments: t.string, }), ]), + unredactions: t.array(unredactionRt), }), ]), }); diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/chunk_text.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/chunk_text.ts new file mode 100644 index 0000000000000..0484edc71a9c5 --- /dev/null +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/chunk_text.ts @@ -0,0 +1,29 @@ +/* + * 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 { InferenceChunk } from '../../../common/types'; + +/** + * Splits text into chunks of specified maximum size + * + * Used to prepare text for ML model inference by breaking it into + * smaller pieces that the model can handle efficiently. + * + * @param text - Text to be chunked + * @param maxChars - Maximum characters per chunk (default: 1000) + * @returns Array of chunks with their original character offsets + */ +export function chunkText(text: string, maxChars = 1_000): InferenceChunk[] { + const chunks: InferenceChunk[] = []; + for (let i = 0; i < text.length; i += maxChars) { + chunks.push({ + chunkText: text.slice(i, i + maxChars), + charStartOffset: i, + }); + } + return chunks; +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/deanonymize_text.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/deanonymize_text.ts new file mode 100644 index 0000000000000..f45337611a11d --- /dev/null +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/deanonymize_text.ts @@ -0,0 +1,59 @@ +/* + * 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 { HASH_REGEX } from '../../../common/utils/anonymization/redaction'; +import { DetectedEntity, DetectedEntityType } from '../../../common/types'; + +/** + * Replaces hash placeholders in text with their original values + * + * Takes text containing hash placeholders like "{hash123}" and replaces + * them with their original values from the provided hash map. Also generates + * entity metadata for each replaced value. + * + * @param contentWithHashes - Text containing hash placeholders + * @param hashMap - Map of hash values to their original entity information + * @returns Object containing deanonymized text and detected entities + */ +export function deanonymizeText( + contentWithHashes: string, + hashMap: Map +) { + const detectedEntities: DetectedEntity[] = []; + let unhashedText = ''; + let cursor = 0; + + for (const match of contentWithHashes.matchAll(HASH_REGEX)) { + const hash = match[0]; + const rep = hashMap.get(hash); + if (!rep) { + continue; // keep unknown hash as‑is + } + + // copy segment before the hash + unhashedText += contentWithHashes.slice(cursor, match.index!); + + // insert real value & capture span + const start = unhashedText.length; + unhashedText += rep.value; + const end = unhashedText.length; + + detectedEntities.push({ + entity: rep.value, + class_name: rep.class_name, + start_pos: start, + end_pos: end, + type: rep.type, + hash, + }); + + cursor = match.index! + hash.length; + } + + unhashedText += contentWithHashes.slice(cursor); + return { unhashedText, detectedEntities }; +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/detect_regex_entities.test.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/detect_regex_entities.test.ts new file mode 100644 index 0000000000000..ef9597d8fb872 --- /dev/null +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/detect_regex_entities.test.ts @@ -0,0 +1,100 @@ +/* + * 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 { detectRegexEntities } from './detect_regex_entities'; +import { getEntityHash } from './get_entity_hash'; +import type { AnonymizationRule } from './detect_regex_entities'; + +describe('getEntityHash', () => { + it('returns the same hash for differently cased emails when normalize=true', () => { + const lower = getEntityHash('KATY@GMAIL.COM', 'EMAIL', true); + const upper = getEntityHash('katy@gmail.com', 'EMAIL', true); + expect(lower).toEqual(upper); + }); + + it('returns different hashes for differently cased emails when normalize=false', () => { + const lower = getEntityHash('KATY@GMAIL.COM', 'EMAIL'); + const upper = getEntityHash('katy@gmail.com', 'EMAIL'); + expect(lower).not.toEqual(upper); + }); + + it('defaults normalize=false when not passed', () => { + const withExplicitFalse = getEntityHash('Test', 'CUSTOM'); + const withDefault = getEntityHash('Test', 'CUSTOM'); + expect(withExplicitFalse).toEqual(withDefault); + }); +}); + +describe('detectRegexEntities', () => { + // Mock logger + const mockLogger = { + error: jest.fn(), + } as any; + + // Test rules - similar to what would be in the anonymization.spec.ts + const testRules: AnonymizationRule[] = [ + { + id: 'email-rule', + entityClass: 'EMAIL', + type: 'regex', + pattern: '\\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Z|a-z]{2,}\\b', + enabled: true, + builtIn: true, + description: 'Email detection', + normalize: true, + }, + { + id: 'url-rule', + entityClass: 'URL', + type: 'regex', + pattern: '\\bhttps?://[^\\s]+\\b', + enabled: true, + builtIn: true, + description: 'URL detection', + }, + { + id: 'ip-rule', + entityClass: 'IP', + type: 'regex', + pattern: '\\b(?:\\d{1,3}\\.){3}\\d{1,3}\\b', + enabled: true, + builtIn: true, + description: 'IP address detection', + }, + ]; + + it('detects and hashes an email address', () => { + const content = 'Contact me at TEST@Example.Com'; + const entities = detectRegexEntities(content, testRules, mockLogger); + expect(entities).toHaveLength(1); + expect(entities[0].entity).toBe('TEST@Example.Com'); + expect(entities[0].class_name).toBe('EMAIL'); + + // Confirm normalization by comparing hash to expected + const expectedHash = getEntityHash('test@example.com', 'EMAIL', true); + expect(entities[0].hash).toBe(expectedHash); + }); + + it('detects URL, IP, and email all in one string', () => { + const content = + 'Check https://kibana.elastic.co, reach me at user@elastic.co, or ping 192.168.1.1'; + const entities = detectRegexEntities(content, testRules, mockLogger); + + const classes = entities.map((e) => e.class_name); + expect(classes).toContain('URL'); + expect(classes).toContain('EMAIL'); + expect(classes).toContain('IP'); + }); + + it('computes correct start and end positions', () => { + const content = 'Email: hello@example.com'; + const entities = detectRegexEntities(content, testRules, mockLogger); + const emailEntity = entities.find((e) => e.class_name === 'EMAIL'); + expect(emailEntity).toBeDefined(); + expect(content.slice(emailEntity!.start_pos, emailEntity!.end_pos)).toBe(emailEntity!.entity); + }); +}); diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/detect_regex_entities.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/detect_regex_entities.ts new file mode 100644 index 0000000000000..443c0eb71dda0 --- /dev/null +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/detect_regex_entities.ts @@ -0,0 +1,70 @@ +/* + * 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 { Logger } from '@kbn/core/server'; +import { DetectedEntity } from '../../../common/types'; +import { getEntityHash } from './get_entity_hash'; + +function getMatches( + content: string, + regex: RegExp, + className: string, + normalize: boolean = false +): DetectedEntity[] { + const result: DetectedEntity[] = []; + let match: RegExpExecArray | null; + + while ((match = regex.exec(content)) !== null) { + const entityText = match[0]; + const start = match.index; + const end = start + entityText.length; + const hash = getEntityHash(entityText, className, normalize); + result.push({ + entity: entityText, + class_name: className, + start_pos: start, + end_pos: end, + hash, + type: 'regex', + }); + } + return result; +} + +export interface AnonymizationRule { + id: string; + entityClass: string; + type: 'regex' | 'ner'; + pattern?: string; + enabled: boolean; + builtIn: boolean; + description?: string; + normalize?: boolean; +} + +export function detectRegexEntities( + content: string, + rules: AnonymizationRule[] = [], + logger: Logger +): DetectedEntity[] { + const results: DetectedEntity[] = []; + + // Filter for enabled regex rules + const regexRules = rules.filter((rule) => rule.type === 'regex' && rule.enabled && rule.pattern); + + // Apply each regex rule + for (const rule of regexRules) { + try { + const regex = new RegExp(rule.pattern!, 'g'); + results.push(...getMatches(content, regex, rule.entityClass, rule.normalize ?? false)); + } catch (error) { + // Skip invalid regex patterns + logger.error(`Invalid regex pattern in rule ${rule.id}: ${rule.pattern}`, error); + } + } + + return results; +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/get_entity_hash.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/get_entity_hash.ts new file mode 100644 index 0000000000000..617cb457ea6dd --- /dev/null +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/get_entity_hash.ts @@ -0,0 +1,15 @@ +/* + * 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 objectHash from 'object-hash'; +export function getEntityHash( + entity: string, + className: string, + normalize: boolean = false +): string { + const textForHash = normalize ? entity.toLowerCase() : entity; + return objectHash({ entity: textForHash, class_name: className }); +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/get_redactable_message_parts.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/get_redactable_message_parts.ts new file mode 100644 index 0000000000000..7cb446713f87b --- /dev/null +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/get_redactable_message_parts.ts @@ -0,0 +1,33 @@ +/* + * 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 { ChatCompletionMessageEvent } from '@kbn/inference-common'; +import { type Message, MessageRole } from '../../../common/types'; + +// TODO: to use in redactMessages when we update NER to work with JSON string +export function getRedactableMessageParts(message: Message) { + if (message.message.role === MessageRole.Assistant) { + return { + // we might want to consider not running detection on assistant responses (content) + // as they're already coming from the LLM + content: message.message.content, + function_call: message.message.function_call?.arguments, + }; + } + + return { + content: message.message.content, + }; +} +export function getRedactableMessageEventParts(event: ChatCompletionMessageEvent) { + return { + content: event.content, + toolCalls: event.toolCalls?.map((toolCall) => ({ + function: toolCall.function, + })), + }; +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/index.ts new file mode 100644 index 0000000000000..1e5ecc6da4abf --- /dev/null +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/anonymization/index.ts @@ -0,0 +1,288 @@ +/* + * 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 { withInferenceSpan } from '@kbn/inference-tracing'; +import { ElasticsearchClient } from '@kbn/core/server'; +import pLimit from 'p-limit'; +import { OperatorFunction, map } from 'rxjs'; +import type { Logger } from '@kbn/core/server'; +import { chunk } from 'lodash'; +import { ChatCompletionEvent, ChatCompletionEventType } from '@kbn/inference-common'; +import { ChatCompletionUnredactedMessageEvent } from '@kbn/inference-common/src/chat_complete/events'; +import { unhashString, redactEntities } from '../../../common/utils/anonymization/redaction'; +import { detectRegexEntities } from './detect_regex_entities'; +import { deanonymizeText } from './deanonymize_text'; +import { chunkText } from './chunk_text'; +import { getRedactableMessageEventParts } from './get_redactable_message_parts'; +import { + type DetectedEntity, + DetectedEntityType, + type InferenceChunk, + type Message, + type AnonymizationRule, +} from '../../../common/types'; +import { getEntityHash } from './get_entity_hash'; + +const NER_MODEL_ID = 'elastic__distilbert-base-uncased-finetuned-conll03-english'; +const DEFAULT_MAX_CONCURRENT_REQUESTS = 5; +export interface Dependencies { + esClient: { + asCurrentUser: ElasticsearchClient; + }; + logger: Logger; + anonymizationRules: AnonymizationRule[]; +} + +export class AnonymizationService { + private readonly esClient: { asCurrentUser: ElasticsearchClient }; + private readonly logger: Logger; + private rules: AnonymizationRule[]; + + private currentHashMap: Map< + string, + { value: string; class_name: string; type: DetectedEntityType } + > = new Map(); + + constructor({ esClient, logger, anonymizationRules }: Dependencies) { + this.esClient = esClient; + this.logger = logger; + this.rules = anonymizationRules; + } + + private async detectNamedEntities(chunks: InferenceChunk[]): Promise { + this.logger.debug(`Detecting named entities in ${chunks.length} text chunks`); + + // Maximum number of concurrent requests to the ML model + const limiter = pLimit(DEFAULT_MAX_CONCURRENT_REQUESTS); + + // Batch size - number of documents to send in each request + const BATCH_SIZE = 10; + + // Create batches of chunks for the inference request + const batches = chunk(chunks, BATCH_SIZE); + this.logger.debug(`Processing ${batches.length} batches of up to ${BATCH_SIZE} chunks each`); + + const tasks = batches.map((batchChunks) => + limiter(async () => + withInferenceSpan('infer_ner', async () => { + let response; + try { + response = await this.esClient.asCurrentUser.ml.inferTrainedModel({ + model_id: NER_MODEL_ID, + docs: batchChunks.map((batchChunk) => ({ text_field: batchChunk.chunkText })), + }); + } catch (error) { + throw new Error('NER inference failed', { cause: error }); + } + // Process results from all documents in the batch + const batchResults: DetectedEntity[] = []; + const inferenceResults = response?.inference_results || []; + + if (inferenceResults.length !== batchChunks.length) { + this.logger.warn( + `NER returned ${inferenceResults.length} results for ${batchChunks.length} docs in batch` + ); + } + + // Match results with their original chunks to maintain offsets + inferenceResults.forEach((result, index) => { + const batchChunk = batchChunks[index]; + const entities = result.entities || []; + + batchResults.push( + ...entities.map((e) => ({ + ...e, + start_pos: e.start_pos + batchChunk.charStartOffset, + end_pos: e.end_pos + batchChunk.charStartOffset, + type: 'ner' as const, + hash: getEntityHash(e.entity, e.class_name), + })) + ); + }); + + return batchResults; + }) + ) + ); + + const results = await Promise.all(tasks); + const flatResults = results.flat(); + this.logger.debug(`Total entities detected: ${flatResults.length}`); + return flatResults; + } + + private async detectEntities(content: string): Promise { + // Skip detection if there's no content + if (!content || !content.trim()) { + return []; + } + + this.logger.debug(`Detecting entities in text content`); + + // Filter rules by type + const nerRules = this.rules.filter((rule) => rule.type === 'ner' && rule.enabled); + const regexRules = this.rules.filter((rule) => rule.type === 'regex' && rule.enabled); + + // Only run NER if we have NER rules enabled + let nerEntities: DetectedEntity[] = []; + if (nerRules.length > 0) { + // Detect entities using NER + const chunks = chunkText(content); + nerEntities = await this.detectNamedEntities(chunks); + } + + // Detect entities using regex patterns + const regexEntities = detectRegexEntities(content, regexRules, this.logger); + + // Combine and deduplicate entities + const combined = [...nerEntities, ...regexEntities]; + + // Give precedence to regex entities over overlapping NER entities + const deduped = combined.filter((ent) => + // Regex entities take precedence over NER entities + ent.type === 'regex' + ? true + : // check for intersecting ranges with regex entities + !regexEntities.some((re) => ent.start_pos < re.end_pos && ent.end_pos > re.start_pos) + ); + + this.logger.debug( + `Detected ${nerEntities.length} NER entities and ${regexEntities.length} regex entities, ${deduped.length} after deduplication` + ); + + return deduped; + } + + /** + * Redacts all user messages by replacing detected entities with {hash} placeholders + */ + async redactMessages(messages: Message[]): Promise<{ redactedMessages: Message[] }> { + if (!this.rules.length) { + return { redactedMessages: messages }; + } + + for (const msg of messages) { + // we may want to ignore assistant responses in the future + if (!msg.message.content) { + continue; + } + + const entities = await this.detectEntities(msg.message.content); + + if (entities.length) { + msg.message.content = redactEntities(msg.message.content, entities); + + // Update hashMap + entities.forEach((e) => { + this.currentHashMap.set(e.hash, { + value: e.entity, + class_name: e.class_name, + type: e.type, + }); + }); + } + } + + // Redact entity values inside any function_call.arguments JSON strings + for (const msg of messages) { + const argsStr = msg.message.function_call?.arguments; + if (!argsStr) continue; + + let redactedArgs = argsStr; + // Replace every known entity value with its hash + this.currentHashMap.forEach((info, hash) => { + redactedArgs = redactedArgs.split(info.value).join(hash); + }); + msg.message.function_call!.arguments = redactedArgs; + } + return { redactedMessages: messages }; + } + + /** + * Restores all {hash} placeholders in-place and attaches `unredactions` array + * for UI highlighting (content only). + */ + unredactMessages(messages: Message[]): { unredactedMessages: Message[] } { + for (const msg of messages) { + const content = msg.message.content; + if (content) { + const { unhashedText, detectedEntities } = deanonymizeText(content, this.currentHashMap); + + msg.message.content = unhashedText; + if (detectedEntities.length > 0) { + msg.message.unredactions = detectedEntities.map(({ hash, ...rest }) => rest); + } + } + + // also unhash function_call.arguments if present + if (msg.message.function_call?.arguments) { + msg.message.function_call.arguments = unhashString( + msg.message.function_call.arguments, + this.currentHashMap + ); + } + } + return { unredactedMessages: messages }; + } + unredactChatCompletionEvent(): OperatorFunction< + ChatCompletionEvent, + ChatCompletionEvent | ChatCompletionUnredactedMessageEvent + > { + return (source$) => { + return source$.pipe( + map((event): ChatCompletionEvent | ChatCompletionUnredactedMessageEvent => { + if (event.type === ChatCompletionEventType.ChatCompletionMessage) { + const redacted = getRedactableMessageEventParts(event); + const contentUnredaction = + 'content' in redacted && redacted.content && typeof redacted.content === 'string' + ? deanonymizeText(redacted.content, this.currentHashMap) + : undefined; + const unredaction = deanonymizeText(JSON.stringify(redacted), this.currentHashMap); + const unredactedObj = JSON.parse(unredaction.unhashedText); + + // Ensure tool call arguments are always strings, even if they're objects in the JSON + if (unredactedObj.toolCalls) { + unredactedObj.toolCalls = unredactedObj.toolCalls.map( + (toolCall: { + function?: { + name?: string; + arguments?: any; + }; + }) => { + if (toolCall.function && typeof toolCall.function.arguments === 'object') { + // Convert object arguments to strings to maintain compatibility in redactMessages and unredactMessages + return { + ...toolCall, + function: { + ...toolCall.function, + arguments: JSON.stringify(toolCall.function.arguments), + }, + }; + } + return toolCall; + } + ); + } + + const redactedEvent: ChatCompletionUnredactedMessageEvent = { + ...event, + ...unredactedObj, + }; + if (contentUnredaction && contentUnredaction.detectedEntities.length > 0) { + redactedEvent.unredactions = contentUnredaction.detectedEntities; + // TODO: not being passed through due to concatenateChatCompletionChunks filtering out non chunk events + // causing knowledge base entities to be stored with redactions + // and having to call undreactMessages outside chat + } + return redactedEvent; + } + return event; + }) + ); + }; + } +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.test.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.test.ts index f8eeb58194f1a..dac384ae5706a 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.test.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.test.ts @@ -31,6 +31,7 @@ import type { KnowledgeBaseService } from '../knowledge_base_service'; import { observableIntoStream } from '../util/observable_into_stream'; import type { ObservabilityAIAssistantConfig } from '../../config'; import type { ObservabilityAIAssistantPluginStartDependencies } from '../../types'; +import { AnonymizationService } from '../anonymization'; interface ChunkDelta { content?: string | undefined; @@ -188,6 +189,13 @@ describe('Observability AI Assistant client', () => { }, inferenceClient: inferenceClientMock, knowledgeBaseService: knowledgeBaseServiceMock, + anonymizationService: new AnonymizationService({ + esClient: { + asCurrentUser: currentUserEsClientMock, + }, + anonymizationRules: [], + logger: loggerMock, + }), logger: loggerMock, namespace: 'default', user: { @@ -839,8 +847,8 @@ describe('Observability AI Assistant client', () => { }); }); - it('sends the function response back to the llm', () => { - expect(inferenceClientMock.chatComplete).toHaveBeenCalledTimes(2); + it('sends the function response back to the llm', async () => { + await waitFor(() => expect(inferenceClientMock.chatComplete).toHaveBeenCalledTimes(2)); expect(inferenceClientMock.chatComplete.mock.lastCall!).toEqual([ { @@ -864,6 +872,11 @@ describe('Observability AI Assistant client', () => { describe('and the assistant replies without a function request', () => { beforeEach(async () => { + // 1) wait for the follow-up chatComplete + await waitFor(() => expect(inferenceClientMock.chatComplete).toHaveBeenCalledTimes(2)); + + // 2) wait until llmSimulator has been initialised + await waitFor(() => expect(llmSimulator).toBeDefined()); await llmSimulator.chunk({ content: 'I am done here' }); await llmSimulator.next({ content: 'I am done here' }); await llmSimulator.complete(); diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts index 914c384746db0..67a3ec6b03c87 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts @@ -62,6 +62,7 @@ import { import { CONTEXT_FUNCTION_NAME } from '../../functions/context'; import type { ChatFunctionClient } from '../chat_function_client'; import { KnowledgeBaseService, RecalledEntry } from '../knowledge_base_service'; +import { AnonymizationService } from '../anonymization'; import { getAccessQuery } from '../util/get_access_query'; import { getSystemMessageFromInstructions } from '../util/get_system_message_from_instructions'; import { failOnNonExistingFunctionCall } from './operators/fail_on_non_existing_function_call'; @@ -102,6 +103,7 @@ export class ObservabilityAIAssistantClient { }; knowledgeBaseService: KnowledgeBaseService; scopes: AssistantScope[]; + anonymizationService: AnonymizationService; } ) {} @@ -205,14 +207,11 @@ export class ObservabilityAIAssistantClient { }): Observable> => { return withInferenceSpan('run_tools', () => { const isConversationUpdate = persist && !!predefinedConversationId; - const conversationId = persist ? predefinedConversationId || v4() : ''; if (persist && !isConversationUpdate && kibanaPublicUrl) { functionClient.registerInstruction( - `This conversation will be persisted in Kibana and available at this url: ${ - kibanaPublicUrl + `/app/observabilityAIAssistant/conversations/${conversationId}` - }.` + `This conversation will be persisted in Kibana and available at this url: ${kibanaPublicUrl}/app/observabilityAIAssistant/conversations/${conversationId}.` ); } @@ -242,18 +241,18 @@ export class ObservabilityAIAssistantClient { ), }).pipe(shareReplay()); - const systemMessage$ = kbUserInstructions$.pipe( - map((kbUserInstructions) => { - return getSystemMessageFromInstructions({ - applicationInstructions: functionClient.getInstructions(), - kbUserInstructions, - apiUserInstructions, - availableFunctionNames: disableFunctions - ? [] - : functionClient.getFunctions().map((fn) => fn.definition.name), - }); - }), - shareReplay() + const systemMessage$ = kbUserInstructions$.pipe( + map((kbUserInstructions) => + getSystemMessageFromInstructions({ + applicationInstructions: functionClient.getInstructions(), + kbUserInstructions, + apiUserInstructions, + availableFunctionNames: disableFunctions + ? [] + : functionClient.getFunctions().map((fn) => fn.definition.name), + }) + ), + shareReplay() ); // we continue the conversation here, after resolving both the materialized @@ -333,9 +332,11 @@ export class ObservabilityAIAssistantClient { systemMessage$, ]).pipe( switchMap(([addedMessages, title, systemMessage]) => { - const initialMessagesWithAddedMessages = initialMessages.concat(addedMessages); - - const lastMessage = last(initialMessagesWithAddedMessages); + const { unredactedMessages } = + this.dependencies.anonymizationService.unredactMessages( + initialMessages.concat(addedMessages) + ); + const lastMessage = last(unredactedMessages); // if a function request is at the very end, close the stream to consumer // without persisting or updating the conversation. we need to wait @@ -358,7 +359,7 @@ export class ObservabilityAIAssistantClient { omit(conversation._source, 'messages'), // update messages and system message - { messages: initialMessagesWithAddedMessages, systemMessage }, + { messages: unredactedMessages, systemMessage }, // update title { @@ -389,7 +390,7 @@ export class ObservabilityAIAssistantClient { labels: {}, numeric_labels: {}, systemMessage, - messages: initialMessagesWithAddedMessages, + messages: unredactedMessages, archived: false, }) ).pipe( @@ -418,13 +419,11 @@ export class ObservabilityAIAssistantClient { () => `Added message: ${JSON.stringify(event.message)}` ); break; - case StreamingChatResponseEventType.ConversationCreate: this.dependencies.logger.debug( () => `Created conversation: ${JSON.stringify(event.conversation)}` ); break; - case StreamingChatResponseEventType.ConversationUpdate: this.dependencies.logger.debug( () => `Updated conversation: ${JSON.stringify(event.conversation)}` @@ -501,10 +500,26 @@ export class ObservabilityAIAssistantClient { if (stream) { return defer(() => - this.dependencies.inferenceClient.chatComplete({ - ...options, - stream: true, - }) + from(this.dependencies.anonymizationService.redactMessages(messages)).pipe( + switchMap(({ redactedMessages }) => { + this.dependencies.logger.debug( + () => + `Calling inference client for name: "${name}" with options: ${JSON.stringify( + options + )}` + ); + return ( + this.dependencies.inferenceClient + .chatComplete({ + ...options, + stream: true, + messages: convertMessagesForInference(redactedMessages, this.dependencies.logger), + }) + // unredact complete assistant response event + .pipe(this.dependencies.anonymizationService.unredactChatCompletionEvent()) + ); + }) + ) ).pipe( convertInferenceEventsToStreamingEvents(), failOnNonExistingFunctionCall({ functions }), @@ -522,6 +537,7 @@ export class ObservabilityAIAssistantClient { } else { return this.dependencies.inferenceClient.chatComplete({ ...options, + messages: convertMessagesForInference(messages, this.dependencies.logger), stream: false, }) as TStream extends true ? never : Promise; } diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index.ts index 72c21eb05b343..944e8e91f1853 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index.ts @@ -18,6 +18,9 @@ import { KnowledgeBaseService } from './knowledge_base_service'; import type { RegistrationCallback, RespondFunctionResources } from './types'; import { ObservabilityAIAssistantConfig } from '../config'; import { createOrUpdateConversationIndexAssets } from './index_assets/create_or_update_conversation_index_assets'; +import { AnonymizationService } from './anonymization'; +import { aiAssistantAnonymizationRules } from '../../common'; +import type { AnonymizationRule } from '../../common/types'; export function getResourceName(resource: string) { return `.kibana-observability-ai-assistant-${resource}`; @@ -93,6 +96,11 @@ export class ObservabilityAIAssistantService { const user = plugins.security.authc.getCurrentUser(request); const soClient = coreStart.savedObjects.getScopedClient(request); + const uiSettingsClient = coreStart.uiSettings.asScopedToClient(soClient); + + // Read anonymization rules from advanced settings + const anonymizationRules = + (await uiSettingsClient.get(aiAssistantAnonymizationRules)) ?? []; const basePath = coreStart.http.basePath.get(request); @@ -100,6 +108,7 @@ export class ObservabilityAIAssistantService { const inferenceClient = plugins.inference.getClient({ request }); const { asInternalUser } = coreStart.elasticsearch.client; + const { asCurrentUser } = coreStart.elasticsearch.client.asScoped(request); const kbService = new KnowledgeBaseService({ core: this.core, @@ -109,17 +118,25 @@ export class ObservabilityAIAssistantService { asInternalUser, }, }); + const anonymizationService = new AnonymizationService({ + logger: this.logger.get('anonymization'), + esClient: { + asCurrentUser, + }, + anonymizationRules, + }); return new ObservabilityAIAssistantClient({ core: this.core, config: this.config, actionsClient: await plugins.actions.getActionsClientWithRequest(request), - uiSettingsClient: coreStart.uiSettings.asScopedToClient(soClient), + uiSettingsClient, namespace: spaceId, esClient: { asInternalUser, - asCurrentUser: coreStart.elasticsearch.client.asScoped(request).asCurrentUser, + asCurrentUser, }, + anonymizationService, inferenceClient, logger: this.logger, user: user diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/templates/conversation_component_template.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/templates/conversation_component_template.ts index a0490cb960811..e43d08f9b926e 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/templates/conversation_component_template.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/templates/conversation_component_template.ts @@ -83,6 +83,10 @@ export const conversationComponentTemplate: ClusterComponentTemplate['component_ trigger: keyword, }, }, + unredactions: { + type: 'object', + enabled: false, + }, }, }, }, diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_management/common/ui_settings.ts b/x-pack/solutions/observability/plugins/observability_ai_assistant_management/common/ui_settings.ts index 007c9ed2d498c..3d9a08b6362dc 100644 --- a/x-pack/solutions/observability/plugins/observability_ai_assistant_management/common/ui_settings.ts +++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_management/common/ui_settings.ts @@ -11,6 +11,7 @@ import { i18n } from '@kbn/i18n'; import { aiAssistantSimulatedFunctionCalling, aiAssistantSearchConnectorIndexPattern, + aiAssistantAnonymizationRules, } from '@kbn/observability-ai-assistant-plugin/common'; export const uiSettings: Record = { @@ -57,4 +58,41 @@ export const uiSettings: Record = { requiresPageReload: true, solution: 'oblt', }, + [aiAssistantAnonymizationRules]: { + category: ['observability'], + name: i18n.translate( + 'xpack.observabilityAiAssistantManagement.settingsTab.anonymizationRulesLabel', + { defaultMessage: 'Anonymization Rules' } + ), + value: [], // Default is an empty array, which disables all anonymization rules. + description: i18n.translate( + 'xpack.observabilityAiAssistantManagement.settingsPage.anonymizationRulesDescription', + { + defaultMessage: + 'JSON array of anonymization rules. Each rule is an object with properties:\n' + + '- id: unique string identifier\n' + + '- entityClass: class of entity (e.g., PER, ORG, EMAIL, URL)\n' + + '- type: "ner" or "regex"\n' + + '- pattern: (for regex rules) the regex string to match\n' + + '- enabled: boolean flag to turn the rule on or off\n' + + '- builtIn: boolean indicating this is a built‑in rule\n' + + '- description: optional human‑readable description\n' + + 'Default is an empty array, which disables all anonymization rules.', + } + ), + schema: schema.arrayOf( + schema.object({ + id: schema.string(), + entityClass: schema.string(), + type: schema.oneOf([schema.literal('ner'), schema.literal('regex')]), + pattern: schema.string(), + enabled: schema.boolean(), + builtIn: schema.boolean(), + description: schema.maybe(schema.string()), + }) + ), + type: 'json', + requiresPageReload: true, + solution: 'oblt', + }, }; diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_management/public/routes/components/settings_page.test.tsx b/x-pack/solutions/observability/plugins/observability_ai_assistant_management/public/routes/components/settings_page.test.tsx index c4051b9665b57..bd39621d767ab 100644 --- a/x-pack/solutions/observability/plugins/observability_ai_assistant_management/public/routes/components/settings_page.test.tsx +++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_management/public/routes/components/settings_page.test.tsx @@ -16,7 +16,11 @@ const useKnowledgeBaseMock = useKnowledgeBase as jest.Mock; describe('Settings Page', () => { const appContextValue = { - config: { spacesEnabled: true, visibilityEnabled: true, logSourcesEnabled: true }, + config: { + spacesEnabled: true, + visibilityEnabled: true, + logSourcesEnabled: true, + }, setBreadcrumbs: () => {}, }; useKnowledgeBaseMock.mockReturnValue({ diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_management/server/config.ts b/x-pack/solutions/observability/plugins/observability_ai_assistant_management/server/config.ts index 559904ce3a126..3a20dbecb08f0 100644 --- a/x-pack/solutions/observability/plugins/observability_ai_assistant_management/server/config.ts +++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_management/server/config.ts @@ -18,5 +18,9 @@ export type ObservabilityAIAssistantManagementConfig = TypeOf = { schema: configSchema, - exposeToBrowser: { logSourcesEnabled: true, spacesEnabled: true, visibilityEnabled: true }, + exposeToBrowser: { + logSourcesEnabled: true, + spacesEnabled: true, + visibilityEnabled: true, + }, }; diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_management/server/plugin.ts b/x-pack/solutions/observability/plugins/observability_ai_assistant_management/server/plugin.ts index 8690213b645cc..3bb40498ac8e1 100644 --- a/x-pack/solutions/observability/plugins/observability_ai_assistant_management/server/plugin.ts +++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_management/server/plugin.ts @@ -6,6 +6,7 @@ */ import { CoreSetup, CoreStart, Plugin } from '@kbn/core/server'; +import { aiAssistantAnonymizationRules } from '@kbn/observability-ai-assistant-plugin/common'; import { uiSettings } from '../common/ui_settings'; export type ObservabilityPluginSetup = ReturnType; @@ -20,7 +21,8 @@ export class AiAssistantManagementPlugin implements Plugin, plugins: PluginSetup) { - core.uiSettings.register(uiSettings); + const { [aiAssistantAnonymizationRules]: anonymizationRules, ...restSettings } = uiSettings; + core.uiSettings.register(restSettings); return {}; } diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/anonymization/anonymization.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/anonymization/anonymization.spec.ts new file mode 100644 index 0000000000000..e6af4f0a0deec --- /dev/null +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/anonymization/anonymization.spec.ts @@ -0,0 +1,158 @@ +/* + * 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 expect from '@kbn/expect'; +import { MessageRole, type Message } from '@kbn/observability-ai-assistant-plugin/common'; +import { + LlmProxy, + createLlmProxy, +} from '../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy'; +import { setAdvancedSettings } from '../utils/advanced_settings'; +import type { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context'; +import { clearConversations } from '../utils/conversation'; + +export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) { + const log = getService('log'); + const supertest = getService('supertest'); + const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); + const es = getService('es'); + + let proxy: LlmProxy; + let connectorId: string; + + const conversationsIndex = '.kibana-observability-ai-assistant-conversations-000001'; + + describe('anonymization', function () { + const userText1 = 'My name is Claudia and my email is claudia@example.com'; + const userText2 = 'my website is http://claudia.is'; + // LLM proxy is not working on MKI + this.tags(['failsOnMKI']); + + before(async () => { + await clearConversations(es); + proxy = await createLlmProxy(log); + connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({ + port: proxy.getPort(), + }); + + // configure anonymization rules for these tests + await setAdvancedSettings(supertest, { + 'observability:aiAssistantAnonymizationRules': [ + { + id: 'email_regex', + entityClass: 'EMAIL', + type: 'regex', + pattern: '[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}', + enabled: true, + builtIn: true, + description: 'Anonymize email addresses', + normalize: true, + }, + { + id: 'url_regex', + entityClass: 'URL', + type: 'regex', + pattern: 'https?://[^\\s]+', + enabled: true, + builtIn: true, + description: 'Anonymize URLs', + normalize: true, + }, + ], + }); + }); + + after(async () => { + proxy.close(); + await observabilityAIAssistantAPIClient.deleteActionConnector({ actionId: connectorId }); + await clearConversations(es); + await setAdvancedSettings(supertest, { + 'observability:aiAssistantAnonymizationRules': [], + }); + }); + + it('does not send detected entities to the LLM via chat/complete', async () => { + void proxy.interceptTitle('Title for a new conversation'); + const simulatorPromise = proxy.interceptWithResponse('ok'); + + const res = await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat/complete', + params: { + body: { + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { role: MessageRole.User, content: userText1 }, + } as Message, + { + '@timestamp': new Date().toISOString(), + message: { role: MessageRole.User, content: userText2 }, + } as Message, + ], + connectorId, + persist: true, + screenContexts: [], + scopes: ['all'], + }, + }, + }); + + expect(res.status).to.be(200); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + + const userMsgsReq = simulator.requestBody.messages.filter((m: any) => m.role === 'user'); + expect(userMsgsReq).to.have.length(2); + // First message + const firstMsgReq = userMsgsReq[0].content; + expect(firstMsgReq).to.not.contain('claudia@example.com'); + expect( + typeof firstMsgReq === 'string' && (firstMsgReq.match(/[0-9a-f]{40}/g) || []).length + ).to.be(1); + // Second message + const secMsgReq = userMsgsReq[1].content; + expect(secMsgReq).to.not.contain('http://claudia.is'); + expect( + typeof secMsgReq === 'string' && (secMsgReq.match(/[0-9a-f]{40}/g) || []).length + ).to.be(1); + }); + + it('stores original content and detected entities in Elasticsearch', async () => { + // Refresh the index to make sure our document is searchable + await es.indices.refresh({ + index: conversationsIndex, + }); + const searchRes = await es.search({ + index: conversationsIndex, + size: 1, + sort: '@timestamp:desc', + }); + const hit: any = searchRes.hits.hits[0]._source; + // Find the stored user messages + const storedUserMsgs = hit.messages + .filter( + (m: any) => + m.message.role === 'user' && + (m.message.content === userText1 || m.message.content === userText2) + ) + .map((m: any) => m.message); + expect(storedUserMsgs).to.have.length(2); + + // First stored message + const firstSavedMsg = storedUserMsgs[0]; + expect(firstSavedMsg.unredactions).to.have.length(1); + expect(firstSavedMsg.unredactions[0].entity).to.eql('claudia@example.com'); + expect(firstSavedMsg.unredactions[0].class_name).to.eql('EMAIL'); + + // Second stored message + const secSavedMsg = storedUserMsgs[1]; + expect(secSavedMsg.unredactions).to.have.length(1); + expect(secSavedMsg.unredactions[0].entity).to.eql('http://claudia.is'); + expect(secSavedMsg.unredactions[0].class_name).to.eql('URL'); + }); + }); +} diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index.ts index 675248c90b09a..6a294e1e02421 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index.ts @@ -43,6 +43,7 @@ export default function aiAssistantApiIntegrationTests({ loadTestFile(require.resolve('./index_assets/index_assets.spec.ts')); loadTestFile(require.resolve('./connectors/connectors.spec.ts')); loadTestFile(require.resolve('./conversations/conversations.spec.ts')); + loadTestFile(require.resolve('./anonymization/anonymization.spec.ts')); // public endpoints loadTestFile(require.resolve('./public_complete/public_complete.spec.ts')); diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/advanced_settings.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/advanced_settings.ts index 98fe191096253..a0a60ced56f62 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/advanced_settings.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/advanced_settings.ts @@ -14,7 +14,7 @@ import type SuperTest from 'supertest'; export const setAdvancedSettings = async ( supertest: SuperTest.Agent, - settings: Record + settings: Record ) => { return supertest .post('/internal/kibana/settings') From 8dd10221a9d51f77fa7d2c0b11bdfdf0a899a03b Mon Sep 17 00:00:00 2001 From: Sandra Gonzales Date: Mon, 16 Jun 2025 12:42:03 -0400 Subject: [PATCH 2/3] fix conflict --- .../observability_ai_assistant/server/service/client/index.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts index 67a3ec6b03c87..dbfdfb7e9d0f9 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts @@ -513,7 +513,7 @@ export class ObservabilityAIAssistantClient { .chatComplete({ ...options, stream: true, - messages: convertMessagesForInference(redactedMessages, this.dependencies.logger), + messages: convertMessagesForInference(redactedMessages), }) // unredact complete assistant response event .pipe(this.dependencies.anonymizationService.unredactChatCompletionEvent()) @@ -537,7 +537,7 @@ export class ObservabilityAIAssistantClient { } else { return this.dependencies.inferenceClient.chatComplete({ ...options, - messages: convertMessagesForInference(messages, this.dependencies.logger), + messages: convertMessagesForInference(messages), stream: false, }) as TStream extends true ? never : Promise; } From 533a03ff00d5ad266cbdb60306e9cadac8f061d4 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Mon, 16 Jun 2025 17:08:28 +0000 Subject: [PATCH 3/3] [CI] Auto-commit changed files from 'node scripts/eslint --no-cache --fix' --- .../server/service/client/index.ts | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts index dbfdfb7e9d0f9..3ed905df50169 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts @@ -241,18 +241,18 @@ export class ObservabilityAIAssistantClient { ), }).pipe(shareReplay()); - const systemMessage$ = kbUserInstructions$.pipe( - map((kbUserInstructions) => - getSystemMessageFromInstructions({ - applicationInstructions: functionClient.getInstructions(), - kbUserInstructions, - apiUserInstructions, - availableFunctionNames: disableFunctions - ? [] - : functionClient.getFunctions().map((fn) => fn.definition.name), - }) - ), - shareReplay() + const systemMessage$ = kbUserInstructions$.pipe( + map((kbUserInstructions) => + getSystemMessageFromInstructions({ + applicationInstructions: functionClient.getInstructions(), + kbUserInstructions, + apiUserInstructions, + availableFunctionNames: disableFunctions + ? [] + : functionClient.getFunctions().map((fn) => fn.definition.name), + }) + ), + shareReplay() ); // we continue the conversation here, after resolving both the materialized