Archives/langchain
2
0
Fork 1
mirror of https://github.com/hwchase17/langchain synced 2024-11-13 19:10:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

standard-tests: test that only one chunk sets input_tokens (#27177)

Browse Source
This commit is contained in:
Erick Friis 2024-10-08 11:35:32 -07:00 committed by GitHub
parent 9b7bdf1a26
commit b84e00283f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
libs/standard-tests/langchain_standard_tests/integration_tests/chat_models.py
View File

@ -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 @@ class ChatModelIntegrationTests(ChatModelTests):
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)