Refactor file content streaming implementation

- Removed unused imports and streamlined type hints in `litellm/utils.py` and `litellm/files/main.py`.
- Moved `FileContentStreamingResult` to a new `litellm/files/types.py` for better organization.
- Updated `FileContentStreamingResponse` in `litellm/files/streaming.py` to include asynchronous close methods and improved logging capabilities.
- Enhanced tests to ensure proper closure of streaming iterators in `tests/test_litellm/llms/openai/test_openai_file_content_streaming.py` and `tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py`.

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

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,
             },
         )

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]

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

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(

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,

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

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