diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index 3699c893816..c2a1f071f3a 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -38,15 +38,24 @@ _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", de _resent_on: Final[ContextVar[Callable[[Exception], bool] | None]] = ContextVar("resent_on", default=None) +_logged_by_parent: Final[ContextVar[bool]] = ContextVar("logged_by_parent", default=False) + @contextmanager def internal_sub_call() -> Generator[None]: """A nested call the parent call logs and bills, so the parent must fold this call's cost into its own.""" - token: Final = is_internal_call.set(True) + internal_token: Final = is_internal_call.set(True) + logged_token: Final = _logged_by_parent.set(True) try: yield finally: - is_internal_call.reset(token) + _logged_by_parent.reset(logged_token) + is_internal_call.reset(internal_token) + + +def logged_by_parent() -> bool: + """Set only by ``internal_sub_call``. Other internal calls still log their own cache hits and sync callbacks.""" + return _logged_by_parent.get() @contextmanager diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 063faa70f2f..2088efabda5 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError import litellm -from litellm._internal_context import is_internal_call, post_response_phase +from litellm._internal_context import logged_by_parent, post_response_phase from litellm._logging import print_verbose, verbose_logger from litellm.caching import InMemoryCache from litellm.caching.caching import S3Cache, response_cache_phase @@ -511,7 +511,7 @@ class LLMCachingHandler: custom_llm_provider=custom_llm_provider, ) - if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks( + if not logged_by_parent() and not _should_defer_streaming_cache_hit_callbacks( cached_result=cached_result ): logging_obj.handle_sync_cache_hit_success_callbacks( @@ -823,9 +823,9 @@ class LLMCachingHandler: end_time (datetime): The end time of the operation. cache_hit (bool): Whether it was a cache hit. - An internal sub-call's hit fires nothing, since the parent call that folds it in logs it. + A hit inside ``internal_sub_call`` fires nothing, since the parent call that folds it in logs it. """ - if is_internal_call.get(): + if logged_by_parent(): return from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER diff --git a/litellm/utils.py b/litellm/utils.py index ea8726b7f69..64e1d5c8232 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -51,7 +51,7 @@ import litellm.litellm_core_utils # audio_utils.utils is lazy-loaded - only imported when needed for transcription calls import litellm.litellm_core_utils.json_validation_rule -from litellm._internal_context import is_internal_call, is_resent_failure +from litellm._internal_context import is_internal_call, is_resent_failure, logged_by_parent from litellm._lazy_imports import ( _get_modified_max_tokens, get_default_encoding, @@ -1375,16 +1375,17 @@ def _dispatch_success_logging( is_completion_with_fallbacks: bool, is_litellm_internal_call: bool, ) -> None: - if is_litellm_internal_call: + if logged_by_parent(): return - _schedule_async_success_logging( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) + if not is_litellm_internal_call: + _schedule_async_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 1ff368ae010..4b501ec1c63 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -3092,11 +3092,12 @@ async def test_async_get_cache_answers_400_for_a_single_object_embedding_input() @pytest.mark.asyncio -@pytest.mark.parametrize("internal", [False, True]) -async def test_a_cache_hit_inside_an_internal_sub_call_fires_no_callbacks_of_its_own( - monkeypatch: pytest.MonkeyPatch, internal: bool +@pytest.mark.parametrize("sub_call", ["none", "rag_sub_call", "internal_sub_call"]) +async def test_a_cache_hit_fires_no_callbacks_of_its_own_only_inside_an_internal_sub_call( + monkeypatch: pytest.MonkeyPatch, sub_call: str ) -> None: from litellm.integrations.custom_logger import CustomLogger + from litellm.rag.main import _suppressed_sub_call_billing logged: Final[list[tuple[str, object]]] = [] @@ -3130,7 +3131,12 @@ async def test_a_cache_hit_inside_an_internal_sub_call_fires_no_callbacks_of_its async_logging: Final = _build_logging_obj(CallTypes.acompletion.value, stream=False) sync_logging: Final = _build_logging_obj(CallTypes.completion.value, stream=False) - with internal_sub_call() if internal else nullcontext(): + scope: Final = { + "none": nullcontext, + "rag_sub_call": _suppressed_sub_call_billing, + "internal_sub_call": internal_sub_call, + }[sub_call] + with scope(): async_hit: Final = await LLMCachingHandler( original_function=litellm.acompletion, request_kwargs=kwargs, start_time=_FIXED_START ).async_get_cache( @@ -3155,7 +3161,7 @@ async def test_a_cache_hit_inside_an_internal_sub_call_fires_no_callbacks_of_its expected: Final = ( Counter() - if internal + if sub_call == "internal_sub_call" else Counter({("async", async_logging.litellm_call_id): 1, ("sync", sync_logging.litellm_call_id): 1}) ) for _ in range(100): diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a9e3aa7da5d..81e935e7a47 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -7844,6 +7844,88 @@ async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: assert counting_logger.logged_results[0] is inner_result +@pytest.mark.asyncio +@pytest.mark.parametrize("logged_by_parent", [False, True]) +async def test_an_internal_call_keeps_its_sync_success_callbacks_unless_its_parent_logs_it( + monkeypatch: pytest.MonkeyPatch, logged_by_parent: bool +) -> None: + from litellm._internal_context import internal_sub_call + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.rag.main import _suppressed_sub_call_billing + from litellm.utils import _dispatch_success_logging + + sync_handled: Final[list[object]] = [] + async_logged: Final[list[object]] = [] + + class RecordingLogging(Logging): + def success_handler( + self, + result: object = None, + start_time: datetime | None = None, + end_time: datetime | None = None, + cache_hit: bool | None = None, + *, + for_async_request: bool = False, + **kwargs: object, + ) -> None: + sync_handled.append(result) + + class AsyncRecorder(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + async_logged.append(response_obj) + + def function_success_callback( + kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + pass + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "success_callback", [function_success_callback]) + monkeypatch.setattr(litellm, "_async_success_callback", [AsyncRecorder()]) + messages: Final = [{"role": "user", "content": "hello"}] + logging_obj: Final = RecordingLogging( + model="gpt-5.4", + messages=messages, + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc).timestamp(), + litellm_call_id="sub-call-id", + function_id="sub-call-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={"acompletion": True}, + optional_params={}, + model="gpt-5.4", + custom_llm_provider="openai", + input=messages, + ) + result: Final = ModelResponse(id="sub-call") + expected: Final = [] if logged_by_parent else [result] + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + with internal_sub_call() if logged_by_parent else _suppressed_sub_call_billing(): + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=now, + end_time=now, + is_completion_with_fallbacks=False, + is_litellm_internal_call=True, + ) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + for _ in range(100): + if len(sync_handled) >= len(expected): + break + await asyncio.sleep(0.05) + await asyncio.sleep(0.3) + + assert async_logged == [] + assert sync_handled == expected + + def test_basic_trimming(): litellm.turn_on_debug() messages = [