diff --git a/newrelic/agent.py b/newrelic/agent.py index b433f8e317..bc6cdbbd3a 100644 --- a/newrelic/agent.py +++ b/newrelic/agent.py @@ -153,7 +153,10 @@ 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 get_llm_message_ids as __get_llm_message_ids +from newrelic.api.ml_model import ( + record_llm_feedback_event as __record_llm_feedback_event, +) 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 @@ -341,4 +344,5 @@ 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") +get_llm_message_ids = __wrap_api_call(__get_llm_message_ids, "get_llm_message_ids") +record_llm_feedback_event = __wrap_api_call(__record_llm_feedback_event, "record_llm_feedback_event") diff --git a/newrelic/api/ml_model.py b/newrelic/api/ml_model.py index 91f3656849..d01042b359 100644 --- a/newrelic/api/ml_model.py +++ b/newrelic/api/ml_model.py @@ -13,6 +13,7 @@ # limitations under the License. import sys +import uuid import warnings from newrelic.api.transaction import current_transaction @@ -37,7 +38,7 @@ def wrap_mlmodel(model, name=None, version=None, feature_names=None, label_names model._nr_wrapped_metadata = metadata -def get_ai_message_ids(response_id=None): +def get_llm_message_ids(response_id=None): transaction = current_transaction() if response_id and transaction: nr_message_ids = getattr(transaction, "_nr_message_ids", {}) @@ -50,5 +51,34 @@ def get_ai_message_ids(response_id=None): 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.") + warnings.warn("No message ids found. get_llm_message_ids must be called within the scope of a transaction.") return [] + + +def record_llm_feedback_event( + message_id, rating, conversation_id=None, request_id=None, category=None, message=None, metadata=None +): + transaction = current_transaction() + if not transaction: + warnings.warn( + "No message feedback events will be recorded. record_llm_feedback_event must be called within the " + "scope of a transaction." + ) + return + + feedback_message_id = str(uuid.uuid4()) + metadata = metadata or {} + + feedback_message_event = { + "id": feedback_message_id, + "message_id": message_id, + "rating": rating, + "conversation_id": conversation_id or "", + "request_id": request_id or "", + "category": category or "", + "message": message or "", + "ingest_source": "Python", + } + feedback_message_event.update(metadata) + + transaction.record_ml_event("LlmFeedbackMessage", feedback_message_event) diff --git a/tests/agent_features/test_record_llm_feedback_event.py b/tests/agent_features/test_record_llm_feedback_event.py new file mode 100644 index 0000000000..59921ff400 --- /dev/null +++ b/tests/agent_features/test_record_llm_feedback_event.py @@ -0,0 +1,95 @@ +# 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. + +from testing_support.fixtures import reset_core_stats_engine +from testing_support.validators.validate_ml_event_count import validate_ml_event_count +from testing_support.validators.validate_ml_events import validate_ml_events + +from newrelic.api.background_task import background_task +from newrelic.api.ml_model import record_llm_feedback_event + + +@reset_core_stats_engine() +def test_record_llm_feedback_event_all_args_supplied(): + llm_feedback_all_args_recorded_events = [ + ( + {"type": "LlmFeedbackMessage"}, + { + "id": None, + "category": "informative", + "rating": 1, + "message_id": "message_id", + "request_id": "request_id", + "conversation_id": "conversation_id", + "ingest_source": "Python", + "message": "message", + "foo": "bar", + }, + ), + ] + + @validate_ml_events(llm_feedback_all_args_recorded_events) + @background_task() + def _test(): + record_llm_feedback_event( + rating=1, + message_id="message_id", + category="informative", + request_id="request_id", + conversation_id="conversation_id", + message="message", + metadata={"foo": "bar"}, + ) + + _test() + + +@reset_core_stats_engine() +def test_record_llm_feedback_event_required_args_supplied(): + llm_feedback_required_args_recorded_events = [ + ( + {"type": "LlmFeedbackMessage"}, + { + "id": None, + "category": "", + "rating": "Good", + "message_id": "message_id", + "request_id": "", + "conversation_id": "", + "ingest_source": "Python", + "message": "", + }, + ), + ] + + @validate_ml_events(llm_feedback_required_args_recorded_events) + @background_task() + def _test(): + record_llm_feedback_event(message_id="message_id", rating="Good") + + _test() + + +@reset_core_stats_engine() +@validate_ml_event_count(count=0) +def test_record_llm_feedback_event_outside_txn(): + record_llm_feedback_event( + rating="Good", + message_id="message_id", + category="informative", + request_id="request_id", + conversation_id="conversation_id", + message="message", + metadata={"foo": "bar"}, + ) diff --git a/tests/mlmodel_openai/test_get_ai_message_ids.py b/tests/mlmodel_openai/test_get_llm_message_ids.py similarity index 76% rename from tests/mlmodel_openai/test_get_ai_message_ids.py rename to tests/mlmodel_openai/test_get_llm_message_ids.py index 1a5c29f878..e20245128e 100644 --- a/tests/mlmodel_openai/test_get_ai_message_ids.py +++ b/tests/mlmodel_openai/test_get_llm_message_ids.py @@ -14,9 +14,10 @@ import openai from testing_support.fixtures import reset_core_stats_engine +from testing_support.validators.validate_ml_event_count import validate_ml_event_count from newrelic.api.background_task import background_task -from newrelic.api.ml_model import get_ai_message_ids +from newrelic.api.ml_model import get_llm_message_ids, record_llm_feedback_event from newrelic.api.transaction import add_custom_attribute, current_transaction _test_openai_chat_completion_messages_1 = ( @@ -100,20 +101,20 @@ @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") +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_ai_message_ids_outside_transaction(): - message_ids = get_ai_message_ids("request-id-1") +def test_get_llm_message_ids_outside_transaction(): + message_ids = get_llm_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): +def test_get_llm_message_ids_mulitple_async(loop, set_trace_info): set_trace_info() add_custom_attribute("conversation_id", "my-awesome-id") @@ -128,10 +129,10 @@ async def _run(): results = loop.run_until_complete(_run()) - message_ids = [m for m in get_ai_message_ids(results[0].id)] + message_ids = [m for m in get_llm_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)] + message_ids = [m for m in get_llm_message_ids(results[1].id)] assert message_ids == expected_message_ids_2 # Make sure we aren't causing a memory leak. @@ -141,7 +142,7 @@ async def _run(): @reset_core_stats_engine() @background_task() -def test_get_ai_message_ids_mulitple_async_no_conversation_id(loop, set_trace_info): +def test_get_llm_message_ids_mulitple_async_no_conversation_id(loop, set_trace_info): set_trace_info() async def _run(): @@ -155,10 +156,10 @@ async def _run(): results = loop.run_until_complete(_run()) - message_ids = [m for m in get_ai_message_ids(results[0].id)] + message_ids = [m for m in get_llm_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)] + message_ids = [m for m in get_llm_message_ids(results[1].id)] assert message_ids == expected_message_ids_2_no_conversation_id # Make sure we aren't causing a memory leak. @@ -167,21 +168,33 @@ async def _run(): @reset_core_stats_engine() +# Three chat completion messages and one chat completion summary for each create call (8 in total) +# Three feedback events for the first create call +@validate_ml_event_count(11) @background_task() -def test_get_ai_message_ids_mulitple_sync(set_trace_info): +def test_get_llm_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)] + message_ids = [m for m in get_llm_message_ids(results.id)] assert message_ids == expected_message_ids_1 + for message_id in message_ids: + record_llm_feedback_event( + category="informative", + rating=1, + message_id=message_id.get("message_id"), + request_id=message_id.get("request_id"), + conversation_id=message_id.get("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)] + message_ids = [m for m in get_llm_message_ids(results.id)] assert message_ids == expected_message_ids_2 # Make sure we aren't causing a memory leak. @@ -190,20 +203,30 @@ def test_get_ai_message_ids_mulitple_sync(set_trace_info): @reset_core_stats_engine() +@validate_ml_event_count(11) @background_task() -def test_get_ai_message_ids_mulitple_sync_no_conversation_id(set_trace_info): +def test_get_llm_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)] + message_ids = [m for m in get_llm_message_ids(results.id)] assert message_ids == expected_message_ids_1_no_conversation_id + for message_id in message_ids: + record_llm_feedback_event( + category="informative", + rating=1, + message_id=message_id.get("message_id"), + request_id=message_id.get("request_id"), + conversation_id=message_id.get("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)] + message_ids = [m for m in get_llm_message_ids(results.id)] assert message_ids == expected_message_ids_2_no_conversation_id # Make sure we aren't causing a memory leak.