From 4d1c7b87aceaa10b8c9d10644a3eb163f75b144d Mon Sep 17 00:00:00 2001 From: dschulmeist <62301265+dschulmeist@users.noreply.github.com> Date: Wed, 22 Apr 2026 14:42:20 +0200 Subject: [PATCH] fix(proxy): skip redundant tiktoken recount when provider supplies reasoning_tokens (#26193) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stream_chunk_builder unconditionally called ChunkProcessor.count_reasoning_tokens, which recomputes reasoning_tokens via tiktoken.encode() — a C extension that holds the GIL for tens of seconds on large reasoning responses (Claude extended thinking, OpenAI o1/o3, Gemini thinking). The computed value is already discarded by calculate_usage when the provider supplied reasoning_tokens in a streaming usage chunk (see lines applying reasoning_tokens only when completion_tokens_details.reasoning_tokens is None). All major reasoning providers do supply this value, so in the common case the expensive recount runs purely to be thrown away — while blocking the asyncio event loop long enough to fail K8s liveness probes and get pods killed. Adds ChunkProcessor.chunks_have_reasoning_tokens(chunks) which inspects the already-parsed streaming chunks for any completion_tokens_details.reasoning_tokens entry. Short-circuits the call site in stream_chunk_builder when one is found. The slow fallback path via count_reasoning_tokens is preserved for providers that emit reasoning_content without reporting usage. Tests cover: - chunks_have_reasoning_tokens returns True when provider reports the value - returns False when provider omits it (fallback path preserved) - handles plain-dict chunks as well as ModelResponseStream objects - end-to-end: stream_chunk_builder does not invoke count_reasoning_tokens when the provider already supplies reasoning_tokens Fixes #26193 --- .../streaming_chunk_builder_utils.py | 46 ++++++++ litellm/main.py | 10 +- .../test_streaming_chunk_builder_utils.py | 107 ++++++++++++++++++ 3 files changed, 162 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index c808cfca457..f53300b7b3c 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -508,6 +508,52 @@ class ChunkProcessor: "prompt_tokens_details": prompt_tokens_details, } + def chunks_have_reasoning_tokens( + self, + chunks: List[Union[Dict[str, Any], ModelResponse]], + ) -> bool: + """ + Return True if any streaming usage chunk already reports + completion_tokens_details.reasoning_tokens. + + When the provider supplies this value, calculate_usage uses it + directly and discards any locally recomputed value (see + calculate_usage: reasoning_tokens is only applied when + completion_tokens_details.reasoning_tokens is None). The local + recount via count_reasoning_tokens calls tiktoken.encode() which + holds the GIL and blocks the asyncio event loop on large reasoning + responses, so skipping it in this common case avoids liveness-probe + timeouts without changing any observable output. + """ + for chunk in chunks: + usage: Any = None + if isinstance(chunk, dict): + usage = chunk.get("usage") + else: + usage = getattr(chunk, "usage", None) + if usage is None and hasattr(chunk, "_hidden_params"): + hidden = chunk._hidden_params + if isinstance(hidden, dict): + usage = hidden.get("usage") + if usage is None: + continue + + details: Any = None + if isinstance(usage, dict): + details = usage.get("completion_tokens_details") + elif hasattr(usage, "completion_tokens_details"): + details = usage.completion_tokens_details + if details is None: + continue + + if isinstance(details, dict): + reasoning_tokens = details.get("reasoning_tokens") + else: + reasoning_tokens = getattr(details, "reasoning_tokens", None) + if reasoning_tokens is not None: + return True + return False + def count_reasoning_tokens(self, response: ModelResponse) -> Optional[int]: reasoning_tokens: Optional[int] = None for choice in response.choices: diff --git a/litellm/main.py b/litellm/main.py index 73db4a11cb7..94d0a6474c2 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7619,7 +7619,15 @@ def stream_chunk_builder( # noqa: PLR0915 completion_output = get_content_from_model_response(response) - reasoning_tokens = processor.count_reasoning_tokens(response) + # Skip the local tiktoken recount when the provider already reports + # reasoning_tokens in a streaming usage chunk. calculate_usage would + # discard the recomputed value anyway, and tiktoken.encode() holds + # the GIL for tens of seconds on large reasoning responses — long + # enough to fail liveness probes and get pods killed. + if processor.chunks_have_reasoning_tokens(chunks): + reasoning_tokens = None + else: + reasoning_tokens = processor.count_reasoning_tokens(response) usage = processor.calculate_usage( chunks=chunks, diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd9..47a7fbd3b3a 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProces from litellm.types.utils import ( ChatCompletionDeltaToolCall, ChatCompletionMessageToolCall, + CompletionTokensDetailsWrapper, Delta, Function, ModelResponseStream, @@ -613,3 +614,109 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert ( response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" ) + + +def _make_reasoning_chunk( + reasoning_tokens, + completion_tokens=10, + prompt_tokens=5, + reasoning_content="some reasoning content", +): + """Build a minimal ModelResponseStream with optional reasoning_tokens usage.""" + completion_details = ( + CompletionTokensDetailsWrapper(reasoning_tokens=reasoning_tokens) + if reasoning_tokens is not None + else None + ) + chunk = ModelResponseStream( + id="chatcmpl-reasoning-1", + created=1, + model="o1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="final answer", role="assistant"), + ) + ], + stream_options={"include_usage": True}, + usage=Usage( + completion_tokens=completion_tokens, + prompt_tokens=prompt_tokens, + total_tokens=completion_tokens + prompt_tokens, + completion_tokens_details=completion_details, + prompt_tokens_details=None, + ), + ) + # Attach reasoning_content at the message level (how stream_chunk_builder + # sees the assembled response). + chunk.choices[0].delta.provider_specific_fields = { + "reasoning_content": reasoning_content + } + return chunk + + +def test_chunks_have_reasoning_tokens_true_when_provider_supplies(): + """Regression test for https://github.com/BerriAI/litellm/issues/26193. + + When any streaming usage chunk already reports reasoning_tokens, + chunks_have_reasoning_tokens must return True so callers can skip the + expensive local tiktoken recount (which would be discarded downstream + by calculate_usage and which holds the GIL long enough to fail K8s + liveness probes on large reasoning responses). + """ + chunks = [_make_reasoning_chunk(reasoning_tokens=12345)] + processor = ChunkProcessor(chunks=chunks) + assert processor.chunks_have_reasoning_tokens(chunks) is True + + +def test_chunks_have_reasoning_tokens_false_when_provider_omits(): + """When no streaming usage chunk reports reasoning_tokens, the caller + must still fall back to count_reasoning_tokens (the existing slow path) + so providers that don't supply usage are unaffected by the short-circuit. + """ + chunks = [_make_reasoning_chunk(reasoning_tokens=None)] + processor = ChunkProcessor(chunks=chunks) + assert processor.chunks_have_reasoning_tokens(chunks) is False + + +def test_chunks_have_reasoning_tokens_handles_dict_chunks(): + """stream_chunk_builder accepts plain-dict chunks (e.g. from proxies). + The presence check must handle both ModelResponseStream objects and + dicts without raising.""" + chunks = [ + { + "id": "chatcmpl-reasoning-dict", + "choices": [{"finish_reason": "stop", "index": 0, "delta": {}}], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 10, + "total_tokens": 15, + "completion_tokens_details": {"reasoning_tokens": 999}, + }, + } + ] + processor = ChunkProcessor(chunks=chunks) + assert processor.chunks_have_reasoning_tokens(chunks) is True + + +def test_stream_chunk_builder_skips_count_reasoning_tokens_when_usage_present(): + """End-to-end: stream_chunk_builder must not invoke the expensive + count_reasoning_tokens path when the provider already supplied + reasoning_tokens in a streaming usage chunk.""" + from unittest.mock import patch + + chunks = [_make_reasoning_chunk(reasoning_tokens=42)] + + with patch.object( + ChunkProcessor, + "count_reasoning_tokens", + autospec=True, + ) as mock_count: + response = stream_chunk_builder(chunks=chunks) + + assert response is not None + mock_count.assert_not_called() + # Provider-supplied reasoning_tokens is preserved on the final response. + assert response.usage.completion_tokens_details.reasoning_tokens == 42