Skip to content

Commit

Permalink
Add attr val limit exceptions for LLMs
Browse files Browse the repository at this point in the history
  • Loading branch information
lrafeei committed Mar 8, 2024
1 parent 6899e3d commit 74ed3f5
Show file tree
Hide file tree
Showing 2 changed files with 112 additions and 2 deletions.
10 changes: 9 additions & 1 deletion newrelic/core/custom_event.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand Down
104 changes: 103 additions & 1 deletion tests/agent_features/test_custom_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,19 @@

import time

import pytest
from testing_support.fixtures import (
function_not_called,
override_application_settings,
reset_core_stats_engine,
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()
Expand Down Expand Up @@ -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()

0 comments on commit 74ed3f5

Please sign in to comment.