mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(anthropic): preserve tool blocks in messages streaming
Keep Anthropic /v1/messages streaming aligned with parallel and interleaved tool call chunks, including dict-backed deltas from current ModelResponse streams. Add regression coverage so tool_result payloads can continue matching the emitted tool_use blocks. Made-with: Cursor
This commit is contained in:
CJYLZS
parent
2f22a1293e
commit
eafed0fca7
4 changed files with 633 additions and 330 deletions
|
|
@ -3,7 +3,16 @@
|
|||
import json
|
||||
import traceback
|
||||
from collections import deque
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, Literal, Optional
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Dict,
|
||||
Iterator,
|
||||
Literal,
|
||||
Optional,
|
||||
cast,
|
||||
)
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -28,21 +37,6 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
TextBlock,
|
||||
)
|
||||
|
||||
sent_first_chunk: bool = False
|
||||
sent_content_block_start: bool = False
|
||||
sent_content_block_finish: bool = False
|
||||
current_content_block_type: Literal["text", "tool_use", "thinking"] = "text"
|
||||
sent_last_message: bool = False
|
||||
holding_chunk: Optional[Any] = None
|
||||
holding_stop_reason_chunk: Optional[Any] = None
|
||||
queued_usage_chunk: bool = False
|
||||
current_content_block_index: int = 0
|
||||
current_content_block_start: ContentBlockContentBlockDict = TextBlock(
|
||||
type="text",
|
||||
text="",
|
||||
)
|
||||
chunk_queue: deque = deque() # Queue for buffering multiple chunks
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
completion_stream: Any,
|
||||
|
|
@ -51,6 +45,23 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
):
|
||||
super().__init__(completion_stream)
|
||||
self.model = model
|
||||
self.sent_first_chunk = False
|
||||
self.sent_content_block_start = False
|
||||
self.sent_content_block_finish = False
|
||||
self.current_content_block_type: Literal["text", "tool_use", "thinking"] = (
|
||||
"text"
|
||||
)
|
||||
self.sent_last_message = False
|
||||
self.holding_chunk = None
|
||||
self.holding_stop_reason_chunk = None
|
||||
self.queued_usage_chunk = False
|
||||
self.current_content_block_index = 0
|
||||
self.current_content_block_start = {"type": "text", "text": ""}
|
||||
self.chunk_queue: deque = deque()
|
||||
self.active_tool_call_index: Optional[int] = None
|
||||
self.active_tool_call_id: Optional[str] = None
|
||||
self.active_tool_call_name: Optional[str] = None
|
||||
self.active_tool_content_block_index: Optional[int] = None
|
||||
# Mapping of truncated tool names to original names (for OpenAI's 64-char limit)
|
||||
self.tool_name_mapping = tool_name_mapping or {}
|
||||
|
||||
|
|
@ -76,138 +87,26 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
)
|
||||
|
||||
def __next__(self):
|
||||
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
|
||||
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
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": self.current_content_block_index,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
}
|
||||
)
|
||||
if self._queue_initial_chunks_if_needed():
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
for chunk in self.completion_stream:
|
||||
if chunk == "None" or chunk is None:
|
||||
raise Exception
|
||||
|
||||
should_start_new_block = self._should_start_new_content_block(chunk)
|
||||
if should_start_new_block:
|
||||
self._increment_content_block_index()
|
||||
|
||||
processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic(
|
||||
response=chunk,
|
||||
current_content_block_index=self.current_content_block_index,
|
||||
)
|
||||
|
||||
if should_start_new_block and not self.sent_content_block_finish:
|
||||
# Queue the sequence: content_block_stop -> content_block_start
|
||||
# For text blocks the trigger chunk is not emitted as a separate
|
||||
# delta because content_block_start carries the information.
|
||||
# For tool_use blocks we must also emit the trigger chunk's delta
|
||||
# when it carries input_json_delta data, because some providers
|
||||
# (e.g. xAI, Gemini) include tool arguments in the same streaming
|
||||
# chunk as the function name/id.
|
||||
|
||||
# 1. Stop current content block
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_stop",
|
||||
"index": max(self.current_content_block_index - 1, 0),
|
||||
}
|
||||
)
|
||||
|
||||
# 2. Start new content block
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": self.current_content_block_index,
|
||||
"content_block": self.current_content_block_start,
|
||||
}
|
||||
)
|
||||
|
||||
# 3. If the trigger chunk carries tool argument data, queue it
|
||||
# so the input_json_delta is not silently dropped.
|
||||
if (
|
||||
processed_chunk.get("type") == "content_block_delta"
|
||||
and isinstance(processed_chunk.get("delta"), dict)
|
||||
and processed_chunk["delta"].get("type") == "input_json_delta"
|
||||
and processed_chunk["delta"].get("partial_json")
|
||||
):
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
|
||||
self.sent_content_block_finish = False
|
||||
self._process_stream_chunk(chunk)
|
||||
if self.chunk_queue:
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
if (
|
||||
processed_chunk["type"] == "message_delta"
|
||||
and self.sent_content_block_finish is False
|
||||
):
|
||||
# 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
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
return self.chunk_queue.popleft()
|
||||
elif self.holding_chunk is not None:
|
||||
self.chunk_queue.append(self.holding_chunk)
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
self.holding_chunk = None
|
||||
return self.chunk_queue.popleft()
|
||||
else:
|
||||
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:
|
||||
self.chunk_queue.append(self.holding_chunk)
|
||||
self.holding_chunk = None
|
||||
|
||||
if not self.sent_last_message:
|
||||
self.sent_last_message = True
|
||||
self.chunk_queue.append({"type": "message_stop"})
|
||||
|
||||
self._flush_stream_end()
|
||||
if self.chunk_queue:
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
raise StopIteration
|
||||
except StopIteration:
|
||||
self._flush_stream_end()
|
||||
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"}
|
||||
raise StopIteration
|
||||
except Exception as e:
|
||||
verbose_logger.error(
|
||||
|
|
@ -216,216 +115,27 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
raise StopAsyncIteration
|
||||
|
||||
async def __anext__(self): # noqa: PLR0915
|
||||
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
|
||||
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
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": self.current_content_block_index,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
}
|
||||
)
|
||||
if self._queue_initial_chunks_if_needed():
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
async for chunk in self.completion_stream:
|
||||
if chunk == "None" or chunk is None:
|
||||
raise Exception
|
||||
|
||||
# Check if we need to start a new content block
|
||||
should_start_new_block = self._should_start_new_content_block(chunk)
|
||||
if should_start_new_block:
|
||||
self._increment_content_block_index()
|
||||
|
||||
processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic(
|
||||
response=chunk,
|
||||
current_content_block_index=self.current_content_block_index,
|
||||
)
|
||||
|
||||
# Check if this is a usage chunk and we have a held stop_reason chunk
|
||||
if (
|
||||
self.holding_stop_reason_chunk is not None
|
||||
and getattr(chunk, "usage", None) is not None
|
||||
):
|
||||
# Merge usage into the held stop_reason chunk
|
||||
merged_chunk = self.holding_stop_reason_chunk.copy()
|
||||
if "delta" not in merged_chunk:
|
||||
merged_chunk["delta"] = {}
|
||||
|
||||
# Add usage to the held chunk
|
||||
uncached_input_tokens = chunk.usage.prompt_tokens or 0
|
||||
if (
|
||||
hasattr(chunk.usage, "prompt_tokens_details")
|
||||
and chunk.usage.prompt_tokens_details
|
||||
):
|
||||
cached_tokens = (
|
||||
getattr(
|
||||
chunk.usage.prompt_tokens_details, "cached_tokens", 0
|
||||
)
|
||||
or 0
|
||||
)
|
||||
uncached_input_tokens -= cached_tokens
|
||||
|
||||
usage_dict: UsageDelta = {
|
||||
"input_tokens": uncached_input_tokens,
|
||||
"output_tokens": chunk.usage.completion_tokens or 0,
|
||||
}
|
||||
# Add cache tokens if available (for prompt caching support)
|
||||
if (
|
||||
hasattr(chunk.usage, "_cache_creation_input_tokens")
|
||||
and chunk.usage._cache_creation_input_tokens > 0
|
||||
):
|
||||
usage_dict["cache_creation_input_tokens"] = (
|
||||
chunk.usage._cache_creation_input_tokens
|
||||
)
|
||||
if (
|
||||
hasattr(chunk.usage, "_cache_read_input_tokens")
|
||||
and chunk.usage._cache_read_input_tokens > 0
|
||||
):
|
||||
usage_dict["cache_read_input_tokens"] = (
|
||||
chunk.usage._cache_read_input_tokens
|
||||
)
|
||||
merged_chunk["usage"] = usage_dict
|
||||
|
||||
# Queue the merged chunk and reset
|
||||
self.chunk_queue.append(merged_chunk)
|
||||
self.queued_usage_chunk = True
|
||||
self.holding_stop_reason_chunk = None
|
||||
self._process_stream_chunk(chunk)
|
||||
if self.chunk_queue:
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
# Check if this processed chunk has a stop_reason - hold it for next chunk
|
||||
|
||||
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
|
||||
# For text blocks the trigger chunk is not emitted as a separate
|
||||
# delta because content_block_start carries the information.
|
||||
# For tool_use blocks we must also emit the trigger chunk's delta
|
||||
# when it carries input_json_delta data, because some providers
|
||||
# (e.g. xAI, Gemini) include tool arguments in the same streaming
|
||||
# chunk as the function name/id.
|
||||
|
||||
# 1. Stop current content block
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_stop",
|
||||
"index": max(self.current_content_block_index - 1, 0),
|
||||
}
|
||||
)
|
||||
|
||||
# 2. Start new content block
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": self.current_content_block_index,
|
||||
"content_block": self.current_content_block_start,
|
||||
}
|
||||
)
|
||||
|
||||
# 3. If the trigger chunk carries tool argument data, queue it
|
||||
# so the input_json_delta is not silently dropped.
|
||||
if (
|
||||
processed_chunk.get("type") == "content_block_delta"
|
||||
and isinstance(processed_chunk.get("delta"), dict)
|
||||
and processed_chunk["delta"].get("type")
|
||||
== "input_json_delta"
|
||||
and processed_chunk["delta"].get("partial_json")
|
||||
):
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
|
||||
# Reset state for new block
|
||||
self.sent_content_block_finish = False
|
||||
|
||||
# Return the first queued item
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
if (
|
||||
processed_chunk["type"] == "message_delta"
|
||||
and self.sent_content_block_finish is False
|
||||
):
|
||||
# Queue both the content_block_stop and the holding chunk
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_stop",
|
||||
"index": self.current_content_block_index,
|
||||
}
|
||||
)
|
||||
self.sent_content_block_finish = True
|
||||
if (
|
||||
processed_chunk.get("delta", {}).get("stop_reason")
|
||||
is not None
|
||||
):
|
||||
self.holding_stop_reason_chunk = processed_chunk
|
||||
else:
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
return self.chunk_queue.popleft()
|
||||
elif self.holding_chunk is not None:
|
||||
# Queue both chunks
|
||||
self.chunk_queue.append(self.holding_chunk)
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
self.holding_chunk = None
|
||||
return self.chunk_queue.popleft()
|
||||
else:
|
||||
# Queue the current chunk
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
# Handle any remaining held chunks after stream ends
|
||||
if not self.queued_usage_chunk:
|
||||
if self.holding_stop_reason_chunk is not None:
|
||||
self.chunk_queue.append(self.holding_stop_reason_chunk)
|
||||
self.holding_stop_reason_chunk = None
|
||||
|
||||
if self.holding_chunk is not None:
|
||||
self.chunk_queue.append(self.holding_chunk)
|
||||
self.holding_chunk = None
|
||||
|
||||
if not self.sent_last_message:
|
||||
self.sent_last_message = True
|
||||
self.chunk_queue.append({"type": "message_stop"})
|
||||
|
||||
# Return queued items if any
|
||||
self._flush_stream_end()
|
||||
if self.chunk_queue:
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
raise StopIteration
|
||||
|
||||
except StopIteration:
|
||||
# Handle any remaining queued chunks before stopping
|
||||
self._flush_stream_end()
|
||||
if self.chunk_queue:
|
||||
return self.chunk_queue.popleft()
|
||||
# Handle any held stop_reason chunk
|
||||
if self.holding_stop_reason_chunk is not None:
|
||||
return self.holding_stop_reason_chunk
|
||||
if not self.sent_last_message:
|
||||
self.sent_last_message = True
|
||||
return {"type": "message_stop"}
|
||||
raise StopAsyncIteration
|
||||
|
||||
def anthropic_sse_wrapper(self) -> Iterator[bytes]:
|
||||
|
|
@ -461,6 +171,262 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
def _increment_content_block_index(self):
|
||||
self.current_content_block_index += 1
|
||||
|
||||
def _queue_initial_chunks_if_needed(self) -> bool:
|
||||
if self.sent_first_chunk is False:
|
||||
self.sent_first_chunk = True
|
||||
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 True
|
||||
|
||||
if self.sent_content_block_start is False:
|
||||
self.sent_content_block_start = True
|
||||
self.sent_content_block_finish = False
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": self.current_content_block_index,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
}
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _chunk_has_tool_calls(self, chunk: "ModelResponseStream") -> bool:
|
||||
for choice in chunk.choices:
|
||||
delta = choice.delta
|
||||
tool_calls = (
|
||||
delta.get("tool_calls")
|
||||
if isinstance(delta, dict)
|
||||
else getattr(delta, "tool_calls", None)
|
||||
)
|
||||
if tool_calls is not None and len(tool_calls) > 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _queue_content_block_stop(self, index: int) -> None:
|
||||
self.chunk_queue.append({"type": "content_block_stop", "index": index})
|
||||
|
||||
def _close_active_tool_call(self) -> None:
|
||||
if self.active_tool_content_block_index is None:
|
||||
return
|
||||
|
||||
self._queue_content_block_stop(self.active_tool_content_block_index)
|
||||
self.sent_content_block_finish = True
|
||||
self.active_tool_call_index = None
|
||||
self.active_tool_call_id = None
|
||||
self.active_tool_call_name = None
|
||||
self.active_tool_content_block_index = None
|
||||
|
||||
def _build_usage_dict(self, chunk: "ModelResponseStream") -> UsageDelta:
|
||||
chunk_usage = getattr(chunk, "usage", None)
|
||||
assert chunk_usage is not None
|
||||
|
||||
uncached_input_tokens = chunk_usage.prompt_tokens or 0
|
||||
if (
|
||||
hasattr(chunk_usage, "prompt_tokens_details")
|
||||
and chunk_usage.prompt_tokens_details
|
||||
):
|
||||
cached_tokens = (
|
||||
getattr(chunk_usage.prompt_tokens_details, "cached_tokens", 0) or 0
|
||||
)
|
||||
uncached_input_tokens -= cached_tokens
|
||||
|
||||
usage_dict: UsageDelta = {
|
||||
"input_tokens": uncached_input_tokens,
|
||||
"output_tokens": chunk_usage.completion_tokens or 0,
|
||||
}
|
||||
if (
|
||||
hasattr(chunk_usage, "_cache_creation_input_tokens")
|
||||
and chunk_usage._cache_creation_input_tokens > 0
|
||||
):
|
||||
usage_dict["cache_creation_input_tokens"] = (
|
||||
chunk_usage._cache_creation_input_tokens
|
||||
)
|
||||
if (
|
||||
hasattr(chunk_usage, "_cache_read_input_tokens")
|
||||
and chunk_usage._cache_read_input_tokens > 0
|
||||
):
|
||||
usage_dict["cache_read_input_tokens"] = chunk_usage._cache_read_input_tokens
|
||||
return usage_dict
|
||||
|
||||
def _process_tool_call_chunk(self, chunk: "ModelResponseStream") -> None:
|
||||
from .transformation import LiteLLMAnthropicMessagesAdapter
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
tool_calls = adapter.iter_streaming_tool_calls(chunk.choices) # type: ignore[arg-type]
|
||||
if len(tool_calls) == 0:
|
||||
return
|
||||
|
||||
if (
|
||||
self.current_content_block_type != "tool_use"
|
||||
and not self.sent_content_block_finish
|
||||
):
|
||||
self._queue_content_block_stop(self.current_content_block_index)
|
||||
self.sent_content_block_finish = True
|
||||
|
||||
self.current_content_block_type = "tool_use"
|
||||
|
||||
for tool_call in tool_calls:
|
||||
tool_call_index = tool_call.index or 0
|
||||
tool_call_id = tool_call.id
|
||||
tool_call_name = None
|
||||
if tool_call.function is not None:
|
||||
tool_call_name = tool_call.function.name
|
||||
tool_call_starts_new_block = bool(tool_call_name)
|
||||
|
||||
tool_call_matches_active = (
|
||||
self.active_tool_call_index == tool_call_index
|
||||
and (tool_call_id is None or tool_call_id == self.active_tool_call_id)
|
||||
and (
|
||||
tool_call_name is None
|
||||
or tool_call_name == self.active_tool_call_name
|
||||
)
|
||||
)
|
||||
|
||||
if (
|
||||
self.active_tool_content_block_index is None
|
||||
or tool_call_starts_new_block
|
||||
or not tool_call_matches_active
|
||||
):
|
||||
if self.active_tool_content_block_index is not None:
|
||||
self._close_active_tool_call()
|
||||
|
||||
self._increment_content_block_index()
|
||||
anthropic_index = self.current_content_block_index
|
||||
self.active_tool_call_index = tool_call_index
|
||||
self.active_tool_call_id = tool_call_id
|
||||
self.active_tool_call_name = tool_call_name
|
||||
self.active_tool_content_block_index = anthropic_index
|
||||
content_block = (
|
||||
adapter.translate_streaming_tool_call_to_anthropic_content_block(
|
||||
tool_call
|
||||
)
|
||||
)
|
||||
if tool_call_name:
|
||||
content_block_dict = cast(Dict[str, Any], content_block)
|
||||
content_block_dict["name"] = self.tool_name_mapping.get(
|
||||
tool_call_name, tool_call_name
|
||||
)
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": anthropic_index,
|
||||
"content_block": content_block,
|
||||
}
|
||||
)
|
||||
else:
|
||||
anthropic_index = self.active_tool_content_block_index
|
||||
|
||||
content_block_delta = (
|
||||
adapter.translate_streaming_tool_call_to_anthropic_delta(tool_call)
|
||||
)
|
||||
if content_block_delta is not None:
|
||||
self.chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": anthropic_index,
|
||||
"delta": content_block_delta,
|
||||
}
|
||||
)
|
||||
|
||||
self.sent_content_block_finish = False
|
||||
|
||||
def _process_non_tool_chunk(self, chunk: "ModelResponseStream") -> None:
|
||||
from .transformation import LiteLLMAnthropicMessagesAdapter
|
||||
|
||||
if self.active_tool_content_block_index is not None:
|
||||
self._close_active_tool_call()
|
||||
|
||||
should_start_new_block = self._should_start_new_content_block(chunk)
|
||||
if should_start_new_block:
|
||||
self._increment_content_block_index()
|
||||
|
||||
processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic(
|
||||
response=cast(Any, chunk),
|
||||
current_content_block_index=self.current_content_block_index,
|
||||
)
|
||||
|
||||
if (
|
||||
self.holding_stop_reason_chunk is not None
|
||||
and getattr(chunk, "usage", None) is not None
|
||||
):
|
||||
merged_chunk = self.holding_stop_reason_chunk.copy()
|
||||
if "delta" not in merged_chunk:
|
||||
merged_chunk["delta"] = {}
|
||||
merged_chunk["usage"] = self._build_usage_dict(chunk)
|
||||
self.chunk_queue.append(merged_chunk)
|
||||
self.queued_usage_chunk = True
|
||||
self.holding_stop_reason_chunk = None
|
||||
return
|
||||
|
||||
if should_start_new_block:
|
||||
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
|
||||
|
||||
if processed_chunk["type"] == "message_delta":
|
||||
if self.sent_content_block_finish is False:
|
||||
self._queue_content_block_stop(self.current_content_block_index)
|
||||
self.sent_content_block_finish = True
|
||||
|
||||
if (
|
||||
processed_chunk.get("delta", {}).get("stop_reason") is not None
|
||||
and getattr(chunk, "usage", None) is None
|
||||
):
|
||||
self.holding_stop_reason_chunk = processed_chunk
|
||||
else:
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
return
|
||||
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
|
||||
def _process_stream_chunk(self, chunk: "ModelResponseStream") -> None:
|
||||
if chunk == "None" or chunk is None:
|
||||
raise Exception
|
||||
|
||||
self.queued_usage_chunk = False
|
||||
|
||||
if self._chunk_has_tool_calls(chunk):
|
||||
self._process_tool_call_chunk(chunk)
|
||||
return
|
||||
|
||||
self._process_non_tool_chunk(chunk)
|
||||
|
||||
def _flush_stream_end(self) -> None:
|
||||
if self.active_tool_content_block_index is not None:
|
||||
self._close_active_tool_call()
|
||||
|
||||
if self.holding_stop_reason_chunk is not None:
|
||||
self.chunk_queue.append(self.holding_stop_reason_chunk)
|
||||
self.holding_stop_reason_chunk = None
|
||||
|
||||
if self.holding_chunk is not None:
|
||||
self.chunk_queue.append(self.holding_chunk)
|
||||
self.holding_chunk = None
|
||||
|
||||
if not self.sent_last_message:
|
||||
self.sent_last_message = True
|
||||
self.chunk_queue.append({"type": "message_stop"})
|
||||
|
||||
def _should_start_new_content_block(self, chunk: "ModelResponseStream") -> bool:
|
||||
"""
|
||||
Determine if we should start a new content block based on the processed chunk.
|
||||
|
|
|
|||
|
|
@ -117,7 +117,13 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionToolParamFunctionChunk,
|
||||
ChatCompletionUserMessage,
|
||||
)
|
||||
from litellm.types.utils import Choices, ModelResponse, StreamingChoices, Usage
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionDeltaToolCall,
|
||||
Choices,
|
||||
ModelResponse,
|
||||
StreamingChoices,
|
||||
Usage,
|
||||
)
|
||||
|
||||
from .streaming_iterator import AnthropicStreamWrapper
|
||||
|
||||
|
|
@ -1406,6 +1412,53 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
|
||||
return "text", TextBlock(type="text", text="")
|
||||
|
||||
def iter_streaming_tool_calls(
|
||||
self, choices: List[Union[OpenAIStreamingChoice, StreamingChoices]]
|
||||
) -> List[ChatCompletionDeltaToolCall]:
|
||||
tool_calls: List[ChatCompletionDeltaToolCall] = []
|
||||
for choice in choices:
|
||||
delta = choice.delta
|
||||
choice_tool_calls = (
|
||||
delta.get("tool_calls")
|
||||
if isinstance(delta, dict)
|
||||
else getattr(delta, "tool_calls", None)
|
||||
)
|
||||
if choice_tool_calls is not None:
|
||||
for tool_call in choice_tool_calls:
|
||||
if isinstance(tool_call, dict):
|
||||
tool_calls.append(ChatCompletionDeltaToolCall(**tool_call))
|
||||
else:
|
||||
tool_calls.append(tool_call)
|
||||
return tool_calls
|
||||
|
||||
def translate_streaming_tool_call_to_anthropic_content_block(
|
||||
self, tool_call: ChatCompletionDeltaToolCall
|
||||
) -> "ContentBlockContentBlockDict":
|
||||
from litellm._uuid import uuid
|
||||
from litellm.types.llms.anthropic import ToolUseBlock
|
||||
|
||||
function_name = ""
|
||||
if tool_call.function is not None and tool_call.function.name is not None:
|
||||
function_name = tool_call.function.name
|
||||
|
||||
return ToolUseBlock(
|
||||
type="tool_use",
|
||||
id=tool_call.id or str(uuid.uuid4()),
|
||||
name=function_name,
|
||||
input={}, # type: ignore[typeddict-item]
|
||||
)
|
||||
|
||||
def translate_streaming_tool_call_to_anthropic_delta(
|
||||
self, tool_call: ChatCompletionDeltaToolCall
|
||||
) -> Optional[ContentJsonBlockDelta]:
|
||||
if tool_call.function is None or tool_call.function.arguments in (None, ""):
|
||||
return None
|
||||
|
||||
return ContentJsonBlockDelta(
|
||||
type="input_json_delta",
|
||||
partial_json=tool_call.function.arguments,
|
||||
)
|
||||
|
||||
def _translate_streaming_openai_chunk_to_anthropic(
|
||||
self, choices: List[Union[OpenAIStreamingChoice, StreamingChoices]]
|
||||
) -> Tuple[
|
||||
|
|
|
|||
|
|
@ -74,6 +74,86 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block():
|
|||
}
|
||||
|
||||
|
||||
def test_iter_streaming_tool_calls_multiple_tool_calls():
|
||||
choices = [
|
||||
StreamingChoices(
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
delta=Delta(
|
||||
provider_specific_fields=None,
|
||||
content=None,
|
||||
role="assistant",
|
||||
function_call=None,
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call_weather",
|
||||
function=Function(
|
||||
arguments='{"location": "Boston"}', name="get_weather"
|
||||
),
|
||||
type="function",
|
||||
index=0,
|
||||
),
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call_time",
|
||||
function=Function(
|
||||
arguments='{"timezone": "UTC"}', name="get_time"
|
||||
),
|
||||
type="function",
|
||||
index=1,
|
||||
),
|
||||
],
|
||||
audio=None,
|
||||
),
|
||||
logprobs=None,
|
||||
)
|
||||
]
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
tool_calls = adapter.iter_streaming_tool_calls(choices=choices)
|
||||
|
||||
assert [tool_call.id for tool_call in tool_calls] == ["call_weather", "call_time"]
|
||||
assert [tool_call.index for tool_call in tool_calls] == [0, 1]
|
||||
|
||||
|
||||
def test_translate_streaming_tool_call_helpers():
|
||||
tool_call = ChatCompletionDeltaToolCall(
|
||||
id="call_time",
|
||||
function=Function(arguments='{"timezone":"UTC"}', name="get_time"),
|
||||
type="function",
|
||||
index=1,
|
||||
)
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
content_block = adapter.translate_streaming_tool_call_to_anthropic_content_block(
|
||||
tool_call
|
||||
)
|
||||
content_delta = adapter.translate_streaming_tool_call_to_anthropic_delta(tool_call)
|
||||
|
||||
assert content_block == {
|
||||
"type": "tool_use",
|
||||
"id": "call_time",
|
||||
"name": "get_time",
|
||||
"input": {},
|
||||
}
|
||||
assert content_delta == {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": '{"timezone":"UTC"}',
|
||||
}
|
||||
|
||||
|
||||
def test_translate_streaming_tool_call_to_anthropic_delta_empty_arguments():
|
||||
tool_call = ChatCompletionDeltaToolCall(
|
||||
id="call_weather",
|
||||
function=Function(arguments="", name="get_weather"),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
|
||||
assert adapter.translate_streaming_tool_call_to_anthropic_delta(tool_call) is None
|
||||
|
||||
|
||||
def test_translate_streaming_openai_chunk_to_anthropic_thinking_content_block():
|
||||
choices = [
|
||||
StreamingChoices(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import os
|
||||
import sys
|
||||
from typing import List
|
||||
from typing import List, Sequence
|
||||
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
|
@ -10,6 +10,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterato
|
|||
)
|
||||
from litellm.types.utils import (
|
||||
Delta,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
Usage,
|
||||
|
|
@ -19,7 +20,7 @@ from litellm.types.utils import (
|
|||
|
||||
|
||||
class MockCompletionStream:
|
||||
def __init__(self, responses: List[ModelResponseStream]):
|
||||
def __init__(self, responses: Sequence[ModelResponseStream | ModelResponse]):
|
||||
self.responses = responses
|
||||
self.index = 0
|
||||
|
||||
|
|
@ -103,6 +104,35 @@ def construct_split_tool_call(
|
|||
]
|
||||
|
||||
|
||||
def construct_parallel_tool_call_start(
|
||||
tool_calls: List[tuple[str, str, str]],
|
||||
) -> ModelResponse:
|
||||
return ModelResponse(
|
||||
stream=True,
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id=tool_id,
|
||||
function=Function(arguments=arguments, name=function_name),
|
||||
index=tool_index,
|
||||
type="function",
|
||||
)
|
||||
for tool_index, (
|
||||
tool_id,
|
||||
function_name,
|
||||
arguments,
|
||||
) in enumerate(tool_calls)
|
||||
]
|
||||
),
|
||||
index=0,
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_anthropic_stream_wrapper_single_tool_call():
|
||||
responses = [
|
||||
*construct_split_tool_call("tooluse_foo", "get_weather", ['{"city":', '"NY"}']),
|
||||
|
|
@ -279,6 +309,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
|
|||
"content_block_delta", # "NY"}
|
||||
"content_block_stop", # End of first tool_use content block
|
||||
"content_block_start", # "The weather is nice today"
|
||||
"content_block_delta", # "The weather is nice today."
|
||||
"content_block_stop",
|
||||
"content_block_start", # Start of second tool_use content block
|
||||
"content_block_delta", # {"city":
|
||||
|
|
@ -289,6 +320,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
|
|||
"content_block_delta", # " CHI"}
|
||||
"content_block_stop", # End of third tool_use content block
|
||||
"content_block_start", # "The weather is not so nice today"
|
||||
"content_block_delta", # "The weather is not so nice today."
|
||||
"content_block_stop",
|
||||
"message_delta", # Stop reason with merged usage
|
||||
"message_stop", # Final message stop
|
||||
|
|
@ -307,3 +339,175 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
|
|||
get_weather_calls += 1
|
||||
|
||||
assert get_weather_calls == 3
|
||||
|
||||
|
||||
def test_anthropic_stream_wrapper_parallel_tool_calls_same_chunk():
|
||||
responses = [
|
||||
construct_parallel_tool_call_start(
|
||||
[
|
||||
("tooluse_foo", "get_weather", '{"city":"NY"}'),
|
||||
("tooluse_bar", "get_time", '{"timezone":"UTC"}'),
|
||||
]
|
||||
),
|
||||
ModelResponse(
|
||||
stream=True,
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content="", stop_reason="tool_calls"),
|
||||
index=0,
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
usage=Usage(prompt_tokens=230, completion_tokens=65, total_tokens=295),
|
||||
),
|
||||
]
|
||||
|
||||
wrapper = AnthropicStreamWrapper(
|
||||
completion_stream=MockCompletionStream(responses),
|
||||
model="sonnet-4-5",
|
||||
)
|
||||
|
||||
chunks = []
|
||||
chunk_types = []
|
||||
|
||||
for chunk in wrapper:
|
||||
chunks.append(chunk)
|
||||
chunk_types.append(chunk.get("type"))
|
||||
|
||||
expected_types = [
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_stop",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
]
|
||||
|
||||
assert expected_types == chunk_types
|
||||
|
||||
tool_starts = [
|
||||
chunk
|
||||
for chunk in chunks
|
||||
if chunk.get("type") == "content_block_start"
|
||||
and chunk["content_block"]["type"] == "tool_use"
|
||||
]
|
||||
|
||||
assert [chunk["content_block"]["id"] for chunk in tool_starts] == [
|
||||
"tooluse_foo",
|
||||
"tooluse_bar",
|
||||
]
|
||||
assert [chunk["content_block"]["name"] for chunk in tool_starts] == [
|
||||
"get_weather",
|
||||
"get_time",
|
||||
]
|
||||
|
||||
tool_deltas = [
|
||||
chunk
|
||||
for chunk in chunks
|
||||
if chunk.get("type") == "content_block_delta"
|
||||
and chunk.get("delta", {}).get("type") == "input_json_delta"
|
||||
]
|
||||
assert [chunk["delta"]["partial_json"] for chunk in tool_deltas] == [
|
||||
'{"city":"NY"}',
|
||||
'{"timezone":"UTC"}',
|
||||
]
|
||||
|
||||
|
||||
def test_anthropic_stream_wrapper_parallel_tool_call_index_continuations():
|
||||
responses = [
|
||||
construct_parallel_tool_call_start(
|
||||
[
|
||||
("tooluse_foo", "get_weather", '{"city":"NY"}'),
|
||||
("tooluse_bar", "get_time", ""),
|
||||
]
|
||||
),
|
||||
ModelResponse(
|
||||
stream=True,
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id=None,
|
||||
function=Function(
|
||||
arguments='{"timezone":"UTC"}', name=None
|
||||
),
|
||||
index=1,
|
||||
type=None,
|
||||
)
|
||||
]
|
||||
),
|
||||
index=0,
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
),
|
||||
ModelResponse(
|
||||
stream=True,
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id=None,
|
||||
function=Function(arguments="", name=None),
|
||||
index=1,
|
||||
type=None,
|
||||
)
|
||||
]
|
||||
),
|
||||
index=0,
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
),
|
||||
ModelResponse(
|
||||
stream=True,
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content="", stop_reason="tool_calls"),
|
||||
index=0,
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
usage=Usage(prompt_tokens=230, completion_tokens=65, total_tokens=295),
|
||||
),
|
||||
]
|
||||
|
||||
wrapper = AnthropicStreamWrapper(
|
||||
completion_stream=MockCompletionStream(responses),
|
||||
model="sonnet-4-5",
|
||||
)
|
||||
|
||||
chunks = list(wrapper)
|
||||
|
||||
tool_starts = [
|
||||
chunk
|
||||
for chunk in chunks
|
||||
if chunk.get("type") == "content_block_start"
|
||||
and chunk["content_block"]["type"] == "tool_use"
|
||||
]
|
||||
|
||||
assert [chunk["content_block"]["id"] for chunk in tool_starts] == [
|
||||
"tooluse_foo",
|
||||
"tooluse_bar",
|
||||
]
|
||||
|
||||
tool_deltas = [
|
||||
chunk
|
||||
for chunk in chunks
|
||||
if chunk.get("type") == "content_block_delta"
|
||||
and chunk.get("delta", {}).get("type") == "input_json_delta"
|
||||
]
|
||||
|
||||
assert [
|
||||
(chunk["index"], chunk["delta"]["partial_json"]) for chunk in tool_deltas
|
||||
] == [
|
||||
(1, '{"city":"NY"}'),
|
||||
(2, '{"timezone":"UTC"}'),
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue