Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 24 additions & 17 deletions src/any_llm/any_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
MessageStopEvent,
MessageStreamEvent,
ParsedMessage,
StopReason,
)
from any_llm.types.provider import ProviderMetadata
from any_llm.types.responses import (
Expand Down Expand Up @@ -785,29 +786,35 @@ async def _amessages(

async def convert_stream() -> AsyncIterator[MessageStreamEvent]:
state = StreamingState()
emitted_stop = False
async for chunk in result:
events = chat_completion_chunk_to_message_stream_events(chunk, state)
for event in events:
if isinstance(event, MessageStopEvent):
emitted_stop = True
yield event
# Some providers don't send a final chunk with finish_reason,
# so ensure the stream always ends with message_stop.
if state.started and not emitted_stop:
if state.current_block_type is not None:
yield ContentBlockStopEvent(
type="content_block_stop",
index=state.current_block_index,
)
yield MessageDeltaEvent(

def usage_delta(stop_reason: StopReason | None) -> MessageDeltaEvent:
return MessageDeltaEvent(
type="message_delta",
delta=MessageDelta(stop_reason="end_turn"),
delta=MessageDelta(stop_reason=stop_reason),
usage=MessageDeltaUsage(
output_tokens=state.output_tokens,
input_tokens=state.input_tokens,
cache_read_input_tokens=state.cache_read_input_tokens or None,
),
)

try:
async for chunk in result:
for event in chat_completion_chunk_to_message_stream_events(chunk, state):
yield event
except Exception:
# Flush the usage accumulated so far before re-raising, so a mid-stream failure still reports tokens.
if state.started:
yield usage_delta(state.stop_reason)
raise
# Emit the closing events after the full stream is consumed so trailing-chunk usage is included.
if state.started:
if state.current_block_type is not None:
yield ContentBlockStopEvent(
type="content_block_stop",
index=state.current_block_index,
)
yield usage_delta(state.stop_reason or "end_turn")
yield MessageStopEvent(type="message_stop")

return convert_stream()
Expand Down
46 changes: 9 additions & 37 deletions src/any_llm/utils/messages_compat.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,8 @@
ContentBlockStartEvent,
ContentBlockStopEvent,
InputJSONDelta,
MessageDelta,
MessageDeltaEvent,
MessageDeltaUsage,
MessageResponse,
MessageStartEvent,
MessageStopEvent,
MessageUsage,
StopReason,
TextBlock,
Expand Down Expand Up @@ -316,41 +312,32 @@ def __init__(self) -> None:
self.model = "unknown"
self.input_tokens = 0
self.output_tokens = 0
self.cache_read_input_tokens = 0
self.stop_reason: StopReason | None = None
self.tool_call_id: str | None = None
self.tool_call_name: str | None = None


def chat_completion_chunk_to_message_stream_events(
chunk: ChatCompletionChunk,
state: StreamingState,
) -> list[
MessageStartEvent
| MessageDeltaEvent
| MessageStopEvent
| ContentBlockStartEvent
| ContentBlockDeltaEvent
| ContentBlockStopEvent
]:
) -> list[MessageStartEvent | ContentBlockStartEvent | ContentBlockDeltaEvent | ContentBlockStopEvent]:
"""Convert a ChatCompletionChunk to a list of MessageStreamEvents.

This is stateful: it tracks the current content block index and type to emit
the correct lifecycle events (start/delta/stop).
"""
events: list[
MessageStartEvent
| MessageDeltaEvent
| MessageStopEvent
| ContentBlockStartEvent
| ContentBlockDeltaEvent
| ContentBlockStopEvent
] = []
events: list[MessageStartEvent | ContentBlockStartEvent | ContentBlockDeltaEvent | ContentBlockStopEvent] = []
state.model = chunk.model

if chunk.usage:
if chunk.usage.prompt_tokens:
state.input_tokens = chunk.usage.prompt_tokens
if chunk.usage.completion_tokens:
state.output_tokens = chunk.usage.completion_tokens
prompt_details = chunk.usage.prompt_tokens_details
if prompt_details and prompt_details.cached_tokens:
state.cache_read_input_tokens = prompt_details.cached_tokens
Comment on lines 333 to +340

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

Preserve zero usage values instead of treating them as absent.

Line 334, Line 336 and Line 339 skip valid zero values, so a later usage chunk can leave stale non-zero state behind. Check for None instead.

Proposed fix
-        if chunk.usage.prompt_tokens:
+        if chunk.usage.prompt_tokens is not None:
             state.input_tokens = chunk.usage.prompt_tokens
-        if chunk.usage.completion_tokens:
+        if chunk.usage.completion_tokens is not None:
             state.output_tokens = chunk.usage.completion_tokens
         prompt_details = chunk.usage.prompt_tokens_details
-        if prompt_details and prompt_details.cached_tokens:
+        if prompt_details and prompt_details.cached_tokens is not None:
             state.cache_read_input_tokens = prompt_details.cached_tokens
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if chunk.usage:
if chunk.usage.prompt_tokens:
state.input_tokens = chunk.usage.prompt_tokens
if chunk.usage.completion_tokens:
state.output_tokens = chunk.usage.completion_tokens
prompt_details = chunk.usage.prompt_tokens_details
if prompt_details and prompt_details.cached_tokens:
state.cache_read_input_tokens = prompt_details.cached_tokens
if chunk.usage:
if chunk.usage.prompt_tokens is not None:
state.input_tokens = chunk.usage.prompt_tokens
if chunk.usage.completion_tokens is not None:
state.output_tokens = chunk.usage.completion_tokens
prompt_details = chunk.usage.prompt_tokens_details
if prompt_details and prompt_details.cached_tokens is not None:
state.cache_read_input_tokens = prompt_details.cached_tokens
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/any_llm/utils/messages_compat.py` around lines 333 - 340, Preserve valid
zero token usage values in the message usage handling logic by updating the
checks in the chunk processing block inside messages_compat.py to distinguish
None from 0. In the code that updates state.input_tokens, state.output_tokens,
and state.cache_read_input_tokens, use explicit None checks on
chunk.usage.prompt_tokens, chunk.usage.completion_tokens, and
prompt_details.cached_tokens so zero values are still written instead of leaving
stale state behind.


if not state.started:
state.started = True
Expand Down Expand Up @@ -444,29 +431,14 @@ def chat_completion_chunk_to_message_stream_events(

if choice.finish_reason:
_close_current_block(state, events)
stop_reason = _finish_reason_to_stop_reason(choice.finish_reason)
events.append(
MessageDeltaEvent(
type="message_delta",
delta=MessageDelta(stop_reason=stop_reason),
usage=MessageDeltaUsage(output_tokens=state.output_tokens, input_tokens=state.input_tokens),
)
)
events.append(MessageStopEvent(type="message_stop"))
state.stop_reason = _finish_reason_to_stop_reason(choice.finish_reason)

return events


def _close_current_block(
state: StreamingState,
events: list[
MessageStartEvent
| MessageDeltaEvent
| MessageStopEvent
| ContentBlockStartEvent
| ContentBlockDeltaEvent
| ContentBlockStopEvent
],
events: list[MessageStartEvent | ContentBlockStartEvent | ContentBlockDeltaEvent | ContentBlockStopEvent],
) -> None:
"""Emit a content_block_stop event for the current block if one is open."""
if state.current_block_type is not None:
Expand Down
Loading