From 3e262334ce2813e940c7fdf83258412a3edbd385 Mon Sep 17 00:00:00 2001 From: Chester Curme Date: Thu, 8 May 2025 21:06:08 -0400 Subject: [PATCH] take usage from end of stream --- libs/partners/anthropic/langchain_anthropic/chat_models.py | 4 +--- .../anthropic/tests/integration_tests/test_chat_models.py | 4 ++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/libs/partners/anthropic/langchain_anthropic/chat_models.py b/libs/partners/anthropic/langchain_anthropic/chat_models.py index 25d52e21d3c..10695bc216f 100644 --- a/libs/partners/anthropic/langchain_anthropic/chat_models.py +++ b/libs/partners/anthropic/langchain_anthropic/chat_models.py @@ -1759,15 +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) + 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 ( diff --git a/libs/partners/anthropic/tests/integration_tests/test_chat_models.py b/libs/partners/anthropic/tests/integration_tests/test_chat_models.py index 5c41c3f878e..befdca0ec1c 100644 --- a/libs/partners/anthropic/tests/integration_tests/test_chat_models.py +++ b/libs/partners/anthropic/tests/integration_tests/test_chat_models.py @@ -46,7 +46,7 @@ def test_stream() -> None: if token.usage_metadata is not None: if token.usage_metadata.get("input_tokens"): chunks_with_input_token_counts += 1 - elif token.usage_metadata.get("output_tokens"): + if token.usage_metadata.get("output_tokens"): chunks_with_output_token_counts += 1 chunks_with_model_name += int("model_name" in token.response_metadata) if chunks_with_input_token_counts != 1 or chunks_with_output_token_counts != 1: @@ -85,7 +85,7 @@ async def test_astream() -> None: if token.usage_metadata is not None: if token.usage_metadata.get("input_tokens"): chunks_with_input_token_counts += 1 - elif token.usage_metadata.get("output_tokens"): + if token.usage_metadata.get("output_tokens"): chunks_with_output_token_counts += 1 if chunks_with_input_token_counts != 1 or chunks_with_output_token_counts != 1: raise AssertionError(