Merge a37d98fe75 into 2a01fe6e5c

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(

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(

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()

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:

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"]

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"

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()