diff --git a/litellm/files/main.py b/litellm/files/main.py index 8716cf3b0f..13abdfd92f 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -36,10 +36,8 @@ FileContentProvider = Literal[ import litellm from litellm import get_secret_str -from litellm.files.streaming import ( - FileContentStreamingResponse, - FileContentStreamingResult, -) +from litellm.files.streaming import FileContentStreamingResponse +from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.common_utils import get_azure_credentials diff --git a/litellm/files/streaming.py b/litellm/files/streaming.py index ab149cdb4c..dab190b9b3 100644 --- a/litellm/files/streaming.py +++ b/litellm/files/streaming.py @@ -1,23 +1,20 @@ import datetime import traceback -from typing import AsyncIterator, Dict, Iterator, Literal, NamedTuple, Optional, Union, cast +from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, Literal, Optional, Union, cast -from litellm.litellm_core_utils.litellm_logging import ( - Logging as LiteLLMLoggingObj, - get_standard_logging_object_payload, -) -from litellm.types.utils import StandardLoggingHiddenParams, StandardLoggingPayload +import anyio + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.types.utils import StandardLoggingHiddenParams, StandardLoggingPayload FileContentProvider = Literal[ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus" ] -class FileContentStreamingResult(NamedTuple): - stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]] - headers: Dict[str, str] - - class FileContentStreamingResponse: """ Iterator wrapper for file content streaming that carries LiteLLM metadata @@ -30,18 +27,17 @@ class FileContentStreamingResponse: file_id: str, model: Optional[str], custom_llm_provider: Optional[Union[FileContentProvider, str]], - logging_obj: Optional[LiteLLMLoggingObj], + logging_obj: Optional["LiteLLMLoggingObj"], ) -> None: self.stream_iterator = stream_iterator self.file_id = file_id self.model = model self.custom_llm_provider = custom_llm_provider self.logging_obj = logging_obj - self.standard_logging_object: Optional[StandardLoggingPayload] = None - self._hidden_params: StandardLoggingHiddenParams = cast( - StandardLoggingHiddenParams, {} - ) + self.standard_logging_object: Optional["StandardLoggingPayload"] = None + self._hidden_params: Dict[str, Any] = {} self._logging_completed = False + self._close_completed = False self._start_time = ( logging_obj.start_time if logging_obj is not None and getattr(logging_obj, "start_time", None) @@ -84,6 +80,37 @@ class FileContentStreamingResponse: await self._log_failure_async(e) raise + async def aclose(self) -> None: + if self._close_completed: + return + + self._close_completed = True + self._logging_completed = True + stream_to_close = self.stream_iterator + self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(())) + + # Shield cleanup from request cancellation so upstream HTTP connections + # are released promptly on client disconnects. + with anyio.CancelScope(shield=True): + if hasattr(stream_to_close, "aclose"): + await cast(AsyncIterator[bytes], stream_to_close).aclose() # type: ignore[attr-defined] + elif hasattr(stream_to_close, "close"): + result = cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined] + if result is not None: + await result + + def close(self) -> None: + if self._close_completed: + return + + self._close_completed = True + self._logging_completed = True + stream_to_close = self.stream_iterator + self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(())) + + if hasattr(stream_to_close, "close"): + cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined] + def _build_logging_response(self) -> Dict[str, str]: response = { "id": self.file_id, @@ -112,13 +139,17 @@ class FileContentStreamingResponse: def _build_standard_logging_object( self, end_time: datetime.datetime, - ) -> Optional[StandardLoggingPayload]: + ) -> Optional["StandardLoggingPayload"]: if self.standard_logging_object is not None: return self.standard_logging_object if self.logging_obj is None: return None + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + self._sync_hidden_params() payload = get_standard_logging_object_payload( kwargs=self.logging_obj.model_call_details, @@ -132,11 +163,9 @@ class FileContentStreamingResponse: return None merged_hidden_params = cast( - StandardLoggingHiddenParams, + "StandardLoggingHiddenParams", { - **cast( - StandardLoggingHiddenParams, payload.get("hidden_params") or {} - ), + **cast(Dict[str, Any], payload.get("hidden_params") or {}), **self._hidden_params, }, ) diff --git a/litellm/files/types.py b/litellm/files/types.py new file mode 100644 index 0000000000..2357bd997f --- /dev/null +++ b/litellm/files/types.py @@ -0,0 +1,6 @@ +from typing import AsyncIterator, Dict, Iterator, NamedTuple, Union + + +class FileContentStreamingResult(NamedTuple): + stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]] + headers: Dict[str, str] diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 086ee84852..76f300ec9c 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -32,7 +32,7 @@ import litellm from litellm import LlmProviders from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RETRIES -from litellm.files.streaming import FileContentStreamingResult +from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index acb2ffc3fe..468e0baaa4 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -93,6 +93,9 @@ async def _stream_file_content_with_logging( request_data=data, ) raise + finally: + if hasattr(stream_iterator, "aclose"): + await stream_iterator.aclose() # type: ignore[attr-defined] async def _get_streaming_file_content_response( diff --git a/litellm/utils.py b/litellm/utils.py index f0f0e231c1..d8e14de637 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -243,8 +243,6 @@ from typing import ( get_args, ) -from openai import OpenAIError as OriginalError - # These are lazy loaded via __getattr__ from litellm.llms.base_llm.base_utils import ( BaseLLMModelInfo, diff --git a/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py b/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py index fcc0d9a97b..66c6964021 100644 --- a/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py +++ b/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py @@ -2,7 +2,8 @@ import pytest from typing import AsyncIterator, cast from litellm.files import main as files_main -from litellm.files.streaming import FileContentStreamingResult +from litellm.files.streaming import FileContentStreamingResponse +from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -109,3 +110,30 @@ async def test_afile_content_streaming_builds_standard_logging_object_on_complet captured_standard_logging_object["hidden_params"]["api_base"] == "https://api.openai.com/v1" ) + + +@pytest.mark.asyncio +async def test_file_content_streaming_response_aclose_closes_underlying_async_generator(): + close_called = False + + async def _mock_stream(): + nonlocal close_called + try: + yield b"hello" + yield b"world" + finally: + close_called = True + + stream = FileContentStreamingResponse( + stream_iterator=_mock_stream(), + file_id="file-abc123", + model="gpt-4o", + custom_llm_provider="openai", + logging_obj=None, + ) + + assert await stream.__anext__() == b"hello" + + await stream.aclose() + + assert close_called is True diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 31a6b37437..3736380124 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1,7 +1,7 @@ import json import os import sys -from unittest.mock import ANY +from unittest.mock import ANY, AsyncMock import pytest import respx @@ -14,7 +14,10 @@ sys.path.insert( import litellm from litellm import Router -from litellm.files.streaming import FileContentStreamingResult +from litellm.files.types import FileContentStreamingResult +from litellm.proxy.openai_files_endpoints.files_endpoints import ( + _stream_file_content_with_logging, +) from litellm.proxy._types import LiteLLM_UserTableFiltered, UserAPIKeyAuth from litellm.proxy.hooks import get_proxy_hook from litellm.proxy.management_endpoints.internal_user_endpoints import ui_view_users @@ -78,6 +81,40 @@ def setup_proxy_logging_object(monkeypatch, llm_router: Router) -> ProxyLogging: return proxy_logging_object +@pytest.mark.asyncio +async def test_stream_file_content_with_logging_closes_inner_iterator_on_early_exit(): + class MockStreamIterator: + def __init__(self) -> None: + self._chunks = iter([b"hello", b"world"]) + self.aclose = AsyncMock() + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._chunks) + except StopIteration: + raise StopAsyncIteration + + stream_iterator = MockStreamIterator() + proxy_logging_obj = AsyncMock() + + generator = _stream_file_content_with_logging( + stream_iterator=stream_iterator, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=AsyncMock(), + data={"litellm_call_id": "call-123"}, + ) + + assert await generator.__anext__() == b"hello" + + await generator.aclose() + + stream_iterator.aclose.assert_awaited_once() + proxy_logging_obj.update_request_status.assert_not_called() + + def test_invalid_purpose(mocker: MockerFixture, monkeypatch, llm_router: Router): """ Asserts 'create_file' is called with the correct arguments