diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index defe1e63b86..6677b7327dc 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -153,7 +153,9 @@ from litellm.utils import ( ProviderConfigManager, TextCompletionResponse, TranscriptionResponse, + cached_cost_map_key, cached_get_model_info_helper, + mapped_model_fill_missing, model_info_with_pricing_overrides, runtime_model_cost_registrations, token_counter, @@ -962,21 +964,17 @@ def select_model_name_for_cost_calc( """ 1. If custom pricing is true, return received model name 2. If base_model is set (e.g. for azure models), return that - 3. If completion response has model set return that + 3. If completion response has model set return that, unless only the requested model's cost-map row has rates 4. Check if model is passed in return that """ return_model: str | None = None custom_llm_provider = _get_provider_for_cost_calc(model=model, custom_llm_provider=custom_llm_provider) - completion_response_model: str | None = None - if completion_response is not None: - if isinstance(completion_response, BaseModel): - completion_response_model = getattr(completion_response, "model", None) - elif isinstance(completion_response, dict): - completion_response_model = completion_response.get("model", None) + completion_response_model: Final = _get_response_model(completion_response) hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None) provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model") + requested_model: Final = model or _get_hidden_str_for_cost_calc(hidden_params, "litellm_model_name") explicit_pricing: Final = custom_pricing is True or base_model is not None configured_model: Final = model or ( completion_response_model if isinstance(completion_response_model, str) else None @@ -1002,35 +1000,72 @@ def select_model_name_for_cost_calc( else: return_model = model - elif base_model is not None or bedrock_pricing_model is not None or provider_response_model is not None: - return_model = base_model if base_model is not None else bedrock_pricing_model or provider_response_model + elif base_model is not None or bedrock_pricing_model is not None: + return_model = base_model if base_model is not None else bedrock_pricing_model + + elif provider_response_model is not None: + return _served_or_priced_requested_model( + provider_response_model, requested_model, custom_llm_provider, priced_region + ) elif completion_response_model is None and hidden_params is not None: if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0: return_model = hidden_params.get("model", model) if return_model is None and completion_response_model is not None: - return_model = completion_response_model + return _served_or_priced_requested_model( + completion_response_model, requested_model, custom_llm_provider, priced_region + ) if return_model is None and model is not None: return_model = model - if ( - return_model is not None - and custom_llm_provider is not None - and not _model_contains_known_llm_provider(return_model) - ): # add provider prefix if not already present, to match model_cost - provider_prefix: Final = ( - custom_llm_provider if priced_region is None else f"{custom_llm_provider}/{priced_region}" - ) - return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", priced_region) - - return return_model + if return_model is None: + return None + return _cost_map_model_name(return_model, custom_llm_provider, priced_region) _select_model_name_for_cost_calc = select_model_name_for_cost_calc +def _cost_map_model_name(model: str, custom_llm_provider: str | None, region_name: str | None) -> str: + if custom_llm_provider is None or _model_contains_known_llm_provider(model): + return model + provider_prefix: Final = custom_llm_provider if region_name is None else f"{custom_llm_provider}/{region_name}" + return _strip_unregistered_leading_segments(f"{provider_prefix}/{model}", region_name) + + +def _rule_declares_rates(model: str, custom_llm_provider: str | None) -> bool: + rule_info: Final = unmapped_model_generalization(model, custom_llm_provider) + return rule_info is not None and _cost_map_entry_prices_anything(rule_info) + + +def _resolves_to_rates(model: str, custom_llm_provider: str | None) -> bool: + key: Final = cached_cost_map_key(model, custom_llm_provider) + if key is None: + return False + entry: Final = _raw_cost_map_entry(key) + if entry is None: + return _rule_declares_rates(model, custom_llm_provider) + fill_missing: Final = mapped_model_fill_missing(key, entry) + return _cost_map_entry_prices_anything(entry if fill_missing is None else {**fill_missing, **entry}) + + +def _served_or_priced_requested_model( + served_model: str, requested_model: str | None, custom_llm_provider: str | None, priced_region: str | None +) -> str: + served: Final = _cost_map_model_name(served_model, custom_llm_provider, priced_region) + if requested_model is None: + return served + requested: Final = _cost_map_model_name(requested_model, custom_llm_provider, None) + if served in (requested, requested.partition("/")[2]): + return served + pricing_provider: Final = _get_provider_for_cost_calc(requested, custom_llm_provider) + if _resolves_to_rates(served, pricing_provider): + return served + return requested if _resolves_to_rates(requested, pricing_provider) else served + + def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str: """Resolve a provider-prefixed slash alias like "vertex_ai/vertex/claude-opus-5" to the registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already @@ -3516,7 +3551,7 @@ def _cost_map_entry_declares_pricing(model_name: str, custom_llm_provider: str) lookups cannot strip that prefix, so a zero-rated override read as declaring nothing, and a session that should bill nothing fell through to the public rates instead. """ - resolved: Final = _get_model_info_or_none(model_name, custom_llm_provider) + resolved: Final = _cost_map_model_info(model_name, custom_llm_provider) entries: Final = ( litellm.model_cost.get(resolved.get("key")) if resolved is not None else None, litellm.model_cost.get(model_name), @@ -3712,9 +3747,9 @@ def handle_realtime_transcription_cost_calculation( return 0.0 model_name: Final = _get_transcription_model_name_from_results(results) or litellm_model_name - model_info: Final = _get_model_info_or_none(model_name, custom_llm_provider) + model_info: Final = _cost_map_model_info(model_name, custom_llm_provider) override_info: Final = ( - _get_model_info_or_none(custom_pricing_model, custom_llm_provider) if custom_pricing_model is not None else None + _cost_map_model_info(custom_pricing_model, custom_llm_provider) if custom_pricing_model is not None else None ) total_cost = 0.0 @@ -3745,13 +3780,6 @@ def _get_transcription_model_name_from_results( return None -def _get_model_info_or_none(model: str, custom_llm_provider: str) -> ModelInfo | None: - try: - return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - return None - - _TranscriptionRateKey: TypeAlias = Literal[ "input_cost_per_second", "input_cost_per_audio_token", "input_cost_per_token", "output_cost_per_token" ] diff --git a/litellm/utils.py b/litellm/utils.py index d951b80d49b..41a3fddf3a3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -279,7 +279,7 @@ from litellm.types.utils import ( ) _CALL_TYPE_ENUM_MAP: Final[dict] = {ct.value: ct for ct in CallTypes} -_BACKFILL_MODES: Final = frozenset({"chat", "responses"}) +_BACKFILL_MODES: Final[frozenset[str]] = frozenset({"chat", "responses"}) # +-----------------------------------------------+ # | | @@ -6530,6 +6530,7 @@ def _invalidate_model_cost_lowercase_map() -> None: # Clear LRU caches that depend on model_cost data _cached_get_model_info.cache_clear() cached_get_model_info_helper.cache_clear() + cached_cost_map_key.cache_clear() def _rebuild_model_cost_lowercase_map() -> dict[str, str]: @@ -6755,6 +6756,14 @@ def unmapped_model_generalization(model: str, custom_llm_provider: str | None) - return generalization[1] if generalization is not None else None +def mapped_model_fill_missing(key: str, model_info: Mapping[str, object]) -> Mapping[str, object] | None: + mode: Final = model_info.get("mode", "chat") + provider: Final = model_info.get("litellm_provider", "") + if not isinstance(mode, str) or mode not in _BACKFILL_MODES or not isinstance(provider, str): + return None + return match_fill_missing_generalizations(key, provider) + + def _strip_mantle_region_prefix(model: str) -> str: from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix @@ -6883,6 +6892,14 @@ def cached_get_model_info_helper( _cached_get_model_info_helper = cached_get_model_info_helper +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def cached_cost_map_key(model: str, custom_llm_provider: str | None) -> str | None: + try: + return cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)["key"] + except Exception: + return None + + def get_provider_info(model: str, custom_llm_provider: str | None) -> ProviderSpecificModelInfo | None: ## PROVIDER-SPECIFIC INFORMATION # if custom_llm_provider == "predibase": @@ -7040,8 +7057,8 @@ def get_model_info_helper( key: str | None = lookup[0] _model_info: dict[str, Any] | None = lookup[1] - if _model_info is not None and key is not None and _model_info.get("mode", "chat") in _BACKFILL_MODES: - fill_missing: Final = match_fill_missing_generalizations(key, _model_info.get("litellm_provider", "")) + if _model_info is not None and key is not None: + fill_missing: Final = mapped_model_fill_missing(key, _model_info) if fill_missing is not None: _model_info = { **{k: v for k, v in fill_missing.items() if k not in _model_info}, diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index c320cb36d73..b3ae3bf6240 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -12320,6 +12320,49 @@ def test_get_model_cost_information_custom_pricing_uses_base_model(): assert result["model_map_key"] != "invoke_test_claude" +@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") +def test_get_model_cost_information_maps_requested_model_when_response_names_an_unmapped_display_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requested_model: Final = "fireworks_ai/accounts/fireworks/models/sol-display-test" + monkeypatch.setitem( + litellm.model_cost, + requested_model, + {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + response: Final = ModelResponse(model="Sol Display Max", choices=[]) + response._hidden_params["provider_response_model"] = "Sol Display Max" + response._hidden_params["litellm_model_name"] = requested_model + + result: Final = StandardLoggingPayloadSetup.get_model_cost_information( + base_model=None, custom_pricing=False, custom_llm_provider="fireworks_ai", init_response_obj=response + ) + + assert result["model_map_key"] == requested_model + assert result["model_map_value"] is not None + assert result["model_map_value"]["input_cost_per_token"] == litellm.model_cost[requested_model]["input_cost_per_token"] + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") +def test_get_model_cost_information_keeps_a_priced_response_model_that_echoes_without_its_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + served_row: Final = {"litellm_provider": "mistral", "mode": "chat", "input_cost_per_token": 2e-6, "output_cost_per_token": 5e-6} + requested_row: Final = {"litellm_provider": "mistral", "mode": "chat", "input_cost_per_token": 9e-6, "output_cost_per_token": 9e-6} + monkeypatch.setitem(litellm.model_cost, "mistral/sol-snapshot-test", served_row) + monkeypatch.setitem(litellm.model_cost, "mistral/sol-alias-test", requested_row) + response: Final = ModelResponse(model="sol-snapshot-test", choices=[]) + response._hidden_params["litellm_model_name"] = "mistral/sol-alias-test" + + result: Final = StandardLoggingPayloadSetup.get_model_cost_information( + base_model=None, custom_pricing=False, custom_llm_provider="mistral", init_response_obj=response + ) + + assert result["model_map_key"] == "sol-snapshot-test" + assert result["model_map_value"] is not None + assert result["model_map_value"]["input_cost_per_token"] == served_row["input_cost_per_token"] + + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_standard_logging_payload_uses_deployment_when_no_base_model(): """metadata["deployment"] is used for cost-map lookup when base_model is not set.""" diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 84f6dab9eff..f10baccd1cb 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -37,7 +37,7 @@ from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import CacheCreationTokenDetails, CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject -from litellm.utils import supports_prompt_caching +from litellm.utils import _invalidate_model_cost_lowercase_map, supports_prompt_caching import os import threading from unittest.mock import MagicMock, patch @@ -424,6 +424,313 @@ def test_completion_cost_prices_rule_only_base_model_from_the_response_model( assert cost == pytest.approx(usage.prompt_tokens * 0.0000111 + usage.completion_tokens * 0.0000222) +def _token_rates_cost(row: Mapping[str, object], usage: Usage) -> float: + input_rate: Final = row["input_cost_per_token"] + output_rate: Final = row["output_cost_per_token"] + assert isinstance(input_rate, float) and isinstance(output_rate, float), row + return usage.prompt_tokens * input_rate + usage.completion_tokens * output_rate + + +@pytest.mark.parametrize( + "requested_model", + ["accounts/fireworks/models/sol-display-test", "fireworks_ai/accounts/fireworks/models/sol-display-test"], +) +def test_completion_cost_bills_requested_model_when_response_names_an_unmapped_display_name( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, requested_model: str +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "fireworks_ai/accounts/fireworks/models/sol-display-test", + {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + cost: Final = completion_cost( + completion_response=ModelResponse(model="Sol Display Max", choices=[], usage=usage), + model=requested_model, + custom_llm_provider="fireworks_ai", + ) + + assert cost == pytest.approx( + _token_rates_cost(litellm.model_cost["fireworks_ai/accounts/fireworks/models/sol-display-test"], usage) + ) + + +def test_completion_cost_bills_requested_model_when_stream_reports_an_unmapped_display_name( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + requested_model: Final = "fireworks_ai/accounts/fireworks/models/sol-display-test" + monkeypatch.setitem( + litellm.model_cost, + requested_model, + {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + response: Final = ModelResponse(model="accounts/fireworks/models/sol-display-test", choices=[], usage=usage) + response._hidden_params["provider_response_model"] = "Sol Display Max" + + cost: Final = completion_cost( + completion_response=response, model=requested_model, custom_llm_provider="fireworks_ai" + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[requested_model], usage)) + + +def test_completion_cost_bills_requested_model_when_response_row_has_no_rates( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "openai/sol-requested-test", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + monkeypatch.setitem( + litellm.model_cost, + "openai/sol-served-snapshot-test", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": None, "output_cost_per_token": None}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + cost: Final = completion_cost( + completion_response=ModelResponse(model="sol-served-snapshot-test", choices=[], usage=usage), + model="openai/sol-requested-test", + custom_llm_provider="openai", + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["openai/sol-requested-test"], usage)) + + +def test_completion_cost_bills_priced_response_model_over_priced_requested_model( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "openai/sol-requested-test", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + monkeypatch.setitem( + litellm.model_cost, + "openai/sol-requested-test-2099-01-01", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 5e-6, "output_cost_per_token": 11e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + cost: Final = completion_cost( + completion_response=ModelResponse(model="sol-requested-test-2099-01-01", choices=[], usage=usage), + model="openai/sol-requested-test", + custom_llm_provider="openai", + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["openai/sol-requested-test-2099-01-01"], usage)) + + +def test_completion_cost_keeps_response_pricing_when_requested_alias_has_no_rates( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + size_bucket: Final = get_base_model_for_pricing("Sol 70B Display") + monkeypatch.setitem( + litellm.model_cost, + size_bucket, + {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + cost: Final = completion_cost( + completion_response=ModelResponse(model="Sol 70B Display", choices=[], usage=usage), + model="sol-team-alias", + custom_llm_provider="fireworks_ai", + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[size_bucket], usage)) + + +def test_completion_cost_bills_requested_model_priced_only_by_a_rule_when_response_row_has_no_rates( + _local_model_cost_map: None, _priced_claude_family_generalization: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "anthropic/sol-served-snapshot-test", + {"litellm_provider": "anthropic", "mode": "chat", "input_cost_per_token": None, "output_cost_per_token": None}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + cost: Final = completion_cost( + completion_response=ModelResponse(model="sol-served-snapshot-test", choices=[], usage=usage), + model="anthropic/claude-sol-9", + custom_llm_provider="anthropic", + ) + + assert cost == pytest.approx(usage.prompt_tokens * 0.0000321 + usage.completion_tokens * 0.0000765) + + +def test_completion_cost_bills_requested_model_from_hidden_params_when_no_model_is_passed( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + requested_model: Final = "fireworks_ai/accounts/fireworks/models/sol-display-test" + monkeypatch.setitem( + litellm.model_cost, + requested_model, + {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + response: Final = ModelResponse(model="Sol Display Max", choices=[], usage=usage) + response._hidden_params["litellm_model_name"] = requested_model + + cost: Final = completion_cost(completion_response=response) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[requested_model], usage)) + + +def test_completion_cost_bills_requested_model_regional_row_when_response_names_an_unmapped_display_name( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "bedrock/sol-regional-test", + {"litellm_provider": "bedrock", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + monkeypatch.setitem( + litellm.model_cost, + "bedrock/eu-central-1/sol-regional-test", + {"litellm_provider": "bedrock", "mode": "chat", "input_cost_per_token": 5e-6, "output_cost_per_token": 11e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + response: Final = ModelResponse(model="Sol Regional Display", choices=[], usage=usage) + response._hidden_params.update( + {"litellm_model_name": "bedrock/sol-regional-test", "region_name": "eu-central-1", "custom_llm_provider": "bedrock"} + ) + + cost: Final = completion_cost(completion_response=response) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["bedrock/eu-central-1/sol-regional-test"], usage)) + + +def test_completion_cost_keeps_served_row_rates_filled_by_a_rule_over_the_requested_alias( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem(litellm.model_cost, "sol-fill-snapshot-test", {"litellm_provider": "openai", "mode": "chat"}) + monkeypatch.setitem( + litellm.model_cost, + "sol-fill-alias-test", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + filled_rates: Final = {"input_cost_per_token": 5e-6, "output_cost_per_token": 11e-6} + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + response: Final = ModelResponse(model="sol-fill-snapshot-test", choices=[], usage=usage) + response._hidden_params.update({"litellm_model_name": "openai/sol-fill-alias-test", "custom_llm_provider": "openai"}) + previous_rules: Final = list(get_fallback_generalization_rules()) + set_fallback_generalizations( + [ + { + "name": "sol-fill-snapshot-rates", + "pattern": r"^sol-fill-snapshot-test$", + "fill_missing_for_providers": ["openai"], + "model_info": filled_rates, + } + ] + ) + try: + cost: Final = completion_cost(completion_response=response) + finally: + set_fallback_generalizations(previous_rules) + + assert cost == pytest.approx(_token_rates_cost(filled_rates, usage)) + + +_SIZE_BUCKET_CASES: Final = ( + pytest.param( + "fireworks_ai", + "fireworks_ai/accounts/fireworks/models/sol-bucket-requested-test", + "Sol Bucket 405B Display", + "fireworks-ai-above-16b", + id="fireworks_ai", + ), + pytest.param( + "together_ai", + "together_ai/sol-bucket-requested-test", + "meta-llama/Sol-Bucket-8B-Test", + "together-ai-4.1b-8b", + id="together_ai", + ), +) + + +def _size_bucket_response(served_model: str, requested_model: str, custom_llm_provider: str) -> ModelResponse: + response: Final = ModelResponse(model=served_model, choices=[], usage=_SIZE_BUCKET_USAGE) + response._hidden_params.update({"litellm_model_name": requested_model, "custom_llm_provider": custom_llm_provider}) + return response + + +_SIZE_BUCKET_USAGE: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + +@pytest.mark.parametrize(("custom_llm_provider", "requested_model", "served_model", "bucket"), _SIZE_BUCKET_CASES) +def test_completion_cost_bills_requested_listed_rates_over_a_served_size_bucket( + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + custom_llm_provider: str, + requested_model: str, + served_model: str, + bucket: str, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + requested_model, + { + "litellm_provider": custom_llm_provider, + "mode": "chat", + "input_cost_per_token": 3e-6, + "output_cost_per_token": 7e-6, + }, + ) + + cost: Final = completion_cost( + completion_response=_size_bucket_response(served_model, requested_model, custom_llm_provider) + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[requested_model], _SIZE_BUCKET_USAGE)) + assert cost != pytest.approx(_token_rates_cost(litellm.model_cost[bucket], _SIZE_BUCKET_USAGE)) + + +@pytest.mark.parametrize(("custom_llm_provider", "requested_model", "served_model", "bucket"), _SIZE_BUCKET_CASES) +def test_completion_cost_keeps_served_size_bucket_rates_when_the_requested_model_is_unpriced( + _local_model_cost_map: None, + custom_llm_provider: str, + requested_model: str, + served_model: str, + bucket: str, +) -> None: + cost: Final = completion_cost( + completion_response=_size_bucket_response(served_model, requested_model, custom_llm_provider) + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[bucket], _SIZE_BUCKET_USAGE)) + + +def test_completion_cost_bills_a_requested_model_priced_after_an_unpriced_lookup( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + requested_model: Final = "fireworks_ai/accounts/fireworks/models/sol-late-priced-test" + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + def display_name_response() -> ModelResponse: + response: Final = ModelResponse(model="Sol Late Display", choices=[], usage=usage) + response._hidden_params.update({"litellm_model_name": requested_model, "custom_llm_provider": "fireworks_ai"}) + return response + + completion_cost(completion_response=display_name_response()) + monkeypatch.setitem( + litellm.model_cost, + requested_model, + {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + _invalidate_model_cost_lowercase_map() + + cost: Final = completion_cost(completion_response=display_name_response()) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[requested_model], usage)) + + def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests( _local_model_cost_map: None, ) -> None: