From fc4c961f98790c3ad9589da255af0b345212641d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:03:31 -0700 Subject: [PATCH] fix(anthropic): key the chat-completions bridge spend row on the streamed msg_ id Streaming /v1/messages against a model served through the chat-completions bridge (every non-Anthropic provider other than OpenAI) minted its msg_ id inside the stream wrapper, so the spend row landed under the provider's own completion id and the caller could not find the call by the only id it saw. The wrapper now mints the id once in its constructor and hands it to the logging object, the same way the Responses-API bridge does. --- .../adapters/handler.py | 3 + .../adapters/streaming_iterator.py | 9 +- .../adapters/transformation.py | 3 + .../responses_adapters/handler.py | 18 +-- .../experimental_pass_through/utils.py | 13 +- .../test_streaming_iterator_message_id.py | 111 ++++++++++++++++++ 6 files changed, 140 insertions(+), 17 deletions(-) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 9d61701d26d..87a29ca50ba 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -21,6 +21,7 @@ from litellm.llms.anthropic.experimental_pass_through.context_management import ) from litellm.llms.anthropic.experimental_pass_through.utils import ( is_reasoning_auto_summary_enabled, + litellm_logging_obj_from_kwargs, local_model_name, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -621,6 +622,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tool_name_mapping=tool_name_mapping, polyfill_result=polyfill_result, is_async=True, + litellm_logging_obj=litellm_logging_obj_from_kwargs(kwargs), ) if transformed_stream is not None: return transformed_stream @@ -755,6 +757,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tool_name_mapping=tool_name_mapping, polyfill_result=polyfill_result, is_async=False, + litellm_logging_obj=litellm_logging_obj_from_kwargs(kwargs), ) if transformed_stream is not None: return transformed_stream diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index cc5879df56d..78ff83cafbf 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.types.llms.anthropic import ( from litellm.types.utils import AdapterCompletionStreamWrapper, Delta if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.types.utils import ModelResponseStream @@ -287,12 +288,16 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): applied_edits: list[AppliedEdit] | None = None, compaction_block: CompactionBlock | None = None, iterations_usage: list[UsageIteration] | None = None, + litellm_logging_obj: "LiteLLMLoggingObject | None" = None, ): # Wrap the upstream stream so chunks that carry both content and a # finish_reason (fake-streamed providers) are split into two — see # _CombinedChunkSplitter. super().__init__(_CombinedChunkSplitter(completion_stream)) self.model = model + 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) # 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. @@ -507,7 +512,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): { "type": "message_start", "message": { - "id": f"msg_{uuid.uuid4()}", + "id": self._message_id, "type": "message", "role": "assistant", "content": [], @@ -741,7 +746,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): { "type": "message_start", "message": { - "id": f"msg_{uuid.uuid4()}", + "id": self._message_id, "type": "message", "role": "assistant", "content": [], diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 199a8ab77e7..573a461e89e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -174,6 +174,7 @@ from litellm.types.utils import Choices, ModelResponse, StreamingChoices, Usage from .streaming_iterator import AnthropicStreamWrapper if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.types.llms.anthropic import ContentBlockContentBlockDict ToolResultContent: TypeAlias = str | list[ToolMessageContentPart] @@ -264,6 +265,7 @@ class AnthropicAdapter: tool_name_mapping: dict[str, str] | None = None, polyfill_result: PolyfillResult | None = None, is_async: bool = True, + litellm_logging_obj: "LiteLLMLoggingObject | None" = None, ) -> AsyncIterator[bytes] | Iterator[bytes] | None: """ Translate OpenAI streaming response to Anthropic format. @@ -290,6 +292,7 @@ class AnthropicAdapter: applied_edits=applied_edits, compaction_block=compaction_block, iterations_usage=iterations_usage, + litellm_logging_obj=litellm_logging_obj, ) # Return the SSE-wrapped version for proper event formatting. if is_async: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 77296c3416d..0445c23ed8c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -5,7 +5,7 @@ Used when the target model is an OpenAI or Azure model. """ from collections.abc import AsyncIterator, Coroutine, Mapping -from typing import TYPE_CHECKING, Any, Final, TypeAlias +from typing import Any, Final, TypeAlias import litellm from litellm.types.llms.anthropic import ( @@ -20,25 +20,15 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( ) from litellm.types.llms.openai import ResponsesAPIResponse -from ..utils import local_model_name +from ..utils import litellm_logging_obj_from_kwargs, local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper from .transformation import LiteLLMAnthropicToResponsesAPIAdapter -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject - AnthropicRequestMessages: TypeAlias = list[AllAnthropicMessageValues] | list[dict[str, object]] _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter() -def _litellm_logging_obj(responses_kwargs: Mapping[str, object]) -> "LiteLLMLoggingObject | None": - from litellm.litellm_core_utils.litellm_logging import Logging - - candidate: Final = responses_kwargs.get("litellm_logging_obj") - return candidate if isinstance(candidate, Logging) else None - - def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, object]: """The litellm-specific kwargs forwarded verbatim onto the Responses API request.""" return extra_kwargs or {} @@ -198,7 +188,7 @@ class LiteLLMMessagesToResponsesAPIHandler: wrapper: Final = AnthropicResponsesStreamWrapper( responses_stream=result, model=local_model_name(model, kwargs.get("custom_llm_provider")), - litellm_logging_obj=_litellm_logging_obj(responses_kwargs), + litellm_logging_obj=litellm_logging_obj_from_kwargs(responses_kwargs), ) return wrapper.async_anthropic_sse_wrapper() @@ -280,7 +270,7 @@ class LiteLLMMessagesToResponsesAPIHandler: wrapper: Final = AnthropicResponsesStreamWrapper( responses_stream=result, model=local_model_name(model, kwargs.get("custom_llm_provider")), - litellm_logging_obj=_litellm_logging_obj(responses_kwargs), + litellm_logging_obj=litellm_logging_obj_from_kwargs(responses_kwargs), ) return wrapper.async_anthropic_sse_wrapper() diff --git a/litellm/llms/anthropic/experimental_pass_through/utils.py b/litellm/llms/anthropic/experimental_pass_through/utils.py index 716a4f54778..55fe9c47faf 100644 --- a/litellm/llms/anthropic/experimental_pass_through/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/utils.py @@ -1,11 +1,14 @@ import os from collections.abc import Mapping from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final import litellm from litellm.types.utils import ModelInfo +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject + OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH: Final = 64 _EFFORT_DEGRADATION_CHAIN: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( @@ -24,6 +27,14 @@ def prompt_cache_key_from_user_id(user_id: object) -> str | None: return str(user_id)[:OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH] or None +def litellm_logging_obj_from_kwargs(kwargs: Mapping[str, object]) -> "LiteLLMLoggingObject | None": + """The logging object the bridged call logs through, when the caller supplied one.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + candidate: Final = kwargs.get("litellm_logging_obj") + return candidate if isinstance(candidate, Logging) else None + + def local_model_name(model: str, custom_llm_provider: object) -> str: """The id the provider itself knows, for reporting back to the caller in ``message_start``.""" return model.removeprefix(f"{custom_llm_provider}/") if isinstance(custom_llm_provider, str) else model diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py new file mode 100644 index 00000000000..7cd789529c8 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py @@ -0,0 +1,111 @@ +""" +Streaming ``/v1/messages`` against a model that is neither Anthropic nor OpenAI is served by +translating the call onto ``/v1/chat/completions``, and the ``msg_`` id the caller is streamed +is minted right here. It is the only request id such a caller ever sees, so the spend row has +to be keyed on that same value rather than on the provider's own completion id. +""" + +import datetime +import json + +import pytest +import respx + +import litellm +from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( + LiteLLMMessagesToCompletionTransformationHandler, +) +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) + +MESSAGES = [{"role": "user", "content": "hello"}] + +GROQ_CHAT_URL = "https://api.groq.com/openai/v1/chat/completions" + +CHAT_SSE_BODY = ( + b'data: {"id":"chatcmpl-lit6825","object":"chat.completion.chunk","created":1,' + b'"model":"kimi-k2","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},' + b'"finish_reason":null}]}\n\n' + b'data: {"id":"chatcmpl-lit6825","object":"chat.completion.chunk","created":1,' + b'"model":"kimi-k2","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],' + b'"usage":{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7}}\n\n' + b"data: [DONE]\n\n" +) + + +def _logging_obj(call_id: str): + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model="kimi-k2", + messages=MESSAGES, + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime.now(datetime.timezone.utc), + litellm_call_id=call_id, + function_id="1234", + ) + + +def _streamed_message_id(raw_events: list[bytes]) -> str: + events = [json.loads(chunk.decode().split("data: ", 1)[1]) for chunk in raw_events] + message_start = next(e for e in events if e["type"] == "message_start") + return message_start["message"]["id"] + + +@pytest.fixture(autouse=True) +def _intercept_groq(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("GROQ_API_KEY", "gsk-lit6825-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(GROQ_CHAT_URL).respond( + status_code=200, + headers={"Content-Type": "text/event-stream"}, + content=CHAT_SSE_BODY, + ) + + +@pytest.mark.asyncio +async def test_async_streaming_hands_the_logging_object_the_message_id_the_caller_is_streamed(): + logging_obj = _logging_obj("6825beef-0000-4000-8000-000000000010") + + sse = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=1024, + messages=MESSAGES, + model="groq/kimi-k2", + stream=True, + custom_llm_provider="groq", + litellm_logging_obj=logging_obj, + ) + streamed_id = _streamed_message_id([chunk async for chunk in sse]) + + assert streamed_id.startswith("msg_") + assert logging_obj.streamed_anthropic_message_id == streamed_id + + +def test_sync_streaming_hands_the_logging_object_the_message_id_the_caller_is_streamed(): + logging_obj = _logging_obj("6825beef-0000-4000-8000-000000000011") + + sse = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + max_tokens=1024, + messages=MESSAGES, + model="groq/kimi-k2", + stream=True, + custom_llm_provider="groq", + litellm_logging_obj=logging_obj, + ) + streamed_id = _streamed_message_id(list(sse)) + + assert streamed_id.startswith("msg_") + assert logging_obj.streamed_anthropic_message_id == streamed_id + + +def test_concurrent_streams_are_keyed_on_their_own_message_id(): + """Two callers streaming at once must not be handed, or logged under, one another's id.""" + first = AnthropicStreamWrapper(completion_stream=iter([]), model="kimi-k2") + second = AnthropicStreamWrapper(completion_stream=iter([]), model="kimi-k2") + + assert first._message_id != second._message_id + assert _streamed_message_id(list(first.anthropic_sse_wrapper())) == first._message_id + assert _streamed_message_id(list(second.anthropic_sse_wrapper())) == second._message_id