Skip to content

fix(streaming): initialize usage when message_start omits it #1815

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chenlichao wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(streaming): initialize usage when message_start omits it #1815

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
42 changes: 27 additions & 15 deletions src/anthropic/lib/streaming/_beta_messages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from anthropic.types.beta.beta_tool_use_block import BetaToolUseBlock
from anthropic.types.beta.beta_mcp_tool_use_block import BetaMCPToolUseBlock
from anthropic.types.beta.beta_server_tool_use_block import BetaServerToolUseBlock
from anthropic.types.usage import Usage

from ..._types import NOT_GIVEN, NotGiven
from ..._utils import consume_sync_iterator, consume_async_iterator
Expand Down Expand Up @@ -545,21 +546,32 @@ def accumulate_event(
current_snapshot.stop_sequence = event.delta.stop_sequence
if event.delta.stop_details is not None:
current_snapshot.stop_details = event.delta.stop_details
current_snapshot.usage.output_tokens = event.usage.output_tokens
current_snapshot.context_management = event.context_management
if current_snapshot.usage is None: # pyright: ignore[reportUnnecessaryComparison] # noqa: E501
# `message_start` may omit usage (see the streaming docs), in which
# case the snapshot has no usage yet. Initialize it from the delta
# so the final message still carries token counts, and tolerate
# streams that never supply usage.
current_snapshot.usage = Usage(
input_tokens=event.usage.input_tokens or 0,
output_tokens=event.usage.output_tokens,
)
else:
current_snapshot.usage.output_tokens = event.usage.output_tokens

# Update other usage fields if they exist in the event
if event.usage.input_tokens is not None:
current_snapshot.usage.input_tokens = event.usage.input_tokens
if event.usage.cache_creation_input_tokens is not None:
current_snapshot.usage.cache_creation_input_tokens = event.usage.cache_creation_input_tokens
if event.usage.cache_read_input_tokens is not None:
current_snapshot.usage.cache_read_input_tokens = event.usage.cache_read_input_tokens
if event.usage.server_tool_use is not None:
current_snapshot.usage.server_tool_use = event.usage.server_tool_use
if event.usage.iterations is not None:
current_snapshot.usage.iterations = event.usage.iterations
if event.usage.fallback_credit is not None:
current_snapshot.usage.fallback_credit = event.usage.fallback_credit

# Update other usage fields if they exist in the event
if event.usage.input_tokens is not None:
current_snapshot.usage.input_tokens = event.usage.input_tokens
if event.usage.cache_creation_input_tokens is not None:
current_snapshot.usage.cache_creation_input_tokens = event.usage.cache_creation_input_tokens
if event.usage.cache_read_input_tokens is not None:
current_snapshot.usage.cache_read_input_tokens = event.usage.cache_read_input_tokens
if event.usage.server_tool_use is not None:
current_snapshot.usage.server_tool_use = event.usage.server_tool_use
if event.usage.iterations is not None:
current_snapshot.usage.iterations = event.usage.iterations
if event.usage.fallback_credit is not None:
current_snapshot.usage.fallback_credit = event.usage.fallback_credit
current_snapshot.context_management = event.context_management

return current_snapshot
33 changes: 22 additions & 11 deletions src/anthropic/lib/streaming/_messages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

from anthropic.types.tool_use_block import ToolUseBlock
from anthropic.types.server_tool_use_block import ServerToolUseBlock
from anthropic.types.usage import Usage

from ._types import (
TextEvent,
Expand Down Expand Up @@ -505,16 +506,26 @@ def accumulate_event(
current_snapshot.stop_sequence = event.delta.stop_sequence
if event.delta.stop_details is not None:
current_snapshot.stop_details = event.delta.stop_details
current_snapshot.usage.output_tokens = event.usage.output_tokens

# Update other usage fields if they exist in the event
if event.usage.input_tokens is not None:
current_snapshot.usage.input_tokens = event.usage.input_tokens
if event.usage.cache_creation_input_tokens is not None:
current_snapshot.usage.cache_creation_input_tokens = event.usage.cache_creation_input_tokens
if event.usage.cache_read_input_tokens is not None:
current_snapshot.usage.cache_read_input_tokens = event.usage.cache_read_input_tokens
if event.usage.server_tool_use is not None:
current_snapshot.usage.server_tool_use = event.usage.server_tool_use
if current_snapshot.usage is None: # pyright: ignore[reportUnnecessaryComparison] # noqa: E501
# `message_start` may omit usage (see the streaming docs), in which
# case the snapshot has no usage yet. Initialize it from the delta
# so the final message still carries token counts, and tolerate
# streams that never supply usage.
current_snapshot.usage = Usage(
input_tokens=event.usage.input_tokens or 0,
output_tokens=event.usage.output_tokens,
)
else:
current_snapshot.usage.output_tokens = event.usage.output_tokens

# Update other usage fields if they exist in the event
if event.usage.input_tokens is not None:
current_snapshot.usage.input_tokens = event.usage.input_tokens
if event.usage.cache_creation_input_tokens is not None:
current_snapshot.usage.cache_creation_input_tokens = event.usage.cache_creation_input_tokens
if event.usage.cache_read_input_tokens is not None:
current_snapshot.usage.cache_read_input_tokens = event.usage.cache_read_input_tokens
if event.usage.server_tool_use is not None:
current_snapshot.usage.server_tool_use = event.usage.server_tool_use

return current_snapshot
17 changes: 17 additions & 0 deletions tests/lib/streaming/fixtures/usage_omitted_response.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
event: message_start
data: {"type":"message_start","message":{"id":"msg_usage_omitted","type":"message","role":"assistant","content":[],"model":"claude-test","stop_reason":null,"stop_sequence":null}}

event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}

event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}

event: content_block_stop
data: {"type":"content_block_stop","index":0}

event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":12,"output_tokens":6}}

event: message_stop
data: {"type":"message_stop"}
49 changes: 49 additions & 0 deletions tests/lib/streaming/test_messages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,30 @@ def test_refusal_stop_details_propagated(self, respx_mock: MockRouter) -> None:
) as stream:
assert_refusal_response(stream.get_final_message())

@pytest.mark.respx(base_url=base_url)
def test_usage_omitted_at_message_start(self, respx_mock: MockRouter) -> None:
# The streaming docs show a sequence where `message_start` omits
# `usage`; the accumulator should initialize it from `message_delta`
# instead of crashing on the missing value.
respx_mock.post("/v1/messages").mock(
return_value=httpx.Response(200, content=get_response("usage_omitted_response.txt"))
)

# A default (non-strict) client mirrors how the docs' event sequence
# reaches the accumulator without response-validation rejecting it.
client = Anthropic(base_url=base_url, api_key=api_key)

with client.messages.stream(
max_tokens=1024,
messages=[{"role": "user", "content": "Say hello there!"}],
model="claude-test",
) as stream:
message = stream.get_final_message()

assert message.usage is not None
assert message.usage.input_tokens == 12
assert message.usage.output_tokens == 6


class TestAsyncMessages:
@pytest.mark.asyncio
Expand Down Expand Up @@ -305,6 +329,31 @@ async def test_refusal_stop_details_propagated(self, respx_mock: MockRouter) ->
) as stream:
assert_refusal_response(await stream.get_final_message())

@pytest.mark.asyncio
@pytest.mark.respx(base_url=base_url)
async def test_usage_omitted_at_message_start(self, respx_mock: MockRouter) -> None:
# The streaming docs show a sequence where `message_start` omits
# `usage`; the accumulator should initialize it from `message_delta`
# instead of crashing on the missing value.
respx_mock.post("/v1/messages").mock(
return_value=httpx.Response(200, content=to_async_iter(get_response("usage_omitted_response.txt")))
)

# A default (non-strict) client mirrors how the docs' event sequence
# reaches the accumulator without response-validation rejecting it.
client = AsyncAnthropic(base_url=base_url, api_key=api_key)

async with client.messages.stream(
max_tokens=1024,
messages=[{"role": "user", "content": "Say hello there!"}],
model="claude-test",
) as stream:
message = await stream.get_final_message()

assert message.usage is not None
assert message.usage.input_tokens == 12
assert message.usage.output_tokens == 6


@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
def test_stream_method_definition_in_sync(sync: bool) -> None:
Expand Down