Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
8f59b95
Replace polynomial regular expression with constant string manipulation
KDKHD Feb 3, 2025
b03396a
[CI] Auto-commit changed files from 'node scripts/eslint --no-cache -…
kibanamachine Feb 3, 2025
23a2503
Merge branch 'main' into bug/content-references-regex-performance
elasticmachine Feb 3, 2025
edcefc1
Include missing strings in prompts
KDKHD Feb 4, 2025
196a01d
Add test to citations
KDKHD Feb 5, 2025
791b965
[CI] Auto-commit changed files from 'node scripts/eslint --no-cache -…
kibanamachine Feb 5, 2025
b3fec5d
update content reference parser
KDKHD Feb 5, 2025
3778d7d
lint
KDKHD Feb 5, 2025
af99911
[CI] Auto-commit changed files from 'node scripts/eslint --no-cache -…
kibanamachine Feb 5, 2025
647fb00
lint
KDKHD Feb 5, 2025
4e5663e
Merge branch 'bug/content-referencest-regex-performance' of github.co…
KDKHD Feb 5, 2025
a399f4f
[CI] Auto-commit changed files from 'node scripts/eslint --no-cache -…
kibanamachine Feb 5, 2025
dd972ff
test
KDKHD Feb 5, 2025
79ea623
Merge branch 'bug/content-references-regex-performance' of github.com…
KDKHD Feb 5, 2025
92f9d52
Merge branch 'main' into bug/content-references-regex-performance
elasticmachine Feb 5, 2025
4ddac67
Merge branch 'main' into bug/content-references-regex-performance
KDKHD Feb 6, 2025
5ffbcdf
content reference parser - eat empty references to better handle erro…
KDKHD Feb 6, 2025
55ef292
[CI] Auto-commit changed files from 'node scripts/eslint --no-cache -…
kibanamachine Feb 6, 2025
fcfaccd
content reference parser - parse references on new lines correctly
KDKHD Feb 6, 2025
86e9ef3
Update x-pack/solutions/security/plugins/elastic_assistant/server/lib…
KDKHD Feb 6, 2025
e5482f0
code review
KDKHD Feb 6, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
/*
* 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 { removeContentReferences } from './utils';

describe('utils', () => {
it.each([
['this has no content references', 'this has no content references'],
[
'The sky is blue{reference(1234)} and the grass is green{reference(4321)}',
'The sky is blue and the grass is green',
],
['', ''],
['{reference(1234)}', ''],
[' {reference(1234)} ', ' '],
['{reference(1234', '{reference(1234'],
['{reference(1234)', '{reference(1234)'],
['{reference(1234)}{reference(1234)}{reference(1234)}', ''],
['{reference(1234)}reference(1234)}{reference(1234)}', 'reference(1234)}'],
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I observed the following unexpected output while desk testing:

poem

I attempted to reproduce it via a test case, but the following test passes:

    [
      `With a wagging tail and a wet, cold nose,{reference(ccaSI)}
A furry friend, from head to toes.{reference(ccaSI)}
Loyal companion, always near,{reference(ccaSI)}
Chasing squirrels, full of cheer.{reference(ccaSI)}
A paw to hold, a gentle nudge,{reference(ccaSI)}
A furry alarm, a playful judge.{reference(ccaSI)}
From golden retrievers to tiny Chihuahuas,{reference(ccaSI)}
Their love's a gift, that always conquers.{reference(ccaSI)}
So cherish your dog, with all your might,{reference(ccaSI)}
Their love's a beacon, shining bright.{reference(ccaSI)}`,
      `With a wagging tail and a wet, cold nose,
A furry friend, from head to toes.
Loyal companion, always near,
Chasing squirrels, full of cheer.
A paw to hold, a gentle nudge,
A furry alarm, a playful judge.
From golden retrievers to tiny Chihuahuas,
Their love's a gift, that always conquers.
So cherish your dog, with all your might,
Their love's a beacon, shining bright.`,
    ],

The unreplaced references are still visible when the conversation is re-opened, and when the Show citations toggle is clicked, as illustrated by the animated gif below:

loading_poem_convo

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

observation: The poem contains a single entry Knowledge base entry: Favorite Color, however that KB entry does not appear to be related to the poem:

favorite_color

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So there are 2 problems here:

  1. The {reference(...)} strings are not getting parsed correctly. This has been fixed here
  2. Irrelevant references are being added. I will tweak the prompts to fix this in a separate PR.

])('removesContentReferences from "%s"', (input: string, expected: string) => {
const result = removeContentReferences(input);

expect(result).toEqual(expected);
});

// https://github.com/elastic/kibana/security/code-scanning/539
it('removesContentReferences does not run in polynomial time', () => {
const input = `${'{reference('.repeat(100000)}x${')'.repeat(100000)}`;
const startTime = performance.now(); // Start timing

removeContentReferences(input);

const endTime = performance.now(); // End timing
const executionTime = endTime - startTime; // Time in milliseconds

expect(executionTime).toBeLessThan(1000); // Assert under 1 second
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Executes in < 1ms so this won't become a flakey test.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Executes in < 1ms so this won't become a flakey test.

😂 auto reply from CI: hold my 🍺

Jokes aside, to test this test locally, I restored the original implementation of removeContentReferences:

export const removeContentReferences = (content: string) => {
  return content.replaceAll(/\{reference\(.*?\)\}/g, '');
};

and re-ran the test. It failed with:

    ✕ removesContentReferences does not run in polynomial time (41169 ms)

  ● utils › removesContentReferences does not run in polynomial time

    expect(received).toBeLessThan(expected)

    Expected: < 1000
    Received:   41167.3355

      38 |     const executionTime = endTime - startTime; // Time in milliseconds
      39 | 
    > 40 |     expect(executionTime).toBeLessThan(1000); // Assert under 1 second
         |                           ^
      41 |   });
      42 | });
      43 | 

      at Object.toBeLessThan (x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.test.ts:40:27)

After restoring the PR version of removeContentReferences, the test once again passes locally with:

✓ removesContentReferences does not run in polynomial time (2 ms)

});
});
Original file line number Diff line number Diff line change
Expand Up @@ -46,5 +46,28 @@ export const contentReferenceString = (contentReference: ContentReference) => {
* @returns content with content references replaced with ''
*/
export const removeContentReferences = (content: string) => {
return content.replaceAll(/\{reference\(.*?\)\}/g, '');
let result = '';
let i = 0;

while (i < content.length) {
const start = content.indexOf('{reference(', i);
if (start === -1) {
// No more "{reference(" → append the rest of the string
result += content.slice(i);
break;
}

const end = content.indexOf(')}', start);
if (end === -1) {
// If no closing ")}" is found, treat the rest as normal text
result += content.slice(i);
break;
}

// Append everything before "{reference(" and skip the matched part
result += content.slice(i, start);
i = end + 2; // Move index past ")}"
}

return result;
};
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
/*
* 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 {
BEDROCK_SYSTEM_PROMPT,
DEFAULT_SYSTEM_PROMPT,
GEMINI_SYSTEM_PROMPT,
STRUCTURED_SYSTEM_PROMPT,
} from './prompts';

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, 'You are a security analyst', 1],
[GEMINI_SYSTEM_PROMPT, 'You are an assistant', 1],
[BEDROCK_SYSTEM_PROMPT, 'You are a security analyst', 1],
])(
'"%s" contains "%s" %s times',
(prompt: string, containedString: string, expectedCount: number) => {
const regex = new RegExp(containedString, 'g');
expect((prompt.match(regex) || []).length).toBe(expectedCount);
}
);
});
Original file line number Diff line number Diff line change
Expand Up @@ -7,20 +7,15 @@

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 = `In your response, always include citations using the format: \`{reference(...)}\` when information returned by a tool is used. Only use the reference string provided by the tools and do not create reference strings using other information. The reference should be placed after the punctuation marks.
Example citations:
\`\`\`
Your favourite food is pizza. {reference(HMCxq)}
The document was published in 2025. {reference(prSit)}
\`\`\``;
export const INCLUDE_CITATIONS = `\n\nAnnotate your answer with relevant citations. For example: "The sky is blue. {reference(prSit)}"\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}`;
// 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} ${KB_CATCH} {include_citations_prompt_placeholder}`;
export const BEDROCK_SYSTEM_PROMPT = `Use tools as often as possible, as they have access to the latest data and syntax. Never return <thinking> tags in the response, but make sure to include <result> 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_SYSTEM_PROMPT = `${BASE_GEMINI_PROMPT} {include_citations_prompt_placeholder} ${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 <thinking> tags in the response, but make sure to include <result> 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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,9 @@ import type { ContentReferenceNode } from '../content_reference_parser';
const testContentReferenceNode = { contentReferenceId: '1' } as ContentReferenceNode;

jest.mock('../../../../common/lib/kibana', () => ({
useNavigation: jest.fn().mockReturnValue({
navigateTo: jest.fn(),
}),
useKibana: jest.fn().mockReturnValue({
services: {
discover: {
Expand Down Expand Up @@ -154,4 +157,26 @@ describe('contentReferenceComponentFactory', () => {
expect(container).toBeEmptyDOMElement();
expect(screen.queryByText('[1]')).not.toBeInTheDocument();
});

it('renders nothing if contentReferenceId is empty string', async () => {
const Component = contentReferenceComponentFactory({
contentReferences: {
'1': {
id: '1',
type: 'SecurityAlertsPage',
},
} as ContentReferences,
contentReferencesVisible: true,
loading: false,
});

const { container } = render(
<Component
{...({ contentReferenceId: '', contentReferenceCount: -1 } as ContentReferenceNode)}
/>
);

expect(container).toBeEmptyDOMElement();
expect(screen.queryByText('[-1]')).not.toBeInTheDocument();
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ export const contentReferenceComponentFactory = ({
contentReferenceNode: ContentReferenceNode
): React.ReactNode => {
if (!contentReferencesVisible) return null;
if (!contentReferenceNode.contentReferenceId) return null;

const defaultNode = (
<ContentReferenceButton
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ import { EuiLink } from '@elastic/eui';
import type { ContentReferenceNode } from '../content_reference_parser';
import { PopoverReference } from './popover_reference';
import { SECURITY_ALERTS_PAGE_REFERENCE_LABEL } from './translations';
import { useKibana } from '../../../../common/lib/kibana';
import { useNavigateToAlertsPageWithFilters } from '../../../../common/hooks/use_navigate_to_alerts_page_with_filters';
import { FILTER_OPEN, FILTER_ACKNOWLEDGED } from '../../../../../common/types';

interface Props {
contentReferenceNode: ContentReferenceNode;
Expand All @@ -22,17 +23,22 @@ export const SecurityAlertsPageReference: React.FC<Props> = ({
contentReferenceNode,
securityAlertsPageContentReference,
}) => {
const { navigateToApp } = useKibana().services.application;
const openAlertsPageWithFilters = useNavigateToAlertsPageWithFilters();

const onClick = useCallback(
(e: React.MouseEvent) => {
e.preventDefault();
navigateToApp('security', {
path: `alerts`,
openInNewTab: true,
});
openAlertsPageWithFilters(
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for adding these filters!

{
selectedOptions: [FILTER_OPEN, FILTER_ACKNOWLEDGED],
fieldName: 'kibana.alert.workflow_status',
persist: false,
},
true,
'(global:(timerange:(fromStr:now-24h,kind:relative,toStr:now)))'
);
},
[navigateToApp]
[openAlertsPageWithFilters]
);

return (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,62 @@ import type { Parent } from 'mdast';
import { ContentReferenceParser } from './content_reference_parser';

describe('ContentReferenceParser', () => {
it('extracts references from poem', async () => {
const file = unified().use([[markdown, {}], ContentReferenceParser])
.parse(`With a wagging tail and a wet, cold nose,{reference(ccaSI)}
A furry friend, from head to toes.{reference(ccaSI)}
Loyal companion, always near,{reference(ccaSI)}
Chasing squirrels, full of cheer.{reference(ccaSI)}
A paw to hold, a gentle nudge,
{reference(ccaSI)}
A furry alarm, a playful judge.{reference(ccaSI)}
From golden retrievers to tiny Chihuahuas,{reference(ccaSI)}
Their love's a gift, that always conquers.{reference(ccaSI)}
So cherish your dog, with all your might,{reference(ccaSI)}
Their love's a beacon, shining bright.{reference(ccaSI)}`) as Parent;

expect(
(file.children[0] as Parent).children.filter(
(child) => (child.type as string) === 'contentReference'
)
).toHaveLength(10);
expect(file.children[0].children).toEqual(
expect.arrayContaining([
expect.objectContaining({ type: 'text', value: '\nA paw to hold, a gentle nudge,\n' }),
])
);
});

it('extracts reference after linebreak', async () => {
const file = unified().use([[markdown, {}], ContentReferenceParser]).parse(`First line
{reference(FTQJp)}
`) as Parent;

expect(file.children[0].children).toEqual(
expect.arrayContaining([
expect.objectContaining({ type: 'text', value: 'First line\n' }),
expect.objectContaining({ type: 'contentReference' }),
])
);
});

it('eats empty content reference', async () => {
const file = unified()
.use([[markdown, {}], ContentReferenceParser])
.parse('There is an empty content reference.{reference()}') as Parent;

expect(file.children[0].children).toEqual(
expect.arrayContaining([
expect.objectContaining({ type: 'text', value: 'There is an empty content reference.' }),
expect.objectContaining({
type: 'contentReference',
contentReferenceCount: -1,
contentReferenceId: '',
}),
])
);
});

it('eats space preceding content reference', async () => {
const file = unified()
.use([[markdown, {}], ContentReferenceParser])
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@ export interface ContentReferenceNode extends Node {
contentReferenceBlock: ContentReferenceBlock;
}

/**
* Parses `{reference(contentReferenceId)}` or ` {reference(contentReferenceId)}` (notice space prefix) into ContentReferenceNode
*/
/** Matches `{reference` and ` {reference(` */
const REFERENCE_START_PATTERN = '\\u0020?\\{reference';

export const ContentReferenceParser: Plugin = function ContentReferenceParser() {
const Parser = this.Parser;
const tokenizers = Parser.prototype.inlineTokenizers;
Expand All @@ -33,10 +33,9 @@ export const ContentReferenceParser: Plugin = function ContentReferenceParser()
value,
silent
) {
const [match] = value.match(/^\s?{reference/) || [];
if (!match) return false;
const [match] = value.match(new RegExp(`^${REFERENCE_START_PATTERN}`)) || [];

if (value.includes('\n')) return false;
if (!match) return false;

if (value[match.length] !== '(') return false;

Expand Down Expand Up @@ -81,10 +80,6 @@ export const ContentReferenceParser: Plugin = function ContentReferenceParser()
});
}

if (!contentReferenceId) {
return false;
}

if (silent) {
return true;
}
Expand All @@ -95,6 +90,9 @@ export const ContentReferenceParser: Plugin = function ContentReferenceParser()
const contentReferenceBlock: ContentReferenceBlock = `{reference(${contentReferenceId})}`;

const getContentReferenceCount = (id: string) => {
if (!id) {
return -1;
}
if (id in contentReferenceCounts) {
return contentReferenceCounts[id];
}
Expand All @@ -104,18 +102,24 @@ export const ContentReferenceParser: Plugin = function ContentReferenceParser()

const toEat = `${match.startsWith(' ') ? ' ' : ''}${contentReferenceBlock}`;

return eat(toEat)({
const contentReferenceNode: ContentReferenceNode = {
type: 'contentReference',
contentReferenceId,
contentReferenceCount: getContentReferenceCount(contentReferenceId),
contentReferenceBlock,
} as ContentReferenceNode);
};

return eat(toEat)(contentReferenceNode);
};

tokenizeCustomCitation.notInLink = true;

tokenizeCustomCitation.locator = (value, fromIndex) => {
return 1 + (value.substring(fromIndex).match(/\s?{reference/)?.index ?? -2);
const nextIndex = value.substring(fromIndex).match(new RegExp(REFERENCE_START_PATTERN))?.index;
if (nextIndex === undefined) {
return -1;
}
return nextIndex + 1;
};

tokenizers.contentReference = tokenizeCustomCitation;
Expand Down