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:
commit
2c5c38333d
@ -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
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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"),
|
||||
|
||||
@ -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"),
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user