diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/capabilities/index.ts index b7c8e092b2eec..854e7f0d84f74 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -21,5 +21,4 @@ export type AssistantFeatureKey = keyof AssistantFeatures; export const defaultAssistantFeatures = Object.freeze({ assistantModelEvaluation: false, defendInsights: true, - contentReferencesEnabled: false, }); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.ts index 887ccf26bc8a6..e7a65fc928507 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.ts @@ -18,11 +18,10 @@ import { ContentReferencesStore, ContentReferenceBlock } from '../types'; export const pruneContentReferences = ( content: string, contentReferencesStore: ContentReferencesStore -): ContentReferences | undefined => { +): ContentReferences => { const fullStore = contentReferencesStore.getStore(); const prunedStore: Record = {}; const matches = content.matchAll(/\{reference\([0-9a-zA-Z]+\)\}/g); - let isPrunedStoreEmpty = true; for (const match of matches) { const referenceElement = match[0]; @@ -30,15 +29,10 @@ export const pruneContentReferences = ( if (!(referenceId in prunedStore)) { const contentReference = fullStore[referenceId]; if (contentReference) { - isPrunedStoreEmpty = false; prunedStore[referenceId] = contentReference; } } } - if (isPrunedStoreEmpty) { - return undefined; - } - return prunedStore; }; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.ts index 967ceb64b3546..df8bab861907e 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.ts @@ -26,8 +26,11 @@ export const getContentReferenceId = ( * @returns ContentReferenceBlock */ export const contentReferenceBlock = ( - contentReference: ContentReference -): ContentReferenceBlock => { + contentReference: ContentReference | undefined +): ContentReferenceBlock | '' => { + if (!contentReference) { + return ''; + } return `{reference(${contentReference.id})}`; }; @@ -36,7 +39,10 @@ export const contentReferenceBlock = ( * @param contentReference A ContentReference * @returns the string: `Reference: ` */ -export const contentReferenceString = (contentReference: ContentReference) => { +export const contentReferenceString = (contentReference: ContentReference | undefined) => { + if (!contentReference) { + return ''; + } return `Citation: ${contentReferenceBlock(contentReference)}` as const; }; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts index 9ff0c2ebf59e7..8777e8d728279 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts @@ -19,6 +19,5 @@ import { z } from '@kbn/zod'; export type GetCapabilitiesResponse = z.infer; export const GetCapabilitiesResponse = z.object({ assistantModelEvaluation: z.boolean(), - contentReferencesEnabled: z.boolean(), defendInsights: z.boolean(), }); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml index 2dc79cfe3d116..e9b6ca9697256 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml @@ -22,13 +22,10 @@ paths: properties: assistantModelEvaluation: type: boolean - contentReferencesEnabled: - type: boolean defendInsights: type: boolean required: - assistantModelEvaluation - - contentReferencesEnabled - defendInsights '400': description: Generic Error diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/index.tsx index 82765fac41e15..a936eae8ddafa 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/index.tsx @@ -101,7 +101,6 @@ const AssistantComponent: React.FC = ({ showAnonymizedValues, setContentReferencesVisible, setShowAnonymizedValues, - assistantFeatures: { contentReferencesEnabled }, } = useAssistantContext(); const [selectedPromptContexts, setSelectedPromptContexts] = useState< @@ -408,7 +407,6 @@ const AssistantComponent: React.FC = ({ currentUserAvatar, systemPromptContent: currentSystemPrompt?.content, contentReferencesVisible, - contentReferencesEnabled, })} // Avoid comments going off the flyout css={css` @@ -439,7 +437,6 @@ const AssistantComponent: React.FC = ({ contentReferencesVisible, euiTheme.size.l, selectedPromptContextsCount, - contentReferencesEnabled, ] ); @@ -457,9 +454,7 @@ const AssistantComponent: React.FC = ({ return ( <> - {contentReferencesEnabled && ( - - )} + {chatHistoryVisible && ( = React.memo( contentReferencesVisible, showAnonymizedValues, setShowAnonymizedValues, - assistantFeatures: { contentReferencesEnabled }, } = useAssistantContext(); const [isPopoverOpen, setPopover] = useState(false); @@ -256,62 +255,60 @@ export const SettingsContextMenu: React.FC = React.memo( - - {contentReferencesEnabled && ( - - - - ( - - } - > - {children} - - )} - > - + + + ( + + } + > + {children} + + )} + > + + + + + {str}, + }} /> - - - - {str}, - }} - /> - } - > - - - - - - )} + } + > + + + + + + = React.memo( handleShowAlertsModal, knowledgeBase.latestAlerts, showDestroyModal, - contentReferencesEnabled, euiTheme.size.m, euiTheme.size.xs, selectedConversationHasCitations, diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/types.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/types.tsx index ed5b690e9aa57..8f724e1e0b280 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/types.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/types.tsx @@ -83,6 +83,5 @@ export type GetAssistantMessages = (commentArgs: { currentUserAvatar?: UserAvatar; setIsStreaming: (isStreaming: boolean) => void; systemPromptContent?: string; - contentReferencesVisible?: boolean; - contentReferencesEnabled?: boolean; + contentReferencesVisible: boolean; }) => EuiCommentProps[]; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/scripts/draw_graph_script.ts b/x-pack/solutions/security/plugins/elastic_assistant/scripts/draw_graph_script.ts index a5ef8e07a07e3..89c6d7b50c295 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/scripts/draw_graph_script.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/scripts/draw_graph_script.ts @@ -68,7 +68,6 @@ async function getAssistantGraph(logger: Logger): Promise { createLlmInstance, tools: [], replacements: {}, - contentReferencesEnabled: false, savedObjectsClient: savedObjectsClientMock.create(), }); return graph.getGraph(); 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 314c83728bcd0..96b12a9051f20 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 @@ -97,6 +97,16 @@ export const conversationsFieldMap: FieldMap = { array: false, required: false, }, + 'messages.metadata': { + type: 'object', + array: false, + required: false, + }, + 'messages.metadata.content_references': { + type: 'flattened', + array: false, + required: false, + }, replacements: { type: 'object', array: false, @@ -168,18 +178,3 @@ export const conversationsFieldMap: FieldMap = { required: false, }, } as const; - -// Once the `contentReferencesEnabled` feature flag is removed, the properties from the schema bellow should me moved into `conversationsFieldMap` -export const conversationsContentReferencesFieldMap: FieldMap = { - ...conversationsFieldMap, - 'messages.metadata': { - type: 'object', - array: false, - required: false, - }, - 'messages.metadata.content_references': { - type: 'flattened', - array: false, - required: false, - }, -} as const; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index 7be0eda1b299c..f01258daf55b1 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -153,7 +153,7 @@ export const getStructuredToolForIndexEntry = ({ }: { indexEntry: IndexEntry; esClient: ElasticsearchClient; - contentReferencesStore: ContentReferencesStore | undefined; + contentReferencesStore: ContentReferencesStore; logger: Logger; }): DynamicStructuredTool => { const inputSchema = indexEntry.inputSchema?.reduce((prev, input) => { @@ -223,8 +223,7 @@ export const getStructuredToolForIndexEntry = ({ const result = await esClient.search(params); const kbDocs = result.hits.hits.map((hit) => { - const reference = - contentReferencesStore && contentReferencesStore.add((p) => createReference(p.id, hit)); + const reference = contentReferencesStore.add((p) => createReference(p.id, hit)); if (indexEntry.outputFields && indexEntry.outputFields.length > 0) { return indexEntry.outputFields.reduce( @@ -232,13 +231,13 @@ export const getStructuredToolForIndexEntry = ({ // @ts-expect-error return { ...prev, [field]: hit._source[field] }; }, - reference ? { citation: contentReferenceBlock(reference) } : {} + { citation: contentReferenceBlock(reference) } ); } return { text: hit.highlight?.[indexEntry.field].join('\n --- \n'), - ...(reference ? { citation: contentReferenceBlock(reference) } : {}), + citation: contentReferenceBlock(reference), }; }); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 561467f2256ea..a4cad89abdc22 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -817,7 +817,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { contentReferencesStore, esClient, }: { - contentReferencesStore: ContentReferencesStore | undefined; + contentReferencesStore: ContentReferencesStore; esClient: ElasticsearchClient; }): Promise => { const user = this.options.currentUser; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 9aef0e6c8db58..a55a02cc977a9 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -33,10 +33,7 @@ import { errorResult, successResult, } from './create_resource_installation_helper'; -import { - conversationsFieldMap, - conversationsContentReferencesFieldMap, -} from '../ai_assistant_data_clients/conversations/field_maps_configuration'; +import { conversationsFieldMap } from '../ai_assistant_data_clients/conversations/field_maps_configuration'; import { assistantPromptsFieldMap } from '../ai_assistant_data_clients/prompts/field_maps_configuration'; import { assistantAnonymizationFieldsFieldMap } from '../ai_assistant_data_clients/anonymization_fields/field_maps_configuration'; import { AIAssistantDataClient } from '../ai_assistant_data_clients'; @@ -107,9 +104,6 @@ export class AIAssistantService { private isKBSetupInProgress: boolean = false; private hasInitializedV2KnowledgeBase: boolean = false; private productDocManager?: ProductDocBaseStartContract['management']; - // Temporary 'feature flag' to determine if we should initialize the new message metadata mappings, toggled when citations should be enabled. - private contentReferencesEnabled: boolean = false; - private hasInitializedContentReferences: boolean = false; // Temporary 'feature flag' to determine if we should initialize the new knowledge base mappings private assistantDefaultInferenceEndpoint: boolean = false; @@ -228,17 +222,6 @@ export class AIAssistantService { void ensureProductDocumentationInstalled(this.productDocManager, this.options.logger); } - // If contentReferencesEnabled is true, re-install data stream resources for new mappings if it has not been done already - if (this.contentReferencesEnabled && !this.hasInitializedContentReferences) { - this.options.logger.debug(`Creating conversation datastream with content references`); - this.conversationsDataStream = this.createDataStream({ - resource: 'conversations', - kibanaVersion: this.options.kibanaVersion, - fieldMap: conversationsContentReferencesFieldMap, - }); - this.hasInitializedContentReferences = true; - } - await this.conversationsDataStream.install({ esClient, logger: this.options.logger, @@ -494,19 +477,6 @@ export class AIAssistantService { return null; } - // Note: Due to plugin lifecycle and feature flag registration timing, we need to pass in the feature flag here - // Remove this param and initialization when the `contentReferencesEnabled` feature flag is removed - if (opts.contentReferencesEnabled) { - this.contentReferencesEnabled = true; - } - - // If contentReferences are enable but the conversation field mappings with content references have not been initialized, - // then call initializeResources which will create the datastreams with content references field mappings. After they have - // been created, hasInitializedContentReferences will ensure they dont get created again. - if (this.contentReferencesEnabled && !this.hasInitializedContentReferences) { - await this.initializeResources(); - } - return new AIAssistantConversationsDataClient({ logger: this.options.logger, elasticsearchClientPromise: this.options.elasticsearchClientPromise, 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 e1133aa0d9bb0..d81fe718ea3c0 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 @@ -49,7 +49,7 @@ export interface AgentExecutorParams { assistantTools?: AssistantTool[]; connectorId: string; conversationId?: string; - contentReferencesStore: ContentReferencesStore | undefined; + contentReferencesStore: ContentReferencesStore; dataClients?: AssistantDataClients; esClient: ElasticsearchClient; langChainMessages: BaseMessage[]; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts index bafd7c472bdf7..e594412d4ea45 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts @@ -42,7 +42,6 @@ export interface GetDefaultAssistantGraphParams { signal?: AbortSignal; tools: StructuredTool[]; replacements: Replacements; - contentReferencesEnabled: boolean; } export type DefaultAssistantGraph = ReturnType; @@ -58,7 +57,6 @@ export const getDefaultAssistantGraph = ({ signal, tools, replacements, - contentReferencesEnabled = false, }: GetDefaultAssistantGraphParams) => { try { // Default graph state @@ -123,10 +121,6 @@ export const getDefaultAssistantGraph = ({ value: (x: string, y?: string) => y ?? x, default: () => 'English', }, - contentReferencesEnabled: { - value: (x: boolean, y?: boolean) => y ?? x, - default: () => contentReferencesEnabled, - }, provider: { value: (x: string, y?: string) => y ?? x, default: () => '', diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts index 22889419885ec..35c6c258a267c 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts @@ -16,6 +16,7 @@ import { APMTracer } from '@kbn/langchain/server/tracers/apm'; import { TelemetryTracer } from '@kbn/langchain/server/tracers/telemetry'; import { pruneContentReferences, MessageMetadata } from '@kbn/elastic-assistant-common'; import { getPrompt, resolveProviderAndModel } from '@kbn/security-ai-prompts'; +import { isEmpty } from 'lodash'; import { localToolPrompts, promptGroupId as toolsGroupId } from '../../../prompt/tool_prompts'; import { promptGroupId } from '../../../prompt/local_prompt_object'; import { getModelOrOss } from '../../../prompt/helpers'; @@ -228,7 +229,6 @@ export const callAssistantGraph: AgentExecutor = async ({ replacements, // some chat models (bedrock) require a signal to be passed on agent invoke rather than the signal passed to the chat model ...(llmType === 'bedrock' ? { signal: abortSignal } : {}), - contentReferencesEnabled: Boolean(contentReferencesStore), }); const inputs: GraphInputs = { responseLanguage, @@ -263,15 +263,12 @@ export const callAssistantGraph: AgentExecutor = async ({ traceOptions, }); - const contentReferences = - contentReferencesStore && pruneContentReferences(graphResponse.output, contentReferencesStore); + const contentReferences = pruneContentReferences(graphResponse.output, contentReferencesStore); const metadata: MessageMetadata = { - ...(contentReferences ? { contentReferences } : {}), + ...(!isEmpty(contentReferences) ? { contentReferences } : {}), }; - const isMetadataPopulated = !!contentReferences; - return { body: { connector_id: connectorId, @@ -279,7 +276,7 @@ export const callAssistantGraph: AgentExecutor = async ({ trace_data: graphResponse.traceData, replacements, status: 'ok', - ...(isMetadataPopulated ? { metadata } : {}), + ...(!isEmpty(metadata) ? { metadata } : {}), ...(graphResponse.conversationId ? { conversationId: graphResponse.conversationId } : {}), }, headers: { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts index 72b4d6d483e93..16faf3132db2c 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts @@ -9,7 +9,6 @@ import { RunnableConfig } from '@langchain/core/runnables'; import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; import { BaseMessage } from '@langchain/core/messages'; import { removeContentReferences } from '@kbn/elastic-assistant-common'; -import { INCLUDE_CITATIONS } from '../../../../prompt/prompts'; import { promptGroupId } from '../../../../prompt/local_prompt_object'; import { getPrompt, promptDictionary } from '../../../../prompt'; import { AgentState, NodeParamsBase } from '../types'; @@ -70,9 +69,6 @@ export async function runAgent({ ? JSON.stringify(knowledgeHistory.map((e) => e.text)) : NO_KNOWLEDGE_HISTORY }`, - include_citations_prompt_placeholder: state.contentReferencesEnabled - ? INCLUDE_CITATIONS - : '', // prepend any user prompt (gemini) input: `${userPrompt}${state.input}`, chat_history: sanitizeChatHistory(state.messages), // TODO: Message de-dupe with ...state spread diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts index af3c315a189d4..587989667b144 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/types.ts @@ -43,7 +43,6 @@ export interface AgentState extends AgentStateBase { connectorId: string; conversation: ConversationResponse | undefined; conversationId: string; - contentReferencesEnabled: boolean; } export interface NodeParamsBase { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.test.ts index 3ab7c5e8351d0..97f3118475aca 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.test.ts @@ -14,10 +14,10 @@ import { describe('prompts', () => { it.each([ - [DEFAULT_SYSTEM_PROMPT, '{include_citations_prompt_placeholder}', 1], - [GEMINI_SYSTEM_PROMPT, '{include_citations_prompt_placeholder}', 1], - [BEDROCK_SYSTEM_PROMPT, '{include_citations_prompt_placeholder}', 1], - [STRUCTURED_SYSTEM_PROMPT, '{include_citations_prompt_placeholder}', 1], + [DEFAULT_SYSTEM_PROMPT, 'Annotate your answer with relevant citations', 1], + [GEMINI_SYSTEM_PROMPT, 'Annotate your answer with relevant citations', 1], + [BEDROCK_SYSTEM_PROMPT, 'Annotate your answer with relevant citations', 1], + [STRUCTURED_SYSTEM_PROMPT, 'Annotate your answer with relevant citations', 1], [DEFAULT_SYSTEM_PROMPT, 'You are a security analyst', 1], [GEMINI_SYSTEM_PROMPT, 'You are an assistant', 1], [BEDROCK_SYSTEM_PROMPT, 'You are a security analyst', 1], diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts index dd39e650f2bf5..170f3e03d3895 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts @@ -7,18 +7,18 @@ export const KNOWLEDGE_HISTORY = 'If available, use the Knowledge History provided to try and answer the question. If not provided, you can try and query for additional knowledge via the KnowledgeBaseRetrievalTool.'; -export const INCLUDE_CITATIONS = `\n\nAnnotate your answer with relevant citations. Here are some example responses with citations: \n1. "Machine learning is increasingly used in cyber threat detection. {reference(prSit)}" \n2. "The alert has a risk score of 72. {reference(OdRs2)}"\n\nOnly use the citations returned by tools\n\n`; -export const DEFAULT_SYSTEM_PROMPT = `You are a security analyst and expert in resolving security incidents. Your role is to assist by answering questions about Elastic Security. Do not answer questions unrelated to Elastic Security. ${KNOWLEDGE_HISTORY} {include_citations_prompt_placeholder}`; +export const INCLUDE_CITATIONS = `\n\nAnnotate your answer with relevant citations. Here are some example responses with citations: \n1. "Machine learning is increasingly used in cyber threat detection. {{reference(prSit)}}" \n2. "The alert has a risk score of 72. {{reference(OdRs2)}}"\n\nOnly use the citations returned by tools\n\n`; +export const DEFAULT_SYSTEM_PROMPT = `You are a security analyst and expert in resolving security incidents. Your role is to assist by answering questions about Elastic Security. Do not answer questions unrelated to Elastic Security. ${KNOWLEDGE_HISTORY} ${INCLUDE_CITATIONS}`; // system prompt from @afirstenberg const BASE_GEMINI_PROMPT = 'You are an assistant that is an expert at using tools and Elastic Security, doing your best to use these tools to answer questions or follow instructions. It is very important to use tools to answer the question or follow the instructions rather than coming up with your own answer. Tool calls are good. Sometimes you may need to make several tool calls to accomplish the task or get an answer to the question that was asked. Use as many tool calls as necessary.'; const KB_CATCH = 'If the knowledge base tool gives empty results, do your best to answer the question from the perspective of an expert security analyst.'; -export const GEMINI_SYSTEM_PROMPT = `${BASE_GEMINI_PROMPT} {include_citations_prompt_placeholder} ${KB_CATCH}`; +export const GEMINI_SYSTEM_PROMPT = `${BASE_GEMINI_PROMPT} ${INCLUDE_CITATIONS} ${KB_CATCH}`; export const BEDROCK_SYSTEM_PROMPT = `${DEFAULT_SYSTEM_PROMPT} Use tools as often as possible, as they have access to the latest data and syntax. Never return tags in the response, but make sure to include tags content in the response. Do not reflect on the quality of the returned search results in your response. ALWAYS return the exact response from NaturalLanguageESQLTool verbatim in the final response, without adding further description.`; export const GEMINI_USER_PROMPT = `Now, always using the tools at your disposal, step by step, come up with a response to this request:\n\n`; -export const STRUCTURED_SYSTEM_PROMPT = `Respond to the human as helpfully and accurately as possible. ${KNOWLEDGE_HISTORY} {include_citations_prompt_placeholder} You have access to the following tools: +export const STRUCTURED_SYSTEM_PROMPT = `Respond to the human as helpfully and accurately as possible. ${KNOWLEDGE_HISTORY} ${INCLUDE_CITATIONS} You have access to the following tools: {tools} 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 1dc44448f3e38..3d0d64bb74b73 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 @@ -27,7 +27,6 @@ import { buildResponse } from '../../lib/build_response'; import { appendAssistantMessageToConversation, createConversationWithUserInput, - DEFAULT_PLUGIN_NAME, getIsKnowledgeBaseInstalled, langChainExecute, performChecks, @@ -87,15 +86,8 @@ export const chatCompleteRoute = ( return checkResponse.response; } - const contentReferencesEnabled = - ctx.elasticAssistant.getRegisteredFeatures( - DEFAULT_PLUGIN_NAME - ).contentReferencesEnabled; - const conversationsDataClient = - await ctx.elasticAssistant.getAIAssistantConversationsDataClient({ - contentReferencesEnabled, - }); + await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); const anonymizationFieldsDataClient = await ctx.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient(); @@ -186,9 +178,7 @@ export const chatCompleteRoute = ( })); } - const contentReferencesStore = contentReferencesEnabled - ? newContentReferencesStore() - : undefined; + const contentReferencesStore = newContentReferencesStore(); const onLlmResponse = async ( content: string, @@ -196,8 +186,7 @@ export const chatCompleteRoute = ( isError = false ): Promise => { if (newConversation?.id && conversationsDataClient) { - const contentReferences = - contentReferencesStore && pruneContentReferences(content, contentReferencesStore); + const contentReferences = pruneContentReferences(content, contentReferencesStore); await appendAssistantMessageToConversation({ conversationId: newConversation?.id, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index bfd7d6df95a08..226b7602b9d76 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -290,13 +290,7 @@ export const postEvaluateRoute = ( }, }; - const contentReferencesEnabled = - assistantContext.getRegisteredFeatures( - DEFAULT_PLUGIN_NAME - ).contentReferencesEnabled; - const contentReferencesStore = contentReferencesEnabled - ? newContentReferencesStore() - : undefined; + const contentReferencesStore = newContentReferencesStore(); // Fetch any applicable tools that the source plugin may have registered const assistantToolParams: AssistantToolParams = { @@ -395,7 +389,6 @@ export const postEvaluateRoute = ( savedObjectsClient, tools, replacements: {}, - contentReferencesEnabled: Boolean(contentReferencesStore), }), }; }) 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 d97a6b77f93b3..346d17eb7bdfc 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 @@ -34,6 +34,7 @@ import { AssistantFeatureKey } from '@kbn/elastic-assistant-common/impl/capabili import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import type { InferenceServerStart } from '@kbn/inference-plugin/server'; import type { LlmTasksPluginStart } from '@kbn/llm-tasks-plugin/server'; +import { isEmpty } from 'lodash'; import { INVOKE_ASSISTANT_SUCCESS_EVENT } from '../lib/telemetry/event_based_telemetry'; import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; import { FindResponse } from '../ai_assistant_data_clients/find'; @@ -175,7 +176,7 @@ export interface AppendAssistantMessageToConversationParams { messageContent: string; replacements: Replacements; conversationId: string; - contentReferences?: ContentReferences | false; + contentReferences: ContentReferences; isError?: boolean; traceData?: Message['traceData']; } @@ -194,11 +195,9 @@ export const appendAssistantMessageToConversation = async ({ } const metadata: MessageMetadata = { - ...(contentReferences ? { contentReferences } : {}), + ...(!isEmpty(contentReferences) ? { contentReferences } : {}), }; - const isMetadataPopulated = Boolean(contentReferences) !== false; - await conversationsDataClient.appendConversationMessages({ existingConversation: conversation, messages: [ @@ -207,7 +206,7 @@ export const appendAssistantMessageToConversation = async ({ messageContent, replacements, }), - metadata: isMetadataPopulated ? metadata : undefined, + metadata: !isEmpty(metadata) ? metadata : undefined, traceData, isError, }), @@ -232,7 +231,7 @@ export interface LangChainExecuteParams { telemetry: AnalyticsServiceSetup; actionTypeId: string; connectorId: string; - contentReferencesStore: ContentReferencesStore | undefined; + contentReferencesStore: ContentReferencesStore; llmTasks?: LlmTasksPluginStart; inference: InferenceServerStart; isOssModel?: boolean; 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 685863cfa97b1..41ecb19c989ed 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 @@ -25,7 +25,6 @@ import { buildResponse } from '../lib/build_response'; import { ElasticAssistantRequestHandlerContext, GetElser } from '../types'; import { appendAssistantMessageToConversation, - DEFAULT_PLUGIN_NAME, getIsKnowledgeBaseInstalled, getSystemPromptFromUserConversation, langChainExecute, @@ -110,18 +109,11 @@ export const postActionsConnectorExecuteRoute = ( const connector = connectors.length > 0 ? connectors[0] : undefined; const isOssModel = isOpenSourceModel(connector); - const contentReferencesEnabled = - assistantContext.getRegisteredFeatures(DEFAULT_PLUGIN_NAME).contentReferencesEnabled; - const conversationsDataClient = - await assistantContext.getAIAssistantConversationsDataClient({ - contentReferencesEnabled, - }); + await assistantContext.getAIAssistantConversationsDataClient(); const promptsDataClient = await assistantContext.getAIAssistantPromptsDataClient(); - const contentReferencesStore = contentReferencesEnabled - ? newContentReferencesStore() - : undefined; + const contentReferencesStore = newContentReferencesStore(); onLlmResponse = async ( content: string, @@ -129,8 +121,7 @@ export const postActionsConnectorExecuteRoute = ( isError = false ): Promise => { if (conversationsDataClient && conversationId) { - const contentReferences = - contentReferencesStore && pruneContentReferences(content, contentReferencesStore); + const contentReferences = pruneContentReferences(content, contentReferencesStore); await appendAssistantMessageToConversation({ conversationId, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts index cc07747725ee7..3295dad4ea5bf 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts @@ -21,7 +21,7 @@ import { ElasticAssistantPluginRouter } from '../../types'; import { buildResponse } from '../utils'; import { EsConversationSchema } from '../../ai_assistant_data_clients/conversations/types'; import { transformESSearchToConversations } from '../../ai_assistant_data_clients/conversations/transforms'; -import { DEFAULT_PLUGIN_NAME, performChecks } from '../helpers'; +import { performChecks } from '../helpers'; export const findUserConversationsRoute = (router: ElasticAssistantPluginRouter) => { router.versioned @@ -58,14 +58,7 @@ export const findUserConversationsRoute = (router: ElasticAssistantPluginRouter) return checkResponse.response; } - const contentReferencesEnabled = - ctx.elasticAssistant.getRegisteredFeatures( - DEFAULT_PLUGIN_NAME - ).contentReferencesEnabled; - - const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient({ - contentReferencesEnabled, - }); + const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); const currentUser = checkResponse.currentUser; const additionalFilter = query.filter ? ` AND ${query.filter}` : ''; diff --git a/x-pack/solutions/security/plugins/security_solution/common/experimental_features.ts b/x-pack/solutions/security/plugins/security_solution/common/experimental_features.ts index 7955243b2cf80..44deffa56d6e5 100644 --- a/x-pack/solutions/security/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/solutions/security/plugins/security_solution/common/experimental_features.ts @@ -114,11 +114,6 @@ export const allowedExperimentalValues = Object.freeze({ */ assistantModelEvaluation: false, - /** - * Enables content references (citations) in the AI Assistant - */ - contentReferencesEnabled: false, - /** * Enables the Managed User section inside the new user details flyout. */ diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.test.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.test.tsx index 9671ba49d762b..a8b8b46212251 100644 --- a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.test.tsx +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.test.tsx @@ -38,6 +38,7 @@ const testProps = { isFetchingResponse: false, currentConversation, showAnonymizedValues, + contentReferencesVisible: true, }; describe('getComments', () => { it('Does not add error state message has no error', () => { diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.tsx index 45f6c4009f72b..a2e5e646f6ae5 100644 --- a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.tsx +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.tsx @@ -62,7 +62,6 @@ export const getComments: GetAssistantMessages = ({ setIsStreaming, systemPromptContent, contentReferencesVisible, - contentReferencesEnabled, }) => { if (!currentConversation) return []; @@ -87,6 +86,7 @@ export const getComments: GetAssistantMessages = ({ refetchCurrentConversation={refetchCurrentConversation} regenerateMessage={regenerateMessageOfConversation} setIsStreaming={setIsStreaming} + contentReferencesVisible={contentReferencesVisible} transformMessage={() => ({ content: '' } as unknown as ContentMessage)} contentReferences={null} isFetching @@ -133,6 +133,7 @@ export const getComments: GetAssistantMessages = ({ regenerateMessage={regenerateMessageOfConversation} setIsStreaming={setIsStreaming} contentReferences={null} + contentReferencesVisible={contentReferencesVisible} transformMessage={() => ({ content: '' } as unknown as ContentMessage)} // we never need to append to a code block in the system comment, which is what this index is used for index={999} @@ -180,7 +181,6 @@ export const getComments: GetAssistantMessages = ({ abortStream={abortStream} contentReferences={null} contentReferencesVisible={contentReferencesVisible} - contentReferencesEnabled={contentReferencesEnabled} index={index} isControlsEnabled={isControlsEnabled} isError={message.isError} @@ -206,7 +206,6 @@ export const getComments: GetAssistantMessages = ({ content={transformedMessage.content} contentReferences={message.metadata?.contentReferences} contentReferencesVisible={contentReferencesVisible} - contentReferencesEnabled={contentReferencesEnabled} index={index} isControlsEnabled={isControlsEnabled} isError={message.isError} diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx index e46e949e05d54..908d6723b5b9b 100644 --- a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx @@ -48,6 +48,7 @@ const testProps = { setIsStreaming: jest.fn(), transformMessage: jest.fn(), contentReferences: undefined, + contentReferencesVisible: true, }; const mockReader = jest.fn() as unknown as ReadableStreamDefaultReader; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.tsx index fd49909ccda4f..c1dd855388827 100644 --- a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.tsx +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.tsx @@ -19,8 +19,7 @@ interface Props { abortStream: () => void; content?: string; contentReferences: StreamingOrFinalContentReferences; - contentReferencesVisible?: boolean; - contentReferencesEnabled?: boolean; + contentReferencesVisible: boolean; isError?: boolean; isFetching?: boolean; isControlsEnabled?: boolean; @@ -36,8 +35,7 @@ export const StreamComment = ({ abortStream, content, contentReferences, - contentReferencesVisible = true, - contentReferencesEnabled = false, + contentReferencesVisible, index, isControlsEnabled = false, isError = false, @@ -114,7 +112,6 @@ export const StreamComment = ({ data-test-subj={isError ? 'errorComment' : undefined} content={message} contentReferences={contentReferences} - contentReferencesEnabled={contentReferencesEnabled} index={index} contentReferencesVisible={contentReferencesVisible} loading={isAnythingLoading} diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx index 60020f8b1abff..de7697f1bad8a 100644 --- a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx @@ -30,7 +30,6 @@ interface Props { content: string; contentReferences: StreamingOrFinalContentReferences; contentReferencesVisible: boolean; - contentReferencesEnabled: boolean; index: number; loading: boolean; ['data-test-subj']?: string; @@ -107,13 +106,11 @@ const loadingCursorPlugin = () => { interface GetPluginDependencies { contentReferences: StreamingOrFinalContentReferences; contentReferencesVisible: boolean; - contentReferencesEnabled: boolean; } const getPluginDependencies = ({ contentReferences, contentReferencesVisible, - contentReferencesEnabled, }: GetPluginDependencies) => { const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); @@ -123,18 +120,14 @@ const getPluginDependencies = ({ processingPlugins[1][1].components = { ...components, - ...(contentReferencesEnabled - ? { - contentReference: (contentReferenceNode) => { - return ( - - ); - }, - } - : {}), + contentReference: (contentReferenceNode) => { + return ( + + ); + }, cursor: Cursor, customCodeBlock: (props) => { return ( @@ -170,7 +163,7 @@ const getPluginDependencies = ({ loadingCursorPlugin, customCodeBlockLanguagePlugin, ...parsingPlugins, - ...(contentReferencesEnabled ? [contentReferenceParser({ contentReferences })] : []), + contentReferenceParser({ contentReferences }), ], processingPluginList: processingPlugins, }; @@ -181,7 +174,6 @@ export function MessageText({ content, contentReferences, contentReferencesVisible, - contentReferencesEnabled, index, 'data-test-subj': dataTestSubj, }: Props) { @@ -194,9 +186,8 @@ export function MessageText({ getPluginDependencies({ contentReferences, contentReferencesVisible, - contentReferencesEnabled, }), - [contentReferences, contentReferencesVisible, contentReferencesEnabled] + [contentReferences, contentReferencesVisible] ); return ( diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts index 61ddceab607cf..28bd8970fbcb3 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts @@ -184,21 +184,6 @@ describe('AlertCountsTool', () => { expect(result).toContain('Citation: {reference(exampleContentReferenceId)}'); }); - it('does not include citations when contentReferencesStore is false', async () => { - const tool: DynamicTool = ALERT_COUNTS_TOOL.getTool({ - alertsIndexPattern, - esClient, - replacements, - request, - ...rest, - contentReferencesStore: undefined, - }) as DynamicTool; - - const result = await tool.func(''); - - expect(result).not.toContain('Citation:'); - }); - it('returns null when the alertsIndexPattern is undefined', () => { const tool = ALERT_COUNTS_TOOL.getTool({ // alertsIndexPattern is undefined diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts index e1150948fde1e..801b4054a8ef3 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts @@ -43,13 +43,11 @@ export const ALERT_COUNTS_TOOL: AssistantTool = { func: async () => { const query = getAlertsCountQuery(alertsIndexPattern); const result = await esClient.search(query); - const alertsCountReference = - contentReferencesStore && - contentReferencesStore.add((p) => securityAlertsPageReference(p.id)); + const alertsCountReference = contentReferencesStore?.add((p) => + securityAlertsPageReference(p.id) + ); - const reference = alertsCountReference - ? `\n${contentReferenceString(alertsCountReference)}` - : ''; + const reference = `\n${contentReferenceString(alertsCountReference)}`; return `${JSON.stringify(result)}${reference}`; }, diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts index 6218c4c441a44..f5f50821abeb5 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts @@ -64,26 +64,5 @@ describe('KnowledgeBaseRetievalTool', () => { expect(result).toContain('citation":"{reference(exampleContentReferenceId)}"'); }); - - it('does not include citations if contentReferenceStore is false', async () => { - const tool = KNOWLEDGE_BASE_RETRIEVAL_TOOL.getTool({ - ...defaultArgs, - contentReferencesStore: undefined, - }) as DynamicStructuredTool; - - getKnowledgeBaseDocumentEntries.mockResolvedValue([ - new Document({ - id: 'exampleId', - pageContent: 'text', - metadata: { - name: 'exampleName', - }, - }), - ] as Document[]); - - const result = await tool.func({ query: 'What is my favourite food' }); - - expect(result).not.toContain('citation'); - }); }); }); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts index 689fb7f189ce8..beddd4efeadb9 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts @@ -58,11 +58,7 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { required: false, }); - if (contentReferencesStore) { - return JSON.stringify(docs.map(enrichDocument(contentReferencesStore))); - } - - return JSON.stringify(docs); + return JSON.stringify(docs.map(enrichDocument(contentReferencesStore))); }, tags: ['knowledge-base'], // TODO: Remove after ZodAny is fixed https://github.com/langchain-ai/langchainjs/blob/main/langchain-core/src/tools.ts @@ -70,9 +66,9 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { }, }; -function enrichDocument(contentReferencesStore: ContentReferencesStore) { +function enrichDocument(contentReferencesStore: ContentReferencesStore | undefined) { return (document: Document>) => { - if (document.id == null) { + if (document.id == null || contentReferencesStore == null) { return document; } const documentId = document.id; diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts index aeda113abc2f0..d2fc1888e2ae8 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts @@ -265,29 +265,6 @@ describe('OpenAndAcknowledgedAlertsTool', () => { expect(result).toContain('Citation,{reference(exampleContentReferenceId)}'); }); - it('does not include citations if content references store is false', async () => { - const tool: DynamicTool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ - alertsIndexPattern, - anonymizationFields, - onNewReplacements: jest.fn(), - replacements, - request, - size: request.body.size, - ...rest, - contentReferencesStore: undefined, - }) as DynamicTool; - - (esClient.search as jest.Mock).mockResolvedValue({ - hits: { - hits: [{ _id: 4 }], - }, - }); - - const result = await tool.func(''); - - expect(result).not.toContain('Citation'); - }); - it('returns null when alertsIndexPattern is undefined', () => { const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ // alertsIndexPattern is undefined diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts index c7ca6972f5476..d73bc266239d5 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts @@ -86,21 +86,20 @@ export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = { }; return JSON.stringify( - result.hits?.hits?.map((x) => { + result.hits?.hits?.map((hit) => { const transformed = transformRawData({ anonymizationFields, currentReplacements: localReplacements, // <-- the latest local replacements getAnonymizedValue, onNewReplacements: localOnNewReplacements, // <-- the local callback - rawData: getRawDataOrDefault(x.fields), + rawData: getRawDataOrDefault(hit.fields), }); - const hitId = x._id; - const citation = - hitId && - contentReferencesStore && - `\nCitation,${contentReferenceBlock( - contentReferencesStore.add((p) => securityAlertReference(p.id, hitId)) - )}`; + + const hitId = hit._id; + const reference = hitId + ? contentReferencesStore?.add((p) => securityAlertReference(p.id, hitId)) + : undefined; + const citation = reference && `\nCitation,${contentReferenceBlock(reference)}`; return `${transformed}${citation ?? ''}`; }) diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts index 8d803f3bbaf25..b4dfd08af5e48 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts @@ -140,38 +140,5 @@ describe('ProductDocumentationTool', () => { }, }); }); - - it('does not include citations if contentReferencesStore is false', async () => { - const tool = PRODUCT_DOCUMENTATION_TOOL.getTool({ - ...defaultArgs, - contentReferencesStore: undefined, - }) as DynamicStructuredTool; - - (retrieveDocumentation as jest.Mock).mockResolvedValue({ - documents: [ - { - title: 'exampleTitle', - url: 'exampleUrl', - content: 'exampleContent', - summarized: false, - }, - ] as RetrieveDocumentationResultDoc[], - }); - - const result = await tool.func({ query: 'What is Kibana Security?', product: 'kibana' }); - - expect(result).toEqual({ - content: { - documents: [ - { - content: 'exampleContent', - title: 'exampleTitle', - url: 'exampleUrl', - summarized: false, - }, - ], - }, - }); - }); }); }); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts index bcfcadf1ca076..30dd2a1ec50cb 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts @@ -75,19 +75,11 @@ export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = { functionCalling: 'auto', }); - if (contentReferencesStore) { - const enrichedDocuments = response.documents.map(enrichDocument(contentReferencesStore)); - - return { - content: { - documents: enrichedDocuments, - }, - }; - } + const enrichedDocuments = response.documents.map(enrichDocument(contentReferencesStore)); return { content: { - documents: response.documents, + documents: enrichedDocuments, }, }; }, @@ -98,11 +90,14 @@ export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = { }; type EnrichedDocument = RetrieveDocumentationResultDoc & { - citation: string; + citation?: string; }; -const enrichDocument = (contentReferencesStore: ContentReferencesStore) => { +const enrichDocument = (contentReferencesStore: ContentReferencesStore | undefined) => { return (document: RetrieveDocumentationResultDoc): EnrichedDocument => { + if (contentReferencesStore == null) { + return document; + } const reference = contentReferencesStore.add((p) => productDocumentationReference(p.id, document.title, document.url) ); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts index 4e7df37d4f904..9800a7f6dd038 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts @@ -50,16 +50,5 @@ describe('SecurityLabsTool', () => { expect(result).toContain('Citation: {reference(exampleContentReferenceId)}'); }); - - it('does not include citations when contentReferencesStore is false', async () => { - const tool = SECURITY_LABS_KNOWLEDGE_BASE_TOOL.getTool({ - ...defaultArgs, - contentReferencesStore: undefined, - }) as DynamicStructuredTool; - - const result = await tool.func({ query: 'What is Kibana Security?', product: 'kibana' }); - - expect(result).not.toContain('Citation:'); - }); }); }); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts index 140cc09a670c6..2faad9ba71c06 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts @@ -51,17 +51,15 @@ export const SECURITY_LABS_KNOWLEDGE_BASE_TOOL: AssistantTool = { query: input.question, }); - const reference = - contentReferencesStore && - contentReferencesStore.add((p) => - knowledgeBaseReference(p.id, 'Elastic Security Labs content', 'securityLabsId') - ); + const reference = contentReferencesStore?.add((p) => + knowledgeBaseReference(p.id, 'Elastic Security Labs content', 'securityLabsId') + ); // TODO: Token pruning const result = JSON.stringify(docs).substring(0, 20000); - const citation = reference ? `\n${contentReferenceString(reference)}` : ''; - return `${result}${citation}`; + const citation = contentReferenceString(reference); + return `${result}\n${citation}`; }, tags: ['security-labs', 'knowledge-base'], // TODO: Remove after ZodAny is fixed https://github.com/langchain-ai/langchainjs/blob/main/langchain-core/src/tools.ts diff --git a/x-pack/solutions/security/plugins/security_solution/server/plugin.ts b/x-pack/solutions/security/plugins/security_solution/server/plugin.ts index a655b777760e1..904b178527723 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/plugin.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/plugin.ts @@ -602,7 +602,6 @@ export class Plugin implements ISecuritySolutionPlugin { plugins.elasticAssistant.registerTools(APP_UI_ID, assistantTools); const features = { assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, - contentReferencesEnabled: config.experimentalFeatures.contentReferencesEnabled, }; plugins.elasticAssistant.registerFeatures(APP_UI_ID, features); plugins.elasticAssistant.registerFeatures('management', features);