mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
cf491d1df9
* ci: run the unit_selection.sh shard files on every event instead of only fork pull requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rename fork-flag to unit-flag now that it applies on every event * test: move tests/test_litellm root and small trees into tests/unit Pure renames, no content changes. Follow-up commits in this PR fix references, merge the three files that already existed in tests/unit, keep live-provider tests in tests/test_litellm and wire CI. * test: carry tests/test_litellm conftest isolation into tests/unit Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS, proxy-URL and keychain env, and session-end client cleanup now reset for unit tests too. The environment isolation owns its MonkeyPatch so a test's own monkeypatch is undone before the model-cost teardown runs. * test: merge, split and prune the moved root and small-tree tests Merge batches/test_batch_utils.py and the chat_completions and messages dispatch tests into the files that already existed in tests/unit. Keep the live Gemini interactions tests, the async image-fetch format test and the OpenAI embedding scorer test in tests/test_litellm since they need real network or keys. Put test_router.py under tests/unit/test_router so the existing package no longer shadows it. Delete eight tests the audit found superseded by stronger ones kept in this move. * ci: run the moved root and small-tree tests under their legacy flags Add the misc and responses-caching-types flags to unit_selection.sh and CircleCI, extend enterprise-routing and mcp-integration, and point the legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest and change classifier at the new paths. * test: make the new tests/unit directories packages tests/unit/test_package_layout.py requires every directory to carry an __init__.py, and without one the moved and retained test_litellm_responses_bridge.py modules collide on import. * test: scope the unit socket block to tests/unit in shared sessions The GHA shards collect the legacy test-path and the unit selection in one pytest session. The unit conftest's loopback-only block leaked into legacy modules that reach the network at import. The legacy conftest now lifts the restriction at collect and setup time, and the unit conftest re-applies it when collecting its own modules. * test: move tests/test_litellm/llms into tests/unit/llms Rename-only. Moves the provider tests and the fine-tuning fixtures they load, mirroring the old paths. Follow-up commits merge, split and wire them. * test: merge, split and prune the moved llms tests Merges the Databricks chat transformation tests into the existing unit file, keeps the tests that need real keys or the network in tests/test_litellm, deletes the audited tests a stronger unit test already covers, and points imports at tests.unit.llms. * ci: run the moved llms tests under their legacy flags The Vertex AI and All Other Providers shards keep their legacy test-path for the retained files and add the llm-vertex-ai and llm-other-providers unit selections. CircleCI gets matching unit jobs. * test: make the tests/unit/llms directories packages Adds __init__.py to the moved dirs and drops the legacy ones whose directories no longer hold tests. * test: drop script runners and path hacks the llms split left dangling The __main__ runners in the split openai_like files and the Databricks e2e runner called tests that now live in the other half of the split or were deleted. The retained legacy halves also no longer need sys.path edits. * test: give the shard-script tests their own GITHUB_OUTPUT They only passed where the runner set it. The CircleCI unit job's env allowlist drops it, so the script's redirect failed there. * test: point the router and module-deletion checks at tests/unit router_code_coverage and code_qa_check_tests only searched tests/test_litellm, so the moved router tests no longer counted. The two silent-experiment tests the audit deleted were the only direct callers of those methods; they are replaced with tests that assert the forwarded shadow request and the recursion guard. * test: move tests/test_litellm integrations and secret_managers into tests/unit Rename-only. Mirrors the old paths, including the directory conftests and the prompt and JSON fixtures. Follow-up commits prune and wire them. * test: prune and repoint the moved integrations tests Deletes the 7 audited tests a stronger test in the same tree already covers, imports the TLS sink helpers from their new conftest path, and restores os.environ after each integrations test. Some presets write OTEL_EXPORTER_OTLP_HEADERS straight into os.environ, and without the legacy tree's test ordering that header leaked into the AgentOps tests. * ci: run the moved integrations tests under their legacy flag The integrations GHA shard and a new CircleCI job run the integrations unit selection. secret_managers joins the misc selection. * docs: point integrations and secret_managers references at tests/unit * test: make the moved integrations directories packages * test: keep the Databricks manual e2e runner and fix the SageMaker Nova run path The Databricks e2e file is a manual script whose main() calls the tests that were pruned, so pruning them broke the documented run. It is back to its main version. The SageMaker Nova docstring now points at the file's real location in tests/local_testing. * test: keep the job's UNIT_FLAG out of the shard-script tests --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
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
|