From 4a58784fa2dfef85308e3cc77f36c54f93ff28f3 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 06:46:52 +0000 Subject: [PATCH] fix(cost): price exact cost-map names before price-free capability rules, and fix SCIM and MCP integration tests on CircleCI (#47037) * fix(test-harness): preserve sibling auth modules in isolated worker configs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): expect servers in MCP connect-flow descriptions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): expect ptu_hours in daily activity metrics after #43043 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): kill every serving worker once the SCIM push blocks on the row lock * fix(cost): price exact cost-map names before price-free capability rule matches * test(streaming): price a priced alias stream at its deployment model and keep the unpriceable case unmapped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost): price an explicit base_model under its own provider after candidate reordering Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost): retain priced generalized candidates and skip explicit-pricing lookups Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): assert custom pricing makes no remote model-info requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(cost): keep the per-candidate pricing provider final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): pin a Bedrock application profile prompt for the one-hour ttl its base model sends Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert the Nova Sonic tool-call argument delta carries the full arguments * fix(cost): sort cost candidates from the static cost map and rules, without provider lookups * test(integration): bill route-priced Bedrock chat streams and responses, and pin application profiles for an hour on messages and cache hits --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Yuneng Jiang --- litellm/cost_calculator.py | 71 +++- .../prompt_caching_deployment_check.py | 35 +- litellm/utils.py | 9 + .../cost_calculation/cases/misc.py | 27 ++ .../management/test_scim_group_bulk_roster.py | 31 +- tests/integration/mcp/test_mcp_oauth_flows.py | 2 + ..._bedrock_nova_sonic_realtime_tool_calls.py | 2 + .../routing/test_prompt_cache_pin_ttl.py | 28 +- .../spend/test_daily_rollup_retry.py | 1 + .../unit/integration_support/test_process.py | 26 ++ .../test_prompt_caching_deployment_check.py | 17 + tests/unit/test_cost_calculator.py | 307 +++++++++++++++++- tests/unit/test_main.py | 37 ++- 13 files changed, 540 insertions(+), 53 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index fef6d618e9d..3976612c6e1 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -7,7 +7,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast from httpx import Response from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -150,6 +150,7 @@ from litellm.utils import ( TranscriptionResponse, cached_get_model_info_helper, token_counter, + unmapped_model_generalization, ) if TYPE_CHECKING: @@ -875,6 +876,16 @@ def _base_model_pricing_provider(base_model: str, custom_llm_provider: str | Non return own_model_info["litellm_provider"] +def _pricing_provider_for_cost_candidate( + candidate_provider: str | None, + request_provider: str | None, + current_provider: object, +) -> str | None: + if candidate_provider != request_provider: + return candidate_provider + return current_provider if isinstance(current_provider, str) else None + + def select_model_name_for_cost_calc( model: str | None, completion_response: object | None, @@ -1544,14 +1555,20 @@ def completion_cost( ) requested_model: Final = model - potential_model_names: Final = [ - selected_model, - _get_response_model(completion_response), - ] - if model is not None: - potential_model_names.append(model) + request_pricing_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None + raw_candidates: Final = ( + ( + selected_model, + base_model_provider if base_model_provider is not None else request_pricing_provider, + ), + (_get_response_model(completion_response), request_pricing_provider), + *(((model, request_pricing_provider),) if model is not None else ()), + ) + potential_model_names: Final = ( + raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates) + ) - for idx, model in enumerate(potential_model_names): + for idx, (model, pricing_provider) in enumerate(potential_model_names): try: if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug("selected model name for cost calculation: %s", model) @@ -1896,13 +1913,6 @@ def completion_cost( if litellm_logging_obj is not None: request_model_for_cost = litellm_logging_obj.model - pricing_provider: str | None = ( - base_model_provider - if base_model_provider is not None and idx == 0 - else custom_llm_provider - if isinstance(custom_llm_provider, str) - else None - ) ( prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar, @@ -1910,7 +1920,11 @@ def completion_cost( model=model, prompt_tokens=prompt_tokens or 0, completion_tokens=completion_tokens or 0, - custom_llm_provider=pricing_provider, + custom_llm_provider=_pricing_provider_for_cost_candidate( + candidate_provider=pricing_provider, + request_provider=request_pricing_provider, + current_provider=custom_llm_provider, + ), response_time_ms=response_time_ms, region_name=None if explicit_pricing else region_name, custom_cost_per_second=custom_cost_per_second, @@ -2009,7 +2023,11 @@ def completion_cost( if cost_per_token_usage_object is not None and model: _token_type_breakdown = get_token_type_cost_breakdown( model=model, - custom_llm_provider=pricing_provider, + custom_llm_provider=_pricing_provider_for_cost_candidate( + candidate_provider=pricing_provider, + request_provider=request_pricing_provider, + current_provider=custom_llm_provider, + ), usage=cost_per_token_usage_object, service_tier=service_tier, data_residency=data_residency, @@ -2271,6 +2289,25 @@ def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelIn return None +_CostCandidate: TypeAlias = tuple[str | None, str | None] + + +def _is_unpriced_generalized_candidate(candidate: _CostCandidate) -> bool: + model, pricing_provider = candidate + if model is None: + return False + rule_info: Final = unmapped_model_generalization(model, pricing_provider) + return rule_info is not None and not _cost_map_entry_prices_anything(rule_info) + + +def _priced_candidates_before_unpriced_rules(candidates: tuple[_CostCandidate, ...]) -> tuple[_CostCandidate, ...]: + unpriced: Final = tuple(_is_unpriced_generalized_candidate(candidate) for candidate in candidates) + return ( + *(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if not is_unpriced), + *(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if is_unpriced), + ) + + def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None: raw_entry: Final = litellm.model_cost.get(key) return raw_entry if isinstance(raw_entry, Mapping) else None diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 091472c1ea3..b0da3ed264a 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -4,8 +4,10 @@ Check if prompt caching is valid for a given deployment Route to previously cached model id, if valid """ -from collections.abc import Callable -from typing import Final, cast +from collections.abc import Callable, Mapping +from typing import Final, TypedDict, cast + +from typing_extensions import ReadOnly from litellm import verbose_logger from litellm.caching.dual_cache import DualCache @@ -16,7 +18,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( ) from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.token_counter import offload_token_count -from litellm.llms.bedrock.common_utils import supports_bedrock_prompt_cache_ttl +from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -24,6 +26,11 @@ from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_p from ..prompt_caching_cache import PromptCachingCache +class _PromptCachingSuccessEventKwargs(TypedDict, total=False): + standard_logging_object: ReadOnly[StandardLoggingPayload | None] + litellm_params: ReadOnly[Mapping[str, object] | None] + + def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int: """ Returns the lowest minimum cacheable prefix across a model group. @@ -49,9 +56,13 @@ def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int ) -def _honors_one_hour_cache_ttl(standard_logging_object: StandardLoggingPayload) -> bool: +def _honors_one_hour_cache_ttl( + standard_logging_object: StandardLoggingPayload, litellm_params: Mapping[str, object] +) -> bool: if standard_logging_object["custom_llm_provider"] == "bedrock": - return supports_bedrock_prompt_cache_ttl(standard_logging_object["model"]) + deployment_model: Final = standard_logging_object["model"] + capability_model: Final = bedrock_capability_model(deployment_model, litellm_params.get("base_model")) + return bedrock_cache_point_ttl_supported(capability_model, deployment_model) model_map_value: Final = standard_logging_object["model_map_information"]["model_map_value"] return model_map_value is not None and (model_map_value.get("cache_creation_input_token_cost_above_1hr") or 0) > 0 @@ -121,7 +132,13 @@ class PromptCachingDeploymentCheck(CustomLogger): return healthy_deployments - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event( + self, + kwargs: _PromptCachingSuccessEventKwargs, + response_obj: object, + start_time: object, + end_time: object, + ) -> None: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: @@ -163,11 +180,15 @@ class PromptCachingDeploymentCheck(CustomLogger): cache: Final = PromptCachingCache( cache=self.cache, ) + raw_litellm_params: Final = kwargs.get("litellm_params") + litellm_params: Final[Mapping[str, object]] = ( + raw_litellm_params if isinstance(raw_litellm_params, Mapping) else {} + ) await cache.async_add_model_id( model_id=model_id, messages=prompt_messages, tools=None, # [TODO]: add tools once standard_logging_object supports it - honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object), + honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object, litellm_params), ) return diff --git a/litellm/utils.py b/litellm/utils.py index 608aa9fd112..fc3acd3065d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5940,6 +5940,15 @@ def _get_model_info_from_generalization( return None +def unmapped_model_generalization(model: str, custom_llm_provider: str | None) -> Mapping[str, object] | None: + generalization: Final = _get_model_info_from_generalization( + model=model, + potential_model_names=get_potential_model_names(model, custom_llm_provider), + custom_llm_provider=custom_llm_provider, + ) + return generalization[1] if generalization is not None else None + + def _strip_mantle_region_prefix(model: str) -> str: from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index f98dda5fb0c..1a8b2fcb66b 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -782,6 +782,31 @@ BEDROCK_INVOKE_HAIKU_STREAM: Final = CostTrackingTestCase( ), ) +BEDROCK_INVOKE_HAIKU_CHAT_STREAM: Final = CostTrackingTestCase( + name="bedrock-invoke-haiku-chat-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "bedrock-invoke-haiku-chat-stream"}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": 412, + }, + response=BEDROCK_INVOKE_HAIKU_STREAM.response, + expected=BEDROCK_INVOKE_HAIKU_STREAM.expected, +) + +BEDROCK_INVOKE_HAIKU_RESPONSES: Final = CostTrackingTestCase( + name="bedrock-invoke-haiku-responses", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + endpoint="/v1/responses", + request={"model": "$MODEL", "input": "bedrock-invoke-haiku-responses", "max_output_tokens": 412}, + response=BEDROCK_INVOKE_HAIKU_JSON.response, + expected=BEDROCK_INVOKE_HAIKU_JSON.expected, +) + # azure_ai/gpt-5.4-mini-2026-03-17 AZURE_AI_GPT_5_4_MINI_LATEST: Final = CostTrackingTestCase( @@ -2031,6 +2056,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( XAI_GROK_5_PROVIDER_REPORTED_COST, BEDROCK_INVOKE_HAIKU_JSON, BEDROCK_INVOKE_HAIKU_STREAM, + BEDROCK_INVOKE_HAIKU_CHAT_STREAM, + BEDROCK_INVOKE_HAIKU_RESPONSES, AZURE_AI_GPT_5_4_MINI_LATEST, AZURE_AI_GPT_5_4_MINI_LATEST_STREAM, AZURE_AI_GROK_4_6_REASONING_FOLDED_INTO_COMPLETION, diff --git a/tests/integration/management/test_scim_group_bulk_roster.py b/tests/integration/management/test_scim_group_bulk_roster.py index 8387beca958..2af6ae0a83d 100644 --- a/tests/integration/management/test_scim_group_bulk_roster.py +++ b/tests/integration/management/test_scim_group_bulk_roster.py @@ -18,7 +18,6 @@ from pathlib import Path from typing import Final, LiteralString import httpx -import psutil import psycopg import pytest from integration._support.budget import case_timeout_seconds @@ -55,9 +54,6 @@ LOCK_KEY_ROW_SQL: Final = 'SELECT token FROM "LiteLLM_VerificationToken" WHERE t SESSIONS_BLOCKED_BY_SQL: Final = ( "SELECT count(*)::int AS waiting FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))" ) -CLIENT_PORTS_BLOCKED_BY_SQL: Final = ( - "SELECT client_port FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid)) AND client_port > 0" -) TRANSACTIONS_OPEN_SINCE_SQL: Final = """ SELECT count(*)::int AS open FROM pg_stat_activity WHERE datname = current_database() AND backend_type = 'client backend' AND pid <> pg_backend_pid() @@ -296,22 +292,6 @@ def _transactions_open_since(moment: str) -> int: return _int_value(read_rows(TRANSACTIONS_OPEN_SINCE_SQL, (moment,)), "open") -def _client_ports_blocked_by(backend_pid: int) -> frozenset[int]: - rows: Final = read_rows(CLIENT_PORTS_BLOCKED_BY_SQL, (str(backend_pid),)) - return frozenset(_int_value((row,), "client_port") for row in rows) - - -def _local_tcp_ports(pid: int) -> Iterator[int]: - worker: Final = psutil.Process(pid) - for process in (worker, *worker.children(recursive=True)): - yield from (connection.laddr.port for connection in process.net_connections(kind="tcp")) - - -def _worker_holding(owned: OwnedProxy, client_ports: frozenset[int]) -> int: - (worker,) = tuple(pid for pid in owned.worker_pids() if client_ports & frozenset(_local_tcp_ports(pid))) - return worker - - def _roster_ids(team: str) -> tuple[str, ...]: rows: Final = read_rows(TEAM_ROW_SQL, (team,)) assert len(rows) == 1, rows @@ -508,11 +488,16 @@ def test_a_worker_killed_mid_push_leaves_the_old_roster_and_a_retry_lands_the_ne pinned_key: Final = _team_key_for(scenario, team, max(leaving)) with _holding_row_lock(LOCK_KEY_ROW_SQL, sha256(pinned_key.encode()).hexdigest()) as holder_pid: push: Final = pool.submit(_push_outcome, candidate, team, arriving) - blocked_ports: Final = eventually( - lambda: _client_ports_blocked_by(holder_pid), bool, seconds=UNFIXED_PUSH_SECONDS + _ = eventually( + lambda: _sessions_blocked_by(holder_pid), + lambda waiting: waiting > 0, + seconds=UNFIXED_PUSH_SECONDS, ) killed_at: Final = _database_now() - _ = owned.kill_worker(_worker_holding(owned, blocked_ports)) + victims: Final = tuple(sorted(owned.worker_pids())) + assert victims, "expected at least one serving worker" + for pid in victims: + _ = owned.kill_worker(pid) record_property("interrupted_push", push.result()) _ = eventually(lambda: _transactions_open_since(killed_at), lambda open: open == 0, seconds=ROLLBACK_SECONDS) settled: Final = frozenset(_roster_ids(team)) diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 9d15ea0afa9..6f59f694352 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -860,6 +860,7 @@ def test_spec_client_discovers_the_aggregate_authorization_server_and_signs_in_t "client_origin": "http://127.0.0.1:9", "server_id": None, "server_name": None, + "servers": None, "connected": None, }, described.text completed: Final = gateway.client.post("/authorize/complete", data={"flow": connect_flow}, cookies=flow_cookies) @@ -969,6 +970,7 @@ def test_per_server_resource_scopes_the_session_bearer_to_the_named_server(gatew "client_origin": "http://127.0.0.1:9", "server_id": server_a, "server_name": alias_a, + "servers": None, "connected": True, }, described.text completed: Final = gateway.client.post("/authorize/complete", data={"flow": connect_flow}, cookies=flow_cookies) diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py index b36da940a42..1c8850b9bfa 100644 --- a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py @@ -313,6 +313,8 @@ def test_nested_tool_call_before_assistant_block(candidate: Gateway, cert: tuple "response.output_item.done", "response.done", ), events + (delta,) = tuple(event for event in events if event.get("type") == "response.function_call_arguments.delta") + assert string_value(delta["delta"]) == '{"location":"New York"}' (call,) = _function_calls(events) assert string_value(call["call_id"]) == "tooluse_A" assert string_value(call["name"]) == "get_weather" diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl.py b/tests/integration/routing/test_prompt_cache_pin_ttl.py index d4328cd0e31..1ca92505191 100644 --- a/tests/integration/routing/test_prompt_cache_pin_ttl.py +++ b/tests/integration/routing/test_prompt_cache_pin_ttl.py @@ -84,7 +84,7 @@ _FAMILIES: Final = ( _Family("pin-bedrock-converse-1h", "pin-bedrock-converse-1h", one_hour=True, wire_one_hour=True), _Family("pin-bedrock-converse-5m", "pin-bedrock-converse-5m", one_hour=False, wire_one_hour=False), _Family("pin-bedrock-invoke-1h", "pin-bedrock-invoke-1h", one_hour=True, wire_one_hour=True), - _Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=False, wire_one_hour=False), + _Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=True, wire_one_hour=True), _Family("pin-bedrock-arn", "pin-bedrock-arn", one_hour=False, wire_one_hour=False), _Family("pin-azure-ai", "pin-azure-ai", one_hour=True, wire_one_hour=True), _Family("pin-vertex", "pin-vertex", one_hour=True, wire_one_hour=True), @@ -449,19 +449,39 @@ def test_provider_family_pin_lifetime_agrees_with_the_ttl_the_wire_carries(rig: spend_row_lands(response_id) -def test_messages_through_bedrock_invoke_pins_an_hour_and_sends_the_hour(rig: _Rig) -> None: +@pytest.mark.parametrize("group", ["pin-bedrock-invoke-1h", "pin-bedrock-arn-base"]) +def test_messages_through_bedrock_pins_an_hour_and_sends_the_hour(rig: _Rig, group: str) -> None: before: Final = _before(rig) marker: Final = new_marker() response: Final = rig.owned.gateway.request( - "POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": "pin-bedrock-invoke-1h"} + "POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": group} ) assert response.status_code == 200, response.text _, pin = one_new_pin(rig.redis, before) - assert pin.model_id == "pin-bedrock-invoke-1h", pin + assert pin.model_id == group, pin assert lives_one_hour(pin), pin assert [wire_carries_one_hour(request) for request in sent_with((rig.bedrock,), marker)] == [True] +def test_response_cache_hit_on_an_application_profile_rewrites_the_pin_for_its_base_model_hour(rig: _Rig) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + body: Final = chat_body("pin-bedrock-arn-base", _first_turn(new_marker(), marker, ONE_HOUR_MARK)) + first, served = post_chat(rig.owned.gateway, body) + key, pin = one_new_pin(rig.redis, before) + assert served == pin.model_id == "pin-bedrock-arn-base", (served, pin) + assert lives_one_hour(pin), pin + assert [wire_carries_one_hour(request) for request in sent_with(rig.wires(), marker)] == [True] + spend_row_lands(first) + rig.redis.delete(key) + second, _ = post_chat(rig.owned.gateway, body) + assert sent_with(rig.wires(), marker) == () + assert second == first + rewritten: Final = eventually(lambda: pin_at(rig.redis, key), lambda found: found is not None, seconds=30) + assert rewritten is not None and rewritten.model_id == "pin-bedrock-arn-base", rewritten + assert lives_one_hour(rewritten), rewritten + + @pytest.mark.parametrize("ttl", _MALFORMED_TTLS, ids=["int", "2h", "upper-1H", "empty", "null", "list", "5kb"]) def test_malformed_ttl_answers_and_pins_five_minutes_without_a_wire_ttl(rig: _Rig, ttl: JsonValue) -> None: before: Final = _before(rig) diff --git a/tests/integration/spend/test_daily_rollup_retry.py b/tests/integration/spend/test_daily_rollup_retry.py index 5b48d841f75..9ed712cdd06 100644 --- a/tests/integration/spend/test_daily_rollup_retry.py +++ b/tests/integration/spend/test_daily_rollup_retry.py @@ -142,6 +142,7 @@ def test_failed_daily_user_rollup_commit_is_retried_so_spend_report_and_daily_ac assert metrics == { "spend": pytest.approx(0.06), "flat_cost": pytest.approx(0.0), + "ptu_hours": pytest.approx(0.0), "prompt_tokens": 20, "completion_tokens": 20, "cache_read_input_tokens": 0, diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 6d1e0af5624..4b053fb5a97 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -37,6 +37,32 @@ from tests.integration._support.process import ( TESTS_DIR: Final = Path(__file__).resolve().parents[2] +@pytest.mark.parametrize("config_contents", ("[]\n", "general_settings: []\n", "general_settings: {}\n")) +def test_worker_configuration_preserves_sibling_auth_modules( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + config_contents: str, +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + monkeypatch.setenv("INTEGRATION_PROVIDER_PORT", "8203") + monkeypatch.setenv("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8202") + source: Final = tmp_path / "source.yaml" + source.write_text(config_contents) + auth_state: Final = tmp_path / "sibling_auth.json" + auth_state.write_text("authenticated") + (tmp_path / "sibling_auth.py").write_text( + "from pathlib import Path\n" + "def authenticate():\n" + " return Path(__file__).with_suffix('.json').read_text()\n" + ) + + configured: Final = worker_configuration(source, tmp_path / "worker") + + assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "authenticated" + auth_state.write_text("updated") + assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "updated" + + def test_worker_configuration_routes_stock_edges_to_its_own_servers_and_preserves_other_targets( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 018fa808f70..e5565a3df16 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -15,10 +15,12 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, + _honors_one_hour_cache_ttl, _get_min_token_count_for_deployments, ) from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt, token_counter MODEL_GROUP_ALIAS = "my-claude-group" @@ -126,6 +128,21 @@ def test_get_min_token_count_for_deployments_falls_back_to_default_for_empty_gro assert _get_min_token_count_for_deployments([]) == DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT +def test_bedrock_application_profile_uses_its_base_model_for_one_hour_ttl() -> None: + standard_logging_object: Final = cast( + StandardLoggingPayload, + { + "custom_llm_provider": "bedrock", + "model": "bedrock/converse/arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/pinprofile", + }, + ) + + assert _honors_one_hour_cache_ttl( + standard_logging_object, + {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + ) is True + + @pytest.mark.asyncio async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minimum(): """ diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 2fbe6ec294b..2d13829bf33 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -2,7 +2,7 @@ import json import asyncio import datetime import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from types import MappingProxyType, SimpleNamespace from typing import Final, Literal, cast @@ -17,6 +17,10 @@ from litellm import TranscriptionResponse, model_cost from litellm.constants import REALTIME_METERED_USAGE_KEY from litellm.cost_calculator import BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, batch_cost_calculator, completion_cost, cost_per_token, handle_realtime_stream_cost_calculation, response_cost_calculator from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.fallback_generalizations import ( + get_fallback_generalization_rules, + set_fallback_generalizations, +) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -75,6 +79,48 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) +@pytest.fixture +def _claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-baseline", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": {"mode": "chat", "max_input_tokens": 200000, "max_output_tokens": 64000}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +@pytest.fixture +def _priced_claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-priced", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": { + "mode": "chat", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "input_cost_per_token": 0.0000321, + "output_cost_per_token": 0.0000765, + }, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + def test_cost_per_token_duplicate_openai_prefix_matches_model_cost(monkeypatch): """ Router/proxy configs may use deployment ids like openai/openai/. Cost lookup must @@ -149,6 +195,265 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co assert cost > 0, "Cost should be calculated using response model" +def test_completion_cost_prefers_exact_route_price_over_capability_generalization( + _local_model_cost_map: None, + _claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + priced_entry: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + monkeypatch.setattr(litellm, "model_cost", {priced_model: priced_entry}) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=usage, + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + ) + + input_rate: Final = priced_entry["input_cost_per_token"] + output_rate: Final = priced_entry["output_cost_per_token"] + assert isinstance(input_rate, (int, float)), priced_entry + assert isinstance(output_rate, (int, float)), priced_entry + assert cost == usage.prompt_tokens * input_rate + usage.completion_tokens * output_rate + + +def test_completion_cost_keeps_priced_generalization_before_later_exact_candidate( + _local_model_cost_map: None, + _priced_claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + exact_entry: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + monkeypatch.setattr(litellm, "model_cost", {priced_model: exact_entry}) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=usage, + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx(usage.prompt_tokens * 0.0000321 + usage.completion_tokens * 0.0000765) + + +@pytest.fixture +def _per_second_priced_claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-per-second", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": {"cost_per_second": 0.25}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +def test_completion_cost_keeps_per_second_priced_generalization_before_later_exact_candidate( + _local_model_cost_map: None, + _per_second_priced_claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + monkeypatch.setattr( + litellm, + "model_cost", + { + priced_model: { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + }, + ) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + total_time=2000, + ) + + assert cost == pytest.approx(0.25 * 2) + + +@pytest.fixture +def _price_free_gpt_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "gpt-family-baseline", + "pattern": r"gpt-\d+(?:\.\d+)?", + "model_info": {"mode": "chat", "supports_reasoning": True}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +def test_completion_cost_prices_rule_only_response_model_at_the_azure_deployment_rate( + _local_model_cost_map: None, + _price_free_gpt_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + deployment_model: Final = "azure/gpt-4o-2024-11-20" + monkeypatch.setattr( + litellm, + "model_cost", + { + deployment_model: { + "litellm_provider": "azure", + "mode": "chat", + "input_cost_per_token": 0.0000111, + "output_cost_per_token": 0.0000222, + } + }, + ) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse(model="gpt-9.9", choices=[], usage=usage) + + cost: Final = completion_cost( + completion_response=response, + model=deployment_model, + custom_llm_provider="azure", + ) + + assert cost == pytest.approx(usage.prompt_tokens * 0.0000111 + usage.completion_tokens * 0.0000222) + + +def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests( + _local_model_cost_map: None, +) -> None: + model: Final = "lazy-candidate-regression-org/lazy-candidate-model" + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model=model, + choices=[], + usage=usage, + ) + + with respx.mock(assert_all_called=False) as huggingface: + huggingface.get(host="huggingface.co").mock( + return_value=httpx.Response(200, json={"max_position_embeddings": 4096}) + ) + cost: Final = completion_cost( + completion_response=response, + model=f"huggingface/{model}", + custom_llm_provider="huggingface", + ) + + assert len(huggingface.calls) == 1 + assert cost == 0.0 + + +def test_completion_cost_skips_huggingface_config_request_with_custom_pricing( + _local_model_cost_map: None, +) -> None: + model: Final = "unmapped-cost-regression-org/unmapped-model" + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model=model, + choices=[], + usage=usage, + ) + + with respx.mock(assert_all_called=False) as huggingface: + huggingface.get(host="huggingface.co").mock( + return_value=httpx.Response(200, json={"max_position_embeddings": 4096}) + ) + cost: Final = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="huggingface", + custom_pricing=True, + custom_cost_per_token={"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + + assert len(huggingface.calls) == 0 + assert cost == pytest.approx(usage.prompt_tokens * 0.001 + usage.completion_tokens * 0.002) + + +def test_completion_cost_returns_zero_for_capability_only_generalization( + _local_model_cost_map: None, + _claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model: Final = "bedrock/invoke/anthropic.claude-haiku-9-9-20991231-v1:0" + monkeypatch.setattr(litellm, "model_cost", {}) + response: Final = ModelResponse( + model="anthropic.claude-haiku-9-9-20991231-v1:0", + choices=[], + usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), + ) + + cost: Final = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + ) + + assert cost == 0.0 + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", "hosted_vllm"]) +def test_completion_cost_prices_base_model_under_its_own_provider( + custom_llm_provider: str, _local_model_cost_map: None +) -> None: + response: Final = ModelResponse( + model="my-deploy", + choices=[], + usage=Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000), + ) + + cost: Final = completion_cost( + completion_response=response, + model=f"{custom_llm_provider}/my-deploy", + custom_llm_provider=custom_llm_provider, + base_model="claude-sonnet-4-5", + ) + + info: Final = litellm.get_model_info("claude-sonnet-4-5", custom_llm_provider="anthropic") + assert cost == pytest.approx( + 1000 * info["input_cost_per_token"] + 1000 * info["output_cost_per_token"] + ) + + def test_completion_cost_strips_dated_azure_snapshot_model(_local_model_cost_map: None) -> None: dated_response = ModelResponse( model="gpt-5.6-luna-2099-01-01", diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index bb0e8db7f1c..0293af05c84 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4563,7 +4563,7 @@ def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging logging_obj: Final = LiteLLMLogging( - model="us.anthropic.claude-opus-5", + model="us.anthropic.claude-unpriced-999", messages=[{"role": "user", "content": "hi"}], stream=True, call_type="completion", @@ -4589,6 +4589,41 @@ def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): assert response._hidden_params.get("response_cost") is None +def test_stream_chunk_builder_prices_alias_response_at_its_deployment_model(): + import time as time_module + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + deployment_model: Final = "us.anthropic.claude-opus-5" + logging_obj: Final = LiteLLMLogging( + model=deployment_model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time_module.time(), + litellm_call_id="stream-builder-alias-priced", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.optional_params = {} + usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") + usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) + chunks: Final = [ + _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + entry: Final = litellm.model_cost[deployment_model] + expected_cost: Final = 40 * entry["input_cost_per_token"] + 5 * entry["output_cost_per_token"] + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(expected_cost) + assert response._hidden_params.get("response_cost") == pytest.approx(expected_cost) + + def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5)