[Security Assistant] Automatically Install Knowledge Base #182763

spong · 2024-05-06T23:19:09Z

Summary

This PR is Phase 1 of the Knowledge Base work for 8.15, which includes automatically setting up the Knowledge base (this PR), introducing new generic KB tools for recall/retrieval, a CRUD API for for managing Knowledge Base Entries, and a basic UI for Knowledge Base Entry management (all captured in this issue). Once complete, this will also provide the opportunity to remove the !isEnabledKnowledgeBase code paths, directing all interactions through our LangChain Agent pipeline.

This PR sets the ground work for all of the above by moving ELSER setup and Knowledge Base data management to use the new AssistantDataClient architecture used for Conversations, AnonymizationFields and Prompts.

This feature is currently behind the assistantKnowledgeBaseByDefault experimental feature flag, which can be enabled by adding the following to your kibana.dev.yml:

xpack.securitySolution.enableExperimental:
  - 'assistantKnowledgeBaseByDefault'

Once enabled, an Install Knowledge Base button will be shown when starting a new conversation. Note: UX is still under development.

Useful Dev Tools Queries

The new assistantKnowledgeBaseByDefault flows are quite resilient, and so everything should function as expected even if one piece of the puzzle is missing or incomplete. Here are some dev tool queries to check and delete individual resources, which is nice for testing. For instance, you can nuke the ingest pipeline, or ELSER, and the Install KB button will appear and function as intended.

Note

Since the existing API's were used, with forked logic for the assistantKnowledgeBaseByDefault FF, the existing KB Settings UI still functions as expected, and can be used for deleting and re-initializing the KB. This functionality will most likely go away with updates to the KB UI, but is nice for testing in the interim.

Useful Dev Tools Queries

// New KB
GET /_ingest/pipeline/.kibana-elastic-ai-assistant-ingest-pipeline-knowledge-base
GET /_index_template/.kibana-elastic-ai-assistant-index-template-knowledge-base
GET /_data_stream/.kibana-elastic-ai-assistant-knowledge-base-default/
GET .kibana-elastic-ai-assistant-knowledge-base-default/_count
GET .kibana-elastic-ai-assistant-knowledge-base-default/_mapping
GET .kibana-elastic-ai-assistant-knowledge-base-default/_search
{ "size": 1000 }

// MSearch for if ES|QLKB docs exists
GET .kibana-elastic-ai-assistant-knowledge-base-default/_msearch
{}
{"query":{"bool":{"must_not":[{"term":{"metadata.kbResource":"esql"}},{"term":{"metadata.required":true}}],"must":[{"text_expansion":{"vector.tokens":{"model_id":".elser_model_2","model_text":"You can chain processing commands, separated by a pipe character: `|`."}}}]}},"size":10}
{}
{"query":{"bool":{"must":[{"term":{"metadata.kbResource":"esql"}},{"term":{"metadata.required":true}}]}},"size":10000}


// Other DataClient Assets
GET .kibana-elastic-ai-assistant-anonymization-fields-default/_search
{ "size": 1000 }
GET .kibana-elastic-ai-assistant-conversations-default/_search
{ "size": 1000 }
GET .kibana-elastic-ai-assistant-prompts-default/_search
{ "size": 1000 }
GET /_ingest/pipeline/.kibana-elastic-ai-assistant-ingest-pipeline-knowledge-base



// Delete them all!
// Data Streams
DELETE /_data_stream/.kibana-elastic-ai-assistant-anonymization-fields-default
DELETE /_data_stream/.kibana-elastic-ai-assistant-conversations-default
DELETE /_data_stream/.kibana-elastic-ai-assistant-knowledge-base-default
DELETE /_data_stream/.kibana-elastic-ai-assistant-prompts-default
// Index Templates
DELETE /_index_template/.kibana-elastic-ai-assistant-index-template-anonymization-fields
DELETE /_index_template/.kibana-elastic-ai-assistant-index-template-conversations
DELETE /_index_template/.kibana-elastic-ai-assistant-index-template-knowledge-base
DELETE /_index_template/.kibana-elastic-ai-assistant-index-template-prompts
// Pipelines
DELETE /_ingest/pipeline/.kibana-elastic-ai-assistant-ingest-pipeline-knowledge-base

New Features:

Plumbed through new assistantKnowledgeBaseByDefault experimental feature flag and exposed through assistantCapabilities API
Cleaned up assistantFeatures made available in AssistantProvider and tests (no need to individually plumb new features)
Introduced new AIAssistantDataClient for creating Data Streams, Component Templates, and Ingest Pipeline
Use AIAssistantDataClient to automatically install ELSER via installElasticModel() API, then deploy via TrainedModelsAPI
Plumb through addKnowledgeBaseDocuments() for creating KB entries from LangChain Documents within AIAssistantDataClient
Update ElasticsearchStore to take a kbDataClient for use in document adding/retrieval

Changes not behind FF:

Updated getELSER() helper function to be called by internalUser to prevent ml privilege requirements to setup Knowledge Base once a privileged user has enabled ELSER
Updated get/post/delete knowledge base routes to create esStore as internalUser to enable all assistant users the ability to enable the KB once ELSER has been installed (currently they 503 if currentUser doesn't have read_ingest and manage/read for the .kibana-elastic-ai-assistant-kb index)
Updated get/post/delete knowledge base routes with assistant API access controls: tags: ['access:elasticAssistant'],
Relaxed id validation from UUID to UUID or NonEmptyString to support ES generated id's, see: ddf93a8

Checklist

Delete any items that are not applicable to this PR.

Any text added follows EUI's writing guidelines, uses sentence case text and includes i18n support
Documentation was added for features that require explanation or tutorials
- TBD once FF is removed. Tracked in [Request] 8.16 Security Assistant Custom Knowledge Bases security-docs#5337.
Unit or functional tests were updated or added to match the most common scenarios
- Existing tests updated with new functionality, new tests to be fleshed out as feature stabilizes before removing FF

…fault

…ipeline and esql docs via KB data client

…fault

…fix'

…fault

…tore

…fault

…atus checks when FF is enabled

…fault

… generated id's

…fault

spong · 2024-05-15T17:08:26Z

x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts

+          // Do not pre-gen _id for bulk create operations to avoid `version_conflict_engine_exception`
+          // TODO: Breaks Attack Discovery anonymization-field creation :(
+          { create: { _index: this.options.index } },


Letting ES generate ID's ended up breaking Attack Discovery as it validated anonymizationFields.id to be UUID.

After speaking with @YulNaumenko, we decided to relax this constraint, so it and all other ID's that ES may have a chance at generating have been updated to be UUID or NonEmptyString in this commit: ddf93a8

For reference, as mentioned in the source comment above, we removed the explicit id generation _id: uuid.v4() as this was causing version_conflict_engine_exception errors during the bulk document load of KB entries.

Incidentally, all the documents were actually loaded without issue, however we couldn't determine why the errors were thrown (which ultimately mess with application retry/error logic), so opted for having ES generate the id's. I have reached out to the ES folks (internal slack) to see if perhaps this is related to the ELSER ingest_pipeline as none of the other DataClient assets like anonymization-fields have this issue when bulk creating.

…fault

andrew-goldstein · 2024-05-17T13:52:00Z

...tic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts

+  return {
+    '@timestamp': createdAt,
+    created_at: createdAt,
+    created_by: user.profile_uid ?? 'unknown',


Is unknown ever a valid profile_uid?

It is not, just using it as a sentinel value at the moment as the authenticatedUser can be null on the kbDataClient (if for some reason it was not instantiated from a request, or an unauthenticated request).

@YulNaumenko has started putting auth checks on routes (like here in post_actions_execute), so I think either that or individual checks in the kbDataClient methods that require an auth'd user as you mentioned below would be good here as we discussed. Or perhaps we can make authenticatedUser mandatory on the client and just return a null client (which it can be now) if the authc.getCurrentUser() call returns null.

If the user is mandatory, how we will make "shared"/public knowledge?

Yeah, good point. We could use a skeleton 'system authenticatedUser', but I'm thinking it might be better to leave optional and do null checks as needed. I'll see what this looks like when I get closer to fleshing out the RBAC portions of the API. Thanks Yuliia!

...lastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts

x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts

andrew-goldstein · 2024-05-17T16:42:17Z

...lastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts

+      seq_no_primary_term: true,
+    });
+    const conversation = transformESSearchToKnowledgeBaseEntry(response);
+    return conversation[0] ?? null;


Is this an exceptional case, and if so, should it throw instead of returning null?

I was following the same pattern as @YulNaumenko put in place with the conversation data client here:

kibana/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.ts

Line 77 in 1c26791

return conversation[0] ?? null;

Let's discuss further with @YulNaumenko, as when you and I discussed offline, there are a few other failure points in this get action to consider, including any potential latency in writing/reading on serverless as well.

Looks like I need to update that variable name though, so I'll go ahead and do that 😅

andrew-goldstein · 2024-05-17T16:44:36Z

...lastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts

+                  {
+                    match: user.profile_uid
+                      ? { 'users.id': user.profile_uid }
+                      : { 'users.name': user.username },


Is user.username a valid fallback when user.profile_uid does not exist?

That's a good question... Again, I was following the existing pattern put in place with conversations here:

kibana/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.ts

Lines 38 to 40 in 1c26791

match: user.profile_uid

? { 'users.id': user.profile_uid }

: { 'users.name': user.username },

I'm not sure here, so let's discuss with @YulNaumenko as there may be some cloud/ESS/Serverless things to consider around identity/users (as this seems to indicate profile_uid is not present everywhere, by maybe username is).

…fault

andrew-goldstein · 2024-05-17T19:50:28Z

x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.ts

@@ -21,7 +21,7 @@ export interface AIAssistantDataClientParams {
  kibanaVersion: string;
  spaceId: string;
  logger: Logger;
-  indexPatternsResorceName: string;
+  indexPatternsResourceName: string;


nice find! 🙏

x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts

x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts

andrew-goldstein · 2024-05-17T20:35:04Z

...ck/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts

+      return response.map((doc) => doc.id);
+    } catch (e) {
+      this.logger.error(`Error loading data into KB\n ${e}`);
+      return [];


Consider letting this error propagate, because if it doesn't, the caller must compare the length of the input documents with the length of the returned response here, in order to detect errors.

Going to leave this one as-is for now since it matches the current behavior of the addDocuments() function. But will revisit this as I work through the tool creation next week and (possibly) move document creation out of esStore.

x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts

andrew-goldstein

Thanks @spong for carefully slotting in this improvement to the KB setup UX while ensuring the current path continues to work 🙏
✅ Desk tested locally
LGTM 🚀

…fault

kibana-ci · 2024-05-20T17:27:39Z

💚 Build Succeeded

Buildkite Build
Commit: c91ea58

Metrics [docs]

Module Count

Fewer modules leads to a faster build time

id	before	after	diff
`securitySolution`	5492	5497	+5

Public APIs missing comments

Total count of every public API that lacks a comment. Target amount is 0. Run node scripts/build_api_docs --plugin [yourplugin] --stats comments for more detailed information.

id	before	after	diff
`@kbn/elastic-assistant`	147	138	-9
`@kbn/elastic-assistant-common`	217	275	+58
total			+49

Async chunks

Total size of all lazy-loaded chunks that will be downloaded as the user navigates the app

id	before	after	diff
`securitySolution`	15.1MB	15.1MB	+2.4KB

Page load bundle

Size of the bundles that are downloaded on every page load. Target size is below 100kb

id	before	after	diff
`securitySolution`	83.5KB	83.6KB	+35.0B

Unknown metric groups

API count

id	before	after	diff
`@kbn/elastic-assistant`	174	165	-9
`@kbn/elastic-assistant-common`	232	292	+60
total			+51

ESLint disabled line counts

id	before	after	diff
`elasticAssistant`	38	39	+1

Total ESLint disabled count

id	before	after	diff
`elasticAssistant`	38	39	+1

History

💔 Build #211013 failed 297b9bf
💚 Build #210882 succeeded 239f834
💚 Build #210848 succeeded 17f633e
💚 Build #210261 succeeded bef08b7
💔 Build #210137 failed 146aafb

To update your PR or re-run it, just comment with:
@elasticmachine merge upstream

cc @spong

YulNaumenko

LGTM! Exited about the next steps of making KB so easy to use!

spong added 6 commits May 1, 2024 10:14

Adds feature flag and cleaned up 'feature capabilities' touch points

2e8b823

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

97e2c35

…fault

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

cbcc5de

…fault

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

86e5ccc

…fault

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

7d0acb0

…fault

Use _inference API to automatically setup ELSER, and install ingest p…

4e4e1e7

…ipeline and esql docs via KB data client

spong added release_note:skip Skip the PR/issue when compiling release notes Team: SecuritySolution Security Solutions Team working on SIEM, Endpoint, Timeline, Resolver, etc. Feature:Security Assistant Security Assistant Team:Security Generative AI Security Generative AI v8.15.0 labels May 6, 2024

spong self-assigned this May 6, 2024

spong and others added 16 commits May 7, 2024 14:53

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

8fd6b2b

…fault

Fix settings tests

316e8c9

Switch to using installElasticModel()

71f4dd0

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

8843e87

…fault

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

d2c27ac

…fault

Resouce check note update

49a9c5a

[CI] Auto-commit changed files from 'node scripts/lint_ts_projects --…

512806a

…fix'

Creates 'kb entries' OAS and standardize routes and constants

dbc96e4

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

26006ff

…fault

Plumb through kbDataClient create and add kbDataClient support to EsS…

acdc0e7

…tore

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

cd91732

…fault

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

0eac540

…fault

Adds ingest pipeline to index template and fixes initial setup and st…

49986ea

…atus checks when FF is enabled

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

7354ffc

…fault

Fix kb document loading and auto elser setup

5b3d172

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

60ed588

…fault

spong marked this pull request as ready for review May 14, 2024 19:43

spong requested a review from a team as a code owner May 14, 2024 19:43

spong added 3 commits May 15, 2024 10:49

Fix isSetupInProgress update not propagating

252b8e8

Relax id validation from UUID to UUID or NonEmptyString to support ES…

ddf93a8

… generated id's

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

bef08b7

…fault

spong commented May 15, 2024

View reviewed changes

Merge branch 'main' of github.com:elastic/kibana into knowledge-by-de…

0038901

…fault