Merge pull request #22701 from BerriAI/fix/streaming-and-azure-gpt5-test-failures

Fix Anthropic streaming sync and Azure GPT-5.1 logprobs tests
This commit is contained in:
Julio Quinteros Pro 2026-03-03 19:36:44 -03:00 committed by GitHub
commit 2c5c38333d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 105 additions and 96 deletions

View File

@ -41,7 +41,6 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
type="text",
text="",
)
pending_new_content_block: bool = False
chunk_queue: deque = deque() # Queue for buffering multiple chunks
def __init__(
@ -80,38 +79,40 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
from .transformation import LiteLLMAnthropicMessagesAdapter
try:
# Always return queued chunks first
if self.chunk_queue:
return self.chunk_queue.popleft()
# Queue initial chunks if not sent yet
if self.sent_first_chunk is False:
self.sent_first_chunk = True
return {
"type": "message_start",
"message": {
"id": "msg_{}".format(uuid.uuid4()),
"type": "message",
"role": "assistant",
"content": [],
"model": self.model,
"stop_reason": None,
"stop_sequence": None,
"usage": self._create_initial_usage_delta(),
},
}
self.chunk_queue.append(
{
"type": "message_start",
"message": {
"id": "msg_{}".format(uuid.uuid4()),
"type": "message",
"role": "assistant",
"content": [],
"model": self.model,
"stop_reason": None,
"stop_sequence": None,
"usage": self._create_initial_usage_delta(),
},
}
)
return self.chunk_queue.popleft()
if self.sent_content_block_start is False:
self.sent_content_block_start = True
return {
"type": "content_block_start",
"index": self.current_content_block_index,
"content_block": {"type": "text", "text": ""},
}
# Handle pending new content block start
if self.pending_new_content_block:
self.pending_new_content_block = False
self.sent_content_block_finish = False # Reset for new block
return {
"type": "content_block_start",
"index": self.current_content_block_index,
"content_block": self.current_content_block_start,
}
self.chunk_queue.append(
{
"type": "content_block_start",
"index": self.current_content_block_index,
"content_block": {"type": "text", "text": ""},
}
)
return self.chunk_queue.popleft()
for chunk in self.completion_stream:
if chunk == "None" or chunk is None:
@ -126,45 +127,65 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
current_content_block_index=self.current_content_block_index,
)
# Check if we need to start a new content block
# This is where you'd add your logic to detect when a new content block should start
# For example, if the chunk indicates a tool call or different content type
if should_start_new_block and not self.sent_content_block_finish:
# End current content block and prepare for new one
self.holding_chunk = processed_chunk
self.sent_content_block_finish = True
self.pending_new_content_block = True
return {
"type": "content_block_stop",
"index": max(self.current_content_block_index - 1, 0),
}
# Queue the sequence: content_block_stop -> content_block_start
# The trigger chunk itself is not emitted as a delta since the
# content_block_start already carries the relevant information.
self.chunk_queue.append(
{
"type": "content_block_stop",
"index": max(self.current_content_block_index - 1, 0),
}
)
self.chunk_queue.append(
{
"type": "content_block_start",
"index": self.current_content_block_index,
"content_block": self.current_content_block_start,
}
)
self.sent_content_block_finish = False
return self.chunk_queue.popleft()
if (
processed_chunk["type"] == "message_delta"
and self.sent_content_block_finish is False
):
self.holding_chunk = processed_chunk
# Queue both the content_block_stop and the message_delta
self.chunk_queue.append(
{
"type": "content_block_stop",
"index": self.current_content_block_index,
}
)
self.sent_content_block_finish = True
return {
"type": "content_block_stop",
"index": self.current_content_block_index,
}
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
elif self.holding_chunk is not None:
return_chunk = self.holding_chunk
self.holding_chunk = processed_chunk
return return_chunk
self.chunk_queue.append(self.holding_chunk)
self.chunk_queue.append(processed_chunk)
self.holding_chunk = None
return self.chunk_queue.popleft()
else:
return processed_chunk
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
# Handle any remaining held chunks after stream ends
if self.holding_chunk is not None:
return_chunk = self.holding_chunk
self.chunk_queue.append(self.holding_chunk)
self.holding_chunk = None
return return_chunk
if self.sent_last_message is False:
if not self.sent_last_message:
self.sent_last_message = True
return {"type": "message_stop"}
self.chunk_queue.append({"type": "message_stop"})
if self.chunk_queue:
return self.chunk_queue.popleft()
raise StopIteration
except StopIteration:
if self.chunk_queue:
return self.chunk_queue.popleft()
if self.sent_last_message is False:
self.sent_last_message = True
return {"type": "message_stop"}
@ -265,7 +286,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
if not self.queued_usage_chunk:
if should_start_new_block and not self.sent_content_block_finish:
# Queue the sequence: content_block_stop -> content_block_start -> current_chunk
# Queue the sequence: content_block_stop -> content_block_start
# The trigger chunk itself is not emitted as a delta since the
# content_block_start already carries the relevant information.
# 1. Stop current content block
self.chunk_queue.append(
@ -284,9 +307,6 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
}
)
# 3. Queue the current chunk (don't lose it!)
self.chunk_queue.append(processed_chunk)
# Reset state for new block
self.sent_content_block_finish = False

View File

@ -43,8 +43,12 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
if "tool_choice" not in params:
params.append("tool_choice")
# Only gpt-5.2 has been verified to support logprobs on Azure
if self.is_model_gpt_5_2_model(model):
# Only gpt-5.2 has been verified to support logprobs on Azure.
# The base OpenAI class includes logprobs for gpt-5.1+, but Azure
# hasn't verified support for gpt-5.1, so remove them unless gpt-5.2.
if self.is_model_gpt_5_1_model(model) and not self.is_model_gpt_5_2_model(model):
params = [p for p in params if p not in ["logprobs", "top_logprobs"]]
elif self.is_model_gpt_5_2_model(model):
azure_supported_params = ["logprobs", "top_logprobs"]
params.extend(azure_supported_params)

View File

@ -23,7 +23,7 @@ sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import (
AnthropicStreamWrapper,
)
from litellm.types.utils import Delta, ModelResponse, StreamingChoices, Usage
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
class MockCompletionStreamWithContentAfterStopReason:
@ -32,16 +32,14 @@ class MockCompletionStreamWithContentAfterStopReason:
def __init__(self):
self.responses = [
# Initial text content
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="Hello"), index=0, finish_reason=None
)
],
),
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=" world"), index=0, finish_reason=None
@ -49,8 +47,7 @@ class MockCompletionStreamWithContentAfterStopReason:
],
),
# Message delta with stop_reason AND usage (this is how it actually comes from the API)
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=""), index=0, finish_reason="stop"
@ -60,8 +57,7 @@ class MockCompletionStreamWithContentAfterStopReason:
),
# Additional content after the stop_reason - this simulates the scenario
# where there might be additional content blocks after the main response
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=" Additional content"),

View File

@ -10,7 +10,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterato
)
from litellm.types.utils import (
Delta,
ModelResponse,
ModelResponseStream,
StreamingChoices,
Usage,
ChatCompletionDeltaToolCall,
@ -19,7 +19,7 @@ from litellm.types.utils import (
class MockCompletionStream:
def __init__(self, responses: List[ModelResponse]):
def __init__(self, responses: List[ModelResponseStream]):
self.responses = responses
self.index = 0
@ -44,9 +44,8 @@ class MockCompletionStream:
return response
def construct_text_chunk(text: str) -> ModelResponse:
return ModelResponse(
stream=True,
def construct_text_chunk(text: str) -> ModelResponseStream:
return ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=text),
@ -59,11 +58,10 @@ def construct_text_chunk(text: str) -> ModelResponse:
def construct_split_tool_call(
id: str, function_name: str, function_arg_parts: List[str]
) -> List[ModelResponse]:
) -> List[ModelResponseStream]:
return [
# https://platform.openai.com/docs/guides/function-calling#streaming
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(
@ -82,8 +80,7 @@ def construct_split_tool_call(
],
),
*[
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(
@ -109,8 +106,7 @@ def construct_split_tool_call(
def test_anthropic_stream_wrapper_single_tool_call():
responses = [
*construct_split_tool_call("tooluse_foo", "get_weather", ['{"city":', '"NY"}']),
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="", stop_reason="tool_calls"),
@ -172,8 +168,7 @@ def test_anthropic_stream_wrapper_back_to_back_tool_calls():
responses = [
*construct_split_tool_call("tooluse_foo", "get_weather", ['{"city":', '"NY"}']),
*construct_split_tool_call("tooluse_bar", "get_weather", ['{"city":', '"SF"}']),
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="", stop_reason="tool_calls"),
@ -244,8 +239,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
"tooluse_bar", "get_weather", ['{"city":', '"CHI"}']
),
construct_text_chunk("The weather is not so nice today."),
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="", stop_reason="tool_calls"),

View File

@ -9,31 +9,28 @@ sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import (
AnthropicStreamWrapper,
)
from litellm.types.utils import Delta, ModelResponse, StreamingChoices
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
# Create a simple test
class MockCompletionStream:
def __init__(self):
self.responses = [
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="Hello"), index=0, finish_reason=None
)
],
),
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=" World"), index=0, finish_reason=None
)
],
),
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=""), index=0, finish_reason="stop"
@ -109,16 +106,14 @@ async def test_async_anthropic_sse_wrapper():
class AsyncMockCompletionStream:
def __init__(self):
self.responses = [
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="Hello"), index=0, finish_reason=None
)
],
),
ModelResponse(
stream=True,
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=" World"), index=0, finish_reason=None