diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 505b5b09433..93481ba4493 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -6,6 +6,7 @@ from collections.abc import Coroutine, Mapping from typing import Final import litellm +from litellm._internal_context import is_internal_call from litellm.responses.additional_tools import hoist_additional_tools from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, @@ -113,9 +114,13 @@ class LiteLLMCompletionTransformationHandler: acompletion_args.update(litellm_completion_request) acompletion_args["_skip_responses_api_bridge"] = True - litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( - **acompletion_args, - ) + token: Final = is_internal_call.set(True) + try: + litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( + **acompletion_args + ) + finally: + is_internal_call.reset(token) if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( diff --git a/tests/unit/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py index bb374b90f4e..33c52d04289 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_handler.py @@ -21,6 +21,7 @@ import httpx import pytest import litellm +from litellm._internal_context import is_internal_call from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, @@ -55,12 +56,13 @@ def test_sync_fallback_tags_skip_responses_api_bridge(): @pytest.mark.asyncio -async def test_async_fallback_tags_skip_responses_api_bridge(): +async def test_async_fallback_tags_internal_call_and_skip_responses_api_bridge(): handler = LiteLLMCompletionTransformationHandler() captured: dict = {} async def fake_acompletion(**kwargs): captured.update(kwargs) + captured["is_internal_call"] = is_internal_call.get() raise _StopForwarding() with patch("litellm.acompletion", fake_acompletion): @@ -75,6 +77,8 @@ async def test_async_fallback_tags_skip_responses_api_bridge(): await coro assert captured.get("_skip_responses_api_bridge") is True + assert captured["is_internal_call"] is True + assert is_internal_call.get() is False _CODEX_ADDITIONAL_TOOLS_ITEM = {