Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions astrbot/core/provider/sources/anthropic_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -445,15 +445,21 @@ def _extract_usage(self, usage: Usage | None) -> TokenUsage:
if usage is None:
return TokenUsage()
# https://docs.claude.com/en/docs/build-with-claude/prompt-caching#tracking-cache-performance
# Anthropic's input_tokens excludes cache served reads AND writes, so
# cache_creation_input_tokens must be added back into input_other to
# keep total input (and context-occupancy stats) accurate.
return TokenUsage(
input_other=usage.input_tokens or 0,
input_other=(usage.input_tokens or 0)
+ (usage.cache_creation_input_tokens or 0),
input_cached=usage.cache_read_input_tokens or 0,
output=usage.output_tokens or 0,
)

def _update_usage(self, token_usage: TokenUsage, usage: MessageDeltaUsage) -> None:
if usage.input_tokens is not None:
token_usage.input_other = usage.input_tokens
token_usage.input_other = usage.input_tokens + (
usage.cache_creation_input_tokens or 0
)
if usage.cache_read_input_tokens is not None:
token_usage.input_cached = usage.cache_read_input_tokens
if usage.output_tokens is not None:
Expand Down
77 changes: 77 additions & 0 deletions tests/test_anthropic_source.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
from anthropic.types import MessageDeltaUsage, Usage

from astrbot.core.provider.entities import TokenUsage
from astrbot.core.provider.sources.anthropic_source import ProviderAnthropic


def _provider() -> ProviderAnthropic:
return ProviderAnthropic.__new__(ProviderAnthropic)


def test_anthropic_extract_usage_counts_cache_creation_input():
provider = _provider()

usage = provider._extract_usage(
Usage(
input_tokens=10,
cache_read_input_tokens=100,
cache_creation_input_tokens=50,
output_tokens=20,
)
)

# Anthropic's input_tokens excludes cache writes, so cache_creation
# must be folded into input_other to keep total input accurate.
assert usage.input_other == 60
assert usage.input_cached == 100
assert usage.input == 160
assert usage.output == 20


def test_anthropic_extract_usage_without_cache_breakpoints():
provider = _provider()

usage = provider._extract_usage(Usage(input_tokens=30, output_tokens=10))

assert usage.input_other == 30
assert usage.input_cached == 0
assert usage.input == 30
assert usage.output == 10


def test_anthropic_extract_usage_none_returns_empty():
provider = _provider()

assert provider._extract_usage(None) == TokenUsage()


def test_anthropic_update_usage_counts_cache_creation_input():
provider = _provider()
token_usage = TokenUsage(input_other=5, input_cached=0, output=0)

provider._update_usage(
token_usage,
MessageDeltaUsage(
input_tokens=10,
cache_read_input_tokens=100,
cache_creation_input_tokens=50,
output_tokens=20,
),
)

assert token_usage.input_other == 60
assert token_usage.input_cached == 100
assert token_usage.input == 160
assert token_usage.output == 20


def test_anthropic_update_usage_omitted_fields_are_preserved():
provider = _provider()
token_usage = TokenUsage(input_other=5, input_cached=0, output=0)

# message_delta usage only carries output tokens in practice.
provider._update_usage(token_usage, MessageDeltaUsage(output_tokens=7))

assert token_usage.input_other == 5
assert token_usage.input_cached == 0
assert token_usage.output == 7
Loading