Add bedrock feedback support

newrelic/agent.py

@@ -153,6 +153,7 @@ def __asgi_application(*args, **kwargs):
 from newrelic.api.message_transaction import (
     wrap_message_transaction as __wrap_message_transaction,
 )
+from newrelic.api.ml_model import get_llm_message_ids as __get_llm_message_ids
 from newrelic.api.ml_model import wrap_mlmodel as __wrap_mlmodel
 from newrelic.api.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper
 from newrelic.api.profile_trace import profile_trace as __profile_trace
@@ -340,3 +341,4 @@ def __asgi_application(*args, **kwargs):
 insert_html_snippet = __wrap_api_call(__insert_html_snippet, "insert_html_snippet")
 verify_body_exists = __wrap_api_call(__verify_body_exists, "verify_body_exists")
 wrap_mlmodel = __wrap_api_call(__wrap_mlmodel, "wrap_mlmodel")
+get_llm_message_ids = __wrap_api_call(__get_llm_message_ids, "get_llm_message_ids")

newrelic/api/ml_model.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import sys
+import warnings
 
+from newrelic.api.transaction import current_transaction
 from newrelic.common.object_names import callable_name
 from newrelic.hooks.mlmodel_sklearn import _nr_instrument_model
 
@@ -33,3 +35,24 @@ def wrap_mlmodel(model, name=None, version=None, feature_names=None, label_names
         model._nr_wrapped_label_names = label_names
     if metadata:
         model._nr_wrapped_metadata = metadata
+
+
+def get_llm_message_ids(response_id=None):
+    transaction = current_transaction()
+    if transaction:
+        nr_message_ids = getattr(transaction, "_nr_message_ids", {})
+        message_id_info = (
+            nr_message_ids.pop("bedrock_key", ()) if not response_id else nr_message_ids.pop(response_id, ())
+        )
+
+        if not message_id_info:
+            response_id_warning = "." if not response_id else " for %s." % response_id
+            warnings.warn("No message ids found%s" % response_id_warning)
+            return []
+
+        conversation_id, request_id, ids = message_id_info
+
+        return [{"conversation_id": conversation_id, "request_id": request_id, "message_id": _id} for _id in ids]
+
+    warnings.warn("No message ids found. get_llm_message_ids must be called within the scope of a transaction.")
+    return []

newrelic/hooks/external_botocore.py

@@ -91,11 +91,13 @@ def create_chat_completion_message_event(
     if not transaction:
         return
 
+    message_ids = []
     for index, message in enumerate(message_list):
         if response_id:
             id_ = "%s-%d" % (response_id, index)  # Response ID was set, append message index to it.
         else:
             id_ = str(uuid.uuid4())  # No response IDs, use random UUID
+        message_ids.append(id_)
 
         chat_completion_message_dict = {
             "id": id_,
@@ -115,6 +117,8 @@ def create_chat_completion_message_event(
         }
         transaction.record_ml_event("LlmChatCompletionMessage", chat_completion_message_dict)
 
+    return (conversation_id, request_id, message_ids)
+
 
 def extract_bedrock_titan_text_model(request_body, response_body=None):
     request_body = json.loads(request_body)
@@ -319,13 +323,20 @@ def wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
     response_headers = response["ResponseMetadata"]["HTTPHeaders"]
 
     if model.startswith("amazon.titan-embed"):  # Only available embedding models
-        handle_embedding_event(instance, transaction, extractor, model, response_body, response_headers, request_body, ft.duration)
+        handle_embedding_event(
+            instance, transaction, extractor, model, response_body, response_headers, request_body, ft.duration
+        )
     else:
-        handle_chat_completion_event(instance, transaction, extractor, model, response_body, response_headers, request_body, ft.duration)
+        handle_chat_completion_event(
+            instance, transaction, extractor, model, response_body, response_headers, request_body, ft.duration
+        )
 
     return response
 
-def handle_embedding_event(client, transaction, extractor, model, response_body, response_headers, request_body, duration):
+
+def handle_embedding_event(
+    client, transaction, extractor, model, response_body, response_headers, request_body, duration
+):
     embedding_id = str(uuid.uuid4())
     available_metadata = get_trace_linking_metadata()
     span_id = available_metadata.get("span.id", "")
@@ -336,25 +347,29 @@ def handle_embedding_event(client, transaction, extractor, model, response_body,
 
     _, embedding_dict = extractor(request_body, response_body)
 
-    embedding_dict.update({
-        "vendor": "bedrock",
-        "ingest_source": "Python",
-        "id": embedding_id,
-        "appName": settings.app_name,
-        "span_id": span_id,
-        "trace_id": trace_id,
-        "request_id": request_id,
-        "transaction_id": transaction._transaction_id,
-        "api_key_last_four_digits": client._request_signer._credentials.access_key[-4:],
-        "duration": duration,
-        "request.model": model,
-        "response.model": model,
-    })
+    embedding_dict.update(
+        {
+            "vendor": "bedrock",
+            "ingest_source": "Python",
+            "id": embedding_id,
+            "appName": settings.app_name,
+            "span_id": span_id,
+            "trace_id": trace_id,
+            "request_id": request_id,
+            "transaction_id": transaction._transaction_id,
+            "api_key_last_four_digits": client._request_signer._credentials.access_key[-4:],
+            "duration": duration,
+            "request.model": model,
+            "response.model": model,
+        }
+    )
 
     transaction.record_ml_event("LlmEmbedding", embedding_dict)
 
 
-def handle_chat_completion_event(client, transaction, extractor, model, response_body, response_headers, request_body, duration):
+def handle_chat_completion_event(
+    client, transaction, extractor, model, response_body, response_headers, request_body, duration
+):
     custom_attrs_dict = transaction._custom_params
     conversation_id = custom_attrs_dict.get("conversation_id", "")
 
@@ -388,7 +403,7 @@ def handle_chat_completion_event(client, transaction, extractor, model, response
 
     transaction.record_ml_event("LlmChatCompletionSummary", chat_completion_summary_dict)
 
-    create_chat_completion_message_event(
+    message_ids = create_chat_completion_message_event(
         transaction=transaction,
         app_name=settings.app_name,
         message_list=message_list,
@@ -400,6 +415,9 @@ def handle_chat_completion_event(client, transaction, extractor, model, response
         conversation_id=conversation_id,
         response_id=response_id,
     )
+    if not hasattr(transaction, "_nr_message_ids"):
+        transaction._nr_message_ids = {}
+    transaction._nr_message_ids["bedrock_key"] = message_ids
 
 
 CUSTOM_TRACE_POINTS = {

tests/external_botocore/_test_bedrock_chat_completion.py

@@ -1,10 +1,83 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 chat_completion_payload_templates = {
     "amazon.titan-text-express-v1": '{ "inputText": "%s", "textGenerationConfig": {"temperature": %f, "maxTokenCount": %d }}',
     "ai21.j2-mid-v1": '{"prompt": "%s", "temperature": %f, "maxTokens": %d}',
     "anthropic.claude-instant-v1": '{"prompt": "Human: %s Assistant:", "temperature": %f, "max_tokens_to_sample": %d}',
     "cohere.command-text-v14": '{"prompt": "%s", "temperature": %f, "max_tokens": %d}',
 }
 
+chat_completion_get_llm_message_ids = {
+    "amazon.titan-text-express-v1": {
+        "bedrock_key": [
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "03524118-8d77-430f-9e08-63b5c03a40cf",
+                "message_id": None,  # UUID that varies with each run
+            },
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "03524118-8d77-430f-9e08-63b5c03a40cf",
+                "message_id": None,  # UUID that varies with each run
+            },
+        ]
+    },
+    "ai21.j2-mid-v1": {
+        "bedrock_key": [
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "c863d9fc-888b-421c-a175-ac5256baec62",
+                "message_id": "1234-0",
+            },
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "c863d9fc-888b-421c-a175-ac5256baec62",
+                "message_id": "1234-1",
+            },
+        ]
+    },
+    "anthropic.claude-instant-v1": {
+        "bedrock_key": [
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "7b0b37c6-85fb-4664-8f5b-361ca7b1aa18",
+                "message_id": None,  # UUID that varies with each run
+            },
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "7b0b37c6-85fb-4664-8f5b-361ca7b1aa18",
+                "message_id": None,  # UUID that varies with each run
+            },
+        ]
+    },
+    "cohere.command-text-v14": {
+        "bedrock_key": [
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "e77422c8-fbbf-4e17-afeb-c758425c9f97",
+                "message_id": "e77422c8-fbbf-4e17-afeb-c758425c9f97-0",
+            },
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "e77422c8-fbbf-4e17-afeb-c758425c9f97",
+                "message_id": "e77422c8-fbbf-4e17-afeb-c758425c9f97-1",
+            },
+        ]
+    },
+}
+
 chat_completion_expected_events = {
     "amazon.titan-text-express-v1": [
         (

tests/external_botocore/test_bedrock_chat_completion.py

@@ -21,6 +21,7 @@
 import pytest
 from _test_bedrock_chat_completion import (
     chat_completion_expected_events,
+    chat_completion_get_llm_message_ids,
     chat_completion_payload_templates,
     chat_completion_expected_client_errors,
 )
@@ -37,7 +38,7 @@
 )
 
 from newrelic.api.background_task import background_task
-from newrelic.api.time_trace import current_trace
+from newrelic.api.ml_model import get_llm_message_ids
 from newrelic.api.transaction import add_custom_attribute, current_transaction
 
 from newrelic.common.object_names import callable_name
@@ -94,6 +95,11 @@ def expected_events_no_convo_id(model_id):
     return events
 
 
+@pytest.fixture(scope="module")
+def expected_ai_message_ids(model_id):
+    return chat_completion_get_llm_message_ids[model_id]
+
+
 @pytest.fixture(scope="module")
 def expected_client_error(model_id):
     return chat_completion_expected_client_errors[model_id]
@@ -170,6 +176,72 @@ def test_bedrock_chat_completion_disabled_settings(set_trace_info, exercise_mode
     exercise_model(prompt=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100)
 
 
+# Testing get_llm_message_ids:
+
+@reset_core_stats_engine()
+@background_task()
+def test_get_llm_message_ids_when_nr_message_ids_not_set():
+    message_ids = get_llm_message_ids("request-id-1")
+    assert message_ids == []
+
+
+@reset_core_stats_engine()
+def test_get_llm_message_ids_outside_transaction():
+    message_ids = get_llm_message_ids("request-id-1")
+    assert message_ids == []
+
+
+@reset_core_stats_engine()
+def test_get_llm_message_ids_bedrock_chat_completion_in_txn(
+    set_trace_info, exercise_model, expected_ai_message_ids
+):  # noqa: F811
+    @background_task()
+    def _test():
+        set_trace_info()
+        add_custom_attribute("conversation_id", "my-awesome-id")
+        exercise_model(prompt=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100)
+
+        expected_message_ids = [value for value in expected_ai_message_ids.values()][0]
+        message_ids = [m for m in get_llm_message_ids()]
+        for index, message_id_info in enumerate(message_ids):
+            expected_message_id_info = expected_message_ids[index]
+            assert message_id_info["conversation_id"] == expected_message_id_info["conversation_id"]
+            assert message_id_info["request_id"] == expected_message_id_info["request_id"]
+            if expected_message_id_info["message_id"]:
+                assert message_id_info["message_id"] == expected_message_id_info["message_id"]
+            else:
+                # We are checking for the presence of a message_id since in this case it is
+                # a UUID that changes with each run.
+                assert message_id_info["message_id"]
+
+        assert current_transaction()._nr_message_ids == {}
+
+    _test()
+
+
+@reset_core_stats_engine()
+def test_get_llm_message_ids_bedrock_chat_completion_no_convo_id(
+    set_trace_info, exercise_model, expected_ai_message_ids
+):  # noqa: F811
+    @background_task()
+    def _test():
+        set_trace_info()
+        exercise_model(prompt=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100)
+
+        expected_message_ids = [value for value in expected_ai_message_ids.values()][0]
+        message_ids = [m for m in get_llm_message_ids()]
+        for index, message_id_info in enumerate(message_ids):
+            expected_message_id_info = expected_message_ids[index]
+            assert message_id_info["request_id"] == expected_message_id_info["request_id"]
+            if expected_message_id_info["message_id"]:
+                assert message_id_info["message_id"] == expected_message_id_info["message_id"]
+            else:
+                # We are checking for the presence of a message_id since in this case it is
+                # a UUID that changes with each run.
+                assert message_id_info["message_id"]
+            assert message_id_info["conversation_id"] == ""
+
+        assert current_transaction()._nr_message_ids == {}
 
 _client_error = botocore.exceptions.ClientError
 _client_error_name = callable_name(_client_error)