take usage from end of stream

libs/partners/anthropic/langchain_anthropic/chat_models.py

@@ -1759,21 +1759,13 @@ def _make_message_chunk_from_anthropic_event(
     """
     message_chunk: Optional[AIMessageChunk] = None
     # See https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/lib/streaming/_messages.py  # noqa: E501
-    if event.type == "message_start" and stream_usage:
-        usage_metadata = _create_usage_metadata(event.message.usage)
-        # We pick up a cumulative count of output_tokens at the end of the stream,
-        # so here we zero out to avoid double counting.
-        usage_metadata["total_tokens"] = (
-            usage_metadata["total_tokens"] - usage_metadata["output_tokens"]
-        )
-        usage_metadata["output_tokens"] = 0
+    if event.type == "message_start":
         if hasattr(event.message, "model"):
             response_metadata = {"model_name": event.message.model}
         else:
             response_metadata = {}
         message_chunk = AIMessageChunk(
             content="" if coerce_content_to_string else [],
-            usage_metadata=usage_metadata,
             response_metadata=response_metadata,
         )
     elif (
@@ -1840,11 +1832,7 @@ def _make_message_chunk_from_anthropic_event(
                 tool_call_chunks=[tool_call_chunk],  # type: ignore
             )
     elif event.type == "message_delta" and stream_usage:
-        usage_metadata = UsageMetadata(
-            input_tokens=0,
-            output_tokens=event.usage.output_tokens,
-            total_tokens=event.usage.output_tokens,
-        )
+        usage_metadata = _create_usage_metadata(event.usage)
         message_chunk = AIMessageChunk(
             content="",
             usage_metadata=usage_metadata,

libs/partners/anthropic/tests/integration_tests/test_chat_models.py

@@ -134,9 +134,6 @@ async def test_stream_usage() -> None:
     async for token in model.astream("hi"):
         assert isinstance(token, AIMessageChunk)
         assert token.usage_metadata is None
-
-
-async def test_stream_usage_override() -> None:
     # check we override with kwarg
     model = ChatAnthropic(model_name=MODEL_NAME)  # type: ignore[call-arg]
     assert model.stream_usage