diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index d869136d8e2..6e685e21991 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -24,6 +24,7 @@ from litellm.llms.anthropic.pass_through.context_management import ( apply_context_management, ) from litellm.llms.anthropic.pass_through.convertible_blocks import reject_malformed_tool_result_text +from litellm.llms.anthropic.pass_through.held_stream_completion import hold_completion_of from litellm.llms.anthropic.pass_through.safeguards import ( SafeguardsEvaluator, build_safeguards_evaluator, @@ -814,6 +815,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: is_async=True, litellm_logging_obj=logging_obj, safeguards_evaluator=safeguards_evaluator, + held_completion=None if planned_chat is None else hold_completion_of(completion_response, logging_obj), ) if transformed_stream is not None: return transformed_stream @@ -962,6 +964,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: polyfill_result=polyfill_result, is_async=False, litellm_logging_obj=logging_obj, + held_completion=None if planned_chat is None else hold_completion_of(completion_response, logging_obj), ) if transformed_stream is not None: return transformed_stream diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 52b836c789a..cbefb2dbdbe 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -45,6 +45,7 @@ from litellm.types.utils import AdapterCompletionStreamWrapper, ChatCompletionDe if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject + from litellm.llms.anthropic.pass_through.held_stream_completion import HeldStreamCompletion from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponseStream @@ -518,6 +519,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): iterations_usage: list[UsageIteration] | None = None, litellm_logging_obj: "LiteLLMLoggingObject | None" = None, safeguards_evaluator: "SafeguardsEvaluator | None" = None, + held_completion: "HeldStreamCompletion | None" = None, ): # Wrap the upstream stream so chunks that carry both content and a # finish_reason (fake-streamed providers) are split into two — see @@ -529,6 +531,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id) self.litellm_logging_obj = litellm_logging_obj self.safeguards_evaluator: Final = safeguards_evaluator + self._held_completion: Final = held_completion # Mapping of truncated tool names to original names (for OpenAI's 64-char limit) self.tool_name_mapping = tool_name_mapping or {} # Polyfill applied_edits on final message_delta. @@ -1211,14 +1214,18 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): This wrapper ensures dict chunks are SSE formatted with both event and data lines. """ - for chunk in self: - if isinstance(chunk, dict): - event_type: str = str(chunk.get("type", "message")) - payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n" - yield payload.encode() - else: - # For non-dict chunks, forward the original value unchanged - yield chunk + try: + for chunk in self: + if isinstance(chunk, dict): + event_type: str = str(chunk.get("type", "message")) + payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n" + yield payload.encode() + else: + # For non-dict chunks, forward the original value unchanged + yield chunk + finally: + if self._held_completion is not None: + self._held_completion.log(is_async=False) async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]: """ @@ -1237,6 +1244,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): else: yield anthropic_sse_frame(await safeguard_results.observe(_SSE_EVENT.validate_python(chunk))) except Exception as e: # noqa: BLE001 # boundary before the socket: any upstream failure becomes an Anthropic error event + if self._held_completion is not None: + self._held_completion.discard() verbose_logger.exception("Anthropic Adapter - mid-stream error: %s", e) logging_obj: Final = self.litellm_logging_obj provider_error: Final = _provider_error(e) @@ -1257,6 +1266,12 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): failure_handler_error, ) yield _mid_stream_error_sse_event(e) + else: + if self._held_completion is not None: + self._held_completion.log(is_async=True) + finally: + if self._held_completion is not None: + self._held_completion.log(is_async=True, client_detached=True) def _increment_content_block_index(self): self.current_content_block_index += 1 diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 0152931ebc3..ead35d727f7 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -210,6 +210,7 @@ from .streaming_iterator import AnthropicSSEStream, AnthropicStreamWrapper if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject + from litellm.llms.anthropic.pass_through.held_stream_completion import HeldStreamCompletion from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator from litellm.types.llms.anthropic import ContentBlockContentBlockDict @@ -322,6 +323,7 @@ class AnthropicAdapter: is_async: bool = True, litellm_logging_obj: "LiteLLMLoggingObject | None" = None, safeguards_evaluator: "SafeguardsEvaluator | None" = None, + held_completion: "HeldStreamCompletion | None" = None, ) -> AsyncIterator[bytes] | Iterator[bytes] | None: """ Translate OpenAI streaming response to Anthropic format. @@ -350,6 +352,7 @@ class AnthropicAdapter: iterations_usage=iterations_usage, litellm_logging_obj=litellm_logging_obj, safeguards_evaluator=safeguards_evaluator, + held_completion=held_completion, ) # Return the SSE-wrapped version for proper event formatting. if is_async: diff --git a/litellm/llms/anthropic/pass_through/held_stream_completion.py b/litellm/llms/anthropic/pass_through/held_stream_completion.py new file mode 100644 index 00000000000..8b84fd6df3e --- /dev/null +++ b/litellm/llms/anthropic/pass_through/held_stream_completion.py @@ -0,0 +1,64 @@ +from collections.abc import Callable +from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable + +from litellm.responses.streaming_iterator import dispatch_stream_success + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +@runtime_checkable +class _HandsOffCompletion(Protocol): + def hand_completion_to(self, on_complete: Callable[[object], None]) -> None: ... + + +class _NothingHeld: + pass + + +_NOTHING_HELD: Final = _NothingHeld() + + +class HeldStreamCompletion: + """The executor stream's assembled response, logged by the client-facing Messages stream when it ends. + + A stream that ended with no content hands over ``None``; it is still logged, as the success + handler logs an empty stream when it is not held. + """ + + def __init__(self, logging_obj: "LiteLLMLoggingObj") -> None: + self._logging_obj: Final = logging_obj + self._response: object = _NOTHING_HELD + + def take(self, response: object) -> None: + self._response = response + + def _release(self) -> object: + response: Final = self._response + self._response = _NOTHING_HELD + return response + + def discard(self) -> None: + self._release() + + def log(self, *, is_async: bool, client_detached: bool = False) -> None: + response: Final = self._release() + if response is _NOTHING_HELD: + return + dispatch_stream_success( + self._logging_obj, + response, + start_time=None, + end_time=None, + cache_hit=None, + is_async=is_async, + client_detached=client_detached, + ) + + +def hold_completion_of(stream: object, logging_obj: "LiteLLMLoggingObj | None") -> HeldStreamCompletion | None: + if logging_obj is None or not isinstance(stream, _HandsOffCompletion): + return None + held: Final = HeldStreamCompletion(logging_obj) + stream.hand_completion_to(held.take) + return held diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index fed76d8bbd5..e60ead0cd95 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -17,6 +17,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import cache_control_as_prompt_cache_breakpoint +from litellm.llms.anthropic.pass_through.held_stream_completion import hold_completion_of from litellm.responses.bridge_executor import anative_responses, native_logging_obj, native_responses from litellm.responses.utils import auto_reasoning_summary from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS @@ -310,6 +311,7 @@ class LiteLLMMessagesToResponsesAPIHandler: model=local_model_name(model, kwargs.get("custom_llm_provider")), litellm_logging_obj=logging_obj, safeguards_evaluator=safeguards_evaluator, + held_completion=None if native_logging is None else hold_completion_of(result, native_logging), ) return wrapper.async_anthropic_sse_wrapper() @@ -399,6 +401,7 @@ class LiteLLMMessagesToResponsesAPIHandler: responses_stream=result, model=local_model_name(model, kwargs.get("custom_llm_provider")), litellm_logging_obj=logging_obj, + held_completion=None if native_logging is None else hold_completion_of(result, native_logging), ) return wrapper.async_anthropic_sse_wrapper() diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py index e628a30fdd6..9a80a23c6e6 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py @@ -35,6 +35,7 @@ from .transformation import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject + from litellm.llms.anthropic.pass_through.held_stream_completion import HeldStreamCompletion from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator from litellm.types.llms.openai import ResponseIncompleteEvent @@ -127,10 +128,12 @@ class AnthropicResponsesStreamWrapper: model: str, litellm_logging_obj: "LiteLLMLoggingObject | None" = None, safeguards_evaluator: "SafeguardsEvaluator | None" = None, + held_completion: "HeldStreamCompletion | None" = None, ) -> None: self.responses_stream = responses_stream self.model = model self.safeguards_evaluator: Final = safeguards_evaluator + self._held_completion: Final = held_completion self._message_id: str = f"msg_{uuid.uuid4()}" if litellm_logging_obj is not None: litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id) @@ -498,8 +501,22 @@ class AnthropicResponsesStreamWrapper: async def _sse_frames(self) -> AsyncGenerator[bytes, None]: safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator) - async for chunk in self: - yield anthropic_sse_frame(await safeguard_results.observe(chunk)) + held: Final = self._held_completion + try: + async for chunk in self: + yield anthropic_sse_frame(await safeguard_results.observe(chunk)) + except Exception: + if held is not None: + held.discard() + raise + else: + if held is not None and self._stream_failed: + held.discard() + if held is not None: + held.log(is_async=True) + finally: + if held is not None: + held.log(is_async=True, client_detached=True) def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": return partial_responses_incomplete_event_on_disconnect(self.responses_stream) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 71abef9e4b1..586b546039d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -48,6 +48,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_r from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.litellm_completion_transformation.transformation import ( @@ -363,6 +364,7 @@ def dispatch_stream_success( end_time: datetime | None, cache_hit: bool | None, is_async: bool, + client_detached: bool = False, ) -> None: if not is_async: run_async_function( @@ -387,6 +389,9 @@ def dispatch_stream_success( cache_hit=cache_hit, prefer_async_handlers=True, ) + if client_detached: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + return if getattr(logging_obj, "_on_deferred_stream_complete", None) is not None: # End-of-stream guardrail scans write guardrail_information after # the terminal event; dispatching now would snapshot metadata early. diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_chat_executor.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_chat_executor.py index 39351cbacfd..c0d77f441f8 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_chat_executor.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_chat_executor.py @@ -1,17 +1,20 @@ +import asyncio import datetime -from collections.abc import Mapping, Sequence +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import Final, TypeAlias +from typing import Final, TypeAlias, cast import httpx import pytest import respx +from pydantic import TypeAdapter import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.anthropic.pass_through.adapters.handler import LiteLLMMessagesToCompletionTransformationHandler +from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import AnthropicSSEStream from litellm.types.llms.openai import AllMessageValues from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.utils import StandardCallbackDynamicParams @@ -28,10 +31,8 @@ CHAT_RESPONSE: Final = { class _DeploymentHookCounter(CustomLogger): - def __init__(self) -> None: - super().__init__() - self.pre_call = 0 - self.post_call = 0 + pre_call: int = 0 + post_call: int = 0 async def async_pre_call_deployment_hook(self, kwargs: Mapping[str, object], call_type: object) -> None: self.pre_call += 1 @@ -142,6 +143,18 @@ async def test_bridge_maps_an_upstream_429_to_rate_limit_error( await _ask() +def _sync_bridge(logging_obj: Logging, *, stream: bool) -> object: + return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + max_tokens=10, + messages=[{"role": "user", "content": "hi"}], + model="hosted_vllm/served-model", + stream=stream, + api_base=API_BASE, + api_key="sk-chat", + litellm_logging_obj=logging_obj, + ) + + def test_sync_bridge_answers_from_one_upstream_call(upstream: respx.Route): logging_obj: Final = Logging( model="hosted_vllm/served-model", @@ -153,14 +166,7 @@ def test_sync_bridge_answers_from_one_upstream_call(upstream: respx.Route): function_id="1234", ) - response: Final = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( - max_tokens=10, - messages=[{"role": "user", "content": "hi"}], - model="hosted_vllm/served-model", - api_base=API_BASE, - api_key="sk-chat", - litellm_logging_obj=logging_obj, - ) + response: Final = _sync_bridge(logging_obj, stream=False) assert isinstance(response, dict) assert response["content"] == [{"type": "text", "text": "from upstream"}] @@ -202,3 +208,173 @@ async def test_bridge_goes_through_a_wrapped_acompletion(upstream: respx.Route, assert response["content"] == [{"type": "text", "text": "from upstream"}] assert seen == ["hosted_vllm/served-model"] assert upstream.call_count == 1 + + +class _SuccessRowRecorder(CustomLogger): + call_types: tuple[str, ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.call_types = (*self.call_types, str(kwargs.get("call_type"))) + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.call_types = (*self.call_types, str(kwargs.get("call_type"))) + + +_PARKED_COMPLETION: Final = TypeAdapter(tuple[object]) + + +async def _release(parked: Coroutine[object, object, object]) -> None: + await parked + + +async def _settled_rows(recorder: _SuccessRowRecorder) -> tuple[str, ...]: + for _ in range(200): + if recorder.call_types: + break + await asyncio.sleep(0.01) + await asyncio.sleep(0.05) + return recorder.call_types + + +STREAM_CHUNK: Final = '{"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":"streamed"},"finish_reason":null}]}' +STREAM_DONE: Final = '{"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + + +@pytest.fixture +def stream_recorder(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> _SuccessRowRecorder: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{API_BASE}/chat/completions").mock( + return_value=httpx.Response( + 200, + content=f"data: {STREAM_CHUNK}\n\ndata: {STREAM_DONE}\n\ndata: [DONE]\n\n".encode(), + headers={"content-type": "text/event-stream"}, + ) + ) + recorder: Final = _SuccessRowRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + return recorder + + +async def _read_through_message_stop(stream: AsyncIterator[bytes]) -> None: + async for frame in stream: + if frame.startswith(b"event: message_stop"): + return + + +@pytest.mark.asyncio +async def test_bridged_stream_logs_one_row_when_the_client_stream_ends(stream_recorder: _SuccessRowRecorder): + stream: Final = await _ask({"stream": True}) + assert isinstance(stream, AnthropicSSEStream) + + await _read_through_message_stop(stream) + rows_while_the_client_stream_is_open: Final = await _settled_rows(stream_recorder) + remaining: Final[list[bytes]] = [frame async for frame in stream] + + assert rows_while_the_client_stream_is_open == () + assert remaining == [] + assert await _settled_rows(stream_recorder) == ("anthropic_messages",) + + +@pytest.mark.asyncio +async def test_bridged_stream_still_logs_when_the_client_leaves_after_message_stop( + stream_recorder: _SuccessRowRecorder, +): + stream: Final = await _ask({"stream": True}) + assert isinstance(stream, AnthropicSSEStream) + + await _read_through_message_stop(stream) + await stream.aclose() + + assert await _settled_rows(stream_recorder) == ("anthropic_messages",) + + +@pytest.mark.asyncio +async def test_bridged_stream_holds_its_row_for_the_proxy_end_of_stream_guardrails( + stream_recorder: _SuccessRowRecorder, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "_async_success_callback", [stream_recorder]) + logging_obj: Final = Logging( + model="hosted_vllm/served-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime.now(datetime.timezone.utc), + litellm_call_id="6825beef-0000-4000-8000-000000000007", + function_id="1234", + ) + + async def release(*held: object) -> None: + return None + + setattr(logging_obj, "_on_deferred_stream_complete", release) # noqa: B010 # the proxy arms this slot dynamically + stream: Final = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=10, + messages=[{"role": "user", "content": "hi"}], + model="hosted_vllm/served-model", + stream=True, + api_base=API_BASE, + api_key="sk-chat", + litellm_logging_obj=logging_obj, + ) + assert isinstance(stream, AnthropicSSEStream) + frames: Final = [frame async for frame in stream] + rows_before_release: Final = await _settled_rows(stream_recorder) + held: Final[object] = getattr(logging_obj, "_deferred_stream_complete_args", None) + + assert frames[-1].startswith(b"event: message_stop") + assert rows_before_release == () + (parked,) = _PARKED_COMPLETION.validate_python(held) + assert asyncio.iscoroutine(parked) + await _release(parked) + assert stream_recorder.call_types == ("anthropic_messages",) + + +def test_sync_bridged_stream_logs_one_row(stream_recorder: _SuccessRowRecorder, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "success_callback", [stream_recorder]) + logging_obj: Final = Logging( + model="hosted_vllm/served-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime.now(datetime.timezone.utc), + litellm_call_id="6825beef-0000-4000-8000-000000000006", + function_id="1234", + ) + + stream: Final = _sync_bridge(logging_obj, stream=True) + assert isinstance(stream, Iterator) + frames: Final = list(cast(Iterator[bytes], stream)) # cast-ok: the sync bridge streams SSE frames as bytes + + assert frames[-1].startswith(b"event: message_stop") + assert asyncio.run(_settled_rows(stream_recorder)) == ("anthropic_messages",) + + +@pytest.mark.asyncio +async def test_bridged_empty_stream_still_logs_one_row(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{API_BASE}/chat/completions").mock( + return_value=httpx.Response(200, content=b"data: [DONE]\n\n", headers={"content-type": "text/event-stream"}) + ) + recorder: Final = _SuccessRowRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + + # deepinfra sends no stream_options, so the empty stream assembles no response to log + stream: Final = await litellm.anthropic.messages.acreate( + model="deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct", + api_base=API_BASE, + api_key="sk-chat", + max_tokens=10, + messages=[{"role": "user", "content": "hi"}], + stream=True, + ) + assert isinstance(stream, AnthropicSSEStream) + frames: Final = [frame async for frame in stream] + + assert frames[-1].startswith(b"event: message_stop") + assert await _settled_rows(recorder) == ("anthropic_messages",) diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py index 785e1f1096b..f11427f9e60 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py @@ -19,9 +19,9 @@ import json import os import sys import threading +from collections.abc import AsyncIterator, Callable, Iterator, Mapping from datetime import datetime -from collections.abc import Callable -from typing import List, Optional +from typing import Final from unittest.mock import MagicMock import pytest @@ -37,11 +37,12 @@ from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, _mid_stream_error_sse_event, ) +from litellm.llms.anthropic.pass_through.held_stream_completion import HeldStreamCompletion from litellm.llms.bedrock.common_utils import BedrockError -from litellm.types.utils import Delta, StreamingChoices +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices -def _make_chunk(delta: Delta, finish_reason: Optional[str] = None) -> MagicMock: +def _make_chunk(delta: Delta, finish_reason: str | None = None) -> MagicMock: chunk = MagicMock() chunk.choices = [ StreamingChoices(finish_reason=finish_reason, index=0, delta=delta, logprobs=None) @@ -55,7 +56,7 @@ class _AsyncStreamThenRaise: """Yields the given chunks, then raises ``exc`` (mimics a provider stream that terminates mid-response).""" - def __init__(self, items: List[MagicMock], exc: BaseException): + def __init__(self, items: list[MagicMock], exc: BaseException): self._it = iter(items) self._exc = exc @@ -75,7 +76,7 @@ def _parse_sse(raw: bytes) -> tuple[str, dict]: return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: ")) -async def _drain_sse(wrapper: AnthropicStreamWrapper) -> List[bytes]: +async def _drain_sse(wrapper: AnthropicStreamWrapper) -> list[bytes]: return [event async for event in wrapper.async_anthropic_sse_wrapper()] @@ -318,3 +319,51 @@ async def test_mid_stream_error_emits_error_event_without_logging_obj(): events = await _drain_sse(_failing_wrapper(None)) assert _parse_sse(events[-1])[0] == "error" + + +class _AsyncSuccessRecorder(CustomLogger): + calls: tuple[object, ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.calls = (*self.calls, response_obj) + + +class _CompletesThenRaises: + def __init__(self, exc: Exception, held: HeldStreamCompletion) -> None: + self._exc: Final = exc + self._held: Final = held + + def __iter__(self) -> Iterator[ModelResponseStream]: + return iter(()) + + def __aiter__(self) -> AsyncIterator[ModelResponseStream]: + return self._chunks() + + async def _chunks(self) -> AsyncIterator[ModelResponseStream]: + yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="partial"))]) + self._held.take(litellm.ModelResponse()) + raise self._exc + + +@pytest.mark.asyncio +async def test_a_stream_that_fails_after_the_executor_completed_logs_no_success_row(): + success_recorder: Final = _AsyncSuccessRecorder() + async_recorder: Final = _AsyncFailureRecorder() + logging_obj: Final = _make_logging_obj("fails-after-complete", async_recorder, _SyncFailureRecorder()) + logging_obj.dynamic_async_success_callbacks = [success_recorder] + held: Final = HeldStreamCompletion(logging_obj) + wrapper: Final = AnthropicStreamWrapper( + completion_stream=_CompletesThenRaises(_bedrock_drop(), held), + model="bedrock-converse-sonnet-4-6", + litellm_logging_obj=logging_obj, + held_completion=held, + ) + + events: Final = await _drain_sse(wrapper) + await asyncio.sleep(0.2) + + assert _parse_sse(events[-1])[0] == "error" + assert [str(exc) for exc in async_recorder.exceptions] == ["ConverseStream ended without messageStop"] + assert success_recorder.calls == () diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index d6a8a3e781e..226228670a7 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -1,18 +1,21 @@ +import asyncio import datetime import json import os import sys -from collections.abc import Callable, Coroutine, Mapping +from collections.abc import AsyncIterator, Callable, Coroutine, Mapping from typing import Final, cast from unittest.mock import AsyncMock, patch import httpx import pytest import respx +from pydantic import TypeAdapter sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -23,6 +26,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.handler import ( LiteLLMMessagesToResponsesAPIHandler, _build_responses_request, ) +from litellm.llms.anthropic.pass_through.responses_adapters.streaming_iterator import AnthropicResponsesSSEByteStream from litellm.router import Router from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.openai import AllMessageValues @@ -952,3 +956,125 @@ def test_sync_bridge_raises_the_litellm_rate_limit_error_for_an_upstream_429( monkeypatch, httpx.Response(429, json={"error": {"message": "slow down", "type": "rate_limit_error"}}), ) + + +class _SuccessRowRecorder(CustomLogger): + call_types: tuple[str, ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.call_types = (*self.call_types, str(kwargs.get("call_type"))) + + +_PARKED_COMPLETION: Final = TypeAdapter(tuple[object]) + + +async def _release(parked: Coroutine[object, object, object]) -> None: + await parked + + +async def _settled_rows(recorder: _SuccessRowRecorder) -> tuple[str, ...]: + for _ in range(200): + if recorder.call_types: + break + await asyncio.sleep(0.01) + await asyncio.sleep(0.05) + return recorder.call_types + + +async def _bridged_stream( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, logging_obj: Logging +) -> tuple[AnthropicResponsesSSEByteStream, _SuccessRowRecorder]: + monkeypatch.setenv("OPENAI_API_KEY", "sk-native-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + recorder: Final = _SuccessRowRecorder() + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=RESPONSES_SSE_BODY) + ) + stream: Final = await _untyped_bridge()( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + stream=True, + custom_llm_provider="openai", + litellm_logging_obj=logging_obj, + ) + assert isinstance(stream, AnthropicResponsesSSEByteStream) + return stream, recorder + + +@pytest.mark.asyncio +async def test_bridged_stream_opens_with_the_cost_breakdown_headers( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import cost_headers + + logging_obj: Final = _messages_logging_obj(stream=True) + stream, _ = await _bridged_stream(respx_mock, monkeypatch, logging_obj) + + headers: Final = cost_headers(None, logging_obj) + await stream.aclose() + + assert headers["x-litellm-response-cost-output"] == "0.0" + assert headers["x-litellm-response-cost-discount-amount"] == "0.0" + assert headers["x-litellm-response-cost-margin-amount"] == "0.0" + + +async def _read_through_message_stop(stream: AsyncIterator[bytes]) -> None: + async for frame in stream: + if frame.startswith(b"event: message_stop"): + return + + +@pytest.mark.asyncio +async def test_bridged_stream_logs_one_row_when_the_client_stream_ends( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + stream, recorder = await _bridged_stream(respx_mock, monkeypatch, _messages_logging_obj(stream=True)) + + await _read_through_message_stop(stream) + rows_while_the_client_stream_is_open: Final = await _settled_rows(recorder) + remaining: Final[list[bytes]] = [frame async for frame in stream] + + assert rows_while_the_client_stream_is_open == () + assert remaining == [] + assert await _settled_rows(recorder) == ("anthropic_messages",) + + +@pytest.mark.asyncio +async def test_bridged_stream_still_logs_when_the_client_leaves_after_message_stop( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + stream, recorder = await _bridged_stream(respx_mock, monkeypatch, _messages_logging_obj(stream=True)) + + await _read_through_message_stop(stream) + await stream.aclose() + + assert await _settled_rows(recorder) == ("anthropic_messages",) + + +@pytest.mark.asyncio +async def test_bridged_stream_holds_its_row_for_the_proxy_end_of_stream_guardrails( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + logging_obj: Final = _messages_logging_obj(stream=True) + + async def release(*held: object) -> None: + return None + + setattr(logging_obj, "_on_deferred_stream_complete", release) # noqa: B010 # the proxy arms this slot dynamically + stream, recorder = await _bridged_stream(respx_mock, monkeypatch, logging_obj) + + frames: Final = [frame async for frame in stream] + rows_before_release: Final = await _settled_rows(recorder) + held: Final[object] = getattr(logging_obj, "_deferred_stream_complete_args", None) + + assert frames[-1].startswith(b"event: message_stop") + assert rows_before_release == () + (parked,) = _PARKED_COMPLETION.validate_python(held) + assert asyncio.iscoroutine(parked) + await _release(parked) + assert recorder.call_types == ("anthropic_messages",) diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index e1dded214bf..369a10577d2 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -7,6 +7,8 @@ import asyncio import json import os import sys +from collections.abc import AsyncIterator, Mapping +from datetime import datetime from types import SimpleNamespace import pytest @@ -15,9 +17,12 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../. import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) +from litellm.llms.anthropic.pass_through.held_stream_completion import HeldStreamCompletion from litellm.llms.anthropic.pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE from litellm.llms.anthropic.pass_through.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, @@ -622,3 +627,56 @@ class TestUpstreamFailureEndsStreamWithErrorEvent: chunks = _collect(iter(())) assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] assert chunks[1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE} + + +class _AsyncSuccessRecorder(CustomLogger): + calls: tuple[object, ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.calls = (*self.calls, response_obj) + + +def test_a_stream_that_fails_after_the_executor_completed_logs_no_success_row(): + recorder = _AsyncSuccessRecorder() + logging_obj = Logging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="fails-after-complete", + function_id="fails-after-complete", + dynamic_async_success_callbacks=[recorder], + ) + held = HeldStreamCompletion(logging_obj) + completed = ResponsesAPIResponse( + id="resp_1", + created_at=1, + status="completed", + model="m", + output=[], + parallel_tool_calls=True, + tool_choice="auto", + tools=[], + ) + + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + held.take(completed) + yield {"type": "response.completed", "response": completed} + raise ConnectionResetError("Response payload is not completed") + + async def _run() -> list[bytes]: + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=_gen(), model="m", litellm_logging_obj=logging_obj, held_completion=held + ) + frames: list[bytes] = [frame async for frame in wrapper.async_anthropic_sse_wrapper()] + await asyncio.sleep(0.2) + return frames + + frames = asyncio.run(_run()) + + assert frames[-1].startswith(b"event: error") + assert recorder.calls == ()