diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 64508ec93a6..9520aa4f657 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -4,9 +4,6 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy.pass_through_endpoints.success_handler import ( - PassThroughEndpointLogging, -) from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType if TYPE_CHECKING: @@ -16,8 +13,6 @@ if TYPE_CHECKING: else: BaseGoogleGenAIGenerateContentConfig = Any -GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() - def _encode_google_genai_sse_event(event_lines: list[str]) -> bytes: return ("\n".join(event_lines) + "\n\n").encode("utf-8") @@ -87,6 +82,9 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) + from litellm.proxy.pass_through_endpoints.success_handler import ( + GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + ) end_time: Final = datetime.now() asyncio.create_task( diff --git a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py index f37e5190850..547dd098b55 100644 --- a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py @@ -17,15 +17,10 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP from litellm.llms.anthropic.pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE -from litellm.proxy.pass_through_endpoints.success_handler import ( - PassThroughEndpointLogging, -) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType from litellm.types.utils import GenericStreamingChunk, ModelResponseStream -GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() - _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks _DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains @@ -414,6 +409,9 @@ class BaseAnthropicMessagesStreamingIterator: from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) + from litellm.proxy.pass_through_endpoints.success_handler import ( + GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + ) end_time: Final = datetime.now() # Set completion_start_time so TTFT is calculated from the first @@ -452,6 +450,9 @@ class BaseAnthropicMessagesStreamingIterator: from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) + from litellm.proxy.pass_through_endpoints.success_handler import ( + GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + ) # Use the existing streaming handler for Anthropic return PassThroughStreamingHandler.chunk_processor( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 9c6c3111209..9362d192f1b 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -677,3 +677,6 @@ class PassThroughEndpointLogging: logging_obj.model_call_details["response_cost"] = passthrough_logging_payload.get("cost_per_request") return kwargs + + +GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() diff --git a/litellm/rust_bridge/callbacks_legacy_python.py b/litellm/rust_bridge/callbacks_legacy_python.py index bfd16b29bf6..151502b77ed 100644 --- a/litellm/rust_bridge/callbacks_legacy_python.py +++ b/litellm/rust_bridge/callbacks_legacy_python.py @@ -360,10 +360,8 @@ def stream_success( end: datetime.datetime, first_chunk: datetime.datetime | None, ) -> None: - from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( - GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, - ) from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + from litellm.proxy.pass_through_endpoints.success_handler import GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType if first_chunk is not None: diff --git a/tests/unit/google_genai/test_google_genai_streaming_iterator.py b/tests/unit/google_genai/test_google_genai_streaming_iterator.py index e8ec2848233..3363553573b 100644 --- a/tests/unit/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/unit/google_genai/test_google_genai_streaming_iterator.py @@ -1,5 +1,6 @@ +import asyncio import json -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -8,6 +9,8 @@ from litellm.google_genai.streaming_iterator import ( GoogleGenAIGenerateContentStreamingIterator, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy.pass_through_endpoints import success_handler +from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType @@ -154,3 +157,27 @@ async def test_async_streaming_iterator_forwards_sse_comment_events(): chunk = await iterator.__anext__() assert chunk == b": keepalive\n\n" + + +@pytest.mark.asyncio +async def test_async_streaming_logging_bills_through_the_proxy_pass_through_success_handler(): # test-quality-ok: no injection seam, the handler object identity passed to the router is the contract + iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( + response=MagicMock(), + model="gemini-3.1-flash-image", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + iterator.collected_chunks.append(b"data: {}\n\n") + + with patch.object(PassThroughStreamingHandler, "route_streaming_logging_to_handler", new=AsyncMock()) as route: + await iterator._handle_async_streaming_logging() + await asyncio.sleep(0) + + route.assert_awaited_once() + assert route.await_args.kwargs["passthrough_success_handler_obj"] is ( + success_handler.GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ + ) + assert route.await_args.kwargs["endpoint_type"] is EndpointType.GEMINI + assert route.await_args.kwargs["raw_bytes"] == [b"data: {}\n\n"] diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py index 1912876a453..35aa9829f2f 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py @@ -1,8 +1,9 @@ import asyncio import json from datetime import datetime -from unittest.mock import patch +from unittest.mock import AsyncMock, patch +import httpx import pytest @@ -22,6 +23,8 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( is_provider_error_chunk, parse_anthropic_error_event, ) +from litellm.proxy.pass_through_endpoints import success_handler +from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler class _RecordingLoggingIterator(BaseAnthropicMessagesStreamingIterator): @@ -1523,3 +1526,30 @@ async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_co assert len(worker.enqueued) == 1, "teardown billing enqueued alongside the deferred dispatch" await worker.enqueued[0] assert deferred_fired.is_set() + + +@pytest.mark.asyncio +async def test_async_streaming_response_iterator_relays_bytes_and_bills_through_proxy_success_handler(): + body = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' + response = httpx.Response(200, content=body) + logging_obj = _make_logging_obj("test_async_streaming_response_iterator_relays_bytes") + iterator = BaseAnthropicMessagesStreamingIterator(litellm_logging_obj=logging_obj, request_body={}) + + with patch.object(PassThroughStreamingHandler, "route_streaming_logging_to_handler", new=AsyncMock()) as route: + relayed = [ + chunk + async for chunk in iterator.get_async_streaming_response_iterator( + httpx_response=response, request_body={}, litellm_logging_obj=logging_obj + ) + ] + for _ in range(100): + if route.await_count: + break + await asyncio.sleep(0.01) + + assert b"".join(relayed) == body + route.assert_awaited_once() + assert route.await_args.kwargs["passthrough_success_handler_obj"] is ( + success_handler.GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ + ) + assert route.await_args.kwargs["url_route"] == "/v1/messages" diff --git a/tests/unit/rust_bridge/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/test_callbacks_legacy_python.py index c78fbb3f246..2a331683901 100644 --- a/tests/unit/rust_bridge/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/test_callbacks_legacy_python.py @@ -10,6 +10,7 @@ from pydantic import TypeAdapter from litellm._internal_context import is_internal_call from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy.pass_through_endpoints import success_handler from litellm.rust_bridge import callbacks_legacy_python as legacy from litellm.rust_bridge.callbacks_legacy_python import failure_handler, setup from litellm.types.utils import ModelResponse @@ -174,3 +175,29 @@ def test_the_rust_contract_matches_the_shim_signatures() -> None: contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} + + +def test_stream_success_defers_billing_through_the_proxy_pass_through_success_handler() -> None: + logger: Final = _supplied_logger() + logger._on_deferred_stream_complete = lambda: None # pyright: ignore[reportAttributeAccessIssue] # the proxy's deferred stream release sets this slot + first_chunk: Final = datetime.datetime.now() + + legacy.stream_success( + logger, + url_route="/v1/messages", + endpoint_type="anthropic", + request_body={}, + chunks=[b'event: message_stop\ndata: {"type": "message_stop"}\n\n'], + start=first_chunk, + end=datetime.datetime.now(), + first_chunk=first_chunk, + ) + + assert logger.completion_start_time == first_chunk + (deferred,) = logger._deferred_stream_complete_args # pyright: ignore[reportAttributeAccessIssue] # set by stream_success for the deferred release + try: + assert inspect.getcoroutinelocals(deferred)["passthrough_success_handler_obj"] is ( + success_handler.GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ + ) + finally: + deferred.close()