fix(logging): mark mirrored calls through a context variable so router usage counts them once

This commit is contained in:
mateo-berri 2026-10-28 13:04:04 -07:00
parent 4e248208a5
commit e2b020d6ce
7 changed files with 57 additions and 80 deletions

View file

@ -36,6 +36,8 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families"
_service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None)
_mirrored_call: Final[ContextVar[bool]] = ContextVar("mirrored_call", default=False)
@contextmanager
def internal_sub_call() -> Generator[None]:
@ -123,6 +125,20 @@ def current_service_caller() -> str | None:
return _service_caller.get()
@contextmanager
def mirrored_call() -> Generator[None]:
"""A silent_model mirror of a request: every litellm call inside, its retries and fallbacks included."""
token: Final = _mirrored_call.set(True)
try:
yield
finally:
_restore(_mirrored_call, token)
def in_mirrored_call() -> bool:
return _mirrored_call.get()
@contextmanager
def pinned_billing_time(moment: datetime) -> Generator[None]:
"""Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read."""

View file

@ -1663,7 +1663,6 @@ ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags"
VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata"
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
SILENT_EXPERIMENT_METADATA_KEY: Final = "_litellm_silent_experiment"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt"
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"

View file

@ -23,7 +23,7 @@ from httpx import Response, Timeout
from pydantic import BaseModel, JsonValue, TypeAdapter
import litellm
from litellm._internal_context import post_response_phase
from litellm._internal_context import in_mirrored_call, post_response_phase
from litellm._logging import (
is_debugging_on,
redact_string,
@ -47,7 +47,6 @@ from litellm.constants import (
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
PROVIDER_REQUEST_ID_HEADERS,
REDACTED_BY_LITELLM,
SILENT_EXPERIMENT_METADATA_KEY,
)
from litellm.cost_calculator import (
RealtimeAPITokenUsageProcessor,
@ -163,7 +162,7 @@ from litellm.types.llms.openai import (
from litellm.types.mcp import MCPPostCallResponseObject
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.rerank import RerankResponse
from litellm.types.router import RouterRateLimitError, SilentExperimentFlag, SilentExperimentMarker
from litellm.types.router import RouterRateLimitError, SilentExperimentFlag
from litellm.types.utils import (
DEPLOYMENT_SCOPED_PRICING_FIELDS,
CachingDetails,
@ -317,7 +316,6 @@ _in_memory_loggers: Final[list[CustomLogger]] = []
_STANDARD_LOGGING_METADATA_RESOLVED_KEYS: Final[frozenset[str]] = frozenset(
("used_client_oauth_token", "usage_object", KUBERNETES_POD_ROUTING_KEY, "is_silent_experiment")
)
_SILENT_EXPERIMENT_BUCKETS: Final = ("metadata", "litellm_metadata")
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = (
frozenset(StandardLoggingMetadata.__annotations__.keys()) - _STANDARD_LOGGING_METADATA_RESOLVED_KEYS
)
@ -810,7 +808,7 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_params = scrub_sensitive_keys_in_metadata(litellm_params)
self.litellm_params = litellm_params
self.is_silent_experiment = False
self.is_silent_experiment: Final = in_mirrored_call()
# Initialize cost breakdown field
self.cost_breakdown: CostBreakdown | None = None
@ -1126,7 +1124,6 @@ class Logging(LiteLLMLoggingBaseClass):
if model is not None:
self.model = model
self.user = user
self.is_silent_experiment = self.is_silent_experiment or is_silent_experiment_call(litellm_params)
self.litellm_params = {
**self.litellm_params,
**scrub_sensitive_keys_in_metadata(litellm_params),
@ -7727,9 +7724,7 @@ def get_standard_logging_metadata(
return clean_metadata
def scrub_sensitive_keys_in_metadata(
litellm_params: dict[str, Any] | None, # mutable-ok: the langfuse masking function is moved onto these params
) -> dict[str, Any]: # mutable-ok: loggers write into the returned params and their metadata
def scrub_sensitive_keys_in_metadata(litellm_params: dict | None):
if litellm_params is None:
litellm_params = {}
@ -7744,24 +7739,7 @@ def scrub_sensitive_keys_in_metadata(
litellm_params["_langfuse_masking_function"] = masking_fn
litellm_params["metadata"] = metadata
if not is_silent_experiment_call(litellm_params):
return litellm_params
return {
name: (
{key: value for key, value in bucket.items() if key != SILENT_EXPERIMENT_METADATA_KEY}
if name in _SILENT_EXPERIMENT_BUCKETS and _carries_silent_experiment_marker(bucket)
else bucket
)
for name, bucket in litellm_params.items()
}
def _carries_silent_experiment_marker(bucket: object) -> bool:
return isinstance(bucket, dict) and bucket.get(SILENT_EXPERIMENT_METADATA_KEY) is SilentExperimentMarker.MIRRORED
def is_silent_experiment_call(litellm_params: Mapping[str, object]) -> bool:
return any(_carries_silent_experiment_marker(litellm_params.get(name)) for name in _SILENT_EXPERIMENT_BUCKETS)
return litellm_params
# integration helper function

View file

@ -52,7 +52,7 @@ from typing_extensions import assert_never, overload
import litellm
import litellm.litellm_core_utils.exception_mapping_utils
from litellm import get_secret_str
from litellm._internal_context import service_target, with_service_target
from litellm._internal_context import mirrored_call, service_target, with_service_target
from litellm._logging import verbose_router_logger
from litellm._uuid import uuid
from litellm.caching.caching import (
@ -79,7 +79,6 @@ from litellm.constants import (
ROUTING_REQUEST_TAGS_METADATA_KEY,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
SILENT_EXPERIMENT_METADATA_KEY,
)
from litellm.integrations.custom_guardrail import CustomGuardrail, is_guardrail_intervention
from litellm.integrations.custom_logger import CustomLogger
@ -424,7 +423,6 @@ from litellm.types.router import (
RoutingPlugin,
RoutingStrategy,
SearchToolTypedDict,
SilentExperimentMarker,
TaggedPreRoutingStrategy,
holds_secret_pointer,
)
@ -3294,12 +3292,6 @@ class Router:
silent_kwargs["metadata"].pop("litellm_parent_otel_span", None)
silent_kwargs["metadata"]["is_silent_experiment"] = True
silent_kwargs["metadata"][SILENT_EXPERIMENT_METADATA_KEY] = SilentExperimentMarker.MIRRORED
if "litellm_metadata" in silent_kwargs and isinstance(silent_kwargs["litellm_metadata"], dict):
silent_kwargs["litellm_metadata"] = {
**silent_kwargs["litellm_metadata"],
SILENT_EXPERIMENT_METADATA_KEY: SilentExperimentMarker.MIRRORED,
}
# Pop logging objects and call IDs to ensure a fresh logging context
# This prevents collisions in the Proxy's database (spend_logs)
@ -3316,16 +3308,17 @@ class Router:
remaining_kwargs: Final = MappingProxyType(
{key: value for key, value in silent_kwargs.items() if key != "stream"}
)
response: Final = await self.acompletion(
model=silent_model,
messages=cast(list[AllMessageValues], messages),
stream=bool(silent_kwargs.get("stream", False)),
**remaining_kwargs,
)
if not isinstance(response, CustomStreamWrapper):
return
async for _ in response:
pass
with mirrored_call():
response: Final = await self.acompletion(
model=silent_model,
messages=cast(list[AllMessageValues], messages),
stream=bool(silent_kwargs.get("stream", False)),
**remaining_kwargs,
)
if not isinstance(response, CustomStreamWrapper):
return
async for _ in response:
pass
def _silent_experiment_completion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
"""

View file

@ -66,10 +66,6 @@ from .utils import (
AllowedModelRegion: TypeAlias = Literal["eu", "us"]
class SilentExperimentMarker(str, enum.Enum):
MIRRORED = "mirrored"
class SilentExperimentFlag(TypedDict, total=False):
is_silent_experiment: ReadOnly[bool]

View file

@ -30,7 +30,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS, SILENT_EXPERIMENT_METADATA_KEY
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -63,7 +63,6 @@ from litellm.proxy.pass_through_endpoints.success_handler import (
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
from litellm.types import utils as types_utils
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import SilentExperimentMarker
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
@ -10439,9 +10438,8 @@ def test_init_kwargs_no_authenticated_identity_field_is_client_settable(mock_req
assert survived == {}, f"client-supplied values survived for: {sorted(survived)}"
@pytest.mark.parametrize("planted", [True, SilentExperimentMarker.MIRRORED.value])
def test_init_kwargs_client_metadata_cannot_mark_a_pass_through_call_as_a_silent_experiment(
planted: object, mock_user_api_key_dict: UserAPIKeyAuth
mock_user_api_key_dict: UserAPIKeyAuth,
) -> None:
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
@ -10461,8 +10459,8 @@ def test_init_kwargs_client_metadata_cannot_mark_a_pass_through_call_as_a_silent
function_id="test-function-id",
),
_parsed_body={
"litellm_metadata": {SILENT_EXPERIMENT_METADATA_KEY: planted, "client_key": "kept"},
"metadata": {SILENT_EXPERIMENT_METADATA_KEY: planted},
"litellm_metadata": {"is_silent_experiment": True, "client_key": "kept"},
"metadata": {"is_silent_experiment": True},
},
)

View file

@ -12,13 +12,11 @@ import pytest
import respx
import litellm
from litellm.constants import SILENT_EXPERIMENT_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.router import Router
from litellm.router import _silent_experiment_kwargs_snapshot
from litellm.router import _silent_experiment_targets
from litellm.types.router import SilentExperimentMarker
class _RecordingLogger(CustomLogger):
@ -349,26 +347,27 @@ async def test_spend_rows_mark_only_the_mirrored_calls_when_the_caller_sends_lit
@pytest.mark.asyncio
async def test_logging_callbacks_never_receive_the_router_mirror_marker(recording_logger: _RecordingLogger) -> None:
router: Final = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
litellm_metadata={"user_api_key": "test-key"},
async def test_router_counts_a_mirrored_call_once_against_its_target_rate_limits(
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-b"},
},
{
"model_name": "shadow-b",
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake-key", "rpm": 100, "tpm": 100000},
},
]
)
await _wait_for_all_successes(recording_logger, expected=3)
await router.acompletion(model="primary-model", messages=[{"role": "user", "content": "hi"}], mock_response="pong")
await _wait_for_shadow_successes(recording_logger, expected=1)
await GLOBAL_LOGGING_WORKER.flush()
request_metadata_keys: Final = {
call["standard_logging_object"]["model_group"]: {
*call["litellm_params"]["metadata"],
*call["litellm_params"]["litellm_metadata"],
}
for call in recording_logger.success_kwargs
}
assert sorted(request_metadata_keys) == ["primary-model", "shadow-a", "shadow-b"]
assert all(SILENT_EXPERIMENT_METADATA_KEY not in keys for keys in request_metadata_keys.values())
assert all("is_silent_experiment" in request_metadata_keys[group] for group in ("shadow-a", "shadow-b"))
shadow_tokens: Final = recording_logger.shadow_successes()[0]["standard_logging_object"]["total_tokens"]
assert await router.get_model_group_usage("shadow-b") == (shadow_tokens, 1)
@pytest.mark.asyncio
@ -376,9 +375,7 @@ async def test_logging_callbacks_never_receive_the_router_mirror_marker(recordin
("bucket", "planted"),
[
("metadata", {"is_silent_experiment": True}),
("metadata", {SILENT_EXPERIMENT_METADATA_KEY: SilentExperimentMarker.MIRRORED.value}),
("litellm_metadata", {SILENT_EXPERIMENT_METADATA_KEY: SilentExperimentMarker.MIRRORED.value}),
("litellm_metadata", {SILENT_EXPERIMENT_METADATA_KEY: True}),
("litellm_metadata", {"is_silent_experiment": True}),
],
)
async def test_a_caller_sent_silent_experiment_flag_does_not_mark_a_real_call_as_mirrored(