From ee8d66644ee1fb7a992ee7505a08bc171ca29594 Mon Sep 17 00:00:00 2001 From: addy79 <125811947+addy79@users.noreply.github.com> Date: Wed, 21 Oct 2026 19:36:10 +0000 Subject: [PATCH] style: apply ruff format and clean env key reference --- litellm/__init__.py | 2 +- litellm/cost_calculator.py | 10 +- litellm/litellm_core_utils/litellm_logging.py | 6 +- .../test_register_model_custom_pricing.py | 100 +++++------------- 4 files changed, 42 insertions(+), 76 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 6adc8b45f1d..c4df163b643 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -245,7 +245,7 @@ use_chat_completions_url_for_anthropic_messages: bool = bool( # Or via `litellm_settings.strip_anthropic_total_tokens: true` in # config.yaml. strip_anthropic_total_tokens: bool = False -custom_pricing: bool = bool(os.getenv("LITELLM_CUSTOM_PRICING", False)) +custom_pricing: bool = False anthropic_sse_ping_interval_seconds: float = 15.0 sse_keepalive_ping_interval_seconds: float | None = None route_all_chat_openai_to_responses: bool = ( diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index d32c44d9bad..e01fe95dfe4 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -805,7 +805,9 @@ def cost_per_token( service_tier=service_tier, ) else: - resolved_fallback_model_info: Final = model_info or cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + resolved_fallback_model_info: Final = model_info or cached_get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider + ) if _has_token_or_tiered_pricing(resolved_fallback_model_info): return generic_cost_per_token( model=model, @@ -2061,7 +2063,11 @@ def completion_cost( model_info=( deployment_model_info if ( - (deployment_model_info := _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id)) + ( + deployment_model_info := _deployment_model_info( + litellm_logging_obj, custom_pricing, router_model_id + ) + ) and ( call_type in _RERANK_CALL_TYPES or deployment_model_info.get("input_cost_per_token") is not None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 039d2c37dd4..208e87b5d5e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2765,7 +2765,11 @@ class Logging(LiteLLMLoggingBaseClass): elif "response_cost" in hidden_params and not use_custom_pricing_for_model(self): self.model_call_details["response_cost"] = hidden_params["response_cost"] self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"]) - elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0 and not use_custom_pricing_for_model(self): + elif ( + (existing_cost := self.model_call_details.get("response_cost")) is not None + and existing_cost != 0 + and not use_custom_pricing_for_model(self) + ): # Preserve response_cost if already calculated (e.g., by pass-through # handlers like Gemini/Vertex which call completion_cost directly). # Do not preserve 0 from failure_handler on intermediate router retries. diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index a84c2d52e3f..11f1582994b 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -236,10 +236,7 @@ def test_register_model_strips_none_litellm_provider(): registered = litellm.model_cost.get(model_key) assert registered is not None, f"{model_key} should be in model_cost" # The key may be absent entirely, but if present it must not be None. - assert ( - "litellm_provider" not in registered - or registered["litellm_provider"] is not None - ) + assert "litellm_provider" not in registered or registered["litellm_provider"] is not None # Downstream consumers must accept this entry for any provider, # mirroring what the cost calculator does. assert _check_provider_match(registered, "openai") is True @@ -297,9 +294,7 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp # The strip must have removed the None-valued provider that # ``get_model_info`` returned. The key may be absent entirely, but # it must never be present with value ``None``. - assert "litellm_provider" not in registered or ( - registered["litellm_provider"] is not None - ), ( + assert "litellm_provider" not in registered or (registered["litellm_provider"] is not None), ( "register_model failed to strip litellm_provider=None returned " f"by get_model_info, got {registered.get('litellm_provider')!r}" ) @@ -360,14 +355,8 @@ def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(monkeypa ) registered = litellm.model_cost[registered_key] - assert ( - registered.get("cache_creation_input_token_cost") - == builtin["cache_creation_input_token_cost"] - ) - assert ( - registered.get("cache_read_input_token_cost") - == builtin["cache_read_input_token_cost"] - ) + assert registered.get("cache_creation_input_token_cost") == builtin["cache_creation_input_token_cost"] + assert registered.get("cache_read_input_token_cost") == builtin["cache_read_input_token_cost"] assert registered["litellm_provider"] == "bedrock" usage = Usage( @@ -431,8 +420,7 @@ def test_register_model_warns_when_no_builtin_match_for_cache_pricing(caplog): ) assert any( - registered_key in record.message - and "cache_creation_input_token_cost" in record.message + registered_key in record.message and "cache_creation_input_token_cost" in record.message for record in caplog.records ), "expected a warning naming the unmapped key and the cache cost fields" finally: @@ -619,17 +607,14 @@ def test_register_model_router_add_deployment_custom_pricing_applies(): # the deployment's provider. from litellm.utils import _check_provider_match - registered_keys = [ - k for k in (deployment_model, model_key) if k in litellm.model_cost - ] + registered_keys = [k for k in (deployment_model, model_key) if k in litellm.model_cost] assert registered_keys, ( - "Router.add_deployment did not register custom pricing for " - f"{model_key} / {deployment_model}" + f"Router.add_deployment did not register custom pricing for {model_key} / {deployment_model}" ) for k in registered_keys: - assert ( - _check_provider_match(litellm.model_cost[k], "openai") is True - ), f"custom pricing for {k} was dropped by _check_provider_match" + assert _check_provider_match(litellm.model_cost[k], "openai") is True, ( + f"custom pricing for {k} was dropped by _check_provider_match" + ) finally: litellm.model_cost.pop(model_key, None) litellm.model_cost.pop(deployment_model, None) @@ -646,12 +631,8 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing(): """ shared_key = "openai/text-embedding-3-small" deployment_id = "lit3991-wildcard-embed-zero" - snapshot = _snapshot_model_cost_entries( - [shared_key, "text-embedding-3-small", deployment_id] - ) - builtin_input_cost = litellm.get_model_info(model=shared_key)[ - "input_cost_per_token" - ] + snapshot = _snapshot_model_cost_entries([shared_key, "text-embedding-3-small", deployment_id]) + builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"] assert builtin_input_cost > 0 try: @@ -666,10 +647,9 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing(): mock_response=[0.1, 0.2], ) - assert ( - litellm.get_model_info(model=shared_key)["input_cost_per_token"] - == builtin_input_cost - ), "wildcard deployment's zero pricing leaked into the shared model_cost key" + assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost, ( + "wildcard deployment's zero pricing leaked into the shared model_cost key" + ) assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.0 assert litellm.model_cost[deployment_id]["output_cost_per_token"] == 0.0 @@ -679,9 +659,7 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing(): api_key="fake-key", mock_response=[0.1, 0.2], ) - sibling_cost = litellm.completion_cost( - completion_response=sibling_response, call_type="embedding" - ) + sibling_cost = litellm.completion_cost(completion_response=sibling_response, call_type="embedding") assert sibling_cost == pytest.approx(10 * builtin_input_cost) finally: _restore_model_cost_entries(snapshot) @@ -695,12 +673,8 @@ def test_embedding_router_custom_pricing_costs_request_via_deployment_id(): shared_key = "openai/text-embedding-3-small" deployment_id = "lit3991-wildcard-embed-custom" override_input_cost = 5e-05 - snapshot = _snapshot_model_cost_entries( - [shared_key, "text-embedding-3-small", deployment_id] - ) - builtin_input_cost = litellm.get_model_info(model=shared_key)[ - "input_cost_per_token" - ] + snapshot = _snapshot_model_cost_entries([shared_key, "text-embedding-3-small", deployment_id]) + builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"] assert builtin_input_cost != override_input_cost try: @@ -724,10 +698,7 @@ def test_embedding_router_custom_pricing_costs_request_via_deployment_id(): router_model_id=deployment_id, ) assert request_cost == pytest.approx(10 * override_input_cost) - assert ( - litellm.get_model_info(model=shared_key)["input_cost_per_token"] - == builtin_input_cost - ) + assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost finally: _restore_model_cost_entries(snapshot) @@ -736,12 +707,8 @@ def test_completion_router_zero_pricing_does_not_clobber_builtin_pricing(): """Same isolation as the embedding path, exercised through completion().""" shared_key = "openai/gpt-4o-mini" deployment_id = "lit3991-wildcard-chat-zero" - snapshot = _snapshot_model_cost_entries( - [shared_key, "gpt-4o-mini", deployment_id] - ) - builtin_input_cost = litellm.get_model_info(model=shared_key)[ - "input_cost_per_token" - ] + snapshot = _snapshot_model_cost_entries([shared_key, "gpt-4o-mini", deployment_id]) + builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"] assert builtin_input_cost > 0 try: @@ -756,10 +723,9 @@ def test_completion_router_zero_pricing_does_not_clobber_builtin_pricing(): mock_response="hello back", ) - assert ( - litellm.get_model_info(model=shared_key)["input_cost_per_token"] - == builtin_input_cost - ), "wildcard deployment's zero pricing leaked into the shared model_cost key" + assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost, ( + "wildcard deployment's zero pricing leaked into the shared model_cost key" + ) assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.0 finally: _restore_model_cost_entries(snapshot) @@ -782,10 +748,7 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key(): mock_response=[0.1, 0.2], ) - assert ( - litellm.model_cost[model_key]["input_cost_per_token"] - == override_input_cost - ) + assert litellm.model_cost[model_key]["input_cost_per_token"] == override_input_cost cost = litellm.completion_cost( completion_response=response, model=model_key, @@ -1745,9 +1708,7 @@ def test_rerank_custom_pricing_global_and_register_model(monkeypatch: pytest.Mon api_key="fake-cohere-key", ) assert response.response_cost == pytest.approx(0.005) - calc_cost = litellm.completion_cost( - completion_response=response, model=custom_model, call_type="rerank" - ) + calc_cost = litellm.completion_cost(completion_response=response, model=custom_model, call_type="rerank") assert calc_cost == pytest.approx(0.005) finally: _restore_model_cost_entries(snapshot) @@ -1777,9 +1738,7 @@ def test_embedding_custom_pricing_global_and_register_model(monkeypatch: pytest. ) assert response.response_cost is not None assert response.response_cost > 0 - calc_cost = litellm.completion_cost( - completion_response=response, model=custom_model, call_type="embedding" - ) + calc_cost = litellm.completion_cost(completion_response=response, model=custom_model, call_type="embedding") assert calc_cost == pytest.approx(response.response_cost) finally: _restore_model_cost_entries(snapshot) @@ -1814,9 +1773,7 @@ def test_rerank_per_request_custom_pricing(): @pytest.mark.asyncio -async def test_rerank_router_deployment_custom_pricing( - monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter -): +async def test_rerank_router_deployment_custom_pricing(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() try: @@ -1851,4 +1808,3 @@ async def test_rerank_router_deployment_custom_pricing( assert public_response.response_cost == pytest.approx(0.004) finally: litellm.in_memory_llm_clients_cache.flush_cache() -