fix(logging): take the mirror flag in Prometheus labels and Datadog cost tags only from the router

This commit is contained in:
mateo-berri 2026-10-28 23:23:39 -07:00
parent 232e61d6b7
commit 5b45922c63
6 changed files with 158 additions and 5 deletions

View file

@ -214,11 +214,12 @@ class DatadogCostManagementLogger(CustomBatchLogger):
for k, v in metadata.items():
if k in allow and v is not None and not isinstance(v, (dict, list)):
self._set_custom_tag(tags, k, str(v))
client_allow: Final = allow - {"is_silent_experiment"}
for nested_key in ("spend_logs_metadata", "requester_metadata"):
nested = metadata.get(nested_key)
if isinstance(nested, dict):
for k, v in nested.items():
if k in allow and v is not None and not isinstance(v, (dict, list)):
if k in client_allow and v is not None and not isinstance(v, (dict, list)):
self._set_custom_tag(tags, k, str(v))
return tags

View file

@ -5148,6 +5148,12 @@ def _get_combined_custom_metadata_from_standard_logging_payload(
**(requester_metadata if isinstance(requester_metadata, dict) else {}),
**(user_api_key_auth_metadata if isinstance(user_api_key_auth_metadata, dict) else {}),
**(spend_logs_metadata if isinstance(spend_logs_metadata, dict) else {}),
"is_silent_experiment": (
"true"
if "is_silent_experiment" in standard_logging_metadata
and standard_logging_metadata["is_silent_experiment"] is True
else None
),
}

View file

@ -1,4 +1,5 @@
import time
from typing import Final
from unittest.mock import AsyncMock
import pytest
@ -314,6 +315,34 @@ async def test_extract_tags_nested_metadata_allowlisted(clean_env):
assert "ignored" not in tags
@pytest.mark.parametrize(
("metadata", "expected_tag"),
[
({"is_silent_experiment": True, "requester_metadata": {"is_silent_experiment": False}}, "true"),
(
{
"requester_metadata": {"is_silent_experiment": True},
"spend_logs_metadata": {"is_silent_experiment": "true"},
},
None,
),
],
ids=["mirrored call", "direct call with a client flag"],
)
@pytest.mark.asyncio
async def test_extract_tags_takes_the_silent_experiment_flag_only_from_the_router(clean_env, metadata, expected_tag):
logger: Final = DatadogCostManagementLogger(cost_tag_keys=["is_silent_experiment"])
log: Final = StandardLoggingPayload(
custom_llm_provider="openai",
model="gpt-4",
response_cost=0.01,
startTime=time.time(),
metadata=metadata,
)
assert logger._extract_tags(log).get("is_silent_experiment") == expected_tag
@pytest.mark.asyncio
async def test_extract_tags_allowlist_cannot_override_reserved_dimensions(clean_env):
"""

View file

@ -5,6 +5,10 @@ Verifies that metadata from x-litellm-spend-logs-metadata header is available
in Prometheus custom labels via combined_metadata.
"""
from typing import Final
import pytest
from litellm.integrations.prometheus import (
_get_combined_custom_metadata_from_standard_logging_payload,
get_custom_labels_from_metadata,
@ -205,3 +209,32 @@ def test_project_alias_accessible_without_prefix(monkeypatch):
result = get_custom_labels_from_metadata(combined)
assert result == {"user_api_key_project_alias": "hotel-recommendations"}
@pytest.mark.parametrize(
("standard_logging_metadata", "expected_labels"),
[
(
{"is_silent_experiment": True, "requester_metadata": {"is_silent_experiment": False}},
{"metadata_is_silent_experiment": "true"},
),
(
{
"requester_metadata": {"is_silent_experiment": True},
"spend_logs_metadata": {"is_silent_experiment": "true"},
},
{},
),
],
ids=["mirrored call", "direct call with a client flag"],
)
def test_silent_experiment_label_comes_only_from_the_router_flag(
monkeypatch, standard_logging_metadata, expected_labels
):
monkeypatch.setattr("litellm.custom_prometheus_metadata_labels", ["metadata.is_silent_experiment"])
combined: Final = _get_combined_custom_metadata_from_standard_logging_payload(
{"metadata": standard_logging_metadata}
)
assert get_custom_labels_from_metadata(combined) == expected_labels

View file

@ -11979,7 +11979,7 @@ def test_get_standard_logging_metadata(metadata_key, metadata_value):
def test_module_level_get_standard_logging_metadata_drops_only_a_caller_sent_silent_experiment_flag():
from litellm.litellm_core_utils.litellm_logging import get_standard_logging_metadata
caller_metadata = {
caller_metadata: Final = {
"is_silent_experiment": True,
"user_api_key_alias": "alias",
"usage_object": {"total_tokens": 3},

View file

@ -24,8 +24,14 @@ class _RecordingLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.success_kwargs: list[dict[str, object]] = []
self.failure_kwargs: list[dict[str, object]] = []
self._logged: Final = threading.Condition()
def _record(self, calls: list[dict[str, object]], kwargs: dict[str, object]) -> None:
with self._logged:
calls.append(kwargs)
self._logged.notify_all()
async def async_log_success_event(
self,
kwargs: dict[str, object],
@ -33,9 +39,16 @@ class _RecordingLogger(CustomLogger):
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
with self._logged:
self.success_kwargs.append(kwargs)
self._logged.notify_all()
self._record(self.success_kwargs, kwargs)
async def async_log_failure_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
self._record(self.failure_kwargs, kwargs)
def wait_until(self, logged_enough: Callable[[], bool], timeout: float = 30.0) -> None:
with self._logged:
@ -49,6 +62,17 @@ class _RecordingLogger(CustomLogger):
]
class _SyncAndAsyncRecordingLogger(_RecordingLogger):
def log_success_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
self._record(self.success_kwargs, kwargs)
async def _settle_shared_logging_worker() -> None:
try:
await GLOBAL_LOGGING_WORKER.flush()
@ -117,6 +141,11 @@ def recording_logger():
yield from _installed(_RecordingLogger())
@pytest.fixture
def sync_and_async_recording_logger():
yield from _installed(_SyncAndAsyncRecordingLogger())
@pytest.fixture
def judging_logger():
yield from _installed(_JudgingLogger())
@ -356,6 +385,61 @@ async def test_standard_logging_payload_marks_only_the_mirrored_calls_as_silent_
assert "is_silent_experiment" not in metadata_by_group["primary-model"]
def test_standard_logging_payload_marks_only_the_mirrored_calls_on_the_sync_path(
sync_and_async_recording_logger: _SyncAndAsyncRecordingLogger,
) -> None:
router: Final = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
response: Final = router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
)
assert response.choices[0].message.content == "pong"
sync_and_async_recording_logger.wait_until(
lambda: _silent_flag_by_group(sync_and_async_recording_logger).keys() >= {"primary-model", "shadow-a", "shadow-b"}
)
assert _silent_flag_by_group(sync_and_async_recording_logger) == {
"primary-model": None,
"shadow-a": True,
"shadow-b": True,
}
@pytest.mark.asyncio
async def test_a_failed_mirrored_call_reaches_failure_callbacks_marked_as_mirrored(
recording_logger: _RecordingLogger,
) -> None:
router: Final = Router(
model_list=[
{
"model_name": "primary-model",
"litellm_params": {
"model": "openai/gpt-5.4-mini",
"api_key": "fake-key",
"silent_model": "shadow-a",
"mock_response": "pong",
},
},
{
"model_name": "shadow-a",
"litellm_params": {
"model": "openai/gpt-5.4-nano",
"api_key": "fake-key",
"mock_response": "litellm.InternalServerError",
},
},
]
)
await router.acompletion(model="primary-model", messages=[{"role": "user", "content": "hi"}])
await asyncio.to_thread(recording_logger.wait_until, lambda: len(recording_logger.failure_kwargs) >= 1)
assert {
(call["standard_logging_object"]["model_group"], call["standard_logging_object"]["metadata"].get("is_silent_experiment"))
for call in recording_logger.failure_kwargs
} == {("shadow-a", True)}
_SPEND_ROW_TIME: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)