Add bedrock feedback into preview

newrelic/api/ml_model.py

@@ -40,12 +40,15 @@ def wrap_mlmodel(model, name=None, version=None, feature_names=None, label_names
 
 def get_llm_message_ids(response_id=None):
     transaction = current_transaction()
-    if response_id and transaction:
+    if transaction:
         nr_message_ids = getattr(transaction, "_nr_message_ids", {})
-        message_id_info = nr_message_ids.pop(response_id, ())
+        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:
-            warnings.warn("No message ids found for %s" % response_id)
+            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

newrelic/hooks/external_botocore.py

@@ -97,6 +97,7 @@ def create_chat_completion_message_event(
     if not transaction:
         return
 
+    message_ids = []
     for index, message in enumerate(input_message_list):
         if response_id:
             id_ = "%s-%d" % (response_id, index)  # Response ID was set, append message index to it.
@@ -128,6 +129,7 @@ def create_chat_completion_message_event(
             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_,
@@ -147,6 +149,7 @@ def create_chat_completion_message_event(
             "is_response": True,
         }
         transaction.record_custom_event("LlmChatCompletionMessage", chat_completion_message_dict)
+    return (conversation_id, request_id, message_ids)
 
 
 def extract_bedrock_titan_text_model(request_body, response_body=None):
@@ -577,7 +580,7 @@ def handle_chat_completion_event(
 
     transaction.record_custom_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,
         input_message_list=input_message_list,
@@ -591,6 +594,10 @@ def handle_chat_completion_event(
         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 = {
     ("sns", "publish"): message_trace("SNS", "Produce", "Topic", extract(("TopicArn", "TargetArn"), "PhoneNumber")),

tests/external_botocore/_test_bedrock_chat_completion.py

@@ -1,3 +1,17 @@
+# 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}',
@@ -6,6 +20,79 @@
     "meta.llama2-13b-chat-v1": '{"prompt": "%s", "temperature": %f, "max_gen_len": %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",
+            },
+        ]
+    },
+    "meta.llama2-13b-chat-v1": {
+        "bedrock_key": [
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "9a64cdb0-3e82-41c7-873a-c12a77e0143a",
+                "message_id": "9a64cdb0-3e82-41c7-873a-c12a77e0143a-0",
+            },
+            {
+                "conversation_id": "my-awesome-id",
+                "request_id": "9a64cdb0-3e82-41c7-873a-c12a77e0143a",
+                "message_id": "9a64cdb0-3e82-41c7-873a-c12a77e0143a-1",
+            },
+        ]
+    },
+}
+
 chat_completion_expected_events = {
     "amazon.titan-text-express-v1": [
         (

tests/external_botocore/_test_bedrock_embeddings.py

@@ -1,3 +1,17 @@
+# 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.
+
 embedding_payload_templates = {
     "amazon.titan-embed-text-v1": '{ "inputText": "%s" }',
     "amazon.titan-embed-g1-text-02": '{ "inputText": "%s" }',
@@ -68,7 +82,7 @@
                 "request_id": "",
                 "vendor": "bedrock",
                 "ingest_source": "Python",
-                "error": True
+                "error": True,
             },
         ),
     ],
@@ -89,7 +103,7 @@
                 "request_id": "",
                 "vendor": "bedrock",
                 "ingest_source": "Python",
-                "error": True
+                "error": True,
             },
         ),
     ],

tests/external_botocore/test_bedrock_embeddings.py

@@ -19,8 +19,8 @@
 import pytest
 from _test_bedrock_embeddings import (
     embedding_expected_client_errors,
-    embedding_expected_events,
     embedding_expected_error_events,
+    embedding_expected_events,
     embedding_payload_templates,
 )
 from conftest import BOTOCORE_VERSION

tox.ini

@@ -67,19 +67,19 @@ envlist =
     python-mlmodel_sklearn-{py37}-scikitlearn0101,
     python-component_djangorestframework-py27-djangorestframework0300,
     python-component_djangorestframework-{py37,py38,py39,py310,py311}-djangorestframeworklatest,
-    python-component_flask_rest-{py37,py38,py39,pypy38}-flaskrestxlatest,
+    python-component_flask_rest-py37-flaskrestx110,
+    python-component_flask_rest-{py38,py39,py310,py311,pypy38}-flaskrestxlatest,
     python-component_flask_rest-{py27,pypy27}-flaskrestx051,
     python-component_graphqlserver-{py37,py38,py39,py310,py311},
-    python-component_tastypie-{py27,pypy27}-tastypie0143,
-    python-component_tastypie-{py37,py38,py39,pypy38}-tastypie{0143,latest},
+    python-component_tastypie-{py37,py38,py39,pypy38}-tastypielatest,
     python-coroutines_asyncio-{py37,py38,py39,py310,py311,pypy38},
     python-cross_agent-{py27,py37,py38,py39,py310,py311}-{with,without}_extensions,
     python-cross_agent-pypy27-without_extensions,
     postgres-datastore_asyncpg-{py37,py38,py39,py310,py311},
     memcached-datastore_bmemcached-{pypy27,py27,py37,py38,py39,py310,py311}-memcached030,
     elasticsearchserver07-datastore_elasticsearch-{py27,py37,py38,py39,py310,py311,pypy27,pypy38}-elasticsearch07,
     elasticsearchserver08-datastore_elasticsearch-{py37,py38,py39,py310,py311,pypy38}-elasticsearch08,
-    memcached-datastore_memcache-{py27,py37,py38,py39,py310,py311,pypy27,pypy38}-memcached01,
+    memcached-datastore_memcache-{py37,py38,py39,py310,py311,pypy38}-memcached01,
     mysql-datastore_mysql-mysql080023-py27,
     mysql-datastore_mysql-mysqllatest-{py37,py38,py39,py310,py311},
     firestore-datastore_firestore-{py37,py38,py39,py310,py311},
@@ -177,7 +177,8 @@ deps =
     adapter_gunicorn-aiohttp3: aiohttp<4.0
     adapter_gunicorn-gunicorn19: gunicorn<20
     adapter_gunicorn-gunicornlatest: gunicorn
-    adapter_hypercorn-hypercornlatest: hypercorn
+    ; Temporarily pinned.  Needs to be addressed
+    adapter_hypercorn-hypercornlatest: hypercorn<0.16
     adapter_hypercorn-hypercorn0013: hypercorn<0.14
     adapter_hypercorn-hypercorn0012: hypercorn<0.13
     adapter_hypercorn-hypercorn0011: hypercorn<0.12
@@ -204,23 +205,21 @@ deps =
     component_djangorestframework-djangorestframework0300: djangorestframework<3.1
     component_djangorestframework-djangorestframeworklatest: Django
     component_djangorestframework-djangorestframeworklatest: djangorestframework
-    component_flask_rest: flask
     component_flask_rest: flask-restful
     component_flask_rest: jinja2
     component_flask_rest: itsdangerous
+    component_flask_rest-flaskrestxlatest: flask
     component_flask_rest-flaskrestxlatest: flask-restx
-    ; Pin Flask version until flask-restx is updated to support v3
-    component_flask_rest-flaskrestxlatest: flask<3.0
+    ; flask-restx only supports Flask v3 after flask-restx v1.3.0
+    component_flask_rest-flaskrestx110: Flask<3.0
+    component_flask_rest-flaskrestx110: flask-restx<1.2
+    component_flask_rest-flaskrestx051: Flask<3.0
     component_flask_rest-flaskrestx051: flask-restx<1.0
     component_graphqlserver: graphql-server[sanic,flask]==3.0.0b5
     component_graphqlserver: sanic>20
     component_graphqlserver: Flask
     component_graphqlserver: markupsafe<2.1
     component_graphqlserver: jinja2<3.1
-    component_tastypie-tastypie0143: django-tastypie<0.14.4
-    component_tastypie-{py27,pypy27}-tastypie0143: django<1.12
-    component_tastypie-{py37,py38,py39,py310,py311,pypy38}-tastypie0143: django<3.0.1
-    component_tastypie-{py37,py38,py39,py310,py311,pypy38}-tastypie0143: asgiref<3.7.1  # asgiref==3.7.1 only suppport Python 3.10+
     component_tastypie-tastypielatest: django-tastypie
     component_tastypie-tastypielatest: django<4.1
     component_tastypie-tastypielatest: asgiref<3.7.1  # asgiref==3.7.1 only suppport Python 3.10+