mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(logging): mark mirrored calls through a context variable so router usage counts them once
This commit is contained in:
parent
4e248208a5
commit
e2b020d6ce
7 changed files with 57 additions and 80 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue