github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-19 05:13:46 +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 base64
import json import json
from typing import List, Optional from typing import List, Optional, cast
import httpx import httpx
import pytest import pytest
@ -209,10 +209,21 @@ class ChatModelIntegrationTests(ChatModelTests):
def test_usage_metadata_streaming(self, model: BaseChatModel) -> None: def test_usage_metadata_streaming(self, model: BaseChatModel) -> None:
if not self.returns_usage_metadata: if not self.returns_usage_metadata:
pytest.skip("Not implemented.") pytest.skip("Not implemented.")
full: Optional[BaseMessageChunk] = None full: Optional[AIMessageChunk] = None
for chunk in model.stream("Hello"): for chunk in model.stream("Write me 2 haikus. Only include the haikus."):
assert isinstance(chunk, AIMessageChunk) 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 isinstance(full, AIMessageChunk)
assert full.usage_metadata is not None assert full.usage_metadata is not None
assert isinstance(full.usage_metadata["input_tokens"], int) assert isinstance(full.usage_metadata["input_tokens"], int)