Add response token count logic to Bedrock instrumentation. (#1504)

* Add bedrock token counting.

* [MegaLinter] Apply linters fixes

* Add bedrock token counting.

* Add safeguards when grabbing token counts.

* Remove extra None defaults.

* Cleanup default None checks.

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: umaannamalai

newrelic/hooks/external_botocore.py

@@ -23,7 +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_to_events,
+    add_token_counts_to_chat_events,
     disabled_ai_monitoring_record_content_settings,
     disabled_ai_monitoring_settings,
     events_sans_content,
@@ -147,7 +147,7 @@ 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_count_to_events(expected_events))
+    @validate_custom_events(add_token_counts_to_chat_events(chat_completion_expected_events))
     # One summary event, one user message, and one response message from the assistant
     @validate_custom_event_count(count=4)
     @validate_transaction_metrics(
@@ -278,49 +278,6 @@ def _test():
     _test()
 
 
-@reset_core_stats_engine()
-@override_llm_token_callback_settings(llm_token_count_callback)
-def test_bedrock_chat_completion_error_incorrect_access_key_with_token_count(
-    exercise_converse_incorrect_access_key, set_trace_info, expected_metric
-):
-    """
-    A request is made to the server with invalid credentials. botocore will reach out to the server and receive an
-    UnrecognizedClientException as a response. Information from the request will be parsed and reported in customer
-    events. The error response can also be parsed, and will be included as attributes on the recorded exception.
-    """
-
-    @validate_custom_events(add_token_count_to_events(chat_completion_invalid_access_key_error_events))
-    @validate_error_trace_attributes(
-        _client_error_name,
-        exact_attrs={
-            "agent": {},
-            "intrinsic": {},
-            "user": {
-                "http.statusCode": 403,
-                "error.message": "The security token included in the request is invalid.",
-                "error.code": "UnrecognizedClientException",
-            },
-        },
-    )
-    @validate_transaction_metrics(
-        name="test_bedrock_chat_completion_incorrect_access_key_with_token_count",
-        scoped_metrics=[expected_metric],
-        rollup_metrics=[expected_metric],
-        custom_metrics=[(f"Supportability/Python/ML/Bedrock/{BOTOCORE_VERSION}", 1)],
-        background_task=True,
-    )
-    @background_task(name="test_bedrock_chat_completion_incorrect_access_key_with_token_count")
-    def _test():
-        set_trace_info()
-        add_custom_attribute("llm.conversation_id", "my-awesome-id")
-        add_custom_attribute("llm.foo", "bar")
-        add_custom_attribute("non_llm_attr", "python-agent")
-
-        exercise_converse_incorrect_access_key()
-
-    _test()
-
-
 @pytest.fixture
 def exercise_converse_invalid_model(loop, bedrock_converse_server, response_streaming, monkeypatch):
     def _exercise_converse_invalid_model():

tests/external_aiobotocore/test_bedrock_chat_completion_converse.py

@@ -34,7 +34,8 @@
 )
 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_to_events,
+    add_token_count_streaming_events,
+    add_token_counts_to_chat_events,
     disabled_ai_monitoring_record_content_settings,
     disabled_ai_monitoring_settings,
     disabled_ai_monitoring_streaming_settings,
@@ -207,7 +208,7 @@ 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_events, expected_metrics):
-    @validate_custom_events(add_token_count_to_events(expected_events))
+    @validate_custom_events(add_token_counts_to_chat_events(add_token_count_streaming_events(expected_events)))
     # One summary event, one user message, and one response message from the assistant
     @validate_custom_event_count(count=3)
     @validate_transaction_metrics(
@@ -456,51 +457,6 @@ def _test():
     _test()
 
 
-@reset_core_stats_engine()
-@override_llm_token_callback_settings(llm_token_count_callback)
-def test_bedrock_chat_completion_error_incorrect_access_key_with_token(
-    monkeypatch,
-    bedrock_server,
-    exercise_model,
-    set_trace_info,
-    expected_invalid_access_key_error_events,
-    expected_metrics,
-):
-    @validate_custom_events(add_token_count_to_events(expected_invalid_access_key_error_events))
-    @validate_error_trace_attributes(
-        _client_error_name,
-        exact_attrs={
-            "agent": {},
-            "intrinsic": {},
-            "user": {
-                "http.statusCode": 403,
-                "error.message": "The security token included in the request is invalid.",
-                "error.code": "UnrecognizedClientException",
-            },
-        },
-    )
-    @validate_transaction_metrics(
-        name="test_bedrock_chat_completion",
-        scoped_metrics=expected_metrics,
-        rollup_metrics=expected_metrics,
-        custom_metrics=[(f"Supportability/Python/ML/Bedrock/{BOTOCORE_VERSION}", 1)],
-        background_task=True,
-    )
-    @background_task(name="test_bedrock_chat_completion")
-    def _test():
-        monkeypatch.setattr(bedrock_server._request_signer._credentials, "access_key", "INVALID-ACCESS-KEY")
-
-        with pytest.raises(_client_error):  # not sure where this exception actually comes from
-            set_trace_info()
-            add_custom_attribute("llm.conversation_id", "my-awesome-id")
-            add_custom_attribute("llm.foo", "bar")
-            add_custom_attribute("non_llm_attr", "python-agent")
-
-            exercise_model(prompt="Invalid Token", temperature=0.7, max_tokens=100)
-
-    _test()
-
-
 def invoke_model_malformed_request_body(loop, bedrock_server, response_streaming):
     async def _coro():
         with pytest.raises(_client_error):
@@ -799,58 +755,6 @@ async def _test():
     loop.run_until_complete(_test())
 
 
-@reset_core_stats_engine()
-@override_llm_token_callback_settings(llm_token_count_callback)
-@validate_custom_events(add_token_count_to_events(chat_completion_expected_streaming_error_events))
-@validate_custom_event_count(count=2)
-@validate_error_trace_attributes(
-    _event_stream_error_name,
-    exact_attrs={
-        "agent": {},
-        "intrinsic": {},
-        "user": {
-            "error.message": "Malformed input request, please reformat your input and try again.",
-            "error.code": "ValidationException",
-        },
-    },
-    forgone_params={"agent": (), "intrinsic": (), "user": ("http.statusCode")},
-)
-@validate_transaction_metrics(
-    name="test_bedrock_chat_completion",
-    scoped_metrics=[("Llm/completion/Bedrock/invoke_model_with_response_stream", 1)],
-    rollup_metrics=[("Llm/completion/Bedrock/invoke_model_with_response_stream", 1)],
-    custom_metrics=[(f"Supportability/Python/ML/Bedrock/{BOTOCORE_VERSION}", 1)],
-    background_task=True,
-)
-@background_task(name="test_bedrock_chat_completion")
-def test_bedrock_chat_completion_error_streaming_exception_with_token_count(loop, bedrock_server, set_trace_info):
-    """
-    Duplicate of test_bedrock_chat_completion_error_streaming_exception, but with token callback being set.
-
-    See the original test for a description of the error case.
-    """
-
-    async def _test():
-        with pytest.raises(_event_stream_error):
-            model = "amazon.titan-text-express-v1"
-            body = (chat_completion_payload_templates[model] % ("Streaming Exception", 0.7, 100)).encode("utf-8")
-
-            set_trace_info()
-            add_custom_attribute("llm.conversation_id", "my-awesome-id")
-            add_custom_attribute("llm.foo", "bar")
-            add_custom_attribute("non_llm_attr", "python-agent")
-
-            response = await bedrock_server.invoke_model_with_response_stream(
-                body=body, modelId=model, accept="application/json", contentType="application/json"
-            )
-
-            body = response.get("body")
-            async for resp in body:
-                assert resp
-
-    loop.run_until_complete(_test())
-
-
 def test_bedrock_chat_completion_functions_marked_as_wrapped_for_sdk_compatibility(bedrock_server):
     assert bedrock_server._nr_wrapped
 

tests/external_aiobotocore/test_bedrock_chat_completion_invoke_model.py

@@ -28,7 +28,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_to_events,
+    add_token_count_to_embedding_events,
     disabled_ai_monitoring_record_content_settings,
     disabled_ai_monitoring_settings,
     events_sans_content,
@@ -165,7 +165,7 @@ def _test():
 @reset_core_stats_engine()
 @override_llm_token_callback_settings(llm_token_count_callback)
 def test_bedrock_embedding_with_token_count(set_trace_info, exercise_model, expected_events):
-    @validate_custom_events(add_token_count_to_events(expected_events))
+    @validate_custom_events(add_token_count_to_embedding_events(expected_events))
     @validate_custom_event_count(count=1)
     @validate_transaction_metrics(
         name="test_bedrock_embedding",
@@ -290,45 +290,6 @@ def _test():
     _test()
 
 
-@reset_core_stats_engine()
-@override_llm_token_callback_settings(llm_token_count_callback)
-def test_bedrock_embedding_error_incorrect_access_key_with_token_count(
-    monkeypatch, bedrock_server, exercise_model, set_trace_info, expected_invalid_access_key_error_events
-):
-    @validate_custom_events(add_token_count_to_events(expected_invalid_access_key_error_events))
-    @validate_error_trace_attributes(
-        _client_error_name,
-        exact_attrs={
-            "agent": {},
-            "intrinsic": {},
-            "user": {
-                "http.statusCode": 403,
-                "error.message": "The security token included in the request is invalid.",
-                "error.code": "UnrecognizedClientException",
-            },
-        },
-    )
-    @validate_transaction_metrics(
-        name="test_bedrock_embedding",
-        scoped_metrics=[("Llm/embedding/Bedrock/invoke_model", 1)],
-        rollup_metrics=[("Llm/embedding/Bedrock/invoke_model", 1)],
-        background_task=True,
-    )
-    @background_task(name="test_bedrock_embedding")
-    def _test():
-        monkeypatch.setattr(bedrock_server._request_signer._credentials, "access_key", "INVALID-ACCESS-KEY")
-
-        with pytest.raises(_client_error):  # not sure where this exception actually comes from
-            set_trace_info()
-            add_custom_attribute("llm.conversation_id", "my-awesome-id")
-            add_custom_attribute("llm.foo", "bar")
-            add_custom_attribute("non_llm_attr", "python-agent")
-
-            exercise_model(prompt="Invalid Token")
-
-    _test()
-
-
 @reset_core_stats_engine()
 @validate_custom_events(embedding_expected_malformed_request_body_events)
 @validate_custom_event_count(count=1)

tests/external_aiobotocore/test_bedrock_embeddings.py

@@ -29,6 +29,9 @@
             "duration": None,  # Response time varies each test run
             "request.model": "anthropic.claude-3-sonnet-20240229-v1:0",
             "response.model": "anthropic.claude-3-sonnet-20240229-v1:0",
+            "response.usage.prompt_tokens": 26,
+            "response.usage.completion_tokens": 100,
+            "response.usage.total_tokens": 126,
             "request.temperature": 0.7,
             "request.max_tokens": 100,
             "response.choices.finish_reason": "max_tokens",
@@ -51,6 +54,7 @@
             "role": "system",
             "completion_id": None,
             "sequence": 0,
+            "token_count": 0,
             "response.model": "anthropic.claude-3-sonnet-20240229-v1:0",
             "vendor": "bedrock",
             "ingest_source": "Python",
@@ -70,6 +74,7 @@
             "role": "user",
             "completion_id": None,
             "sequence": 1,
+            "token_count": 0,
             "response.model": "anthropic.claude-3-sonnet-20240229-v1:0",
             "vendor": "bedrock",
             "ingest_source": "Python",
@@ -89,6 +94,7 @@
             "role": "assistant",
             "completion_id": None,
             "sequence": 2,
+            "token_count": 0,
             "response.model": "anthropic.claude-3-sonnet-20240229-v1:0",
             "vendor": "bedrock",
             "ingest_source": "Python",