Fix aiobotocore test.

Author: umaannamalai

tests/external_aiobotocore/test_bedrock_chat_completion_converse.py

@@ -23,6 +23,7 @@
 )
 from testing_support.fixtures import override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes
 from testing_support.ml_testing_utils import (
+    add_token_count_streaming_events,
     add_token_counts_to_chat_events,
     disabled_ai_monitoring_record_content_settings,
     disabled_ai_monitoring_settings,
@@ -146,8 +147,14 @@ def _test():
 
 @reset_core_stats_engine()
 @override_llm_token_callback_settings(llm_token_count_callback)
-def test_bedrock_chat_completion_with_token_count(set_trace_info, exercise_model, expected_metric, expected_events):
-    @validate_custom_events(add_token_counts_to_chat_events(chat_completion_expected_events))
+def test_bedrock_chat_completion_with_token_count(
+    set_trace_info, exercise_model, expected_metric, expected_events, response_streaming
+):
+    expected_events = add_token_counts_to_chat_events(expected_events)
+    if response_streaming:
+        expected_events = add_token_count_streaming_events(expected_events)
+
+    @validate_custom_events(expected_events)
     # One summary event, one user message, and one response message from the assistant
     @validate_custom_event_count(count=4)
     @validate_transaction_metrics(