[Security solution] Bedrock streaming and token tracking

[stephmilovic]

Summary

Part 2 of 3 of Streaming in Security AI Assistant (Part 1)

1. Implement streaming for the Bedrock connector
2. Add invokeStream subaction to OpenAI and Bedrock. Handles standardizing our assistant stream request to OpenAI and Bedrock formats, and parses the response from the OpenAI and Bedrock formats to a simple string for the consumer to stream
3. Adds token tracking for Bedrock, streaming (invokeStream) and non streaming (run, test), in action_executor.ts
   • Note: the invokeAI method of Bedrock does not have token tracking. I plan remove that method in part 3 of 3 of the streaming work. At that time I will add a link to the token tracking dashboard form the Bedrock connector edit, like we have for OpenAI

package.json changes

I could not find documentation on how to handle the Bedrock codec. I did find awsdocs/aws-doc-sdk-examples#5541 that we can track for when documentation should be added.

I followed the bedrock extension in the langchainjs repo here: https://github.com/langchain-ai/langchainjs/pull/2367/files#diff-ead6be41360b3bf1c4559ea4d0069a015b69a57d5f8c51d0cce73253fde30aa2R246-R268. This is how I found that we'd need these packages: @smithy/eventstream-codec, @smithy/util-utf

To test:

1. Create Bedrock connector with Bedrock credentials and assign to your security AI assistant conversation
2. Ensure langchain is not enabled. Knowledge base should be off:

3. Send message to the connector. Test that:
   • response streams
   • "Stop generating" button works
   • "Regenerate" button works
4. Save this file as export.ndjson and upload to your saved objects
5. Click into the saved search we just uploaded, "Bedrock token tracking"
6. Ensure you see events with tokens for bedrock
7. Back to the Assistant, create a bedrock connector with bad credentials and ensure the Assistant returns an error when you attempt to use this connector

A note on api_alerting_integration tests

I spent 2 full work days trying to get these working. I’m trying to get the stream response from supertest. I can see the stream running on the backend, but when it comes back to supertest it only returns the readableState, not the stream itself. If there is anyone who can pair, I’d appreciate it. Otherwise, the code is covered in jest by x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.test.ts

[stephmilovic]

shareReplay causes the observable to be hot, which means unsubscribe does not actually stop the stream

[stephmilovic]

I could not find documentation on how to handle their codec. I did find awsdocs/aws-doc-sdk-examples#5541 for when documentation should be added.

I followed the bedrock extension in the langchainjs repo here: https://github.com/langchain-ai/langchainjs/pull/2367/files#diff-ead6be41360b3bf1c4559ea4d0069a015b69a57d5f8c51d0cce73253fde30aa2R246-R268

[elasticmachine]

Pinging @elastic/security-threat-hunting (Team:Threat Hunting)

[elasticmachine]

Pinging @elastic/security-solution (Team: SecuritySolution)

[semd]

Awesome! LGTM. Thanks Steph

[pmuellr]

code LGTM, left a comment about missing (trivial) jest test case

Regarding the functional test, can you use something other than supertest? I saw a thread somewhere of you asking about this, and I think someone else mentioned using node-fetch or such. If that can be made to work, no problem! If not, I guess there's something in our test framework that forces the use of supertest?

Worse case, create some kind of manual test in stack_connectors like this one we have in the actions plugin - so we have something to test with ... https://github.com/elastic/kibana/blob/main/x-pack/plugins/actions/server/manual_tests/forward_proxy.js

[stephmilovic]

My test wasn't working because I was trying to call execute route directly, instead of my special route that handles the stream.. duh! Thanks @dgieselaar for helping me figure this out!

[pmuellr]

LGTM

[kibana-ci]

💚 Build Succeeded

• Buildkite Build
• Commit: 25ffb7e

Metrics [docs]

Async chunks

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

id	before	after	diff
securitySolution	12.8MB	12.8MB	+754.0B

Page load bundle

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

id	before	after	diff
stackConnectors	39.3KB	39.4KB	+62.0B

History

• 💔 Build #176525 failed 5b7cf70
• 💛 Build #176331 was flaky 33d1f92
• 💔 Build #176309 failed ce6668f
• 💔 Build #176277 failed 902e7ba
• 💔 Build #176231 failed ae4e85d

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