Merge pull request #1214 from newrelic/llm-custom-attrs-api

LLM Custom Attributes Context Manager API

Author: umaannamalai

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,7 +159,9 @@ 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.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper
 from newrelic.api.profile_trace import profile_trace as __profile_trace
@@ -251,6 +256,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"
 )

newrelic/api/llm_custom_attributes.py

@@ -0,0 +1,47 @@
+# 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 logging
+
+from newrelic.api.transaction import current_transaction
+
+_logger = logging.getLogger(__name__)
+
+
+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)
+            )
+
+        # 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()}
+
+        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):
+        # Clear out context attributes once we leave the current context
+        if self.transaction:
+            del self.transaction._llm_context_attrs

newrelic/hooks/external_botocore.py

@@ -787,6 +787,10 @@ 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, "_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)

newrelic/hooks/mlmodel_langchain.py

@@ -697,7 +697,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 [])
@@ -708,6 +708,10 @@ 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
 
 

newrelic/hooks/mlmodel_openai.py

@@ -927,8 +927,13 @@ def is_stream(wrapped, args, kwargs):
 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.")}
-    return llm_metadata
+    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
 
 
 def instrument_openai_api_resources_embedding(module):

tests/agent_features/test_llm_custom_attributes.py

@@ -0,0 +1,50 @@
+# 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 newrelic.api.background_task import background_task
+from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes
+from newrelic.api.transaction import current_transaction
+
+
+@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 not hasattr(transaction, "_llm_context_attrs")
+
+
+@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):
+            pass
+
+    assert not hasattr(transaction, "_llm_context_attrs")
+
+
+@background_task()
+def test_llm_custom_attributes_prefixed_attrs():
+    transaction = current_transaction()
+    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")

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,
 )
@@ -58,6 +59,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
 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()
 

tests/external_botocore/test_bedrock_chat_completion_via_langchain.py

@@ -19,14 +19,18 @@
 )
 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 (  # 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 (
     validate_transaction_metrics,
 )
 
 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 +109,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 +128,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()