From c84d105fbc5b1d1c3124fe60c17bf6bbbee3757a Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Tue, 17 Sep 2024 23:07:54 -0700 Subject: [PATCH 01/18] Add new custom attribute context manager. --- newrelic/agent.py | 2 + newrelic/api/llm_custom_attributes.py | 55 +++++++++++++++++++ newrelic/hooks/external_botocore.py | 3 + newrelic/hooks/mlmodel_langchain.py | 7 ++- newrelic/hooks/mlmodel_openai.py | 32 ++++++++--- .../test_llm_custom_attributes.py | 49 +++++++++++++++++ .../test_bedrock_chat_completion.py | 42 +++++++------- tests/mlmodel_langchain/test_tool.py | 5 +- tests/mlmodel_langchain/test_vectorstore.py | 52 ++++++++++-------- tests/mlmodel_openai/test_chat_completion.py | 23 ++++---- .../test_chat_completion_error.py | 17 +++--- .../test_chat_completion_error_v1.py | 23 +++++--- .../test_chat_completion_stream.py | 2 + .../test_chat_completion_stream_error.py | 1 + .../test_chat_completion_stream_error_v1.py | 1 + .../test_chat_completion_stream_v1.py | 29 +++++----- .../mlmodel_openai/test_chat_completion_v1.py | 24 ++++---- tests/testing_support/ml_testing_utils.py | 7 +++ 18 files changed, 273 insertions(+), 101 deletions(-) create mode 100644 newrelic/api/llm_custom_attributes.py create mode 100644 tests/agent_features/test_llm_custom_attributes.py diff --git a/newrelic/agent.py b/newrelic/agent.py index 4c07186268..ddc4824293 100644 --- a/newrelic/agent.py +++ b/newrelic/agent.py @@ -158,6 +158,7 @@ def __asgi_application(*args, **kwargs): ) from newrelic.api.ml_model import set_llm_token_count_callback as __set_llm_token_count_callback from newrelic.api.ml_model import wrap_mlmodel as __wrap_mlmodel +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes as __WithLlmCustomAttributes from newrelic.api.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper from newrelic.api.profile_trace import profile_trace as __profile_trace from newrelic.api.profile_trace import wrap_profile_trace as __wrap_profile_trace @@ -251,6 +252,7 @@ def __asgi_application(*args, **kwargs): record_custom_event = __wrap_api_call(__record_custom_event, "record_custom_event") record_log_event = __wrap_api_call(__record_log_event, "record_log_event") record_ml_event = __wrap_api_call(__record_ml_event, "record_ml_event") +WithLlmCustomAttributes = __wrap_api_call(__WithLlmCustomAttributes, "WithLlmCustomAttributes") accept_distributed_trace_payload = __wrap_api_call( __accept_distributed_trace_payload, "accept_distributed_trace_payload" ) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py new file mode 100644 index 0000000000..0e7b8622c7 --- /dev/null +++ b/newrelic/api/llm_custom_attributes.py @@ -0,0 +1,55 @@ +# 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 functools +import logging + +from newrelic.api.time_trace import TimeTrace, current_trace +from newrelic.api.transaction import current_transaction +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper +from newrelic.common.object_names import callable_name +from newrelic.common.object_wrapper import FunctionWrapper, wrap_object +from newrelic.core.function_node import FunctionNode + +_logger = logging.getLogger(__name__) + + +class WithLlmCustomAttributes(object): + def __init__(self, custom_attr_dict): + transaction = current_transaction() + if not isinstance(custom_attr_dict, dict) or custom_attr_dict is None: + raise TypeError("custom_attr_dict must be a dictionary. Received type: %s" % type(custom_attr_dict)) + + # Add "llm." prefix to all keys in attribute dictionary + prefixed_attr_dict = {} + for k, v in custom_attr_dict.items(): + if not k.startswith("llm."): + _logger.warning("Invalid attribute name %s. Renamed to llm.%s." % (k, k)) + prefixed_attr_dict["llm." + k] = v + + context_attrs = prefixed_attr_dict if prefixed_attr_dict else custom_attr_dict + + self.attr_dict = context_attrs + self.transaction = transaction + + def __enter__(self): + if not self.transaction: + _logger.warning("WithLlmCustomAttributes must be called within the scope of a transaction.") + return self + + self.transaction._llm_context_attrs = self.attr_dict + return self + + def __exit__(self, exc, value, tb): + self.transaction._llm_context_attrs = None diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index e4ade6be4a..515b276626 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -786,6 +786,9 @@ def handle_chat_completion_event(transaction, bedrock_attrs): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata_dict.update(llm_context_attrs) span_id = bedrock_attrs.get("span_id", None) trace_id = bedrock_attrs.get("trace_id", None) diff --git a/newrelic/hooks/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index 54a27693e0..4ff40bfffc 100644 --- a/newrelic/hooks/mlmodel_langchain.py +++ b/newrelic/hooks/mlmodel_langchain.py @@ -698,7 +698,7 @@ def _get_run_manager_info(transaction, run_args, instance, completion_id): # metadata and tags are keys in the config parameter. metadata = {} metadata.update((run_args.get("config") or {}).get("metadata") or {}) - # Do not report intenral nr_completion_id in metadata. + # Do not report internal nr_completion_id in metadata. metadata = {key: value for key, value in metadata.items() if key != "nr_completion_id"} tags = [] tags.extend((run_args.get("config") or {}).get("tags") or []) @@ -709,6 +709,11 @@ def _get_llm_metadata(transaction): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata_dict.update(llm_context_attrs) + return llm_metadata_dict diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 806aa91a63..3f5b2a7724 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -101,7 +101,7 @@ def wrap_chat_completion_sync(wrapped, instance, args, kwargs): try: return_val = wrapped(*args, **kwargs) except Exception as exc: - _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, None, exc) raise _handle_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, return_val) return return_val @@ -423,7 +423,7 @@ async def wrap_chat_completion_async(wrapped, instance, args, kwargs): try: return_val = await wrapped(*args, **kwargs) except Exception as exc: - _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, None, exc) raise _handle_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, return_val) @@ -436,6 +436,9 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages") or [] stream = kwargs.get("stream", False) + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + + # Only if streaming and streaming monitoring is enabled and the response is not empty # do we not exit the function trace. if not stream or not settings.ai_monitoring.streaming.enabled or not return_val: @@ -449,6 +452,7 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # The function trace will be exited when in the final iteration of the response # generator. setattr(return_val, "_nr_ft", ft) + setattr(return_val, "_nr_llm_context_attrs", llm_context_attrs) setattr(return_val, "_nr_openai_attrs", getattr(return_val, "_nr_openai_attrs", {})) return_val._nr_openai_attrs["messages"] = kwargs.get("messages", []) return_val._nr_openai_attrs["temperature"] = kwargs.get("temperature") @@ -475,12 +479,12 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # openai._legacy_response.LegacyAPIResponse response = json.loads(response.http_response.text.strip()) - _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response) + _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, None, response) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response): +def _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, llm_context_attrs, response): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") try: @@ -556,6 +560,8 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa "response.number_of_messages": len(input_message_list) + len(output_message_list), } llm_metadata = _get_llm_attributes(transaction) + if llm_context_attrs: + llm_metadata.update(llm_context_attrs) full_chat_completion_summary_dict.update(llm_metadata) transaction.record_custom_event("LlmChatCompletionSummary", full_chat_completion_summary_dict) @@ -576,7 +582,7 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc): +def _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, llm_context_attrs, exc): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages", None) or [] @@ -641,7 +647,9 @@ def _record_completion_error(transaction, linking_metadata, completion_id, kwarg "error": True, } llm_metadata = _get_llm_attributes(transaction) + llm_metadata.update(llm_context_attrs) error_chat_completion_dict.update(llm_metadata) + transaction.record_custom_event("LlmChatCompletionSummary", error_chat_completion_dict) output_message_list = [] @@ -780,15 +788,15 @@ def _record_events_on_stop_iteration(self, transaction): self._nr_ft.__exit__(None, None, None) try: openai_attrs = getattr(self, "_nr_openai_attrs", {}) - # If there are no openai attrs exit early as there's no data to record. if not openai_attrs: return completion_id = str(uuid.uuid4()) + llm_context_attrs = getattr(self, "nr_llm_context_attrs", None) response_headers = openai_attrs.get("response_headers") or {} _record_completion_success( - transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, response_headers, None + transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, response_headers, llm_context_attrs, None ) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) @@ -811,9 +819,10 @@ def _handle_streaming_completion_error(self, transaction, exc): if not openai_attrs: self._nr_ft.__exit__(*sys.exc_info()) return + llm_context_attrs = getattr(self, "_nr_llm_context_attrs", None) linking_metadata = get_trace_linking_metadata() completion_id = str(uuid.uuid4()) - _record_completion_error(transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, llm_context_attrs, exc) class AsyncGeneratorProxy(ObjectProxy): @@ -883,6 +892,8 @@ def set_attrs_on_generator_proxy(proxy, instance): proxy._nr_response_headers = instance._nr_response_headers if hasattr(instance, "_nr_openai_attrs"): proxy._nr_openai_attrs = instance._nr_openai_attrs + if hasattr(instance, "_nr_llm_context_attrs"): + proxy._nr_llm_context_attrs = instance._nr_llm_context_attrs def wrap_engine_api_resource_create_sync(wrapped, instance, args, kwargs): @@ -928,6 +939,11 @@ def _get_llm_attributes(transaction): """Returns llm.* custom attributes off of the transaction.""" custom_attrs_dict = transaction._custom_params llm_metadata = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata.update(llm_context_attrs) + return llm_metadata diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py new file mode 100644 index 0000000000..27c20a59ee --- /dev/null +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -0,0 +1,49 @@ +# 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 pytest + +from testing_support.fixtures import reset_core_stats_engine +from newrelic.api.background_task import background_task +from newrelic.api.transaction import current_transaction +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes + + +@background_task() +def test_llm_custom_attributes(): + transaction = current_transaction() + with WithLlmCustomAttributes({"test": "attr", "test1": "attr1"}): + assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + + assert transaction._llm_context_attrs is None + + +@pytest.mark.parametrize("context_attrs", (None, "not-a-dict")) +@background_task() +def test_llm_custom_attributes_no_attrs(context_attrs): + transaction = current_transaction() + + with pytest.raises(TypeError): + with WithLlmCustomAttributes(context_attrs): + assert transaction._llm_context_attrs is None + + +@background_task() +def test_llm_custom_attributes_prefixed_attrs(): + transaction = current_transaction() + with WithLlmCustomAttributes({"llm.test": "attr", "llm.test1": "attr1"}): + # Validate API does not prefix attributes that already begin with "llm." + assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + + assert transaction._llm_context_attrs is None \ No newline at end of file diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py index 7cab446348..938a85bb62 100644 --- a/tests/external_botocore/test_bedrock_chat_completion.py +++ b/tests/external_botocore/test_bedrock_chat_completion.py @@ -45,6 +45,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -59,6 +60,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name from newrelic.hooks.external_botocore import MODEL_EXTRACTORS @@ -161,7 +163,7 @@ def expected_invalid_access_key_error_events(model_id): def test_bedrock_chat_completion_in_txn_with_llm_metadata( set_trace_info, exercise_model, expected_events, expected_metrics ): - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) # One summary event, one user message, and one response message from the assistant @validate_custom_event_count(count=3) @validate_transaction_metrics( @@ -180,7 +182,8 @@ def _test(): 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=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100) + with WithLlmCustomAttributes({"context": "attr"}): + exercise_model(prompt=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100) _test() @@ -320,7 +323,7 @@ def _test(): def test_bedrock_chat_completion_error_invalid_model( bedrock_server, set_trace_info, response_streaming, expected_metrics ): - @validate_custom_events(chat_completion_invalid_model_error_events) + @validate_custom_events(events_with_context_attrs(chat_completion_invalid_model_error_events)) @validate_error_trace_attributes( "botocore.errorfactory:ValidationException", exact_attrs={ @@ -350,22 +353,23 @@ def _test(): add_custom_attribute("non_llm_attr", "python-agent") with pytest.raises(_client_error): - if response_streaming: - stream = bedrock_server.invoke_model_with_response_stream( - body=b"{}", - modelId="does-not-exist", - accept="application/json", - contentType="application/json", - ) - for _ in stream: - pass - else: - bedrock_server.invoke_model( - body=b"{}", - modelId="does-not-exist", - accept="application/json", - contentType="application/json", - ) + with WithLlmCustomAttributes({"context": "attr"}): + if response_streaming: + stream = bedrock_server.invoke_model_with_response_stream( + body=b"{}", + modelId="does-not-exist", + accept="application/json", + contentType="application/json", + ) + for _ in stream: + pass + else: + bedrock_server.invoke_model( + body=b"{}", + modelId="does-not-exist", + accept="application/json", + contentType="application/json", + ) _test() diff --git a/tests/mlmodel_langchain/test_tool.py b/tests/mlmodel_langchain/test_tool.py index 86a6716627..ed17cbd0f1 100644 --- a/tests/mlmodel_langchain/test_tool.py +++ b/tests/mlmodel_langchain/test_tool.py @@ -25,6 +25,7 @@ from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, + events_with_context_attrs, set_trace_info, ) from testing_support.validators.validate_custom_event import validate_custom_event_count @@ -40,6 +41,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name @@ -108,7 +110,8 @@ def events_sans_content(event): @background_task() def test_langchain_single_arg_tool(set_trace_info, single_arg_tool): set_trace_info() - single_arg_tool.run({"query": "Python Agent"}) + with WithLlmCustomAttributes({"context": "attrs"}): + single_arg_tool.run({"query": "Python Agent"}) @reset_core_stats_engine() diff --git a/tests/mlmodel_langchain/test_vectorstore.py b/tests/mlmodel_langchain/test_vectorstore.py index 41a9dfc146..3c97843d9d 100644 --- a/tests/mlmodel_langchain/test_vectorstore.py +++ b/tests/mlmodel_langchain/test_vectorstore.py @@ -23,6 +23,7 @@ from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, + events_with_context_attrs, set_trace_info, ) from testing_support.validators.validate_custom_event import validate_custom_event_count @@ -36,6 +37,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name @@ -125,7 +127,7 @@ def test_vectorstore_modules_instrumented(): @reset_core_stats_engine() -@validate_custom_events(vectorstore_recorded_events) +@validate_custom_events(events_with_context_attrs(vectorstore_recorded_events)) # Two OpenAI LlmEmbedded, two LangChain LlmVectorSearch @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -145,13 +147,14 @@ def test_pdf_pagesplitter_vectorstore_in_txn(set_trace_info, embedding_openai_cl add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = FAISS.from_documents(docs, embedding_openai_client) - docs = faiss_index.similarity_search("Complete this sentence: Hello", k=1) - assert "Hello world" in docs[0].page_content + faiss_index = FAISS.from_documents(docs, embedding_openai_client) + docs = faiss_index.similarity_search("Complete this sentence: Hello", k=1) + assert "Hello world" in docs[0].page_content @reset_core_stats_engine() @@ -216,7 +219,7 @@ def test_pdf_pagesplitter_vectorstore_ai_monitoring_disabled(set_trace_info, emb @reset_core_stats_engine() -@validate_custom_events(vectorstore_recorded_events) +@validate_custom_events(events_with_context_attrs(vectorstore_recorded_events)) # Two OpenAI LlmEmbedded, two LangChain LlmVectorSearch @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -237,13 +240,14 @@ async def _test(): add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = await FAISS.afrom_documents(docs, embedding_openai_client) - docs = await faiss_index.asimilarity_search("Complete this sentence: Hello", k=1) - return docs + faiss_index = await FAISS.afrom_documents(docs, embedding_openai_client) + docs = await faiss_index.asimilarity_search("Complete this sentence: Hello", k=1) + return docs docs = loop.run_until_complete(_test()) assert "Hello world" in docs[0].page_content @@ -343,7 +347,7 @@ async def _test(): callable_name(AssertionError), required_params={"user": ["vector_store_id"], "intrinsic": [], "agent": []}, ) -@validate_custom_events(vectorstore_error_events) +@validate_custom_events(events_with_context_attrs(vectorstore_error_events)) @validate_transaction_metrics( name="test_vectorstore:test_vectorstore_error", scoped_metrics=[("Llm/vectorstore/LangChain/similarity_search", 1)], @@ -356,13 +360,14 @@ async def _test(): @background_task() def test_vectorstore_error(set_trace_info, embedding_openai_client, loop): with pytest.raises(AssertionError): - set_trace_info() - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + set_trace_info() + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = FAISS.from_documents(docs, embedding_openai_client) - faiss_index.similarity_search(query="Complete this sentence: Hello", k=-1) + faiss_index = FAISS.from_documents(docs, embedding_openai_client) + faiss_index.similarity_search(query="Complete this sentence: Hello", k=-1) @reset_core_stats_engine() @@ -398,7 +403,7 @@ def test_vectorstore_error_no_content(set_trace_info, embedding_openai_client): callable_name(AssertionError), required_params={"user": ["vector_store_id"], "intrinsic": [], "agent": []}, ) -@validate_custom_events(vectorstore_error_events) +@validate_custom_events(events_with_context_attrs(vectorstore_error_events)) @validate_transaction_metrics( name="test_vectorstore:test_async_vectorstore_error", scoped_metrics=[("Llm/vectorstore/LangChain/asimilarity_search", 1)], @@ -422,7 +427,8 @@ async def _test(): return docs with pytest.raises(AssertionError): - loop.run_until_complete(_test()) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(_test()) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index 229ec44272..1d308ca2aa 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -36,6 +37,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, @@ -130,7 +132,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -147,10 +149,10 @@ def test_openai_chat_completion_sync_with_llm_metadata(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") - - openai.ChatCompletion.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -300,7 +302,7 @@ def test_openai_chat_completion_async_stream_monitoring_disabled(loop, set_trace @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion:test_openai_chat_completion_async_with_llm_metadata", @@ -319,11 +321,12 @@ def test_openai_chat_completion_async_with_llm_metadata(loop, set_trace_info): add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - loop.run_until_complete( - openai.ChatCompletion.acreate( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + openai.ChatCompletion.acreate( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion_error.py b/tests/mlmodel_openai/test_chat_completion_error.py index d3ed79bea9..9883c9b75b 100644 --- a/tests/mlmodel_openai/test_chat_completion_error.py +++ b/tests/mlmodel_openai/test_chat_completion_error.py @@ -24,6 +24,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -39,6 +40,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( @@ -120,19 +122,20 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - openai.ChatCompletion.create( - # no model provided, - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - ) + with WithLlmCustomAttributes({"context": "attr"}): + openai.ChatCompletion.create( + # no model provided, + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + ) @dt_enabled diff --git a/tests/mlmodel_openai/test_chat_completion_error_v1.py b/tests/mlmodel_openai/test_chat_completion_error_v1.py index 18c2bb7da6..ae1bb1bfef 100644 --- a/tests/mlmodel_openai/test_chat_completion_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_error_v1.py @@ -23,6 +23,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -38,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( @@ -116,16 +118,17 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info, sync_openai_client): with pytest.raises(TypeError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - sync_openai_client.chat.completions.create( - messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + sync_openai_client.chat.completions.create( + messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -182,18 +185,20 @@ def test_chat_completion_invalid_request_error_no_model_no_content(set_trace_inf rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_info, async_openai_client): with pytest.raises(TypeError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - loop.run_until_complete( - async_openai_client.chat.completions.create( - messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + async_openai_client.chat.completions.create( + messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion_stream.py b/tests/mlmodel_openai/test_chat_completion_stream.py index 3c32dd9f05..619df7f6d6 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream.py +++ b/tests/mlmodel_openai/test_chat_completion_stream.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -35,6 +36,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.api.transaction import add_custom_attribute disabled_custom_insights_settings = {"custom_insights_events.enabled": False} diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error.py b/tests/mlmodel_openai/test_chat_completion_stream_error.py index 2dc0400b7e..1b94b91247 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error.py @@ -39,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py index 33bda04f78..bc4aa41afe 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py @@ -39,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_v1.py index 7d268ced9a..d5c946b3b5 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_v1.py @@ -27,6 +27,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -37,6 +38,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.api.transaction import add_custom_attribute # TODO: Once instrumentation support is added for `.with_streaming_response.` @@ -142,7 +144,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant # @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -160,16 +162,17 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info, sync_open add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - generator = sync_openai_client.chat.completions.create( - model="gpt-3.5-turbo", - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, - ) + with WithLlmCustomAttributes({"context": "attr"}): + generator = sync_openai_client.chat.completions.create( + model="gpt-3.5-turbo", + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) - for resp in generator: - assert resp + for resp in generator: + assert resp @SKIP_IF_NO_OPENAI_WITH_STREAMING_RESPONSE @@ -471,7 +474,7 @@ async def consumer(): @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion_stream_v1:test_openai_chat_completion_async_with_llm_metadata", @@ -500,8 +503,8 @@ async def consumer(): ) async for resp in generator: assert resp - - loop.run_until_complete(consumer()) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(consumer()) @SKIP_IF_NO_OPENAI_WITH_STREAMING_RESPONSE diff --git a/tests/mlmodel_openai/test_chat_completion_v1.py b/tests/mlmodel_openai/test_chat_completion_v1.py index cf0c4f8491..3c8eeea867 100644 --- a/tests/mlmodel_openai/test_chat_completion_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_v1.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -36,6 +37,8 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes + _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, @@ -130,7 +133,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -147,10 +150,10 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info, sync_open add_custom_attribute("llm.conversation_id", "my-awesome-id") add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - - sync_openai_client.chat.completions.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + sync_openai_client.chat.completions.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -324,7 +327,7 @@ def test_openai_chat_completion_async_stream_monitoring_disabled(loop, set_trace @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion_v1:test_openai_chat_completion_async_with_llm_metadata", @@ -343,11 +346,12 @@ def test_openai_chat_completion_async_with_llm_metadata(loop, set_trace_info, as add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - loop.run_until_complete( - async_openai_client.chat.completions.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + async_openai_client.chat.completions.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/testing_support/ml_testing_utils.py b/tests/testing_support/ml_testing_utils.py index a9a74af17a..9d6923f95e 100644 --- a/tests/testing_support/ml_testing_utils.py +++ b/tests/testing_support/ml_testing_utils.py @@ -54,6 +54,13 @@ def events_sans_llm_metadata(expected_events): return events +def events_with_context_attrs(expected_events): + events = copy.deepcopy(expected_events) + for event in events: + event[1]["llm.context"] = "attr" + return events + + @pytest.fixture(scope="session") def set_trace_info(): def _set_trace_info(): From 5a7c183f249846de4a2fec40ad041a46ef4fe648 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Tue, 17 Sep 2024 23:07:54 -0700 Subject: [PATCH 02/18] Add new custom attribute context manager. --- newrelic/agent.py | 2 + newrelic/api/llm_custom_attributes.py | 55 +++++++++++++++++++ newrelic/hooks/external_botocore.py | 3 + newrelic/hooks/mlmodel_langchain.py | 7 ++- newrelic/hooks/mlmodel_openai.py | 32 ++++++++--- .../test_llm_custom_attributes.py | 49 +++++++++++++++++ .../test_bedrock_chat_completion.py | 42 +++++++------- tests/mlmodel_langchain/test_tool.py | 5 +- tests/mlmodel_langchain/test_vectorstore.py | 52 ++++++++++-------- tests/mlmodel_openai/test_chat_completion.py | 23 ++++---- .../test_chat_completion_error.py | 17 +++--- .../test_chat_completion_error_v1.py | 23 +++++--- .../test_chat_completion_stream.py | 2 + .../test_chat_completion_stream_error.py | 1 + .../test_chat_completion_stream_error_v1.py | 1 + .../test_chat_completion_stream_v1.py | 29 +++++----- .../mlmodel_openai/test_chat_completion_v1.py | 24 ++++---- tests/testing_support/ml_testing_utils.py | 7 +++ 18 files changed, 273 insertions(+), 101 deletions(-) create mode 100644 newrelic/api/llm_custom_attributes.py create mode 100644 tests/agent_features/test_llm_custom_attributes.py diff --git a/newrelic/agent.py b/newrelic/agent.py index 4c07186268..ddc4824293 100644 --- a/newrelic/agent.py +++ b/newrelic/agent.py @@ -158,6 +158,7 @@ def __asgi_application(*args, **kwargs): ) from newrelic.api.ml_model import set_llm_token_count_callback as __set_llm_token_count_callback from newrelic.api.ml_model import wrap_mlmodel as __wrap_mlmodel +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes as __WithLlmCustomAttributes from newrelic.api.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper from newrelic.api.profile_trace import profile_trace as __profile_trace from newrelic.api.profile_trace import wrap_profile_trace as __wrap_profile_trace @@ -251,6 +252,7 @@ def __asgi_application(*args, **kwargs): record_custom_event = __wrap_api_call(__record_custom_event, "record_custom_event") record_log_event = __wrap_api_call(__record_log_event, "record_log_event") record_ml_event = __wrap_api_call(__record_ml_event, "record_ml_event") +WithLlmCustomAttributes = __wrap_api_call(__WithLlmCustomAttributes, "WithLlmCustomAttributes") accept_distributed_trace_payload = __wrap_api_call( __accept_distributed_trace_payload, "accept_distributed_trace_payload" ) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py new file mode 100644 index 0000000000..0e7b8622c7 --- /dev/null +++ b/newrelic/api/llm_custom_attributes.py @@ -0,0 +1,55 @@ +# 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 functools +import logging + +from newrelic.api.time_trace import TimeTrace, current_trace +from newrelic.api.transaction import current_transaction +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper +from newrelic.common.object_names import callable_name +from newrelic.common.object_wrapper import FunctionWrapper, wrap_object +from newrelic.core.function_node import FunctionNode + +_logger = logging.getLogger(__name__) + + +class WithLlmCustomAttributes(object): + def __init__(self, custom_attr_dict): + transaction = current_transaction() + if not isinstance(custom_attr_dict, dict) or custom_attr_dict is None: + raise TypeError("custom_attr_dict must be a dictionary. Received type: %s" % type(custom_attr_dict)) + + # Add "llm." prefix to all keys in attribute dictionary + prefixed_attr_dict = {} + for k, v in custom_attr_dict.items(): + if not k.startswith("llm."): + _logger.warning("Invalid attribute name %s. Renamed to llm.%s." % (k, k)) + prefixed_attr_dict["llm." + k] = v + + context_attrs = prefixed_attr_dict if prefixed_attr_dict else custom_attr_dict + + self.attr_dict = context_attrs + self.transaction = transaction + + def __enter__(self): + if not self.transaction: + _logger.warning("WithLlmCustomAttributes must be called within the scope of a transaction.") + return self + + self.transaction._llm_context_attrs = self.attr_dict + return self + + def __exit__(self, exc, value, tb): + self.transaction._llm_context_attrs = None diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index e4ade6be4a..515b276626 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -786,6 +786,9 @@ def handle_chat_completion_event(transaction, bedrock_attrs): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata_dict.update(llm_context_attrs) span_id = bedrock_attrs.get("span_id", None) trace_id = bedrock_attrs.get("trace_id", None) diff --git a/newrelic/hooks/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index 54a27693e0..4ff40bfffc 100644 --- a/newrelic/hooks/mlmodel_langchain.py +++ b/newrelic/hooks/mlmodel_langchain.py @@ -698,7 +698,7 @@ def _get_run_manager_info(transaction, run_args, instance, completion_id): # metadata and tags are keys in the config parameter. metadata = {} metadata.update((run_args.get("config") or {}).get("metadata") or {}) - # Do not report intenral nr_completion_id in metadata. + # Do not report internal nr_completion_id in metadata. metadata = {key: value for key, value in metadata.items() if key != "nr_completion_id"} tags = [] tags.extend((run_args.get("config") or {}).get("tags") or []) @@ -709,6 +709,11 @@ def _get_llm_metadata(transaction): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata_dict.update(llm_context_attrs) + return llm_metadata_dict diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 806aa91a63..3f5b2a7724 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -101,7 +101,7 @@ def wrap_chat_completion_sync(wrapped, instance, args, kwargs): try: return_val = wrapped(*args, **kwargs) except Exception as exc: - _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, None, exc) raise _handle_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, return_val) return return_val @@ -423,7 +423,7 @@ async def wrap_chat_completion_async(wrapped, instance, args, kwargs): try: return_val = await wrapped(*args, **kwargs) except Exception as exc: - _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, None, exc) raise _handle_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, return_val) @@ -436,6 +436,9 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages") or [] stream = kwargs.get("stream", False) + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + + # Only if streaming and streaming monitoring is enabled and the response is not empty # do we not exit the function trace. if not stream or not settings.ai_monitoring.streaming.enabled or not return_val: @@ -449,6 +452,7 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # The function trace will be exited when in the final iteration of the response # generator. setattr(return_val, "_nr_ft", ft) + setattr(return_val, "_nr_llm_context_attrs", llm_context_attrs) setattr(return_val, "_nr_openai_attrs", getattr(return_val, "_nr_openai_attrs", {})) return_val._nr_openai_attrs["messages"] = kwargs.get("messages", []) return_val._nr_openai_attrs["temperature"] = kwargs.get("temperature") @@ -475,12 +479,12 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # openai._legacy_response.LegacyAPIResponse response = json.loads(response.http_response.text.strip()) - _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response) + _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, None, response) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response): +def _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, llm_context_attrs, response): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") try: @@ -556,6 +560,8 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa "response.number_of_messages": len(input_message_list) + len(output_message_list), } llm_metadata = _get_llm_attributes(transaction) + if llm_context_attrs: + llm_metadata.update(llm_context_attrs) full_chat_completion_summary_dict.update(llm_metadata) transaction.record_custom_event("LlmChatCompletionSummary", full_chat_completion_summary_dict) @@ -576,7 +582,7 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc): +def _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, llm_context_attrs, exc): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages", None) or [] @@ -641,7 +647,9 @@ def _record_completion_error(transaction, linking_metadata, completion_id, kwarg "error": True, } llm_metadata = _get_llm_attributes(transaction) + llm_metadata.update(llm_context_attrs) error_chat_completion_dict.update(llm_metadata) + transaction.record_custom_event("LlmChatCompletionSummary", error_chat_completion_dict) output_message_list = [] @@ -780,15 +788,15 @@ def _record_events_on_stop_iteration(self, transaction): self._nr_ft.__exit__(None, None, None) try: openai_attrs = getattr(self, "_nr_openai_attrs", {}) - # If there are no openai attrs exit early as there's no data to record. if not openai_attrs: return completion_id = str(uuid.uuid4()) + llm_context_attrs = getattr(self, "nr_llm_context_attrs", None) response_headers = openai_attrs.get("response_headers") or {} _record_completion_success( - transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, response_headers, None + transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, response_headers, llm_context_attrs, None ) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) @@ -811,9 +819,10 @@ def _handle_streaming_completion_error(self, transaction, exc): if not openai_attrs: self._nr_ft.__exit__(*sys.exc_info()) return + llm_context_attrs = getattr(self, "_nr_llm_context_attrs", None) linking_metadata = get_trace_linking_metadata() completion_id = str(uuid.uuid4()) - _record_completion_error(transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, llm_context_attrs, exc) class AsyncGeneratorProxy(ObjectProxy): @@ -883,6 +892,8 @@ def set_attrs_on_generator_proxy(proxy, instance): proxy._nr_response_headers = instance._nr_response_headers if hasattr(instance, "_nr_openai_attrs"): proxy._nr_openai_attrs = instance._nr_openai_attrs + if hasattr(instance, "_nr_llm_context_attrs"): + proxy._nr_llm_context_attrs = instance._nr_llm_context_attrs def wrap_engine_api_resource_create_sync(wrapped, instance, args, kwargs): @@ -928,6 +939,11 @@ def _get_llm_attributes(transaction): """Returns llm.* custom attributes off of the transaction.""" custom_attrs_dict = transaction._custom_params llm_metadata = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata.update(llm_context_attrs) + return llm_metadata diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py new file mode 100644 index 0000000000..27c20a59ee --- /dev/null +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -0,0 +1,49 @@ +# 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 pytest + +from testing_support.fixtures import reset_core_stats_engine +from newrelic.api.background_task import background_task +from newrelic.api.transaction import current_transaction +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes + + +@background_task() +def test_llm_custom_attributes(): + transaction = current_transaction() + with WithLlmCustomAttributes({"test": "attr", "test1": "attr1"}): + assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + + assert transaction._llm_context_attrs is None + + +@pytest.mark.parametrize("context_attrs", (None, "not-a-dict")) +@background_task() +def test_llm_custom_attributes_no_attrs(context_attrs): + transaction = current_transaction() + + with pytest.raises(TypeError): + with WithLlmCustomAttributes(context_attrs): + assert transaction._llm_context_attrs is None + + +@background_task() +def test_llm_custom_attributes_prefixed_attrs(): + transaction = current_transaction() + with WithLlmCustomAttributes({"llm.test": "attr", "llm.test1": "attr1"}): + # Validate API does not prefix attributes that already begin with "llm." + assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + + assert transaction._llm_context_attrs is None \ No newline at end of file diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py index 0c24fea0cc..ee094f3db5 100644 --- a/tests/external_botocore/test_bedrock_chat_completion.py +++ b/tests/external_botocore/test_bedrock_chat_completion.py @@ -45,6 +45,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -59,6 +60,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name from newrelic.hooks.external_botocore import MODEL_EXTRACTORS @@ -161,7 +163,7 @@ def expected_invalid_access_key_error_events(model_id): def test_bedrock_chat_completion_in_txn_with_llm_metadata( set_trace_info, exercise_model, expected_events, expected_metrics ): - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) # One summary event, one user message, and one response message from the assistant @validate_custom_event_count(count=3) @validate_transaction_metrics( @@ -180,7 +182,8 @@ def _test(): 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=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100) + with WithLlmCustomAttributes({"context": "attr"}): + exercise_model(prompt=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100) _test() @@ -320,7 +323,7 @@ def _test(): def test_bedrock_chat_completion_error_invalid_model( bedrock_server, set_trace_info, response_streaming, expected_metrics ): - @validate_custom_events(chat_completion_invalid_model_error_events) + @validate_custom_events(events_with_context_attrs(chat_completion_invalid_model_error_events)) @validate_error_trace_attributes( "botocore.errorfactory:ValidationException", exact_attrs={ @@ -350,22 +353,23 @@ def _test(): add_custom_attribute("non_llm_attr", "python-agent") with pytest.raises(_client_error): - if response_streaming: - stream = bedrock_server.invoke_model_with_response_stream( - body=b"{}", - modelId="does-not-exist", - accept="application/json", - contentType="application/json", - ) - for _ in stream: - pass - else: - bedrock_server.invoke_model( - body=b"{}", - modelId="does-not-exist", - accept="application/json", - contentType="application/json", - ) + with WithLlmCustomAttributes({"context": "attr"}): + if response_streaming: + stream = bedrock_server.invoke_model_with_response_stream( + body=b"{}", + modelId="does-not-exist", + accept="application/json", + contentType="application/json", + ) + for _ in stream: + pass + else: + bedrock_server.invoke_model( + body=b"{}", + modelId="does-not-exist", + accept="application/json", + contentType="application/json", + ) _test() diff --git a/tests/mlmodel_langchain/test_tool.py b/tests/mlmodel_langchain/test_tool.py index 1d41b41e86..46ec1d21b0 100644 --- a/tests/mlmodel_langchain/test_tool.py +++ b/tests/mlmodel_langchain/test_tool.py @@ -25,6 +25,7 @@ from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, + events_with_context_attrs, set_trace_info, ) from testing_support.validators.validate_custom_event import validate_custom_event_count @@ -40,6 +41,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name @@ -108,7 +110,8 @@ def events_sans_content(event): @background_task() def test_langchain_single_arg_tool(set_trace_info, single_arg_tool): set_trace_info() - single_arg_tool.run({"query": "Python Agent"}) + with WithLlmCustomAttributes({"context": "attrs"}): + single_arg_tool.run({"query": "Python Agent"}) @reset_core_stats_engine() diff --git a/tests/mlmodel_langchain/test_vectorstore.py b/tests/mlmodel_langchain/test_vectorstore.py index 41a9dfc146..3c97843d9d 100644 --- a/tests/mlmodel_langchain/test_vectorstore.py +++ b/tests/mlmodel_langchain/test_vectorstore.py @@ -23,6 +23,7 @@ from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, + events_with_context_attrs, set_trace_info, ) from testing_support.validators.validate_custom_event import validate_custom_event_count @@ -36,6 +37,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name @@ -125,7 +127,7 @@ def test_vectorstore_modules_instrumented(): @reset_core_stats_engine() -@validate_custom_events(vectorstore_recorded_events) +@validate_custom_events(events_with_context_attrs(vectorstore_recorded_events)) # Two OpenAI LlmEmbedded, two LangChain LlmVectorSearch @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -145,13 +147,14 @@ def test_pdf_pagesplitter_vectorstore_in_txn(set_trace_info, embedding_openai_cl add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = FAISS.from_documents(docs, embedding_openai_client) - docs = faiss_index.similarity_search("Complete this sentence: Hello", k=1) - assert "Hello world" in docs[0].page_content + faiss_index = FAISS.from_documents(docs, embedding_openai_client) + docs = faiss_index.similarity_search("Complete this sentence: Hello", k=1) + assert "Hello world" in docs[0].page_content @reset_core_stats_engine() @@ -216,7 +219,7 @@ def test_pdf_pagesplitter_vectorstore_ai_monitoring_disabled(set_trace_info, emb @reset_core_stats_engine() -@validate_custom_events(vectorstore_recorded_events) +@validate_custom_events(events_with_context_attrs(vectorstore_recorded_events)) # Two OpenAI LlmEmbedded, two LangChain LlmVectorSearch @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -237,13 +240,14 @@ async def _test(): add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = await FAISS.afrom_documents(docs, embedding_openai_client) - docs = await faiss_index.asimilarity_search("Complete this sentence: Hello", k=1) - return docs + faiss_index = await FAISS.afrom_documents(docs, embedding_openai_client) + docs = await faiss_index.asimilarity_search("Complete this sentence: Hello", k=1) + return docs docs = loop.run_until_complete(_test()) assert "Hello world" in docs[0].page_content @@ -343,7 +347,7 @@ async def _test(): callable_name(AssertionError), required_params={"user": ["vector_store_id"], "intrinsic": [], "agent": []}, ) -@validate_custom_events(vectorstore_error_events) +@validate_custom_events(events_with_context_attrs(vectorstore_error_events)) @validate_transaction_metrics( name="test_vectorstore:test_vectorstore_error", scoped_metrics=[("Llm/vectorstore/LangChain/similarity_search", 1)], @@ -356,13 +360,14 @@ async def _test(): @background_task() def test_vectorstore_error(set_trace_info, embedding_openai_client, loop): with pytest.raises(AssertionError): - set_trace_info() - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + set_trace_info() + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = FAISS.from_documents(docs, embedding_openai_client) - faiss_index.similarity_search(query="Complete this sentence: Hello", k=-1) + faiss_index = FAISS.from_documents(docs, embedding_openai_client) + faiss_index.similarity_search(query="Complete this sentence: Hello", k=-1) @reset_core_stats_engine() @@ -398,7 +403,7 @@ def test_vectorstore_error_no_content(set_trace_info, embedding_openai_client): callable_name(AssertionError), required_params={"user": ["vector_store_id"], "intrinsic": [], "agent": []}, ) -@validate_custom_events(vectorstore_error_events) +@validate_custom_events(events_with_context_attrs(vectorstore_error_events)) @validate_transaction_metrics( name="test_vectorstore:test_async_vectorstore_error", scoped_metrics=[("Llm/vectorstore/LangChain/asimilarity_search", 1)], @@ -422,7 +427,8 @@ async def _test(): return docs with pytest.raises(AssertionError): - loop.run_until_complete(_test()) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(_test()) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index 229ec44272..1d308ca2aa 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -36,6 +37,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, @@ -130,7 +132,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -147,10 +149,10 @@ def test_openai_chat_completion_sync_with_llm_metadata(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") - - openai.ChatCompletion.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -300,7 +302,7 @@ def test_openai_chat_completion_async_stream_monitoring_disabled(loop, set_trace @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion:test_openai_chat_completion_async_with_llm_metadata", @@ -319,11 +321,12 @@ def test_openai_chat_completion_async_with_llm_metadata(loop, set_trace_info): add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - loop.run_until_complete( - openai.ChatCompletion.acreate( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + openai.ChatCompletion.acreate( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion_error.py b/tests/mlmodel_openai/test_chat_completion_error.py index d3ed79bea9..9883c9b75b 100644 --- a/tests/mlmodel_openai/test_chat_completion_error.py +++ b/tests/mlmodel_openai/test_chat_completion_error.py @@ -24,6 +24,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -39,6 +40,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( @@ -120,19 +122,20 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - openai.ChatCompletion.create( - # no model provided, - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - ) + with WithLlmCustomAttributes({"context": "attr"}): + openai.ChatCompletion.create( + # no model provided, + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + ) @dt_enabled diff --git a/tests/mlmodel_openai/test_chat_completion_error_v1.py b/tests/mlmodel_openai/test_chat_completion_error_v1.py index 18c2bb7da6..ae1bb1bfef 100644 --- a/tests/mlmodel_openai/test_chat_completion_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_error_v1.py @@ -23,6 +23,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -38,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( @@ -116,16 +118,17 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info, sync_openai_client): with pytest.raises(TypeError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - sync_openai_client.chat.completions.create( - messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + sync_openai_client.chat.completions.create( + messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -182,18 +185,20 @@ def test_chat_completion_invalid_request_error_no_model_no_content(set_trace_inf rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_info, async_openai_client): with pytest.raises(TypeError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - loop.run_until_complete( - async_openai_client.chat.completions.create( - messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + async_openai_client.chat.completions.create( + messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion_stream.py b/tests/mlmodel_openai/test_chat_completion_stream.py index 3c32dd9f05..619df7f6d6 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream.py +++ b/tests/mlmodel_openai/test_chat_completion_stream.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -35,6 +36,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.api.transaction import add_custom_attribute disabled_custom_insights_settings = {"custom_insights_events.enabled": False} diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error.py b/tests/mlmodel_openai/test_chat_completion_stream_error.py index 2dc0400b7e..1b94b91247 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error.py @@ -39,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py index 33bda04f78..bc4aa41afe 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py @@ -39,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_v1.py index 7d268ced9a..d5c946b3b5 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_v1.py @@ -27,6 +27,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -37,6 +38,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.api.transaction import add_custom_attribute # TODO: Once instrumentation support is added for `.with_streaming_response.` @@ -142,7 +144,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant # @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -160,16 +162,17 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info, sync_open add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - generator = sync_openai_client.chat.completions.create( - model="gpt-3.5-turbo", - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, - ) + with WithLlmCustomAttributes({"context": "attr"}): + generator = sync_openai_client.chat.completions.create( + model="gpt-3.5-turbo", + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) - for resp in generator: - assert resp + for resp in generator: + assert resp @SKIP_IF_NO_OPENAI_WITH_STREAMING_RESPONSE @@ -471,7 +474,7 @@ async def consumer(): @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion_stream_v1:test_openai_chat_completion_async_with_llm_metadata", @@ -500,8 +503,8 @@ async def consumer(): ) async for resp in generator: assert resp - - loop.run_until_complete(consumer()) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(consumer()) @SKIP_IF_NO_OPENAI_WITH_STREAMING_RESPONSE diff --git a/tests/mlmodel_openai/test_chat_completion_v1.py b/tests/mlmodel_openai/test_chat_completion_v1.py index cf0c4f8491..3c8eeea867 100644 --- a/tests/mlmodel_openai/test_chat_completion_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_v1.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -36,6 +37,8 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes + _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, @@ -130,7 +133,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -147,10 +150,10 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info, sync_open add_custom_attribute("llm.conversation_id", "my-awesome-id") add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - - sync_openai_client.chat.completions.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + sync_openai_client.chat.completions.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -324,7 +327,7 @@ def test_openai_chat_completion_async_stream_monitoring_disabled(loop, set_trace @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion_v1:test_openai_chat_completion_async_with_llm_metadata", @@ -343,11 +346,12 @@ def test_openai_chat_completion_async_with_llm_metadata(loop, set_trace_info, as add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - loop.run_until_complete( - async_openai_client.chat.completions.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + async_openai_client.chat.completions.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/testing_support/ml_testing_utils.py b/tests/testing_support/ml_testing_utils.py index a9a74af17a..9d6923f95e 100644 --- a/tests/testing_support/ml_testing_utils.py +++ b/tests/testing_support/ml_testing_utils.py @@ -54,6 +54,13 @@ def events_sans_llm_metadata(expected_events): return events +def events_with_context_attrs(expected_events): + events = copy.deepcopy(expected_events) + for event in events: + event[1]["llm.context"] = "attr" + return events + + @pytest.fixture(scope="session") def set_trace_info(): def _set_trace_info(): From df1185870443092d000772a89338744b7d53ed99 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Fri, 20 Sep 2024 13:55:38 -0700 Subject: [PATCH 03/18] Add test cases. --- newrelic/api/llm_custom_attributes.py | 3 +- newrelic/hooks/external_botocore.py | 2 + newrelic/hooks/mlmodel_openai.py | 4 +- ...t_bedrock_chat_completion_via_langchain.py | 8 +- tests/mlmodel_langchain/test_chain.py | 62 +++++---- tests/mlmodel_langchain/test_tool.py | 33 ++--- tests/mlmodel_openai/test_chat_completion.py | 121 +++++++++++++++++- .../test_chat_completion_stream.py | 26 ++-- .../test_chat_completion_stream_error.py | 45 ++++--- 9 files changed, 221 insertions(+), 83 deletions(-) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py index 0e7b8622c7..f1be046023 100644 --- a/newrelic/api/llm_custom_attributes.py +++ b/newrelic/api/llm_custom_attributes.py @@ -52,4 +52,5 @@ def __enter__(self): return self def __exit__(self, exc, value, tb): - self.transaction._llm_context_attrs = None + if self.transaction: + self.transaction._llm_context_attrs = None diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 515b276626..cd5f5ae4f5 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -786,6 +786,7 @@ def handle_chat_completion_event(transaction, bedrock_attrs): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) if llm_context_attrs: llm_metadata_dict.update(llm_context_attrs) @@ -821,6 +822,7 @@ def handle_chat_completion_event(transaction, bedrock_attrs): "response.choices.finish_reason": bedrock_attrs.get("response.choices.finish_reason", None), "error": bedrock_attrs.get("error", None), } + chat_completion_summary_dict.update(llm_metadata_dict) chat_completion_summary_dict = {k: v for k, v in chat_completion_summary_dict.items() if v is not None} diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 3f5b2a7724..2e1dbee19e 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -560,6 +560,7 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa "response.number_of_messages": len(input_message_list) + len(output_message_list), } llm_metadata = _get_llm_attributes(transaction) + if llm_context_attrs: llm_metadata.update(llm_context_attrs) full_chat_completion_summary_dict.update(llm_metadata) @@ -647,7 +648,8 @@ def _record_completion_error(transaction, linking_metadata, completion_id, kwarg "error": True, } llm_metadata = _get_llm_attributes(transaction) - llm_metadata.update(llm_context_attrs) + if llm_context_attrs: + llm_metadata.update(llm_context_attrs) error_chat_completion_dict.update(llm_metadata) transaction.record_custom_event("LlmChatCompletionSummary", error_chat_completion_dict) diff --git a/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py b/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py index 00be00e17a..46ad8307eb 100644 --- a/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py +++ b/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py @@ -19,7 +19,7 @@ ) from conftest import BOTOCORE_VERSION # pylint: disable=E0611 from testing_support.fixtures import reset_core_stats_engine, validate_attributes -from testing_support.ml_testing_utils import set_trace_info # noqa: F401 +from testing_support.ml_testing_utils import set_trace_info, events_with_context_attrs # noqa: F401 from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( @@ -27,6 +27,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.api.transaction import add_custom_attribute UNSUPPORTED_LANGCHAIN_MODELS = [ @@ -105,7 +106,7 @@ def test_bedrock_chat_completion_in_txn_with_llm_metadata( expected_metrics, response_streaming, ): - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) # One summary event, one user message, and one response message from the assistant @validate_custom_event_count(count=6) @validate_transaction_metrics( @@ -124,6 +125,7 @@ def _test(): 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="Hi there!") + with WithLlmCustomAttributes({"context": "attr"}): + exercise_model(prompt="Hi there!") _test() diff --git a/tests/mlmodel_langchain/test_chain.py b/tests/mlmodel_langchain/test_chain.py index f3c2ab0859..a8cc4cfe13 100644 --- a/tests/mlmodel_langchain/test_chain.py +++ b/tests/mlmodel_langchain/test_chain.py @@ -31,6 +31,7 @@ disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, events_sans_content, + events_with_context_attrs, set_trace_info, ) from testing_support.validators.validate_custom_event import validate_custom_event_count @@ -47,6 +48,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( @@ -690,7 +692,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events_list_response) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events_list_response)) @validate_custom_event_count(count=7) @validate_transaction_metrics( name="test_chain:test_langchain_chain_list_response", @@ -720,10 +722,11 @@ def test_langchain_chain_list_response(set_trace_info, comma_separated_list_outp ] ) chain = chat_prompt | chat_openai_client | comma_separated_list_output_parser - chain.invoke( - {"text": "colors"}, - config={"metadata": {"id": "123"}}, - ) + with WithLlmCustomAttributes({"context": "attr"}): + chain.invoke( + {"text": "colors"}, + config={"metadata": {"id": "123"}}, + ) @pytest.mark.parametrize( @@ -991,7 +994,7 @@ def test_langchain_chain_error_in_openai( ): @reset_core_stats_engine() @validate_transaction_error_event_count(1) - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) @validate_custom_event_count(count=6) @validate_transaction_metrics( name="test_chain:test_langchain_chain_error_in_openai.._test", @@ -1012,7 +1015,8 @@ def _test(): runnable = create_function(json_schema, chat_openai_client, prompt_openai_error) with pytest.raises(openai.AuthenticationError): - getattr(runnable, call_function)(*call_function_args, **call_function_kwargs) + with WithLlmCustomAttributes({"context": "attr"}): + getattr(runnable, call_function)(*call_function_args, **call_function_kwargs) _test() @@ -1215,7 +1219,7 @@ def test_langchain_chain_ai_monitoring_disabled( @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events_list_response) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events_list_response)) @validate_custom_event_count(count=7) @validate_transaction_metrics( name="test_chain:test_async_langchain_chain_list_response", @@ -1247,15 +1251,15 @@ def test_async_langchain_chain_list_response( ] ) chain = chat_prompt | chat_openai_client | comma_separated_list_output_parser - - loop.run_until_complete( - chain.ainvoke( - {"text": "colors"}, - config={ - "metadata": {"id": "123"}, - }, + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + chain.ainvoke( + {"text": "colors"}, + config={ + "metadata": {"id": "123"}, + }, + ) ) - ) @reset_core_stats_engine() @@ -1495,7 +1499,7 @@ def test_async_langchain_chain_error_in_openai( ): @reset_core_stats_engine() @validate_transaction_error_event_count(1) - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) @validate_custom_event_count(count=6) @validate_transaction_metrics( name="test_chain:test_async_langchain_chain_error_in_openai.._test", @@ -1516,7 +1520,8 @@ def _test(): runnable = create_function(json_schema, chat_openai_client, prompt_openai_error) with pytest.raises(openai.AuthenticationError): - loop.run_until_complete(getattr(runnable, call_function)(*call_function_args, **call_function_kwargs)) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(getattr(runnable, call_function)(*call_function_args, **call_function_kwargs)) _test() @@ -1740,11 +1745,11 @@ def test_multiple_async_langchain_chain( expected_events, loop, ): - call1 = expected_events.copy() + call1 = events_with_context_attrs(expected_events.copy()) call1[0][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" call1[1][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" call1[2][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" - call2 = expected_events.copy() + call2 = events_with_context_attrs(expected_events.copy()) call2[0][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" call2[1][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" call2[2][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" @@ -1781,14 +1786,15 @@ def _test(): add_custom_attribute("non_llm_attr", "python-agent") runnable = create_function(json_schema, chat_openai_client, prompt) - - call1 = asyncio.ensure_future( - getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop - ) - call2 = asyncio.ensure_future( - getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop - ) - loop.run_until_complete(asyncio.gather(call1, call2)) + with WithLlmCustomAttributes({"context": "attr"}): + + call1 = asyncio.ensure_future( + getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop + ) + call2 = asyncio.ensure_future( + getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop + ) + loop.run_until_complete(asyncio.gather(call1, call2)) _test() diff --git a/tests/mlmodel_langchain/test_tool.py b/tests/mlmodel_langchain/test_tool.py index 46ec1d21b0..e421193e6f 100644 --- a/tests/mlmodel_langchain/test_tool.py +++ b/tests/mlmodel_langchain/test_tool.py @@ -95,7 +95,7 @@ def events_sans_content(event): @reset_core_stats_engine() -@validate_custom_events(single_arg_tool_recorded_events) +@validate_custom_events(events_with_context_attrs(single_arg_tool_recorded_events)) @validate_custom_event_count(count=1) @validate_transaction_metrics( name="test_tool:test_langchain_single_arg_tool", @@ -110,7 +110,7 @@ def events_sans_content(event): @background_task() def test_langchain_single_arg_tool(set_trace_info, single_arg_tool): set_trace_info() - with WithLlmCustomAttributes({"context": "attrs"}): + with WithLlmCustomAttributes({"context": "attr"}): single_arg_tool.run({"query": "Python Agent"}) @@ -135,7 +135,7 @@ def test_langchain_single_arg_tool_no_content(set_trace_info, single_arg_tool): @reset_core_stats_engine() -@validate_custom_events(single_arg_tool_recorded_events) +@validate_custom_events(events_with_context_attrs(single_arg_tool_recorded_events)) @validate_custom_event_count(count=1) @validate_transaction_metrics( name="test_tool:test_langchain_single_arg_tool_async", @@ -150,7 +150,8 @@ def test_langchain_single_arg_tool_no_content(set_trace_info, single_arg_tool): @background_task() def test_langchain_single_arg_tool_async(set_trace_info, single_arg_tool, loop): set_trace_info() - loop.run_until_complete(single_arg_tool.arun({"query": "Python Agent"})) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(single_arg_tool.arun({"query": "Python Agent"})) @reset_core_stats_engine() @@ -279,7 +280,7 @@ def test_langchain_multi_arg_tool_async(set_trace_info, multi_arg_tool, loop): "user": {}, }, ) -@validate_custom_events(multi_arg_error_recorded_events) +@validate_custom_events(events_with_context_attrs(multi_arg_error_recorded_events)) @validate_custom_event_count(count=1) @validate_transaction_metrics( name="test_tool:test_langchain_error_in_run", @@ -295,9 +296,10 @@ def test_langchain_error_in_run(set_trace_info, multi_arg_tool): with pytest.raises(pydantic_core._pydantic_core.ValidationError): set_trace_info() # Only one argument is provided while the tool expects two to create an error - multi_arg_tool.run( - {"first_num": 53}, tags=["test_tags", "python"], metadata={"test_run": True, "test": "langchain"} - ) + with WithLlmCustomAttributes({"context": "attr"}): + multi_arg_tool.run( + {"first_num": 53}, tags=["test_tags", "python"], metadata={"test_run": True, "test": "langchain"} + ) @reset_core_stats_engine() @@ -342,7 +344,7 @@ def test_langchain_error_in_run_no_content(set_trace_info, multi_arg_tool): "user": {}, }, ) -@validate_custom_events(multi_arg_error_recorded_events) +@validate_custom_events(events_with_context_attrs(multi_arg_error_recorded_events)) @validate_custom_event_count(count=1) @validate_transaction_metrics( name="test_tool:test_langchain_error_in_run_async", @@ -356,13 +358,14 @@ def test_langchain_error_in_run_no_content(set_trace_info, multi_arg_tool): @background_task() def test_langchain_error_in_run_async(set_trace_info, multi_arg_tool, loop): with pytest.raises(pydantic_core._pydantic_core.ValidationError): - set_trace_info() - # Only one argument is provided while the tool expects two to create an error - loop.run_until_complete( - multi_arg_tool.arun( - {"first_num": 53}, tags=["test_tags", "python"], metadata={"test_run": True, "test": "langchain"} + with WithLlmCustomAttributes({"context": "attr"}): + set_trace_info() + # Only one argument is provided while the tool expects two to create an error + loop.run_until_complete( + multi_arg_tool.arun( + {"first_num": 53}, tags=["test_tags", "python"], metadata={"test_run": True, "test": "langchain"} + ) ) - ) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index 1d308ca2aa..b14ea1a04b 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -35,10 +35,103 @@ validate_transaction_metrics, ) -from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute +from newrelic.api.background_task import background_task, BackgroundTask +from newrelic.api.transaction import add_custom_attribute, current_transaction from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.application import application_instance as application + +_test_openai_chat_completion_messages = ( + {"role": "system", "content": "You are a scientist."}, + {"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"}, +) + +testing = [ + ( + {"type": "LlmChatCompletionSummary"}, + { + "id": None, # UUID that varies with each run + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "span_id": None, + "trace_id": None, + "request_id": None, + "duration": None, # Response time varies each test run + "request.model": "gpt-3.5-turbo", + "response.model": "gpt-3.5-turbo-0613", + "response.organization": "new-relic-nkmd8b", + "request.temperature": 0.7, + "request.max_tokens": 100, + "response.choices.finish_reason": "stop", + "response.headers.llmVersion": "2020-10-01", + "response.headers.ratelimitLimitRequests": 200, + "response.headers.ratelimitLimitTokens": 40000, + "response.headers.ratelimitResetTokens": "90ms", + "response.headers.ratelimitResetRequests": "7m12s", + "response.headers.ratelimitRemainingTokens": 39940, + "response.headers.ratelimitRemainingRequests": 199, + "vendor": "openai", + "ingest_source": "Python", + "response.number_of_messages": 3, + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "You are a scientist.", + "role": "system", + "completion_id": None, + "sequence": 0, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "ingest_source": "Python", + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "What is 212 degrees Fahrenheit converted to Celsius?", + "role": "user", + "completion_id": None, + "sequence": 1, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "ingest_source": "Python", + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "212 degrees Fahrenheit is equal to 100 degrees Celsius.", + "role": "assistant", + "completion_id": None, + "sequence": 2, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "is_response": True, + "ingest_source": "Python", + }, + ), +] + _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, {"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"}, @@ -155,6 +248,30 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info): ) +@reset_core_stats_engine() +@validate_custom_events(testing + testing) +# One summary event, one system message, one user message, and one response message from the assistant +@background_task() +def test_openai_chat_completion_exit(set_trace_info): + set_trace_info() + transaction = current_transaction() + add_custom_attribute("llm.conversation_id", "my-awesome-id") + add_custom_attribute("llm.foo", "bar") + add_custom_attribute("non_llm_attr", "python-agent") + with WithLlmCustomAttributes({"context": "attr"}): + openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) + transaction.__exit__(None, None, None) + with BackgroundTask(application(), "fg") as txn: + add_custom_attribute("llm.conversation_id", "my-awesome-id") + add_custom_attribute("llm.foo", "bar") + add_custom_attribute("non_llm_attr", "python-agent") + openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) + + @reset_core_stats_engine() @disabled_ai_monitoring_record_content_settings @validate_custom_events(events_sans_content(chat_completion_recorded_events)) diff --git a/tests/mlmodel_openai/test_chat_completion_stream.py b/tests/mlmodel_openai/test_chat_completion_stream.py index 619df7f6d6..f7977b164f 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream.py +++ b/tests/mlmodel_openai/test_chat_completion_stream.py @@ -134,7 +134,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -152,15 +152,16 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info): add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - generator = openai.ChatCompletion.create( - model="gpt-3.5-turbo", - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, - ) - for resp in generator: - assert resp + with WithLlmCustomAttributes({"context": "attr"}): + generator = openai.ChatCompletion.create( + model="gpt-3.5-turbo", + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) + for resp in generator: + assert resp @reset_core_stats_engine() @@ -325,7 +326,7 @@ async def consumer(): @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion_stream:test_openai_chat_completion_async_with_llm_metadata", @@ -355,7 +356,8 @@ async def consumer(): async for resp in generator: assert resp - loop.run_until_complete(consumer()) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(consumer()) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error.py b/tests/mlmodel_openai/test_chat_completion_stream_error.py index 1b94b91247..2961728edd 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error.py @@ -24,6 +24,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -120,22 +121,23 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - generator = openai.ChatCompletion.create( - # no model provided, - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, - ) - for resp in generator: - assert resp + with WithLlmCustomAttributes({"context": "attr"}): + generator = openai.ChatCompletion.create( + # no model provided, + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) + for resp in generator: + assert resp @dt_enabled @@ -490,22 +492,23 @@ def test_chat_completion_wrong_api_key_error(monkeypatch, set_trace_info): rollup_metrics=[("Llm/completion/OpenAI/acreate", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_info): with pytest.raises(openai.InvalidRequestError): - set_trace_info() - add_custom_attribute("llm.conversation_id", "my-awesome-id") - loop.run_until_complete( - openai.ChatCompletion.acreate( - # no model provided, - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, + with WithLlmCustomAttributes({"context": "attr"}): + set_trace_info() + add_custom_attribute("llm.conversation_id", "my-awesome-id") + loop.run_until_complete( + openai.ChatCompletion.acreate( + # no model provided, + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) ) - ) @dt_enabled From 977f1b7f6d9bbb70e296e215461af0da6835871a Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Fri, 20 Sep 2024 14:45:38 -0700 Subject: [PATCH 04/18] Cleanup files. --- newrelic/hooks/external_botocore.py | 4 ---- tests/agent_features/test_llm_custom_attributes.py | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index fa19e5b155..93c4cfb2b3 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -790,10 +790,6 @@ def handle_chat_completion_event(transaction, bedrock_attrs): if llm_context_attrs: llm_metadata_dict.update(llm_context_attrs) - llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) - if llm_context_attrs: - llm_metadata_dict.update(llm_context_attrs) - span_id = bedrock_attrs.get("span_id", None) trace_id = bedrock_attrs.get("trace_id", None) request_id = bedrock_attrs.get("request_id", None) diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py index 27c20a59ee..204cd37909 100644 --- a/tests/agent_features/test_llm_custom_attributes.py +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -46,4 +46,4 @@ def test_llm_custom_attributes_prefixed_attrs(): # Validate API does not prefix attributes that already begin with "llm." assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} - assert transaction._llm_context_attrs is None \ No newline at end of file + assert transaction._llm_context_attrs is None From f51c6213ee8c6079bfc314e405e677279380a1c0 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Tue, 17 Sep 2024 23:07:54 -0700 Subject: [PATCH 05/18] Add new custom attribute context manager. --- newrelic/agent.py | 2 + newrelic/api/llm_custom_attributes.py | 55 +++++++++++++++++++ newrelic/hooks/external_botocore.py | 3 + newrelic/hooks/mlmodel_langchain.py | 7 ++- newrelic/hooks/mlmodel_openai.py | 32 ++++++++--- .../test_llm_custom_attributes.py | 49 +++++++++++++++++ .../test_bedrock_chat_completion.py | 42 +++++++------- tests/mlmodel_langchain/test_tool.py | 5 +- tests/mlmodel_langchain/test_vectorstore.py | 52 ++++++++++-------- tests/mlmodel_openai/test_chat_completion.py | 23 ++++---- .../test_chat_completion_error.py | 17 +++--- .../test_chat_completion_error_v1.py | 23 +++++--- .../test_chat_completion_stream.py | 2 + .../test_chat_completion_stream_error.py | 1 + .../test_chat_completion_stream_error_v1.py | 1 + .../test_chat_completion_stream_v1.py | 29 +++++----- .../mlmodel_openai/test_chat_completion_v1.py | 24 ++++---- tests/testing_support/ml_testing_utils.py | 7 +++ 18 files changed, 273 insertions(+), 101 deletions(-) create mode 100644 newrelic/api/llm_custom_attributes.py create mode 100644 tests/agent_features/test_llm_custom_attributes.py diff --git a/newrelic/agent.py b/newrelic/agent.py index 4c07186268..ddc4824293 100644 --- a/newrelic/agent.py +++ b/newrelic/agent.py @@ -158,6 +158,7 @@ def __asgi_application(*args, **kwargs): ) from newrelic.api.ml_model import set_llm_token_count_callback as __set_llm_token_count_callback from newrelic.api.ml_model import wrap_mlmodel as __wrap_mlmodel +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes as __WithLlmCustomAttributes from newrelic.api.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper from newrelic.api.profile_trace import profile_trace as __profile_trace from newrelic.api.profile_trace import wrap_profile_trace as __wrap_profile_trace @@ -251,6 +252,7 @@ def __asgi_application(*args, **kwargs): record_custom_event = __wrap_api_call(__record_custom_event, "record_custom_event") record_log_event = __wrap_api_call(__record_log_event, "record_log_event") record_ml_event = __wrap_api_call(__record_ml_event, "record_ml_event") +WithLlmCustomAttributes = __wrap_api_call(__WithLlmCustomAttributes, "WithLlmCustomAttributes") accept_distributed_trace_payload = __wrap_api_call( __accept_distributed_trace_payload, "accept_distributed_trace_payload" ) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py new file mode 100644 index 0000000000..0e7b8622c7 --- /dev/null +++ b/newrelic/api/llm_custom_attributes.py @@ -0,0 +1,55 @@ +# 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 functools +import logging + +from newrelic.api.time_trace import TimeTrace, current_trace +from newrelic.api.transaction import current_transaction +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper +from newrelic.common.object_names import callable_name +from newrelic.common.object_wrapper import FunctionWrapper, wrap_object +from newrelic.core.function_node import FunctionNode + +_logger = logging.getLogger(__name__) + + +class WithLlmCustomAttributes(object): + def __init__(self, custom_attr_dict): + transaction = current_transaction() + if not isinstance(custom_attr_dict, dict) or custom_attr_dict is None: + raise TypeError("custom_attr_dict must be a dictionary. Received type: %s" % type(custom_attr_dict)) + + # Add "llm." prefix to all keys in attribute dictionary + prefixed_attr_dict = {} + for k, v in custom_attr_dict.items(): + if not k.startswith("llm."): + _logger.warning("Invalid attribute name %s. Renamed to llm.%s." % (k, k)) + prefixed_attr_dict["llm." + k] = v + + context_attrs = prefixed_attr_dict if prefixed_attr_dict else custom_attr_dict + + self.attr_dict = context_attrs + self.transaction = transaction + + def __enter__(self): + if not self.transaction: + _logger.warning("WithLlmCustomAttributes must be called within the scope of a transaction.") + return self + + self.transaction._llm_context_attrs = self.attr_dict + return self + + def __exit__(self, exc, value, tb): + self.transaction._llm_context_attrs = None diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index f281c96097..a60ad04367 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -786,6 +786,9 @@ def handle_chat_completion_event(transaction, bedrock_attrs): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata_dict.update(llm_context_attrs) span_id = bedrock_attrs.get("span_id", None) trace_id = bedrock_attrs.get("trace_id", None) diff --git a/newrelic/hooks/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index 48d281a246..9c29514559 100644 --- a/newrelic/hooks/mlmodel_langchain.py +++ b/newrelic/hooks/mlmodel_langchain.py @@ -698,7 +698,7 @@ def _get_run_manager_info(transaction, run_args, instance, completion_id): # metadata and tags are keys in the config parameter. metadata = {} metadata.update((run_args.get("config") or {}).get("metadata") or {}) - # Do not report intenral nr_completion_id in metadata. + # Do not report internal nr_completion_id in metadata. metadata = {key: value for key, value in metadata.items() if key != "nr_completion_id"} tags = [] tags.extend((run_args.get("config") or {}).get("tags") or []) @@ -709,6 +709,11 @@ def _get_llm_metadata(transaction): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata_dict.update(llm_context_attrs) + return llm_metadata_dict diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 96228fd853..e6ee1e9225 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -101,7 +101,7 @@ def wrap_chat_completion_sync(wrapped, instance, args, kwargs): try: return_val = wrapped(*args, **kwargs) except Exception as exc: - _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, None, exc) raise _handle_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, return_val) return return_val @@ -423,7 +423,7 @@ async def wrap_chat_completion_async(wrapped, instance, args, kwargs): try: return_val = await wrapped(*args, **kwargs) except Exception as exc: - _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, None, exc) raise _handle_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, return_val) @@ -436,6 +436,9 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages") or [] stream = kwargs.get("stream", False) + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + + # Only if streaming and streaming monitoring is enabled and the response is not empty # do we not exit the function trace. if not stream or not settings.ai_monitoring.streaming.enabled or not return_val: @@ -449,6 +452,7 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # The function trace will be exited when in the final iteration of the response # generator. setattr(return_val, "_nr_ft", ft) + setattr(return_val, "_nr_llm_context_attrs", llm_context_attrs) setattr(return_val, "_nr_openai_attrs", getattr(return_val, "_nr_openai_attrs", {})) return_val._nr_openai_attrs["messages"] = kwargs.get("messages", []) return_val._nr_openai_attrs["temperature"] = kwargs.get("temperature") @@ -475,12 +479,12 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # openai._legacy_response.LegacyAPIResponse response = json.loads(response.http_response.text.strip()) - _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response) + _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, None, response) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response): +def _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, llm_context_attrs, response): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") try: @@ -556,6 +560,8 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa "response.number_of_messages": len(input_message_list) + len(output_message_list), } llm_metadata = _get_llm_attributes(transaction) + if llm_context_attrs: + llm_metadata.update(llm_context_attrs) full_chat_completion_summary_dict.update(llm_metadata) transaction.record_custom_event("LlmChatCompletionSummary", full_chat_completion_summary_dict) @@ -576,7 +582,7 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc): +def _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, llm_context_attrs, exc): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages", None) or [] @@ -641,7 +647,9 @@ def _record_completion_error(transaction, linking_metadata, completion_id, kwarg "error": True, } llm_metadata = _get_llm_attributes(transaction) + llm_metadata.update(llm_context_attrs) error_chat_completion_dict.update(llm_metadata) + transaction.record_custom_event("LlmChatCompletionSummary", error_chat_completion_dict) output_message_list = [] @@ -780,15 +788,15 @@ def _record_events_on_stop_iteration(self, transaction): self._nr_ft.__exit__(None, None, None) try: openai_attrs = getattr(self, "_nr_openai_attrs", {}) - # If there are no openai attrs exit early as there's no data to record. if not openai_attrs: return completion_id = str(uuid.uuid4()) + llm_context_attrs = getattr(self, "nr_llm_context_attrs", None) response_headers = openai_attrs.get("response_headers") or {} _record_completion_success( - transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, response_headers, None + transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, response_headers, llm_context_attrs, None ) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) @@ -811,9 +819,10 @@ def _handle_streaming_completion_error(self, transaction, exc): if not openai_attrs: self._nr_ft.__exit__(*sys.exc_info()) return + llm_context_attrs = getattr(self, "_nr_llm_context_attrs", None) linking_metadata = get_trace_linking_metadata() completion_id = str(uuid.uuid4()) - _record_completion_error(transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, exc) + _record_completion_error(transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, llm_context_attrs, exc) class AsyncGeneratorProxy(ObjectProxy): @@ -883,6 +892,8 @@ def set_attrs_on_generator_proxy(proxy, instance): proxy._nr_response_headers = instance._nr_response_headers if hasattr(instance, "_nr_openai_attrs"): proxy._nr_openai_attrs = instance._nr_openai_attrs + if hasattr(instance, "_nr_llm_context_attrs"): + proxy._nr_llm_context_attrs = instance._nr_llm_context_attrs def wrap_engine_api_resource_create_sync(wrapped, instance, args, kwargs): @@ -928,6 +939,11 @@ def _get_llm_attributes(transaction): """Returns llm.* custom attributes off of the transaction.""" custom_attrs_dict = transaction._custom_params llm_metadata = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata.update(llm_context_attrs) + return llm_metadata diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py new file mode 100644 index 0000000000..27c20a59ee --- /dev/null +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -0,0 +1,49 @@ +# 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 pytest + +from testing_support.fixtures import reset_core_stats_engine +from newrelic.api.background_task import background_task +from newrelic.api.transaction import current_transaction +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes + + +@background_task() +def test_llm_custom_attributes(): + transaction = current_transaction() + with WithLlmCustomAttributes({"test": "attr", "test1": "attr1"}): + assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + + assert transaction._llm_context_attrs is None + + +@pytest.mark.parametrize("context_attrs", (None, "not-a-dict")) +@background_task() +def test_llm_custom_attributes_no_attrs(context_attrs): + transaction = current_transaction() + + with pytest.raises(TypeError): + with WithLlmCustomAttributes(context_attrs): + assert transaction._llm_context_attrs is None + + +@background_task() +def test_llm_custom_attributes_prefixed_attrs(): + transaction = current_transaction() + with WithLlmCustomAttributes({"llm.test": "attr", "llm.test1": "attr1"}): + # Validate API does not prefix attributes that already begin with "llm." + assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + + assert transaction._llm_context_attrs is None \ No newline at end of file diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py index 8cc1fdaa8a..496d97e2c1 100644 --- a/tests/external_botocore/test_bedrock_chat_completion.py +++ b/tests/external_botocore/test_bedrock_chat_completion.py @@ -45,6 +45,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -59,6 +60,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name from newrelic.hooks.external_botocore import MODEL_EXTRACTORS @@ -161,7 +163,7 @@ def expected_invalid_access_key_error_events(model_id): def test_bedrock_chat_completion_in_txn_with_llm_metadata( set_trace_info, exercise_model, expected_events, expected_metrics ): - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) # One summary event, one user message, and one response message from the assistant @validate_custom_event_count(count=3) @validate_transaction_metrics( @@ -180,7 +182,8 @@ def _test(): 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=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100) + with WithLlmCustomAttributes({"context": "attr"}): + exercise_model(prompt=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100) _test() @@ -320,7 +323,7 @@ def _test(): def test_bedrock_chat_completion_error_invalid_model( bedrock_server, set_trace_info, response_streaming, expected_metrics ): - @validate_custom_events(chat_completion_invalid_model_error_events) + @validate_custom_events(events_with_context_attrs(chat_completion_invalid_model_error_events)) @validate_error_trace_attributes( "botocore.errorfactory:ValidationException", exact_attrs={ @@ -350,22 +353,23 @@ def _test(): add_custom_attribute("non_llm_attr", "python-agent") with pytest.raises(_client_error): - if response_streaming: - stream = bedrock_server.invoke_model_with_response_stream( - body=b"{}", - modelId="does-not-exist", - accept="application/json", - contentType="application/json", - ) - for _ in stream: - pass - else: - bedrock_server.invoke_model( - body=b"{}", - modelId="does-not-exist", - accept="application/json", - contentType="application/json", - ) + with WithLlmCustomAttributes({"context": "attr"}): + if response_streaming: + stream = bedrock_server.invoke_model_with_response_stream( + body=b"{}", + modelId="does-not-exist", + accept="application/json", + contentType="application/json", + ) + for _ in stream: + pass + else: + bedrock_server.invoke_model( + body=b"{}", + modelId="does-not-exist", + accept="application/json", + contentType="application/json", + ) _test() diff --git a/tests/mlmodel_langchain/test_tool.py b/tests/mlmodel_langchain/test_tool.py index a153c8200c..a2631e7240 100644 --- a/tests/mlmodel_langchain/test_tool.py +++ b/tests/mlmodel_langchain/test_tool.py @@ -25,6 +25,7 @@ from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, + events_with_context_attrs, set_trace_info, ) from testing_support.validators.validate_custom_event import validate_custom_event_count @@ -40,6 +41,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name @@ -108,7 +110,8 @@ def events_sans_content(event): @background_task() def test_langchain_single_arg_tool(set_trace_info, single_arg_tool): set_trace_info() - single_arg_tool.run({"query": "Python Agent"}) + with WithLlmCustomAttributes({"context": "attrs"}): + single_arg_tool.run({"query": "Python Agent"}) @reset_core_stats_engine() diff --git a/tests/mlmodel_langchain/test_vectorstore.py b/tests/mlmodel_langchain/test_vectorstore.py index d406277f22..ec1d25b8e8 100644 --- a/tests/mlmodel_langchain/test_vectorstore.py +++ b/tests/mlmodel_langchain/test_vectorstore.py @@ -23,6 +23,7 @@ from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, + events_with_context_attrs, set_trace_info, ) from testing_support.validators.validate_custom_event import validate_custom_event_count @@ -36,6 +37,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name @@ -123,7 +125,7 @@ def test_vectorstore_modules_instrumented(): @reset_core_stats_engine() -@validate_custom_events(vectorstore_recorded_events) +@validate_custom_events(events_with_context_attrs(vectorstore_recorded_events)) # Two OpenAI LlmEmbedded, two LangChain LlmVectorSearch @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -143,13 +145,14 @@ def test_pdf_pagesplitter_vectorstore_in_txn(set_trace_info, embedding_openai_cl add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = FAISS.from_documents(docs, embedding_openai_client) - docs = faiss_index.similarity_search("Complete this sentence: Hello", k=1) - assert "Hello world" in docs[0].page_content + faiss_index = FAISS.from_documents(docs, embedding_openai_client) + docs = faiss_index.similarity_search("Complete this sentence: Hello", k=1) + assert "Hello world" in docs[0].page_content @reset_core_stats_engine() @@ -214,7 +217,7 @@ def test_pdf_pagesplitter_vectorstore_ai_monitoring_disabled(set_trace_info, emb @reset_core_stats_engine() -@validate_custom_events(vectorstore_recorded_events) +@validate_custom_events(events_with_context_attrs(vectorstore_recorded_events)) # Two OpenAI LlmEmbedded, two LangChain LlmVectorSearch @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -235,13 +238,14 @@ async def _test(): add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = await FAISS.afrom_documents(docs, embedding_openai_client) - docs = await faiss_index.asimilarity_search("Complete this sentence: Hello", k=1) - return docs + faiss_index = await FAISS.afrom_documents(docs, embedding_openai_client) + docs = await faiss_index.asimilarity_search("Complete this sentence: Hello", k=1) + return docs docs = loop.run_until_complete(_test()) assert "Hello world" in docs[0].page_content @@ -341,7 +345,7 @@ async def _test(): callable_name(AssertionError), required_params={"user": ["vector_store_id"], "intrinsic": [], "agent": []}, ) -@validate_custom_events(vectorstore_error_events) +@validate_custom_events(events_with_context_attrs(vectorstore_error_events)) @validate_transaction_metrics( name="test_vectorstore:test_vectorstore_error", scoped_metrics=[("Llm/vectorstore/LangChain/similarity_search", 1)], @@ -354,13 +358,14 @@ async def _test(): @background_task() def test_vectorstore_error(set_trace_info, embedding_openai_client, loop): with pytest.raises(AssertionError): - set_trace_info() - script_dir = os.path.dirname(__file__) - loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) - docs = loader.load() + with WithLlmCustomAttributes({"context": "attr"}): + set_trace_info() + script_dir = os.path.dirname(__file__) + loader = PyPDFLoader(os.path.join(script_dir, "hello.pdf")) + docs = loader.load() - faiss_index = FAISS.from_documents(docs, embedding_openai_client) - faiss_index.similarity_search(query="Complete this sentence: Hello", k=-1) + faiss_index = FAISS.from_documents(docs, embedding_openai_client) + faiss_index.similarity_search(query="Complete this sentence: Hello", k=-1) @reset_core_stats_engine() @@ -396,7 +401,7 @@ def test_vectorstore_error_no_content(set_trace_info, embedding_openai_client): callable_name(AssertionError), required_params={"user": ["vector_store_id"], "intrinsic": [], "agent": []}, ) -@validate_custom_events(vectorstore_error_events) +@validate_custom_events(events_with_context_attrs(vectorstore_error_events)) @validate_transaction_metrics( name="test_vectorstore:test_async_vectorstore_error", scoped_metrics=[("Llm/vectorstore/LangChain/asimilarity_search", 1)], @@ -420,7 +425,8 @@ async def _test(): return docs with pytest.raises(AssertionError): - loop.run_until_complete(_test()) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(_test()) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index cbeb9cdd0d..6a0927d108 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -36,6 +37,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, @@ -130,7 +132,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -147,10 +149,10 @@ def test_openai_chat_completion_sync_with_llm_metadata(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") - - openai.ChatCompletion.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -300,7 +302,7 @@ def test_openai_chat_completion_async_stream_monitoring_disabled(loop, set_trace @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion:test_openai_chat_completion_async_with_llm_metadata", @@ -319,11 +321,12 @@ def test_openai_chat_completion_async_with_llm_metadata(loop, set_trace_info): add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - loop.run_until_complete( - openai.ChatCompletion.acreate( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + openai.ChatCompletion.acreate( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion_error.py b/tests/mlmodel_openai/test_chat_completion_error.py index d3ed79bea9..9883c9b75b 100644 --- a/tests/mlmodel_openai/test_chat_completion_error.py +++ b/tests/mlmodel_openai/test_chat_completion_error.py @@ -24,6 +24,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -39,6 +40,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( @@ -120,19 +122,20 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - openai.ChatCompletion.create( - # no model provided, - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - ) + with WithLlmCustomAttributes({"context": "attr"}): + openai.ChatCompletion.create( + # no model provided, + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + ) @dt_enabled diff --git a/tests/mlmodel_openai/test_chat_completion_error_v1.py b/tests/mlmodel_openai/test_chat_completion_error_v1.py index 18c2bb7da6..ae1bb1bfef 100644 --- a/tests/mlmodel_openai/test_chat_completion_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_error_v1.py @@ -23,6 +23,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -38,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( @@ -116,16 +118,17 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info, sync_openai_client): with pytest.raises(TypeError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - sync_openai_client.chat.completions.create( - messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + sync_openai_client.chat.completions.create( + messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -182,18 +185,20 @@ def test_chat_completion_invalid_request_error_no_model_no_content(set_trace_inf rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_info, async_openai_client): with pytest.raises(TypeError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - loop.run_until_complete( - async_openai_client.chat.completions.create( - messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + async_openai_client.chat.completions.create( + messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion_stream.py b/tests/mlmodel_openai/test_chat_completion_stream.py index 32420c78f1..1e3470c1ef 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream.py +++ b/tests/mlmodel_openai/test_chat_completion_stream.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -35,6 +36,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.api.transaction import add_custom_attribute disabled_custom_insights_settings = {"custom_insights_events.enabled": False} diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error.py b/tests/mlmodel_openai/test_chat_completion_stream_error.py index 2dc0400b7e..1b94b91247 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error.py @@ -39,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py index 33bda04f78..bc4aa41afe 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py @@ -39,6 +39,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_v1.py index c94cbef558..40811794b3 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_v1.py @@ -27,6 +27,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -37,6 +38,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.api.transaction import add_custom_attribute # TODO: Once instrumentation support is added for `.with_streaming_response.` @@ -142,7 +144,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant # @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -160,16 +162,17 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info, sync_open add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - generator = sync_openai_client.chat.completions.create( - model="gpt-3.5-turbo", - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, - ) + with WithLlmCustomAttributes({"context": "attr"}): + generator = sync_openai_client.chat.completions.create( + model="gpt-3.5-turbo", + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) - for resp in generator: - assert resp + for resp in generator: + assert resp @SKIP_IF_NO_OPENAI_WITH_STREAMING_RESPONSE @@ -471,7 +474,7 @@ async def consumer(): @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion_stream_v1:test_openai_chat_completion_async_with_llm_metadata", @@ -500,8 +503,8 @@ async def consumer(): ) async for resp in generator: assert resp - - loop.run_until_complete(consumer()) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(consumer()) @SKIP_IF_NO_OPENAI_WITH_STREAMING_RESPONSE diff --git a/tests/mlmodel_openai/test_chat_completion_v1.py b/tests/mlmodel_openai/test_chat_completion_v1.py index cbf631d550..420574f06b 100644 --- a/tests/mlmodel_openai/test_chat_completion_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_v1.py @@ -25,6 +25,7 @@ disabled_ai_monitoring_streaming_settings, events_sans_content, events_sans_llm_metadata, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -36,6 +37,8 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes + _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, @@ -130,7 +133,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -147,10 +150,10 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info, sync_open add_custom_attribute("llm.conversation_id", "my-awesome-id") add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - - sync_openai_client.chat.completions.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + sync_openai_client.chat.completions.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) @reset_core_stats_engine() @@ -324,7 +327,7 @@ def test_openai_chat_completion_async_stream_monitoring_disabled(loop, set_trace @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion_v1:test_openai_chat_completion_async_with_llm_metadata", @@ -343,11 +346,12 @@ def test_openai_chat_completion_async_with_llm_metadata(loop, set_trace_info, as add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - loop.run_until_complete( - async_openai_client.chat.completions.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + async_openai_client.chat.completions.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) ) - ) @reset_core_stats_engine() diff --git a/tests/testing_support/ml_testing_utils.py b/tests/testing_support/ml_testing_utils.py index a9a74af17a..9d6923f95e 100644 --- a/tests/testing_support/ml_testing_utils.py +++ b/tests/testing_support/ml_testing_utils.py @@ -54,6 +54,13 @@ def events_sans_llm_metadata(expected_events): return events +def events_with_context_attrs(expected_events): + events = copy.deepcopy(expected_events) + for event in events: + event[1]["llm.context"] = "attr" + return events + + @pytest.fixture(scope="session") def set_trace_info(): def _set_trace_info(): From f6a8c5da6f0c8f09778865ece0b64a7784ac76b3 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Fri, 20 Sep 2024 13:55:38 -0700 Subject: [PATCH 06/18] Add test cases. --- newrelic/api/llm_custom_attributes.py | 3 +- newrelic/hooks/external_botocore.py | 2 + newrelic/hooks/mlmodel_openai.py | 4 +- ...t_bedrock_chat_completion_via_langchain.py | 8 +- tests/mlmodel_langchain/test_chain.py | 62 +++++---- tests/mlmodel_langchain/test_tool.py | 33 ++--- tests/mlmodel_openai/test_chat_completion.py | 121 +++++++++++++++++- .../test_chat_completion_stream.py | 26 ++-- .../test_chat_completion_stream_error.py | 45 ++++--- 9 files changed, 221 insertions(+), 83 deletions(-) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py index 0e7b8622c7..f1be046023 100644 --- a/newrelic/api/llm_custom_attributes.py +++ b/newrelic/api/llm_custom_attributes.py @@ -52,4 +52,5 @@ def __enter__(self): return self def __exit__(self, exc, value, tb): - self.transaction._llm_context_attrs = None + if self.transaction: + self.transaction._llm_context_attrs = None diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index a60ad04367..f66e3b3056 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -786,6 +786,7 @@ def handle_chat_completion_event(transaction, bedrock_attrs): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) if llm_context_attrs: llm_metadata_dict.update(llm_context_attrs) @@ -821,6 +822,7 @@ def handle_chat_completion_event(transaction, bedrock_attrs): "response.choices.finish_reason": bedrock_attrs.get("response.choices.finish_reason", None), "error": bedrock_attrs.get("error", None), } + chat_completion_summary_dict.update(llm_metadata_dict) chat_completion_summary_dict = {k: v for k, v in chat_completion_summary_dict.items() if v is not None} diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index e6ee1e9225..ec9c0c655e 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -560,6 +560,7 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa "response.number_of_messages": len(input_message_list) + len(output_message_list), } llm_metadata = _get_llm_attributes(transaction) + if llm_context_attrs: llm_metadata.update(llm_context_attrs) full_chat_completion_summary_dict.update(llm_metadata) @@ -647,7 +648,8 @@ def _record_completion_error(transaction, linking_metadata, completion_id, kwarg "error": True, } llm_metadata = _get_llm_attributes(transaction) - llm_metadata.update(llm_context_attrs) + if llm_context_attrs: + llm_metadata.update(llm_context_attrs) error_chat_completion_dict.update(llm_metadata) transaction.record_custom_event("LlmChatCompletionSummary", error_chat_completion_dict) diff --git a/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py b/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py index 5f9b87b82f..afe88f4bf8 100644 --- a/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py +++ b/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py @@ -19,7 +19,7 @@ ) from conftest import BOTOCORE_VERSION # pylint: disable=E0611 from testing_support.fixtures import reset_core_stats_engine, validate_attributes -from testing_support.ml_testing_utils import set_trace_info # noqa: F401 +from testing_support.ml_testing_utils import set_trace_info, events_with_context_attrs # noqa: F401 from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( @@ -27,6 +27,7 @@ ) from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.api.transaction import add_custom_attribute UNSUPPORTED_LANGCHAIN_MODELS = [ @@ -105,7 +106,7 @@ def test_bedrock_chat_completion_in_txn_with_llm_metadata( expected_metrics, response_streaming, ): - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) # One summary event, one user message, and one response message from the assistant @validate_custom_event_count(count=6) @validate_transaction_metrics( @@ -124,6 +125,7 @@ def _test(): 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="Hi there!") + with WithLlmCustomAttributes({"context": "attr"}): + exercise_model(prompt="Hi there!") _test() diff --git a/tests/mlmodel_langchain/test_chain.py b/tests/mlmodel_langchain/test_chain.py index 6d8b2943d5..147f690647 100644 --- a/tests/mlmodel_langchain/test_chain.py +++ b/tests/mlmodel_langchain/test_chain.py @@ -31,6 +31,7 @@ disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, events_sans_content, + events_with_context_attrs, set_trace_info, ) from testing_support.validators.validate_custom_event import validate_custom_event_count @@ -47,6 +48,7 @@ from newrelic.api.background_task import background_task from newrelic.api.transaction import add_custom_attribute +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( @@ -690,7 +692,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events_list_response) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events_list_response)) @validate_custom_event_count(count=7) @validate_transaction_metrics( name="test_chain:test_langchain_chain_list_response", @@ -720,10 +722,11 @@ def test_langchain_chain_list_response(set_trace_info, comma_separated_list_outp ] ) chain = chat_prompt | chat_openai_client | comma_separated_list_output_parser - chain.invoke( - {"text": "colors"}, - config={"metadata": {"id": "123"}}, - ) + with WithLlmCustomAttributes({"context": "attr"}): + chain.invoke( + {"text": "colors"}, + config={"metadata": {"id": "123"}}, + ) @pytest.mark.parametrize( @@ -991,7 +994,7 @@ def test_langchain_chain_error_in_openai( ): @reset_core_stats_engine() @validate_transaction_error_event_count(1) - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) @validate_custom_event_count(count=6) @validate_transaction_metrics( name="test_chain:test_langchain_chain_error_in_openai.._test", @@ -1012,7 +1015,8 @@ def _test(): runnable = create_function(json_schema, chat_openai_client, prompt_openai_error) with pytest.raises(openai.AuthenticationError): - getattr(runnable, call_function)(*call_function_args, **call_function_kwargs) + with WithLlmCustomAttributes({"context": "attr"}): + getattr(runnable, call_function)(*call_function_args, **call_function_kwargs) _test() @@ -1215,7 +1219,7 @@ def test_langchain_chain_ai_monitoring_disabled( @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events_list_response) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events_list_response)) @validate_custom_event_count(count=7) @validate_transaction_metrics( name="test_chain:test_async_langchain_chain_list_response", @@ -1247,15 +1251,15 @@ def test_async_langchain_chain_list_response( ] ) chain = chat_prompt | chat_openai_client | comma_separated_list_output_parser - - loop.run_until_complete( - chain.ainvoke( - {"text": "colors"}, - config={ - "metadata": {"id": "123"}, - }, + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete( + chain.ainvoke( + {"text": "colors"}, + config={ + "metadata": {"id": "123"}, + }, + ) ) - ) @reset_core_stats_engine() @@ -1495,7 +1499,7 @@ def test_async_langchain_chain_error_in_openai( ): @reset_core_stats_engine() @validate_transaction_error_event_count(1) - @validate_custom_events(expected_events) + @validate_custom_events(events_with_context_attrs(expected_events)) @validate_custom_event_count(count=6) @validate_transaction_metrics( name="test_chain:test_async_langchain_chain_error_in_openai.._test", @@ -1516,7 +1520,8 @@ def _test(): runnable = create_function(json_schema, chat_openai_client, prompt_openai_error) with pytest.raises(openai.AuthenticationError): - loop.run_until_complete(getattr(runnable, call_function)(*call_function_args, **call_function_kwargs)) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(getattr(runnable, call_function)(*call_function_args, **call_function_kwargs)) _test() @@ -1740,11 +1745,11 @@ def test_multiple_async_langchain_chain( expected_events, loop, ): - call1 = expected_events.copy() + call1 = events_with_context_attrs(expected_events.copy()) call1[0][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" call1[1][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" call1[2][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" - call2 = expected_events.copy() + call2 = events_with_context_attrs(expected_events.copy()) call2[0][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" call2[1][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" call2[2][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" @@ -1781,14 +1786,15 @@ def _test(): add_custom_attribute("non_llm_attr", "python-agent") runnable = create_function(json_schema, chat_openai_client, prompt) - - call1 = asyncio.ensure_future( - getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop - ) - call2 = asyncio.ensure_future( - getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop - ) - loop.run_until_complete(asyncio.gather(call1, call2)) + with WithLlmCustomAttributes({"context": "attr"}): + + call1 = asyncio.ensure_future( + getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop + ) + call2 = asyncio.ensure_future( + getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop + ) + loop.run_until_complete(asyncio.gather(call1, call2)) _test() diff --git a/tests/mlmodel_langchain/test_tool.py b/tests/mlmodel_langchain/test_tool.py index a2631e7240..0714e84aa6 100644 --- a/tests/mlmodel_langchain/test_tool.py +++ b/tests/mlmodel_langchain/test_tool.py @@ -95,7 +95,7 @@ def events_sans_content(event): @reset_core_stats_engine() -@validate_custom_events(single_arg_tool_recorded_events) +@validate_custom_events(events_with_context_attrs(single_arg_tool_recorded_events)) @validate_custom_event_count(count=1) @validate_transaction_metrics( name="test_tool:test_langchain_single_arg_tool", @@ -110,7 +110,7 @@ def events_sans_content(event): @background_task() def test_langchain_single_arg_tool(set_trace_info, single_arg_tool): set_trace_info() - with WithLlmCustomAttributes({"context": "attrs"}): + with WithLlmCustomAttributes({"context": "attr"}): single_arg_tool.run({"query": "Python Agent"}) @@ -135,7 +135,7 @@ def test_langchain_single_arg_tool_no_content(set_trace_info, single_arg_tool): @reset_core_stats_engine() -@validate_custom_events(single_arg_tool_recorded_events) +@validate_custom_events(events_with_context_attrs(single_arg_tool_recorded_events)) @validate_custom_event_count(count=1) @validate_transaction_metrics( name="test_tool:test_langchain_single_arg_tool_async", @@ -150,7 +150,8 @@ def test_langchain_single_arg_tool_no_content(set_trace_info, single_arg_tool): @background_task() def test_langchain_single_arg_tool_async(set_trace_info, single_arg_tool, loop): set_trace_info() - loop.run_until_complete(single_arg_tool.arun({"query": "Python Agent"})) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(single_arg_tool.arun({"query": "Python Agent"})) @reset_core_stats_engine() @@ -279,7 +280,7 @@ def test_langchain_multi_arg_tool_async(set_trace_info, multi_arg_tool, loop): "user": {}, }, ) -@validate_custom_events(multi_arg_error_recorded_events) +@validate_custom_events(events_with_context_attrs(multi_arg_error_recorded_events)) @validate_custom_event_count(count=1) @validate_transaction_metrics( name="test_tool:test_langchain_error_in_run", @@ -295,9 +296,10 @@ def test_langchain_error_in_run(set_trace_info, multi_arg_tool): with pytest.raises(pydantic_core._pydantic_core.ValidationError): set_trace_info() # Only one argument is provided while the tool expects two to create an error - multi_arg_tool.run( - {"first_num": 53}, tags=["test_tags", "python"], metadata={"test_run": True, "test": "langchain"} - ) + with WithLlmCustomAttributes({"context": "attr"}): + multi_arg_tool.run( + {"first_num": 53}, tags=["test_tags", "python"], metadata={"test_run": True, "test": "langchain"} + ) @reset_core_stats_engine() @@ -342,7 +344,7 @@ def test_langchain_error_in_run_no_content(set_trace_info, multi_arg_tool): "user": {}, }, ) -@validate_custom_events(multi_arg_error_recorded_events) +@validate_custom_events(events_with_context_attrs(multi_arg_error_recorded_events)) @validate_custom_event_count(count=1) @validate_transaction_metrics( name="test_tool:test_langchain_error_in_run_async", @@ -356,13 +358,14 @@ def test_langchain_error_in_run_no_content(set_trace_info, multi_arg_tool): @background_task() def test_langchain_error_in_run_async(set_trace_info, multi_arg_tool, loop): with pytest.raises(pydantic_core._pydantic_core.ValidationError): - set_trace_info() - # Only one argument is provided while the tool expects two to create an error - loop.run_until_complete( - multi_arg_tool.arun( - {"first_num": 53}, tags=["test_tags", "python"], metadata={"test_run": True, "test": "langchain"} + with WithLlmCustomAttributes({"context": "attr"}): + set_trace_info() + # Only one argument is provided while the tool expects two to create an error + loop.run_until_complete( + multi_arg_tool.arun( + {"first_num": 53}, tags=["test_tags", "python"], metadata={"test_run": True, "test": "langchain"} + ) ) - ) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index 6a0927d108..f371b1cecd 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -35,10 +35,103 @@ validate_transaction_metrics, ) -from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute +from newrelic.api.background_task import background_task, BackgroundTask +from newrelic.api.transaction import add_custom_attribute, current_transaction from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.application import application_instance as application + +_test_openai_chat_completion_messages = ( + {"role": "system", "content": "You are a scientist."}, + {"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"}, +) + +testing = [ + ( + {"type": "LlmChatCompletionSummary"}, + { + "id": None, # UUID that varies with each run + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "span_id": None, + "trace_id": None, + "request_id": None, + "duration": None, # Response time varies each test run + "request.model": "gpt-3.5-turbo", + "response.model": "gpt-3.5-turbo-0613", + "response.organization": "new-relic-nkmd8b", + "request.temperature": 0.7, + "request.max_tokens": 100, + "response.choices.finish_reason": "stop", + "response.headers.llmVersion": "2020-10-01", + "response.headers.ratelimitLimitRequests": 200, + "response.headers.ratelimitLimitTokens": 40000, + "response.headers.ratelimitResetTokens": "90ms", + "response.headers.ratelimitResetRequests": "7m12s", + "response.headers.ratelimitRemainingTokens": 39940, + "response.headers.ratelimitRemainingRequests": 199, + "vendor": "openai", + "ingest_source": "Python", + "response.number_of_messages": 3, + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "You are a scientist.", + "role": "system", + "completion_id": None, + "sequence": 0, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "ingest_source": "Python", + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "What is 212 degrees Fahrenheit converted to Celsius?", + "role": "user", + "completion_id": None, + "sequence": 1, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "ingest_source": "Python", + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "212 degrees Fahrenheit is equal to 100 degrees Celsius.", + "role": "assistant", + "completion_id": None, + "sequence": 2, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "is_response": True, + "ingest_source": "Python", + }, + ), +] + _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, {"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"}, @@ -155,6 +248,30 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info): ) +@reset_core_stats_engine() +@validate_custom_events(testing + testing) +# One summary event, one system message, one user message, and one response message from the assistant +@background_task() +def test_openai_chat_completion_exit(set_trace_info): + set_trace_info() + transaction = current_transaction() + add_custom_attribute("llm.conversation_id", "my-awesome-id") + add_custom_attribute("llm.foo", "bar") + add_custom_attribute("non_llm_attr", "python-agent") + with WithLlmCustomAttributes({"context": "attr"}): + openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) + transaction.__exit__(None, None, None) + with BackgroundTask(application(), "fg") as txn: + add_custom_attribute("llm.conversation_id", "my-awesome-id") + add_custom_attribute("llm.foo", "bar") + add_custom_attribute("non_llm_attr", "python-agent") + openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 + ) + + @reset_core_stats_engine() @disabled_ai_monitoring_record_content_settings @validate_custom_events(events_sans_content(chat_completion_recorded_events)) diff --git a/tests/mlmodel_openai/test_chat_completion_stream.py b/tests/mlmodel_openai/test_chat_completion_stream.py index 1e3470c1ef..fce533350d 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream.py +++ b/tests/mlmodel_openai/test_chat_completion_stream.py @@ -134,7 +134,7 @@ @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) # One summary event, one system message, one user message, and one response message from the assistant @validate_custom_event_count(count=4) @validate_transaction_metrics( @@ -152,15 +152,16 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info): add_custom_attribute("llm.foo", "bar") add_custom_attribute("non_llm_attr", "python-agent") - generator = openai.ChatCompletion.create( - model="gpt-3.5-turbo", - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, - ) - for resp in generator: - assert resp + with WithLlmCustomAttributes({"context": "attr"}): + generator = openai.ChatCompletion.create( + model="gpt-3.5-turbo", + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) + for resp in generator: + assert resp @reset_core_stats_engine() @@ -325,7 +326,7 @@ async def consumer(): @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(events_with_context_attrs(chat_completion_recorded_events)) @validate_custom_event_count(count=4) @validate_transaction_metrics( "test_chat_completion_stream:test_openai_chat_completion_async_with_llm_metadata", @@ -355,7 +356,8 @@ async def consumer(): async for resp in generator: assert resp - loop.run_until_complete(consumer()) + with WithLlmCustomAttributes({"context": "attr"}): + loop.run_until_complete(consumer()) @reset_core_stats_engine() diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error.py b/tests/mlmodel_openai/test_chat_completion_stream_error.py index 1b94b91247..2961728edd 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error.py @@ -24,6 +24,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -120,22 +121,23 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - generator = openai.ChatCompletion.create( - # no model provided, - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, - ) - for resp in generator: - assert resp + with WithLlmCustomAttributes({"context": "attr"}): + generator = openai.ChatCompletion.create( + # no model provided, + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) + for resp in generator: + assert resp @dt_enabled @@ -490,22 +492,23 @@ def test_chat_completion_wrong_api_key_error(monkeypatch, set_trace_info): rollup_metrics=[("Llm/completion/OpenAI/acreate", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_info): with pytest.raises(openai.InvalidRequestError): - set_trace_info() - add_custom_attribute("llm.conversation_id", "my-awesome-id") - loop.run_until_complete( - openai.ChatCompletion.acreate( - # no model provided, - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - stream=True, + with WithLlmCustomAttributes({"context": "attr"}): + set_trace_info() + add_custom_attribute("llm.conversation_id", "my-awesome-id") + loop.run_until_complete( + openai.ChatCompletion.acreate( + # no model provided, + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + stream=True, + ) ) - ) @dt_enabled From 9cd6d6bbe88db1f7ca915ef71f197bcaaffd74c7 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Thu, 26 Sep 2024 11:58:51 -0700 Subject: [PATCH 07/18] Add new custom attribute context manager. --- newrelic/api/llm_custom_attributes.py | 20 +++++++++----------- newrelic/hooks/external_botocore.py | 3 +++ newrelic/hooks/mlmodel_openai.py | 1 - tests/mlmodel_openai/test_chat_completion.py | 6 ------ 4 files changed, 12 insertions(+), 18 deletions(-) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py index f1be046023..7e62450d5b 100644 --- a/newrelic/api/llm_custom_attributes.py +++ b/newrelic/api/llm_custom_attributes.py @@ -12,17 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -import functools +import contextvars import logging -from newrelic.api.time_trace import TimeTrace, current_trace from newrelic.api.transaction import current_transaction -from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper -from newrelic.common.object_names import callable_name -from newrelic.common.object_wrapper import FunctionWrapper, wrap_object -from newrelic.core.function_node import FunctionNode _logger = logging.getLogger(__name__) +custom_attr_context_var = contextvars.ContextVar("custom_attr_context_var", default={}) class WithLlmCustomAttributes(object): @@ -38,9 +34,9 @@ def __init__(self, custom_attr_dict): _logger.warning("Invalid attribute name %s. Renamed to llm.%s." % (k, k)) prefixed_attr_dict["llm." + k] = v - context_attrs = prefixed_attr_dict if prefixed_attr_dict else custom_attr_dict + finalized_attrs = prefixed_attr_dict if prefixed_attr_dict else custom_attr_dict - self.attr_dict = context_attrs + self.attr_dict = finalized_attrs self.transaction = transaction def __enter__(self): @@ -48,9 +44,11 @@ def __enter__(self): _logger.warning("WithLlmCustomAttributes must be called within the scope of a transaction.") return self - self.transaction._llm_context_attrs = self.attr_dict - return self + token = custom_attr_context_var.set(self.attr_dict) + self.transaction._custom_attr_context_var = custom_attr_context_var + return token def __exit__(self, exc, value, tb): if self.transaction: - self.transaction._llm_context_attrs = None + custom_attr_context_var.set(None) + self.transaction._custom_attr_context_var = custom_attr_context_var diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index f66e3b3056..a2e7a5976e 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -786,6 +786,9 @@ def handle_chat_completion_event(transaction, bedrock_attrs): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + if llm_context_attrs: + llm_metadata_dict.update(llm_context_attrs) llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) if llm_context_attrs: diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index ec9c0c655e..02e56993d2 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -560,7 +560,6 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa "response.number_of_messages": len(input_message_list) + len(output_message_list), } llm_metadata = _get_llm_attributes(transaction) - if llm_context_attrs: llm_metadata.update(llm_context_attrs) full_chat_completion_summary_dict.update(llm_metadata) diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index f371b1cecd..de69f5a7b4 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -34,7 +34,6 @@ from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, ) - from newrelic.api.background_task import background_task, BackgroundTask from newrelic.api.transaction import add_custom_attribute, current_transaction from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes @@ -132,11 +131,6 @@ ), ] -_test_openai_chat_completion_messages = ( - {"role": "system", "content": "You are a scientist."}, - {"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"}, -) - chat_completion_recorded_events = [ ( {"type": "LlmChatCompletionSummary"}, From db2ac31e9651b9469bd73d74ffafaf6977ae9833 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Fri, 20 Sep 2024 14:45:38 -0700 Subject: [PATCH 08/18] Cleanup files. --- newrelic/hooks/external_botocore.py | 4 ---- tests/agent_features/test_llm_custom_attributes.py | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index a2e7a5976e..a92e54a15f 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -790,10 +790,6 @@ def handle_chat_completion_event(transaction, bedrock_attrs): if llm_context_attrs: llm_metadata_dict.update(llm_context_attrs) - llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) - if llm_context_attrs: - llm_metadata_dict.update(llm_context_attrs) - span_id = bedrock_attrs.get("span_id", None) trace_id = bedrock_attrs.get("trace_id", None) request_id = bedrock_attrs.get("request_id", None) diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py index 27c20a59ee..204cd37909 100644 --- a/tests/agent_features/test_llm_custom_attributes.py +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -46,4 +46,4 @@ def test_llm_custom_attributes_prefixed_attrs(): # Validate API does not prefix attributes that already begin with "llm." assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} - assert transaction._llm_context_attrs is None \ No newline at end of file + assert transaction._llm_context_attrs is None From ba684ed1593a347162bd3b1838fb52fc3e6a3ff1 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Thu, 26 Sep 2024 12:00:11 -0700 Subject: [PATCH 09/18] Add contextvars. --- newrelic/hooks/external_botocore.py | 5 +- newrelic/hooks/mlmodel_langchain.py | 5 +- newrelic/hooks/mlmodel_openai.py | 27 +++-- .../test_llm_custom_attributes.py | 11 +- tests/mlmodel_openai/test_chat_completion.py | 104 ++++++++++++++++-- 5 files changed, 128 insertions(+), 24 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index a92e54a15f..ef70940f4d 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -786,9 +786,10 @@ def handle_chat_completion_event(transaction, bedrock_attrs): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} - llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + llm_context_attrs = getattr(transaction, "_custom_attr_context_var", None) if llm_context_attrs: - llm_metadata_dict.update(llm_context_attrs) + context_attrs = llm_context_attrs.get() + llm_metadata_dict.update(context_attrs) span_id = bedrock_attrs.get("span_id", None) trace_id = bedrock_attrs.get("trace_id", None) diff --git a/newrelic/hooks/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index 9c29514559..8bd6720fa3 100644 --- a/newrelic/hooks/mlmodel_langchain.py +++ b/newrelic/hooks/mlmodel_langchain.py @@ -710,9 +710,10 @@ def _get_llm_metadata(transaction): custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} - llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + llm_context_attrs = getattr(transaction, "_custom_attr_context_var", None) if llm_context_attrs: - llm_metadata_dict.update(llm_context_attrs) + context_attrs = llm_context_attrs.get() + llm_metadata_dict.update(context_attrs) return llm_metadata_dict diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 02e56993d2..8146871ed7 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -438,7 +438,6 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa stream = kwargs.get("stream", False) llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) - # Only if streaming and streaming monitoring is enabled and the response is not empty # do we not exit the function trace. if not stream or not settings.ai_monitoring.streaming.enabled or not return_val: @@ -479,12 +478,16 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # openai._legacy_response.LegacyAPIResponse response = json.loads(response.http_response.text.strip()) - _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, None, response) + _record_completion_success( + transaction, linking_metadata, completion_id, kwargs, ft, response_headers, None, response + ) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, llm_context_attrs, response): +def _record_completion_success( + transaction, linking_metadata, completion_id, kwargs, ft, response_headers, llm_context_attrs, response +): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") try: @@ -797,7 +800,14 @@ def _record_events_on_stop_iteration(self, transaction): llm_context_attrs = getattr(self, "nr_llm_context_attrs", None) response_headers = openai_attrs.get("response_headers") or {} _record_completion_success( - transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, response_headers, llm_context_attrs, None + transaction, + linking_metadata, + completion_id, + openai_attrs, + self._nr_ft, + response_headers, + llm_context_attrs, + None, ) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) @@ -823,7 +833,9 @@ def _handle_streaming_completion_error(self, transaction, exc): llm_context_attrs = getattr(self, "_nr_llm_context_attrs", None) linking_metadata = get_trace_linking_metadata() completion_id = str(uuid.uuid4()) - _record_completion_error(transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, llm_context_attrs, exc) + _record_completion_error( + transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, llm_context_attrs, exc + ) class AsyncGeneratorProxy(ObjectProxy): @@ -941,9 +953,10 @@ def _get_llm_attributes(transaction): custom_attrs_dict = transaction._custom_params llm_metadata = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} - llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) + llm_context_attrs = getattr(transaction, "_custom_attr_context_var", None) if llm_context_attrs: - llm_metadata.update(llm_context_attrs) + context_attrs = llm_context_attrs.get() + llm_metadata.update(context_attrs) return llm_metadata diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py index 204cd37909..3ce3b92082 100644 --- a/tests/agent_features/test_llm_custom_attributes.py +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -14,7 +14,6 @@ import pytest -from testing_support.fixtures import reset_core_stats_engine from newrelic.api.background_task import background_task from newrelic.api.transaction import current_transaction from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes @@ -24,9 +23,9 @@ def test_llm_custom_attributes(): transaction = current_transaction() with WithLlmCustomAttributes({"test": "attr", "test1": "attr1"}): - assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + assert transaction._custom_attr_context_var.get() == {"llm.test": "attr", "llm.test1": "attr1"} - assert transaction._llm_context_attrs is None + assert transaction._custom_attr_context_var.get() is None @pytest.mark.parametrize("context_attrs", (None, "not-a-dict")) @@ -36,7 +35,7 @@ def test_llm_custom_attributes_no_attrs(context_attrs): with pytest.raises(TypeError): with WithLlmCustomAttributes(context_attrs): - assert transaction._llm_context_attrs is None + assert transaction._custom_attr_context_var is None @background_task() @@ -44,6 +43,6 @@ def test_llm_custom_attributes_prefixed_attrs(): transaction = current_transaction() with WithLlmCustomAttributes({"llm.test": "attr", "llm.test1": "attr1"}): # Validate API does not prefix attributes that already begin with "llm." - assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + assert transaction._custom_attr_context_var.get() == {"llm.test": "attr", "llm.test1": "attr1"} - assert transaction._llm_context_attrs is None + assert transaction._custom_attr_context_var.get() is None diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index de69f5a7b4..45ee4d7b56 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -131,6 +131,92 @@ ), ] +testing = [ + ( + {"type": "LlmChatCompletionSummary"}, + { + "id": None, # UUID that varies with each run + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "span_id": None, + "trace_id": None, + "request_id": None, + "duration": None, # Response time varies each test run + "request.model": "gpt-3.5-turbo", + "response.model": "gpt-3.5-turbo-0613", + "response.organization": "new-relic-nkmd8b", + "request.temperature": 0.7, + "request.max_tokens": 100, + "response.choices.finish_reason": "stop", + "response.headers.llmVersion": "2020-10-01", + "response.headers.ratelimitLimitRequests": 200, + "response.headers.ratelimitLimitTokens": 40000, + "response.headers.ratelimitResetTokens": "90ms", + "response.headers.ratelimitResetRequests": "7m12s", + "response.headers.ratelimitRemainingTokens": 39940, + "response.headers.ratelimitRemainingRequests": 199, + "vendor": "openai", + "ingest_source": "Python", + "response.number_of_messages": 3, + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "You are a scientist.", + "role": "system", + "completion_id": None, + "sequence": 0, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "ingest_source": "Python", + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "What is 212 degrees Fahrenheit converted to Celsius?", + "role": "user", + "completion_id": None, + "sequence": 1, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "ingest_source": "Python", + }, + ), + ( + {"type": "LlmChatCompletionMessage"}, + { + "id": None, + "llm.conversation_id": "my-awesome-id", + "llm.foo": "bar", + "request_id": None, + "span_id": None, + "trace_id": None, + "content": "212 degrees Fahrenheit is equal to 100 degrees Celsius.", + "role": "assistant", + "completion_id": None, + "sequence": 2, + "response.model": "gpt-3.5-turbo-0613", + "vendor": "openai", + "is_response": True, + "ingest_source": "Python", + }, + ), +] + chat_completion_recorded_events = [ ( {"type": "LlmChatCompletionSummary"}, @@ -243,7 +329,7 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info): @reset_core_stats_engine() -@validate_custom_events(testing + testing) +@validate_custom_events(events_with_context_attrs(testing) + events_with_context_attrs(testing)) # One summary event, one system message, one user message, and one response message from the assistant @background_task() def test_openai_chat_completion_exit(set_trace_info): @@ -258,12 +344,16 @@ def test_openai_chat_completion_exit(set_trace_info): ) transaction.__exit__(None, None, None) with BackgroundTask(application(), "fg") as txn: - add_custom_attribute("llm.conversation_id", "my-awesome-id") - add_custom_attribute("llm.foo", "bar") - add_custom_attribute("non_llm_attr", "python-agent") - openai.ChatCompletion.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) + with WithLlmCustomAttributes({"context": "attr"}): + add_custom_attribute("llm.conversation_id", "my-awesome-id") + add_custom_attribute("llm.foo", "bar") + add_custom_attribute("non_llm_attr", "python-agent") + openai.ChatCompletion.create( + model="gpt-3.5-turbo", + messages=_test_openai_chat_completion_messages, + temperature=0.7, + max_tokens=100, + ) @reset_core_stats_engine() From 34ae93e5556457c31d77f226ddde0724dabb11c4 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Mon, 30 Sep 2024 09:58:51 -0700 Subject: [PATCH 10/18] Cleanup. --- newrelic/hooks/mlmodel_openai.py | 23 +-- tests/mlmodel_openai/test_chat_completion.py | 200 ------------------- 2 files changed, 6 insertions(+), 217 deletions(-) diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 8146871ed7..8442eebec3 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -101,7 +101,7 @@ def wrap_chat_completion_sync(wrapped, instance, args, kwargs): try: return_val = wrapped(*args, **kwargs) except Exception as exc: - _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, None, exc) + _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc) raise _handle_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, return_val) return return_val @@ -423,7 +423,7 @@ async def wrap_chat_completion_async(wrapped, instance, args, kwargs): try: return_val = await wrapped(*args, **kwargs) except Exception as exc: - _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, None, exc) + _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc) raise _handle_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, return_val) @@ -436,7 +436,6 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages") or [] stream = kwargs.get("stream", False) - llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) # Only if streaming and streaming monitoring is enabled and the response is not empty # do we not exit the function trace. @@ -451,7 +450,6 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # The function trace will be exited when in the final iteration of the response # generator. setattr(return_val, "_nr_ft", ft) - setattr(return_val, "_nr_llm_context_attrs", llm_context_attrs) setattr(return_val, "_nr_openai_attrs", getattr(return_val, "_nr_openai_attrs", {})) return_val._nr_openai_attrs["messages"] = kwargs.get("messages", []) return_val._nr_openai_attrs["temperature"] = kwargs.get("temperature") @@ -479,14 +477,14 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa response = json.loads(response.http_response.text.strip()) _record_completion_success( - transaction, linking_metadata, completion_id, kwargs, ft, response_headers, None, response + transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response ) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) def _record_completion_success( - transaction, linking_metadata, completion_id, kwargs, ft, response_headers, llm_context_attrs, response + transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response ): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") @@ -563,8 +561,6 @@ def _record_completion_success( "response.number_of_messages": len(input_message_list) + len(output_message_list), } llm_metadata = _get_llm_attributes(transaction) - if llm_context_attrs: - llm_metadata.update(llm_context_attrs) full_chat_completion_summary_dict.update(llm_metadata) transaction.record_custom_event("LlmChatCompletionSummary", full_chat_completion_summary_dict) @@ -585,7 +581,7 @@ def _record_completion_success( _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, llm_context_attrs, exc): +def _record_completion_error(transaction, linking_metadata, completion_id, kwargs, ft, exc): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages", None) or [] @@ -650,8 +646,6 @@ def _record_completion_error(transaction, linking_metadata, completion_id, kwarg "error": True, } llm_metadata = _get_llm_attributes(transaction) - if llm_context_attrs: - llm_metadata.update(llm_context_attrs) error_chat_completion_dict.update(llm_metadata) transaction.record_custom_event("LlmChatCompletionSummary", error_chat_completion_dict) @@ -797,7 +791,6 @@ def _record_events_on_stop_iteration(self, transaction): return completion_id = str(uuid.uuid4()) - llm_context_attrs = getattr(self, "nr_llm_context_attrs", None) response_headers = openai_attrs.get("response_headers") or {} _record_completion_success( transaction, @@ -806,7 +799,6 @@ def _record_events_on_stop_iteration(self, transaction): openai_attrs, self._nr_ft, response_headers, - llm_context_attrs, None, ) except Exception: @@ -830,11 +822,10 @@ def _handle_streaming_completion_error(self, transaction, exc): if not openai_attrs: self._nr_ft.__exit__(*sys.exc_info()) return - llm_context_attrs = getattr(self, "_nr_llm_context_attrs", None) linking_metadata = get_trace_linking_metadata() completion_id = str(uuid.uuid4()) _record_completion_error( - transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, llm_context_attrs, exc + transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, exc ) @@ -905,8 +896,6 @@ def set_attrs_on_generator_proxy(proxy, instance): proxy._nr_response_headers = instance._nr_response_headers if hasattr(instance, "_nr_openai_attrs"): proxy._nr_openai_attrs = instance._nr_openai_attrs - if hasattr(instance, "_nr_llm_context_attrs"): - proxy._nr_llm_context_attrs = instance._nr_llm_context_attrs def wrap_engine_api_resource_create_sync(wrapped, instance, args, kwargs): diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index 45ee4d7b56..1d17e3c739 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -45,178 +45,6 @@ {"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"}, ) -testing = [ - ( - {"type": "LlmChatCompletionSummary"}, - { - "id": None, # UUID that varies with each run - "llm.conversation_id": "my-awesome-id", - "llm.foo": "bar", - "span_id": None, - "trace_id": None, - "request_id": None, - "duration": None, # Response time varies each test run - "request.model": "gpt-3.5-turbo", - "response.model": "gpt-3.5-turbo-0613", - "response.organization": "new-relic-nkmd8b", - "request.temperature": 0.7, - "request.max_tokens": 100, - "response.choices.finish_reason": "stop", - "response.headers.llmVersion": "2020-10-01", - "response.headers.ratelimitLimitRequests": 200, - "response.headers.ratelimitLimitTokens": 40000, - "response.headers.ratelimitResetTokens": "90ms", - "response.headers.ratelimitResetRequests": "7m12s", - "response.headers.ratelimitRemainingTokens": 39940, - "response.headers.ratelimitRemainingRequests": 199, - "vendor": "openai", - "ingest_source": "Python", - "response.number_of_messages": 3, - }, - ), - ( - {"type": "LlmChatCompletionMessage"}, - { - "id": None, - "llm.conversation_id": "my-awesome-id", - "llm.foo": "bar", - "request_id": None, - "span_id": None, - "trace_id": None, - "content": "You are a scientist.", - "role": "system", - "completion_id": None, - "sequence": 0, - "response.model": "gpt-3.5-turbo-0613", - "vendor": "openai", - "ingest_source": "Python", - }, - ), - ( - {"type": "LlmChatCompletionMessage"}, - { - "id": None, - "llm.conversation_id": "my-awesome-id", - "llm.foo": "bar", - "request_id": None, - "span_id": None, - "trace_id": None, - "content": "What is 212 degrees Fahrenheit converted to Celsius?", - "role": "user", - "completion_id": None, - "sequence": 1, - "response.model": "gpt-3.5-turbo-0613", - "vendor": "openai", - "ingest_source": "Python", - }, - ), - ( - {"type": "LlmChatCompletionMessage"}, - { - "id": None, - "llm.conversation_id": "my-awesome-id", - "llm.foo": "bar", - "request_id": None, - "span_id": None, - "trace_id": None, - "content": "212 degrees Fahrenheit is equal to 100 degrees Celsius.", - "role": "assistant", - "completion_id": None, - "sequence": 2, - "response.model": "gpt-3.5-turbo-0613", - "vendor": "openai", - "is_response": True, - "ingest_source": "Python", - }, - ), -] - -testing = [ - ( - {"type": "LlmChatCompletionSummary"}, - { - "id": None, # UUID that varies with each run - "llm.conversation_id": "my-awesome-id", - "llm.foo": "bar", - "span_id": None, - "trace_id": None, - "request_id": None, - "duration": None, # Response time varies each test run - "request.model": "gpt-3.5-turbo", - "response.model": "gpt-3.5-turbo-0613", - "response.organization": "new-relic-nkmd8b", - "request.temperature": 0.7, - "request.max_tokens": 100, - "response.choices.finish_reason": "stop", - "response.headers.llmVersion": "2020-10-01", - "response.headers.ratelimitLimitRequests": 200, - "response.headers.ratelimitLimitTokens": 40000, - "response.headers.ratelimitResetTokens": "90ms", - "response.headers.ratelimitResetRequests": "7m12s", - "response.headers.ratelimitRemainingTokens": 39940, - "response.headers.ratelimitRemainingRequests": 199, - "vendor": "openai", - "ingest_source": "Python", - "response.number_of_messages": 3, - }, - ), - ( - {"type": "LlmChatCompletionMessage"}, - { - "id": None, - "llm.conversation_id": "my-awesome-id", - "llm.foo": "bar", - "request_id": None, - "span_id": None, - "trace_id": None, - "content": "You are a scientist.", - "role": "system", - "completion_id": None, - "sequence": 0, - "response.model": "gpt-3.5-turbo-0613", - "vendor": "openai", - "ingest_source": "Python", - }, - ), - ( - {"type": "LlmChatCompletionMessage"}, - { - "id": None, - "llm.conversation_id": "my-awesome-id", - "llm.foo": "bar", - "request_id": None, - "span_id": None, - "trace_id": None, - "content": "What is 212 degrees Fahrenheit converted to Celsius?", - "role": "user", - "completion_id": None, - "sequence": 1, - "response.model": "gpt-3.5-turbo-0613", - "vendor": "openai", - "ingest_source": "Python", - }, - ), - ( - {"type": "LlmChatCompletionMessage"}, - { - "id": None, - "llm.conversation_id": "my-awesome-id", - "llm.foo": "bar", - "request_id": None, - "span_id": None, - "trace_id": None, - "content": "212 degrees Fahrenheit is equal to 100 degrees Celsius.", - "role": "assistant", - "completion_id": None, - "sequence": 2, - "response.model": "gpt-3.5-turbo-0613", - "vendor": "openai", - "is_response": True, - "ingest_source": "Python", - }, - ), -] - chat_completion_recorded_events = [ ( {"type": "LlmChatCompletionSummary"}, @@ -328,34 +156,6 @@ def test_openai_chat_completion_sync_with_llm_metadata(set_trace_info): ) -@reset_core_stats_engine() -@validate_custom_events(events_with_context_attrs(testing) + events_with_context_attrs(testing)) -# One summary event, one system message, one user message, and one response message from the assistant -@background_task() -def test_openai_chat_completion_exit(set_trace_info): - set_trace_info() - transaction = current_transaction() - add_custom_attribute("llm.conversation_id", "my-awesome-id") - add_custom_attribute("llm.foo", "bar") - add_custom_attribute("non_llm_attr", "python-agent") - with WithLlmCustomAttributes({"context": "attr"}): - openai.ChatCompletion.create( - model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 - ) - transaction.__exit__(None, None, None) - with BackgroundTask(application(), "fg") as txn: - with WithLlmCustomAttributes({"context": "attr"}): - add_custom_attribute("llm.conversation_id", "my-awesome-id") - add_custom_attribute("llm.foo", "bar") - add_custom_attribute("non_llm_attr", "python-agent") - openai.ChatCompletion.create( - model="gpt-3.5-turbo", - messages=_test_openai_chat_completion_messages, - temperature=0.7, - max_tokens=100, - ) - - @reset_core_stats_engine() @disabled_ai_monitoring_record_content_settings @validate_custom_events(events_sans_content(chat_completion_recorded_events)) From b38c89b9a5695c1c70d5c78318e10b235887f839 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Wed, 2 Oct 2024 15:04:51 -0400 Subject: [PATCH 11/18] Update error handling. --- newrelic/api/llm_custom_attributes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py index 7b5eea29a2..b4f0ffc025 100644 --- a/newrelic/api/llm_custom_attributes.py +++ b/newrelic/api/llm_custom_attributes.py @@ -24,8 +24,8 @@ class WithLlmCustomAttributes(object): def __init__(self, custom_attr_dict): transaction = current_transaction() - if not isinstance(custom_attr_dict, dict) or custom_attr_dict is None: - raise TypeError("custom_attr_dict must be a dictionary. Received type: %s" % type(custom_attr_dict)) + if not custom_attr_dict or not isinstance(custom_attr_dict, dict): + raise TypeError("custom_attr_dict must be a non-empty dictionary. Received type: %s" % type(custom_attr_dict)) # Add "llm." prefix to all keys in attribute dictionary context_attrs = {k if k.startswith("llm.") else f"llm.{k}": v for k, v in custom_attr_dict.items()} From 18278db0e321ef1b80e01b89511158536ee1867c Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Mon, 7 Oct 2024 13:39:21 -0700 Subject: [PATCH 12/18] Undo unnecessary formatting changes --- newrelic/hooks/external_botocore.py | 1 - newrelic/hooks/mlmodel_openai.py | 24 +++++------------------- 2 files changed, 5 insertions(+), 20 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 3ff7ece19a..4f17cca917 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -823,7 +823,6 @@ def handle_chat_completion_event(transaction, bedrock_attrs): "response.choices.finish_reason": bedrock_attrs.get("response.choices.finish_reason", None), "error": bedrock_attrs.get("error", None), } - chat_completion_summary_dict.update(llm_metadata_dict) chat_completion_summary_dict = {k: v for k, v in chat_completion_summary_dict.items() if v is not None} diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 41b6a30390..b2085c790f 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -436,7 +436,6 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa trace_id = linking_metadata.get("trace.id") request_message_list = kwargs.get("messages") or [] stream = kwargs.get("stream", False) - # Only if streaming and streaming monitoring is enabled and the response is not empty # do we not exit the function trace. if not stream or not settings.ai_monitoring.streaming.enabled or not return_val: @@ -476,16 +475,12 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa # openai._legacy_response.LegacyAPIResponse response = json.loads(response.http_response.text.strip()) - _record_completion_success( - transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response - ) + _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) -def _record_completion_success( - transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response -): +def _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response): span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") try: @@ -647,7 +642,6 @@ def _record_completion_error(transaction, linking_metadata, completion_id, kwarg } llm_metadata = _get_llm_attributes(transaction) error_chat_completion_dict.update(llm_metadata) - transaction.record_custom_event("LlmChatCompletionSummary", error_chat_completion_dict) output_message_list = [] @@ -786,6 +780,7 @@ def _record_events_on_stop_iteration(self, transaction): self._nr_ft.__exit__(None, None, None) try: openai_attrs = getattr(self, "_nr_openai_attrs", {}) + # If there are no openai attrs exit early as there's no data to record. if not openai_attrs: return @@ -793,13 +788,7 @@ def _record_events_on_stop_iteration(self, transaction): completion_id = str(uuid.uuid4()) response_headers = openai_attrs.get("response_headers") or {} _record_completion_success( - transaction, - linking_metadata, - completion_id, - openai_attrs, - self._nr_ft, - response_headers, - None + transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, response_headers, None ) except Exception: _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info())) @@ -824,10 +813,7 @@ def _handle_streaming_completion_error(self, transaction, exc): return linking_metadata = get_trace_linking_metadata() completion_id = str(uuid.uuid4()) - - _record_completion_error( - transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, exc - ) + _record_completion_error(transaction, linking_metadata, completion_id, openai_attrs, self._nr_ft, exc) class AsyncGeneratorProxy(ObjectProxy): From 169484458a3e2e0024067179eb4a0db9d060e8a3 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Tue, 8 Oct 2024 09:50:25 -0700 Subject: [PATCH 13/18] Revert to transaction based storage of attrs. --- newrelic/api/llm_custom_attributes.py | 13 +++----- newrelic/hooks/external_botocore.py | 5 ++- newrelic/hooks/mlmodel_langchain.py | 5 ++- newrelic/hooks/mlmodel_openai.py | 5 ++- .../test_llm_custom_attributes.py | 13 ++++---- tests/mlmodel_openai/test_chat_completion.py | 3 +- .../test_chat_completion_stream_error_v1.py | 31 ++++++++++--------- 7 files changed, 35 insertions(+), 40 deletions(-) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py index b4f0ffc025..5387318b44 100644 --- a/newrelic/api/llm_custom_attributes.py +++ b/newrelic/api/llm_custom_attributes.py @@ -12,13 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import contextvars -import functools import logging from newrelic.api.transaction import current_transaction _logger = logging.getLogger(__name__) -custom_attr_context_var = contextvars.ContextVar("custom_attr_context_var", default={}) class WithLlmCustomAttributes(object): @@ -38,12 +35,10 @@ def __enter__(self): _logger.warning("WithLlmCustomAttributes must be called within the scope of a transaction.") return self - token = custom_attr_context_var.set(self.attr_dict) - self.transaction._custom_attr_context_var = custom_attr_context_var - return token + self.transaction._llm_context_attrs = self.attr_dict + return self def __exit__(self, exc, value, tb): + # Clear out context attributes once we leave the current context if self.transaction: - custom_attr_context_var.set(None) - self.transaction._custom_attr_context_var = custom_attr_context_var - + self.transaction._llm_context_attrs = None diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 3ff7ece19a..f66e3b3056 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -787,10 +787,9 @@ def handle_chat_completion_event(transaction, bedrock_attrs): custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} - llm_context_attrs = getattr(transaction, "_custom_attr_context_var", None) + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) if llm_context_attrs: - context_attrs = llm_context_attrs.get() - llm_metadata_dict.update(context_attrs) + llm_metadata_dict.update(llm_context_attrs) span_id = bedrock_attrs.get("span_id", None) trace_id = bedrock_attrs.get("trace_id", None) diff --git a/newrelic/hooks/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index c6a84c8ace..86a8b71cab 100644 --- a/newrelic/hooks/mlmodel_langchain.py +++ b/newrelic/hooks/mlmodel_langchain.py @@ -709,10 +709,9 @@ def _get_llm_metadata(transaction): # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} - llm_context_attrs = getattr(transaction, "_custom_attr_context_var", None) + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) if llm_context_attrs: - context_attrs = llm_context_attrs.get() - llm_metadata_dict.update(context_attrs) + llm_metadata_dict.update(llm_context_attrs) return llm_metadata_dict diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 41b6a30390..d1098a9f66 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -943,10 +943,9 @@ def _get_llm_attributes(transaction): custom_attrs_dict = transaction._custom_params llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")} - llm_context_attrs = getattr(transaction, "_custom_attr_context_var", None) + llm_context_attrs = getattr(transaction, "_llm_context_attrs", None) if llm_context_attrs: - context_attrs = llm_context_attrs.get() - llm_metadata_dict.update(context_attrs) + llm_metadata_dict.update(llm_context_attrs) return llm_metadata_dict diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py index 520d9c4c20..aa1b0927ff 100644 --- a/tests/agent_features/test_llm_custom_attributes.py +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -1,3 +1,4 @@ + # Copyright 2010 New Relic, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,9 +24,9 @@ def test_llm_custom_attributes(): transaction = current_transaction() with WithLlmCustomAttributes({"test": "attr", "test1": "attr1"}): - assert transaction._custom_attr_context_var.get() == {"llm.test": "attr", "llm.test1": "attr1"} + assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} - assert transaction._custom_attr_context_var.get() is None + assert transaction._llm_context_attrs is None @pytest.mark.parametrize("context_attrs", (None, "not-a-dict")) @@ -35,14 +36,14 @@ def test_llm_custom_attributes_no_attrs(context_attrs): with pytest.raises(TypeError): with WithLlmCustomAttributes(context_attrs): - assert transaction._custom_attr_context_var.get() is None + assert transaction._llm_context_attrs is None @background_task() def test_llm_custom_attributes_prefixed_attrs(): transaction = current_transaction() - with WithLlmCustomAttributes({"llm.test": "attr", "llm.test1": "attr1"}): + with WithLlmCustomAttributes({"llm.test": "attr", "test1": "attr1"}): # Validate API does not prefix attributes that already begin with "llm." - assert transaction._custom_attr_context_var.get() == {"llm.test": "attr", "llm.test1": "attr1"} + assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} - assert transaction._custom_attr_context_var.get() is None + assert transaction._llm_context_attrs is None diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index cdcebc6c0a..f111036660 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -35,10 +35,9 @@ validate_transaction_metrics, ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute, current_transaction +from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes -from newrelic.api.application import application_instance as application _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py index bc4aa41afe..dfdd9b47d6 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py @@ -24,6 +24,7 @@ add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, + events_with_context_attrs, llm_token_count_callback, set_trace_info, ) @@ -118,18 +119,19 @@ rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model(set_trace_info, sync_openai_client): with pytest.raises(TypeError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") - generator = sync_openai_client.chat.completions.create( - messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100, stream=True - ) - for resp in generator: - assert resp + with WithLlmCustomAttributes({"context": "attr"}): + generator = sync_openai_client.chat.completions.create( + messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100, stream=True + ) + for resp in generator: + assert resp @dt_enabled @@ -189,22 +191,23 @@ def test_chat_completion_invalid_request_error_no_model_no_content(set_trace_inf rollup_metrics=[("Llm/completion/OpenAI/create", 1)], background_task=True, ) -@validate_custom_events(expected_events_on_no_model_error) +@validate_custom_events(events_with_context_attrs(expected_events_on_no_model_error)) @validate_custom_event_count(count=3) @background_task() def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_info, async_openai_client): with pytest.raises(TypeError): set_trace_info() add_custom_attribute("llm.conversation_id", "my-awesome-id") + with WithLlmCustomAttributes({"context": "attr"}): - async def consumer(): - generator = await async_openai_client.chat.completions.create( - messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100, stream=True - ) - async for resp in generator: - assert resp + async def consumer(): + generator = await async_openai_client.chat.completions.create( + messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100, stream=True + ) + async for resp in generator: + assert resp - loop.run_until_complete(consumer()) + loop.run_until_complete(consumer()) @dt_enabled From 76e7640d2a5982bdb6445241a8121f5cbc81c26c Mon Sep 17 00:00:00 2001 From: umaannamalai Date: Tue, 8 Oct 2024 16:52:11 +0000 Subject: [PATCH 14/18] [Mega-Linter] Apply linters fixes --- newrelic/agent.py | 8 ++++++-- newrelic/api/llm_custom_attributes.py | 5 ++++- newrelic/hooks/mlmodel_langchain.py | 4 +--- tests/agent_features/test_llm_custom_attributes.py | 3 +-- tests/external_botocore/test_bedrock_chat_completion.py | 2 +- .../test_bedrock_chat_completion_via_langchain.py | 5 ++++- tests/mlmodel_langchain/test_chain.py | 2 +- tests/mlmodel_langchain/test_vectorstore.py | 2 +- tests/mlmodel_openai/test_chat_completion.py | 4 ++-- tests/mlmodel_openai/test_chat_completion_error.py | 2 +- tests/mlmodel_openai/test_chat_completion_error_v1.py | 2 +- tests/mlmodel_openai/test_chat_completion_stream_error.py | 2 +- .../test_chat_completion_stream_error_v1.py | 2 +- tests/mlmodel_openai/test_chat_completion_stream_v1.py | 1 + tests/mlmodel_openai/test_chat_completion_v1.py | 3 +-- 15 files changed, 27 insertions(+), 20 deletions(-) diff --git a/newrelic/agent.py b/newrelic/agent.py index ddc4824293..76f02b8e0d 100644 --- a/newrelic/agent.py +++ b/newrelic/agent.py @@ -139,6 +139,9 @@ def __asgi_application(*args, **kwargs): from newrelic.api.html_insertion import verify_body_exists as __verify_body_exists from newrelic.api.lambda_handler import LambdaHandlerWrapper as __LambdaHandlerWrapper from newrelic.api.lambda_handler import lambda_handler as __lambda_handler +from newrelic.api.llm_custom_attributes import ( + WithLlmCustomAttributes as __WithLlmCustomAttributes, +) from newrelic.api.message_trace import MessageTrace as __MessageTrace from newrelic.api.message_trace import MessageTraceWrapper as __MessageTraceWrapper from newrelic.api.message_trace import message_trace as __message_trace @@ -156,9 +159,10 @@ def __asgi_application(*args, **kwargs): from newrelic.api.ml_model import ( record_llm_feedback_event as __record_llm_feedback_event, ) -from newrelic.api.ml_model import set_llm_token_count_callback as __set_llm_token_count_callback +from newrelic.api.ml_model import ( + set_llm_token_count_callback as __set_llm_token_count_callback, +) from newrelic.api.ml_model import wrap_mlmodel as __wrap_mlmodel -from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes as __WithLlmCustomAttributes from newrelic.api.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper from newrelic.api.profile_trace import profile_trace as __profile_trace from newrelic.api.profile_trace import wrap_profile_trace as __wrap_profile_trace diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py index 5387318b44..8eaaddd8c2 100644 --- a/newrelic/api/llm_custom_attributes.py +++ b/newrelic/api/llm_custom_attributes.py @@ -13,6 +13,7 @@ # limitations under the License. import logging + from newrelic.api.transaction import current_transaction _logger = logging.getLogger(__name__) @@ -22,7 +23,9 @@ class WithLlmCustomAttributes(object): def __init__(self, custom_attr_dict): transaction = current_transaction() if not custom_attr_dict or not isinstance(custom_attr_dict, dict): - raise TypeError("custom_attr_dict must be a non-empty dictionary. Received type: %s" % type(custom_attr_dict)) + raise TypeError( + "custom_attr_dict must be a non-empty dictionary. Received type: %s" % type(custom_attr_dict) + ) # Add "llm." prefix to all keys in attribute dictionary context_attrs = {k if k.startswith("llm.") else f"llm.{k}": v for k, v in custom_attr_dict.items()} diff --git a/newrelic/hooks/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index 86a8b71cab..7855f9486a 100644 --- a/newrelic/hooks/mlmodel_langchain.py +++ b/newrelic/hooks/mlmodel_langchain.py @@ -452,9 +452,7 @@ def _record_tool_success( try: result = str(response) except Exception: - _logger.debug( - f"Failed to convert tool response into a string.\n{traceback.format_exception(*sys.exc_info())}" - ) + _logger.debug(f"Failed to convert tool response into a string.\n{traceback.format_exception(*sys.exc_info())}") if settings.ai_monitoring.record_content.enabled: full_tool_event_dict.update( { diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py index aa1b0927ff..2ed8adfebb 100644 --- a/tests/agent_features/test_llm_custom_attributes.py +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -1,4 +1,3 @@ - # Copyright 2010 New Relic, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -16,8 +15,8 @@ import pytest from newrelic.api.background_task import background_task -from newrelic.api.transaction import current_transaction from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.transaction import current_transaction @background_task() diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py index 496d97e2c1..be0226e55c 100644 --- a/tests/external_botocore/test_bedrock_chat_completion.py +++ b/tests/external_botocore/test_bedrock_chat_completion.py @@ -59,8 +59,8 @@ ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.transaction import add_custom_attribute from newrelic.common.object_names import callable_name from newrelic.hooks.external_botocore import MODEL_EXTRACTORS diff --git a/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py b/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py index afe88f4bf8..3bd18764fa 100644 --- a/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py +++ b/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py @@ -19,7 +19,10 @@ ) from conftest import BOTOCORE_VERSION # pylint: disable=E0611 from testing_support.fixtures import reset_core_stats_engine, validate_attributes -from testing_support.ml_testing_utils import set_trace_info, events_with_context_attrs # noqa: F401 +from testing_support.ml_testing_utils import ( # noqa: F401 + events_with_context_attrs, + set_trace_info, +) from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( diff --git a/tests/mlmodel_langchain/test_chain.py b/tests/mlmodel_langchain/test_chain.py index 147f690647..9a372f78dc 100644 --- a/tests/mlmodel_langchain/test_chain.py +++ b/tests/mlmodel_langchain/test_chain.py @@ -47,8 +47,8 @@ ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.transaction import add_custom_attribute from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_langchain/test_vectorstore.py b/tests/mlmodel_langchain/test_vectorstore.py index ec1d25b8e8..4a9188e749 100644 --- a/tests/mlmodel_langchain/test_vectorstore.py +++ b/tests/mlmodel_langchain/test_vectorstore.py @@ -36,8 +36,8 @@ ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.transaction import add_custom_attribute from newrelic.common.object_names import callable_name diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index f111036660..e7985013dd 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -34,10 +34,10 @@ from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, ) + from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes - +from newrelic.api.transaction import add_custom_attribute _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, diff --git a/tests/mlmodel_openai/test_chat_completion_error.py b/tests/mlmodel_openai/test_chat_completion_error.py index 9883c9b75b..fab701d1c7 100644 --- a/tests/mlmodel_openai/test_chat_completion_error.py +++ b/tests/mlmodel_openai/test_chat_completion_error.py @@ -39,8 +39,8 @@ ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.transaction import add_custom_attribute from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_error_v1.py b/tests/mlmodel_openai/test_chat_completion_error_v1.py index ae1bb1bfef..32147691d1 100644 --- a/tests/mlmodel_openai/test_chat_completion_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_error_v1.py @@ -38,8 +38,8 @@ ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.transaction import add_custom_attribute from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error.py b/tests/mlmodel_openai/test_chat_completion_stream_error.py index 2961728edd..56135c5d05 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error.py @@ -39,8 +39,8 @@ ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.transaction import add_custom_attribute from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py index dfdd9b47d6..0c4978c97a 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py @@ -39,8 +39,8 @@ ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.api.transaction import add_custom_attribute from newrelic.common.object_names import callable_name _test_openai_chat_completion_messages = ( diff --git a/tests/mlmodel_openai/test_chat_completion_stream_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_v1.py index 40811794b3..5e60fd8880 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_v1.py @@ -503,6 +503,7 @@ async def consumer(): ) async for resp in generator: assert resp + with WithLlmCustomAttributes({"context": "attr"}): loop.run_until_complete(consumer()) diff --git a/tests/mlmodel_openai/test_chat_completion_v1.py b/tests/mlmodel_openai/test_chat_completion_v1.py index 420574f06b..dededb840d 100644 --- a/tests/mlmodel_openai/test_chat_completion_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_v1.py @@ -36,9 +36,8 @@ ) from newrelic.api.background_task import background_task -from newrelic.api.transaction import add_custom_attribute from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes - +from newrelic.api.transaction import add_custom_attribute _test_openai_chat_completion_messages = ( {"role": "system", "content": "You are a scientist."}, From beeb67f0f575dd5ea3db12856349c382513178fa Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Tue, 8 Oct 2024 15:46:03 -0700 Subject: [PATCH 15/18] Delete context attrs instead of nullifying. --- newrelic/api/llm_custom_attributes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/newrelic/api/llm_custom_attributes.py b/newrelic/api/llm_custom_attributes.py index 5387318b44..4bdb9a265b 100644 --- a/newrelic/api/llm_custom_attributes.py +++ b/newrelic/api/llm_custom_attributes.py @@ -41,4 +41,4 @@ def __enter__(self): def __exit__(self, exc, value, tb): # Clear out context attributes once we leave the current context if self.transaction: - self.transaction._llm_context_attrs = None + del self.transaction._llm_context_attrs From 46d552112a9b53d0a30c16251ab7fa0cdb998178 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Tue, 8 Oct 2024 17:32:43 -0700 Subject: [PATCH 16/18] Update tests. --- tests/agent_features/test_llm_custom_attributes.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py index 2ed8adfebb..6abba23f40 100644 --- a/tests/agent_features/test_llm_custom_attributes.py +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -25,7 +25,6 @@ def test_llm_custom_attributes(): with WithLlmCustomAttributes({"test": "attr", "test1": "attr1"}): assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} - assert transaction._llm_context_attrs is None @pytest.mark.parametrize("context_attrs", (None, "not-a-dict")) @@ -44,5 +43,3 @@ def test_llm_custom_attributes_prefixed_attrs(): with WithLlmCustomAttributes({"llm.test": "attr", "test1": "attr1"}): # Validate API does not prefix attributes that already begin with "llm." assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} - - assert transaction._llm_context_attrs is None From 7160696b42adb108336ccf6ef58c37e323c26828 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Tue, 8 Oct 2024 20:55:10 -0700 Subject: [PATCH 17/18] Remove newline. --- tests/agent_features/test_llm_custom_attributes.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py index 6abba23f40..04bd564dbb 100644 --- a/tests/agent_features/test_llm_custom_attributes.py +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -26,7 +26,6 @@ def test_llm_custom_attributes(): assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} - @pytest.mark.parametrize("context_attrs", (None, "not-a-dict")) @background_task() def test_llm_custom_attributes_no_attrs(context_attrs): From ef40ca6ba3b0cc620f674e94d4f9f5d227895d34 Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Wed, 9 Oct 2024 14:07:50 -0700 Subject: [PATCH 18/18] Assert attribute deletion was successful. --- tests/agent_features/test_llm_custom_attributes.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/agent_features/test_llm_custom_attributes.py b/tests/agent_features/test_llm_custom_attributes.py index 04bd564dbb..1f02c231c0 100644 --- a/tests/agent_features/test_llm_custom_attributes.py +++ b/tests/agent_features/test_llm_custom_attributes.py @@ -25,6 +25,8 @@ def test_llm_custom_attributes(): with WithLlmCustomAttributes({"test": "attr", "test1": "attr1"}): assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + assert not hasattr(transaction, "_llm_context_attrs") + @pytest.mark.parametrize("context_attrs", (None, "not-a-dict")) @background_task() @@ -33,7 +35,9 @@ def test_llm_custom_attributes_no_attrs(context_attrs): with pytest.raises(TypeError): with WithLlmCustomAttributes(context_attrs): - assert transaction._llm_context_attrs is None + pass + + assert not hasattr(transaction, "_llm_context_attrs") @background_task() @@ -42,3 +46,5 @@ def test_llm_custom_attributes_prefixed_attrs(): with WithLlmCustomAttributes({"llm.test": "attr", "test1": "attr1"}): # Validate API does not prefix attributes that already begin with "llm." assert transaction._llm_context_attrs == {"llm.test": "attr", "llm.test1": "attr1"} + + assert not hasattr(transaction, "_llm_context_attrs")