mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-21 05:15:06 +00:00
Merge 43e20f5a3a into 3ba5266ab3
This commit is contained in:
commit
66449c9ee9
7 changed files with 343 additions and 22 deletions
27
.github/workflows/test_server_root_path.yml
vendored
27
.github/workflows/test_server_root_path.yml
vendored
|
|
@ -36,13 +36,26 @@ jobs:
|
|||
uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
|
||||
|
||||
- name: Build Docker image
|
||||
uses: docker/build-push-action@0adf9959216b96bec444f325f1e493d4aa344497 # v6.14.0
|
||||
with:
|
||||
context: .
|
||||
file: ./docker/Dockerfile.non_root
|
||||
tags: litellm-test:${{ github.sha }}
|
||||
load: true
|
||||
push: false
|
||||
run: |
|
||||
retry() {
|
||||
local attempt=1
|
||||
local max_attempts=4
|
||||
until "$@"; do
|
||||
if [ "$attempt" -ge "$max_attempts" ]; then
|
||||
echo "Command failed after $attempt attempts: $*"
|
||||
return 1
|
||||
fi
|
||||
echo "Attempt $attempt failed: $*. Retrying in $((attempt * 15))s..."
|
||||
sleep $((attempt * 15))
|
||||
attempt=$((attempt + 1))
|
||||
done
|
||||
}
|
||||
|
||||
retry docker buildx build \
|
||||
--file ./docker/Dockerfile.non_root \
|
||||
--tag litellm-test:${{ github.sha }} \
|
||||
--load \
|
||||
.
|
||||
|
||||
- name: Start LiteLLM container with SERVER_ROOT_PATH
|
||||
run: |
|
||||
|
|
|
|||
|
|
@ -39,3 +39,4 @@
|
|||
- {id: quota_management.spend_tracking.failure.writes_failure_row, module: quota_management, tier: P1, behavior: spend_tracking, variant: failure, assertions: [writes_failure_row], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_log_error_logger.py", rationale: "A failed call writes a failure-status spend row"}
|
||||
- {id: quota_management.spend_tracking.spend_calculate.returns_cost, module: quota_management, tier: P2, behavior: spend_tracking, variant: spend_calculate, assertions: [returns_cost], exercised_on: [spend_calculate], source: "proxy/spend_tracking/spend_management_endpoints.py", rationale: "/spend/calculate prices a hypothetical request at nonzero cost"}
|
||||
- {id: quota_management.spend_tracking.pagination.keeps_total, module: quota_management, tier: P2, behavior: spend_tracking, variant: pagination, assertions: [keeps_total], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_management_endpoints.py", rationale: "Spend-logs v2 pagination caps page size without losing the total"}
|
||||
- {id: quota_management.spend_tracking.non_negative.never_negative, module: quota_management, tier: P1, behavior: spend_tracking, variant: non_negative, assertions: [never_negative], exercised_on: [chat_completions, embeddings], source: "litellm_core_utils/llm_cost_calc/utils.py", fail_before_fix: unproven, rationale: "No provider (gemini/anthropic/openai/bedrock) writes a negative spend row across streaming/non-streaming chat and embeddings; negative spend credits a budget back, letting a key outspend its cap. Broad invariant sweep on uncached traffic: it does NOT cover the bedrock cache-token path (#25846), which needs cache_hit > prompt_tokens to fire, so it stays unproven until a test drives that path"}
|
||||
|
|
|
|||
|
|
@ -29,6 +29,8 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`.
|
|||
|
||||
| Modality | Existing | Status | Live e2e |
|
||||
|----------|----------|--------|----------|
|
||||
| Non-negative spend (all providers, uncached traffic) | `test_anthropic_claude3_transformation.py`, `llm_cost_calc/utils.py` | partial | yes (`test_no_provider_logs_negative_spend`) |
|
||||
| Negative cost from bedrock cache-token split (#25846) | `test_llm_cost_calc_utils.py::test_inconsistent_cache_usage_never_prices_negative`, `test_anthropic_claude3_transformation.py::test_bedrock_sse_wrapper_bills_cache_rate_when_only_message_start_has_breakdown` (both fail-before-fix proven) | unsupported (live) | no - not reproducible on commercial bedrock; see note below. Covered offline instead |
|
||||
| Chat (non-stream) | `test_cost_calculator.py`, `local_testing/test_completion_cost.py` | covered | yes (`test_chat_completion_writes_nonzero_spend_row`) |
|
||||
| Chat (streaming) | `test_streaming_interrupt_spend_tracking.py` | partial | yes (`test_streaming_chat_completion_tracks_spend`) |
|
||||
| Embedding | `test_cost_calculator.py` (#29956) | partial | yes (`test_embedding_writes_nonzero_spend_row`) |
|
||||
|
|
@ -68,6 +70,7 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`.
|
|||
| `test_tag_spend_matches_sum_of_tagged_logs` | `/spend/tags` SUM/COUNT == tagged rows |
|
||||
| `test_end_user_spend_attributed_on_row` | `end_user` attributed + costed |
|
||||
| `test_each_model_on_a_shared_key_gets_its_own_row` | per-model/provider rows, correct model + cost, distinct request_ids matching response id |
|
||||
| `test_no_provider_logs_negative_spend` | no provider (gemini/anthropic/openai/bedrock) writes spend < 0 across streaming/non-streaming chat + embeddings; each provider still logs a positive row (non-vacuous). Uncached traffic only, so it does not reach the #25846 cache-token path; `fail_before_fix=unproven` |
|
||||
| `test_failure_call_writes_failure_status_row` | failed call -> `status=failure`, `spend=0` |
|
||||
| `test_spend_calculate_returns_nonzero_cost` | cost-map smoke (no batch wait) |
|
||||
| `test_spend_logs_endpoint_returns_spend` | `/spend/logs` returns 200 + the key's spend, never a 5xx (intermittent-500 regression) |
|
||||
|
|
@ -75,6 +78,30 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`.
|
|||
| `test_spend_logs_v2_pagination_caps_pages_and_keeps_total` | `/spend/logs/v2` page cap, stable total on out-of-range page, zero total on no-match filter |
|
||||
| `test_spend_routes.py` (23) | no spend route 404s or 5xxs |
|
||||
|
||||
## Why the #25846 cache-token path is `unsupported` live, not a gap
|
||||
|
||||
Probed against real bedrock (`us.anthropic.claude-haiku-4-5`, commercial `us-east-1`) on
|
||||
2026-07-14 with a 14713-token cached prefix, streaming `/v1/messages`, cache write then
|
||||
cache read. Findings, so nobody re-runs this:
|
||||
|
||||
- `message_stop` carries **no** cache breakdown here, only `{input_tokens: 12, output_tokens: N}`.
|
||||
So the "cache fields land on message_stop" reading of the code comment does not hold on
|
||||
commercial bedrock.
|
||||
- The raw `message_delta` **already** carries the full breakdown
|
||||
(`input_tokens: 12, cache_read_input_tokens: 14713`), so it is self-consistent on its own.
|
||||
`prompt_tokens` becomes `12 + 14713 = 14725` and `text_tokens = 14725 - 14713 = 12`, never
|
||||
below zero.
|
||||
- Confirmed by reverting the fix (`patched_stream = completion_stream`) and re-driving the
|
||||
same cached call: spend stayed **positive** (0.00166463, 0.00165363; prompt_tokens 14725),
|
||||
zero negative rows. `_merge_message_start_cache_into_delta_usage` is a no-op on this
|
||||
deployment.
|
||||
|
||||
The bug needs a deployment whose `message_delta` omits the cache breakdown (GovCloud, per the
|
||||
transformation docstring / LIT-2411). A live commercial call cannot produce that payload, so
|
||||
this cell is `unsupported` for live e2e (excluded from the denominator) rather than a gap to
|
||||
chase. It is covered offline instead by the two tests named above, both proven fail-before-fix
|
||||
by removing the real fix. Revisit only if GovCloud is in scope.
|
||||
|
||||
## Design + timing
|
||||
|
||||
`proxy_batch_write_at` (~60s) means rows land late; every read polls to a deadline.
|
||||
|
|
|
|||
|
|
@ -5,13 +5,14 @@ live in the parent tests/e2e/conftest.py. SpendClient exposes the shared Gateway
|
|||
(GatewayProvider), so the `resources` fixture cleans up keys and customers this
|
||||
suite creates.
|
||||
|
||||
The suite drives real calls through three deployments. On the stage gateway they
|
||||
are baked into the proxy config; on a local dev proxy they usually are not, so
|
||||
`driver_models` registers whichever are missing via /model/new and deletes only
|
||||
the ones it created, never a config-baked deployment. Each registration carries
|
||||
the provider key from the test runner's env when set (so a local proxy whose
|
||||
container env lacks the key still works); otherwise it falls back to an
|
||||
os.environ reference resolved from the proxy's own env, the stage convention.
|
||||
The suite drives real calls through four deployments (gemini, anthropic, bedrock,
|
||||
openai embeddings). On the stage gateway they are baked into the proxy config; on a
|
||||
local dev proxy they usually are not, so `driver_models` registers whichever are
|
||||
missing via /model/new and deletes only the ones it created, never a config-baked
|
||||
deployment. Each registration carries the provider credential from the test runner's
|
||||
env when set (so a local proxy whose container env lacks it still works); otherwise
|
||||
it falls back to an os.environ reference resolved from the proxy's own env, the stage
|
||||
convention. Bedrock is credentialed with AWS keys rather than a single api_key.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
|
@ -23,17 +24,37 @@ from models import LiteLLMParamsBody
|
|||
from spend_e2e_client import SpendClient, build_client
|
||||
|
||||
|
||||
def _driver_params(provider_model: str, env_var: str) -> LiteLLMParamsBody:
|
||||
def _env_ref(name: str) -> str:
|
||||
"""The runner's value when set (so a local proxy whose container env lacks it
|
||||
still works), else an os.environ reference the proxy resolves from its own env."""
|
||||
return os.environ.get(name) or f"os.environ/{name}"
|
||||
|
||||
|
||||
def _api_key_params(provider_model: str, env_var: str) -> LiteLLMParamsBody:
|
||||
return LiteLLMParamsBody(model=provider_model, api_key=_env_ref(env_var))
|
||||
|
||||
|
||||
def _bedrock_params(provider_model: str) -> LiteLLMParamsBody:
|
||||
"""Bedrock authenticates with AWS creds, not a single api_key."""
|
||||
return LiteLLMParamsBody(
|
||||
model=provider_model,
|
||||
api_key=os.environ.get(env_var) or f"os.environ/{env_var}",
|
||||
aws_access_key_id=_env_ref("AWS_ACCESS_KEY_ID"),
|
||||
aws_secret_access_key=_env_ref("AWS_SECRET_ACCESS_KEY"),
|
||||
aws_region_name=os.environ.get("AWS_REGION_NAME") or _env_ref("AWS_REGION"),
|
||||
)
|
||||
|
||||
|
||||
DRIVER_MODELS: tuple[tuple[str, str, str], ...] = (
|
||||
("gemini-2.5-flash", "gemini/gemini-2.5-flash", "GEMINI_API_KEY"),
|
||||
("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"),
|
||||
("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"),
|
||||
DRIVER_MODELS: tuple[tuple[str, LiteLLMParamsBody], ...] = (
|
||||
("gemini-2.5-flash", _api_key_params("gemini/gemini-2.5-flash", "GEMINI_API_KEY")),
|
||||
("claude-haiku-4-5", _api_key_params("anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY")),
|
||||
(
|
||||
"bedrock-claude-haiku-4-5",
|
||||
_bedrock_params("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"),
|
||||
),
|
||||
(
|
||||
"openai-text-embedding-3-small",
|
||||
_api_key_params("openai/text-embedding-3-small", "OPENAI_API_KEY"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -46,8 +67,8 @@ def client() -> SpendClient:
|
|||
def driver_models(client: SpendClient) -> Iterator[None]:
|
||||
existing = frozenset(entry.model_name for entry in client.gateway.model_info())
|
||||
created = tuple(
|
||||
client.gateway.create_model(name, _driver_params(provider_model, env_var))
|
||||
for name, provider_model, env_var in DRIVER_MODELS
|
||||
client.gateway.create_model(name, params)
|
||||
for name, params in DRIVER_MODELS
|
||||
if name not in existing
|
||||
)
|
||||
yield
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ fails the test; a pricing or token-count drift does not.
|
|||
import time
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -430,6 +431,112 @@ def test_each_model_on_a_shared_key_gets_its_own_row(
|
|||
), f"claude row request_id {claude_row.request_id} != response id {claude.id}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SweepModel:
|
||||
"""A deployment the negative-spend guard drives. `call` is the model_name sent to
|
||||
the proxy; `row_marker` is the substring the resulting spend row's `model` must
|
||||
contain (rows carry the resolved provider/model, so bedrock's alias shows up as
|
||||
`bedrock/us.anthropic...`). The two claude markers are prefix-qualified so a
|
||||
bedrock-anthropic row can't stand in for anthropic-direct, or vice versa."""
|
||||
|
||||
call: str
|
||||
row_marker: str
|
||||
|
||||
|
||||
# Chat providers this guard sweeps when the proxy exposes them, each non-streaming and
|
||||
# streaming: gemini, anthropic, openai, and bedrock. gemini/anthropic/bedrock are
|
||||
# registered by the suite's driver_models; gpt-5.5 is only wired on the stage/docker
|
||||
# gateways. Adding a provider is one line.
|
||||
_CHAT_SWEEP: tuple[_SweepModel, ...] = (
|
||||
_SweepModel("gemini-2.5-flash", "gemini-2.5-flash"),
|
||||
_SweepModel("claude-haiku-4-5", "anthropic/claude-haiku-4-5"),
|
||||
_SweepModel("gpt-5.5", "gpt-5.5"),
|
||||
_SweepModel("bedrock-claude-haiku-4-5", "us.anthropic.claude-haiku-4-5"),
|
||||
)
|
||||
_EMBEDDING_SWEEP = _SweepModel("openai-text-embedding-3-small", "text-embedding-3-small")
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.non_negative.never_negative")
|
||||
def test_no_provider_logs_negative_spend(
|
||||
client: SpendClient, scoped_key: str
|
||||
) -> None:
|
||||
"""No provider ever writes a negative spend row. One key sweeps every chat provider
|
||||
the proxy exposes, both non-streaming and streaming, plus an embedding; then every
|
||||
logged row is asserted non-negative. Negative spend is a billing bug in any form: it
|
||||
credits the customer's budget back, so a key can outspend its cap without ever
|
||||
tripping it. The per-provider positive-row check keeps the guard non-vacuous - a
|
||||
pipeline that silently logs 0 (or drops the row) fails instead of sliding past a
|
||||
bare `>= 0`.
|
||||
|
||||
Scope, deliberately stated: this is a broad invariant sweep, NOT a regression test
|
||||
for the bedrock cache-token bug (#25846). That bug needs prompt_tokens to exclude
|
||||
cache tokens so `prompt_tokens - cache_hit` goes below zero; the traffic here is
|
||||
uncached, so `cache_hit == 0` and that subtraction can't go negative. Reverting
|
||||
#25846's clamp leaves this test green, so the cell is fail_before_fix=unproven.
|
||||
|
||||
Nor could that be fixed by adding a cached bedrock leg here: probing real bedrock
|
||||
(commercial us-east-1) on 2026-07-14 showed its raw message_delta already carries
|
||||
the cache breakdown, so usage stays self-consistent and spend stays positive even
|
||||
with the fix reverted. The bug needs a deployment whose message_delta omits it
|
||||
(GovCloud / LIT-2411), which a live commercial call cannot produce. That path is
|
||||
therefore unsupported live and guarded by the mock unit test instead; see
|
||||
SPEND_TRACKING_COVERAGE_MATRIX.md."""
|
||||
present = frozenset(entry.model_name for entry in client.gateway.model_info())
|
||||
chat = tuple(m for m in _CHAT_SWEEP if m.call in present)
|
||||
assert len(chat) >= 2, (
|
||||
f"need >=2 chat providers for a cross-provider guard; "
|
||||
f"proxy exposes {sorted(present)}"
|
||||
)
|
||||
assert _EMBEDDING_SWEEP.call in present, (
|
||||
f"embedding deployment {_EMBEDDING_SWEEP.call!r} not registered; "
|
||||
f"proxy exposes {sorted(present)}"
|
||||
)
|
||||
|
||||
for model in chat:
|
||||
_ = unwrap(
|
||||
client.chat(
|
||||
scoped_key, model.call, f"one word {unique_marker()}", max_tokens=16
|
||||
)
|
||||
)
|
||||
stream = client.chat_stream(
|
||||
scoped_key, model.call, f"count to three {unique_marker()}", max_tokens=32
|
||||
)
|
||||
assert stream.ok, (
|
||||
f"{model.call} stream failed (status {stream.status_code}): {stream.body[:300]}"
|
||||
)
|
||||
_ = unwrap(
|
||||
client.embed(
|
||||
scoped_key, _EMBEDDING_SWEEP.call, f"vectorize {unique_marker()}"
|
||||
)
|
||||
)
|
||||
|
||||
expected = (*chat, _EMBEDDING_SWEEP)
|
||||
|
||||
def every_provider_costed(rows: list[SpendLogRow]) -> bool:
|
||||
return all(
|
||||
any(m.row_marker in (r.model or "") and (r.spend or 0) > 0 for r in rows)
|
||||
for m in expected
|
||||
)
|
||||
|
||||
# min_rows waits for the streaming rows too (non-stream + stream per chat model,
|
||||
# plus the embedding), so the streaming spend is actually observed and checked.
|
||||
rows = client.poll_logs_for_key(
|
||||
scoped_key,
|
||||
min_rows=2 * len(chat) + 1,
|
||||
predicate=every_provider_costed,
|
||||
)
|
||||
|
||||
negative = [r for r in rows if (r.spend or 0) < 0]
|
||||
assert not negative, (
|
||||
f"provider logged negative spend (billing regression): {_summarize(negative)}"
|
||||
)
|
||||
|
||||
for m in expected:
|
||||
assert any(
|
||||
m.row_marker in (r.model or "") and (r.spend or 0) > 0 for r in rows
|
||||
), f"no positive spend row for {m.row_marker!r}; guard would be vacuous: {_summarize(rows)}"
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.failure.writes_failure_row")
|
||||
def test_failure_call_writes_failure_status_row(
|
||||
client: SpendClient, scoped_key: str
|
||||
|
|
|
|||
|
|
@ -400,6 +400,49 @@ def test_generic_cost_per_token_above_200k_tokens():
|
|||
)
|
||||
|
||||
|
||||
def test_inconsistent_cache_usage_never_prices_negative():
|
||||
"""An inconsistent usage payload must never produce a negative cost.
|
||||
|
||||
generic_cost_per_token treats prompt_tokens as cache-INCLUSIVE and derives the
|
||||
uncached remainder by subtraction: text_tokens = prompt_tokens - cache_hit. A
|
||||
provider stream can hand us a payload where prompt_tokens is cache-EXCLUSIVE while
|
||||
cached_tokens is populated (bedrock reports the cache breakdown on one SSE event and
|
||||
uncached input_tokens on another; if the two are read from different events they
|
||||
disagree). The subtraction then underflows and a negative token count times a
|
||||
positive rate bills the customer a negative amount, which credits their budget back
|
||||
and lets a key outspend its cap. See BerriAI/litellm#25846 / LIT-2411.
|
||||
|
||||
Fails without the zero-clamp in llm_cost_calc/utils.py: prompt_cost lands at
|
||||
about -0.0146 instead of a positive number.
|
||||
"""
|
||||
model = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
custom_llm_provider = "bedrock"
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
||||
# cached_tokens exceeds prompt_tokens: the payload the clamp exists to absorb.
|
||||
usage = Usage(
|
||||
prompt_tokens=12,
|
||||
completion_tokens=6,
|
||||
total_tokens=18,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=14713),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model,
|
||||
usage=usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
assert prompt_cost >= 0, f"negative input cost from inconsistent usage: {prompt_cost}"
|
||||
assert completion_cost >= 0, f"negative output cost: {completion_cost}"
|
||||
|
||||
# The cached tokens must still be billed at the cache-read rate rather than dropped,
|
||||
# so the clamp cannot be satisfied by zeroing the whole input cost.
|
||||
model_cost_map = litellm.model_cost[model]
|
||||
assert prompt_cost >= model_cost_map["cache_read_input_token_cost"] * 14713
|
||||
|
||||
|
||||
def test_get_token_base_cost_picks_highest_crossed_tier():
|
||||
"""Regression test for #30345.
|
||||
|
||||
|
|
|
|||
|
|
@ -240,6 +240,115 @@ async def test_bedrock_sse_wrapper_keeps_usage_in_message_start_and_message_delt
|
|||
assert delta_json["usage"]["output_tokens"] == 8
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_sse_wrapper_bills_cache_rate_when_only_message_start_has_breakdown(
|
||||
local_model_cost_map,
|
||||
):
|
||||
"""A stream that reports the cache breakdown ONLY on message_start must still bill
|
||||
the cached tokens at the cache-read rate.
|
||||
|
||||
Bedrock usually repeats the cache breakdown on message_stop, but some deployments
|
||||
(GovCloud) put it only on message_start and then have message_delta / message_stop
|
||||
repeat uncached input_tokens alone. message_delta is the event the cost path reads,
|
||||
so without the message_start merge it carries input_tokens=12 and no cache fields:
|
||||
the 14713 cached tokens vanish from the bill and the call is charged ~36x too
|
||||
little. Worse, a payload that pairs a cache-exclusive prompt_tokens with a populated
|
||||
cached_tokens drives the cost calculator negative (see
|
||||
test_inconsistent_cache_usage_never_prices_negative). See BerriAI/litellm#25846 /
|
||||
LIT-2411.
|
||||
|
||||
Asserts the resulting cost, not just the usage dict, because the dict being merged
|
||||
is the mechanism while the bill is the promise. Fails without
|
||||
_merge_message_start_cache_into_delta_usage: cost drops to ~4.62e-05 from ~1.66e-03.
|
||||
"""
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
model = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
uncached_input, cached_read, output = 12, 14713, 6
|
||||
|
||||
async def _govcloud_stream():
|
||||
yield {
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"model": model,
|
||||
# the ONLY event carrying the cache breakdown
|
||||
"usage": {
|
||||
"input_tokens": uncached_input,
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": cached_read,
|
||||
"output_tokens": 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
# uncached input_tokens only, no cache fields
|
||||
yield {
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"input_tokens": uncached_input, "output_tokens": output},
|
||||
}
|
||||
yield {
|
||||
"type": "message_stop",
|
||||
"usage": {"input_tokens": uncached_input, "output_tokens": output},
|
||||
}
|
||||
|
||||
collected: list[bytes] = []
|
||||
async for chunk in AmazonAnthropicClaudeMessagesConfig().bedrock_sse_wrapper(
|
||||
_govcloud_stream(),
|
||||
litellm_logging_obj=LiteLLMLoggingObj(
|
||||
model=f"bedrock/invoke/{model}",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
stream=True,
|
||||
call_type="chat",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test_bedrock_sse_wrapper_govcloud_cache_on_start",
|
||||
function_id="test_bedrock_sse_wrapper_govcloud_cache_on_start",
|
||||
),
|
||||
request_body={},
|
||||
):
|
||||
collected.append(chunk)
|
||||
|
||||
delta_chunk = next(c for c in collected if b"event: message_delta\n" in c)
|
||||
delta_json = json.loads(delta_chunk.decode("utf-8").split("data: ", 1)[1].strip())
|
||||
|
||||
# the delta the cost path reads must be self-consistent on its own
|
||||
delta_usage = delta_json["usage"]
|
||||
assert delta_usage.get("cache_read_input_tokens") == cached_read, (
|
||||
f"message_delta lost the cache breakdown, so {cached_read} cached tokens go "
|
||||
f"unbilled: {delta_usage}"
|
||||
)
|
||||
assert delta_usage.get("input_tokens") == uncached_input, delta_usage
|
||||
|
||||
usage = AnthropicConfig().calculate_usage(
|
||||
usage_object=dict(delta_usage), reasoning_content=None
|
||||
)
|
||||
# prompt_tokens must come out cache-INCLUSIVE, which is what the cost calculator
|
||||
# assumes when it derives the uncached remainder by subtraction
|
||||
assert usage.prompt_tokens == uncached_input + cached_read
|
||||
assert usage.prompt_tokens_details.cached_tokens == cached_read
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model, usage=usage, custom_llm_provider="bedrock"
|
||||
)
|
||||
|
||||
cost_map = litellm.model_cost[model]
|
||||
expected_prompt_cost = (
|
||||
uncached_input * cost_map["input_cost_per_token"]
|
||||
+ cached_read * cost_map["cache_read_input_token_cost"]
|
||||
)
|
||||
assert prompt_cost == pytest.approx(expected_prompt_cost, rel=1e-6), (
|
||||
f"cached tokens not billed at the cache rate: {prompt_cost} != {expected_prompt_cost}"
|
||||
)
|
||||
assert completion_cost == pytest.approx(
|
||||
output * cost_map["output_cost_per_token"], rel=1e-6
|
||||
)
|
||||
|
||||
|
||||
def test_chunk_parser_usage_transformation():
|
||||
"""Ensure Bedrock invocation metrics are transformed to Anthropic usage keys."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue