diff --git a/src/anthropic/lib/streaming/_beta_messages.py b/src/anthropic/lib/streaming/_beta_messages.py index 9c4ca2bd..1084d418 100644 --- a/src/anthropic/lib/streaming/_beta_messages.py +++ b/src/anthropic/lib/streaming/_beta_messages.py @@ -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 @@ -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 diff --git a/src/anthropic/lib/streaming/_messages.py b/src/anthropic/lib/streaming/_messages.py index 5c0da999..c17d1649 100644 --- a/src/anthropic/lib/streaming/_messages.py +++ b/src/anthropic/lib/streaming/_messages.py @@ -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, @@ -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 diff --git a/tests/lib/streaming/fixtures/usage_omitted_response.txt b/tests/lib/streaming/fixtures/usage_omitted_response.txt new file mode 100644 index 00000000..c6dc0f55 --- /dev/null +++ b/tests/lib/streaming/fixtures/usage_omitted_response.txt @@ -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"} diff --git a/tests/lib/streaming/test_messages.py b/tests/lib/streaming/test_messages.py index b86a3906..8fdfea02 100644 --- a/tests/lib/streaming/test_messages.py +++ b/tests/lib/streaming/test_messages.py @@ -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 @@ -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: