standard-tests: test that only one chunk sets input_tokens

libs/standard-tests/langchain_standard_tests/integration_tests/chat_models.py

@@ -1,6 +1,6 @@
 import base64
 import json
-from typing import List, Optional
+from typing import List, Optional, cast
 
 import httpx
 import pytest
@@ -209,10 +209,21 @@ def test_usage_metadata(self, model: BaseChatModel) -> None:
     def test_usage_metadata_streaming(self, model: BaseChatModel) -> None:
         if not self.returns_usage_metadata:
             pytest.skip("Not implemented.")
-        full: Optional[BaseMessageChunk] = None
-        for chunk in model.stream("Hello"):
+        full: Optional[AIMessageChunk] = None
+        for chunk in model.stream("Write me 2 haikus. Only include the haikus."):
             assert isinstance(chunk, AIMessageChunk)
-            full = chunk if full is None else full + chunk
+            # only one chunk is allowed to set usage_metadata.input_tokens
+            # if multiple do, it's likely a bug that will result in overcounting
+            # input tokens
+            if full and full.usage_metadata and full.usage_metadata["input_tokens"]:
+                assert (
+                    not chunk.usage_metadata or not chunk.usage_metadata["input_tokens"]
+                ), (
+                    "Only one chunk should set input_tokens,"
+                    " the rest should be 0 or None"
+                )
+            full = chunk if full is None else cast(AIMessageChunk, full + chunk)
+
         assert isinstance(full, AIMessageChunk)
         assert full.usage_metadata is not None
         assert isinstance(full.usage_metadata["input_tokens"], int)