diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index 65edfd98768..c2243b90fbf 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -139,6 +139,16 @@ def in_mirrored_call() -> bool: return _mirrored_call.get() +@contextmanager +def outside_mirrored_call() -> Generator[None]: + """Logging a mirrored call: the calls its callbacks start are not mirrors themselves.""" + token: Final = _mirrored_call.set(False) + try: + yield + finally: + _restore(_mirrored_call, token) + + @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.""" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 59956e0a5eb..d86090b183e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -23,7 +23,7 @@ from httpx import Response, Timeout from pydantic import BaseModel, JsonValue, TypeAdapter import litellm -from litellm._internal_context import in_mirrored_call, post_response_phase +from litellm._internal_context import in_mirrored_call, outside_mirrored_call, post_response_phase from litellm._logging import ( is_debugging_on, redact_string, @@ -3246,7 +3246,7 @@ class Logging(LiteLLMLoggingBaseClass): """Restores trace_id/session_id contextvars once this attempt's own success logging (including any nested calls its callbacks trigger) is fully done.""" try: - with post_response_phase(): + with post_response_phase(), outside_mirrored_call(): return self._success_handler_body( result=result, start_time=start_time, @@ -3746,7 +3746,7 @@ class Logging(LiteLLMLoggingBaseClass): """Restores trace_id/session_id contextvars once this attempt's own success logging (including any nested calls its callbacks trigger) is fully done.""" try: - with post_response_phase(): + with post_response_phase(), outside_mirrored_call(): return await self._async_success_handler_body( result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs ) @@ -4263,13 +4263,14 @@ class Logging(LiteLLMLoggingBaseClass): ``for_async_request`` marks the sync run of a failure whose async run also logs it, so a CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook.""" try: - return self._failure_handler_body( - exception=exception, - traceback_exception=traceback_exception, - start_time=start_time, - end_time=end_time, - for_async_request=for_async_request, - ) + with outside_mirrored_call(): + return self._failure_handler_body( + exception=exception, + traceback_exception=traceback_exception, + start_time=start_time, + end_time=end_time, + for_async_request=for_async_request, + ) finally: self._restore_correlation_context() @@ -4462,12 +4463,13 @@ class Logging(LiteLLMLoggingBaseClass): """Restores trace_id/session_id contextvars once this attempt's own failure logging (including any nested calls its callbacks trigger) is fully done.""" try: - return await self._async_failure_handler_body( - exception=exception, - traceback_exception=traceback_exception, - start_time=start_time, - end_time=end_time, - ) + with outside_mirrored_call(): + return await self._async_failure_handler_body( + exception=exception, + traceback_exception=traceback_exception, + start_time=start_time, + end_time=end_time, + ) finally: self._restore_correlation_context() @@ -7714,9 +7716,8 @@ def get_standard_logging_metadata( ) if isinstance(metadata, dict): # Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields - for key in StandardLoggingMetadata.__annotations__: - if key in metadata: - clean_metadata[key] = metadata[key] + for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS: + clean_metadata[key] = metadata[key] if metadata.get("user_api_key") is not None: if is_valid_sha256_hash(str(metadata.get("user_api_key"))): diff --git a/litellm/router.py b/litellm/router.py index a94ab430898..a5bdb02aadb 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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 mirrored_call, service_target, with_service_target +from litellm._internal_context import in_mirrored_call, mirrored_call, service_target, with_service_target from litellm._logging import verbose_router_logger from litellm._uuid import uuid from litellm.caching.caching import ( @@ -702,7 +702,7 @@ def _every_deployment_refused( def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]: - if silent_model is None: + if silent_model is None or in_mirrored_call(): return () try: targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model) @@ -3325,10 +3325,6 @@ class Router: Run a silent experiment in the background (thread). """ try: - # Prevent infinite recursion if silent model also has a silent model - if kwargs.get("metadata", {}).get("is_silent_experiment", False): - return - messages = copy.deepcopy(messages) verbose_router_logger.info("Starting silent experiment for model %s", silent_model) @@ -4513,10 +4509,6 @@ class Router: Run a silent experiment in the background. """ try: - # Prevent infinite recursion if silent model also has a silent model - if kwargs.get("metadata", {}).get("is_silent_experiment", False): - return - messages = copy.deepcopy(messages) verbose_router_logger.info("Starting silent experiment for model %s", silent_model) diff --git a/tests/integration/sdk/test_router_silent_model_shadow_sdk.py b/tests/integration/sdk/test_router_silent_model_shadow_sdk.py index 821e4ddd0b7..be5998e205e 100644 --- a/tests/integration/sdk/test_router_silent_model_shadow_sdk.py +++ b/tests/integration/sdk/test_router_silent_model_shadow_sdk.py @@ -82,9 +82,9 @@ class _ShadowRecorder(CustomLogger): async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - metadata: Final = _OBJECT.validate_python(_OBJECT.validate_python(kwargs["litellm_params"])["metadata"] or {}) - if metadata.get("is_silent_experiment") is True: - self.shadow_groups.put(str(metadata.get("model_group"))) + payload: Final = _OBJECT.validate_python(kwargs["standard_logging_object"]) + if _OBJECT.validate_python(payload["metadata"]).get("is_silent_experiment") is True: + self.shadow_groups.put(str(payload["model_group"])) def drained(self) -> tuple[str, ...]: return tuple(self.shadow_groups.get_nowait() for _ in range(self.shadow_groups.qsize())) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 4c65613b0c8..2e48c16c2b5 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -11975,6 +11975,16 @@ def test_get_standard_logging_metadata(metadata_key, metadata_value): assert standard_logging_metadata[metadata_key] == metadata_value +@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") +def test_module_level_get_standard_logging_metadata_ignores_a_caller_sent_silent_experiment_flag(): + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_metadata + + result = get_standard_logging_metadata({"is_silent_experiment": True, "user_api_key_alias": "alias"}) + + assert "is_silent_experiment" not in result + assert result["user_api_key_alias"] == "alias" + + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_user_api_key_hash(): valid_hash = "a" * 64 # 64 character string diff --git a/tests/unit/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py index 69b06525792..35b902c089b 100644 --- a/tests/unit/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -2,7 +2,7 @@ import asyncio import datetime import json import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Generator, Mapping from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -42,15 +42,43 @@ async def _settle_shared_logging_worker() -> None: await GLOBAL_LOGGING_WORKER.stop() -@pytest.fixture -def recording_logger(): +class _JudgingLogger(_RecordingLogger): + async def _judge(self, kwargs: Mapping[str, object]) -> None: + metadata: Final = kwargs["litellm_params"]["metadata"] + if "judge_of" in metadata: + return + await litellm.acompletion( + model="openai/gpt-5.4-nano", + api_key="fake-key", + messages=[{"role": "user", "content": "score it"}], + mock_response="ok", + metadata={"judge_of": metadata["model_group"]}, + ) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await super().async_log_success_event(kwargs, response_obj, start_time, end_time) + await self._judge(kwargs) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await self._judge(kwargs) + + def judge_flags(self) -> dict[str, object]: + return { + call["litellm_params"]["metadata"]["judge_of"]: call["standard_logging_object"]["metadata"].get( + "is_silent_experiment" + ) + for call in self.success_kwargs + if "judge_of" in call["litellm_params"]["metadata"] + } + + +def _installed(logger: _RecordingLogger) -> Generator[_RecordingLogger]: settle_loop: Final = asyncio.new_event_loop() try: settle_loop.run_until_complete(_settle_shared_logging_worker()) finally: settle_loop.close() original_callbacks: Final = litellm.callbacks - logger: Final = _RecordingLogger() litellm.callbacks = [logger] try: yield logger @@ -58,6 +86,16 @@ def recording_logger(): litellm.callbacks = original_callbacks +@pytest.fixture +def recording_logger(): + yield from _installed(_RecordingLogger()) + + +@pytest.fixture +def judging_logger(): + yield from _installed(_JudgingLogger()) + + async def _wait_for_shadow_successes(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None: deadline: Final = time.monotonic() + timeout while len(logger.shadow_successes()) < expected and time.monotonic() < deadline: @@ -70,6 +108,21 @@ async def _wait_for_all_successes(logger: _RecordingLogger, expected: int, timeo await asyncio.sleep(0.05) +async def _wait_for_judges(logger: _JudgingLogger, judged: set[str], timeout: float = 5.0) -> None: + deadline: Final = time.monotonic() + timeout + while not judged <= logger.judge_flags().keys() and time.monotonic() < deadline: + await asyncio.sleep(0.05) + + +def _silent_flag_by_group(logger: _RecordingLogger) -> dict[str, object]: + return { + call["standard_logging_object"]["model_group"]: call["standard_logging_object"]["metadata"].get( + "is_silent_experiment" + ) + for call in logger.success_kwargs + } + + def _wait_for_shadow_successes_sync(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None: deadline: Final = time.monotonic() + timeout while len(logger.shadow_successes()) < expected and time.monotonic() < deadline: @@ -806,15 +859,62 @@ async def test_aembedding_with_silent_model_sends_provider_body_without_it( request_body: Final = json.loads(route.calls.last.request.read()) assert request_body == {"model": "embed-model", "input": ["black dresses"], "input_type": "query"} assert response.data[0]["embedding"] == [0.1, 0.2] -@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) -def test_silent_experiment_does_not_launch_from_a_shadow_request(run_silent_experiment): - router = Router(model_list=_streaming_model_list(["shadow-a"])) - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: - run_silent_experiment( - router, - silent_model="shadow-a", - messages=[{"role": "user", "content": "hi"}], - metadata={"is_silent_experiment": True}, - ) +@pytest.mark.asyncio +async def test_a_caller_sent_silent_experiment_flag_still_mirrors_the_request_async( + recording_logger: _RecordingLogger, +) -> None: + router: Final = Router(model_list=_streaming_model_list(["shadow-a"])) + await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="pong", + metadata={"is_silent_experiment": True}, + ) + await _wait_for_all_successes(recording_logger, expected=2) - acompletion.assert_not_awaited() + assert _silent_flag_by_group(recording_logger) == {"primary-model": None, "shadow-a": True} + + +def test_a_caller_sent_silent_experiment_flag_still_mirrors_the_request_sync( + recording_logger: _RecordingLogger, +) -> None: + router: Final = Router(model_list=_streaming_model_list(["shadow-a"])) + router.completion( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="pong", + metadata={"is_silent_experiment": True}, + ) + deadline: Final = time.monotonic() + 5.0 + while "shadow-a" not in _silent_flag_by_group(recording_logger) and time.monotonic() < deadline: + time.sleep(0.05) + + assert _silent_flag_by_group(recording_logger) == {"shadow-a": True} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("shadow_outcome", ["pong", "litellm.InternalServerError"]) +async def test_calls_a_callback_starts_while_logging_a_mirrored_call_are_not_marked_as_mirrored( + judging_logger: _JudgingLogger, shadow_outcome: str +) -> 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": shadow_outcome}, + }, + ] + ) + await router.acompletion(model="primary-model", messages=[{"role": "user", "content": "hi"}]) + await _wait_for_judges(judging_logger, {"primary-model", "shadow-a"}) + + assert judging_logger.judge_flags() == {"primary-model": None, "shadow-a": None}