github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2026-07-12 11:21:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(anthropic): encode inference_geo as a key prefix in token details

Browse Source 
Earlier commit added a bare `inference_geo_us` key set to the full input
total, which overlaps with cache_read / cache_creation / priority — the
keys in input_token_details were no longer mutually exclusive, which the
langchainplus cost engine depends on (it sums detail values and treats
the remainder as the implicit non-cache base).

Restructure so inference_geo and service_tier compose into a single key
prefix. Every cache bucket gets the combined prefix, and the non-cache
base input lives under the bare prefix as its own bucket. Buckets sum to
input_tokens / output_tokens regardless of which dimensions are active:

  service_tier=priority, inference_geo=us:
    inference_geo_us_priority_cache_read
    inference_geo_us_priority_cache_creation
    inference_geo_us_priority           (= non-cache input)

  service_tier=standard, inference_geo=us:
    inference_geo_us_cache_read
    inference_geo_us_cache_creation
    inference_geo_us                    (= non-cache input)

  service_tier=priority, inference_geo=global:
    priority_cache_read
    priority_cache_creation
    priority                            (= non-cache input)

Default service_tier=standard + inference_geo=global emits the original
bare `cache_read` / `cache_creation` keys with no non-cache bucket
(matches legacy behavior; non-cache base stays implicit).

Tests now assert sum(details.values()) == input_tokens for every
non-default combination.
This commit is contained in:
Bagatur Askaryan
2026-05-28 16:47:01 -04:00
parent aec13fc781
commit d9d4691cc1
2 changed files with 81 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
libs/partners/anthropic/langchain_anthropic/chat_models.py
View File

@@ -2324,18 +2324,32 @@ def _create_usage_metadata(
if inference_geo is None:
inference_geo = getattr(anthropic_usage, "inference_geo", None)
# Only break tokens out by tier when the tier carries a non-default pricing
# multiplier. `standard` is the default rate; `priority` is ~1.5x; `batch`
# is half-price.
# Only break tokens out when the value carries a non-default pricing
# multiplier. `standard`/`global` are default rates; `priority` is ~1.5x;
# `batch` is half-price; `us` is 1.1x on every category.
if service_tier not in {"priority", "batch"}:
service_tier = None
tier_prefix = f"{service_tier}_" if service_tier else ""
if inference_geo != "us":
inference_geo = None
# Build a combined prefix that encodes geo and tier together. When set, every
# bucket in input_token_details / output_token_details is renamed with this
# prefix so the cost engine can read the multiplier off the key while the
# buckets stay mutually exclusive (their values sum to input_tokens /
# output_tokens).
prefix_parts: list[str] = []
if inference_geo:
prefix_parts.append(f"inference_geo_{inference_geo}")
if service_tier:
prefix_parts.append(service_tier)
bare_prefix = "_".join(prefix_parts)
key_prefix = f"{bare_prefix}_" if bare_prefix else ""
input_token_details: dict = {
f"{tier_prefix}cache_read": getattr(
f"{key_prefix}cache_read": getattr(
anthropic_usage, "cache_read_input_tokens", None
),
f"{tier_prefix}cache_creation": getattr(
f"{key_prefix}cache_creation": getattr(
anthropic_usage, "cache_creation_input_tokens", None
),
}
@@ -2352,19 +2366,19 @@ def _create_usage_metadata(
cache_creation = cache_creation.model_dump()
for k in cache_creation_keys:
specific_cache_creation_tokens += cache_creation.get(k, 0)
input_token_details[f"{tier_prefix}{k}"] = cache_creation.get(k)
input_token_details[f"{key_prefix}{k}"] = cache_creation.get(k)
if not isinstance(specific_cache_creation_tokens, int):
specific_cache_creation_tokens = 0
if specific_cache_creation_tokens > 0:
# Remove generic key to avoid double counting cache creation tokens
input_token_details[f"{tier_prefix}cache_creation"] = 0
input_token_details[f"{key_prefix}cache_creation"] = 0
# Calculate total input tokens: Anthropic's `input_tokens` excludes cached tokens,
# so we need to add them back to get the true total input token count
cache_read_total = input_token_details[f"{tier_prefix}cache_read"] or 0
cache_read_total = input_token_details[f"{key_prefix}cache_read"] or 0
cache_creation_total = (
specific_cache_creation_tokens
or input_token_details[f"{tier_prefix}cache_creation"]
or input_token_details[f"{key_prefix}cache_creation"]
or 0
)
input_tokens = base_non_cache_input + cache_read_total + cache_creation_total
@@ -2372,17 +2386,13 @@ def _create_usage_metadata(
output_token_details: dict = {}
if service_tier is not None:
# Mirror OpenAI: the bare tier key tracks tokens billed at tier rates that
# aren't already covered by a cache breakdown.
input_token_details[service_tier] = base_non_cache_input
output_token_details[service_tier] = output_tokens
# US-only inference applies a 1.1x multiplier to every token category, so the
# full input/output totals are the count subject to the multiplier.
if inference_geo == "us":
input_token_details["us"] = input_tokens
output_token_details["us"] = output_tokens
if bare_prefix:
# Non-cache base input gets a bare bucket so all input_token_details keys
# together partition input_tokens. Same for output_tokens — Anthropic
# doesn't currently split output further (no reasoning sub-bucket), so
# the bare bucket equals output_tokens.
input_token_details[bare_prefix] = base_non_cache_input
output_token_details[bare_prefix] = output_tokens
output_details_filtered = {
k: v for k, v in output_token_details.items() if v is not None

81
libs/partners/anthropic/tests/unit_tests/test_chat_models.py
View File

@@ -1523,6 +1523,10 @@ def test_usage_metadata_standardization() -> None:
def test_usage_metadata_service_tier_and_inference_geo() -> None:
"""`service_tier` and `inference_geo` should drive token-detail breakdowns
so downstream pricing can apply Anthropic's priority and US-only multipliers.
`input_token_details` keys must stay mutually exclusive: their values sum
to `input_tokens`, with the cost engine reading the multiplier off the key
prefix (`inference_geo_us_`, `priority_`, `batch_`, or a combination).
"""
class UsageStandard(BaseModel):
@@ -1534,12 +1538,15 @@ def test_usage_metadata_service_tier_and_inference_geo() -> None:
inference_geo: str = "global"
# standard tier + global geo: no extra keys, behaves like a vanilla response.
# No bare non-cache key (the non-cache portion stays implicit, matching legacy
# behavior on default-priced requests).
result = _create_usage_metadata(UsageStandard())
details = dict(result.get("input_token_details") or {})
assert details == {"cache_read": 5, "cache_creation": 3}
assert not result.get("output_token_details")
# priority tier: cache keys get prefixed, bare `priority` = non-cache input.
# priority tier alone: keys are prefixed with `priority_`, bare `priority`
# holds non-cache input. Mutually exclusive: keys sum to input_tokens.
class UsagePriority(BaseModel):
input_tokens: int = 80
output_tokens: int = 20
@@ -1550,15 +1557,16 @@ def test_usage_metadata_service_tier_and_inference_geo() -> None:
result = _create_usage_metadata(UsagePriority())
details = dict(result.get("input_token_details") or {})
assert details["priority_cache_read"] == 5
assert details["priority_cache_creation"] == 3
assert details["priority"] == 80 # non-cache input
assert "cache_read" not in details
assert "cache_creation" not in details
assert details == {
"priority_cache_read": 5,
"priority_cache_creation": 3,
"priority": 80,
}
assert sum(details.values()) == result["input_tokens"]
assert dict(result.get("output_token_details") or {}) == {"priority": 20}
assert result["input_tokens"] == 88 # 80 + 5 + 3
# batch tier behaves the same as priority (different multiplier downstream).
# batch tier alone: prefixed with `batch_`, bare `batch` = non-cache input.
# cache fields absent on the usage object (None) are filtered out of details.
class UsageBatch(BaseModel):
input_tokens: int = 80
output_tokens: int = 20
@@ -1566,11 +1574,11 @@ def test_usage_metadata_service_tier_and_inference_geo() -> None:
result = _create_usage_metadata(UsageBatch())
details = dict(result.get("input_token_details") or {})
assert details["batch"] == 80
assert details == {"batch": 80}
assert dict(result.get("output_token_details") or {}) == {"batch": 20}
# inference_geo='us': 1.1x applies to all categories, so the bare `us` key
# carries the full input/output totals.
# inference_geo='us' alone (standard tier): every key prefixed with
# `inference_geo_us_`, bare `inference_geo_us` for non-cache input.
class UsageUS(BaseModel):
input_tokens: int = 80
output_tokens: int = 20
@@ -1581,12 +1589,15 @@ def test_usage_metadata_service_tier_and_inference_geo() -> None:
result = _create_usage_metadata(UsageUS())
details = dict(result.get("input_token_details") or {})
assert details["us"] == 88 # full input total incl. cache
assert details["cache_read"] == 5
assert details["cache_creation"] == 3
assert dict(result.get("output_token_details") or {}) == {"us": 20}
assert details == {
"inference_geo_us_cache_read": 5,
"inference_geo_us_cache_creation": 3,
"inference_geo_us": 80,
}
assert sum(details.values()) == result["input_tokens"]
assert dict(result.get("output_token_details") or {}) == {"inference_geo_us": 20}
# Priority + US: both keys present, tier-prefixed cache keys + `us` total.
# priority + us: both multipliers encoded in the prefix.
class UsagePriorityUS(BaseModel):
input_tokens: int = 80
output_tokens: int = 20
@@ -1597,23 +1608,28 @@ def test_usage_metadata_service_tier_and_inference_geo() -> None:
result = _create_usage_metadata(UsagePriorityUS())
details = dict(result.get("input_token_details") or {})
assert details["priority_cache_read"] == 5
assert details["priority_cache_creation"] == 3
assert details["priority"] == 80
assert details["us"] == 88
out_details = dict(result.get("output_token_details") or {})
assert out_details == {"priority": 20, "us": 20}
assert details == {
"inference_geo_us_priority_cache_read": 5,
"inference_geo_us_priority_cache_creation": 3,
"inference_geo_us_priority": 80,
}
assert sum(details.values()) == result["input_tokens"]
assert dict(result.get("output_token_details") or {}) == {
"inference_geo_us_priority": 20
}
# 'not_available' inference_geo (older models) is a no-op.
# 'not_available' inference_geo (older models) is a no-op; same for `global`.
class UsageGeoUnavailable(BaseModel):
input_tokens: int = 10
output_tokens: int = 5
cache_read_input_tokens: int = 2
service_tier: str = "standard"
inference_geo: str = "not_available"
result = _create_usage_metadata(UsageGeoUnavailable())
details = dict(result.get("input_token_details") or {})
assert details == {"cache_read": 2}
assert not result.get("output_token_details")
assert "us" not in dict(result.get("input_token_details") or {})
# Explicit overrides win over the values on the usage object — used by the
# streaming path to inject service_tier/inference_geo from message_start.
@@ -1625,8 +1641,8 @@ def test_usage_metadata_service_tier_and_inference_geo() -> None:
UsageBareDelta(), service_tier="priority", inference_geo="us"
)
details = dict(result.get("input_token_details") or {})
assert details["priority"] == 80
assert details["us"] == 80
assert details == {"inference_geo_us_priority": 80}
assert sum(details.values()) == result["input_tokens"]
def test_usage_metadata_cache_creation_ttl() -> None:
@@ -2219,12 +2235,15 @@ def test_streaming_service_tier_and_inference_geo_propagation() -> None:
assert delta_chunk is not None
assert delta_chunk.usage_metadata is not None
input_details = delta_chunk.usage_metadata["input_token_details"]
assert input_details["priority_cache_read"] == 25
assert input_details["priority"] == 100 # non-cache input
assert input_details["us"] == 125 # total input subject to 1.1x
assert input_details == {
"inference_geo_us_priority_cache_read": 25,
"inference_geo_us_priority_cache_creation": 0,
"inference_geo_us_priority": 100,
}
# Buckets stay mutually exclusive: they sum to input_tokens.
assert sum(input_details.values()) == delta_chunk.usage_metadata["input_tokens"]
output_details = delta_chunk.usage_metadata.get("output_token_details") or {}
assert output_details["priority"] == 50
assert output_details["us"] == 50
assert output_details == {"inference_geo_us_priority": 50}
assert delta_chunk.response_metadata["service_tier"] == "priority"
assert delta_chunk.response_metadata["inference_geo"] == "us"