Revert to transaction based storage of attrs.

Author: umaannamalai

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

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)

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
 

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
 

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

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."},

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