[Security solution] Bedrock streaming and token tracking #170815

stephmilovic · 2023-11-07T23:01:03Z

Summary

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

Implement streaming for the Bedrock connector
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
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:

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

Send message to the connector. Test that:
- response streams
- "Stop generating" button works
- "Regenerate" button works
Save this file as export.ndjson and upload to your saved objects
Click into the saved search we just uploaded, "Bedrock token tracking"
Ensure you see events with tokens for bedrock
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 · 2023-11-09T20:17:28Z

x-pack/plugins/security_solution/public/assistant/get_comments/stream/stream_observable.ts

-    // make sure the request is only triggered once,
-    // even with multiple subscribers
-    shareReplay(1),


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

stephmilovic · 2023-11-09T20:20:22Z

x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts

+  new Transform({
+    transform(chunk, encoding, callback) {
+      const encoder = new TextEncoder();
+      const decoder = new EventStreamCodec(toUtf8, fromUtf8);


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 · 2023-11-14T17:30:27Z

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

elasticmachine · 2023-11-14T17:30:28Z

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

x-pack/plugins/actions/server/lib/get_token_count_from_bedrock_invoke.ts

x-pack/plugins/actions/server/lib/get_token_count_from_invoke_stream.ts

x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts

…esting

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

x-pack/plugins/actions/server/lib/gen_ai_token_tracking.ts

…esting

stephmilovic · 2023-11-15T18:48:02Z

...alerting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts

+              const passThrough = new PassThrough();
+
+              supertest
+                .post(`/internal/elastic_assistant/actions/connector/${bedrockActionId}/_execute`)


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

x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts

x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts

kibana-ci · 2023-11-16T21:26:41Z

💚 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

stephmilovic added 2 commits November 7, 2023 15:21

bedrock magic

a1929d0

cleanup

f49cf23

stephmilovic added release_note:enhancement WIP Work in progress Team:Threat Hunting Security Solution Threat Hunting Team Team: SecuritySolution Security Solutions Team working on SIEM, Endpoint, Timeline, Resolver, etc. Team:Threat Hunting:Explore skip-ci v8.12.0 labels Nov 7, 2023

stephmilovic added 6 commits November 8, 2023 10:23

cleanup and openai wip

8738b15

better

427f6d2

wip

e32ffdb

token tracking for bedrock

237f6fc

cleanup

4604df4

rm

9e69031

stephmilovic commented Nov 9, 2023

View reviewed changes

stephmilovic added 5 commits November 9, 2023 13:25

cleanup

1a396f5

fix api tests

8696210

token tests

2eabb1b

security solution tests

b1154db

stack connector tests

a375387

stephmilovic removed WIP Work in progress skip-ci labels Nov 14, 2023

stephmilovic marked this pull request as ready for review November 14, 2023 17:30

stephmilovic requested review from a team as code owners November 14, 2023 17:30

semd reviewed Nov 15, 2023

View reviewed changes

x-pack/plugins/actions/server/lib/get_token_count_from_bedrock_invoke.ts Show resolved Hide resolved

semd reviewed Nov 15, 2023

View reviewed changes

x-pack/plugins/actions/server/lib/get_token_count_from_invoke_stream.ts Outdated Show resolved Hide resolved

semd reviewed Nov 15, 2023

View reviewed changes

x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts Outdated Show resolved Hide resolved

stephmilovic added 5 commits November 15, 2023 07:28

Merge branch 'bedrock_streaming' into bedrock_streaming_integration_t…

9ab90b4

…esting

Sergi PR changes

b291fe3

Sergi was right

c8957c5

one more!

ae4e85d

Merge branch 'bedrock_streaming' into bedrock_streaming_integration_t…

40bd1c9

…esting

semd approved these changes Nov 15, 2023

View reviewed changes

pmuellr requested changes Nov 15, 2023

View reviewed changes

x-pack/plugins/actions/server/lib/gen_ai_token_tracking.ts Show resolved Hide resolved

stephmilovic added 6 commits November 15, 2023 10:01

fix whoops

3f70494

add tests for shouldTrackGenAiToken

902e7ba

Merge branch 'bedrock_streaming' into bedrock_streaming_integration_t…

10c188b

…esting

commit

325b0ff

done?

b5e8e83

really

ce6668f

stephmilovic commented Nov 15, 2023

View reviewed changes

better error handling

33d1f92

stephmilovic requested a review from pmuellr November 15, 2023 20:30

pmuellr approved these changes Nov 15, 2023

View reviewed changes

tests for line buffer

5b7cf70

stephmilovic commented Nov 16, 2023

View reviewed changes

x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts Outdated Show resolved Hide resolved

stephmilovic commented Nov 16, 2023

View reviewed changes

x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts Outdated Show resolved Hide resolved

stephmilovic commented Nov 16, 2023

View reviewed changes

x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts Outdated Show resolved Hide resolved

stephmilovic added 2 commits November 16, 2023 12:52

fixed

cd6cb1e

Merge branch 'main' into bedrock_streaming

25ffb7e

stephmilovic merged commit d201610 into elastic:main Nov 16, 2023
27 checks passed

kibanamachine added the backport:skip This commit does not require backporting label Nov 16, 2023

jamesspi added this to the 8.12 milestone Nov 16, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Security solution] Bedrock streaming and token tracking #170815

[Security solution] Bedrock streaming and token tracking #170815

stephmilovic commented Nov 7, 2023 •

edited

Loading

stephmilovic Nov 9, 2023

stephmilovic Nov 9, 2023

elasticmachine commented Nov 14, 2023

elasticmachine commented Nov 14, 2023

semd left a comment

pmuellr left a comment

stephmilovic Nov 15, 2023

pmuellr left a comment

kibana-ci commented Nov 16, 2023

[Security solution] Bedrock streaming and token tracking #170815

[Security solution] Bedrock streaming and token tracking #170815

Conversation

stephmilovic commented Nov 7, 2023 • edited Loading

Summary

package.json changes

To test:

A note on api_alerting_integration tests

stephmilovic Nov 9, 2023

Choose a reason for hiding this comment

stephmilovic Nov 9, 2023

Choose a reason for hiding this comment

elasticmachine commented Nov 14, 2023

elasticmachine commented Nov 14, 2023

semd left a comment

Choose a reason for hiding this comment

pmuellr left a comment

Choose a reason for hiding this comment

stephmilovic Nov 15, 2023

Choose a reason for hiding this comment

pmuellr left a comment

Choose a reason for hiding this comment

kibana-ci commented Nov 16, 2023

💚 Build Succeeded

Metrics [docs]

Async chunks

Page load bundle

History

stephmilovic commented Nov 7, 2023 •

edited

Loading

A note on `api_alerting_integration` tests