diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index d975c3551f3..67684a230e3 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -685,7 +685,7 @@ class ChunkProcessor: def _flush_thinking_block() -> None: nonlocal current_thinking_text_parts, current_signature - if len(current_thinking_text_parts) > 0 and current_signature: + if current_signature: thinking_blocks.append( ChatCompletionThinkingBlock( type="thinking", diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 5173cd04a89..21a33c17ab8 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -73,6 +73,11 @@ def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | ) +def _delta_has_signed_thinking_block(delta: object) -> bool: + blocks: Final = getattr(delta, "thinking_blocks", None) or () + return any(isinstance(b, dict) and (b.get("signature") or b.get("data")) for b in blocks) + + class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): """ Async iterator for processing streaming responses from the Responses API. @@ -936,7 +941,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.sent_output_item_added_event = True # Reasoning-first - if hasattr(delta, "reasoning_content") and delta.reasoning_content: + if (hasattr(delta, "reasoning_content") and delta.reasoning_content) or _delta_has_signed_thinking_block(delta): self._reasoning_active = True if self._cached_reasoning_item_id is None: self._cached_reasoning_item_id = f"rs_{uuid.uuid4()}" diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py deleted file mode 100644 index 294d26b2e58..00000000000 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ /dev/null @@ -1,93 +0,0 @@ -import json -from collections.abc import AsyncIterator -from types import SimpleNamespace -from typing import Any, cast - -import httpx -import pytest - -import litellm -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.main import vertex_gemma_chat_completion -from litellm.types.llms.openai import OutputTextDeltaEvent, ResponseCompletedEvent, ResponsesAPIStreamingResponse - -_VERTEX_URL = "https://example.invalid/v1/projects/test/locations/us-central1/endpoints/test:predict" -_MESSAGES = [{"role": "user", "content": "Reply exactly READY"}] -_FAKE_CREDENTIALS = "gemma-test-credentials" - - -def _vertex_response(): - return { - "predictions": { - "id": "chatcmpl-stream-test", - "created": 1759863903, - "model": "google/gemma-3-12b-it", - "object": "chat.completion", - "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "READY"}}], - "usage": {"prompt_tokens": 14, "completion_tokens": 1, "total_tokens": 15}, - } - } - - -@pytest.fixture(autouse=True) -def _cached_access_token(): - """Serve a fake token from the handler's credential cache so no auth round-trip runs.""" - cache = vertex_gemma_chat_completion._credentials_project_mapping - key = (_FAKE_CREDENTIALS, "test") - cache[key] = (SimpleNamespace(token="fake-token", expired=False), "test") - yield - cache.pop(key, None) - - -def test_sync_gemma_stream(): - captured: dict[str, Any] = {} - - def handle(request: httpx.Request) -> httpx.Response: - captured["body"] = json.loads(request.content) - return httpx.Response(200, json=_vertex_response()) - - stream = litellm.completion( - model="vertex_ai/gemma/test-model", - messages=_MESSAGES, - stream=True, - api_base=_VERTEX_URL, - vertex_project="test", - vertex_location="us-central1", - vertex_credentials=_FAKE_CREDENTIALS, - client=httpx.Client(transport=httpx.MockTransport(handle)), - ) - - assert isinstance(stream, CustomStreamWrapper) - chunks = list(stream) - - assert "stream" not in captured["body"]["instances"][0] - assert len(chunks) == 2 - assert chunks[0].choices[0].delta.content == "READY" - assert chunks[1].choices[0].finish_reason == "stop" - - -@pytest.mark.asyncio -async def test_async_gemma_responses_stream(): - captured: dict[str, Any] = {} - - def handle(request: httpx.Request) -> httpx.Response: - captured["body"] = json.loads(request.content) - return httpx.Response(200, json=_vertex_response()) - - response = await litellm.aresponses( - model="vertex_ai/gemma/test-model", - input="Reply exactly READY", - stream=True, - api_base=_VERTEX_URL, - vertex_project="test", - vertex_location="us-central1", - vertex_credentials=_FAKE_CREDENTIALS, - client=httpx.AsyncClient(transport=httpx.MockTransport(handle)), - ) - events = [event async for event in cast(AsyncIterator[ResponsesAPIStreamingResponse], response)] - - assert "stream" not in captured["body"]["instances"][0] - assert "READY" in "".join(event.delta for event in events if isinstance(event, OutputTextDeltaEvent)) - assert isinstance(events[-1], ResponseCompletedEvent) - assert events[-1].response.usage is not None - assert events[-1].response.usage.total_tokens == 15 diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index af763da2d87..aaf877df364 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -236,6 +236,31 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): assert result[2]["signature"] == "sig_block2" +def test_get_combined_thinking_content_keeps_signed_block_without_thinking_text(): + chunks: Final = [ + ModelResponseStream( + id="chatcmpl-123", + object="chat.completion.chunk", + created=1234567890, + model="claude-sonnet-4-20250514", + choices=[ + StreamingChoices( + index=0, + delta=Delta(thinking_blocks=[{"type": "thinking", "thinking": "", "signature": "sig_only"}]), + finish_reason=None, + ) + ], + ) + ] + + result: Final = ChunkProcessor(chunks=chunks).get_combined_thinking_content(chunks) + + assert result is not None + assert [(block["type"], block["thinking"], block["signature"]) for block in result] == [ + ("thinking", "", "sig_only") + ] + + def test_cache_read_input_tokens_retained(): chunk1 = ModelResponseStream( id="chatcmpl-95aabb85-c39f-443d-ae96-0370c404d70c", diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index 97f4f290958..92e684e42c8 100644 --- a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -1303,3 +1303,81 @@ class TestVertexGemmaCompletion: mock_async_post.assert_awaited_once() assert mock_async_post.call_args.kwargs["client"] is None assert response.choices[0].message.content == "default async handler fallback" + + +_GEMMA_VERTEX_URL = "https://example.invalid/v1/projects/test/locations/us-central1/endpoints/test:predict" +_FAKE_GEMMA_CREDENTIALS = "gemma-test-credentials" + + +@pytest.fixture +def _gemma_cached_access_token(): + """Serve a fake token from the handler's credential cache so no auth round-trip runs.""" + from types import SimpleNamespace + + from litellm.main import vertex_gemma_chat_completion + + cache = vertex_gemma_chat_completion._credentials_project_mapping + key = (_FAKE_GEMMA_CREDENTIALS, "test") + cache[key] = (SimpleNamespace(token="fake-token", expired=False), "test") + yield + cache.pop(key, None) + + +def test_sync_gemma_stream(_gemma_cached_access_token): + import httpx + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=_make_gemma_vertex_response(content="READY", total_tokens=15)) + + stream = litellm.completion( + model="vertex_ai/gemma/test-model", + messages=[{"role": "user", "content": "Reply exactly READY"}], + stream=True, + api_base=_GEMMA_VERTEX_URL, + vertex_project="test", + vertex_location="us-central1", + vertex_credentials=_FAKE_GEMMA_CREDENTIALS, + client=httpx.Client(transport=httpx.MockTransport(handle)), + ) + + assert isinstance(stream, CustomStreamWrapper) + chunks = list(stream) + + assert "stream" not in captured["body"]["instances"][0] + assert len(chunks) == 2 + assert chunks[0].choices[0].delta.content == "READY" + assert chunks[1].choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_async_gemma_responses_stream(_gemma_cached_access_token): + import httpx + + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=_make_gemma_vertex_response(content="READY", total_tokens=15)) + + response = await litellm.aresponses( + model="vertex_ai/gemma/test-model", + input="Reply exactly READY", + stream=True, + api_base=_GEMMA_VERTEX_URL, + vertex_project="test", + vertex_location="us-central1", + vertex_credentials=_FAKE_GEMMA_CREDENTIALS, + client=httpx.AsyncClient(transport=httpx.MockTransport(handle)), + ) + events = [event async for event in cast(AsyncIterator[ResponsesAPIStreamingResponse], response)] + + assert "stream" not in captured["body"]["instances"][0] + assert "READY" in "".join(event.delta for event in events if isinstance(event, OutputTextDeltaEvent)) + assert isinstance(events[-1], ResponseCompletedEvent) + assert events[-1].response.usage is not None + assert events[-1].response.usage.total_tokens == 15 diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 8fbba0dbf87..041bcf1b6d7 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -978,6 +978,25 @@ def _reasoning_chunk(reasoning: str, finish_reason: str | None = None) -> ModelR ) +def _signature_only_thinking_chunk(signature: str) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + thinking_blocks=[{"type": "thinking", "thinking": "", "signature": signature}], + ), + finish_reason=None, + ) + ], + ) + + async def _collect_events( iterator: LiteLLMCompletionStreamingIterator, sync_mode: bool ) -> list[BaseLiteLLMOpenAIResponseObject]: @@ -1015,6 +1034,27 @@ async def test_tool_only_stream_emits_no_message_item_events(sync_mode: bool): assert any(getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED for event in events) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_signature_only_thinking_streams_a_replayable_reasoning_item(sync_mode: bool): + iterator: Final = _build_iterator([_signature_only_thinking_chunk("sig_only"), _chunk("4", finish_reason="stop")]) + + events: Final = await _collect_events(iterator, sync_mode) + + added_item_types: Final = [ + event.item.type + for event in events + if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + ] + completed: Final = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + reasoning_items: Final = [item for item in completed.response.output if getattr(item, "type", None) == "reasoning"] + assert added_item_types[0] == "reasoning" + assert len(reasoning_items) == 1 + assert json.loads(reasoning_items[0].encrypted_content)[0]["signature"] == "sig_only" + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode: bool):