mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
1359 lines
52 KiB
Python
1359 lines
52 KiB
Python
import os
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
# newrelic is a proxy-runtime dependency (pyproject.toml) and is not installed
|
|
# in the CI Python environment. Mock it in sys.modules before importing the
|
|
# integration so that deferred `import newrelic.agent` calls inside NewRelicLogger
|
|
# methods resolve to these mocks rather than failing with ModuleNotFoundError.
|
|
_mock_newrelic = MagicMock()
|
|
_mock_newrelic_agent = MagicMock()
|
|
# Explicitly link so _mock_newrelic.agent IS _mock_newrelic_agent. Without this,
|
|
# the first getattr(_mock_newrelic, 'agent') auto-creates a different child mock,
|
|
# causing patch("newrelic.agent.xxx") to patch the wrong object.
|
|
_mock_newrelic.agent = _mock_newrelic_agent
|
|
sys.modules["newrelic"] = _mock_newrelic
|
|
sys.modules["newrelic.agent"] = _mock_newrelic_agent
|
|
|
|
import litellm
|
|
import litellm.integrations.newrelic.newrelic as nr_module
|
|
from litellm.integrations.newrelic.newrelic import NewRelicLogger
|
|
|
|
|
|
def _rendered_log_message(call):
|
|
message = str(call.args[0])
|
|
values = call.args[1:]
|
|
return message % values if values else message
|
|
|
|
|
|
# The module may have been imported before sys.modules was patched (e.g. via
|
|
# litellm's own startup imports), leaving _newrelic_agent=None. Point it at
|
|
# the mock agent so all tests see a non-None agent.
|
|
nr_module._newrelic_agent = _mock_newrelic_agent
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
NR_ENV = {
|
|
"NEW_RELIC_LICENSE_KEY": "test-license-key",
|
|
"NEW_RELIC_APP_NAME": "test-app",
|
|
}
|
|
|
|
|
|
def make_logger(**kwargs) -> NewRelicLogger:
|
|
"""Instantiate NewRelicLogger with NR agent calls mocked out."""
|
|
with patch.dict(os.environ, NR_ENV):
|
|
return NewRelicLogger(**kwargs)
|
|
|
|
|
|
def make_kwargs(
|
|
model="gpt-4",
|
|
provider="openai",
|
|
messages=None,
|
|
optional_params=None,
|
|
traceparent=None,
|
|
) -> dict:
|
|
"""Build a minimal kwargs dict representative of a litellm callback invocation."""
|
|
headers = {}
|
|
if traceparent:
|
|
headers["traceparent"] = traceparent
|
|
|
|
return {
|
|
"model": model,
|
|
"messages": messages or [{"role": "user", "content": "Hello"}],
|
|
"optional_params": optional_params or {},
|
|
"litellm_params": {
|
|
"custom_llm_provider": provider,
|
|
"metadata": {"headers": headers},
|
|
},
|
|
"start_time": 1_000_000.0,
|
|
"end_time": 1_000_001.5,
|
|
"llm_api_duration_ms": 1500.0,
|
|
}
|
|
|
|
|
|
def make_response(
|
|
model="gpt-4",
|
|
response_id="chatcmpl-abc123",
|
|
content="Hello there!",
|
|
finish_reason="stop",
|
|
prompt_tokens=10,
|
|
completion_tokens=20,
|
|
):
|
|
"""Build a minimal ModelResponse-like dict."""
|
|
return {
|
|
"id": response_id,
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"message": {"role": "assistant", "content": content},
|
|
"finish_reason": finish_reason,
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": prompt_tokens + completion_tokens,
|
|
},
|
|
}
|
|
|
|
|
|
def make_slo(**overrides):
|
|
"""Build a StandardLoggingPayload-like dict with sentinel values distinct from
|
|
make_kwargs/make_response defaults, so tests can prove the SLO branch won."""
|
|
base = {
|
|
"trace_id": "slo-trace-abc",
|
|
"custom_llm_provider": "slo-provider",
|
|
"model": "slo-model",
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 200,
|
|
"total_tokens": 300,
|
|
"response_time": 1.5, # seconds; converted to ms by _get_duration
|
|
"model_parameters": {"temperature": 0.7, "max_tokens": 500},
|
|
"startTime": 2_000_000.0,
|
|
"endTime": 2_000_001.5,
|
|
"messages": [{"role": "user", "content": "from-slo"}],
|
|
}
|
|
base.update(overrides)
|
|
return base
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Init / configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestNewRelicLoggerInit:
|
|
def test_disabled_when_license_key_missing(self):
|
|
with patch("newrelic.agent.register_application"):
|
|
with patch.dict(os.environ, {"NEW_RELIC_APP_NAME": "app"}, clear=True):
|
|
logger = NewRelicLogger()
|
|
assert logger.enabled is False
|
|
|
|
def test_disabled_when_app_name_missing(self):
|
|
with patch("newrelic.agent.register_application"):
|
|
with patch.dict(os.environ, {"NEW_RELIC_LICENSE_KEY": "key"}, clear=True):
|
|
logger = NewRelicLogger()
|
|
assert logger.enabled is False
|
|
|
|
def test_enabled_with_valid_env_vars(self):
|
|
with patch("newrelic.agent.register_application"):
|
|
with patch.dict(os.environ, NR_ENV):
|
|
logger = NewRelicLogger()
|
|
assert logger.enabled is True
|
|
|
|
def test_disabled_on_import_error(self):
|
|
with patch.object(
|
|
_mock_newrelic_agent, "register_application", side_effect=ImportError
|
|
):
|
|
with patch.dict(os.environ, NR_ENV):
|
|
logger = NewRelicLogger()
|
|
assert logger.enabled is False
|
|
|
|
def test_disabled_on_agent_startup_error(self):
|
|
with patch.object(
|
|
_mock_newrelic_agent,
|
|
"register_application",
|
|
side_effect=RuntimeError("agent startup failed"),
|
|
):
|
|
with patch.dict(os.environ, NR_ENV):
|
|
logger = NewRelicLogger()
|
|
assert logger.enabled is False
|
|
|
|
def test_disabled_when_agent_package_missing(self):
|
|
with patch.object(nr_module, "_newrelic_agent", None):
|
|
with patch.dict(os.environ, NR_ENV):
|
|
logger = NewRelicLogger()
|
|
assert logger.enabled is False
|
|
|
|
def test_record_content_default_true(self):
|
|
logger = make_logger()
|
|
assert logger.record_content is True
|
|
|
|
def test_record_content_disabled_by_param(self):
|
|
logger = make_logger(turn_off_message_logging=True)
|
|
assert logger.record_content is False
|
|
|
|
def test_record_content_disabled_by_env_var(self):
|
|
with patch("newrelic.agent.register_application"):
|
|
with patch.dict(
|
|
os.environ,
|
|
{**NR_ENV, "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": "false"},
|
|
):
|
|
logger = NewRelicLogger()
|
|
assert logger.record_content is False
|
|
|
|
def test_record_content_requires_both_enabled(self):
|
|
"""param says record, but env var says no — result is False."""
|
|
with patch("newrelic.agent.register_application"):
|
|
with patch.dict(
|
|
os.environ,
|
|
{**NR_ENV, "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": "false"},
|
|
):
|
|
logger = NewRelicLogger(turn_off_message_logging=False)
|
|
assert logger.record_content is False
|
|
|
|
def test_constructor_kwargs_take_priority_over_global_params(self):
|
|
"""Constructor turn_off_message_logging=True must not be overwritten by
|
|
litellm.newrelic_params which defaults turn_off_message_logging to False."""
|
|
from litellm.types.integrations.newrelic import NewRelicInitParams
|
|
|
|
with patch("newrelic.agent.register_application"):
|
|
with patch.dict(os.environ, NR_ENV):
|
|
with patch(
|
|
"litellm.newrelic_params",
|
|
NewRelicInitParams(turn_off_message_logging=False),
|
|
):
|
|
logger = NewRelicLogger(turn_off_message_logging=True)
|
|
assert logger.record_content is False
|
|
|
|
def test_newrelic_params_plain_dict_branch(self):
|
|
"""litellm.newrelic_params can be a plain dict; it should be validated
|
|
through NewRelicInitParams and its values applied to the logger."""
|
|
with patch("newrelic.agent.register_application"):
|
|
with patch.dict(os.environ, NR_ENV):
|
|
with patch(
|
|
"litellm.newrelic_params",
|
|
{"turn_off_message_logging": True},
|
|
):
|
|
logger = NewRelicLogger()
|
|
assert logger.turn_off_message_logging is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _parse_bool_env
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestParseBoolEnv:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
@pytest.mark.parametrize("raw", ["true", "TRUE", "True", "1", "yes", "on", "ON"])
|
|
def test_truthy_values(self, raw):
|
|
with patch.dict(os.environ, {"MY_VAR": raw}):
|
|
assert self.logger._parse_bool_env("MY_VAR") is True
|
|
|
|
@pytest.mark.parametrize("raw", ["false", "FALSE", "0", "no", "off", "Off"])
|
|
def test_falsy_values(self, raw):
|
|
with patch.dict(os.environ, {"MY_VAR": raw}):
|
|
assert self.logger._parse_bool_env("MY_VAR") is False
|
|
|
|
@pytest.mark.parametrize("raw", [" true ", " 1\t", "\nyes"])
|
|
def test_whitespace_tolerance_truthy(self, raw):
|
|
with patch.dict(os.environ, {"MY_VAR": raw}):
|
|
assert self.logger._parse_bool_env("MY_VAR") is True
|
|
|
|
@pytest.mark.parametrize("raw", [" false ", " 0\t", "\nno"])
|
|
def test_whitespace_tolerance_falsy(self, raw):
|
|
with patch.dict(os.environ, {"MY_VAR": raw}):
|
|
assert self.logger._parse_bool_env("MY_VAR") is False
|
|
|
|
def test_missing_uses_default(self):
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
assert self.logger._parse_bool_env("MY_VAR", default=True) is True
|
|
assert self.logger._parse_bool_env("MY_VAR", default=False) is False
|
|
|
|
def test_empty_string_uses_default(self):
|
|
with patch.dict(os.environ, {"MY_VAR": ""}):
|
|
assert self.logger._parse_bool_env("MY_VAR", default=True) is True
|
|
assert self.logger._parse_bool_env("MY_VAR", default=False) is False
|
|
|
|
@pytest.mark.parametrize("raw", ["maybe", "2", "enabled", "tru"])
|
|
def test_unrecognised_value_falls_back_to_default_with_warning(self, raw):
|
|
with (
|
|
patch.dict(os.environ, {"MY_VAR": raw}),
|
|
patch.object(nr_module.verbose_logger, "warning") as mock_warn,
|
|
):
|
|
assert self.logger._parse_bool_env("MY_VAR", default=True) is True
|
|
assert self.logger._parse_bool_env("MY_VAR", default=False) is False
|
|
assert mock_warn.call_count == 2
|
|
# Warning should mention the variable name and the raw value
|
|
for call in mock_warn.call_args_list:
|
|
rendered = _rendered_log_message(call)
|
|
assert "MY_VAR" in rendered
|
|
assert repr(raw) in rendered
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _get_trace_context
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGetTraceContext:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_extracts_trace_id_from_traceparent(self):
|
|
kwargs = make_kwargs(
|
|
traceparent="00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"
|
|
)
|
|
trace_id = self.logger._get_trace_context(kwargs)
|
|
assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736"
|
|
|
|
def test_generates_uuid_when_no_headers(self):
|
|
kwargs = make_kwargs()
|
|
trace_id = self.logger._get_trace_context(kwargs)
|
|
assert trace_id is not None
|
|
assert (
|
|
len(trace_id) == 32
|
|
) # 32-char lowercase hex, matches W3C traceparent format
|
|
|
|
def test_generates_uuid_when_traceparent_malformed(self):
|
|
kwargs = make_kwargs(traceparent="not-valid")
|
|
trace_id = self.logger._get_trace_context(kwargs)
|
|
# Falls back to a 32-char lowercase hex, matching W3C traceparent format
|
|
assert trace_id is not None
|
|
assert len(trace_id) == 32
|
|
|
|
def test_extracts_trace_id_from_mixed_case_traceparent_header(self):
|
|
# Callers passing headers directly may not normalise case; per W3C spec
|
|
# header names are case-insensitive, so "Traceparent" must work too.
|
|
kwargs = make_kwargs()
|
|
kwargs["litellm_params"]["metadata"]["headers"] = {
|
|
"Traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"
|
|
}
|
|
trace_id = self.logger._get_trace_context(kwargs)
|
|
assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736"
|
|
|
|
def test_parse_failure_falls_through_to_synthetic_uuid(self):
|
|
"""When parsing upstream sources raises, emit a synthetic UUID rather
|
|
than dropping the event. NR schema requires every AIM event carry a
|
|
trace_id; this method's contract is to always return a valid string.
|
|
"""
|
|
# Non-dict headers value forces .items() to raise inside the try
|
|
kwargs = {"litellm_params": {"metadata": {"headers": "not-a-dict"}}}
|
|
trace_id = self.logger._get_trace_context(kwargs)
|
|
assert trace_id is not None
|
|
assert len(trace_id) == 32 # 32-char lowercase hex fallback
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _extract_message_content edge cases
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtractMessageContent:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_plain_text(self):
|
|
assert self.logger._extract_message_content({"content": "hello"}) == "hello"
|
|
|
|
def test_none_content_returns_empty_string(self):
|
|
assert self.logger._extract_message_content({"content": None}) == ""
|
|
|
|
def test_missing_content_returns_empty_string(self):
|
|
assert self.logger._extract_message_content({}) == ""
|
|
|
|
def test_tool_calls_serialized_as_json(self):
|
|
msg = {
|
|
"content": None,
|
|
"tool_calls": [{"id": "call_1", "function": {"name": "get_weather"}}],
|
|
}
|
|
result = self.logger._extract_message_content(msg)
|
|
assert "get_weather" in result
|
|
assert "call_1" in result
|
|
|
|
def test_multimodal_list_serialized_as_json(self):
|
|
msg = {
|
|
"content": [
|
|
{"type": "text", "text": "describe this"},
|
|
{"type": "image_url"},
|
|
]
|
|
}
|
|
result = self.logger._extract_message_content(msg)
|
|
assert "describe this" in result
|
|
assert "image_url" in result
|
|
|
|
def test_non_string_content_coerced_to_str(self):
|
|
"""Numeric/bool content passes the None and list guards; final branch coerces to str."""
|
|
assert self.logger._extract_message_content({"content": 123}) == "123"
|
|
assert self.logger._extract_message_content({"content": True}) == "True"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _extract_all_messages — record_content=False path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtractAllMessagesContentDisabled:
|
|
def test_no_content_key_when_recording_disabled(self):
|
|
logger = make_logger(turn_off_message_logging=True)
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "secret"}])
|
|
response = make_response(content="also secret")
|
|
|
|
messages = logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
|
|
for msg in messages:
|
|
assert "content" not in msg
|
|
|
|
|
|
class TestExtractAllMessagesRespectsLitellmRedaction:
|
|
"""Regression tests for the async-streaming redaction bypass.
|
|
|
|
NR-specific switches alone are insufficient: when
|
|
``litellm.turn_off_message_logging=True`` (or the per-request equivalents),
|
|
async streaming callbacks receive an unredacted
|
|
``async_complete_streaming_response``. Without consulting LiteLLM's
|
|
redaction decision the integration would still write generated content
|
|
into NR events.
|
|
"""
|
|
|
|
def _assert_no_content(self, logger, kwargs):
|
|
response = make_response(content="streamed assistant text")
|
|
|
|
messages = logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
|
|
# All extracted messages must carry no content payload
|
|
for msg in messages:
|
|
assert (
|
|
"content" not in msg
|
|
), f"content leaked despite redaction signal: {msg}"
|
|
# And there must actually be at least one user + one assistant entry,
|
|
# otherwise the test would pass vacuously.
|
|
assert any(not m.get("is_response") for m in messages)
|
|
assert any(m.get("is_response") for m in messages)
|
|
|
|
def test_global_turn_off_message_logging_blocks_content(self, monkeypatch):
|
|
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
|
|
logger = make_logger()
|
|
assert logger.record_content is True
|
|
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}])
|
|
self._assert_no_content(logger, kwargs)
|
|
|
|
def test_dynamic_param_turn_off_message_logging_blocks_content(self):
|
|
logger = make_logger()
|
|
assert logger.record_content is True
|
|
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}])
|
|
kwargs["standard_callback_dynamic_params"] = {
|
|
"turn_off_message_logging": True,
|
|
}
|
|
self._assert_no_content(logger, kwargs)
|
|
|
|
def test_enable_redaction_header_blocks_content(self):
|
|
logger = make_logger()
|
|
assert logger.record_content is True
|
|
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}])
|
|
kwargs["litellm_params"]["metadata"]["headers"] = {
|
|
"x-litellm-enable-message-redaction": True,
|
|
}
|
|
self._assert_no_content(logger, kwargs)
|
|
|
|
def test_dynamic_param_explicit_false_overrides_global_redaction(self, monkeypatch):
|
|
"""The dynamic param has higher priority than the global flag (see
|
|
should_redact_message_logging). When a caller explicitly opts back into
|
|
message logging per-request, NR must record content again."""
|
|
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
|
|
logger = make_logger()
|
|
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "ok to log"}])
|
|
kwargs["standard_callback_dynamic_params"] = {
|
|
"turn_off_message_logging": False,
|
|
}
|
|
response = make_response(content="response text")
|
|
|
|
messages = logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
|
|
request_msg = next(m for m in messages if not m.get("is_response"))
|
|
response_msg = next(m for m in messages if m.get("is_response"))
|
|
assert request_msg["content"] == "ok to log"
|
|
assert response_msg["content"] == "response text"
|
|
|
|
|
|
class TestExtractAllMessagesTimestamps:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_input_messages_get_start_time_timestamp(self):
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
|
|
# make_kwargs sets start_time=1_000_000.0 and end_time=1_000_001.5
|
|
response = make_response()
|
|
|
|
messages = self.logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
|
|
input_msg = next(m for m in messages if not m.get("is_response"))
|
|
assert input_msg["timestamp"] == int(1_000_000.0 * 1000.0)
|
|
|
|
def test_output_messages_get_end_time_timestamp(self):
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
|
|
response = make_response()
|
|
|
|
messages = self.logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
|
|
output_msg = next(m for m in messages if m.get("is_response"))
|
|
assert output_msg["timestamp"] == int(1_000_001.5 * 1000.0)
|
|
|
|
def test_timestamp_forwarded_to_event_data(self):
|
|
logger = make_logger()
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
|
|
kwargs = make_kwargs(
|
|
traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01",
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
)
|
|
response = make_response()
|
|
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
logger._process_success(kwargs, response, start_time=1.0, end_time=2.5)
|
|
|
|
calls = mock_app.record_custom_event.call_args_list
|
|
message_events = [
|
|
c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage"
|
|
]
|
|
for event in message_events:
|
|
assert "timestamp" in event
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Streaming response handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def make_streaming_response(
|
|
model="gpt-4",
|
|
response_id="chatcmpl-stream123",
|
|
content="Hello from streaming!",
|
|
finish_reason="stop",
|
|
prompt_tokens=8,
|
|
completion_tokens=15,
|
|
):
|
|
"""Build a streaming-assembled response dict using 'delta' instead of 'message'."""
|
|
return {
|
|
"id": response_id,
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"delta": {"role": "assistant", "content": content},
|
|
"finish_reason": finish_reason,
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": prompt_tokens + completion_tokens,
|
|
},
|
|
}
|
|
|
|
|
|
class TestStreamingResponse:
|
|
"""Verify graceful handling of streaming-assembled responses.
|
|
|
|
When LiteLLM assembles a streaming response, some providers produce a
|
|
final choice dict with a 'delta' key instead of 'message'. The integration
|
|
must extract content from either key without raising.
|
|
"""
|
|
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_extracts_content_from_delta_key(self):
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
|
|
response = make_streaming_response(content="Streamed reply")
|
|
|
|
messages = self.logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
|
|
response_msgs = [m for m in messages if m.get("is_response")]
|
|
assert len(response_msgs) == 1
|
|
assert response_msgs[0]["content"] == "Streamed reply"
|
|
assert response_msgs[0]["role"] == "assistant"
|
|
|
|
def test_streaming_response_records_summary_and_message_events(self):
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
|
|
kwargs = make_kwargs(
|
|
traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01",
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
)
|
|
response = make_streaming_response(
|
|
response_id="chatcmpl-stream123",
|
|
content="Streamed reply",
|
|
finish_reason="stop",
|
|
prompt_tokens=8,
|
|
completion_tokens=15,
|
|
)
|
|
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
self.logger._process_success(kwargs, response, start_time=1.0, end_time=2.0)
|
|
|
|
calls = mock_app.record_custom_event.call_args_list
|
|
event_types = [c[0][0] for c in calls]
|
|
assert "LlmChatCompletionSummary" in event_types
|
|
assert "LlmChatCompletionMessage" in event_types
|
|
|
|
message_events = [
|
|
c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage"
|
|
]
|
|
response_msg = next((e for e in message_events if e.get("is_response")), None)
|
|
assert response_msg is not None
|
|
assert response_msg["content"] == "Streamed reply"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_log_success_event_streaming(self):
|
|
"""async_log_success_event is the primary entry point for streaming calls."""
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
|
|
response = make_streaming_response()
|
|
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
await self.logger.async_log_success_event(
|
|
kwargs, response, start_time=1.0, end_time=2.0
|
|
)
|
|
|
|
calls = mock_app.record_custom_event.call_args_list
|
|
event_types = [c[0][0] for c in calls]
|
|
assert "LlmChatCompletionSummary" in event_types
|
|
assert "LlmChatCompletionMessage" in event_types
|
|
|
|
def test_no_content_when_recording_disabled_streaming(self):
|
|
logger = make_logger(turn_off_message_logging=True)
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "secret"}])
|
|
response = make_streaming_response(content="also secret")
|
|
|
|
messages = logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
|
|
for msg in messages:
|
|
assert "content" not in msg
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Explicit-None defensive tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExplicitNoneValues:
|
|
"""Verify that explicitly None values in kwargs/response don't raise or silently drop events."""
|
|
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
# _get_trace_context — chained dict lookups
|
|
def test_trace_context_litellm_params_none(self):
|
|
kwargs = make_kwargs()
|
|
kwargs["litellm_params"] = None
|
|
trace_id = self.logger._get_trace_context(kwargs)
|
|
assert trace_id is not None # falls back to UUID
|
|
|
|
def test_trace_context_metadata_none(self):
|
|
kwargs = make_kwargs()
|
|
kwargs["litellm_params"] = {"metadata": None}
|
|
trace_id = self.logger._get_trace_context(kwargs)
|
|
assert trace_id is not None
|
|
|
|
def test_trace_context_headers_none(self):
|
|
kwargs = make_kwargs()
|
|
kwargs["litellm_params"] = {"metadata": {"headers": None}}
|
|
trace_id = self.logger._get_trace_context(kwargs)
|
|
assert trace_id is not None
|
|
|
|
# _get_request_params
|
|
def test_request_params_optional_params_none(self):
|
|
assert self.logger._get_request_params({"optional_params": None}) == {}
|
|
|
|
# _get_model_names
|
|
def test_model_names_model_none_in_kwargs(self):
|
|
request_model, _ = self.logger._get_model_names(
|
|
{"model": None}, make_response()
|
|
)
|
|
assert request_model == "unknown"
|
|
|
|
def test_model_names_model_none_in_response(self):
|
|
response = make_response()
|
|
response["model"] = None
|
|
_, response_model = self.logger._get_model_names(make_kwargs(), response)
|
|
assert response_model == "gpt-4" # falls back to request_model from kwargs
|
|
|
|
# _extract_all_messages
|
|
def test_extract_messages_messages_none(self):
|
|
kwargs = make_kwargs()
|
|
kwargs["messages"] = None
|
|
response = make_response()
|
|
messages = self.logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
# No request messages, but response message should still be extracted
|
|
assert any(m.get("is_response") for m in messages)
|
|
|
|
def test_extract_messages_choices_none(self):
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
|
|
response = make_response()
|
|
response["choices"] = None
|
|
messages = self.logger._extract_all_messages(
|
|
kwargs, response, response_model="gpt-4", vendor="openai"
|
|
)
|
|
# No response messages, but request message should still be extracted
|
|
assert any(not m.get("is_response") for m in messages)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helper edge cases
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtractUsage:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_missing_usage_returns_zeros(self):
|
|
response = {"id": "r1", "model": "gpt-4", "choices": []}
|
|
usage = self.logger._extract_usage(response)
|
|
assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
|
|
|
|
def test_explicit_none_token_fields_return_zeros(self):
|
|
response = {
|
|
"usage": {
|
|
"prompt_tokens": None,
|
|
"completion_tokens": None,
|
|
"total_tokens": None,
|
|
}
|
|
}
|
|
usage = self.logger._extract_usage(response)
|
|
assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
|
|
|
|
|
|
class TestGetFinishReason:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_returns_unknown_when_no_choices(self):
|
|
response = {"choices": []}
|
|
assert self.logger._get_finish_reason(response) == "unknown"
|
|
|
|
def test_returns_unknown_when_choices_missing(self):
|
|
assert self.logger._get_finish_reason({}) == "unknown"
|
|
|
|
def test_returns_unknown_when_finish_reason_explicitly_none(self):
|
|
response = {"choices": [{"finish_reason": None}]}
|
|
assert self.logger._get_finish_reason(response) == "unknown"
|
|
|
|
|
|
class TestToEpochMs:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_float_passthrough(self):
|
|
assert self.logger._to_epoch_ms(1.0) == pytest.approx(1000.0)
|
|
|
|
def test_datetime_converted(self):
|
|
dt = datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
|
|
assert self.logger._to_epoch_ms(dt) == pytest.approx(dt.timestamp() * 1000.0)
|
|
|
|
|
|
class TestGetDuration:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_uses_kwargs_value_when_present(self):
|
|
kwargs = {"llm_api_duration_ms": 750.0}
|
|
assert self.logger._get_duration(kwargs, 0.0, 1.0) == 750.0
|
|
|
|
def test_calculates_from_float_timestamps(self):
|
|
kwargs = {}
|
|
result = self.logger._get_duration(kwargs, 1.0, 2.5)
|
|
assert result == pytest.approx(1500.0)
|
|
|
|
def test_calculates_from_datetime_timestamps(self):
|
|
kwargs = {}
|
|
start = datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
|
|
end = datetime(2024, 1, 1, 0, 0, 1, 500000, tzinfo=timezone.utc) # +1.5s
|
|
result = self.logger._get_duration(kwargs, start, end)
|
|
assert result == pytest.approx(1500.0)
|
|
|
|
def test_returns_none_when_nothing_available(self):
|
|
assert self.logger._get_duration({}, None, None) is None
|
|
|
|
|
|
class TestGetRequestParams:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_includes_only_present_params(self):
|
|
kwargs = {"optional_params": {"temperature": 0.7}}
|
|
params = self.logger._get_request_params(kwargs)
|
|
assert params == {"temperature": 0.7}
|
|
assert "max_tokens" not in params
|
|
|
|
def test_empty_when_no_optional_params(self):
|
|
assert self.logger._get_request_params({}) == {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _process_success — comprehensive happy-path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProcessSuccess:
|
|
def test_records_summary_and_message_events(self):
|
|
logger = make_logger()
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
|
|
kwargs = make_kwargs(
|
|
traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
optional_params={"temperature": 0.5, "max_tokens": 100},
|
|
)
|
|
response = make_response(
|
|
response_id="chatcmpl-xyz",
|
|
content="Hi there!",
|
|
finish_reason="stop",
|
|
prompt_tokens=5,
|
|
completion_tokens=10,
|
|
)
|
|
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
logger._process_success(kwargs, response, start_time=1.0, end_time=2.5)
|
|
|
|
calls = mock_app.record_custom_event.call_args_list
|
|
event_types = [c[0][0] for c in calls]
|
|
|
|
assert "LlmChatCompletionSummary" in event_types
|
|
assert "LlmChatCompletionMessage" in event_types
|
|
|
|
# Verify summary event fields
|
|
summary_data = next(
|
|
c[0][1] for c in calls if c[0][0] == "LlmChatCompletionSummary"
|
|
)
|
|
assert summary_data["vendor"] == "openai"
|
|
assert summary_data["request.model"] == "gpt-4"
|
|
assert summary_data["response.model"] == "gpt-4"
|
|
assert summary_data["response.choices.finish_reason"] == "stop"
|
|
assert summary_data["response.usage.prompt_tokens"] == 5
|
|
assert summary_data["response.usage.completion_tokens"] == 10
|
|
assert summary_data["response.usage.total_tokens"] == 15
|
|
assert summary_data["request.temperature"] == 0.5
|
|
assert summary_data["request.max_tokens"] == 100
|
|
assert summary_data["ingest_source"] == "litellm"
|
|
assert summary_data["trace_id"] == "aabbccddeeff00112233445566778899"
|
|
|
|
# Verify message event id format: "{llm_response_id}-{sequence}"
|
|
message_events = [
|
|
c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage"
|
|
]
|
|
assert any(e["id"].startswith("chatcmpl-xyz-") for e in message_events)
|
|
response_msg = next(e for e in message_events if e.get("is_response"))
|
|
assert response_msg["content"] == "Hi there!"
|
|
assert response_msg["role"] == "assistant"
|
|
|
|
def test_skips_when_disabled(self):
|
|
logger = make_logger()
|
|
logger.enabled = False
|
|
|
|
with patch("newrelic.agent.application") as mock_app:
|
|
logger._process_success(make_kwargs(), make_response())
|
|
|
|
mock_app.assert_not_called()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _record_error_metric
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRecordErrorMetric:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_calls_record_custom_metric(self):
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
|
|
with patch.object(self.logger, "_check_and_emit_periodic_metric"):
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
self.logger._record_error_metric()
|
|
|
|
mock_app.record_custom_metric.assert_called_once_with("LLM/LiteLLM/Error", 1)
|
|
|
|
def test_skips_when_app_disabled(self):
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = False
|
|
|
|
with patch.object(self.logger, "_check_and_emit_periodic_metric"):
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
self.logger._record_error_metric()
|
|
|
|
mock_app.record_custom_metric.assert_not_called()
|
|
|
|
def test_calls_check_and_emit_periodic_metric(self):
|
|
with patch.object(
|
|
self.logger, "_check_and_emit_periodic_metric"
|
|
) as mock_periodic:
|
|
with patch("newrelic.agent.application", return_value=MagicMock()):
|
|
self.logger._record_error_metric()
|
|
|
|
mock_periodic.assert_called_once()
|
|
|
|
def test_skips_when_logger_disabled(self):
|
|
self.logger.enabled = False
|
|
with patch("newrelic.agent.application") as mock_app:
|
|
self.logger._record_error_metric()
|
|
mock_app.assert_not_called()
|
|
|
|
def test_handles_exception(self):
|
|
with patch(
|
|
"newrelic.agent.application", side_effect=RuntimeError("agent down")
|
|
):
|
|
self.logger._record_error_metric() # must not raise
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _emit_supportability_metric
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestEmitSupportabilityMetric:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
NewRelicLogger._last_metric_emission_time = 0.0
|
|
|
|
def test_records_metric_with_correct_name_and_value(self):
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
with patch.object(
|
|
self.logger, "_get_litellm_version", return_value="1.80.0"
|
|
):
|
|
self.logger._emit_supportability_metric()
|
|
mock_app.record_custom_metric.assert_called_once_with(
|
|
"Supportability/Python/ML/LiteLLM/1.80.0", 1
|
|
)
|
|
|
|
def test_updates_last_emission_time(self):
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
fake_now = 9_999_999.0
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
with patch(
|
|
"litellm.integrations.newrelic.newrelic.time.time",
|
|
return_value=fake_now,
|
|
):
|
|
self.logger._emit_supportability_metric()
|
|
assert NewRelicLogger._last_metric_emission_time == fake_now
|
|
|
|
def test_skips_when_app_disabled(self):
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = False
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
self.logger._emit_supportability_metric()
|
|
mock_app.record_custom_metric.assert_not_called()
|
|
# Timestamp is still updated to back off lock contention during registration.
|
|
assert NewRelicLogger._last_metric_emission_time != 0.0
|
|
|
|
def test_skips_when_no_app(self):
|
|
with patch("newrelic.agent.application", return_value=None):
|
|
self.logger._emit_supportability_metric()
|
|
# Timestamp is updated even when app is None to back off lock contention
|
|
# if the agent never starts or is slow to initialise.
|
|
assert NewRelicLogger._last_metric_emission_time != 0.0
|
|
|
|
def test_handles_exception(self):
|
|
with patch(
|
|
"newrelic.agent.application", side_effect=RuntimeError("agent down")
|
|
):
|
|
self.logger._emit_supportability_metric() # must not raise
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _check_and_emit_periodic_metric
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCheckAndEmitPeriodicMetric:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
NewRelicLogger._last_metric_emission_time = 0.0
|
|
|
|
def test_emits_on_first_call(self):
|
|
"""_last_metric_emission_time starts at 0.0; any real time satisfies 27-hour window."""
|
|
with patch.object(self.logger, "_emit_supportability_metric") as mock_emit:
|
|
with patch(
|
|
"litellm.integrations.newrelic.newrelic.time.time",
|
|
return_value=100_000.0,
|
|
):
|
|
self.logger._check_and_emit_periodic_metric()
|
|
mock_emit.assert_called_once()
|
|
|
|
def test_does_not_re_emit_within_27_hours(self):
|
|
recent = 1_000_000.0
|
|
NewRelicLogger._last_metric_emission_time = recent
|
|
with patch.object(self.logger, "_emit_supportability_metric") as mock_emit:
|
|
with patch(
|
|
"litellm.integrations.newrelic.newrelic.time.time",
|
|
return_value=recent + 3600, # 1 hour later
|
|
):
|
|
self.logger._check_and_emit_periodic_metric()
|
|
mock_emit.assert_not_called()
|
|
|
|
def test_re_emits_after_27_hours(self):
|
|
old = 1_000_000.0
|
|
NewRelicLogger._last_metric_emission_time = old
|
|
with patch.object(self.logger, "_emit_supportability_metric") as mock_emit:
|
|
with patch(
|
|
"litellm.integrations.newrelic.newrelic.time.time",
|
|
return_value=old + 97201, # 27 hours + 1 second
|
|
):
|
|
self.logger._check_and_emit_periodic_metric()
|
|
mock_emit.assert_called_once()
|
|
|
|
def test_boundary_exactly_27_hours_triggers_emission(self):
|
|
old = 1_000_000.0
|
|
NewRelicLogger._last_metric_emission_time = old
|
|
with patch.object(self.logger, "_emit_supportability_metric") as mock_emit:
|
|
with patch(
|
|
"litellm.integrations.newrelic.newrelic.time.time",
|
|
return_value=old + 97200,
|
|
):
|
|
self.logger._check_and_emit_periodic_metric()
|
|
mock_emit.assert_called_once()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _get_litellm_version
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGetLitellmVersion:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_returns_unknown_on_exception(self):
|
|
with patch("importlib.metadata.version", side_effect=Exception("no package")):
|
|
result = self.logger._get_litellm_version()
|
|
assert result == "unknown"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _record_summary_event — disabled-app and exception paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_USAGE = {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}
|
|
|
|
|
|
class TestRecordSummaryEvent:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def _call(self, **kwargs):
|
|
self.logger._record_summary_event(
|
|
request_id="req-1",
|
|
trace_id="trace-abc",
|
|
request_model="gpt-4",
|
|
response_model="gpt-4",
|
|
vendor="openai",
|
|
finish_reason="stop",
|
|
num_messages=2,
|
|
usage=_USAGE,
|
|
**kwargs,
|
|
)
|
|
|
|
def test_skips_when_app_disabled(self):
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = False
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
self._call()
|
|
mock_app.record_custom_event.assert_not_called()
|
|
|
|
def test_handles_exception(self):
|
|
with patch(
|
|
"newrelic.agent.application", side_effect=RuntimeError("agent down")
|
|
):
|
|
self._call() # must not raise
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _record_message_events — disabled-app and exception paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_MESSAGES = [
|
|
{"role": "user", "sequence": 0, "response.model": "gpt-4", "vendor": "openai"}
|
|
]
|
|
|
|
|
|
class TestRecordMessageEvents:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def _call(self):
|
|
self.logger._record_message_events(
|
|
request_id="req-1",
|
|
llm_response_id="resp-1",
|
|
trace_id="trace-abc",
|
|
messages=_MESSAGES,
|
|
)
|
|
|
|
def test_skips_when_app_disabled(self):
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = False
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
self._call()
|
|
mock_app.record_custom_event.assert_not_called()
|
|
|
|
def test_handles_exception(self):
|
|
with patch(
|
|
"newrelic.agent.application", side_effect=RuntimeError("agent down")
|
|
):
|
|
self._call() # must not raise
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CustomLogger interface entry points
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLogSuccessEvent:
|
|
def test_delegates_to_process_success(self):
|
|
logger = make_logger()
|
|
with patch.object(logger, "_process_success") as mock_process:
|
|
logger.log_success_event(make_kwargs(), make_response(), 1.0, 2.0)
|
|
mock_process.assert_called_once()
|
|
|
|
def test_exception_is_handled(self):
|
|
logger = make_logger()
|
|
with patch.object(logger, "_process_success", side_effect=RuntimeError("boom")):
|
|
logger.log_success_event(make_kwargs(), make_response(), 1.0, 2.0)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_delegates_to_process_success(self):
|
|
logger = make_logger()
|
|
with patch.object(logger, "_process_success") as mock_process:
|
|
await logger.async_log_success_event(
|
|
make_kwargs(), make_response(), 1.0, 2.0
|
|
)
|
|
mock_process.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_exception_is_handled(self):
|
|
logger = make_logger()
|
|
with patch.object(logger, "_process_success", side_effect=RuntimeError("boom")):
|
|
await logger.async_log_success_event(
|
|
make_kwargs(), make_response(), 1.0, 2.0
|
|
)
|
|
|
|
|
|
class TestLogFailureEvent:
|
|
def test_sync_records_error_metric(self):
|
|
logger = make_logger()
|
|
with patch.object(logger, "_record_error_metric") as mock_metric:
|
|
logger.log_failure_event(make_kwargs(), None, 1.0, 2.0)
|
|
mock_metric.assert_called_once()
|
|
|
|
def test_sync_exception_is_handled(self):
|
|
logger = make_logger()
|
|
with patch.object(
|
|
logger, "_record_error_metric", side_effect=RuntimeError("boom")
|
|
):
|
|
logger.log_failure_event(make_kwargs(), None, 1.0, 2.0)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_records_error_metric(self):
|
|
logger = make_logger()
|
|
with patch.object(logger, "_record_error_metric") as mock_metric:
|
|
await logger.async_log_failure_event(make_kwargs(), None, 1.0, 2.0)
|
|
mock_metric.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_exception_is_handled(self):
|
|
logger = make_logger()
|
|
with patch.object(
|
|
logger, "_record_error_metric", side_effect=RuntimeError("boom")
|
|
):
|
|
await logger.async_log_failure_event(make_kwargs(), None, 1.0, 2.0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# async_health_check
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAsyncHealthCheck:
|
|
@pytest.mark.asyncio
|
|
async def test_unhealthy_when_disabled(self):
|
|
logger = make_logger()
|
|
logger.enabled = False
|
|
result = await logger.async_health_check()
|
|
assert result["status"] == "unhealthy"
|
|
assert result["error_message"] is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_healthy_when_app_enabled_records_test_event(self):
|
|
logger = make_logger()
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
result = await logger.async_health_check()
|
|
assert result["status"] == "healthy"
|
|
assert result["error_message"] is None
|
|
|
|
mock_app.record_custom_event.assert_called_once()
|
|
event_type, event_data = mock_app.record_custom_event.call_args[0]
|
|
assert event_type == "LiteLLMConnectionTest"
|
|
assert event_data["is_test_event"] is True
|
|
assert event_data["app_name"] == logger.app_name
|
|
assert event_data["source"] == "litellm-proxy"
|
|
assert isinstance(event_data["timestamp"], float)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unhealthy_when_app_disabled(self):
|
|
logger = make_logger()
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = False
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
result = await logger.async_health_check()
|
|
assert result["status"] == "unhealthy"
|
|
assert result["error_message"] is not None
|
|
mock_app.record_custom_event.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exception_returns_unhealthy(self):
|
|
logger = make_logger()
|
|
with patch(
|
|
"newrelic.agent.application", side_effect=RuntimeError("agent down")
|
|
):
|
|
result = await logger.async_health_check()
|
|
assert result["status"] == "unhealthy"
|
|
assert "agent down" in result["error_message"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_record_custom_event_failure_returns_unhealthy(self):
|
|
logger = make_logger()
|
|
mock_app = MagicMock()
|
|
mock_app.enabled = True
|
|
mock_app.record_custom_event.side_effect = RuntimeError("intake unreachable")
|
|
with patch("newrelic.agent.application", return_value=mock_app):
|
|
result = await logger.async_health_check()
|
|
assert result["status"] == "unhealthy"
|
|
assert "intake unreachable" in result["error_message"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _extract_completion_id fallback chain
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtractCompletionId:
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_uses_litellm_call_id_when_response_has_no_id(self):
|
|
result = self.logger._extract_completion_id(
|
|
kwargs={"litellm_call_id": "call-abc-123"},
|
|
response_obj={},
|
|
)
|
|
assert result == "call-abc-123"
|
|
|
|
def test_generates_uuid_when_neither_id_present(self):
|
|
result = self.logger._extract_completion_id(kwargs={}, response_obj={})
|
|
# UUID4 hex-with-dashes is 36 chars; just confirm shape and uniqueness
|
|
assert isinstance(result, str)
|
|
assert len(result) == 36
|
|
second = self.logger._extract_completion_id(kwargs={}, response_obj={})
|
|
assert result != second
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# StandardLoggingPayload preference across extractors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestStandardLoggingPayloadPreference:
|
|
"""Each extractor that accepts a StandardLoggingPayload must prefer its
|
|
values over the raw kwargs/response fallbacks."""
|
|
|
|
def setup_method(self):
|
|
self.logger = make_logger()
|
|
|
|
def test_trace_context_uses_slo_trace_id_when_no_traceparent(self):
|
|
kwargs = {"litellm_params": {"metadata": {"headers": {}}}}
|
|
trace_id = self.logger._get_trace_context(
|
|
kwargs, standard_logging_object=make_slo()
|
|
)
|
|
assert trace_id == "slo-trace-abc"
|
|
|
|
def test_vendor_from_slo(self):
|
|
# kwargs carries a different provider; SLO must win.
|
|
kwargs = {"litellm_params": {"custom_llm_provider": "kwargs-provider"}}
|
|
assert (
|
|
self.logger._get_vendor(kwargs, standard_logging_object=make_slo())
|
|
== "slo-provider"
|
|
)
|
|
|
|
def test_model_names_uses_slo_model(self):
|
|
request_model, _ = self.logger._get_model_names(
|
|
{"model": "kwargs-model"},
|
|
make_response(model="response-model"),
|
|
standard_logging_object=make_slo(),
|
|
)
|
|
assert request_model == "slo-model"
|
|
|
|
def test_usage_from_slo_when_any_token_field_present(self):
|
|
# make_response defaults to 10/20/30 tokens; SLO sentinels are 100/200/300.
|
|
usage = self.logger._extract_usage(
|
|
make_response(), standard_logging_object=make_slo()
|
|
)
|
|
assert usage == {
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 200,
|
|
"total_tokens": 300,
|
|
}
|
|
|
|
def test_duration_from_slo_response_time_converted_to_ms(self):
|
|
# SLO response_time is 1.5 seconds; expected 1500.0 ms.
|
|
# Pass start/end that would compute a different value to prove SLO won.
|
|
duration = self.logger._get_duration(
|
|
kwargs={"llm_api_duration_ms": 9999.0},
|
|
start_time=1.0,
|
|
end_time=2.0,
|
|
standard_logging_object=make_slo(),
|
|
)
|
|
assert duration == 1500.0
|
|
|
|
def test_request_params_from_slo_model_parameters(self):
|
|
params = self.logger._get_request_params(
|
|
{"optional_params": {"temperature": 0.1}},
|
|
standard_logging_object=make_slo(),
|
|
)
|
|
assert params == {"temperature": 0.7, "max_tokens": 500}
|
|
|
|
def test_extract_all_messages_sources_timestamps_and_messages_from_slo(self):
|
|
"""Covers three SLO branches at once: startTime, endTime, and messages list."""
|
|
kwargs = make_kwargs(messages=[{"role": "user", "content": "from-kwargs"}])
|
|
messages = self.logger._extract_all_messages(
|
|
kwargs,
|
|
make_response(),
|
|
response_model="gpt-4",
|
|
vendor="openai",
|
|
standard_logging_object=make_slo(),
|
|
)
|
|
|
|
request = next(m for m in messages if not m.get("is_response"))
|
|
assert request["content"] == "from-slo" # SLO messages list wins
|
|
assert request["timestamp"] == int(2_000_000.0 * 1000.0) # SLO startTime
|
|
|
|
response = next(m for m in messages if m.get("is_response"))
|
|
assert response["timestamp"] == int(2_000_001.5 * 1000.0) # SLO endTime
|