scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-01 02:02:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 92

refactor(anthropic): extract usage-merge helper; guard empty slice-only compaction result

Browse source 
- Extract the duplicated hold-and-merge usage logic from the sync __next__ and
  async __anext__ paths into a shared _merge_usage_into_held_stop_reason_chunk
  helper so the subtle cache-token / context_management attachment lives in
  exactly one place.
- In the compact_20260112 slice-only path, fall back to _select_last_user_question
  when _strip_compaction_blocks produces an empty list (e.g. messages ending on
  an assistant turn whose only content was the compaction block) so the
  downstream API never receives an empty messages array.

Co-authored-by: Yassin Kortam <yassin@berri.ai>
This commit is contained in:
Cursor Agent 2026-05-26 13:34:44 +00:00
parent b8fd48ae01
commit 6cf22cf9fd
No known key found for this signature in database
2 changed files with 59 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

141
litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py
View file

@ -73,6 +73,55 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
# deque and corrupt each other's SSE event order.
self.chunk_queue: deque = deque()
def _merge_usage_into_held_stop_reason_chunk(self, chunk: Any) -> Dict[str, Any]:
"""Merge usage data from ``chunk`` into the held ``message_delta`` chunk.
Shared by both the sync ``__next__`` and async ``__anext__`` paths so
the subtle hold-and-merge logic (cache tokens, ``context_management``
attachment, ``UsageDelta`` shape) lives in exactly one place.
Caller is responsible for managing ``self.holding_stop_reason_chunk``
and ``self.queued_usage_chunk`` state and for queuing the returned
merged chunk.
"""
assert self.holding_stop_reason_chunk is not None
merged_chunk = self.holding_stop_reason_chunk.copy()
if "delta" not in merged_chunk:
merged_chunk["delta"] = {}
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
merged_chunk["usage"] = usage_dict
if self.applied_edits and "context_management" not in merged_chunk:
merged_chunk["context_management"] = ContextManagementResponse(
applied_edits=list(self.applied_edits)
)
return merged_chunk
def _create_initial_usage_delta(self) -> UsageDelta:
"""
Create the initial UsageDelta for the message_start event.
@ -157,51 +206,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
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
if self.applied_edits and "context_management" not in merged_chunk:
merged_chunk["context_management"] = ContextManagementResponse(
applied_edits=list(self.applied_edits)
)
# Queue the merged chunk and reset
merged_chunk = self._merge_usage_into_held_stop_reason_chunk(chunk)
self.chunk_queue.append(merged_chunk)
self.queued_usage_chunk = True
self.holding_stop_reason_chunk = None
@ -381,51 +386,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
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
if self.applied_edits and "context_management" not in merged_chunk:
merged_chunk["context_management"] = ContextManagementResponse(
applied_edits=list(self.applied_edits)
)
# Queue the merged chunk and reset
merged_chunk = self._merge_usage_into_held_stop_reason_chunk(chunk)
self.chunk_queue.append(merged_chunk)
self.queued_usage_chunk = True
self.holding_stop_reason_chunk = None

8
litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
View file

@ -465,6 +465,14 @@ async def apply_compact_20260112(
if current_tokens <= trigger_tokens:
# Slice-only path. If Phase A fired we still slice + system-prefix.
# Guard against an empty downstream message list: this happens when the
# sliced result was a single assistant turn whose only block was the
# compaction block, which ``_strip_compaction_blocks`` then drops. The
# downstream API rejects an empty messages array, so fall back to
# picking a real user question (or a synthetic continuation prompt) —
# mirroring the full-summary path's behavior below.
if not downstream_messages:
downstream_messages = _select_last_user_question(effective_messages)
return PolyfillResult(
messages=downstream_messages,
system=augmented_system,