From 74ed3f5658c1f1068ceeef617177bbe5d8f4a1f0 Mon Sep 17 00:00:00 2001
From: Lalleh Rafeei <lrafeei@newrelic.com>
Date: Fri, 8 Mar 2024 11:13:06 -0800
Subject: [PATCH] Add attr val limit exceptions for LLMs

---
 newrelic/core/custom_event.py              |  10 +-
 tests/agent_features/test_custom_events.py | 104 ++++++++++++++++++++-
 2 files changed, 112 insertions(+), 2 deletions(-)

diff --git a/newrelic/core/custom_event.py b/newrelic/core/custom_event.py
index 15becf437c..4d541fc6ca 100644
--- a/newrelic/core/custom_event.py
+++ b/newrelic/core/custom_event.py
@@ -30,6 +30,10 @@
 _logger = logging.getLogger(__name__)
 
 EVENT_TYPE_VALID_CHARS_REGEX = re.compile(r"^[a-zA-Z0-9:_ ]+$")
+NO_LIMIT_LLM_EVENT_TYPE = {
+    "LlmChatCompletionMessage": "content",
+    "LlmEmbedding": "input",
+}
 
 
 class NameInvalidCharactersException(Exception):
@@ -115,7 +119,11 @@ def create_custom_event(event_type, params, settings=None, is_ml_event=False):
                 max_length = MAX_ML_ATTRIBUTE_LENGTH
                 max_num_attrs = MAX_NUM_ML_USER_ATTRIBUTES
             else:
-                max_length = settings.custom_insights_events.max_attribute_value
+                max_length = (
+                    settings.custom_insights_events.max_attribute_value
+                    if not (name in NO_LIMIT_LLM_EVENT_TYPE.keys() and NO_LIMIT_LLM_EVENT_TYPE[name] == k)
+                    else None
+                )
                 max_num_attrs = MAX_NUM_USER_ATTRIBUTES
             key, value = process_user_attribute(k, v, max_length=max_length)
             if key:
diff --git a/tests/agent_features/test_custom_events.py b/tests/agent_features/test_custom_events.py
index 0fb9c80bc2..b5b86bcf1e 100644
--- a/tests/agent_features/test_custom_events.py
+++ b/tests/agent_features/test_custom_events.py
@@ -14,6 +14,7 @@
 
 import time
 
+import pytest
 from testing_support.fixtures import (
     function_not_called,
     override_application_settings,
@@ -21,10 +22,11 @@
     validate_custom_event_count,
     validate_custom_event_in_application_stats_engine,
 )
+from testing_support.validators.validate_custom_events import validate_custom_events
 
 from newrelic.api.application import application_instance as application
 from newrelic.api.background_task import background_task
-from newrelic.api.transaction import record_custom_event
+from newrelic.api.transaction import current_transaction, record_custom_event
 from newrelic.core.custom_event import process_event_type
 
 # Test process_event_type()
@@ -212,3 +214,103 @@ def test_transaction_create_custom_event_not_called():
 def test_application_create_custom_event_not_called():
     app = application()
     record_custom_event("FooEvent", _user_params, application=app)
+
+
+# Test completness of LLM content/input despite attribute limits being set
+
+
+LlmChatCompletionMessage = {
+    "content": "A" * 9001,  # Should print out to completion
+    "input": "B" * 5000,  # Should cap out at 255 or 4095
+    "foo": "b" + "a" * 6000 + "r",
+}
+LlmEmbedding = {
+    "content": "A" * 9001,  # Should cap out at 255 or 4095
+    "input": "B" * 5000,  # Should print out to completion
+    "foo": "b" + "a" * 6000 + "r",
+}
+SomeOtherDict = {
+    "content": "A" * 9001,
+    "input": "B" * 5000,
+    "foo": "b" + "a" * 6000 + "r",
+}
+
+
+recorded_events_256 = [
+    (
+        {"type": "LlmChatCompletionMessage"},
+        {
+            "content": "A" * 9001,  # Should print out to completion
+            "input": "B" * 255,  # Should cap out at 255 or 4095
+            "foo": "b" + "a" * 254,
+        },
+    ),
+    (
+        {"type": "LlmEmbedding"},
+        {
+            "content": "A" * 255,  # Should cap out at 255 or 4095
+            "input": "B" * 5000,  # Should print out to completion
+            "foo": "b" + "a" * 254,
+        },
+    ),
+    (
+        {"type": "SomeOtherDict"},
+        {
+            "content": "A" * 255,
+            "input": "B" * 255,
+            "foo": "b" + "a" * 254,
+        },
+    ),
+]
+
+recorded_events_4096 = [
+    (
+        {"type": "LlmChatCompletionMessage"},
+        {
+            "content": "A" * 9001,  # Should print out to completion
+            "input": "B" * 4095,  # Should cap out at 255 or 4095
+            "foo": "b" + "a" * 4094,
+        },
+    ),
+    (
+        {"type": "LlmEmbedding"},
+        {
+            "content": "A" * 4095,  # Should cap out at 255 or 4095
+            "input": "B" * 5000,  # Should print out to completion
+            "foo": "b" + "a" * 4094,
+        },
+    ),
+    (
+        {"type": "SomeOtherDict"},
+        {
+            "content": "A" * 4095,
+            "input": "B" * 4095,
+            "foo": "b" + "a" * 4094,
+        },
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "max_val,expected_events",
+    (
+        (255, recorded_events_256),
+        (4095, recorded_events_4096),
+    ),
+)
+def test_create_custom_event_no_limit(max_val, expected_events):
+    @reset_core_stats_engine()
+    @override_application_settings({"custom_insights_events.max_attribute_value": max_val})
+    @validate_custom_event_count(3)
+    @validate_custom_events(expected_events)
+    @background_task()
+    def _test():
+        transaction = current_transaction()
+        if not transaction:
+            return
+
+        transaction.record_custom_event("LlmChatCompletionMessage", LlmChatCompletionMessage)
+        transaction.record_custom_event("LlmEmbedding", LlmEmbedding)
+        transaction.record_custom_event("SomeOtherDict", SomeOtherDict)
+
+    _test()