From 7312e6c234009eb7687b0e2d13d6e512b54bb025 Mon Sep 17 00:00:00 2001 From: Kenneth Kreindler <42113355+KDKHD@users.noreply.github.com> Date: Wed, 29 Jan 2025 01:47:59 +0000 Subject: [PATCH 1/4] [Security Solution] [AI Assistant] security assistant content references (#206683) > [!note] > Planning to merge before the 9.0 feature freeze. Documentation update issue: https://github.com/elastic/security-docs/issues/6473 > [!tip] > ### Tip for the reviewer > As a starting point to review this PR I would suggest reading the section "How does it work (on a high level)" and viewing the hyperlinked code. The linked code covers the main concepts of this feature and the majority of the remaining changes in the PR are related to API schema updates and tests. ## Summary This PR adds citations to the security AI assistant. Citations are produced when tools are used and they are displayed in the LLM response as numbered superscript elements. A label appears when the user hovers over the numbered elements and clicking on the label opens a new tab that displays the cited data. ## How to test: 1. Enables the feature flags: - Set value to `true` [here](https://github.com/elastic/kibana/pull/206683/files#diff-f55be7c50853801c3933b48064ab9cbf0356e2941cd97c4c365be7a6ded9bffdR125) 2. Populate the security knowledge base with some information (e.g. a document, an index, product documentation, the global threat report, etc...). 3. Open the security assistant 4. Ask it the following questions about data in the knowledge base: - What is an elastic search cold tier? - Make sure product docs are installed - What topics are covered in the security lab content? - Ask it a question about one of your knowledge base documents. - Which platform did my most recent alert happen on? - Make sure you have a recent alert ## How does it work (on a high level)? Citations are stored inside the [ContentReferencesStore](https://github.com/elastic/kibana/pull/206683/files#diff-baf03ce192db4f13999748a38b4d920428358a4ffc62527a1d6ac0d9b234f306R17) object. When tools are called, the tools [add citations to the ContentReferencesStore](https://github.com/elastic/kibana/pull/206683/files#diff-5a333fdd9bf864dced06500263577e495c95c9b32c7dae9074090775df542d22R97-R99) and pass the Id of the [ContentReferences back to the LLM](https://github.com/elastic/kibana/pull/206683/files#diff-5a333fdd9bf864dced06500263577e495c95c9b32c7dae9074090775df542d22R102) along side the result of the tool. The LLM can then use those contentReference IDs in its response by forming a response like: ``` The sky is blue {reference(12345)} ``` The web client [parses out the contentReference](https://github.com/elastic/kibana/pull/206683/files#diff-3a5c8305ac899a9e78903b0b60141dd997ba61e87342de2b9ec377165d99cfe6R23) (`{reference(12345)}`) from the assistant message and[ replaces it with the citation react component](https://github.com/elastic/kibana/pull/206683/files#diff-db928fb87a862e3ebf7247baefc418de539f9c0f3fc5134a2ef56f921a52bdcbR125-R129). ### Tools that are cited: Include citations for the following tools: alert_counts_tool -> cites to alerts page knowledge_base_retrieval_tool -> cites knowledge base management page with specific entry pre-filtered open_and_acknowledged_alerts_tool -> cites to specific alert security_labs_tool -> cites knowledge base management page with specific entry pre-filtered knowledge_base indices -> opens ESQL view selecting the particular document used product_documentation -> cites documentation ### Endpoints impacted - POST /internal/elastic_assistant/actions/connector/{connectorId}/_execute - POST /api/security_ai_assistant/chat/complete - GET /api/security_ai_assistant/current_user/conversations/_find - GET /api/security_ai_assistant/current_user/conversations/:id - PUT /api/security_ai_assistant/current_user/conversations/{id} ### Considerations: - One of the main objectives of this feature was to produce in-text citations to create a great user experience. Multiple approaches were tested to do this reliably. Attempts were made to make the LLM return structured JSON containing the citations however this was unreliable with smaller models. Generation post-processing (issuing an additional LLM call to annotate the response with citations) was also explored however this also had limitations as the second LLM call would not contain enough contextual information to create the citations reliably. Eventually, the approach described in the section above was used alongside few shot promoting. - Instead of using the ContentReferencesStore to store citations, the langGraph state could be used to save the citations. I looked at doing this but currently, there are a few blockers in the langgraph API the prevent this. - Lang graph must be updated to @langchain/langgraph>=0.2.31 to get access to the Command type so that tools can update the graph state. - It seems that DynamicStructuredTools do not support the Command type yet. This is something that we can clarify with the langchain team. Once these blockers have been addressed, ContentReferencesStore could easily be refactored to the graph state. - The feature has been put behind a feature flag so we can test during the feature freeze and sync the release of the documentation update. The only thing that is not behind a feature flag is the new anonymization button in the settings menu (don't think it is necessary and it means a lot more code changes are required). On few occasions, you can nudge the LLM a bit more to include citations by appending "Include citations" to your message. ![image](https://github.com/user-attachments/assets/e87b010b-4c29-48c7-8b2b-f17ad1878b8b) Furthermore, the settings menu has been updated to include anonymized values and citation toggles: ![image](https://github.com/user-attachments/assets/efcbabe5-4325-4b6b-b387-84295cb0fb70) ### Checklist Check the PR satisfies following conditions. Reviewers should verify this PR satisfies this list as well. - [X] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md) - [x] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [X] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios - [X] If a plugin configuration key changed, check if it needs to be allowlisted in the cloud and added to the [docker list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker) - [x] This was checked for breaking HTTP API changes, and any breaking changes have been approved by the breaking-change committee. The `release_note:breaking` label should be applied in these situations. - [X] [Flaky Test Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was used on any tests changed - [x] The PR description includes the appropriate Release Notes section, and the correct `release_note:*` label is applied per the [guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) ### Identify risks Does this PR introduce any risks? For example, consider risks like hard to test bugs, performance regression, potential of data loss. Describe the risk, its severity, and mitigation for each identified risk. Invite stakeholders and evaluate how to proceed before merging. - [ ] [See some risk examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx) - [ ] ... ## Release note Adds in-text citations to security solution AI assistant responses. --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Co-authored-by: Elastic Machine Co-authored-by: Patryk Kopycinski (cherry picked from commit 5f888d00802ea84767fb801ca7fde4faabf412fb) # Conflicts: # x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx # x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/index.tsx # x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx # x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx --- oas_docs/output/kibana.serverless.yaml | 122 +++++++ oas_docs/output/kibana.yaml | 122 +++++++ ...sistant_api_2023_10_31.bundled.schema.yaml | 122 +++++++ ...sistant_api_2023_10_31.bundled.schema.yaml | 122 +++++++ .../impl/capabilities/index.ts | 1 + .../content_references_store.mock.ts | 17 + .../content_references_store_factory.test.ts | 54 +++ .../content_references_store_factory.ts | 52 +++ .../prune_content_references.test.ts | 38 ++ .../prune_content_references.ts | 44 +++ .../impl/content_references/index.ts | 21 ++ .../content_references/references/index.ts | 102 ++++++ .../content_references/references/utils.ts | 50 +++ .../impl/content_references/types.ts | 29 ++ .../get_capabilities_route.gen.ts | 1 + .../get_capabilities_route.schema.yaml | 3 + .../conversations/common_attributes.gen.ts | 140 ++++++++ .../common_attributes.schema.yaml | 135 ++++++++ .../kbn-elastic-assistant-common/index.ts | 18 + .../impl/assistant/assistant_body/index.tsx | 3 +- .../get_anonymization_tooltip/index.test.ts | 42 --- .../get_anonymization_tooltip/index.ts | 22 -- .../assistant/assistant_header/index.test.tsx | 156 +-------- .../impl/assistant/assistant_header/index.tsx | 41 +-- .../assistant_header/translations.ts | 39 ++- .../impl/assistant/index.tsx | 40 ++- .../settings_context_menu.tsx | 99 +++++- .../impl/assistant_context/constants.tsx | 2 + .../impl/assistant_context/index.tsx | 38 +- .../impl/assistant_context/types.tsx | 2 + .../index.test.tsx | 55 ++- .../index.tsx | 9 +- .../kbn-elastic-assistant/tsconfig.json | 1 + .../scripts/draw_graph_script.ts | 1 + .../append_conversation_messages.ts | 12 + .../conversations/field_maps_configuration.ts | 15 + .../conversations/helpers.ts | 6 + .../conversations/index.ts | 9 + .../conversations/transforms.ts | 9 + .../conversations/update_conversation.test.ts | 20 ++ .../conversations/update_conversation.ts | 14 + .../knowledge_base/helpers.test.tsx | 40 ++- .../knowledge_base/helpers.ts | 31 +- .../knowledge_base/index.test.ts | 4 + .../knowledge_base/index.ts | 7 + .../server/ai_assistant_service/index.ts | 39 ++- .../server/lib/langchain/executors/types.ts | 8 +- .../graphs/default_assistant_graph/graph.ts | 6 + .../default_assistant_graph/index.test.ts | 2 + .../graphs/default_assistant_graph/index.ts | 15 + .../nodes/run_agent.ts | 20 +- .../graphs/default_assistant_graph/types.ts | 1 + .../server/lib/prompt/prompts.ts | 11 +- .../server/routes/chat/chat_complete_route.ts | 20 +- .../routes/defend_insights/helpers.test.ts | 2 + .../server/routes/defend_insights/helpers.ts | 5 + .../defend_insights/post_defend_insights.ts | 1 + .../server/routes/evaluate/post_evaluate.ts | 10 + .../server/routes/helpers.ts | 18 + .../routes/post_actions_connector_execute.ts | 18 +- .../server/routes/request_context_factory.ts | 3 +- .../routes/user_conversations/find_route.ts | 12 +- .../plugins/elastic_assistant/server/types.ts | 11 +- .../common/experimental_features.ts | 5 + .../assistant/comment_actions/index.test.tsx | 54 +++ .../assistant/comment_actions/index.tsx | 9 +- .../components/content_reference_button.tsx | 39 +++ ...ntent_reference_component_factory.test.tsx | 157 +++++++++ .../content_reference_component_factory.tsx | 91 +++++ .../components/esql_query_reference.tsx | 57 +++ .../knowledge_base_entry_reference.tsx | 48 +++ .../components/popover_reference.tsx | 55 +++ .../product_documentation_reference.tsx | 33 ++ .../components/security_alert_reference.tsx | 46 +++ .../security_alerts_page_reference.tsx | 46 +++ .../components/translations.ts | 29 ++ .../content_reference_parser.test.ts | 325 ++++++++++++++++++ .../content_reference_parser.ts | 123 +++++++ .../public/assistant/get_comments/index.tsx | 8 + .../assistant/get_comments/stream/index.tsx | 10 + .../get_comments/stream/message_text.tsx | 59 +++- .../alert_counts/alert_counts_tool.test.ts | 41 +++ .../tools/alert_counts/alert_counts_tool.ts | 13 +- .../tools/esql/nl_to_esql_tool.test.ts | 3 + .../knowledge_base_retrieval_tool.test.ts | 89 +++++ .../knowledge_base_retrieval_tool.ts | 29 +- .../open_and_acknowledged_alerts_tool.test.ts | 66 ++++ .../open_and_acknowledged_alerts_tool.ts | 20 +- .../product_documentation_tool.test.ts | 86 ++++- .../product_documentation_tool.ts | 35 +- .../security_labs/security_labs_tool.test.ts | 65 ++++ .../tools/security_labs/security_labs_tool.ts | 15 +- .../security_solution/server/plugin.ts | 1 + .../cypress/screens/ai_assistant.ts | 1 - .../cypress/tasks/assistant.ts | 2 - 95 files changed, 3433 insertions(+), 339 deletions(-) create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock.ts create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.test.ts create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.ts create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.ts create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/index.ts create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/index.ts create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.ts create mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/types.ts delete mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/get_anonymization_tooltip/index.test.ts delete mode 100644 x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/get_anonymization_tooltip/index.ts create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/comment_actions/index.test.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_button.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_component_factory.test.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_component_factory.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/esql_query_reference.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/knowledge_base_entry_reference.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/popover_reference.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/product_documentation_reference.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/security_alert_reference.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/security_alerts_page_reference.tsx create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/translations.ts create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/content_reference_parser.test.ts create mode 100644 x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/content_reference_parser.ts create mode 100644 x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts create mode 100644 x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts diff --git a/oas_docs/output/kibana.serverless.yaml b/oas_docs/output/kibana.serverless.yaml index 1a0ecf8800cdf..5a5fc8bdc0fdb 100644 --- a/oas_docs/output/kibana.serverless.yaml +++ b/oas_docs/output/kibana.serverless.yaml @@ -40006,6 +40006,19 @@ components: required: - connectorId - actionTypeId + Security_AI_Assistant_API_BaseContentReference: + description: The basis of a content reference + type: object + properties: + id: + description: Id of the content reference + type: string + type: + description: Type of the content reference + type: string + required: + - id + - type Security_AI_Assistant_API_BulkCrudActionSummary: type: object properties: @@ -40077,6 +40090,17 @@ components: - user - assistant type: string + Security_AI_Assistant_API_ContentReferences: + additionalProperties: + oneOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_KnowledgeBaseEntryContentReference' + - $ref: '#/components/schemas/Security_AI_Assistant_API_SecurityAlertContentReference' + - $ref: '#/components/schemas/Security_AI_Assistant_API_SecurityAlertsPageContentReference' + - $ref: '#/components/schemas/Security_AI_Assistant_API_ProductDocumentationContentReference' + - $ref: '#/components/schemas/Security_AI_Assistant_API_EsqlContentReference' + additionalProperties: false + description: A union of all content reference types + type: object Security_AI_Assistant_API_ConversationCategory: description: The conversation category. enum: @@ -40315,6 +40339,26 @@ components: required: - id - $ref: '#/components/schemas/Security_AI_Assistant_API_DocumentEntryCreateFields' + Security_AI_Assistant_API_EsqlContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + label: + description: Label of the query + type: string + query: + description: An ESQL query + type: string + type: + enum: + - EsqlQuery + type: string + required: + - type + - query + - label + description: References an ESQL query Security_AI_Assistant_API_FindAnonymizationFieldsSortField: enum: - created_at @@ -40546,6 +40590,26 @@ components: - skipped - succeeded - total + Security_AI_Assistant_API_KnowledgeBaseEntryContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + knowledgeBaseEntryId: + description: Id of the Knowledge Base Entry + type: string + knowledgeBaseEntryName: + description: Name of the knowledge base entry + type: string + type: + enum: + - KnowledgeBaseEntry + type: string + required: + - type + - knowledgeBaseEntryId + - knowledgeBaseEntryName + description: References a knowledge base entry Security_AI_Assistant_API_KnowledgeBaseEntryCreateProps: anyOf: - $ref: '#/components/schemas/Security_AI_Assistant_API_DocumentEntryCreateFields' @@ -40604,6 +40668,9 @@ components: isError: description: Is error message. type: boolean + metadata: + $ref: '#/components/schemas/Security_AI_Assistant_API_MessageMetadata' + description: metadata reader: $ref: '#/components/schemas/Security_AI_Assistant_API_Reader' description: Message content. @@ -40623,6 +40690,13 @@ components: Security_AI_Assistant_API_MessageData: additionalProperties: true type: object + Security_AI_Assistant_API_MessageMetadata: + description: Message metadata + type: object + properties: + contentReferences: + $ref: '#/components/schemas/Security_AI_Assistant_API_ContentReferences' + description: Data refered to by the message content. Security_AI_Assistant_API_MessageRole: description: Message role. enum: @@ -40686,6 +40760,26 @@ components: - message - status_code - prompts + Security_AI_Assistant_API_ProductDocumentationContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + title: + description: Title of the documentation + type: string + type: + enum: + - ProductDocumentation + type: string + url: + description: URL to the documentation + type: string + required: + - type + - title + - url + description: References the product documentation Security_AI_Assistant_API_PromptCreateProps: type: object properties: @@ -40897,6 +40991,34 @@ components: - createdBy - updatedAt - updatedBy + Security_AI_Assistant_API_SecurityAlertContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + alertId: + description: ID of the Alert + type: string + type: + enum: + - SecurityAlert + type: string + required: + - type + - alertId + description: References a security alert + Security_AI_Assistant_API_SecurityAlertsPageContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + type: + enum: + - SecurityAlertsPage + type: string + required: + - type + description: References the security alerts page Security_AI_Assistant_API_SortOrder: enum: - asc diff --git a/oas_docs/output/kibana.yaml b/oas_docs/output/kibana.yaml index 00fd588f6890c..0252c9abd5253 100644 --- a/oas_docs/output/kibana.yaml +++ b/oas_docs/output/kibana.yaml @@ -28196,6 +28196,19 @@ components: required: - connectorId - actionTypeId + Security_AI_Assistant_API_BaseContentReference: + description: The basis of a content reference + type: object + properties: + id: + description: Id of the content reference + type: string + type: + description: Type of the content reference + type: string + required: + - id + - type Security_AI_Assistant_API_BulkCrudActionSummary: type: object properties: @@ -28267,6 +28280,17 @@ components: - user - assistant type: string + Security_AI_Assistant_API_ContentReferences: + additionalProperties: + oneOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_KnowledgeBaseEntryContentReference' + - $ref: '#/components/schemas/Security_AI_Assistant_API_SecurityAlertContentReference' + - $ref: '#/components/schemas/Security_AI_Assistant_API_SecurityAlertsPageContentReference' + - $ref: '#/components/schemas/Security_AI_Assistant_API_ProductDocumentationContentReference' + - $ref: '#/components/schemas/Security_AI_Assistant_API_EsqlContentReference' + additionalProperties: false + description: A union of all content reference types + type: object Security_AI_Assistant_API_ConversationCategory: description: The conversation category. enum: @@ -28505,6 +28529,26 @@ components: required: - id - $ref: '#/components/schemas/Security_AI_Assistant_API_DocumentEntryCreateFields' + Security_AI_Assistant_API_EsqlContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + label: + description: Label of the query + type: string + query: + description: An ESQL query + type: string + type: + enum: + - EsqlQuery + type: string + required: + - type + - query + - label + description: References an ESQL query Security_AI_Assistant_API_FindAnonymizationFieldsSortField: enum: - created_at @@ -28736,6 +28780,26 @@ components: - skipped - succeeded - total + Security_AI_Assistant_API_KnowledgeBaseEntryContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + knowledgeBaseEntryId: + description: Id of the Knowledge Base Entry + type: string + knowledgeBaseEntryName: + description: Name of the knowledge base entry + type: string + type: + enum: + - KnowledgeBaseEntry + type: string + required: + - type + - knowledgeBaseEntryId + - knowledgeBaseEntryName + description: References a knowledge base entry Security_AI_Assistant_API_KnowledgeBaseEntryCreateProps: anyOf: - $ref: '#/components/schemas/Security_AI_Assistant_API_DocumentEntryCreateFields' @@ -28794,6 +28858,9 @@ components: isError: description: Is error message. type: boolean + metadata: + $ref: '#/components/schemas/Security_AI_Assistant_API_MessageMetadata' + description: metadata reader: $ref: '#/components/schemas/Security_AI_Assistant_API_Reader' description: Message content. @@ -28813,6 +28880,13 @@ components: Security_AI_Assistant_API_MessageData: additionalProperties: true type: object + Security_AI_Assistant_API_MessageMetadata: + description: Message metadata + type: object + properties: + contentReferences: + $ref: '#/components/schemas/Security_AI_Assistant_API_ContentReferences' + description: Data refered to by the message content. Security_AI_Assistant_API_MessageRole: description: Message role. enum: @@ -28876,6 +28950,26 @@ components: - message - status_code - prompts + Security_AI_Assistant_API_ProductDocumentationContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + title: + description: Title of the documentation + type: string + type: + enum: + - ProductDocumentation + type: string + url: + description: URL to the documentation + type: string + required: + - type + - title + - url + description: References the product documentation Security_AI_Assistant_API_PromptCreateProps: type: object properties: @@ -29087,6 +29181,34 @@ components: - createdBy - updatedAt - updatedBy + Security_AI_Assistant_API_SecurityAlertContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + alertId: + description: ID of the Alert + type: string + type: + enum: + - SecurityAlert + type: string + required: + - type + - alertId + description: References a security alert + Security_AI_Assistant_API_SecurityAlertsPageContentReference: + allOf: + - $ref: '#/components/schemas/Security_AI_Assistant_API_BaseContentReference' + - type: object + properties: + type: + enum: + - SecurityAlertsPage + type: string + required: + - type + description: References the security alerts page Security_AI_Assistant_API_SortOrder: enum: - asc diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml index aeb6c63215dc2..7e8583ac718ea 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml @@ -1056,6 +1056,19 @@ components: required: - connectorId - actionTypeId + BaseContentReference: + description: The basis of a content reference + type: object + properties: + id: + description: Id of the content reference + type: string + type: + description: Type of the content reference + type: string + required: + - id + - type BulkCrudActionSummary: type: object properties: @@ -1127,6 +1140,17 @@ components: - user - assistant type: string + ContentReferences: + additionalProperties: + oneOf: + - $ref: '#/components/schemas/KnowledgeBaseEntryContentReference' + - $ref: '#/components/schemas/SecurityAlertContentReference' + - $ref: '#/components/schemas/SecurityAlertsPageContentReference' + - $ref: '#/components/schemas/ProductDocumentationContentReference' + - $ref: '#/components/schemas/EsqlContentReference' + additionalProperties: false + description: A union of all content reference types + type: object ConversationCategory: description: The conversation category. enum: @@ -1375,6 +1399,26 @@ components: required: - id - $ref: '#/components/schemas/DocumentEntryCreateFields' + EsqlContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + label: + description: Label of the query + type: string + query: + description: An ESQL query + type: string + type: + enum: + - EsqlQuery + type: string + required: + - type + - query + - label + description: References an ESQL query FindAnonymizationFieldsSortField: enum: - created_at @@ -1620,6 +1664,26 @@ components: - skipped - succeeded - total + KnowledgeBaseEntryContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + knowledgeBaseEntryId: + description: Id of the Knowledge Base Entry + type: string + knowledgeBaseEntryName: + description: Name of the knowledge base entry + type: string + type: + enum: + - KnowledgeBaseEntry + type: string + required: + - type + - knowledgeBaseEntryId + - knowledgeBaseEntryName + description: References a knowledge base entry KnowledgeBaseEntryCreateProps: anyOf: - $ref: '#/components/schemas/DocumentEntryCreateFields' @@ -1678,6 +1742,9 @@ components: isError: description: Is error message. type: boolean + metadata: + $ref: '#/components/schemas/MessageMetadata' + description: metadata reader: $ref: '#/components/schemas/Reader' description: Message content. @@ -1697,6 +1764,13 @@ components: MessageData: additionalProperties: true type: object + MessageMetadata: + description: Message metadata + type: object + properties: + contentReferences: + $ref: '#/components/schemas/ContentReferences' + description: Data refered to by the message content. MessageRole: description: Message role. enum: @@ -1760,6 +1834,26 @@ components: - message - status_code - prompts + ProductDocumentationContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + title: + description: Title of the documentation + type: string + type: + enum: + - ProductDocumentation + type: string + url: + description: URL to the documentation + type: string + required: + - type + - title + - url + description: References the product documentation PromptCreateProps: type: object properties: @@ -1971,6 +2065,34 @@ components: - createdBy - updatedAt - updatedBy + SecurityAlertContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + alertId: + description: ID of the Alert + type: string + type: + enum: + - SecurityAlert + type: string + required: + - type + - alertId + description: References a security alert + SecurityAlertsPageContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + type: + enum: + - SecurityAlertsPage + type: string + required: + - type + description: References the security alerts page SortOrder: enum: - asc diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml index cfc5c2f310dc9..5c56d8a871c01 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml @@ -1056,6 +1056,19 @@ components: required: - connectorId - actionTypeId + BaseContentReference: + description: The basis of a content reference + type: object + properties: + id: + description: Id of the content reference + type: string + type: + description: Type of the content reference + type: string + required: + - id + - type BulkCrudActionSummary: type: object properties: @@ -1127,6 +1140,17 @@ components: - user - assistant type: string + ContentReferences: + additionalProperties: + oneOf: + - $ref: '#/components/schemas/KnowledgeBaseEntryContentReference' + - $ref: '#/components/schemas/SecurityAlertContentReference' + - $ref: '#/components/schemas/SecurityAlertsPageContentReference' + - $ref: '#/components/schemas/ProductDocumentationContentReference' + - $ref: '#/components/schemas/EsqlContentReference' + additionalProperties: false + description: A union of all content reference types + type: object ConversationCategory: description: The conversation category. enum: @@ -1375,6 +1399,26 @@ components: required: - id - $ref: '#/components/schemas/DocumentEntryCreateFields' + EsqlContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + label: + description: Label of the query + type: string + query: + description: An ESQL query + type: string + type: + enum: + - EsqlQuery + type: string + required: + - type + - query + - label + description: References an ESQL query FindAnonymizationFieldsSortField: enum: - created_at @@ -1620,6 +1664,26 @@ components: - skipped - succeeded - total + KnowledgeBaseEntryContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + knowledgeBaseEntryId: + description: Id of the Knowledge Base Entry + type: string + knowledgeBaseEntryName: + description: Name of the knowledge base entry + type: string + type: + enum: + - KnowledgeBaseEntry + type: string + required: + - type + - knowledgeBaseEntryId + - knowledgeBaseEntryName + description: References a knowledge base entry KnowledgeBaseEntryCreateProps: anyOf: - $ref: '#/components/schemas/DocumentEntryCreateFields' @@ -1678,6 +1742,9 @@ components: isError: description: Is error message. type: boolean + metadata: + $ref: '#/components/schemas/MessageMetadata' + description: metadata reader: $ref: '#/components/schemas/Reader' description: Message content. @@ -1697,6 +1764,13 @@ components: MessageData: additionalProperties: true type: object + MessageMetadata: + description: Message metadata + type: object + properties: + contentReferences: + $ref: '#/components/schemas/ContentReferences' + description: Data refered to by the message content. MessageRole: description: Message role. enum: @@ -1760,6 +1834,26 @@ components: - message - status_code - prompts + ProductDocumentationContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + title: + description: Title of the documentation + type: string + type: + enum: + - ProductDocumentation + type: string + url: + description: URL to the documentation + type: string + required: + - type + - title + - url + description: References the product documentation PromptCreateProps: type: object properties: @@ -1971,6 +2065,34 @@ components: - createdBy - updatedAt - updatedBy + SecurityAlertContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + alertId: + description: ID of the Alert + type: string + type: + enum: + - SecurityAlert + type: string + required: + - type + - alertId + description: References a security alert + SecurityAlertsPageContentReference: + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + properties: + type: + enum: + - SecurityAlertsPage + type: string + required: + - type + description: References the security alerts page SortOrder: enum: - asc 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 35f4b88ef7174..adf474a991dbe 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 @@ -22,4 +22,5 @@ export const defaultAssistantFeatures = Object.freeze({ assistantModelEvaluation: false, attackDiscoveryAlertFiltering: false, defendInsights: false, + contentReferencesEnabled: false, }); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock.ts new file mode 100644 index 0000000000000..bbc19d0566d70 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock.ts @@ -0,0 +1,17 @@ +/* + * 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 { ContentReferencesStore } from '../../types'; + +export const contentReferencesStoreFactoryMock: () => ContentReferencesStore = jest + .fn() + .mockReturnValue({ + add: jest.fn().mockImplementation((creator: Parameters[0]) => { + return creator({ id: 'exampleContentReferenceId' }); + }), + getStore: jest.fn().mockReturnValue({}), + }); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.test.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.test.ts new file mode 100644 index 0000000000000..c3e848ee66d5c --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.test.ts @@ -0,0 +1,54 @@ +/* + * 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 { contentReferencesStoreFactory } from './content_references_store_factory'; +import { securityAlertsPageReference } from '../references'; +import { ContentReferencesStore } from '../types'; + +describe('contentReferencesStoreFactory', () => { + let contentReferencesStore: ContentReferencesStore; + beforeEach(() => { + contentReferencesStore = contentReferencesStoreFactory(); + }); + + it('adds multiple content reference', async () => { + const alertsPageReference1 = contentReferencesStore.add((p) => + securityAlertsPageReference(p.id) + ); + const alertsPageReference2 = contentReferencesStore.add((p) => + securityAlertsPageReference(p.id) + ); + const alertsPageReference3 = contentReferencesStore.add((p) => + securityAlertsPageReference(p.id) + ); + + const store = contentReferencesStore.getStore(); + + const keys = Object.keys(store); + + expect(keys.length).toEqual(3); + expect(store[alertsPageReference1.id]).toEqual(alertsPageReference1); + expect(store[alertsPageReference2.id]).toEqual(alertsPageReference2); + expect(store[alertsPageReference3.id]).toEqual(alertsPageReference3); + }); + + it('referenceIds are unique', async () => { + const numberOfReferencesToCreate = 50; + + const referenceIds = new Set( + [...new Array(numberOfReferencesToCreate)] + .map(() => contentReferencesStore.add((p) => securityAlertsPageReference(p.id))) + .map((alertsPageReference) => alertsPageReference.id) + ); + + const store = contentReferencesStore.getStore(); + const keys = Object.keys(store); + + expect(referenceIds.size).toEqual(numberOfReferencesToCreate); + expect(keys.length).toEqual(numberOfReferencesToCreate); + }); +}); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.ts new file mode 100644 index 0000000000000..b1b78263a31f3 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.ts @@ -0,0 +1,52 @@ +/* + * 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 { ContentReference } from '../../schemas'; +import { ContentReferencesStore } from '../types'; + +const CONTENT_REFERENCE_ID_ALPHABET = + 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'; + +/** + * Creates a new ContentReferencesStore used for storing references (also known as citations) + */ +export const contentReferencesStoreFactory: () => ContentReferencesStore = () => { + const store = new Map(); + + const add: ContentReferencesStore['add'] = (creator) => { + const entry = creator({ id: generateUnsecureId() }); + store.set(entry.id, entry); + return entry; + }; + + const getStore: ContentReferencesStore['getStore'] = () => { + return Object.fromEntries(store); + }; + + /** + * Generates an ID that does not exist in the store yet. This is not cryptographically secure. + * @param size Size of ID to generate + * @returns an unsecure Id + */ + const generateUnsecureId = (size = 5): string => { + let id = ''; + for (let i = 0; i < size; i++) { + id += CONTENT_REFERENCE_ID_ALPHABET.charAt( + Math.floor(Math.random() * CONTENT_REFERENCE_ID_ALPHABET.length) + ); + } + if (store.has(id)) { + return generateUnsecureId(size + 1); + } + return id; + }; + + return { + add, + getStore, + }; +}; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts new file mode 100644 index 0000000000000..5f6184144315c --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts @@ -0,0 +1,38 @@ +/* + * 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 { pruneContentReferences } from './prune_content_references'; +import { securityAlertsPageReference } from '../references'; +import { contentReferenceBlock } from '../references/utils'; +import { ContentReferencesStore } from '../types'; +import { contentReferencesStoreFactory } from './content_references_store_factory'; + +describe('pruneContentReferences', () => { + let contentReferencesStore: ContentReferencesStore; + beforeEach(() => { + contentReferencesStore = contentReferencesStoreFactory(); + }); + + it('prunes content references correctly', async () => { + const alertsPageReference1 = contentReferencesStore.add((p) => + securityAlertsPageReference(p.id) + ); + const alertsPageReference2 = contentReferencesStore.add((p) => + securityAlertsPageReference(p.id) + ); + contentReferencesStore.add((p) => securityAlertsPageReference(p.id)); // this one should get pruned + + const content = `Example ${contentReferenceBlock( + alertsPageReference1 + )} example ${contentReferenceBlock(alertsPageReference2)}`; + + const prunedContentReferences = pruneContentReferences(content, contentReferencesStore); + + const keys = Object.keys(prunedContentReferences!); + expect(keys.sort()).toEqual([alertsPageReference1.id, alertsPageReference2.id].sort()); + }); +}); 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 new file mode 100644 index 0000000000000..887ccf26bc8a6 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.ts @@ -0,0 +1,44 @@ +/* + * 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 { ContentReferences, ContentReference } from '../../schemas'; +import { getContentReferenceId } from '../references/utils'; +import { ContentReferencesStore, ContentReferenceBlock } from '../types'; + +/** + * Returnes a pruned copy of the ContentReferencesStore. + * @param content The content that may contain references to data within the ContentReferencesStore. + * @param contentReferencesStore The ContentReferencesStore contain the contentReferences. + * @returns a new record only containing the ContentReferences that are referenced to by the content. + */ +export const pruneContentReferences = ( + content: string, + contentReferencesStore: ContentReferencesStore +): ContentReferences | undefined => { + 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]; + const referenceId = getContentReferenceId(referenceElement as ContentReferenceBlock); + 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/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/index.ts new file mode 100644 index 0000000000000..ca9afcbb76bbe --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/index.ts @@ -0,0 +1,21 @@ +/* + * 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. + */ + +export { contentReferencesStoreFactory } from './content_references_store/content_references_store_factory'; +export { pruneContentReferences } from './content_references_store/prune_content_references'; +export { + securityAlertReference, + knowledgeBaseReference, + securityAlertsPageReference, + productDocumentationReference, + esqlQueryReference, +} from './references'; +export { + contentReferenceString, + contentReferenceBlock, + removeContentReferences, +} from './references/utils'; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/index.ts new file mode 100644 index 0000000000000..d44d33fdc28e6 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/index.ts @@ -0,0 +1,102 @@ +/* + * 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 { + SecurityAlertContentReference, + SecurityAlertsPageContentReference, + KnowledgeBaseEntryContentReference, + ProductDocumentationContentReference, + EsqlContentReference, +} from '../../schemas'; +import { ContentReferenceId } from '../types'; + +/** + * Generates a contentReference for the alerts count tool. + * @param id id of the contentReference + * @returns AlertsCountReference + */ +export const securityAlertsPageReference = (id: string): SecurityAlertsPageContentReference => { + return { + type: 'SecurityAlertsPage', + id, + }; +}; + +/** + * Generates a contentReference for when a specific alert is referenced. + * @param id id of the contentReference + * @param alertId id of the alert that is referenced + * @returns AlertReference + */ +export const securityAlertReference = ( + id: ContentReferenceId, + alertId: string +): SecurityAlertContentReference => { + return { + type: 'SecurityAlert', + id, + alertId, + }; +}; + +/** + * Generates a contentReference for when a knowledge base entry is referenced. + * @param id id of the contentReference + * @param knowledgeBaseEntryName name of the knowledge base entry + * @param knowledgeBaseEntryId id of the knowledge base entry + * @returns KnowledgeBaseReference + */ +export const knowledgeBaseReference = ( + id: ContentReferenceId, + knowledgeBaseEntryName: string, + knowledgeBaseEntryId: string +): KnowledgeBaseEntryContentReference => { + return { + type: 'KnowledgeBaseEntry', + id, + knowledgeBaseEntryName, + knowledgeBaseEntryId, + }; +}; + +/** + * Generates a contentReference for when a ESQL query is referenced. + * @param id id of the contentReference + * @param query the ESQL query + * @param label content reference label + * @returns KnowledgeBaseReference + */ +export const esqlQueryReference = ( + id: ContentReferenceId, + query: string, + label: string +): EsqlContentReference => { + return { + type: 'EsqlQuery', + id, + label, + query, + }; +}; + +/** + * Generates a contentReference for the alerts count tool. + * @param id id of the contentReference + * @returns AlertsCountReference + */ +export const productDocumentationReference = ( + id: ContentReferenceId, + title: string, + url: string +): ProductDocumentationContentReference => { + return { + type: 'ProductDocumentation', + id, + title, + url, + }; +}; 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 new file mode 100644 index 0000000000000..bfb64b13a9612 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.ts @@ -0,0 +1,50 @@ +/* + * 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 { ContentReference } from '../../schemas'; +import { ContentReferenceBlock, ContentReferenceId } from '../types'; + +/** + * Returns "Arid2" from "{reference(Arid2)}" + * @param contentReference A ContentReferenceBlock + * @returns ContentReferenceId + */ +export const getContentReferenceId = ( + contentReferenceBlock: ContentReferenceBlock +): ContentReferenceId => { + return contentReferenceBlock.replace('{reference(', '').replace(')}', ''); +}; + +/** + * Returns a contentReferenceBlock for a given ContentReference. A ContentReferenceBlock may be provided + * to an LLM alongside grounding documents allowing the LLM to reference the documents in its output. + * @param contentReference A ContentReference + * @returns ContentReferenceBlock + */ +export const contentReferenceBlock = ( + contentReference: ContentReference +): ContentReferenceBlock => { + return `{reference(${contentReference.id})}`; +}; + +/** + * Simplifies passing a contentReferenceBlock alongside grounding documents. + * @param contentReference A ContentReference + * @returns the string: `Reference: ` + */ +export const contentReferenceString = (contentReference: ContentReference) => { + return `Citation: ${contentReferenceBlock(contentReference)}` as const; +}; + +/** + * Removed content references from conent. + * @param content content to remove content references from + * @returns content with content references replaced with '' + */ +export const removeContentReferences = (content: string) => { + return content.replaceAll(/\{reference\(.*?\)\}/g, ''); +}; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/types.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/types.ts new file mode 100644 index 0000000000000..c6d99d7d03a17 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/types.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 { ContentReference, ContentReferences } from '../schemas'; + +export type ContentReferenceId = string; +export type ContentReferenceTypes = ContentReference['type']; +export type ContentReferenceBlock = `{reference(${string})}`; + +export interface ContentReferencesStore { + /** + * Adds a content reference into the ContentReferencesStore. + * @param generator A function that returns a new ContentReference. + * @param generator.params Generator parameters that may be used to generate a new ContentReference. + * @param generator.params.id An ID that is guaranteed to not exist in the store. Intended to be used as the Id of the ContentReference but not required. + * @returns the new ContentReference + */ + add: (generator: (params: { id: ContentReferenceId }) => T) => T; + + /** + * Used to read the content reference store. + * @returns a record that contains all of the ContentReference that have been added . + */ + getStore: () => ContentReferences; +} 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 5fc79ab1ce132..78ee4a6c3e605 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 @@ -20,5 +20,6 @@ export type GetCapabilitiesResponse = z.infer; export const GetCapabilitiesResponse = z.object({ assistantModelEvaluation: z.boolean(), attackDiscoveryAlertFiltering: 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 e1861f01a4d93..684ff6f020793 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 @@ -24,11 +24,14 @@ paths: type: boolean attackDiscoveryAlertFiltering: type: boolean + contentReferencesEnabled: + type: boolean defendInsights: type: boolean required: - assistantModelEvaluation - attackDiscoveryAlertFiltering + - contentReferencesEnabled - defendInsights '400': description: Generic Error diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts index 1dad26e1628db..82a76f850e3bb 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts @@ -33,6 +33,142 @@ export const TraceData = z.object({ traceId: z.string().optional(), }); +/** + * The basis of a content reference + */ +export type BaseContentReference = z.infer; +export const BaseContentReference = z.object({ + /** + * Id of the content reference + */ + id: z.string(), + /** + * Type of the content reference + */ + type: z.string(), +}); + +/** + * References a knowledge base entry + */ +export type KnowledgeBaseEntryContentReference = z.infer; +export const KnowledgeBaseEntryContentReference = BaseContentReference.merge( + z.object({ + type: z.literal('KnowledgeBaseEntry'), + /** + * Id of the Knowledge Base Entry + */ + knowledgeBaseEntryId: z.string(), + /** + * Name of the knowledge base entry + */ + knowledgeBaseEntryName: z.string(), + }) +); + +/** + * References an ESQL query + */ +export type EsqlContentReference = z.infer; +export const EsqlContentReference = BaseContentReference.merge( + z.object({ + type: z.literal('EsqlQuery'), + /** + * An ESQL query + */ + query: z.string(), + /** + * Label of the query + */ + label: z.string(), + }) +); + +/** + * References a security alert + */ +export type SecurityAlertContentReference = z.infer; +export const SecurityAlertContentReference = BaseContentReference.merge( + z.object({ + type: z.literal('SecurityAlert'), + /** + * ID of the Alert + */ + alertId: z.string(), + }) +); + +/** + * References the security alerts page + */ +export type SecurityAlertsPageContentReference = z.infer; +export const SecurityAlertsPageContentReference = BaseContentReference.merge( + z.object({ + type: z.literal('SecurityAlertsPage'), + }) +); + +/** + * References the product documentation + */ +export type ProductDocumentationContentReference = z.infer< + typeof ProductDocumentationContentReference +>; +export const ProductDocumentationContentReference = BaseContentReference.merge( + z.object({ + type: z.literal('ProductDocumentation'), + /** + * Title of the documentation + */ + title: z.string(), + /** + * URL to the documentation + */ + url: z.string(), + }) +); + +/** + * A content reference + */ +export const ContentReferenceInternal = z.union([ + KnowledgeBaseEntryContentReference, + SecurityAlertContentReference, + SecurityAlertsPageContentReference, + ProductDocumentationContentReference, + EsqlContentReference, +]); + +export type ContentReference = z.infer; +export const ContentReference = ContentReferenceInternal as z.ZodType; + +/** + * A union of all content reference types + */ +export type ContentReferences = z.infer; +export const ContentReferences = z + .object({}) + .catchall( + z.union([ + KnowledgeBaseEntryContentReference, + SecurityAlertContentReference, + SecurityAlertsPageContentReference, + ProductDocumentationContentReference, + EsqlContentReference, + ]) + ); + +/** + * Message metadata + */ +export type MessageMetadata = z.infer; +export const MessageMetadata = z.object({ + /** + * Data refered to by the message content. + */ + contentReferences: ContentReferences.optional(), +}); + /** * Replacements object used to anonymize/deanomymize messsages */ @@ -103,6 +239,10 @@ export const Message = z.object({ * trace Data */ traceData: TraceData.optional(), + /** + * metadata + */ + metadata: MessageMetadata.optional(), }); export type ApiConfig = z.infer; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml index 20423236f7423..b651ddd3ca660 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml @@ -18,6 +18,138 @@ components: type: string description: Could be any string, not necessarily a UUID + BaseContentReference: + type: object + description: The basis of a content reference + required: + - 'id' + - 'type' + properties: + id: + type: string + description: Id of the content reference + type: + type: string + description: Type of the content reference + + KnowledgeBaseEntryContentReference: + description: References a knowledge base entry + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + required: + - 'type' + - 'knowledgeBaseEntryId' + - 'knowledgeBaseEntryName' + properties: + type: + type: string + enum: [KnowledgeBaseEntry] + knowledgeBaseEntryId: + description: Id of the Knowledge Base Entry + type: string + knowledgeBaseEntryName: + description: Name of the knowledge base entry + type: string + + EsqlContentReference: + description: References an ESQL query + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + required: + - 'type' + - 'query' + - 'label' + properties: + type: + type: string + enum: [EsqlQuery] + query: + description: An ESQL query + type: string + label: + description: Label of the query + type: string + + SecurityAlertContentReference: + description: References a security alert + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + required: + - 'type' + - 'alertId' + properties: + type: + type: string + enum: [SecurityAlert] + alertId: + description: ID of the Alert + type: string + + SecurityAlertsPageContentReference: + description: References the security alerts page + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + required: + - 'type' + properties: + type: + type: string + enum: [SecurityAlertsPage] + + ProductDocumentationContentReference: + description: References the product documentation + allOf: + - $ref: '#/components/schemas/BaseContentReference' + - type: object + required: + - 'type' + - 'title' + - 'url' + properties: + type: + type: string + enum: [ProductDocumentation] + title: + description: Title of the documentation + type: string + url: + description: URL to the documentation + type: string + + ContentReference: + description: A content reference + oneOf: + - $ref: '#/components/schemas/KnowledgeBaseEntryContentReference' + - $ref: '#/components/schemas/SecurityAlertContentReference' + - $ref: '#/components/schemas/SecurityAlertsPageContentReference' + - $ref: '#/components/schemas/ProductDocumentationContentReference' + - $ref: '#/components/schemas/EsqlContentReference' + additionalProperties: false + + ContentReferences: + description: A union of all content reference types + additionalProperties: + oneOf: + - $ref: '#/components/schemas/KnowledgeBaseEntryContentReference' + - $ref: '#/components/schemas/SecurityAlertContentReference' + - $ref: '#/components/schemas/SecurityAlertsPageContentReference' + - $ref: '#/components/schemas/ProductDocumentationContentReference' + - $ref: '#/components/schemas/EsqlContentReference' + additionalProperties: false + type: object + + MessageMetadata: + type: object + description: Message metadata + properties: + contentReferences: + $ref: '#/components/schemas/ContentReferences' + description: Data refered to by the message content. + Replacements: type: object additionalProperties: @@ -85,6 +217,9 @@ components: traceData: $ref: '#/components/schemas/TraceData' description: trace Data + metadata: + $ref: '#/components/schemas/MessageMetadata' + description: metadata ApiConfig: type: object diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts index 19221ee4c27ac..6c5e52df671e1 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts @@ -22,6 +22,24 @@ export { replaceOriginalValuesWithUuidValues, } from './impl/data_anonymization/helpers'; +export { + contentReferencesStoreFactory, + securityAlertReference, + knowledgeBaseReference, + securityAlertsPageReference, + productDocumentationReference, + esqlQueryReference, + contentReferenceString, + contentReferenceBlock, + removeContentReferences, + pruneContentReferences, +} from './impl/content_references'; + +export type { + ContentReferencesStore, + ContentReferenceBlock, +} from './impl/content_references/types'; + export { transformRawData } from './impl/data_anonymization/transform_raw_data'; export { parseBedrockBuffer, handleBedrockChunk } from './impl/utils/bedrock'; export * from './constants'; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_body/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_body/index.tsx index 6fc10341864f3..8ce728f339db9 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_body/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_body/index.tsx @@ -91,7 +91,8 @@ export const AssistantBody: FunctionComponent = ({ ( commentsContainerRef.current?.childNodes[0].childNodes[0] as HTMLElement ).lastElementChild?.scrollIntoView(); - }); + // currentConversation is required in the dependency array to keep the scroll at the bottom. + }, [currentConversation]); // End Scrolling if (!isAssistantEnabled) { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/get_anonymization_tooltip/index.test.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/get_anonymization_tooltip/index.test.ts deleted file mode 100644 index 35bc314abc813..0000000000000 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/get_anonymization_tooltip/index.test.ts +++ /dev/null @@ -1,42 +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 { getAnonymizationTooltip } from '.'; -import { - SHOW_ANONYMIZED, - SHOW_REAL_VALUES, - THIS_CONVERSATION_DOES_NOT_INCLUDE_ANONYMIZED_FIELDS, -} from '../translations'; - -describe('getAnonymizationTooltip', () => { - it('returns the expected tooltip when conversationHasReplacements is false', () => { - const result = getAnonymizationTooltip({ - conversationHasReplacements: false, // <-- false - showAnonymizedValuesChecked: false, - }); - - expect(result).toBe(THIS_CONVERSATION_DOES_NOT_INCLUDE_ANONYMIZED_FIELDS); - }); - - it('returns SHOW_REAL_VALUES when showAnonymizedValuesChecked is true', () => { - const result = getAnonymizationTooltip({ - conversationHasReplacements: true, - showAnonymizedValuesChecked: true, // <-- true - }); - - expect(result).toBe(SHOW_REAL_VALUES); - }); - - it('returns SHOW_ANONYMIZED when showAnonymizedValuesChecked is false', () => { - const result = getAnonymizationTooltip({ - conversationHasReplacements: true, - showAnonymizedValuesChecked: false, // <-- false - }); - - expect(result).toBe(SHOW_ANONYMIZED); - }); -}); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/get_anonymization_tooltip/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/get_anonymization_tooltip/index.ts deleted file mode 100644 index 6534c6f8b0302..0000000000000 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/get_anonymization_tooltip/index.ts +++ /dev/null @@ -1,22 +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 * as i18n from '../translations'; - -export const getAnonymizationTooltip = ({ - conversationHasReplacements, - showAnonymizedValuesChecked, -}: { - conversationHasReplacements: boolean; - showAnonymizedValuesChecked: boolean; -}): string => { - if (!conversationHasReplacements) { - return i18n.THIS_CONVERSATION_DOES_NOT_INCLUDE_ANONYMIZED_FIELDS; - } - - return showAnonymizedValuesChecked ? i18n.SHOW_REAL_VALUES : i18n.SHOW_ANONYMIZED; -}; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.test.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.test.tsx index f91de230f7fb3..fe4b0de4d2149 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.test.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.test.tsx @@ -6,20 +6,14 @@ */ import React from 'react'; -import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'; -import userEvent, { PointerEventsCheckLevel } from '@testing-library/user-event'; +import { act, fireEvent, render, screen } from '@testing-library/react'; import { AssistantHeader } from '.'; import { TestProviders } from '../../mock/test_providers/test_providers'; -import { alertConvo, emptyWelcomeConvo, welcomeConvo } from '../../mock/conversation'; +import { alertConvo, welcomeConvo } from '../../mock/conversation'; import { useLoadConnectors } from '../../connectorland/use_load_connectors'; import { mockConnectors } from '../../mock/connectors'; -import { - CLOSE, - SHOW_ANONYMIZED, - SHOW_REAL_VALUES, - THIS_CONVERSATION_DOES_NOT_INCLUDE_ANONYMIZED_FIELDS, -} from './translations'; +import { CLOSE } from './translations'; const onConversationSelected = jest.fn(); const mockConversations = { @@ -49,6 +43,8 @@ const testProps = { anonymizationFields: { total: 0, page: 1, perPage: 1000, data: [] }, refetchAnonymizationFieldsResults: jest.fn(), allPrompts: [], + contentReferencesVisible: true, + setContentReferencesVisible: jest.fn(), }; jest.mock('../../connectorland/use_load_connectors', () => ({ @@ -79,59 +75,6 @@ describe('AssistantHeader', () => { beforeEach(() => { jest.clearAllMocks(); }); - it('showAnonymizedValues is not checked when selectedConversation.replacements is null', () => { - const { getByText, getByTestId } = render(, { - wrapper: TestProviders, - }); - expect(getByText(welcomeConvo.title)).toBeInTheDocument(); - expect(getByTestId('showAnonymizedValues').firstChild).toHaveAttribute( - 'data-euiicon-type', - 'eyeClosed' - ); - }); - - it('showAnonymizedValues is not checked when selectedConversation.replacements is empty', () => { - const { getByText, getByTestId } = render( - , - { - wrapper: TestProviders, - } - ); - expect(getByText(welcomeConvo.title)).toBeInTheDocument(); - expect(getByTestId('showAnonymizedValues').firstChild).toHaveAttribute( - 'data-euiicon-type', - 'eyeClosed' - ); - }); - - it('showAnonymizedValues is not checked when selectedConversation.replacements has values and showAnonymizedValues is false', () => { - const { getByTestId } = render( - , - { - wrapper: TestProviders, - } - ); - expect(getByTestId('showAnonymizedValues').firstChild).toHaveAttribute( - 'data-euiicon-type', - 'eyeClosed' - ); - }); - - it('showAnonymizedValues is checked when selectedConversation.replacements has values and showAnonymizedValues is true', () => { - const { getByTestId } = render( - , - { - wrapper: TestProviders, - } - ); - expect(getByTestId('showAnonymizedValues').firstChild).toHaveAttribute( - 'data-euiicon-type', - 'eye' - ); - }); it('Conversation is updated when connector change occurs', async () => { const { getByTestId } = render(, { @@ -157,93 +100,4 @@ describe('AssistantHeader', () => { expect(screen.getByRole('button', { name: CLOSE })).toBeInTheDocument(); }); - - it('disables the anonymization toggle button when there are NO replacements', () => { - render( - , - { - wrapper: TestProviders, - } - ); - - expect(screen.getByTestId('showAnonymizedValues')).toBeDisabled(); - }); - - it('displays the expected anonymization toggle button tooltip when there are NO replacements', async () => { - render( - , - { - wrapper: TestProviders, - } - ); - - await userEvent.hover(screen.getByTestId('showAnonymizedValues'), { - pointerEventsCheck: PointerEventsCheckLevel.Never, - }); - - await waitFor(() => { - expect(screen.getByTestId('showAnonymizedValuesTooltip')).toHaveTextContent( - THIS_CONVERSATION_DOES_NOT_INCLUDE_ANONYMIZED_FIELDS - ); - }); - }); - - it('enables the anonymization toggle button when there are replacements', () => { - render( - , // <-- conversation with replacements - { - wrapper: TestProviders, - } - ); - - expect(screen.getByTestId('showAnonymizedValues')).toBeEnabled(); - }); - - it('displays the SHOW_ANONYMIZED toggle button tooltip when there are replacements and showAnonymizedValues is false', async () => { - render( - , - { - wrapper: TestProviders, - } - ); - - await userEvent.hover(screen.getByTestId('showAnonymizedValues'), { - pointerEventsCheck: PointerEventsCheckLevel.Never, - }); - - await waitFor(() => { - expect(screen.getByTestId('showAnonymizedValuesTooltip')).toHaveTextContent(SHOW_ANONYMIZED); - }); - }); - - it('displays the SHOW_REAL_VALUES toggle button tooltip when there are replacements and showAnonymizedValues is true', async () => { - render( - , - { - wrapper: TestProviders, - } - ); - - await userEvent.hover(screen.getByTestId('showAnonymizedValues'), { - pointerEventsCheck: PointerEventsCheckLevel.Never, - }); - - await waitFor(() => { - expect(screen.getByTestId('showAnonymizedValuesTooltip')).toHaveTextContent(SHOW_REAL_VALUES); - }); - }); }); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx index 8c3b866ee128b..d99fcdc96e0a4 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx @@ -12,7 +12,6 @@ import { EuiFlexItem, EuiButtonIcon, EuiPanel, - EuiToolTip, EuiSkeletonTitle, } from '@elastic/eui'; import { css } from '@emotion/react'; @@ -24,10 +23,9 @@ import { AssistantTitle } from '../assistant_title'; import { ConnectorSelectorInline } from '../../connectorland/connector_selector_inline/connector_selector_inline'; import { FlyoutNavigation } from '../assistant_overlay/flyout_navigation'; import { AssistantSettingsModal } from '../settings/assistant_settings_modal'; -import * as i18n from './translations'; import { AIConnector } from '../../connectorland/connector_selector'; -import { getAnonymizationTooltip } from './get_anonymization_tooltip'; import { SettingsContextMenu } from '../settings/settings_context_menu/settings_context_menu'; +import * as i18n from './translations'; interface OwnProps { selectedConversation: Conversation | undefined; @@ -35,9 +33,7 @@ interface OwnProps { isDisabled: boolean; isLoading: boolean; isSettingsModalVisible: boolean; - onToggleShowAnonymizedValues: () => void; setIsSettingsModalVisible: React.Dispatch>; - showAnonymizedValues: boolean; onChatCleared: () => void; onCloseFlyout?: () => void; chatHistoryVisible?: boolean; @@ -65,9 +61,7 @@ export const AssistantHeader: React.FC = ({ isDisabled, isLoading, isSettingsModalVisible, - onToggleShowAnonymizedValues, setIsSettingsModalVisible, - showAnonymizedValues, onChatCleared, chatHistoryVisible, setChatHistoryVisible, @@ -80,13 +74,7 @@ export const AssistantHeader: React.FC = ({ isAssistantEnabled, refetchPrompts, }) => { - const showAnonymizedValuesChecked = useMemo( - () => - selectedConversation?.replacements != null && - Object.keys(selectedConversation?.replacements).length > 0 && - showAnonymizedValues, - [selectedConversation?.replacements, showAnonymizedValues] - ); + const { euiTheme } = useEuiTheme(); const selectedConnectorId = useMemo( () => selectedConversation?.apiConfig?.connectorId, @@ -103,12 +91,6 @@ export const AssistantHeader: React.FC = ({ [onConversationSelected] ); - const conversationHasReplacements = !isEmpty(selectedConversation?.replacements); - const anonymizationTooltip = getAnonymizationTooltip({ - conversationHasReplacements, - showAnonymizedValuesChecked, - }); - return ( <> = ({ onConnectorSelected={onConversationChange} /> - - - - - diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/translations.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/translations.ts index 7d105e1ee69a6..5589b6a274531 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/translations.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/translations.ts @@ -63,10 +63,43 @@ export const SHOW_REAL_VALUES = i18n.translate( } ); -export const THIS_CONVERSATION_DOES_NOT_INCLUDE_ANONYMIZED_FIELDS = i18n.translate( - 'xpack.elasticAssistant.assistant.settings.thisConversationDoesNotIncludeAnonymizedFieldsTooltip', +export const ANONYMIZE_VALUES = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.anonymizeValues', { - defaultMessage: 'This conversation does not include anonymized fields', + defaultMessage: 'Show anonymize values', + } +); + +export const SHOW_CITATIONS = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.showCitationsLabel', + { + defaultMessage: 'Show citations', + } +); + +export const CHAT_OPTIONS = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.chatOptions.label', + { + defaultMessage: 'Chat options', + } +); + +const isMac = navigator.platform.toLowerCase().indexOf('mac') >= 0; + +export const ANONYMIZE_VALUES_TOOLTIP = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.anonymizeValues.tooltip', + { + values: { keyboardShortcut: isMac ? '⌥ a' : 'Alt a' }, + defaultMessage: + 'Toggle to reveal or hide field values in your chat stream. The data sent to the LLM is still anonymized based on settings in the Anonymization panel. Keyboard shortcut: {keyboardShortcut}', + } +); + +export const SHOW_CITATIONS_TOOLTIP = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.showCitationsLabel.tooltip', + { + values: { keyboardShortcut: isMac ? '⌥ c' : 'Alt c' }, + defaultMessage: 'Keyboard shortcut: {keyboardShortcut}', } ); 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 b20122f822164..7860f197bd832 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 @@ -30,6 +30,7 @@ import { css } from '@emotion/react'; import styled from '@emotion/styled'; import { isEmpty } from 'lodash'; +import useEvent from 'react-use/lib/useEvent'; import { AssistantBody } from './assistant_body'; import { useCurrentConversation } from './use_current_conversation'; import { useDataStreamApis } from './use_data_stream_apis'; @@ -90,6 +91,11 @@ const AssistantComponent: React.FC = ({ promptContexts, currentUserAvatar, setLastConversationId, + contentReferencesVisible, + showAnonymizedValues, + setContentReferencesVisible, + setShowAnonymizedValues, + assistantFeatures: { contentReferencesEnabled }, } = useAssistantContext(); const [selectedPromptContexts, setSelectedPromptContexts] = useState< @@ -203,7 +209,6 @@ const AssistantComponent: React.FC = ({ ]); const [autoPopulatedOnce, setAutoPopulatedOnce] = useState(false); - const [showAnonymizedValues, setShowAnonymizedValues] = useState(false); const [messageCodeBlocks, setMessageCodeBlocks] = useState(); const [_, setCodeBlockControlsVisible] = useState(false); @@ -216,6 +221,28 @@ const AssistantComponent: React.FC = ({ } }, [augmentMessageCodeBlocks, currentConversation, showAnonymizedValues]); + // Keyboard shortcuts to toggle the visibility of content references and anonymized values + const onKeyDown = useCallback( + (event: KeyboardEvent) => { + if (event.altKey && event.code === 'KeyC') { + event.preventDefault(); + setContentReferencesVisible(!contentReferencesVisible); + } + if (event.altKey && event.code === 'KeyA') { + event.preventDefault(); + setShowAnonymizedValues(!showAnonymizedValues); + } + }, + [ + setContentReferencesVisible, + contentReferencesVisible, + setShowAnonymizedValues, + showAnonymizedValues, + ] + ); + + useEvent('keydown', onKeyDown); + // Show missing connector callout if no connectors are configured const showMissingConnectorCallout = useMemo(() => { @@ -265,10 +292,6 @@ const AssistantComponent: React.FC = ({ codeBlockContainers.forEach((e) => (e.style.minHeight = '85px')); //// - const onToggleShowAnonymizedValues = useCallback(() => { - setShowAnonymizedValues((prevValue) => !prevValue); - }, [setShowAnonymizedValues]); - const { abortStream, handleOnChatCleared, @@ -375,6 +398,8 @@ const AssistantComponent: React.FC = ({ setIsStreaming, currentUserAvatar, systemPromptContent: currentSystemPrompt?.content, + contentReferencesVisible, + contentReferencesEnabled, })} // Avoid comments going off the flyout css={css` @@ -402,7 +427,10 @@ const AssistantComponent: React.FC = ({ setIsStreaming, currentUserAvatar, currentSystemPrompt?.content, + contentReferencesVisible, + euiThemeVars.euiSizeL, selectedPromptContextsCount, + contentReferencesEnabled, ] ); @@ -461,9 +489,7 @@ const AssistantComponent: React.FC = ({ defaultConnector={defaultConnector} isDisabled={isDisabled || isLoadingChatSend} isSettingsModalVisible={isSettingsModalVisible} - onToggleShowAnonymizedValues={onToggleShowAnonymizedValues} setIsSettingsModalVisible={setIsSettingsModalVisible} - showAnonymizedValues={showAnonymizedValues} onCloseFlyout={onCloseFlyout} onChatCleared={handleOnChatCleared} chatHistoryVisible={chatHistoryVisible} diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx index 4282955716c7a..8efea263b445e 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx @@ -15,6 +15,12 @@ import { EuiNotificationBadge, EuiPopover, EuiButtonIcon, + EuiSwitch, + EuiPanel, + EuiTitle, + EuiHorizontalRule, + EuiToolTip, + EuiSwitchEvent, } from '@elastic/eui'; import { css } from '@emotion/react'; import { euiThemeVars } from '@kbn/ui-theme'; @@ -33,7 +39,15 @@ interface Params { export const SettingsContextMenu: React.FC = React.memo( ({ isDisabled = false, onChatCleared }: Params) => { - const { navigateToApp, knowledgeBase } = useAssistantContext(); + const { + navigateToApp, + knowledgeBase, + setContentReferencesVisible, + contentReferencesVisible, + showAnonymizedValues, + setShowAnonymizedValues, + assistantFeatures: { contentReferencesEnabled }, + } = useAssistantContext(); const [isPopoverOpen, setPopover] = useState(false); @@ -88,6 +102,20 @@ export const SettingsContextMenu: React.FC = React.memo( closePopover(); }, [closePopover, showAlertSettingsModal]); + const onChangeContentReferencesVisible = useCallback( + (e: EuiSwitchEvent) => { + setContentReferencesVisible(e.target.checked); + }, + [setContentReferencesVisible] + ); + + const onChangeShowAnonymizedValues = useCallback( + (e: EuiSwitchEvent) => { + setShowAnonymizedValues(e.target.checked); + }, + [setShowAnonymizedValues] + ); + const items = useMemo( () => [ = React.memo( {i18n.KNOWLEDGE_BASE} , @@ -133,7 +161,58 @@ export const SettingsContextMenu: React.FC = React.memo( , - + +

{i18n.CHAT_OPTIONS}

+ + + + + + + + {contentReferencesEnabled && ( + + + + + + )} + + = React.memo( `} > {i18n.RESET_CONVERSATION} - , + + , ], - [ + contentReferencesVisible, + onChangeContentReferencesVisible, + showAnonymizedValues, + onChangeShowAnonymizedValues, + euiThemeVars.euiColorDanger, handleNavigateToAnonymization, handleNavigateToKnowledgeBase, handleNavigateToSettings, handleShowAlertsModal, knowledgeBase.latestAlerts, showDestroyModal, + contentReferencesEnabled, + euiThemeVars.euiSizeM, + euiThemeVars.euiSizeXS, ] ); @@ -185,7 +272,7 @@ export const SettingsContextMenu: React.FC = React.memo( diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/constants.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/constants.tsx index c8b08f90f14bc..f181b08fe7268 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/constants.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/constants.tsx @@ -19,6 +19,8 @@ export const QUERY_LOCAL_STORAGE_KEY = 'query'; export const SHOW_SETTINGS_TOUR_LOCAL_STORAGE_KEY = 'showSettingsTour'; export const START_LOCAL_STORAGE_KEY = 'start'; export const STREAMING_LOCAL_STORAGE_KEY = 'streaming'; +export const CONTENT_REFERENCES_VISIBLE_LOCAL_STORAGE_KEY = 'contentReferencesVisible'; +export const SHOW_ANONYMIZED_VALUES_LOCAL_STORAGE_KEY = 'showAnonymizedValues'; export const TRACE_OPTIONS_SESSION_STORAGE_KEY = 'traceOptions'; export const CONVERSATION_TABLE_SESSION_STORAGE_KEY = 'conversationTable'; export const QUICK_PROMPT_TABLE_SESSION_STORAGE_KEY = 'quickPromptTable'; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/index.tsx index ebf85e0f86a90..6ef7a5ef7f04b 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -34,10 +34,12 @@ import { CodeBlockDetails } from '../assistant/use_conversation/helpers'; import { PromptContextTemplate } from '../assistant/prompt_context/types'; import { KnowledgeBaseConfig, TraceOptions } from '../assistant/types'; import { + CONTENT_REFERENCES_VISIBLE_LOCAL_STORAGE_KEY, DEFAULT_ASSISTANT_NAMESPACE, DEFAULT_KNOWLEDGE_BASE_SETTINGS, KNOWLEDGE_BASE_LOCAL_STORAGE_KEY, LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY, + SHOW_ANONYMIZED_VALUES_LOCAL_STORAGE_KEY, STREAMING_LOCAL_STORAGE_KEY, TRACE_OPTIONS_SESSION_STORAGE_KEY, } from './constants'; @@ -115,6 +117,10 @@ export interface UseAssistantContext { nameSpace: string; registerPromptContext: RegisterPromptContext; selectedSettingsTab: SettingsTabs | null; + contentReferencesVisible: boolean; + showAnonymizedValues: boolean; + setShowAnonymizedValues: React.Dispatch>; + setContentReferencesVisible: React.Dispatch>; setAssistantStreamingEnabled: React.Dispatch>; setKnowledgeBase: React.Dispatch>; setLastConversationId: React.Dispatch>; @@ -197,6 +203,24 @@ export const AssistantProvider: React.FC = ({ true ); + /** + * Local storage for content references configuration, prefixed by assistant nameSpace + */ + // can be undefined from localStorage, if not defined, default to true + const [contentReferencesVisible, setContentReferencesVisible] = useLocalStorage( + `${nameSpace}.${CONTENT_REFERENCES_VISIBLE_LOCAL_STORAGE_KEY}`, + true + ); + + /** + * Local storage for anonymized values, prefixed by assistant nameSpace + */ + // can be undefined from localStorage, if not defined, default to false + const [showAnonymizedValues, setShowAnonymizedValues] = useLocalStorage( + `${nameSpace}.${SHOW_ANONYMIZED_VALUES_LOCAL_STORAGE_KEY}`, + false + ); + /** * Prompt contexts are used to provide components a way to register and make their data available to the assistant. */ @@ -273,7 +297,7 @@ export const AssistantProvider: React.FC = ({ // Fetch assistant capabilities const { data: assistantFeatures } = useCapabilities({ http, toasts }); - const value = useMemo( + const value: UseAssistantContext = useMemo( () => ({ actionTypeRegistry, alertsIndexPattern, @@ -302,6 +326,14 @@ export const AssistantProvider: React.FC = ({ assistantStreamingEnabled: localStorageStreaming ?? true, setAssistantStreamingEnabled: setLocalStorageStreaming, setKnowledgeBase: setLocalStorageKnowledgeBase, + contentReferencesVisible: contentReferencesVisible ?? true, + setContentReferencesVisible: setContentReferencesVisible as React.Dispatch< + React.SetStateAction + >, + showAnonymizedValues: showAnonymizedValues ?? false, + setShowAnonymizedValues: setShowAnonymizedValues as React.Dispatch< + React.SetStateAction + >, setSelectedSettingsTab, setShowAssistantOverlay, setTraceOptions: setSessionStorageTraceOptions, @@ -342,6 +374,10 @@ export const AssistantProvider: React.FC = ({ localStorageStreaming, setLocalStorageStreaming, setLocalStorageKnowledgeBase, + showAnonymizedValues, + setShowAnonymizedValues, + contentReferencesVisible, + setContentReferencesVisible, setSessionStorageTraceOptions, showAssistantOverlay, title, 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 80996bbf80d68..ed5b690e9aa57 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,4 +83,6 @@ export type GetAssistantMessages = (commentArgs: { currentUserAvatar?: UserAvatar; setIsStreaming: (isStreaming: boolean) => void; systemPromptContent?: string; + contentReferencesVisible?: boolean; + contentReferencesEnabled?: boolean; }) => EuiCommentProps[]; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.test.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.test.tsx index 4900a6b0966e3..a660d08b46eef 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.test.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.test.tsx @@ -25,6 +25,8 @@ import { useAssistantContext } from '../../..'; import { I18nProvider } from '@kbn/i18n-react'; import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import { useKnowledgeBaseIndices } from '../../assistant/api/knowledge_base/use_knowledge_base_indices'; +import { Router } from '@kbn/shared-ux-router'; +import { createMemoryHistory, History } from 'history'; const mockContext = { basePromptContexts: MOCK_QUICK_PROMPTS, @@ -59,9 +61,17 @@ const mockDataViews = { getExistingIndices: jest.fn().mockResolvedValue(['index-2']), } as unknown as DataViewsContract; const queryClient = new QueryClient(); -const wrapper = (props: { children: React.ReactNode }) => ( +const Wrapper = ({ + children, + history = createMemoryHistory(), +}: { + children: React.ReactNode; + history?: History; +}) => ( - {props.children} + + {children} + ); describe('KnowledgeBaseSettingsManagement', () => { @@ -177,7 +187,7 @@ describe('KnowledgeBaseSettingsManagement', () => { it('renders loading spinner when data is not fetched', () => { (useKnowledgeBaseStatus as jest.Mock).mockReturnValue({ data: {}, isFetched: false }); render(, { - wrapper, + wrapper: Wrapper, }); expect(screen.getByTestId('spinning')).toBeInTheDocument(); @@ -195,7 +205,7 @@ describe('KnowledgeBaseSettingsManagement', () => { }); (isKnowledgeBaseSetup as jest.Mock).mockReturnValue(false); render(, { - wrapper, + wrapper: Wrapper, }); expect(screen.getByTestId('setup-knowledge-base-button')).toBeInTheDocument(); @@ -203,7 +213,7 @@ describe('KnowledgeBaseSettingsManagement', () => { it('renders knowledge base table with entries', async () => { render(, { - wrapper, + wrapper: Wrapper, }); waitFor(() => { expect(screen.getByTestId('knowledge-base-entries-table')).toBeInTheDocument(); @@ -221,7 +231,7 @@ describe('KnowledgeBaseSettingsManagement', () => { }); render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => { @@ -233,6 +243,23 @@ describe('KnowledgeBaseSettingsManagement', () => { expect(openFlyoutMock).toHaveBeenCalled(); }); + it('uses entry_search_term as default query', async () => { + const rawHistory = createMemoryHistory({ + initialEntries: ['/example?entry_search_term=testQuery'], + }); + const { container } = render(, { + wrapper: (props) => {props.children}, + }); + waitFor(() => { + expect(screen.getByTestId('knowledge-base-entries-table')).toBeInTheDocument(); + expect( + container + .querySelector('input[type=search][placeholder="Search for an entry"]') + ?.getAttribute('value') + ).toEqual('testQuery'); + }); + }); + it('refreshes table on refresh button click', async () => { const refetchMock = jest.fn(); (useKnowledgeBaseEntries as jest.Mock).mockReturnValue({ @@ -242,7 +269,7 @@ describe('KnowledgeBaseSettingsManagement', () => { }); render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => { @@ -259,7 +286,7 @@ describe('KnowledgeBaseSettingsManagement', () => { closeFlyout: closeFlyoutMock, }); render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => { @@ -282,7 +309,7 @@ describe('KnowledgeBaseSettingsManagement', () => { it('handles delete confirmation modal actions', async () => { render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => { @@ -302,7 +329,7 @@ describe('KnowledgeBaseSettingsManagement', () => { closeFlyout: jest.fn(), }); render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => { @@ -342,7 +369,7 @@ describe('KnowledgeBaseSettingsManagement', () => { closeFlyout: jest.fn(), }); render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => { @@ -382,7 +409,7 @@ describe('KnowledgeBaseSettingsManagement', () => { closeFlyout: jest.fn(), }); render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => { @@ -422,7 +449,7 @@ describe('KnowledgeBaseSettingsManagement', () => { closeFlyout: closeFlyoutMock, }); render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => { @@ -472,7 +499,7 @@ describe('KnowledgeBaseSettingsManagement', () => { it('shows warning icon for index entries with missing indices', async () => { render(, { - wrapper, + wrapper: Wrapper, }); await waitFor(() => expect(screen.getByTestId('missing-index-icon')).toBeInTheDocument()); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx index b47c7649dcefd..72a4b487794d6 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx @@ -31,6 +31,7 @@ import { import { css } from '@emotion/react'; import { DataViewsContract } from '@kbn/data-views-plugin/public'; import useAsync from 'react-use/lib/useAsync'; +import { useSearchParams } from 'react-router-dom-v5-compat'; import { ProductDocumentationManagement } from '../../assistant/settings/product_documentation'; import { KnowledgeBaseTour } from '../../tour/knowledge_base'; import { AlertsSettingsManagement } from '../../assistant/settings/alerts_settings/alerts_settings_management'; @@ -83,6 +84,11 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ d enabled: isAssistantEnabled, }); const isKbSetup = isKnowledgeBaseSetup(kbStatus); + const [searchParams] = useSearchParams(); + const initialSearchTerm = useMemo( + () => (searchParams.get('entry_search_term') as string) ?? undefined, + [searchParams] + ); const [deleteKBItem, setDeleteKBItem] = useState(null); const [duplicateKBItem, setDuplicateKBItem] = useState( @@ -286,8 +292,9 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ d placeholder: i18n.SEARCH_PLACEHOLDER, }, filters: [], + defaultQuery: initialSearchTerm, }), - [isFetchingEntries, handleRefreshTable, onDocumentClicked, onIndexClicked] + [isFetchingEntries, handleRefreshTable, onDocumentClicked, onIndexClicked, initialSearchTerm] ); const flyoutTitle = useMemo(() => { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/tsconfig.json b/x-pack/platform/packages/shared/kbn-elastic-assistant/tsconfig.json index b3be925d269df..a96a677bd1de5 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/tsconfig.json +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/tsconfig.json @@ -39,5 +39,6 @@ "@kbn/ai-assistant-icon", "@kbn/product-doc-base-plugin", "@kbn/spaces-plugin", + "@kbn/shared-ux-router", ] } 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 89c6d7b50c295..a5ef8e07a07e3 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,6 +68,7 @@ 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/append_conversation_messages.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.ts index 7bfc05993d43d..48b64764ccf85 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/append_conversation_messages.ts @@ -59,6 +59,9 @@ export const appendConversationMessages = async ({ if (message.trace_data != null) { newMessage.trace_data = message.trace_data; } + if (message.metadata != null) { + newMessage.metadata = message.metadata; + } messages.add(newMessage); } ctx._source.messages = messages; @@ -100,6 +103,15 @@ export const transformToUpdateScheme = (updatedAt: string, messages: Message[]) is_error: message.isError, reader: message.reader, role: message.role, + ...(message.metadata + ? { + metadata: { + ...(message.metadata.contentReferences + ? { content_references: message.metadata.contentReferences } + : {}), + }, + } + : {}), ...(message.traceData ? { trace_data: { 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 f23d462b84efc..314c83728bcd0 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 @@ -168,3 +168,18 @@ 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/conversations/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts index bdd1107942cc1..cc25af9f96782 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts @@ -65,6 +65,12 @@ export const getUpdateScript = ({ if (message.trace_data != null) { newMessage.trace_data = message.trace_data; } + if (message.metadata != null) { + newMessage.metadata = [:]; + if (message.metadata.content_references != null) { + newMessage.metadata.content_references = message.metadata.content_references; + } + } messages.add(newMessage); } ctx._source.messages = messages; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts index 7c4f9708862a5..dcddac1400e08 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts @@ -18,6 +18,15 @@ import { getConversation } from './get_conversation'; import { deleteConversation } from './delete_conversation'; import { appendConversationMessages } from './append_conversation_messages'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; + +/** + * Params for when creating ConversationDataClient in Request Context Factory. Useful if needing to modify + * configuration after initial plugin start + */ +export interface GetAIAssistantConversationsDataClientParams { + contentReferencesEnabled?: boolean; +} + export class AIAssistantConversationsDataClient extends AIAssistantDataClient { constructor(public readonly options: AIAssistantDataClientParams) { super(options); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts index 39798aeb2fd5e..c11d872e9d3ec 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts @@ -66,6 +66,15 @@ export const transformESSearchToConversations = ( ...(message.is_error ? { isError: message.is_error } : {}), ...(message.reader ? { reader: message.reader } : {}), role: message.role, + ...(message.metadata + ? { + metadata: { + ...(message.metadata.content_references + ? { contentReferences: message.metadata.content_references } + : {}), + }, + } + : {}), ...(message.trace_data ? { traceData: { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts index baeea677b1a66..cdf630640b452 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts @@ -145,6 +145,16 @@ describe('transformToUpdateScheme', () => { traceId: 'something', transactionId: 'something', }, + metadata: { + contentReferences: { + zm3i5: { + knowledgeBaseEntryName: 'Favorite_Food', + knowledgeBaseEntryId: '1c53565d-c6f1-45ab-9f4b-80b604dba8f3', + id: 'zm3i5', + type: 'KnowledgeBaseEntry', + }, + }, + }, }, { content: 'Message 4', @@ -177,6 +187,16 @@ describe('transformToUpdateScheme', () => { trace_id: 'something', transaction_id: 'something', }, + metadata: { + content_references: { + zm3i5: { + knowledgeBaseEntryName: 'Favorite_Food', + knowledgeBaseEntryId: '1c53565d-c6f1-45ab-9f4b-80b604dba8f3', + id: 'zm3i5', + type: 'KnowledgeBaseEntry', + }, + }, + }, }, { '@timestamp': '2011-10-06T14:48:00.000Z', diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts index 7e9ee336f6fe1..d8abaeaeb2595 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts @@ -14,6 +14,7 @@ import { MessageRole, ConversationSummary, UUID, + ContentReferences, } from '@kbn/elastic-assistant-common'; import { getConversation } from './get_conversation'; import { getUpdateScript } from './helpers'; @@ -33,6 +34,9 @@ export interface UpdateConversationSchema { transaction_id?: string; trace_id?: string; }; + metadata?: { + content_references?: ContentReferences; + }; }>; api_config?: { action_type_id?: string; @@ -66,6 +70,7 @@ export const updateConversation = async ({ }: UpdateConversationParams): Promise => { const updatedAt = new Date().toISOString(); const params = transformToUpdateScheme(updatedAt, conversationUpdateProps); + try { const response = await esClient.updateByQuery({ conflicts: 'proceed', @@ -139,6 +144,15 @@ export const transformToUpdateScheme = ( is_error: message.isError, reader: message.reader, role: message.role, + ...(message.metadata + ? { + metadata: { + ...(message.metadata.contentReferences + ? { content_references: message.metadata.contentReferences } + : {}), + }, + } + : {}), ...(message.traceData ? { trace_data: { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx index d90e1e2a05309..d38bcfef0e755 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx @@ -15,19 +15,15 @@ import { } from './helpers'; import { authenticatedUser } from '../../__mocks__/user'; import { getCreateKnowledgeBaseEntrySchemaMock } from '../../__mocks__/knowledge_base_entry_schema.mock'; -import { IndexEntry } from '@kbn/elastic-assistant-common'; +import { + ContentReferencesStore, + EsqlContentReference, + IndexEntry, +} from '@kbn/elastic-assistant-common'; +import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; // Mock dependencies jest.mock('@elastic/elasticsearch'); -jest.mock('@kbn/zod', () => ({ - z: { - string: jest.fn().mockReturnValue({ describe: (str: string) => str }), - number: jest.fn().mockReturnValue({ describe: (str: string) => str }), - boolean: jest.fn().mockReturnValue({ describe: (str: string) => str }), - object: jest.fn().mockReturnValue({ describe: (str: string) => str }), - any: jest.fn().mockReturnValue({ describe: (str: string) => str }), - }, -})); jest.mock('lodash'); describe('isModelAlreadyExistsError', () => { @@ -153,12 +149,14 @@ describe('getStructuredToolForIndexEntry', () => { const mockEsClient = {} as ElasticsearchClient; const mockIndexEntry = getCreateKnowledgeBaseEntrySchemaMock({ type: 'index' }) as IndexEntry; + const contentReferencesStore = contentReferencesStoreFactoryMock(); it('should return a DynamicStructuredTool with correct name and schema', () => { const tool = getStructuredToolForIndexEntry({ indexEntry: mockIndexEntry, esClient: mockEsClient, logger: mockLogger, + contentReferencesStore, }); expect(tool).toBeInstanceOf(DynamicStructuredTool); @@ -176,6 +174,8 @@ describe('getStructuredToolForIndexEntry', () => { hits: { hits: [ { + _index: 'exampleIndex', + _id: 'exampleId', _source: { field1: 'value1', field2: 2, @@ -194,13 +194,28 @@ describe('getStructuredToolForIndexEntry', () => { indexEntry: mockIndexEntry, esClient: mockEsClient, logger: mockLogger, + contentReferencesStore, }); + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + const reference = creator({ id: 'exampleContentReferenceId' }); + expect(reference.type).toEqual('EsqlQuery'); + expect((reference as EsqlContentReference).label).toEqual('exampleIndex'); + expect((reference as EsqlContentReference).query).toEqual( + 'FROM exampleIndex METADATA _id\n | WHERE _id == "exampleId"' + ); + return reference; + } + ); + const input = { query: 'testQuery', field1: 'value1', field2: 2 }; const result = await tool.invoke(input, {}); expect(result).toContain('Below are all relevant documents in JSON format'); - expect(result).toContain('"text":"Inner text 1\\n --- \\nInner text 2"'); + expect(result).toContain( + '"text":"Inner text 1\\n --- \\nInner text 2","citation":"{reference(exampleContentReferenceId)}"' + ); }); it('should log an error and return error message on Elasticsearch error', async () => { @@ -211,6 +226,7 @@ describe('getStructuredToolForIndexEntry', () => { indexEntry: mockIndexEntry, esClient: mockEsClient, logger: mockLogger, + contentReferencesStore, }); const input = { query: 'testQuery', field1: 'value1', field2: 2 }; @@ -230,6 +246,7 @@ describe('getStructuredToolForIndexEntry', () => { }) as IndexEntry, esClient: mockEsClient, logger: mockLogger, + contentReferencesStore, }); const nameRegex = /^[a-zA-Z0-9_-]+$/; @@ -244,6 +261,7 @@ describe('getStructuredToolForIndexEntry', () => { }) as IndexEntry, esClient: mockEsClient, logger: mockLogger, + contentReferencesStore, }); expect(tool.lc_kwargs.name).toMatch('testing'); 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 0a5830a538148..7d612cd53ae74 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 @@ -10,7 +10,12 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { errors } from '@elastic/elasticsearch'; import { QueryDslQueryContainer, SearchRequest } from '@elastic/elasticsearch/lib/api/types'; import { AuthenticatedUser } from '@kbn/core-security-common'; -import { IndexEntry } from '@kbn/elastic-assistant-common'; +import { + contentReferenceBlock, + ContentReferencesStore, + esqlQueryReference, + IndexEntry, +} from '@kbn/elastic-assistant-common'; import { ElasticsearchClient, Logger } from '@kbn/core/server'; export const isModelAlreadyExistsError = (error: Error) => { @@ -138,10 +143,12 @@ export const getKBVectorSearchQuery = ({ export const getStructuredToolForIndexEntry = ({ indexEntry, esClient, + contentReferencesStore, logger, }: { indexEntry: IndexEntry; esClient: ElasticsearchClient; + contentReferencesStore: ContentReferencesStore | false; logger: Logger; }): DynamicStructuredTool => { const inputSchema = indexEntry.inputSchema?.reduce((prev, input) => { @@ -210,15 +217,27 @@ export const getStructuredToolForIndexEntry = ({ const result = await esClient.search(params); const kbDocs = result.hits.hits.map((hit) => { + const esqlQuery = `FROM ${hit._index} ${ + hit._id ? `METADATA _id\n | WHERE _id == "${hit._id}"` : '' + }`; + + const reference = + contentReferencesStore && + contentReferencesStore.add((p) => esqlQueryReference(p.id, esqlQuery, hit._index)); + if (indexEntry.outputFields && indexEntry.outputFields.length > 0) { - return indexEntry.outputFields.reduce((prev, field) => { - // @ts-expect-error - return { ...prev, [field]: hit._source[field] }; - }, {}); + return indexEntry.outputFields.reduce( + (prev, field) => { + // @ts-expect-error + return { ...prev, [field]: hit._source[field] }; + }, + reference ? { citation: contentReferenceBlock(reference) } : {} + ); } return { text: hit.highlight?.[indexEntry.field].join('\n --- \n'), + ...(reference ? { citation: contentReferenceBlock(reference) } : {}), }; }); @@ -228,7 +247,7 @@ export const getStructuredToolForIndexEntry = ({ return `###\nBelow are all relevant documents in JSON format:\n${JSON.stringify( kbDocs - )}\n###`; + )}###`; } catch (e) { logger.error(`Error performing IndexEntry KB Similarity Search: ${e.message}`); return `I'm sorry, but I was unable to find any information in the knowledge base. Perhaps this error would be useful to deliver to the user. Be sure to print it below your response and in a codeblock so it is rendered nicely: ${e.message}`; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts index cf67d763e3d23..13599516631de 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts @@ -27,6 +27,7 @@ import { getSecurityLabsDocsCount, } from '../../lib/langchain/content_loaders/security_labs_loader'; import { DynamicStructuredTool } from '@langchain/core/tools'; +import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; jest.mock('../../lib/langchain/content_loaders/security_labs_loader'); jest.mock('p-retry'); const date = '2023-03-28T22:27:28.159Z'; @@ -520,6 +521,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { const result = await client.getAssistantTools({ esClient: esClientMock, + contentReferencesStore: contentReferencesStoreFactoryMock(), }); expect(result).toHaveLength(1); @@ -534,6 +536,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { const result = await client.getAssistantTools({ esClient: esClientMock, + contentReferencesStore: contentReferencesStoreFactoryMock(), }); expect(result).toEqual([]); @@ -546,6 +549,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { const result = await client.getAssistantTools({ esClient: esClientMock, + contentReferencesStore: contentReferencesStoreFactoryMock(), }); expect(result).toEqual([]); 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 ab14236df5392..40e107473ff22 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 @@ -22,6 +22,7 @@ import { KnowledgeBaseEntryCreateProps, KnowledgeBaseEntryResponse, Metadata, + ContentReferencesStore, KnowledgeBaseEntryUpdateProps, } from '@kbn/elastic-assistant-common'; import pRetry from 'p-retry'; @@ -531,11 +532,14 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { const results = result.hits.hits.map((hit) => { const metadata = { + name: hit?._source?.name, + index: hit?._index, source: hit?._source?.source, required: hit?._source?.required, kbResource: hit?._source?.kb_resource, }; return new Document({ + id: hit?._id, pageContent: hit?._source?.text ?? '', metadata, }); @@ -768,8 +772,10 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { * is scoped to system user. */ public getAssistantTools = async ({ + contentReferencesStore, esClient, }: { + contentReferencesStore: ContentReferencesStore | false; esClient: ElasticsearchClient; }): Promise => { const user = this.options.currentUser; @@ -809,6 +815,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { indexEntry, esClient, logger: this.options.logger, + contentReferencesStore, }); }) ); 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 ff0f95340d466..4cca23a3300c1 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 @@ -17,7 +17,10 @@ import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/fie import { defendInsightsFieldMap } from '../ai_assistant_data_clients/defend_insights/field_maps_configuration'; import { getDefaultAnonymizationFields } from '../../common/anonymization'; import { AssistantResourceNames, GetElser } from '../types'; -import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; +import { + AIAssistantConversationsDataClient, + GetAIAssistantConversationsDataClientParams, +} from '../ai_assistant_data_clients/conversations'; import { InitializationPromise, ResourceInstallationHelper, @@ -25,7 +28,10 @@ import { errorResult, successResult, } from './create_resource_installation_helper'; -import { conversationsFieldMap } from '../ai_assistant_data_clients/conversations/field_maps_configuration'; +import { + conversationsFieldMap, + conversationsContentReferencesFieldMap, +} 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'; @@ -95,6 +101,9 @@ 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; constructor(private readonly options: AIAssistantServiceOpts) { this.initialized = false; @@ -203,6 +212,17 @@ 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, @@ -355,7 +375,7 @@ export class AIAssistantService { } public async createAIAssistantConversationsDataClient( - opts: CreateAIAssistantClientParams + opts: CreateAIAssistantClientParams & GetAIAssistantConversationsDataClientParams ): Promise { const res = await this.checkResourcesInstallation(opts); @@ -363,6 +383,19 @@ 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 deb737d9ad36a..0ccc7b6453684 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 @@ -11,7 +11,12 @@ import { BaseMessage } from '@langchain/core/messages'; import { Logger } from '@kbn/logging'; import { KibanaRequest, KibanaResponseFactory, ResponseHeaders } from '@kbn/core-http-server'; import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain'; -import { ExecuteConnectorRequestBody, Message, Replacements } from '@kbn/elastic-assistant-common'; +import { + ContentReferencesStore, + ExecuteConnectorRequestBody, + Message, + Replacements, +} from '@kbn/elastic-assistant-common'; import { StreamResponseWithHeaders } from '@kbn/ml-response-stream/server'; import { PublicMethodsOf } from '@kbn/utility-types'; import type { InferenceServerStart } from '@kbn/inference-plugin/server'; @@ -44,6 +49,7 @@ export interface AgentExecutorParams { assistantTools?: AssistantTool[]; connectorId: string; conversationId?: string; + contentReferencesStore: ContentReferencesStore | false; 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 e2fa213aabac6..3cf2aa0ba92c5 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,6 +42,7 @@ export interface GetDefaultAssistantGraphParams { signal?: AbortSignal; tools: StructuredTool[]; replacements: Replacements; + contentReferencesEnabled: boolean; } export type DefaultAssistantGraph = ReturnType; @@ -57,6 +58,7 @@ export const getDefaultAssistantGraph = ({ signal, tools, replacements, + contentReferencesEnabled = false, }: GetDefaultAssistantGraphParams) => { try { // Default graph state @@ -121,6 +123,10 @@ export const getDefaultAssistantGraph = ({ value: (x: string, y?: string) => y ?? x, default: () => 'English', }, + contentReferencesEnabled: { + value: (x: boolean, y?: boolean) => y ?? x, + default: () => contentReferencesEnabled, + }, }; // Default node parameters diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts index fa512b76b3b26..f2327130b6fe6 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts @@ -18,6 +18,7 @@ import { createStructuredChatAgent, createToolCallingAgent, } from 'langchain/agents'; +import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; import { savedObjectsClientMock } from '@kbn/core-saved-objects-api-server-mocks'; jest.mock('./graph'); jest.mock('./helpers'); @@ -75,6 +76,7 @@ describe('callAssistantGraph', () => { telemetryParams: {}, traceOptions: {}, responseLanguage: 'English', + contentReferencesStore: contentReferencesStoreFactoryMock(), } as unknown as AgentExecutorParams; beforeEach(() => { 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 c68c6f176e47d..06cc55791f5f1 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 @@ -14,6 +14,7 @@ import { } from 'langchain/agents'; 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 { promptGroupId } from '../../../prompt/local_prompt_object'; import { getModelOrOss } from '../../../prompt/helpers'; import { getPrompt, promptDictionary } from '../../../prompt'; @@ -33,6 +34,7 @@ export const callAssistantGraph: AgentExecutor = async ({ alertsIndexPattern, assistantTools = [], connectorId, + contentReferencesStore, conversationId, dataClients, esClient, @@ -107,6 +109,7 @@ export const callAssistantGraph: AgentExecutor = async ({ alertsIndexPattern, anonymizationFields, connectorId, + contentReferencesStore, esClient, inference, isEnabledKnowledgeBase, @@ -128,6 +131,7 @@ export const callAssistantGraph: AgentExecutor = async ({ if (isEnabledKnowledgeBase) { const kbTools = await dataClients?.kbDataClient?.getAssistantTools({ esClient, + contentReferencesStore, }); if (kbTools) { tools.push(...kbTools); @@ -191,6 +195,7 @@ 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, @@ -224,6 +229,15 @@ export const callAssistantGraph: AgentExecutor = async ({ traceOptions, }); + const contentReferences = + contentReferencesStore && pruneContentReferences(graphResponse.output, contentReferencesStore); + + const metadata: MessageMetadata = { + ...(contentReferences ? { contentReferences } : {}), + }; + + const isMetadataPopulated = !!contentReferences; + return { body: { connector_id: connectorId, @@ -231,6 +245,7 @@ export const callAssistantGraph: AgentExecutor = async ({ trace_data: graphResponse.traceData, replacements, status: 'ok', + ...(isMetadataPopulated ? { 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 830a120669661..72b4d6d483e93 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 @@ -7,6 +7,9 @@ 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'; @@ -67,9 +70,12 @@ 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: state.messages, // TODO: Message de-dupe with ...state spread + chat_history: sanitizeChatHistory(state.messages), // TODO: Message de-dupe with ...state spread }, config ); @@ -79,3 +85,15 @@ export async function runAgent({ lastNode: NodeType.AGENT, }; } + +/** + * Removes content references from chat history + */ +const sanitizeChatHistory = (messages: BaseMessage[]): BaseMessage[] => { + return messages.map((message) => { + if (!Array.isArray(message.content)) { + message.content = removeContentReferences(message.content); + } + return message; + }); +}; 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 6a30ef54a1f3b..a196229a2481e 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 @@ -41,6 +41,7 @@ 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.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts index 8bccb54653f03..a9a2f166399e8 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,14 +7,19 @@ 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 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}`; +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 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}`; +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. Always return value from NaturalLanguageESQLTool as is. 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.`; 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`; 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 bda4da8413120..2623f3ab60b75 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 @@ -16,6 +16,8 @@ import { transformRawData, getAnonymizedValue, ConversationResponse, + contentReferencesStoreFactory, + pruneContentReferences, } from '@kbn/elastic-assistant-common'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; @@ -25,6 +27,7 @@ import { buildResponse } from '../../lib/build_response'; import { appendAssistantMessageToConversation, createConversationWithUserInput, + DEFAULT_PLUGIN_NAME, getIsKnowledgeBaseInstalled, langChainExecute, performChecks, @@ -84,8 +87,15 @@ export const chatCompleteRoute = ( return checkResponse.response; } + const contentReferencesEnabled = + ctx.elasticAssistant.getRegisteredFeatures( + DEFAULT_PLUGIN_NAME + ).contentReferencesEnabled; + const conversationsDataClient = - await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); + await ctx.elasticAssistant.getAIAssistantConversationsDataClient({ + contentReferencesEnabled, + }); const anonymizationFieldsDataClient = await ctx.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient(); @@ -176,12 +186,18 @@ export const chatCompleteRoute = ( })); } + const contentReferencesStore = + contentReferencesEnabled && contentReferencesStoreFactory(); + const onLlmResponse = async ( content: string, traceData: Message['traceData'] = {}, isError = false ): Promise => { if (newConversation?.id && conversationsDataClient) { + const contentReferences = + contentReferencesStore && pruneContentReferences(content, contentReferencesStore); + await appendAssistantMessageToConversation({ conversationId: newConversation?.id, conversationsDataClient, @@ -189,6 +205,7 @@ export const chatCompleteRoute = ( replacements: latestReplacements, isError, traceData, + contentReferences, }); } }; @@ -209,6 +226,7 @@ export const chatCompleteRoute = ( onLlmResponse, onNewReplacements, replacements: latestReplacements, + contentReferencesStore, request: { ...request, // TODO: clean up after empty tools will be available to use diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts index 22e89202e638b..aa422393a7ccc 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts @@ -10,6 +10,7 @@ import moment from 'moment'; import { + ContentReferencesStore, DEFEND_INSIGHTS_TOOL_ID, DefendInsightStatus, DefendInsightType, @@ -58,6 +59,7 @@ describe('defend insights route helpers', () => { latestReplacements: {}, onNewReplacements: jest.fn(), request: {} as any, + contentReferencesStore: {} as ContentReferencesStore, }; const result = getAssistantToolParams(params); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts index e67f00ef6514c..db7c2d058069f 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts @@ -16,6 +16,7 @@ import type { import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import type { ApiConfig, + ContentReferencesStore, DefendInsight, DefendInsightGenerationInterval, DefendInsightsPostRequestBody, @@ -107,6 +108,7 @@ export function getAssistantToolParams({ langSmithProject, langSmithApiKey, logger, + contentReferencesStore, latestReplacements, onNewReplacements, request, @@ -122,6 +124,7 @@ export function getAssistantToolParams({ langSmithProject?: string; langSmithApiKey?: string; logger: Logger; + contentReferencesStore: ContentReferencesStore | false; latestReplacements: Replacements; onNewReplacements: (newReplacements: Replacements) => void; request: KibanaRequest; @@ -133,6 +136,7 @@ export function getAssistantToolParams({ langChainTimeout: number; llm: ActionsClientLlm; logger: Logger; + contentReferencesStore: ContentReferencesStore | false; replacements: Replacements; onNewReplacements: (newReplacements: Replacements) => void; request: KibanaRequest; @@ -169,6 +173,7 @@ export function getAssistantToolParams({ langChainTimeout, llm, logger, + contentReferencesStore, onNewReplacements, request, modelExists: false, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts index 21d6e775e7739..fb5adfa05bee0 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts @@ -149,6 +149,7 @@ export const postDefendInsightsRoute = (router: IRouter { @@ -112,6 +117,7 @@ export const getMessageFromRawResponse = ({ role: 'assistant', content: rawContent, timestamp: dateTimeString, + metadata, isError, traceData, }; @@ -169,6 +175,7 @@ export interface AppendAssistantMessageToConversationParams { messageContent: string; replacements: Replacements; conversationId: string; + contentReferences?: ContentReferences | false; isError?: boolean; traceData?: Message['traceData']; } @@ -177,6 +184,7 @@ export const appendAssistantMessageToConversation = async ({ messageContent, replacements, conversationId, + contentReferences, isError = false, traceData = {}, }: AppendAssistantMessageToConversationParams) => { @@ -185,6 +193,12 @@ export const appendAssistantMessageToConversation = async ({ return; } + const metadata: MessageMetadata = { + ...(contentReferences ? { contentReferences } : {}), + }; + + const isMetadataPopulated = Boolean(contentReferences) !== false; + await conversationsDataClient.appendConversationMessages({ existingConversation: conversation, messages: [ @@ -193,6 +207,7 @@ export const appendAssistantMessageToConversation = async ({ messageContent, replacements, }), + metadata: isMetadataPopulated ? metadata : undefined, traceData, isError, }), @@ -217,6 +232,7 @@ export interface LangChainExecuteParams { telemetry: AnalyticsServiceSetup; actionTypeId: string; connectorId: string; + contentReferencesStore: ContentReferencesStore | false; llmTasks?: LlmTasksPluginStart; inference: InferenceServerStart; isOssModel?: boolean; @@ -247,6 +263,7 @@ export const langChainExecute = async ({ telemetry, actionTypeId, connectorId, + contentReferencesStore, isOssModel, context, actionsClient, @@ -306,6 +323,7 @@ export const langChainExecute = async ({ assistantTools, conversationId, connectorId, + contentReferencesStore, esClient, llmTasks, inference, 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 8cd6dd24c24c7..e6df1f2ba3d08 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 @@ -12,9 +12,11 @@ import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; import { schema } from '@kbn/config-schema'; import { API_VERSIONS, + contentReferencesStoreFactory, ExecuteConnectorRequestBody, Message, Replacements, + pruneContentReferences, } from '@kbn/elastic-assistant-common'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { INVOKE_ASSISTANT_ERROR_EVENT } from '../lib/telemetry/event_based_telemetry'; @@ -23,6 +25,7 @@ import { buildResponse } from '../lib/build_response'; import { ElasticAssistantRequestHandlerContext, GetElser } from '../types'; import { appendAssistantMessageToConversation, + DEFAULT_PLUGIN_NAME, getIsKnowledgeBaseInstalled, getSystemPromptFromUserConversation, langChainExecute, @@ -107,16 +110,27 @@ 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(); + await assistantContext.getAIAssistantConversationsDataClient({ + contentReferencesEnabled, + }); const promptsDataClient = await assistantContext.getAIAssistantPromptsDataClient(); + const contentReferencesStore = + contentReferencesEnabled && contentReferencesStoreFactory(); + onLlmResponse = async ( content: string, traceData: Message['traceData'] = {}, isError = false ): Promise => { if (conversationsDataClient && conversationId) { + const contentReferences = + contentReferencesStore && pruneContentReferences(content, contentReferencesStore); + await appendAssistantMessageToConversation({ conversationId, conversationsDataClient, @@ -124,6 +138,7 @@ export const postActionsConnectorExecuteRoute = ( replacements: latestReplacements, isError, traceData, + contentReferences, }); } }; @@ -141,6 +156,7 @@ export const postActionsConnectorExecuteRoute = ( actionsClient, actionTypeId, connectorId, + contentReferencesStore, isOssModel, conversationId, context: ctx, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts index 7580fc8c0989d..25935c784e43b 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -145,13 +145,14 @@ export class RequestContextFactory implements IRequestContextFactory { }); }), - getAIAssistantConversationsDataClient: memoize(async () => { + getAIAssistantConversationsDataClient: memoize(async (params) => { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantConversationsDataClient({ spaceId: getSpaceId(), licensing: context.licensing, logger: this.logger, currentUser, + contentReferencesEnabled: params?.contentReferencesEnabled, }); }), }; 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 b485125cc41e3..cc07747725ee7 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 { performChecks } from '../helpers'; +import { DEFAULT_PLUGIN_NAME, performChecks } from '../helpers'; export const findUserConversationsRoute = (router: ElasticAssistantPluginRouter) => { router.versioned @@ -57,7 +57,15 @@ export const findUserConversationsRoute = (router: ElasticAssistantPluginRouter) if (!checkResponse.isSuccess) { return checkResponse.response; } - const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); + + const contentReferencesEnabled = + ctx.elasticAssistant.getRegisteredFeatures( + DEFAULT_PLUGIN_NAME + ).contentReferencesEnabled; + + const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient({ + contentReferencesEnabled, + }); const currentUser = checkResponse.currentUser; const additionalFilter = query.filter ? ` AND ${query.filter}` : ''; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts index b99479591b877..981b00d68e643 100755 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts @@ -33,6 +33,7 @@ import { AssistantFeatures, ExecuteConnectorRequestBody, Replacements, + ContentReferencesStore, } from '@kbn/elastic-assistant-common'; import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; import { @@ -51,7 +52,10 @@ import type { InferenceServerStart } from '@kbn/inference-plugin/server'; import { ProductDocBaseStartContract } from '@kbn/product-doc-base-plugin/server'; import type { GetAIAssistantKnowledgeBaseDataClientParams } from './ai_assistant_data_clients/knowledge_base'; import { AttackDiscoveryDataClient } from './lib/attack_discovery/persistence'; -import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/conversations'; +import { + AIAssistantConversationsDataClient, + GetAIAssistantConversationsDataClientParams, +} from './ai_assistant_data_clients/conversations'; import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; import { AIAssistantDataClient } from './ai_assistant_data_clients'; import { AIAssistantKnowledgeBaseDataClient } from './ai_assistant_data_clients/knowledge_base'; @@ -131,7 +135,9 @@ export interface ElasticAssistantApiRequestHandlerContext { getServerBasePath: () => string; getSpaceId: () => string; getCurrentUser: () => AuthenticatedUser | null; - getAIAssistantConversationsDataClient: () => Promise; + getAIAssistantConversationsDataClient: ( + params?: GetAIAssistantConversationsDataClientParams + ) => Promise; getAIAssistantKnowledgeBaseDataClient: ( params?: GetAIAssistantKnowledgeBaseDataClientParams ) => Promise; @@ -233,6 +239,7 @@ export interface AssistantToolParams { inference?: InferenceServerStart; isEnabledKnowledgeBase: boolean; connectorId?: string; + contentReferencesStore: ContentReferencesStore | false; esClient: ElasticsearchClient; kbDataClient?: AIAssistantKnowledgeBaseDataClient; langChainTimeout?: number; 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 7581e3e447b1b..3a7313dc92b6a 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 @@ -119,6 +119,11 @@ export const allowedExperimentalValues = Object.freeze({ */ attackDiscoveryAlertFiltering: 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/comment_actions/index.test.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/comment_actions/index.test.tsx new file mode 100644 index 0000000000000..24f9e28923c0f --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/comment_actions/index.test.tsx @@ -0,0 +1,54 @@ +/* + * 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 React from 'react'; +import { render } from '@testing-library/react'; +import type { ClientMessage } from '@kbn/elastic-assistant'; +import { createMockStore, mockGlobalState, TestProviders } from '../../common/mock'; +import { EuiCopy } from '@elastic/eui'; +import { CommentActions } from '.'; + +jest.mock('@elastic/eui', () => ({ + ...jest.requireActual('@elastic/eui'), + EuiCopy: jest.fn(), +})); + +const Wrapper: React.FC = ({ children }) => { + const store = createMockStore(mockGlobalState); + + return {children}; +}; + +describe('CommentActions', () => { + beforeEach(() => { + (EuiCopy as unknown as jest.Mock).mockClear(); + }); + + it.each([ + [`Only this should be copied!{reference(exampleReferenceId)}`, 'Only this should be copied!'], + [ + `Only this.{reference(exampleReferenceId)} should be copied!{reference(exampleReferenceId)}`, + 'Only this. should be copied!', + ], + [`{reference(exampleReferenceId)}`, ''], + ])("textToCopy is correct when input is '%s'", async (input, expected) => { + (EuiCopy as unknown as jest.Mock).mockReturnValue(null); + const message: ClientMessage = { + content: input, + role: 'assistant', + timestamp: '2025-01-08T10:47:34.578Z', + }; + render(, { wrapper: Wrapper }); + + expect(EuiCopy).toHaveBeenCalledWith( + expect.objectContaining({ + textToCopy: expected, + }), + expect.anything() + ); + }); +}); diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/comment_actions/index.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/comment_actions/index.tsx index f145deb9febc7..4afcd9ad80738 100644 --- a/x-pack/solutions/security/plugins/security_solution/public/assistant/comment_actions/index.tsx +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/comment_actions/index.tsx @@ -12,6 +12,7 @@ import React, { useCallback } from 'react'; import { useDispatch } from 'react-redux'; import { useAssistantContext } from '@kbn/elastic-assistant/impl/assistant_context'; +import { removeContentReferences } from '@kbn/elastic-assistant-common'; import { useKibana, useToasts } from '../../common/lib/kibana'; import type { Note } from '../../common/lib/note'; import { appActions } from '../../common/store/actions'; @@ -25,6 +26,10 @@ interface Props { message: ClientMessage; } +function getTextToCopy(content: string): string { + return removeContentReferences(content); +} + const CommentActionsComponent: React.FC = ({ message }) => { const toasts = useToasts(); const { cases } = useKibana().services; @@ -91,6 +96,8 @@ const CommentActionsComponent: React.FC = ({ message }) => { // ? `${basePath}/app/apm/services/kibana/transactions/view?kuery=&rangeFrom=now-1y&rangeTo=now&environment=ENVIRONMENT_ALL&serviceGroup=&comparisonEnabled=true&traceId=${message.traceData.traceId}&transactionId=${message.traceData.transactionId}&transactionName=POST%20/internal/elastic_assistant/actions/connector/?/_execute&transactionType=request&offset=1d&latencyAggregationType=avg` // : undefined; + const textToCopy = getTextToCopy(content); + return ( // APM Trace support is currently behind the Model Evaluation feature flag until wider testing is performed @@ -129,7 +136,7 @@ const CommentActionsComponent: React.FC = ({ message }) => { - + {(copy) => ( & { + contentReferenceCount: number; +}; + +export const ContentReferenceButton: React.FC = ({ + contentReferenceCount, + ...euiButtonEmptyProps +}) => { + return ( + + {`[${contentReferenceCount}]`} + + ); +}; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_component_factory.test.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_component_factory.test.tsx new file mode 100644 index 0000000000000..ca238dfc71892 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_component_factory.test.tsx @@ -0,0 +1,157 @@ +/* + * 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 { ContentReferences } from '@kbn/elastic-assistant-common'; +import { contentReferenceComponentFactory } from './content_reference_component_factory'; +import React from 'react'; +import { render, screen } from '@testing-library/react'; +import type { ContentReferenceNode } from '../content_reference_parser'; + +const testContentReferenceNode = { contentReferenceId: '1' } as ContentReferenceNode; + +jest.mock('../../../../common/lib/kibana', () => ({ + useKibana: jest.fn().mockReturnValue({ + services: { + discover: { + locator: jest.fn(), + }, + application: { + navigateToApp: jest.fn(), + }, + }, + }), +})); + +describe('contentReferenceComponentFactory', () => { + it.each([ + [ + 'EsqlQueryReference', + { + '1': { + id: '1', + type: 'EsqlQuery', + query: '', + label: '', + }, + } as ContentReferences, + testContentReferenceNode, + ], + [ + 'KnowledgeBaseEntryReference', + { + '1': { + id: '1', + type: 'KnowledgeBaseEntry', + knowledgeBaseEntryId: '', + knowledgeBaseEntryName: '', + }, + } as ContentReferences, + testContentReferenceNode, + ], + [ + 'ProductDocumentationReference', + { + '1': { + id: '1', + type: 'ProductDocumentation', + title: '', + url: '', + }, + } as ContentReferences, + testContentReferenceNode, + ], + [ + 'SecurityAlertReference', + { + '1': { + id: '1', + type: 'SecurityAlert', + alertId: '', + }, + } as ContentReferences, + testContentReferenceNode, + ], + [ + 'SecurityAlertsPageReference', + { + '1': { + id: '1', + type: 'SecurityAlertsPage', + }, + } as ContentReferences, + testContentReferenceNode, + ], + ])( + "Renders component: '%s'", + async ( + testId: string, + contentReferences: ContentReferences, + contentReferenceNode: ContentReferenceNode + ) => { + const Component = contentReferenceComponentFactory({ + contentReferences, + contentReferencesVisible: true, + loading: false, + }); + + render(); + + expect(screen.getByTestId(testId)).toBeInTheDocument(); + } + ); + + it('renders nothing when specific contentReference is undefined', async () => { + const Component = contentReferenceComponentFactory({ + contentReferences: {}, + contentReferencesVisible: true, + loading: false, + }); + + const { container } = render( + + ); + + expect(container).toBeEmptyDOMElement(); + expect(screen.queryByText('[1]')).not.toBeInTheDocument(); + }); + + it('renders placeholder if contentReferences are undefined and is loading', async () => { + const Component = contentReferenceComponentFactory({ + contentReferences: undefined, + contentReferencesVisible: true, + loading: true, + }); + + render( + + ); + + expect(screen.getByTestId('ContentReferenceButton')).toBeInTheDocument(); + expect(screen.getByText('[1]')).toBeInTheDocument(); + }); + + it('renders nothing if contentReferences are undefined and is not loading', async () => { + const Component = contentReferenceComponentFactory({ + contentReferences: undefined, + contentReferencesVisible: true, + loading: false, + }); + + const { container } = render( + + ); + + expect(container).toBeEmptyDOMElement(); + expect(screen.queryByText('[1]')).not.toBeInTheDocument(); + }); +}); diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_component_factory.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_component_factory.tsx new file mode 100644 index 0000000000000..710c948b449bf --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/content_reference_component_factory.tsx @@ -0,0 +1,91 @@ +/* + * 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 { ContentReferences } from '@kbn/elastic-assistant-common'; +import React from 'react'; +import type { ContentReferenceNode } from '../content_reference_parser'; +import { KnowledgeBaseEntryReference } from './knowledge_base_entry_reference'; +import { SecurityAlertReference } from './security_alert_reference'; +import { SecurityAlertsPageReference } from './security_alerts_page_reference'; +import { ContentReferenceButton } from './content_reference_button'; +import { ProductDocumentationReference } from './product_documentation_reference'; +import { EsqlQueryReference } from './esql_query_reference'; + +export interface ContentReferenceComponentFactory { + contentReferences?: ContentReferences; + contentReferencesVisible: boolean; + loading: boolean; +} + +export const contentReferenceComponentFactory = ({ + contentReferences, + contentReferencesVisible, + loading, +}: ContentReferenceComponentFactory) => { + const ContentReferenceComponent = ( + contentReferenceNode: ContentReferenceNode + ): React.ReactNode => { + if (!contentReferencesVisible) return null; + + const defaultNode = ( + + ); + + if (!contentReferences && loading) return defaultNode; + + const contentReference = contentReferences?.[contentReferenceNode.contentReferenceId]; + + if (!contentReference) return null; + + switch (contentReference.type) { + case 'KnowledgeBaseEntry': + return ( + + ); + case 'SecurityAlert': + return ( + + ); + case 'SecurityAlertsPage': + return ( + + ); + case 'ProductDocumentation': + return ( + + ); + case 'EsqlQuery': + return ( + + ); + default: + return defaultNode; + } + }; + + ContentReferenceComponent.displayName = 'ContentReferenceComponent'; + + return ContentReferenceComponent; +}; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/esql_query_reference.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/esql_query_reference.tsx new file mode 100644 index 0000000000000..0241b8750882d --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/esql_query_reference.tsx @@ -0,0 +1,57 @@ +/* + * 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 { EsqlContentReference } from '@kbn/elastic-assistant-common'; +import React, { useCallback } from 'react'; +import { EuiLink } from '@elastic/eui'; +import type { ContentReferenceNode } from '../content_reference_parser'; +import { PopoverReference } from './popover_reference'; +import { useKibana } from '../../../../common/lib/kibana'; + +interface Props { + contentReferenceNode: ContentReferenceNode; + esqlContentReference: EsqlContentReference; +} + +export const EsqlQueryReference: React.FC = ({ + contentReferenceNode, + esqlContentReference, +}) => { + const { + discover: { locator }, + application: { navigateToApp }, + } = useKibana().services; + + const onClick = useCallback( + async (e: React.MouseEvent) => { + e.preventDefault(); + if (!locator) { + return; + } + const url = await locator.getLocation({ + query: { + esql: esqlContentReference.query, + }, + }); + + navigateToApp(url.app, { + path: url.path, + openInNewTab: true, + }); + }, + [locator, esqlContentReference.query, navigateToApp] + ); + + return ( + + {esqlContentReference.label} + + ); +}; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/knowledge_base_entry_reference.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/knowledge_base_entry_reference.tsx new file mode 100644 index 0000000000000..c5b7e7586bdb6 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/knowledge_base_entry_reference.tsx @@ -0,0 +1,48 @@ +/* + * 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 { KnowledgeBaseEntryContentReference } from '@kbn/elastic-assistant-common'; +import React, { useCallback } from 'react'; +import { EuiLink } from '@elastic/eui'; +import { KNOWLEDGE_BASE_ENTRY_REFERENCE_LABEL } from './translations'; +import type { ContentReferenceNode } from '../content_reference_parser'; +import { PopoverReference } from './popover_reference'; +import { useKibana } from '../../../../common/lib/kibana'; + +interface Props { + contentReferenceNode: ContentReferenceNode; + knowledgeBaseEntryContentReference: KnowledgeBaseEntryContentReference; +} + +export const KnowledgeBaseEntryReference: React.FC = ({ + contentReferenceNode, + knowledgeBaseEntryContentReference, +}) => { + const { navigateToApp } = useKibana().services.application; + + const onClick = useCallback( + (e: React.MouseEvent) => { + e.preventDefault(); + navigateToApp('management', { + path: `kibana/securityAiAssistantManagement?tab=knowledge_base&entry_search_term=${knowledgeBaseEntryContentReference.knowledgeBaseEntryId}`, + openInNewTab: true, + }); + }, + [navigateToApp, knowledgeBaseEntryContentReference] + ); + + return ( + + + {`${KNOWLEDGE_BASE_ENTRY_REFERENCE_LABEL}: ${knowledgeBaseEntryContentReference.knowledgeBaseEntryName}`} + + + ); +}; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/popover_reference.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/popover_reference.tsx new file mode 100644 index 0000000000000..3b8383bba7037 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/popover_reference.tsx @@ -0,0 +1,55 @@ +/* + * 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 { EuiPopover } from '@elastic/eui'; +import React, { useCallback, useMemo, useState } from 'react'; +import { css } from '@emotion/css'; +import { ContentReferenceButton } from './content_reference_button'; + +interface Props { + contentReferenceCount: number; + 'data-test-subj'?: string; +} + +export const PopoverReference: React.FC> = ({ + contentReferenceCount, + children, + ...rest +}) => { + const [isPopoverOpen, setIsPopoverOpen] = useState(false); + + const togglePopover = useCallback(() => setIsPopoverOpen((prev) => !prev), []); + const closePopover = useCallback(() => setIsPopoverOpen(false), []); + const openPopover = useCallback(() => setIsPopoverOpen(true), []); + + const button = useMemo( + () => ( + + ), + [contentReferenceCount, openPopover, togglePopover] + ); + + return ( + + {children} + + ); +}; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/product_documentation_reference.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/product_documentation_reference.tsx new file mode 100644 index 0000000000000..215d2db00a46e --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/product_documentation_reference.tsx @@ -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 type { ProductDocumentationContentReference } from '@kbn/elastic-assistant-common'; +import React from 'react'; +import { EuiLink } from '@elastic/eui'; +import type { ContentReferenceNode } from '../content_reference_parser'; +import { PopoverReference } from './popover_reference'; + +interface Props { + contentReferenceNode: ContentReferenceNode; + productDocumentationContentReference: ProductDocumentationContentReference; +} + +export const ProductDocumentationReference: React.FC = ({ + contentReferenceNode, + productDocumentationContentReference, +}) => { + return ( + + + {productDocumentationContentReference.title} + + + ); +}; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/security_alert_reference.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/security_alert_reference.tsx new file mode 100644 index 0000000000000..8accfd4fd0a2b --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/security_alert_reference.tsx @@ -0,0 +1,46 @@ +/* + * 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 { SecurityAlertContentReference } from '@kbn/elastic-assistant-common'; +import React, { useCallback } from 'react'; +import { EuiLink } from '@elastic/eui'; +import { SECURITY_ALERT_REFERENCE_LABEL } from './translations'; +import type { ContentReferenceNode } from '../content_reference_parser'; +import { PopoverReference } from './popover_reference'; +import { useKibana } from '../../../../common/lib/kibana'; + +interface Props { + contentReferenceNode: ContentReferenceNode; + securityAlertContentReference: SecurityAlertContentReference; +} + +export const SecurityAlertReference: React.FC = ({ + contentReferenceNode, + securityAlertContentReference, +}) => { + const { navigateToApp } = useKibana().services.application; + + const onClick = useCallback( + (e: React.MouseEvent) => { + e.preventDefault(); + navigateToApp('security', { + path: `alerts/redirect/${securityAlertContentReference.alertId}`, + openInNewTab: true, + }); + }, + [navigateToApp, securityAlertContentReference] + ); + + return ( + + {SECURITY_ALERT_REFERENCE_LABEL} + + ); +}; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/security_alerts_page_reference.tsx b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/security_alerts_page_reference.tsx new file mode 100644 index 0000000000000..e79333d4b15e8 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/security_alerts_page_reference.tsx @@ -0,0 +1,46 @@ +/* + * 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 { SecurityAlertsPageContentReference } from '@kbn/elastic-assistant-common'; +import React, { useCallback } from 'react'; +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'; + +interface Props { + contentReferenceNode: ContentReferenceNode; + securityAlertsPageContentReference: SecurityAlertsPageContentReference; +} + +export const SecurityAlertsPageReference: React.FC = ({ + contentReferenceNode, + securityAlertsPageContentReference, +}) => { + const { navigateToApp } = useKibana().services.application; + + const onClick = useCallback( + (e: React.MouseEvent) => { + e.preventDefault(); + navigateToApp('security', { + path: `alerts`, + openInNewTab: true, + }); + }, + [navigateToApp] + ); + + return ( + + {SECURITY_ALERTS_PAGE_REFERENCE_LABEL} + + ); +}; diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/translations.ts b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/translations.ts new file mode 100644 index 0000000000000..c5fe48787d4e9 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/components/translations.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 { i18n } from '@kbn/i18n'; + +export const SECURITY_ALERT_REFERENCE_LABEL = i18n.translate( + 'xpack.securitySolution.assistant.contentReferences.securityAlertReference.label', + { + defaultMessage: 'View alert', + } +); + +export const SECURITY_ALERTS_PAGE_REFERENCE_LABEL = i18n.translate( + 'xpack.securitySolution.assistant.contentReferences.securityAlertsPageReference.label', + { + defaultMessage: 'View alerts', + } +); + +export const KNOWLEDGE_BASE_ENTRY_REFERENCE_LABEL = i18n.translate( + 'xpack.securitySolution.assistant.contentReferences.knowledgeBaseEntryReference.label', + { + defaultMessage: 'Knowledge base entry', + } +); diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/content_reference_parser.test.ts b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/content_reference_parser.test.ts new file mode 100644 index 0000000000000..83c54e5d274a8 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/content_reference_parser.test.ts @@ -0,0 +1,325 @@ +/* + * 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 unified from 'unified'; +import markdown from 'remark-parse-no-trim'; +import type { Parent } from 'mdast'; +import { ContentReferenceParser } from './content_reference_parser'; + +describe('ContentReferenceParser', () => { + it('eats space preceding content reference', async () => { + const file = unified() + .use([[markdown, {}], ContentReferenceParser]) + .parse('Delete space after punctuation. {reference(example)}') as Parent; + + expect(file.children[0].children).toEqual( + expect.arrayContaining([ + expect.objectContaining({ type: 'text', value: 'Delete space after punctuation.' }), + expect.objectContaining({ type: 'contentReference' }), + ]) + ); + }); + + it('parses when there is no space preceding the content reference', async () => { + const file = unified() + .use([[markdown, {}], ContentReferenceParser]) + .parse('No preceding space.{reference(example)}') as Parent; + + expect(file.children[0].children).toEqual( + expect.arrayContaining([ + expect.objectContaining({ type: 'text', value: 'No preceding space.' }), + expect.objectContaining({ type: 'contentReference' }), + ]) + ); + }); + + it('handles single citation block', async () => { + const file = unified() + .use([[markdown, {}], ContentReferenceParser]) + .parse('Hello world {reference(example)} hello wolrd') as Parent; + + expect(file.children[0].children).toEqual([ + { + position: { + end: { + column: 12, + line: 1, + offset: 11, + }, + indent: [], + start: { + column: 1, + line: 1, + offset: 0, + }, + }, + type: 'text', + value: 'Hello world', + }, + { + contentReferenceBlock: '{reference(example)}', + contentReferenceCount: 1, + contentReferenceId: 'example', + position: { + end: { + column: 33, + line: 1, + offset: 32, + }, + indent: [], + start: { + column: 12, + line: 1, + offset: 11, + }, + }, + type: 'contentReference', + }, + { + position: { + end: { + column: 45, + line: 1, + offset: 44, + }, + indent: [], + start: { + column: 33, + line: 1, + offset: 32, + }, + }, + type: 'text', + value: ' hello wolrd', + }, + ]); + }); + + it('handles multiple citation blocks with different referenceIds', async () => { + const file = unified() + .use([[markdown, {}], ContentReferenceParser]) + .parse('Hello world {reference(example)} hello world {reference(example2)}') as Parent; + + expect(file.children[0].children).toEqual([ + { + position: { + end: { + column: 12, + line: 1, + offset: 11, + }, + indent: [], + start: { + column: 1, + line: 1, + offset: 0, + }, + }, + type: 'text', + value: 'Hello world', + }, + { + contentReferenceBlock: '{reference(example)}', + contentReferenceCount: 1, + contentReferenceId: 'example', + position: { + end: { + column: 33, + line: 1, + offset: 32, + }, + indent: [], + start: { + column: 12, + line: 1, + offset: 11, + }, + }, + type: 'contentReference', + }, + { + position: { + end: { + column: 45, + line: 1, + offset: 44, + }, + indent: [], + start: { + column: 33, + line: 1, + offset: 32, + }, + }, + type: 'text', + value: ' hello world', + }, + { + contentReferenceBlock: '{reference(example2)}', + contentReferenceCount: 2, + contentReferenceId: 'example2', + position: { + end: { + column: 67, + line: 1, + offset: 66, + }, + indent: [], + start: { + column: 45, + line: 1, + offset: 44, + }, + }, + type: 'contentReference', + }, + ]); + }); + + it('handles multiple citation blocks with same referenceIds', async () => { + const file = unified() + .use([[markdown, {}], ContentReferenceParser]) + .parse('Hello world {reference(example)} hello world {reference(example)}') as Parent; + + expect(file.children[0].children).toEqual([ + { + position: { + end: { + column: 12, + line: 1, + offset: 11, + }, + indent: [], + start: { + column: 1, + line: 1, + offset: 0, + }, + }, + type: 'text', + value: 'Hello world', + }, + { + contentReferenceBlock: '{reference(example)}', + contentReferenceCount: 1, + contentReferenceId: 'example', + position: { + end: { + column: 33, + line: 1, + offset: 32, + }, + indent: [], + start: { + column: 12, + line: 1, + offset: 11, + }, + }, + type: 'contentReference', + }, + { + position: { + end: { + column: 45, + line: 1, + offset: 44, + }, + indent: [], + start: { + column: 33, + line: 1, + offset: 32, + }, + }, + type: 'text', + value: ' hello world', + }, + { + contentReferenceBlock: '{reference(example)}', + contentReferenceCount: 1, + contentReferenceId: 'example', + position: { + end: { + column: 66, + line: 1, + offset: 65, + }, + indent: [], + start: { + column: 45, + line: 1, + offset: 44, + }, + }, + type: 'contentReference', + }, + ]); + }); + + it('handles partial citation blocks', async () => { + const file = unified() + .use([[markdown, {}], ContentReferenceParser]) + .parse('Hello world {reference(example)} hello world {reference(') as Parent; + + expect(file.children[0].children).toEqual([ + { + position: { + end: { + column: 12, + line: 1, + offset: 11, + }, + indent: [], + start: { + column: 1, + line: 1, + offset: 0, + }, + }, + type: 'text', + value: 'Hello world', + }, + { + contentReferenceBlock: '{reference(example)}', + contentReferenceCount: 1, + contentReferenceId: 'example', + position: { + end: { + column: 33, + line: 1, + offset: 32, + }, + indent: [], + start: { + column: 12, + line: 1, + offset: 11, + }, + }, + type: 'contentReference', + }, + { + position: { + end: { + column: 57, + line: 1, + offset: 56, + }, + indent: [], + start: { + column: 33, + line: 1, + offset: 32, + }, + }, + type: 'text', + value: ' hello world {reference(', + }, + ]); + }); +}); diff --git a/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/content_reference_parser.ts b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/content_reference_parser.ts new file mode 100644 index 0000000000000..74ff7803ce3e0 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/content_reference/content_reference_parser.ts @@ -0,0 +1,123 @@ +/* + * 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 { RemarkTokenizer } from '@elastic/eui'; +import type { ContentReferenceBlock } from '@kbn/elastic-assistant-common'; +import type { Plugin } from 'unified'; +import type { Node } from 'unist'; + +export interface ContentReferenceNode extends Node { + type: 'contentReference'; + contentReferenceId: string; + contentReferenceCount: number; + contentReferenceBlock: ContentReferenceBlock; +} + +/** + * Parses `{reference(contentReferenceId)}` or ` {reference(contentReferenceId)}` (notice space prefix) into ContentReferenceNode + */ +export const ContentReferenceParser: Plugin = function ContentReferenceParser() { + const Parser = this.Parser; + const tokenizers = Parser.prototype.inlineTokenizers; + const methods = Parser.prototype.inlineMethods; + + let currentContentReferenceCount = 1; + const contentReferenceCounts: Record = {}; + + const tokenizeCustomCitation: RemarkTokenizer = function tokenizeCustomCitation( + eat, + value, + silent + ) { + const [match] = value.match(/^\s?{reference/) || []; + if (!match) return false; + + if (value.includes('\n')) return false; + + if (value[match.length] !== '(') return false; + + let index = match.length; + + function readArg(open: string, close: string) { + if (value[index] !== open) return ''; + index++; + + let body = ''; + let openBrackets = 0; + + for (; index < value.length; index++) { + const char = value[index]; + if (char === close && openBrackets === 0) { + index++; + + return body; + } else if (char === close) { + openBrackets--; + } else if (char === open) { + openBrackets++; + } + + body += char; + } + + return ''; + } + + const contentReferenceId = readArg('(', ')'); + + const closeChar = value[index++]; + if (closeChar !== '}') return false; + + const now = eat.now(); + + if (!contentReferenceId) { + this.file.info('No content reference id found', { + line: now.line, + column: now.column + match.length + 1, + }); + } + + if (!contentReferenceId) { + return false; + } + + if (silent) { + return true; + } + + now.column += match.length + 1; + now.offset += match.length + 1; + + const contentReferenceBlock: ContentReferenceBlock = `{reference(${contentReferenceId})}`; + + const getContentReferenceCount = (id: string) => { + if (id in contentReferenceCounts) { + return contentReferenceCounts[id]; + } + contentReferenceCounts[id] = currentContentReferenceCount++; + return contentReferenceCounts[id]; + }; + + const toEat = `${match.startsWith(' ') ? ' ' : ''}${contentReferenceBlock}`; + + return eat(toEat)({ + type: 'contentReference', + contentReferenceId, + contentReferenceCount: getContentReferenceCount(contentReferenceId), + contentReferenceBlock, + } as ContentReferenceNode); + }; + + tokenizeCustomCitation.notInLink = true; + + tokenizeCustomCitation.locator = (value, fromIndex) => { + return 1 + (value.substring(fromIndex).match(/\s?{reference/)?.index ?? -2); + }; + + tokenizers.contentReference = tokenizeCustomCitation; + methods.splice(methods.indexOf('text'), 0, 'contentReference'); +}; 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 8b9dc2f2c80f4..58634c9c53d1c 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 @@ -61,6 +61,8 @@ export const getComments: GetAssistantMessages = ({ currentUserAvatar, setIsStreaming, systemPromptContent, + contentReferencesVisible, + contentReferencesEnabled, }) => { if (!currentConversation) return []; @@ -174,6 +176,9 @@ export const getComments: GetAssistantMessages = ({ children: ( void; content?: string; + contentReferences?: ContentReferences; + contentReferencesVisible?: boolean; + contentReferencesEnabled?: boolean; isError?: boolean; isFetching?: boolean; isControlsEnabled?: boolean; @@ -31,6 +35,9 @@ interface Props { export const StreamComment = ({ abortStream, content, + contentReferences, + contentReferencesVisible = true, + contentReferencesEnabled = false, index, isControlsEnabled = false, isError = false, @@ -106,8 +113,11 @@ export const StreamComment = ({ } error={error ? new Error(error) : undefined} 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 ff526c33497c3..1e1e19e8cf9c8 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 @@ -19,15 +19,21 @@ import { import { css } from '@emotion/css'; import classNames from 'classnames'; import type { Code, InlineCode, Parent, Text } from 'mdast'; -import React from 'react'; +import React, { useMemo } from 'react'; import { euiThemeVars } from '@kbn/ui-theme'; import type { Node } from 'unist'; +import type { ContentReferences } from '@kbn/elastic-assistant-common'; import { customCodeBlockLanguagePlugin } from '../custom_codeblock/custom_codeblock_markdown_plugin'; import { CustomCodeBlock } from '../custom_codeblock/custom_code_block'; +import { ContentReferenceParser } from '../content_reference/content_reference_parser'; +import { contentReferenceComponentFactory } from '../content_reference/components/content_reference_component_factory'; interface Props { content: string; + contentReferences?: ContentReferences; + contentReferencesVisible: boolean; + contentReferencesEnabled: boolean; index: number; loading: boolean; ['data-test-subj']?: string; @@ -99,7 +105,19 @@ const loadingCursorPlugin = () => { }; }; -const getPluginDependencies = () => { +interface GetPluginDependencies { + contentReferences?: ContentReferences; + loading: boolean; + contentReferencesVisible: boolean; + contentReferencesEnabled: boolean; +} + +const getPluginDependencies = ({ + contentReferences, + contentReferencesVisible, + loading, + contentReferencesEnabled, +}: GetPluginDependencies) => { const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); const processingPlugins = getDefaultEuiMarkdownProcessingPlugins(); @@ -108,6 +126,15 @@ const getPluginDependencies = () => { processingPlugins[1][1].components = { ...components, + ...(contentReferencesEnabled + ? { + contentReference: contentReferenceComponentFactory({ + contentReferences, + contentReferencesVisible, + loading, + }), + } + : {}), cursor: Cursor, customCodeBlock: (props) => { return ( @@ -139,17 +166,39 @@ const getPluginDependencies = () => { }; return { - parsingPluginList: [loadingCursorPlugin, customCodeBlockLanguagePlugin, ...parsingPlugins], + parsingPluginList: [ + loadingCursorPlugin, + customCodeBlockLanguagePlugin, + ...parsingPlugins, + ...(contentReferencesEnabled ? [ContentReferenceParser] : []), + ], processingPluginList: processingPlugins, }; }; -export function MessageText({ loading, content, index, 'data-test-subj': dataTestSubj }: Props) { +export function MessageText({ + loading, + content, + contentReferences, + contentReferencesVisible, + contentReferencesEnabled, + index, + 'data-test-subj': dataTestSubj, +}: Props) { const containerClassName = css` overflow-wrap: anywhere; `; - const { parsingPluginList, processingPluginList } = getPluginDependencies(); + const { parsingPluginList, processingPluginList } = useMemo( + () => + getPluginDependencies({ + contentReferences, + contentReferencesVisible, + contentReferencesEnabled, + loading, + }), + [contentReferences, contentReferencesVisible, contentReferencesEnabled, loading] + ); 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 814a00853927f..946447e2815e0 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 @@ -12,6 +12,8 @@ import { loggerMock } from '@kbn/logging-mocks'; import { ALERT_COUNTS_TOOL } from './alert_counts_tool'; import type { RetrievalQAChain } from 'langchain/chains'; import type { ExecuteConnectorRequestBody } from '@kbn/elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen'; +import type { ContentReferencesStore } from '@kbn/elastic-assistant-common'; +import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; describe('AlertCountsTool', () => { const alertsIndexPattern = 'alerts-index'; @@ -30,10 +32,12 @@ describe('AlertCountsTool', () => { const isEnabledKnowledgeBase = true; const chain = {} as unknown as RetrievalQAChain; const logger = loggerMock.create(); + const contentReferencesStore = contentReferencesStoreFactoryMock(); const rest = { isEnabledKnowledgeBase, chain, logger, + contentReferencesStore, }; beforeEach(() => { @@ -158,6 +162,43 @@ describe('AlertCountsTool', () => { }); }); + it('includes citations', async () => { + const tool: DynamicTool = ALERT_COUNTS_TOOL.getTool({ + alertsIndexPattern, + esClient, + replacements, + request, + ...rest, + }) as DynamicTool; + + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + const reference = creator({ id: 'exampleContentReferenceId' }); + expect(reference.type).toEqual('SecurityAlertsPage'); + return reference; + } + ); + + const result = await tool.func(''); + + 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: false, + }) 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 8bc2906d3fde9..0201eeea1578d 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 @@ -10,6 +10,7 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from '@kbn/zod'; import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { contentReferenceString, securityAlertsPageReference } from '@kbn/elastic-assistant-common'; import { getAlertsCountQuery } from './get_alert_counts_query'; import { APP_UI_ID } from '../../../../common'; @@ -30,7 +31,8 @@ export const ALERT_COUNTS_TOOL: AssistantTool = { }, getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { alertsIndexPattern, esClient } = params as AlertCountsToolParams; + const { alertsIndexPattern, esClient, contentReferencesStore } = + params as AlertCountsToolParams; return new DynamicStructuredTool({ name: 'AlertCountsTool', description: ALERT_COUNTS_TOOL_DESCRIPTION, @@ -38,8 +40,15 @@ 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)); - return JSON.stringify(result); + const reference = alertsCountReference + ? `\n${contentReferenceString(alertsCountReference)}` + : ''; + + return `${JSON.stringify(result)}${reference}`; }, tags: ['alerts', 'alerts-count'], }); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.test.ts index 10b1fa21daefe..453cd3e4d27c7 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.test.ts @@ -14,6 +14,7 @@ import type { ExecuteConnectorRequestBody } from '@kbn/elastic-assistant-common/ import { loggerMock } from '@kbn/logging-mocks'; import { getPromptSuffixForOssModel } from './common'; import type { InferenceServerStart } from '@kbn/inference-plugin/server'; +import type { ContentReferencesStore } from '@kbn/elastic-assistant-common'; describe('NaturalLanguageESQLTool', () => { const chain = {} as RetrievalQAChain; @@ -33,6 +34,7 @@ describe('NaturalLanguageESQLTool', () => { const logger = loggerMock.create(); const inference = {} as InferenceServerStart; const connectorId = 'fake-connector'; + const contentReferencesStore = {} as ContentReferencesStore; const rest = { chain, esClient, @@ -41,6 +43,7 @@ describe('NaturalLanguageESQLTool', () => { inference, connectorId, isEnabledKnowledgeBase: true, + contentReferencesStore, }; describe('isSupported', () => { 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 new file mode 100644 index 0000000000000..4a6683624b980 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts @@ -0,0 +1,89 @@ +/* + * 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 { DynamicStructuredTool } from '@langchain/core/tools'; +import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from './knowledge_base_retrieval_tool'; +import type { AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import type { + ContentReferencesStore, + KnowledgeBaseEntryContentReference, +} from '@kbn/elastic-assistant-common'; +import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import { loggerMock } from '@kbn/logging-mocks'; +import { Document } from 'langchain/document'; + +describe('KnowledgeBaseRetievalTool', () => { + const logger = loggerMock.create(); + const contentReferencesStore = contentReferencesStoreFactoryMock(); + const getKnowledgeBaseDocumentEntries = jest.fn(); + const kbDataClient = { getKnowledgeBaseDocumentEntries }; + const defaultArgs = { + isEnabledKnowledgeBase: true, + contentReferencesStore, + kbDataClient, + logger, + } as unknown as AssistantToolParams; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('DynamicStructuredTool', () => { + it('includes citations', async () => { + const tool = KNOWLEDGE_BASE_RETRIEVAL_TOOL.getTool(defaultArgs) as DynamicStructuredTool; + + getKnowledgeBaseDocumentEntries.mockResolvedValue([ + new Document({ + id: 'exampleId', + pageContent: 'text', + metadata: { + name: 'exampleName', + }, + }), + ] as Document[]); + + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + const reference = creator({ id: 'exampleContentReferenceId' }); + expect(reference.type).toEqual('KnowledgeBaseEntry'); + expect((reference as KnowledgeBaseEntryContentReference).knowledgeBaseEntryId).toEqual( + 'exampleId' + ); + expect((reference as KnowledgeBaseEntryContentReference).knowledgeBaseEntryName).toEqual( + 'exampleName' + ); + return reference; + } + ); + + const result = await tool.func({ query: 'What is my favourite food' }); + + 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: false, + }) 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 4369f85a83c25..6bb7455d332ca 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 @@ -9,6 +9,9 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from '@kbn/zod'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; import type { AIAssistantKnowledgeBaseDataClient } from '@kbn/elastic-assistant-plugin/server/ai_assistant_data_clients/knowledge_base'; +import { Document } from 'langchain/document'; +import type { ContentReferencesStore } from '@kbn/elastic-assistant-common'; +import { knowledgeBaseReference, contentReferenceBlock } from '@kbn/elastic-assistant-common'; import { APP_UI_ID } from '../../../../common'; export interface KnowledgeBaseRetrievalToolParams extends AssistantToolParams { @@ -31,7 +34,8 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { kbDataClient, logger } = params as KnowledgeBaseRetrievalToolParams; + const { kbDataClient, logger, contentReferencesStore } = + params as KnowledgeBaseRetrievalToolParams; if (kbDataClient == null) return null; return new DynamicStructuredTool({ @@ -51,6 +55,10 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { required: false, }); + if (contentReferencesStore) { + return JSON.stringify(docs.map(enrichDocument(contentReferencesStore))); + } + return JSON.stringify(docs); }, tags: ['knowledge-base'], @@ -58,3 +66,22 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { }) as unknown as DynamicStructuredTool; }, }; + +function enrichDocument(contentReferencesStore: ContentReferencesStore) { + return (document: Document>) => { + if (document.id == null) { + return document; + } + const documentId = document.id; + const reference = contentReferencesStore.add((p) => + knowledgeBaseReference(p.id, document.metadata.name, documentId) + ); + return new Document({ + ...document, + metadata: { + ...document.metadata, + citation: contentReferenceBlock(reference), + }, + }); + }; +} 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 45587b65f5f4c..5cdc59a5fca94 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 @@ -14,9 +14,19 @@ import type { RetrievalQAChain } from 'langchain/chains'; import { mockAlertsFieldsApi } from '@kbn/elastic-assistant-plugin/server/__mocks__/alerts'; import type { ExecuteConnectorRequestBody } from '@kbn/elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen'; import { loggerMock } from '@kbn/logging-mocks'; +import type { + ContentReferencesStore, + SecurityAlertContentReference, +} from '@kbn/elastic-assistant-common'; +import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; const MAX_SIZE = 10000; +jest.mock('@kbn/elastic-assistant-common', () => ({ + transformRawData: jest.fn(() => 'transformedData'), + ...jest.requireActual('@kbn/elastic-assistant-common'), +})); + describe('OpenAndAcknowledgedAlertsTool', () => { const alertsIndexPattern = 'alerts-index'; const esClient = { @@ -34,11 +44,13 @@ describe('OpenAndAcknowledgedAlertsTool', () => { const isEnabledKnowledgeBase = true; const chain = {} as unknown as RetrievalQAChain; const logger = loggerMock.create(); + const contentReferencesStore = contentReferencesStoreFactoryMock(); const rest = { isEnabledKnowledgeBase, esClient, chain, logger, + contentReferencesStore, }; const anonymizationFields = [ @@ -222,6 +234,60 @@ describe('OpenAndAcknowledgedAlertsTool', () => { }); }); + it('includes citations', async () => { + const tool: DynamicTool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ + alertsIndexPattern, + anonymizationFields, + onNewReplacements: jest.fn(), + replacements, + request, + size: request.body.size, + ...rest, + }) as DynamicTool; + + (esClient.search as jest.Mock).mockResolvedValue({ + hits: { + hits: [{ _id: 4 }], + }, + }); + + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + const reference = creator({ id: 'exampleContentReferenceId' }); + expect(reference.type).toEqual('SecurityAlert'); + expect((reference as SecurityAlertContentReference).alertId).toEqual(4); + return reference; + } + ); + + const result = await tool.func(''); + + 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: false, + }) 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 cab015183f4a2..a8fb0f0064d24 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 @@ -8,11 +8,13 @@ import type { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; import type { Replacements } from '@kbn/elastic-assistant-common'; import { + securityAlertReference, getAnonymizedValue, getOpenAndAcknowledgedAlertsQuery, getRawDataOrDefault, sizeIsOutOfRange, transformRawData, + contentReferenceBlock, } from '@kbn/elastic-assistant-common'; import { DynamicStructuredTool } from '@langchain/core/tools'; import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; @@ -56,6 +58,7 @@ export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = { onNewReplacements, replacements, size, + contentReferencesStore, } = params as OpenAndAcknowledgedAlertsToolParams; return new DynamicStructuredTool({ name: 'OpenAndAcknowledgedAlertsTool', @@ -80,15 +83,24 @@ export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = { }; return JSON.stringify( - result.hits?.hits?.map((x) => - transformRawData({ + result.hits?.hits?.map((x) => { + const transformed = transformRawData({ anonymizationFields, currentReplacements: localReplacements, // <-- the latest local replacements getAnonymizedValue, onNewReplacements: localOnNewReplacements, // <-- the local callback rawData: getRawDataOrDefault(x.fields), - }) - ) + }); + const hitId = x._id; + const citation = + hitId && + contentReferencesStore && + `\nCitation,${contentReferenceBlock( + contentReferencesStore.add((p) => securityAlertReference(p.id, hitId)) + )}`; + + return `${transformed}${citation ?? ''}`; + }) ); }, tags: ['alerts', 'open-and-acknowledged-alerts'], 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 737e82d3837d2..11a95bd29b594 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 @@ -12,7 +12,15 @@ import type { KibanaRequest } from '@kbn/core-http-server'; import type { ExecuteConnectorRequestBody } from '@kbn/elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen'; import { loggerMock } from '@kbn/logging-mocks'; import { PRODUCT_DOCUMENTATION_TOOL } from './product_documentation_tool'; -import type { LlmTasksPluginStart } from '@kbn/llm-tasks-plugin/server'; +import type { + LlmTasksPluginStart, + RetrieveDocumentationResultDoc, +} from '@kbn/llm-tasks-plugin/server'; +import type { + ContentReferencesStore, + ProductDocumentationContentReference, +} from '@kbn/elastic-assistant-common'; +import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; describe('ProductDocumentationTool', () => { const chain = {} as RetrievalQAChain; @@ -27,6 +35,7 @@ describe('ProductDocumentationTool', () => { retrieveDocumentationAvailable: jest.fn(), } as LlmTasksPluginStart; const connectorId = 'fake-connector'; + const contentReferencesStore = contentReferencesStoreFactoryMock(); const defaultArgs = { chain, esClient, @@ -35,6 +44,7 @@ describe('ProductDocumentationTool', () => { llmTasks, connectorId, isEnabledKnowledgeBase: true, + contentReferencesStore, }; beforeEach(() => { jest.clearAllMocks(); @@ -89,5 +99,79 @@ describe('ProductDocumentationTool', () => { functionCalling: 'auto', }); }); + + it('includes citations', async () => { + const tool = PRODUCT_DOCUMENTATION_TOOL.getTool(defaultArgs) as DynamicStructuredTool; + + (retrieveDocumentation as jest.Mock).mockResolvedValue({ + documents: [ + { + title: 'exampleTitle', + url: 'exampleUrl', + content: 'exampleContent', + summarized: false, + }, + ] as RetrieveDocumentationResultDoc[], + }); + + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + const reference = creator({ id: 'exampleContentReferenceId' }); + expect(reference.type).toEqual('ProductDocumentation'); + expect((reference as ProductDocumentationContentReference).title).toEqual('exampleTitle'); + expect((reference as ProductDocumentationContentReference).url).toEqual('exampleUrl'); + return reference; + } + ); + + const result = await tool.func({ query: 'What is Kibana Security?', product: 'kibana' }); + + expect(result).toEqual({ + content: { + documents: [ + { + citation: '{reference(exampleContentReferenceId)}', + content: 'exampleContent', + title: 'exampleTitle', + url: 'exampleUrl', + summarized: false, + }, + ], + }, + }); + }); + + it('does not include citations if contentReferencesStore is false', async () => { + const tool = PRODUCT_DOCUMENTATION_TOOL.getTool({ + ...defaultArgs, + contentReferencesStore: false, + }) 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 9f3dd1c05ce1c..da3b1ba492693 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 @@ -9,6 +9,12 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from '@kbn/zod'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { + contentReferenceBlock, + productDocumentationReference, +} from '@kbn/elastic-assistant-common'; +import type { ContentReferencesStore } from '@kbn/elastic-assistant-common'; +import type { RetrieveDocumentationResultDoc } from '@kbn/llm-tasks-plugin/server'; import { APP_UI_ID } from '../../../../common'; const toolDetails = { @@ -26,7 +32,8 @@ export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { connectorId, llmTasks, request } = params as AssistantToolParams; + const { connectorId, llmTasks, request, contentReferencesStore } = + params as AssistantToolParams; // This check is here in order to satisfy TypeScript if (llmTasks == null || connectorId == null) return null; @@ -65,6 +72,16 @@ export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = { functionCalling: 'auto', }); + if (contentReferencesStore) { + const enrichedDocuments = response.documents.map(enrichDocument(contentReferencesStore)); + + return { + content: { + documents: enrichedDocuments, + }, + }; + } + return { content: { documents: response.documents, @@ -76,3 +93,19 @@ export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = { }) as unknown as DynamicStructuredTool; }, }; + +type EnrichedDocument = RetrieveDocumentationResultDoc & { + citation: string; +}; + +const enrichDocument = (contentReferencesStore: ContentReferencesStore) => { + return (document: RetrieveDocumentationResultDoc): EnrichedDocument => { + const reference = contentReferencesStore.add((p) => + productDocumentationReference(p.id, document.title, document.url) + ); + return { + ...document, + citation: contentReferenceBlock(reference), + }; + }; +}; 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 new file mode 100644 index 0000000000000..c8bba6e3a7113 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts @@ -0,0 +1,65 @@ +/* + * 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 { DynamicStructuredTool } from '@langchain/core/tools'; +import { SECURITY_LABS_KNOWLEDGE_BASE_TOOL } from './security_labs_tool'; +import type { + ContentReferencesStore, + KnowledgeBaseEntryContentReference, +} from '@kbn/elastic-assistant-common'; +import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import type { AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; + +describe('SecurityLabsTool', () => { + const contentReferencesStore = contentReferencesStoreFactoryMock(); + const getKnowledgeBaseDocumentEntries = jest.fn().mockResolvedValue([]); + const kbDataClient = { getKnowledgeBaseDocumentEntries }; + const defaultArgs = { + isEnabledKnowledgeBase: true, + contentReferencesStore, + kbDataClient, + } as unknown as AssistantToolParams; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('DynamicStructuredTool', () => { + it('includes citations', async () => { + const tool = SECURITY_LABS_KNOWLEDGE_BASE_TOOL.getTool(defaultArgs) as DynamicStructuredTool; + + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + const reference = creator({ id: 'exampleContentReferenceId' }); + expect(reference.type).toEqual('KnowledgeBaseEntry'); + expect((reference as KnowledgeBaseEntryContentReference).knowledgeBaseEntryId).toEqual( + 'securityLabsId' + ); + expect((reference as KnowledgeBaseEntryContentReference).knowledgeBaseEntryName).toEqual( + 'Elastic Security Labs content' + ); + return reference; + } + ); + + const result = await tool.func({ query: 'What is Kibana Security?', product: 'kibana' }); + + 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: false, + }) 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 c94b14066947b..b19fa15747e67 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 @@ -10,6 +10,7 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from '@kbn/zod'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; import { SECURITY_LABS_RESOURCE } from '@kbn/elastic-assistant-plugin/server/routes/knowledge_base/constants'; +import { knowledgeBaseReference, contentReferenceString } from '@kbn/elastic-assistant-common'; import { APP_UI_ID } from '../../../../common'; const toolDetails = { @@ -28,7 +29,7 @@ export const SECURITY_LABS_KNOWLEDGE_BASE_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { kbDataClient } = params as AssistantToolParams; + const { kbDataClient, contentReferencesStore } = params as AssistantToolParams; if (kbDataClient == null) return null; return new DynamicStructuredTool({ @@ -46,8 +47,18 @@ export const SECURITY_LABS_KNOWLEDGE_BASE_TOOL: AssistantTool = { kbResource: SECURITY_LABS_RESOURCE, query: input.question, }); + + const reference = + contentReferencesStore && + contentReferencesStore.add((p) => + knowledgeBaseReference(p.id, 'Elastic Security Labs content', 'securityLabsId') + ); + // TODO: Token pruning - return JSON.stringify(docs).substring(0, 20000); + const result = JSON.stringify(docs).substring(0, 20000); + + const citation = reference ? `\n${contentReferenceString(reference)}` : ''; + return `${result}${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 f912a78fa20b8..73cf66816cebd 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/plugin.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/plugin.ts @@ -596,6 +596,7 @@ export class Plugin implements ISecuritySolutionPlugin { const features = { assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, attackDiscoveryAlertFiltering: config.experimentalFeatures.attackDiscoveryAlertFiltering, + contentReferencesEnabled: config.experimentalFeatures.contentReferencesEnabled, }; plugins.elasticAssistant.registerFeatures(APP_UI_ID, features); plugins.elasticAssistant.registerFeatures('management', features); diff --git a/x-pack/test/security_solution_cypress/cypress/screens/ai_assistant.ts b/x-pack/test/security_solution_cypress/cypress/screens/ai_assistant.ts index 6935eb993e3f0..60d9c8b249e05 100644 --- a/x-pack/test/security_solution_cypress/cypress/screens/ai_assistant.ts +++ b/x-pack/test/security_solution_cypress/cypress/screens/ai_assistant.ts @@ -41,7 +41,6 @@ export const QUICK_PROMPT_TITLE_INPUT = export const QUICK_PROMPT_BADGE = (b: string) => `[data-test-subj="quickPrompt-${b}"]`; export const QUICK_PROMPT_BODY_INPUT = '[data-test-subj="quick-prompt-prompt"]'; export const SEND_TO_TIMELINE_BUTTON = '[data-test-subj="sendToTimelineEmptyButton"]'; -export const SHOW_ANONYMIZED_BUTTON = '[data-test-subj="showAnonymizedValues"]'; export const SUBMIT_CHAT = '[data-test-subj="submit-chat"]'; export const SYSTEM_PROMPT = '[data-test-subj="promptSuperSelect"]'; export const SYSTEM_PROMPT_BODY_INPUT = '[data-test-subj="systemPromptModalPromptText"]'; diff --git a/x-pack/test/security_solution_cypress/cypress/tasks/assistant.ts b/x-pack/test/security_solution_cypress/cypress/tasks/assistant.ts index 5f030c61de65a..969e2bb9d2d4a 100644 --- a/x-pack/test/security_solution_cypress/cypress/tasks/assistant.ts +++ b/x-pack/test/security_solution_cypress/cypress/tasks/assistant.ts @@ -41,7 +41,6 @@ import { PROMPT_CONTEXT_SELECTOR, QUICK_PROMPT_BADGE, ADD_NEW_CONNECTOR, - SHOW_ANONYMIZED_BUTTON, SEND_TO_TIMELINE_BUTTON, } from '../screens/ai_assistant'; import { TOASTER } from '../screens/alerts_detection_rules'; @@ -219,7 +218,6 @@ const assertConversationTitleReadOnly = () => { export const assertConversationReadOnly = () => { assertConversationTitleReadOnly(); cy.get(ADD_NEW_CONNECTOR).should('be.disabled'); - cy.get(SHOW_ANONYMIZED_BUTTON).should('be.disabled'); cy.get(CHAT_CONTEXT_MENU).should('be.disabled'); cy.get(FLYOUT_NAV_TOGGLE).should('be.disabled'); cy.get(NEW_CHAT).should('be.disabled'); From f5b11d9e6b2a23b247be3c6055b9360f0f7bfa36 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 29 Jan 2025 02:33:31 +0000 Subject: [PATCH 2/4] [CI] Auto-commit changed files from 'node scripts/eslint --no-cache --fix' --- .../impl/assistant/assistant_header/index.tsx | 8 +------ .../settings_context_menu.tsx | 22 +++++++++---------- 2 files changed, 12 insertions(+), 18 deletions(-) diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx index d99fcdc96e0a4..bb49f791b8b22 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx @@ -7,13 +7,7 @@ import React, { useMemo, useCallback } from 'react'; import { QueryObserverResult, RefetchOptions, RefetchQueryFilters } from '@tanstack/react-query'; -import { - EuiFlexGroup, - EuiFlexItem, - EuiButtonIcon, - EuiPanel, - EuiSkeletonTitle, -} from '@elastic/eui'; +import { EuiFlexGroup, EuiFlexItem, EuiButtonIcon, EuiPanel, EuiSkeletonTitle } from '@elastic/eui'; import { css } from '@emotion/react'; import { euiThemeVars } from '@kbn/ui-theme'; import { isEmpty } from 'lodash'; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx index 8efea263b445e..5bddc794ea24b 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx @@ -213,17 +213,17 @@ export const SettingsContextMenu: React.FC = React.memo( )} - {i18n.RESET_CONVERSATION} - + aria-label={'clear-chat'} + key={'clear-chat'} + onClick={showDestroyModal} + icon={'refresh'} + data-test-subj={'clear-chat'} + css={css` + color: ${euiThemeVars.euiColorDanger}; + `} + > + {i18n.RESET_CONVERSATION} + , ], [ From 703be57b504f34004bc3df3f2cf148880868b8de Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 29 Jan 2025 05:00:54 +0100 Subject: [PATCH 3/4] fix lint --- .../shared/kbn-elastic-assistant/impl/assistant/index.tsx | 1 - .../settings/settings_context_menu/settings_context_menu.tsx | 3 --- 2 files changed, 4 deletions(-) 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 7860f197bd832..07d9097176344 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 @@ -428,7 +428,6 @@ const AssistantComponent: React.FC = ({ currentUserAvatar, currentSystemPrompt?.content, contentReferencesVisible, - euiThemeVars.euiSizeL, selectedPromptContextsCount, contentReferencesEnabled, ] diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx index 5bddc794ea24b..355b65b2e4315 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx @@ -231,7 +231,6 @@ export const SettingsContextMenu: React.FC = React.memo( onChangeContentReferencesVisible, showAnonymizedValues, onChangeShowAnonymizedValues, - euiThemeVars.euiColorDanger, handleNavigateToAnonymization, handleNavigateToKnowledgeBase, handleNavigateToSettings, @@ -239,8 +238,6 @@ export const SettingsContextMenu: React.FC = React.memo( knowledgeBase.latestAlerts, showDestroyModal, contentReferencesEnabled, - euiThemeVars.euiSizeM, - euiThemeVars.euiSizeXS, ] ); From 7cebc01d9a62c33194861395af7f5c9fca94daa4 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 29 Jan 2025 05:37:10 +0100 Subject: [PATCH 4/4] fix types --- .../impl/assistant/assistant_header/index.tsx | 2 -- 1 file changed, 2 deletions(-) diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx index bb49f791b8b22..57e92adead49e 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx @@ -68,8 +68,6 @@ export const AssistantHeader: React.FC = ({ isAssistantEnabled, refetchPrompts, }) => { - const { euiTheme } = useEuiTheme(); - const selectedConnectorId = useMemo( () => selectedConversation?.apiConfig?.connectorId, [selectedConversation?.apiConfig?.connectorId]