mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(anthropic): start a new streamed block for each reasoning item
The backported Responses bridge tags each reasoning item's thinking block with its output index, but this line's Anthropic stream wrapper only opened a new block when the block type changed, so a second reasoning item of the same type was dropped. This ports the index check from #46863, which main already has.
This commit is contained in:
parent
4fc1d5bb0d
commit
ba5d7b83b9
2 changed files with 96 additions and 0 deletions
|
|
@ -24,6 +24,7 @@ from typing_extensions import assert_never
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.hidden_params import set_hidden_params
|
||||
from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults
|
||||
from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame
|
||||
|
|
@ -71,6 +72,14 @@ def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]:
|
|||
return value if value else ()
|
||||
|
||||
|
||||
def _first_thinking_block_index(chunk: "ModelResponseStream") -> int | None:
|
||||
delta: Final = _optional_attr(chunk.choices[0], "delta") if chunk.choices else None
|
||||
blocks: Final = _optional_attr_sequence(delta, "thinking_blocks")
|
||||
first_block: Final = as_str_mapping(blocks[0]) if blocks else None
|
||||
index: Final = first_block.get("index") if first_block is not None else None
|
||||
return index if isinstance(index, int) else None
|
||||
|
||||
|
||||
def _error_status_and_message(exc: Exception) -> tuple[int, str]:
|
||||
if isinstance(exc, (BaseLLMException, MidStreamFallbackError, ContextWindowExceededError)):
|
||||
return exc.status_code, exc.message
|
||||
|
|
@ -464,6 +473,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
holding_stop_reason_chunk: MessageBlockDelta | None = None
|
||||
queued_usage_chunk: bool = False
|
||||
current_content_block_index: int = 0
|
||||
current_thinking_block_index: int | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -1352,11 +1362,24 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
original_name: Final = self.tool_name_mapping.get(truncated_name, truncated_name)
|
||||
tool_block["name"] = original_name
|
||||
|
||||
thinking_block_index: Final = _first_thinking_block_index(chunk)
|
||||
previous_thinking_block_index: Final = self.current_thinking_block_index
|
||||
if thinking_block_index is not None or block_type != self.current_content_block_type:
|
||||
self.current_thinking_block_index = thinking_block_index
|
||||
|
||||
if block_type != self.current_content_block_type:
|
||||
self.current_content_block_type = block_type
|
||||
self.current_content_block_start = content_block_start
|
||||
return True
|
||||
|
||||
if (
|
||||
thinking_block_index is not None
|
||||
and previous_thinking_block_index is not None
|
||||
and thinking_block_index != previous_thinking_block_index
|
||||
):
|
||||
self.current_content_block_start = content_block_start
|
||||
return True
|
||||
|
||||
# For parallel tool calls, we'll necessarily have a new content block
|
||||
# if we get a function name since it signals a new tool call
|
||||
if block_type == "tool_use":
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import base64
|
|||
import functools
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, cast
|
||||
|
||||
|
|
@ -5615,3 +5616,75 @@ async def test_openai_thinking_on_the_chat_adapter_keeps_reasoning_its_responses
|
|||
|
||||
assert route.call_count == 1
|
||||
assert [item["encrypted_content"] for item in _sent_reasoning_items(route.calls[0].request)] == ["gAAA_same_org"]
|
||||
|
||||
|
||||
def _responses_sse(events: Sequence[Mapping[str, object]]) -> bytes:
|
||||
return b"".join(
|
||||
f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode()
|
||||
for number, event in enumerate(events)
|
||||
)
|
||||
|
||||
|
||||
def _reasoning_item(item_id: str, encrypted: str) -> dict[str, object]:
|
||||
return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": encrypted}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_chat_bridge_httpx")
|
||||
async def test_streamed_consecutive_reasoning_items_each_become_their_own_block(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
response: Final = {**_RETRY_REPLY, "id": "resp_two_reasoning", "status": "in_progress", "output": []}
|
||||
message: Final = {
|
||||
"type": "message",
|
||||
"id": "msg_two_reasoning",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "Oslo", "annotations": []}],
|
||||
}
|
||||
first: Final = _reasoning_item("rs_1", "gAAA_first")
|
||||
second: Final = _reasoning_item("rs_2", "gAAA_second")
|
||||
completed: Final = {**response, "status": "completed", "output": [first, second, message]}
|
||||
events: Final = (
|
||||
{"type": "response.created", "response": response},
|
||||
{"type": "response.output_item.added", "output_index": 0, "item": _reasoning_item("rs_1", "")},
|
||||
{"type": "response.output_item.done", "output_index": 0, "item": first},
|
||||
{"type": "response.output_item.added", "output_index": 1, "item": _reasoning_item("rs_2", "")},
|
||||
{"type": "response.output_item.done", "output_index": 1, "item": second},
|
||||
{"type": "response.output_item.added", "output_index": 2, "item": {**message, "content": []}},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"output_index": 2,
|
||||
"item_id": "msg_two_reasoning",
|
||||
"content_index": 0,
|
||||
"delta": "Oslo",
|
||||
},
|
||||
{"type": "response.output_item.done", "output_index": 2, "item": message},
|
||||
{"type": "response.completed", "response": completed},
|
||||
)
|
||||
respx_mock.post(url__regex=_RETRY_URL).mock(
|
||||
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=_responses_sse(events))
|
||||
)
|
||||
|
||||
stream: Final = await litellm.anthropic.messages.acreate(
|
||||
model="azure/gpt-5.4-nano",
|
||||
api_base=_RETRY_API_BASE,
|
||||
api_key="fake-key",
|
||||
api_version="2025-04-01-preview",
|
||||
max_tokens=64,
|
||||
tools=_RETRY_TOOLS,
|
||||
stream=True,
|
||||
messages=[{"role": "user", "content": "Which city is colder?"}],
|
||||
)
|
||||
frames: Final = b"".join([chunk async for chunk in stream]).decode()
|
||||
starts: Final = tuple(
|
||||
json.loads(line.removeprefix("data: "))["content_block"]
|
||||
for line in frames.splitlines()
|
||||
if line.startswith("data: ") and '"content_block_start"' in line
|
||||
)
|
||||
|
||||
assert [(block["type"], block.get("data")) for block in starts] == [
|
||||
("redacted_thinking", encrypted_reasoning_signature("gAAA_first")),
|
||||
("redacted_thinking", encrypted_reasoning_signature("gAAA_second")),
|
||||
("text", None),
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue