Spaces:
Configuration error
Configuration error
File size: 1,908 Bytes
a5a4bde | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from langchain_core.messages import AIMessage
from app.pricing import UsageNumbers
from app.usage import UsageAccumulator, extract_usage
def test_openai_cached_tokens_are_subtracted_from_public_input_bucket() -> None:
message = AIMessage(
content="ok",
response_metadata={
"token_usage": {
"prompt_tokens": 1_000,
"completion_tokens": 100,
"prompt_tokens_details": {"cached_tokens": 600},
}
},
)
accumulator = UsageAccumulator()
accumulator.add(extract_usage(message, "openai"))
assert accumulator.payload() == {
"input_tokens": 400,
"total_input_tokens": 1_000,
"output_tokens": 100,
"cache_read_tokens": 600,
"cache_write_tokens": 0,
"total_tokens": 1_100,
}
def test_anthropic_disjoint_raw_counters_are_normalized_to_total_input() -> None:
message = AIMessage(
content="ok",
response_metadata={
"usage": {
"input_tokens": 300,
"output_tokens": 50,
"cache_read_input_tokens": 500,
"cache_creation_input_tokens": 200,
}
},
)
assert extract_usage(message, "anthropic") == UsageNumbers(
input_tokens=1_000,
output_tokens=50,
cache_read_tokens=500,
cache_write_tokens=200,
)
def test_kimi_direct_cache_hit_counter_is_supported() -> None:
message = AIMessage(
content="ok",
response_metadata={
"token_usage": {
"prompt_tokens": 1_000,
"completion_tokens": 50,
"prompt_cache_hit_tokens": 400,
}
},
)
usage = extract_usage(message, "kimi")
assert usage.input_tokens == 1_000
assert usage.uncached_input_tokens == 600
assert usage.cache_read_tokens == 400
|