Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py: 48%
46 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1from collections.abc import AsyncIterator
2from typing import TYPE_CHECKING, Any, Final, cast
4from fastapi.responses import StreamingResponse
6import litellm
7from litellm.files.types import FileContentProvider, FileContentStreamingResult
8from litellm.types.utils import FILE_CONTENT_STREAMING_PROVIDERS
10if TYPE_CHECKING: 10 ↛ 11line 10 didn't jump to line 11 because the condition on line 10 was never true
11 from litellm.proxy._types import UserAPIKeyAuth
12 from litellm.proxy.utils import ProxyLogging
15class FileContentStreamingHandler:
16 @staticmethod
17 def resolve_streaming_request_params(
18 *,
19 custom_llm_provider: str,
20 file_id: str,
21 data: dict[str, Any],
22 should_route: bool,
23 original_file_id: str | None,
24 credentials: dict[str, Any] | None,
25 ) -> tuple[str, str, dict[str, Any]]:
26 """
27 Resolve the provider, file ID, and request payload to use for streaming.
29 For model-routed requests, this derives the effective provider from
30 credentials, applies `prepare_data_with_credentials()` to a copied
31 payload, swaps in the decoded/original file ID, and removes `model`
32 so `afile_content()` does not re-resolve the provider. This helper
33 does not mutate the passed-in `data` dictionary. Non-routed requests
34 return the original provider, file ID, and data unchanged.
35 """
36 if should_route and credentials is not None: 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true
37 from litellm.proxy.openai_files_endpoints.common_utils import (
38 prepare_data_with_credentials,
39 )
41 resolved_streaming_data = dict(data)
42 prepare_data_with_credentials(
43 data=resolved_streaming_data,
44 credentials=credentials,
45 file_id=original_file_id,
46 include_internal_credentials=True,
47 )
48 resolved_streaming_data.pop("model", None)
49 resolved_streaming_provider: Final = cast(str, credentials["custom_llm_provider"])
50 resolved_custom_llm_provider = resolved_streaming_provider
51 resolved_file_id = cast(str, resolved_streaming_data["file_id"])
52 else:
53 resolved_streaming_data = data
54 resolved_custom_llm_provider = custom_llm_provider
55 resolved_file_id = file_id
57 return (
58 resolved_custom_llm_provider,
59 resolved_file_id,
60 resolved_streaming_data,
61 )
63 @staticmethod
64 def should_stream_file_content(
65 *,
66 custom_llm_provider: str,
67 ) -> bool:
68 return custom_llm_provider in FILE_CONTENT_STREAMING_PROVIDERS
70 @staticmethod
71 async def stream_file_content_with_logging(
72 stream_iterator: AsyncIterator[bytes],
73 proxy_logging_obj: "ProxyLogging",
74 user_api_key_dict: "UserAPIKeyAuth",
75 data: dict[str, Any],
76 ):
77 try:
78 async for chunk in stream_iterator:
79 yield chunk
80 await proxy_logging_obj.update_request_status(
81 litellm_call_id=data.get("litellm_call_id", ""), status="success"
82 )
83 except Exception as e:
84 await proxy_logging_obj.post_call_failure_hook(
85 user_api_key_dict=user_api_key_dict,
86 original_exception=e,
87 request_data=data,
88 )
89 raise
90 finally:
91 if hasattr(stream_iterator, "aclose"):
92 await stream_iterator.aclose()
94 @staticmethod
95 async def get_streaming_file_content_response(
96 *,
97 custom_llm_provider: str,
98 file_id: str,
99 data: dict[str, Any],
100 proxy_logging_obj: "ProxyLogging",
101 user_api_key_dict: "UserAPIKeyAuth",
102 version: str,
103 ) -> StreamingResponse:
104 from litellm.proxy.common_request_processing import (
105 ProxyBaseLLMRequestProcessing,
106 )
108 stream_result: Final = cast(
109 FileContentStreamingResult,
110 await litellm.afile_content(
111 **{
112 "custom_llm_provider": cast(FileContentProvider, custom_llm_provider),
113 "file_id": file_id,
114 "stream": True,
115 **data,
116 }
117 ),
118 )
120 stream_iterator: Final = cast(
121 AsyncIterator[bytes],
122 stream_result.stream_iterator,
123 )
124 hidden_params: Final = getattr(stream_iterator, "_hidden_params", {}) or {}
125 response_headers: Final = {
126 **stream_result.headers,
127 **ProxyBaseLLMRequestProcessing.get_custom_headers(
128 user_api_key_dict=user_api_key_dict,
129 model_id=hidden_params.get("model_id", "") or "",
130 cache_key=hidden_params.get("cache_key", "") or "",
131 api_base=hidden_params.get("api_base", "") or "",
132 version=version,
133 model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
134 ),
135 }
137 return StreamingResponse(
138 FileContentStreamingHandler.stream_file_content_with_logging(
139 stream_iterator=stream_iterator,
140 proxy_logging_obj=proxy_logging_obj,
141 user_api_key_dict=user_api_key_dict,
142 data=data,
143 ),
144 media_type="application/octet-stream",
145 headers=response_headers,
146 )