tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 10:24:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #26719 from BerriAI/litellm_fix-bedrock-stream-interrupt-spend-da73

Browse Source 
fix(passthrough): track spend for interrupted Bedrock streams
This commit is contained in:
Sameer Kankute
2026-04-30 09:09:08 +05:30
committed by GitHub
parent c0211d7d1b 793a35dfe2
commit d3891e6eae
5 changed files with 475 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/passthrough/main.py
+53 -21
View File
@@ -21,6 +21,7 @@ import httpx
from httpx._types import CookieTypes, QueryParamTypes, RequestFiles
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
@@ -390,19 +391,28 @@ def _sync_streaming(
):
from litellm.utils import executor
raw_bytes: List[bytes] = []
flush_scheduled = False
try:
raw_bytes: List[bytes] = []
for chunk in response.iter_bytes(): # type: ignore
raw_bytes.append(chunk)
yield chunk
executor.submit(
litellm_logging_obj.flush_passthrough_collected_chunks,
raw_bytes=raw_bytes,
provider_config=provider_config,
)
except Exception as e:
raise e
finally:
if not flush_scheduled and raw_bytes:
flush_scheduled = True
try:
executor.submit(
litellm_logging_obj.flush_passthrough_collected_chunks,
raw_bytes=raw_bytes,
provider_config=provider_config,
)
except Exception as e:
verbose_logger.exception(
"Failed to schedule passthrough spend-tracking flush "
"in _sync_streaming; %d buffered chunks dropped: %s",
len(raw_bytes),
e,
)
async def _async_streaming(
@@ -411,23 +421,45 @@ async def _async_streaming(
provider_config: "BasePassthroughConfig",
):
iter_response = await response
try:
iter_response.raise_for_status()
raw_bytes: List[bytes] = []
async for chunk in iter_response.aiter_bytes(): # type: ignore
raw_bytes.append(chunk)
yield chunk
asyncio.create_task(
litellm_logging_obj.async_flush_passthrough_collected_chunks(
raw_bytes=raw_bytes,
provider_config=provider_config,
)
)
except Exception:
try:
await iter_response.aclose()
except Exception:
pass
raise
raw_bytes: List[bytes] = []
flush_scheduled = False
try:
async for chunk in iter_response.aiter_bytes(): # type: ignore
raw_bytes.append(chunk)
yield chunk
except Exception:
try:
await iter_response.aclose()
except Exception:
pass
raise
finally:
# GeneratorExit (raised on client disconnect) is not caught by
# `except Exception`; the finally block ensures partial usage
# still gets flushed for spend tracking. See LIT-2642.
if not flush_scheduled and raw_bytes:
flush_scheduled = True
try:
asyncio.create_task(
litellm_logging_obj.async_flush_passthrough_collected_chunks(
raw_bytes=raw_bytes,
provider_config=provider_config,
)
)
except Exception as e:
verbose_logger.exception(
"Failed to schedule passthrough spend-tracking flush "
"in _async_streaming; %d buffered chunks dropped: %s",
len(raw_bytes),
e,
)
litellm/proxy/pass_through_endpoints/streaming_handler.py
+32 -31
View File
@@ -36,21 +36,16 @@ class PassThroughStreamingHandler:
passthrough_success_handler_obj: PassThroughEndpointLogging,
url_route: str,
):
"""
- Yields chunks from the response
- Collect non-empty chunks for post-processing (logging)
- Inject cost into chunks if include_cost_in_streaming_usage is enabled
"""
try:
raw_bytes: List[bytes] = []
# Extract model name for cost injection
model_name = PassThroughStreamingHandler._extract_model_for_cost_injection(
request_body=request_body,
url_route=url_route,
endpoint_type=endpoint_type,
litellm_logging_obj=litellm_logging_obj,
)
raw_bytes: List[bytes] = []
logging_scheduled = False
model_name = PassThroughStreamingHandler._extract_model_for_cost_injection(
request_body=request_body,
url_route=url_route,
endpoint_type=endpoint_type,
litellm_logging_obj=litellm_logging_obj,
)
try:
async for chunk in response.aiter_bytes():
raw_bytes.append(chunk)
if (
@@ -58,7 +53,6 @@ class PassThroughStreamingHandler:
and model_name
):
if endpoint_type == EndpointType.VERTEX_AI:
# Only handle streamRawPredict (uses Anthropic format)
if "streamRawPredict" in url_route or "rawPredict" in url_route:
modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
chunk, model_name
@@ -73,25 +67,32 @@ class PassThroughStreamingHandler:
chunk = modified_chunk
yield chunk
# After all chunks are processed, handle post-processing
end_time = datetime.now()
asyncio.create_task(
PassThroughStreamingHandler._route_streaming_logging_to_handler(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body or {},
endpoint_type=endpoint_type,
start_time=start_time,
raw_bytes=raw_bytes,
end_time=end_time,
)
)
except Exception as e:
verbose_proxy_logger.error(f"Error in chunk_processor: {str(e)}")
raise
finally:
# GeneratorExit (raised on client disconnect) is not caught by
# `except Exception`; the finally block ensures partial usage
# still gets logged for spend tracking. See LIT-2642.
if not logging_scheduled and raw_bytes:
logging_scheduled = True
try:
asyncio.create_task(
PassThroughStreamingHandler._route_streaming_logging_to_handler(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body or {},
endpoint_type=endpoint_type,
start_time=start_time,
raw_bytes=raw_bytes,
end_time=datetime.now(),
)
)
except Exception as e:
verbose_proxy_logger.error(
f"Error scheduling chunk_processor logging: {str(e)}"
)
@staticmethod
async def _route_streaming_logging_to_handler(
tests/otel_tests/test_prometheus.py
+26 -14
View File
@@ -115,6 +115,13 @@ async def test_proxy_failure_metrics():
"litellm_llm_api_failed_requests_metric_total{", # Deprecated but may still be used
]
# Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for
# hash_token(master_key) so the master key (or its hash) never
# propagates into metrics. See PR #26484.
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS
# Check if either pattern is in metrics and contains required fields
found_metric = False
for pattern in expected_patterns:
@@ -125,8 +132,7 @@ async def test_proxy_failure_metrics():
'api_key_alias="None"' in line
and 'exception_class="Openai.RateLimitError"' in line
and 'exception_status="429"' in line
and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
in line
and f'hashed_api_key="{expected_hashed_api_key}"' in line
and 'requested_model="fake-azure-endpoint"' in line
and 'route="/chat/completions"' in line
):
@@ -135,8 +141,7 @@ async def test_proxy_failure_metrics():
# For deprecated llm_api metric, check llm-specific fields
elif "litellm_llm_api_failed_requests_metric_total{" in line:
if (
'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
in line
f'hashed_api_key="{expected_hashed_api_key}"' in line
and 'model="429"' in line
): # The deprecated metric uses the actual model from the request
found_metric = True
@@ -156,8 +161,7 @@ async def test_proxy_failure_metrics():
for line in metrics.split("\n"):
if (
total_requests_pattern in line
and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
in line
and f'hashed_api_key="{expected_hashed_api_key}"' in line
and 'requested_model="fake-azure-endpoint"' in line
and 'status_code="429"' in line
):
@@ -195,6 +199,12 @@ async def test_proxy_success_metrics():
assert END_USER_ID not in metrics
# Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for
# hash_token(master_key) (PR #26484).
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS
# Check if the success metric is present and correct - use flexible matching
# Check for request_total_latency_metric with required fields
# Note: The model can be "gpt-3.5-turbo-0301" or similar depending on what's returned
@@ -203,8 +213,7 @@ async def test_proxy_success_metrics():
if (
"litellm_request_total_latency_metric_bucket{" in line
and 'api_key_alias="None"' in line
and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
in line
and f'hashed_api_key="{expected_hashed_api_key}"' in line
and 'requested_model="fake-openai-endpoint"' in line
and 'le="0.005"' in line
):
@@ -221,8 +230,7 @@ async def test_proxy_success_metrics():
if (
"litellm_llm_api_latency_metric_bucket{" in line
and 'api_key_alias="None"' in line
and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
in line
and f'hashed_api_key="{expected_hashed_api_key}"' in line
and 'requested_model="fake-openai-endpoint"' in line
and 'le="0.005"' in line
):
@@ -298,6 +306,12 @@ async def test_proxy_fallback_metrics():
print("/metrics", metrics)
# Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for
# hash_token(master_key) (PR #26484).
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS
# Check if successful fallback metric is incremented - use flexible matching
found_successful_fallback = False
for line in metrics.split("\n"):
@@ -307,8 +321,7 @@ async def test_proxy_fallback_metrics():
and 'exception_class="Openai.RateLimitError"' in line
and 'exception_status="429"' in line
and 'fallback_model="fake-openai-endpoint"' in line
and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
in line
and f'hashed_api_key="{expected_hashed_api_key}"' in line
and 'requested_model="fake-azure-endpoint"' in line
and "1.0" in line
):
@@ -328,8 +341,7 @@ async def test_proxy_fallback_metrics():
and 'exception_class="Openai.RateLimitError"' in line
and 'exception_status="429"' in line
and 'fallback_model="unknown-model"' in line
and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
in line
and f'hashed_api_key="{expected_hashed_api_key}"' in line
and 'requested_model="fake-azure-endpoint"' in line
and "1.0" in line
):
tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py
+244
View File
@@ -0,0 +1,244 @@
"""Regression tests for LIT-2642 — interrupted streams must still flush usage."""
import asyncio
from typing import List
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
def _make_streaming_response(chunks: List[bytes]):
mock = MagicMock(spec=httpx.Response)
mock.status_code = 200
mock.headers = httpx.Headers({"content-type": "application/vnd.amazon.eventstream"})
mock.raise_for_status = MagicMock(return_value=None)
async def _aiter_bytes():
for chunk in chunks:
yield chunk
mock.aiter_bytes = _aiter_bytes
mock.aclose = AsyncMock()
return mock
def _make_logging_obj():
mock = MagicMock()
mock.async_flush_passthrough_collected_chunks = AsyncMock()
return mock
class _ImmediateExecutor:
def submit(self, fn, *args, **kwargs):
fn(*args, **kwargs)
@pytest.mark.asyncio
async def test_async_streaming_flushes_on_normal_completion():
from litellm.passthrough.main import _async_streaming
chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
mock_response = _make_streaming_response(chunks)
async def response_coro():
return mock_response
mock_logging_obj = _make_logging_obj()
provider_config = MagicMock()
received = []
async for chunk in _async_streaming(
response=response_coro(),
litellm_logging_obj=mock_logging_obj,
provider_config=provider_config,
):
received.append(chunk)
assert received == chunks
await asyncio.sleep(0)
mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
call_kwargs = (
mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs
)
assert call_kwargs["raw_bytes"] == chunks
assert call_kwargs["provider_config"] is provider_config
@pytest.mark.asyncio
async def test_async_streaming_flushes_on_client_disconnect():
from litellm.passthrough.main import _async_streaming
chunks = [
b'{"chunk": 1, "outputTokens": 10}',
b'{"chunk": 2, "outputTokens": 12}',
b'{"chunk": 3, "outputTokens": 8}',
]
mock_response = _make_streaming_response(chunks)
async def response_coro():
return mock_response
mock_logging_obj = _make_logging_obj()
provider_config = MagicMock()
gen = _async_streaming(
response=response_coro(),
litellm_logging_obj=mock_logging_obj,
provider_config=provider_config,
)
received = [await gen.__anext__()]
await gen.aclose()
assert received == [chunks[0]]
await asyncio.sleep(0)
mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
call_kwargs = (
mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs
)
assert call_kwargs["raw_bytes"] == [chunks[0]]
@pytest.mark.asyncio
async def test_async_streaming_does_not_flush_on_4xx():
from litellm.passthrough.main import _async_streaming
err_response = MagicMock(spec=httpx.Response)
err_response.status_code = 429
def _raise():
raise httpx.HTTPStatusError(
"429",
request=httpx.Request("POST", "https://example.com"),
response=httpx.Response(
429, request=httpx.Request("POST", "https://example.com")
),
)
err_response.raise_for_status = _raise
err_response.aclose = AsyncMock()
async def response_coro():
return err_response
mock_logging_obj = _make_logging_obj()
with pytest.raises(httpx.HTTPStatusError):
async for _ in _async_streaming(
response=response_coro(),
litellm_logging_obj=mock_logging_obj,
provider_config=MagicMock(),
):
pass
mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called()
@pytest.mark.asyncio
async def test_async_streaming_flushes_on_upstream_exception_with_partial_data():
from litellm.passthrough.main import _async_streaming
partial_chunks = [b"partial-chunk-1", b"partial-chunk-2"]
mock_response = MagicMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.raise_for_status = MagicMock(return_value=None)
mock_response.aclose = AsyncMock()
async def _aiter_bytes_then_raise():
for c in partial_chunks:
yield c
raise httpx.ReadError("upstream disconnected")
mock_response.aiter_bytes = _aiter_bytes_then_raise
async def response_coro():
return mock_response
mock_logging_obj = _make_logging_obj()
provider_config = MagicMock()
received = []
with pytest.raises(httpx.ReadError):
async for chunk in _async_streaming(
response=response_coro(),
litellm_logging_obj=mock_logging_obj,
provider_config=provider_config,
):
received.append(chunk)
assert received == partial_chunks
await asyncio.sleep(0)
mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
call_kwargs = (
mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs
)
assert call_kwargs["raw_bytes"] == partial_chunks
def test_sync_streaming_flushes_on_normal_completion():
from litellm.passthrough.main import _sync_streaming
chunks = [b"a", b"b", b"c"]
mock_response = MagicMock(spec=httpx.Response)
def _iter_bytes():
yield from chunks
mock_response.iter_bytes = _iter_bytes
mock_logging_obj = MagicMock()
mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
provider_config = MagicMock()
with patch("litellm.utils.executor", _ImmediateExecutor()):
received = list(
_sync_streaming(
response=mock_response,
litellm_logging_obj=mock_logging_obj,
provider_config=provider_config,
)
)
assert received == chunks
mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once()
def test_sync_streaming_flushes_on_early_close():
from litellm.passthrough.main import _sync_streaming
chunks = [b"first", b"second", b"third"]
mock_response = MagicMock(spec=httpx.Response)
def _iter_bytes():
yield from chunks
mock_response.iter_bytes = _iter_bytes
mock_logging_obj = MagicMock()
mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
provider_config = MagicMock()
with patch("litellm.utils.executor", _ImmediateExecutor()):
gen = _sync_streaming(
response=mock_response,
litellm_logging_obj=mock_logging_obj,
provider_config=provider_config,
)
first = next(gen)
gen.close()
assert first == chunks[0]
mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once()
call_kwargs = mock_logging_obj.flush_passthrough_collected_chunks.call_args.kwargs
assert call_kwargs["raw_bytes"] == [chunks[0]]
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
+120
View File
@@ -0,0 +1,120 @@
"""Regression tests for LIT-2642 — interrupted pass-through streams must still log usage."""
import asyncio
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from litellm.proxy.pass_through_endpoints.streaming_handler import (
PassThroughStreamingHandler,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
def _make_streaming_response(chunks):
mock = MagicMock(spec=httpx.Response)
mock.status_code = 200
async def _aiter_bytes():
for c in chunks:
yield c
mock.aiter_bytes = _aiter_bytes
return mock
@pytest.mark.asyncio
async def test_chunk_processor_logs_on_normal_completion():
chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
response = _make_streaming_response(chunks)
mock_logging_obj = MagicMock()
mock_passthrough_handler = MagicMock()
with patch.object(
PassThroughStreamingHandler,
"_route_streaming_logging_to_handler",
new=AsyncMock(),
) as mock_route:
received = []
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=mock_logging_obj,
endpoint_type=EndpointType.GENERIC,
start_time=datetime.now(),
passthrough_success_handler_obj=mock_passthrough_handler,
url_route="/bedrock/model/claude/invoke-with-response-stream",
):
received.append(chunk)
await asyncio.sleep(0)
assert received == chunks
mock_route.assert_called_once()
call_kwargs = mock_route.call_args.kwargs
assert call_kwargs["raw_bytes"] == chunks
@pytest.mark.asyncio
async def test_chunk_processor_logs_on_client_disconnect():
chunks = [b"event-1", b"event-2", b"event-3"]
response = _make_streaming_response(chunks)
mock_logging_obj = MagicMock()
mock_passthrough_handler = MagicMock()
with patch.object(
PassThroughStreamingHandler,
"_route_streaming_logging_to_handler",
new=AsyncMock(),
) as mock_route:
gen = PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=mock_logging_obj,
endpoint_type=EndpointType.GENERIC,
start_time=datetime.now(),
passthrough_success_handler_obj=mock_passthrough_handler,
url_route="/bedrock/model/claude/invoke-with-response-stream",
)
first = await gen.__anext__()
await gen.aclose()
await asyncio.sleep(0)
assert first == chunks[0]
mock_route.assert_called_once()
call_kwargs = mock_route.call_args.kwargs
assert call_kwargs["raw_bytes"] == [chunks[0]]
@pytest.mark.asyncio
async def test_chunk_processor_does_not_schedule_logging_when_no_chunks():
response = _make_streaming_response([])
mock_logging_obj = MagicMock()
mock_passthrough_handler = MagicMock()
with patch.object(
PassThroughStreamingHandler,
"_route_streaming_logging_to_handler",
new=AsyncMock(),
) as mock_route:
received = []
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=mock_logging_obj,
endpoint_type=EndpointType.GENERIC,
start_time=datetime.now(),
passthrough_success_handler_obj=mock_passthrough_handler,
url_route="/bedrock/model/claude/invoke-with-response-stream",
):
received.append(chunk)
assert received == []
mock_route.assert_not_called()