Add openai 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_ai_message_ids as __get_ai_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_ai_message_ids = __wrap_api_call(__get_ai_message_ids, "get_ai_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,20 @@ 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_ai_message_ids(response_id=None):
+    transaction = current_transaction()
+    if response_id and transaction:
+        nr_message_ids = getattr(transaction, "_nr_message_ids", {})
+        message_id_info = nr_message_ids.pop(response_id, ())
+
+        if not message_id_info:
+            warnings.warn("No message ids found for %s" % response_id)
+            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_ai_message_ids must be called within the scope of a transaction.")
+    return []

newrelic/hooks/mlmodel_openai.py

@@ -216,7 +216,7 @@ def wrap_chat_completion_create(wrapped, instance, args, kwargs):
     if choices:
         message_list.extend([choices[0].message])
 
-    create_chat_completion_message_event(
+    message_ids = create_chat_completion_message_event(
         transaction,
         settings.app_name,
         message_list,
@@ -229,6 +229,11 @@ def wrap_chat_completion_create(wrapped, instance, args, kwargs):
         conversation_id,
     )
 
+    # Cache message ids on transaction for retrieval after open ai call completion.
+    if not hasattr(transaction, "_nr_message_ids"):
+        transaction._nr_message_ids = {}
+    transaction._nr_message_ids[response_id] = message_ids
+
     return response
 
 
@@ -260,9 +265,12 @@ def create_chat_completion_message_event(
     request_id,
     conversation_id,
 ):
+    message_ids = []
     for index, message in enumerate(message_list):
+        message_id = "%s-%s" % (response_id, index)
+        message_ids.append(message_id)
         chat_completion_message_dict = {
-            "id": "%s-%s" % (response_id, index),
+            "id": message_id,
             "appName": app_name,
             "conversation_id": conversation_id,
             "request_id": request_id,
@@ -277,6 +285,7 @@ def create_chat_completion_message_event(
             "vendor": "openAI",
         }
         transaction.record_ml_event("LlmChatCompletionMessage", chat_completion_message_dict)
+    return (conversation_id, request_id, message_ids)
 
 
 async def wrap_embedding_acreate(wrapped, instance, args, kwargs):
@@ -456,7 +465,7 @@ async def wrap_chat_completion_acreate(wrapped, instance, args, kwargs):
     if choices:
         message_list.extend([choices[0].message])
 
-    create_chat_completion_message_event(
+    message_ids = create_chat_completion_message_event(
         transaction,
         settings.app_name,
         message_list,
@@ -469,6 +478,11 @@ async def wrap_chat_completion_acreate(wrapped, instance, args, kwargs):
         conversation_id,
     )
 
+    # Cache message ids on transaction for retrieval after open ai call completion.
+    if not hasattr(transaction, "_nr_message_ids"):
+        transaction._nr_message_ids = {}
+    transaction._nr_message_ids[response_id] = message_ids
+
     return response
 
 

tests/mlmodel_openai/_mock_external_openai_server.py

@@ -129,6 +129,40 @@
             "usage": {"completion_tokens": 11, "prompt_tokens": 53, "total_tokens": 64},
         },
     ),
+    "You are a mathematician.": (
+        {
+            "Content-Type": "application/json",
+            "openai-model": "gpt-3.5-turbo-0613",
+            "openai-organization": "new-relic-nkmd8b",
+            "openai-processing-ms": "1469",
+            "openai-version": "2020-10-01",
+            "x-ratelimit-limit-requests": "200",
+            "x-ratelimit-limit-tokens": "40000",
+            "x-ratelimit-remaining-requests": "199",
+            "x-ratelimit-remaining-tokens": "39940",
+            "x-ratelimit-reset-requests": "7m12s",
+            "x-ratelimit-reset-tokens": "90ms",
+            "x-request-id": "49dbbffbd3c3f4612aa48def69059aad",
+        },
+        200,
+        {
+            "choices": [
+                {
+                    "finish_reason": "stop",
+                    "index": 0,
+                    "message": {
+                        "content": "1 plus 2 is 3.",
+                        "role": "assistant",
+                    },
+                }
+            ],
+            "created": 1696888865,
+            "id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat",
+            "model": "gpt-3.5-turbo-0613",
+            "object": "chat.completion",
+            "usage": {"completion_tokens": 11, "prompt_tokens": 53, "total_tokens": 64},
+        },
+    ),
 }
 
 

tests/mlmodel_openai/test_get_ai_message_ids.py

@@ -0,0 +1,211 @@
+# 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.
+
+import openai
+from testing_support.fixtures import reset_core_stats_engine
+
+from newrelic.api.background_task import background_task
+from newrelic.api.ml_model import get_ai_message_ids
+from newrelic.api.transaction import add_custom_attribute, current_transaction
+
+_test_openai_chat_completion_messages_1 = (
+    {"role": "system", "content": "You are a scientist."},
+    {"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"},
+)
+_test_openai_chat_completion_messages_2 = (
+    {"role": "system", "content": "You are a mathematician."},
+    {"role": "user", "content": "What is 1 plus 2?"},
+)
+expected_message_ids_1 = [
+    {
+        "conversation_id": "my-awesome-id",
+        "request_id": "49dbbffbd3c3f4612aa48def69059ccd",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-0",
+    },
+    {
+        "conversation_id": "my-awesome-id",
+        "request_id": "49dbbffbd3c3f4612aa48def69059ccd",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-1",
+    },
+    {
+        "conversation_id": "my-awesome-id",
+        "request_id": "49dbbffbd3c3f4612aa48def69059ccd",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-2",
+    },
+]
+
+expected_message_ids_1_no_conversation_id = [
+    {
+        "conversation_id": "",
+        "request_id": "49dbbffbd3c3f4612aa48def69059ccd",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-0",
+    },
+    {
+        "conversation_id": "",
+        "request_id": "49dbbffbd3c3f4612aa48def69059ccd",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-1",
+    },
+    {
+        "conversation_id": "",
+        "request_id": "49dbbffbd3c3f4612aa48def69059ccd",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-2",
+    },
+]
+expected_message_ids_2 = [
+    {
+        "conversation_id": "my-awesome-id",
+        "request_id": "49dbbffbd3c3f4612aa48def69059aad",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-0",
+    },
+    {
+        "conversation_id": "my-awesome-id",
+        "request_id": "49dbbffbd3c3f4612aa48def69059aad",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-1",
+    },
+    {
+        "conversation_id": "my-awesome-id",
+        "request_id": "49dbbffbd3c3f4612aa48def69059aad",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-2",
+    },
+]
+expected_message_ids_2_no_conversation_id = [
+    {
+        "conversation_id": "",
+        "request_id": "49dbbffbd3c3f4612aa48def69059aad",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-0",
+    },
+    {
+        "conversation_id": "",
+        "request_id": "49dbbffbd3c3f4612aa48def69059aad",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-1",
+    },
+    {
+        "conversation_id": "",
+        "request_id": "49dbbffbd3c3f4612aa48def69059aad",
+        "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-2",
+    },
+]
+
+
+@reset_core_stats_engine()
+@background_task()
+def test_get_ai_message_ids_when_nr_message_ids_not_set():
+    message_ids = get_ai_message_ids("request-id-1")
+    assert message_ids == []
+
+
+@reset_core_stats_engine()
+def test_get_ai_message_ids_outside_transaction():
+    message_ids = get_ai_message_ids("request-id-1")
+    assert message_ids == []
+
+
+@reset_core_stats_engine()
+@background_task()
+def test_get_ai_message_ids_mulitple_async(loop, set_trace_info):
+    set_trace_info()
+    add_custom_attribute("conversation_id", "my-awesome-id")
+
+    async def _run():
+        res1 = await openai.ChatCompletion.acreate(
+            model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100
+        )
+        res2 = await openai.ChatCompletion.acreate(
+            model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_2, temperature=0.7, max_tokens=100
+        )
+        return [res1, res2]
+
+    results = loop.run_until_complete(_run())
+
+    message_ids = [m for m in get_ai_message_ids(results[0].id)]
+    assert message_ids == expected_message_ids_1
+
+    message_ids = [m for m in get_ai_message_ids(results[1].id)]
+    assert message_ids == expected_message_ids_2
+
+    # Make sure we aren't causing a memory leak.
+    transaction = current_transaction()
+    assert not transaction._nr_message_ids
+
+
+@reset_core_stats_engine()
+@background_task()
+def test_get_ai_message_ids_mulitple_async_no_conversation_id(loop, set_trace_info):
+    set_trace_info()
+
+    async def _run():
+        res1 = await openai.ChatCompletion.acreate(
+            model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100
+        )
+        res2 = await openai.ChatCompletion.acreate(
+            model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_2, temperature=0.7, max_tokens=100
+        )
+        return [res1, res2]
+
+    results = loop.run_until_complete(_run())
+
+    message_ids = [m for m in get_ai_message_ids(results[0].id)]
+    assert message_ids == expected_message_ids_1_no_conversation_id
+
+    message_ids = [m for m in get_ai_message_ids(results[1].id)]
+    assert message_ids == expected_message_ids_2_no_conversation_id
+
+    # Make sure we aren't causing a memory leak.
+    transaction = current_transaction()
+    assert not transaction._nr_message_ids
+
+
+@reset_core_stats_engine()
+@background_task()
+def test_get_ai_message_ids_mulitple_sync(set_trace_info):
+    set_trace_info()
+    add_custom_attribute("conversation_id", "my-awesome-id")
+
+    results = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100
+    )
+    message_ids = [m for m in get_ai_message_ids(results.id)]
+    assert message_ids == expected_message_ids_1
+
+    results = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_2, temperature=0.7, max_tokens=100
+    )
+    message_ids = [m for m in get_ai_message_ids(results.id)]
+    assert message_ids == expected_message_ids_2
+
+    # Make sure we aren't causing a memory leak.
+    transaction = current_transaction()
+    assert not transaction._nr_message_ids
+
+
+@reset_core_stats_engine()
+@background_task()
+def test_get_ai_message_ids_mulitple_sync_no_conversation_id(set_trace_info):
+    set_trace_info()
+
+    results = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100
+    )
+    message_ids = [m for m in get_ai_message_ids(results.id)]
+    assert message_ids == expected_message_ids_1_no_conversation_id
+
+    results = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_2, temperature=0.7, max_tokens=100
+    )
+    message_ids = [m for m in get_ai_message_ids(results.id)]
+    assert message_ids == expected_message_ids_2_no_conversation_id
+
+    # Make sure we aren't causing a memory leak.
+    transaction = current_transaction()
+    assert not transaction._nr_message_ids