scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

test: update unit test

Browse source
This commit is contained in:
Krrish Dholakia 2025-08-16 12:39:03 -07:00
parent 1ec0808713
commit 8bb5ee2ba5
1 changed files with 15 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
tests/logging_callback_tests/test_token_counting.py
View file

@ -223,26 +223,31 @@ async def test_stream_token_counting_anthropic_with_include_usage():
json.dumps(all_anthropic_usage_chunks, indent=4, default=str),
)
input_tokens_anthropic_api = sum(
[getattr(usage, "input_tokens", 0) or 0 for usage in all_anthropic_usage_chunks]
)
output_tokens_anthropic_api = sum(
[getattr(usage, "output_tokens", 0) or 0 for usage in all_anthropic_usage_chunks]
)
print("input_tokens_anthropic_api", input_tokens_anthropic_api)
print("output_tokens_anthropic_api", output_tokens_anthropic_api)
# Get the most recent value of input tokens (iterate backwards to find last non-zero value)
anthropic_api_input_tokens = 0
for usage in reversed(all_anthropic_usage_chunks):
if getattr(usage, "input_tokens", 0) > 0:
anthropic_api_input_tokens = getattr(usage, "input_tokens", 0)
break
anthropic_api_output_tokens = 0
for usage in reversed(all_anthropic_usage_chunks):
if getattr(usage, "output_tokens", 0) > 0:
anthropic_api_output_tokens = getattr(usage, "output_tokens", 0)
break
print("input_tokens_anthropic_api", anthropic_api_input_tokens)
print("output_tokens_anthropic_api", anthropic_api_output_tokens)
print("input_tokens_litellm", custom_logger.recorded_usage.prompt_tokens)
print("output_tokens_litellm", custom_logger.recorded_usage.completion_tokens)
## Assert Accuracy of token counting
# input tokens should be exactly the same
assert input_tokens_anthropic_api == custom_logger.recorded_usage.prompt_tokens
assert anthropic_api_input_tokens == custom_logger.recorded_usage.prompt_tokens
# output tokens can have at max abs diff of 10. We can't guarantee the response from two api calls will be exactly the same
assert (
abs(
output_tokens_anthropic_api - custom_logger.recorded_usage.completion_tokens
anthropic_api_output_tokens - custom_logger.recorded_usage.completion_tokens
)
<= 10
)