From 4f53e29bedba4b6f47c88726d6460c9ca2d6e3fc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 2 Nov 2026 12:18:52 -0800 Subject: [PATCH 1/4] fix(prometheus): retain deployment attribution on proxy failures --- litellm/integrations/prometheus.py | 12 ++++-- .../integrations/test_prometheus_labels.py | 42 +++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 42fe0837c29..1c9a443b546 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3036,7 +3036,7 @@ class PrometheusLogger(CustomLogger): async def async_post_call_failure_hook( self, - request_data: dict, + request_data: dict[str, object], original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, traceback_str: str | None = None, @@ -3059,10 +3059,14 @@ class PrometheusLogger(CustomLogger): try: _tags: Final = StandardLoggingPayloadSetup.get_request_tags( litellm_params=request_data, - proxy_server_request=request_data.get("proxy_server_request", {}), + proxy_server_request=_logged_fields(request_data.get("proxy_server_request")), + ) + _metadata: Final = _logged_fields(request_data.get("metadata")) + model_id: Final = ( + _logged_fields(_metadata.get("model_info")).get("id") + or _logged_fields(request_data.get("model_info")).get("id") + or _logged_fields(request_data.get("standard_logging_object")).get("model_id") ) - _metadata: Final = request_data.get("metadata", {}) or {} - model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id") rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception) api_provider: Final = self._extract_api_provider_from_request_data( request_data diff --git a/tests/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index 878f19c66f8..04eb7c3a5a4 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -5,6 +5,7 @@ Unit tests for prometheus metric labels configuration from typing import Final, Literal import pytest +from pydantic import JsonValue import litellm from litellm.types.integrations.prometheus import ( @@ -802,6 +803,47 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric() _clear_prometheus_registry() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata_id", "request_id", "logged_id", "expected"), + ( + (None, None, "logged-deployment", "logged-deployment"), + ("", None, "logged-deployment", "logged-deployment"), + ("metadata-deployment", "request-deployment", "logged-deployment", "metadata-deployment"), + (None, "request-deployment", "logged-deployment", "request-deployment"), + (None, None, None, "None"), + ), +) +async def test_failure_counters_keep_logged_deployment_attribution( + metadata_id: str | None, request_id: str | None, logged_id: str | None, expected: str +) -> None: + from fastapi import HTTPException + + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import UserAPIKeyAuth + + request_data: Final[dict[str, JsonValue]] = { + "model": "metric-test-alias", + "metadata": {"model_info": {"id": metadata_id}}, + "model_info": {"id": request_id}, + "standard_logging_object": {"model_id": logged_id, "custom_llm_provider": "openai"}, + } + _clear_prometheus_registry() + try: + await PrometheusLogger().async_post_call_failure_hook( + request_data=request_data, + original_exception=HTTPException(status_code=400, detail="guardrail rejected"), + user_api_key_dict=UserAPIKeyAuth(token="metric-test-token"), + ) + for metric_name in ("litellm_proxy_failed_requests_metric_total", "litellm_proxy_total_requests_metric_total"): + samples: Final = _collected_samples(metric_name) + assert [(sample.labels["model_id"], sample.value) for sample in samples] == [ + (expected, 1.0) + ] + finally: + _clear_prometheus_registry() + + @pytest.mark.asyncio async def test_failure_hook_labels_a_cli_session_with_the_per_user_alias_not_the_login_token(): from litellm.integrations.prometheus import PrometheusLogger From 322457c62b76e839e100d385e96cb341cd763d3f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 2 Nov 2026 13:09:06 -0800 Subject: [PATCH 2/4] fix(ollama): keep capability and pricing lookups local --- litellm/cost_calculator.py | 12 ++-- litellm/litellm_core_utils/litellm_logging.py | 4 +- .../litellm_core_utils/llm_cost_calc/utils.py | 8 ++- .../llms/base_llm/decisions/transformation.py | 9 ++- litellm/utils.py | 33 ++++++++++ .../llm_cost_calc/test_utils.py | 62 +++++++++++++++++++ .../test_litellm_logging.py | 21 +++++++ .../test_base_decisions_transformation.py | 48 ++++++++++++++ tests/unit/test_utils.py | 24 +++++++ 9 files changed, 210 insertions(+), 11 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index e66be233be4..5b39df887d8 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -155,6 +155,8 @@ from litellm.utils import ( TranscriptionResponse, cached_cost_map_key, cached_get_model_info_helper, + get_model_info_for_pricing, + get_model_info_helper_for_pricing, mapped_model_fill_missing, model_info_with_pricing_overrides, runtime_model_cost_registrations, @@ -379,7 +381,7 @@ def _per_second_pricing_cost( audio_seconds: float = 0.0, ) -> tuple[float, float] | None: try: - model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + model_info: Final = get_model_info_helper_for_pricing(model=model, custom_llm_provider=custom_llm_provider) except Exception: # noqa: BLE001 # the lookup raises plain Exception for an unmapped model return None if _has_token_or_tiered_pricing(model_info) or not _bills_wall_clock_seconds(model_info): @@ -600,7 +602,9 @@ def cost_per_token( model = model_without_prefix if is_decisions_call_type(call_type): - decisions_model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + decisions_model_info: Final = get_model_info_helper_for_pricing( + model=model, custom_llm_provider=custom_llm_provider + ) decisions_overrides: Final = decisions_pricing_fields( litellm.model_cost.get(decisions_model_info["key"], {}) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # model cost map is dynamically loaded ) @@ -808,7 +812,7 @@ def cost_per_token( service_tier=service_tier, ) else: - model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + model_info: Final = get_model_info_helper_for_pricing(model=model, custom_llm_provider=custom_llm_provider) if _has_token_or_tiered_pricing(model_info): return generic_cost_per_token( model=model, @@ -2604,7 +2608,7 @@ def _layered_ocr_pricing(*sources: Mapping[str, object] | None) -> OCRPricing: def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None: try: - return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + return get_model_info_for_pricing(model=model, custom_llm_provider=custom_llm_provider) except Exception: return None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 639003c9a6a..1a303901314 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -211,7 +211,7 @@ from litellm.types.utils import ( batch_chat_usage, ) from litellm.types.videos.main import VideoObject -from litellm.utils import get_base_model_from_metadata, print_verbose +from litellm.utils import get_base_model_from_metadata, get_model_info_for_pricing, print_verbose from ..integrations.argilla import ArgillaLogger from ..integrations.arize.arize_phoenix import ArizePhoenixLogger @@ -7113,7 +7113,7 @@ class StandardLoggingPayloadSetup: model_cost_information = StandardLoggingModelInformation(model_map_key="", model_map_value=None) else: try: - _model_cost_information: Final = litellm.get_model_info( + _model_cost_information: Final = get_model_info_for_pricing( model=pricing_model_name or model_cost_name, custom_llm_provider=custom_llm_provider, api_base=api_base, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 50c26f4d4ca..f8a851ee880 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -35,7 +35,7 @@ from litellm.types.utils import ( Usage, text_tokens_without_nested_reasoning, ) -from litellm.utils import get_model_info +from litellm.utils import get_model_info, get_model_info_for_pricing # Pre-resolved CallTypes enum values for fast membership checks _IMAGE_RESPONSE_CALL_TYPES: Final = frozenset( @@ -1491,7 +1491,9 @@ def generic_cost_per_token( # per-deployment override: those are registered under the deployment id and kept off # the shared model-name key, so resolving from the name here reads the public rate. resolved_model_info: Final = apply_provider_cache_read_default( - get_model_info(model=model, custom_llm_provider=custom_llm_provider) if model_info is None else model_info, + get_model_info_for_pricing(model=model, custom_llm_provider=custom_llm_provider) + if model_info is None + else model_info, custom_llm_provider, ) @@ -1824,7 +1826,7 @@ def get_billed_token_rates( resolved_model_info: Final = apply_provider_cache_read_default( model_info if model_info is not None - else get_model_info(model=model, custom_llm_provider=custom_llm_provider), + else get_model_info_for_pricing(model=model, custom_llm_provider=custom_llm_provider), custom_llm_provider, ) except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index c562c8d0f04..ed1b4011358 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -60,6 +60,7 @@ from litellm.types.decisions import ( systemone_choice_key, ) from litellm.types.utils import is_decisions_model_mode +from litellm.utils import get_model_info_helper _PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object) _SYSTEMONE_RESPONSE_ADAPTER: Final[TypeAdapter[DecisionsResponse]] = TypeAdapter(DecisionsResponse) @@ -500,7 +501,9 @@ def ir_to_systemone_response(response: DecisionsIRResponse, request: DecisionsIR def catalog_mode(model: str, custom_llm_provider: str) -> str | None: try: - key: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)["key"] + key: Final = get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider, catalog_only=True)[ + "key" + ] except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model, which the catalog does not list return None entry: Final = GetModelCostMap.loaded_model_cost_map().get(key) @@ -510,7 +513,9 @@ def catalog_mode(model: str, custom_llm_provider: str) -> str | None: def _is_decisions_model(model: str, custom_llm_provider: str) -> bool: try: - model_mode: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get("mode") + model_mode: Final = get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider, catalog_only=True + ).get("mode") except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model, which has no mode return False return is_decisions_model_mode(model_mode) diff --git a/litellm/utils.py b/litellm/utils.py index b0826f84006..61deb8baa18 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7542,6 +7542,39 @@ get_model_info.cache_clear = _cached_get_model_info.cache_clear get_model_info.cache_info = _cached_get_model_info.cache_info +def get_model_info_for_pricing( + model: str, custom_llm_provider: str | None = None, api_base: str | None = None +) -> ModelInfo: + from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider + + provider: Final = declared_provider(model, custom_llm_provider) + if provider not in ("ollama", "ollama_chat"): + return get_model_info(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base) + try: + catalog: Final = get_model_info_helper(model=model, custom_llm_provider=provider, catalog_only=True) + except ModelNotMappedError: + return ModelInfo( + key=model.removeprefix(f"{provider}/"), + litellm_provider="ollama", + mode="chat", + input_cost_per_token=0.0, + output_cost_per_token=0.0, + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, + supported_openai_params=None, + ) + return ModelInfo(**catalog, supported_openai_params=None) + + +def get_model_info_helper_for_pricing(model: str, custom_llm_provider: str | None) -> ModelInfoBase: + from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider + + if declared_provider(model, custom_llm_provider) in ("ollama", "ollama_chat"): + return get_model_info_for_pricing(model=model, custom_llm_provider=custom_llm_provider) + return cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + + def json_schema_type(python_type_name: str): """Converts standard python types to json schema types diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py index b6c18e0da86..4025b7ce725 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py @@ -9,11 +9,73 @@ import pytest import respx import litellm +from litellm.cost_calculator import cost_per_token, pricing_entry_for_cost_calc from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token, get_billed_token_rates from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.types.llms.base import CachedTokensDetails from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage + +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("registered", (False, True)) +@pytest.mark.parametrize("qualified", (False, True)) +def test_ollama_pricing_uses_registered_rates_without_contacting_the_server( + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + provider: str, + registered: bool, + qualified: bool, +) -> None: + model: Final = "test-ollama-pricing-registered" if registered else "test-ollama-pricing-unmapped" + monkeypatch.setenv("OLLAMA_API_BASE", "https://ambient-pricing.example") + metadata: Final = respx_mock.post("https://ambient-pricing.example/api/show").respond( + json={"model_info": {"test.context_length": 32768}} + ) + if registered: + monkeypatch.setitem( + litellm.model_cost, + f"{provider}/{model}", + { + "litellm_provider": provider, + "mode": "chat", + "input_cost_per_token": 2e-06, + "output_cost_per_token": 5e-06, + "input_cost_per_token_decisions": 3e-06, + "output_cost_per_token_decisions": 7e-06, + }, + ) + usage: Final = Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120) + cost_model: Final = f"{provider}/{model}" if qualified else model + cost_provider: Final = None if qualified else provider + + assert cost_per_token( + model=cost_model, custom_llm_provider=cost_provider, usage_object=usage, call_type="systemone" + ) == pytest.approx((0.0003, 0.00014) if registered else (0.0, 0.0)) + assert generic_cost_per_token(model=cost_model, custom_llm_provider=cost_provider, usage=usage) == pytest.approx( + (0.0002, 0.0001) if registered else (0.0, 0.0) + ) + rates: Final = get_billed_token_rates(model=cost_model, custom_llm_provider=cost_provider, usage=usage) + assert rates is not None + assert (rates.input_cost_per_token, rates.output_cost_per_token) == ( + (2e-06, 5e-06) if registered else (0.0, 0.0) + ) + diagnostic: Final = pricing_entry_for_cost_calc( + model=cost_model, + completion_response=None, + custom_llm_provider=cost_provider, + custom_pricing=False, + base_model=None, + router_model_id=None, + region_name=None, + litellm_logging_obj=None, + ) + assert diagnostic is not None + assert diagnostic[1]["input_cost_per_token"] == (2e-06 if registered else 0.0) + assert not metadata.called + + TIER_MODEL: Final = "tier-priced-test-model" TIER_ROW: Final[Mapping[str, float]] = MappingProxyType( { diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f1854563cd5..272f8993588 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -20,6 +20,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import pytest_asyncio +import respx from mcp.types import AudioContent, CallToolResult, ImageContent, TextContent from openai import AsyncOpenAI from openai._legacy_response import HttpxBinaryResponseContent @@ -12589,6 +12590,26 @@ def test_get_model_cost_information(): assert all(field in result for field in StandardLoggingModelInformation.__annotations__) +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +def test_ollama_cost_information_does_not_request_live_metadata( + local_model_cost_map: None, respx_mock: respx.MockRouter, provider: str +) -> None: + metadata: Final = respx_mock.post("https://cost-information.example/api/show").respond( + json={"model_info": {"test.context_length": 32768}} + ) + info: Final = StandardLoggingPayloadSetup.get_model_cost_information( + base_model="test-ollama-cost-information", + custom_pricing=False, + custom_llm_provider=provider, + init_response_obj=litellm.ModelResponse(model="test-ollama-cost-information"), + api_base="https://cost-information.example", + ) + rates: Final = info["model_map_value"] + assert rates is not None + assert (rates["input_cost_per_token"], rates["output_cost_per_token"]) == (0.0, 0.0) + assert not metadata.called + + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_model_cost_information_custom_pricing_uses_base_model(): result = StandardLoggingPayloadSetup.get_model_cost_information( diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index 041c625eaef..fcac0b0d57d 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -4,6 +4,7 @@ from typing import Final import httpx import pytest +import respx from pydantic import TypeAdapter, ValidationError import litellm @@ -12,6 +13,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( BaseDecisionsConfig, DecisionsInputLimits, + catalog_mode, ir_to_systemone_request, ir_to_systemone_response, parse_systemone_response, @@ -53,6 +55,52 @@ _SYSTEMONE_REQUEST: Final[Mapping[str, object]] = { } +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +def test_decision_capability_checks_do_not_contact_an_ambient_ollama_server( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, provider: str +) -> None: + monkeypatch.setenv("OLLAMA_API_BASE", "https://ambient.example") + monkeypatch.setenv("OLLAMA_API_KEY", "ambient-test-key") + metadata: Final = respx_mock.post("https://ambient.example/api/show").respond( + json={"model_info": {"test.context_length": 32768}} + ) + config: Final = BaseDecisionsConfig() + + assert catalog_mode("unmapped-decisions-test", provider) is None + assert ( + config.supports_box_questions("unmapped-decisions-test", provider, {}, None), + config.supports_audio_input("unmapped-decisions-test", provider, None), + config.supports_video_input("unmapped-decisions-test", provider, None), + ) == (False, False, False) + assert not metadata.called + + +@pytest.mark.parametrize("mode", ("evaluation", "decisions", "chat")) +@pytest.mark.parametrize("supported", (False, True)) +def test_decision_capabilities_preserve_registered_modes_and_flags( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, mode: str, supported: bool +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "bespoke/test-decision-capabilities", + { + "litellm_provider": "bespoke", + "mode": mode, + "supports_box_questions": supported, + "supports_audio_input": supported, + "supports_video_input": supported, + }, + ) + config: Final = BaseDecisionsConfig() + expected: Final = supported and mode in ("evaluation", "decisions") + + assert ( + config.supports_box_questions("test-decision-capabilities", "bespoke", {}, None), + config.supports_audio_input("test-decision-capabilities", "bespoke", None), + config.supports_video_input("test-decision-capabilities", "bespoke", None), + ) == (expected, expected, expected) + + @pytest.mark.parametrize("declared", (4096, 0, -1, True, None)) def test_decision_input_limits_use_only_positive_catalog_limits( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, declared: int | None diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5cf872f4ea7..e0b378fc5e9 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -28,6 +28,8 @@ import respx from jsonschema import Draft7Validator, validate import litellm + + from litellm.exceptions import MidStreamFallbackError from litellm._internal_context import is_internal_call from litellm._logging import ( @@ -11415,3 +11417,25 @@ async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronous assert GLOBAL_LOGGING_WORKER.pending_count > 0, "Tracked enqueue must finish before worker draining starts" await GLOBAL_LOGGING_WORKER.flush() assert logging_obj.model_call_details["has_logged_async_success"] is True + + +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +def test_ollama_pricing_lookup_preserves_live_model_info_discovery( + local_model_cost_map: None, respx_mock: respx.MockRouter, provider: str +) -> None: + from litellm.utils import get_model_info_for_pricing + + metadata: Final = respx_mock.post("https://model-info.example/api/show").respond( + json={"model_info": {"test.context_length": 32768}} + ) + pricing: Final = get_model_info_for_pricing( + model="test-live-discovery", custom_llm_provider=provider, api_base="https://model-info.example" + ) + assert (pricing["input_cost_per_token"], pricing["output_cost_per_token"]) == (0.0, 0.0) + assert not metadata.called + + discovered: Final = litellm.get_model_info( + model="test-live-discovery", custom_llm_provider=provider, api_base="https://model-info.example" + ) + assert discovered["max_input_tokens"] == 32768 + assert metadata.call_count == 1 From a63b543044d622d72eecd5a8e5a2b037f5acd5e4 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 2 Nov 2026 13:14:38 -0800 Subject: [PATCH 3/4] fix(ollama): include provider declaration helper for the RC --- litellm/litellm_core_utils/get_llm_provider_logic.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 38242291b71..94ba02bfbd7 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -146,6 +146,10 @@ def handle_anthropic_text_model_custom_llm_provider( return model, custom_llm_provider +def declared_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: + return custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) + + def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: """The authenticating provider this pair already names, or None. @@ -154,7 +158,7 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str and for a declared pair the resolver's answer is the declaration itself, so metadata callers adopt the declaration instead of resolving. """ - declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) + declared: Final = declared_provider(model, custom_llm_provider) return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None From 0c466b05ce26f6526dc49bb0c3598ebcda9b8685 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 2 Nov 2026 13:29:56 -0800 Subject: [PATCH 4/4] test: use explicit prices in cost calculation checks --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 191 ++++++++---------- tests/unit/llms/reducto/test_cost.py | 98 ++++----- 2 files changed, 131 insertions(+), 158 deletions(-) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 21b527250c4..4540b82a504 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -324,18 +324,15 @@ def test_reasoning_tokens_gemini(_local_model_cost_map): ) -def test_image_tokens_with_custom_pricing(): +def test_image_tokens_with_custom_pricing() -> None: """Test that image_tokens in completion are properly costed with output_cost_per_image_token.""" - from unittest.mock import patch - - # Mock model info with image token pricing - mock_model_info = { + model_info: Final = { "input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6, "output_cost_per_image_token": 5e-6, # Custom pricing for image tokens in output } - usage = Usage( + usage: Final = Usage( completion_tokens=1720, # text_tokens (600) + image_tokens (1120) prompt_tokens=14, total_tokens=1734, @@ -352,36 +349,32 @@ def test_image_tokens_with_custom_pricing(): ), ) - with patch( - "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", - return_value=mock_model_info, - ): - prompt_cost, completion_cost = generic_cost_per_token( - model="test-model", usage=usage, custom_llm_provider="gemini" - ) + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", + usage=usage, + custom_llm_provider="gemini", + model_info=cast(ModelInfo, model_info), + ) # Expected costs: # Prompt: 14 * 1e-6 # Completion: (600 * 2e-6) + (1120 * 5e-6) - expected_prompt_cost = 14 * 1e-6 - expected_completion_cost = (600 * 2e-6) + (1120 * 5e-6) + expected_prompt_cost: Final = 14 * 1e-6 + expected_completion_cost: Final = (600 * 2e-6) + (1120 * 5e-6) assert round(prompt_cost, 12) == round(expected_prompt_cost, 12) assert round(completion_cost, 12) == round(expected_completion_cost, 12) -def test_image_tokens_fallback_to_base_cost(): +def test_image_tokens_fallback_to_base_cost() -> None: """Test that image_tokens fall back to base cost when output_cost_per_image_token is not set.""" - from unittest.mock import patch - - # Mock model info without image token pricing - mock_model_info = { + model_info: Final = { "input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6, # No output_cost_per_image_token defined } - usage = Usage( + usage: Final = Usage( completion_tokens=1720, prompt_tokens=14, total_tokens=1734, @@ -398,19 +391,18 @@ def test_image_tokens_fallback_to_base_cost(): ), ) - with patch( - "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", - return_value=mock_model_info, - ): - prompt_cost, completion_cost = generic_cost_per_token( - model="test-model", usage=usage, custom_llm_provider="gemini" - ) + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", + usage=usage, + custom_llm_provider="gemini", + model_info=cast(ModelInfo, model_info), + ) # Expected costs: # Prompt: 14 * 1e-6 # Completion: (600 * 2e-6) + (1120 * 2e-6) # image_tokens use base cost - expected_prompt_cost = 14 * 1e-6 - expected_completion_cost = (600 * 2e-6) + (1120 * 2e-6) + expected_prompt_cost: Final = 14 * 1e-6 + expected_completion_cost: Final = (600 * 2e-6) + (1120 * 2e-6) assert round(prompt_cost, 12) == round(expected_prompt_cost, 12) assert round(completion_cost, 12) == round(expected_completion_cost, 12) @@ -441,16 +433,14 @@ def test_video_input_tokens_gemini_omni_flash_preview(_local_model_cost_map): ) -def test_video_tokens_fallback_to_base_cost(): +def test_video_tokens_fallback_to_base_cost() -> None: """Video output tokens fall back to the base output rate when output_cost_per_video_token is not set.""" - from unittest.mock import patch - - mock_model_info = { + model_info: Final = { "input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6, } - usage = Usage( + usage: Final = Usage( completion_tokens=1720, prompt_tokens=14, total_tokens=1734, @@ -461,30 +451,27 @@ def test_video_tokens_fallback_to_base_cost(): prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=14), ) - with patch( - "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", - return_value=mock_model_info, - ): - prompt_cost, completion_cost = generic_cost_per_token( - model="test-model", usage=usage, custom_llm_provider="gemini" - ) + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", + usage=usage, + custom_llm_provider="gemini", + model_info=cast(ModelInfo, model_info), + ) assert round(prompt_cost, 12) == round(14 * 1e-6, 12) assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12) -def test_prompt_audio_and_video_tokens_without_modality_rates_bill_at_tiered_input_rate(): +def test_prompt_audio_and_video_tokens_without_modality_rates_bill_at_tiered_input_rate() -> None: """Prompt audio/video tokens with no modality rate are prompt tokens, so a >200k prompt bills them above-200k.""" - from unittest.mock import patch - - mock_model_info = { + model_info: Final = { "input_cost_per_token": 2e-6, "input_cost_per_token_above_200k_tokens": 4e-6, "output_cost_per_token": 1.2e-5, "output_cost_per_token_above_200k_tokens": 1.8e-5, } - usage = Usage( + usage: Final = Usage( prompt_tokens=816032, completion_tokens=4559, total_tokens=820591, @@ -492,23 +479,20 @@ def test_prompt_audio_and_video_tokens_without_modality_rates_bill_at_tiered_inp completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=4559), ) - with patch( - "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", - return_value=mock_model_info, - ): - prompt_cost, completion_cost = generic_cost_per_token( - model="test-model", usage=usage, custom_llm_provider="vertex_ai" - ) + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", + usage=usage, + custom_llm_provider="vertex_ai", + model_info=cast(ModelInfo, model_info), + ) assert round(prompt_cost, 9) == round(816032 * 4e-6, 9) assert round(completion_cost, 9) == round(4559 * 1.8e-5, 9) -def test_prompt_modality_rates_still_win_over_tiered_input_rate(): +def test_prompt_modality_rates_still_win_over_tiered_input_rate() -> None: """An explicit input_cost_per_audio_token / input_cost_per_video_token beats the fallback, even above 200k.""" - from unittest.mock import patch - - mock_model_info = { + model_info: Final = { "input_cost_per_token": 2e-6, "input_cost_per_token_above_200k_tokens": 4e-6, "input_cost_per_audio_token": 3e-6, @@ -516,7 +500,7 @@ def test_prompt_modality_rates_still_win_over_tiered_input_rate(): "output_cost_per_token": 1.2e-5, } - usage = Usage( + usage: Final = Usage( prompt_tokens=300000, completion_tokens=10, total_tokens=300010, @@ -524,11 +508,12 @@ def test_prompt_modality_rates_still_win_over_tiered_input_rate(): completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10), ) - with patch( - "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", - return_value=mock_model_info, - ): - prompt_cost, _ = generic_cost_per_token(model="test-model", usage=usage, custom_llm_provider="vertex_ai") + prompt_cost, _ = generic_cost_per_token( + model="test-model", + usage=usage, + custom_llm_provider="vertex_ai", + model_info=cast(ModelInfo, model_info), + ) assert round(prompt_cost, 9) == round(100000 * 4e-6 + 100000 * 3e-6 + 100000 * 1e-6, 9) @@ -1748,12 +1733,9 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(_local_mo ) -def test_string_cost_values(): +def test_string_cost_values() -> None: """Test that cost values defined as strings are properly converted to floats.""" - from unittest.mock import patch - - # Mock model info with string cost values (as might be read from config.yaml) - mock_model_info = { + model_info: Final = { "input_cost_per_token": "3e-7", # String representation of scientific notation "output_cost_per_token": "6e-7", # String representation of scientific notation "input_cost_per_audio_token": "0.000001", # String representation of decimal @@ -1765,7 +1747,7 @@ def test_string_cost_values(): # Test usage with various token types # Note: prompt_tokens must equal sum of details to avoid double-counting adjustment # text_tokens(700) + audio_tokens(100) + cached_tokens(200) + cache_creation_tokens(150) = 1150 - usage = Usage( + usage: Final = Usage( prompt_tokens=1150, completion_tokens=500, total_tokens=1650, @@ -1785,18 +1767,16 @@ def test_string_cost_values(): ), ) - # Mock get_model_info to return our mock model info - with patch( - "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", - return_value=mock_model_info, - ): - prompt_cost, completion_cost = generic_cost_per_token( - model="test-model", usage=usage, custom_llm_provider="test-provider" - ) + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", + usage=usage, + custom_llm_provider="test-provider", + model_info=cast(ModelInfo, model_info), + ) # Calculate expected costs manually # Prompt cost = text_tokens * input_cost + audio_tokens * audio_cost + cached_tokens * cache_read_cost + cache_creation_tokens * cache_creation_cost - expected_prompt_cost = ( + expected_prompt_cost: Final = ( 700 * 3e-7 # text tokens + 100 * 1e-6 # audio tokens + 200 * 1.5e-8 # cached tokens @@ -1804,7 +1784,7 @@ def test_string_cost_values(): ) # Completion cost = text_tokens * output_cost + audio_tokens * audio_output_cost - expected_completion_cost = 450 * 6e-7 + 50 * 2e-6 # text tokens # audio tokens + expected_completion_cost: Final = 450 * 6e-7 + 50 * 2e-6 # text tokens # audio tokens # Assert costs are calculated correctly assert round(prompt_cost, 12) == round(expected_prompt_cost, 12) @@ -1918,19 +1898,17 @@ def test_calculate_cost_component_with_string_values(): assert cost == 0.0 -def test_string_cost_values_edge_cases(): +def test_string_cost_values_edge_cases() -> None: """Test edge cases for string cost value handling.""" - from unittest.mock import patch - # Test with mixed string and float cost values - mock_model_info = { + model_info: Final = { "input_cost_per_token": "1e-6", # String "output_cost_per_token": 2e-6, # Float "input_cost_per_audio_token": "invalid", # Invalid string "output_cost_per_audio_token": None, # None value } - usage = Usage( + usage: Final = Usage( prompt_tokens=1000, completion_tokens=500, total_tokens=1500, @@ -1946,30 +1924,26 @@ def test_string_cost_values_edge_cases(): ), ) - with patch( - "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", - return_value=mock_model_info, - ): - prompt_cost, completion_cost = generic_cost_per_token( - model="test-model", usage=usage, custom_llm_provider="test-provider" - ) + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", + usage=usage, + custom_llm_provider="test-provider", + model_info=cast(ModelInfo, model_info), + ) # Expected costs: # Prompt: 1000 * 1e-6 + 100 * 0 (invalid string becomes 0) # Completion: 500 * 2e-6 + 50 * 2e-6 (audio tokens fall back to base cost when output_cost_per_audio_token is None) - expected_prompt_cost = 1000 * 1e-6 - expected_completion_cost = 500 * 2e-6 + 50 * 2e-6 + expected_prompt_cost: Final = 1000 * 1e-6 + expected_completion_cost: Final = 500 * 2e-6 + 50 * 2e-6 assert round(prompt_cost, 12) == round(expected_prompt_cost, 12) assert round(completion_cost, 12) == round(expected_completion_cost, 12) -def test_string_cost_values_with_threshold(): +def test_string_cost_values_with_threshold() -> None: """Test that string cost values work correctly with threshold pricing.""" - from unittest.mock import patch - - # Mock model info with string cost values including threshold pricing - mock_model_info = { + model_info: Final = { "input_cost_per_token": "1e-6", # String base cost "output_cost_per_token": "2e-6", # String base cost "input_cost_per_token_above_200k_tokens": "5e-7", # String threshold cost (lower) @@ -1977,23 +1951,22 @@ def test_string_cost_values_with_threshold(): } # Test usage above threshold - usage = Usage( + usage: Final = Usage( prompt_tokens=250000, # Above 200k threshold completion_tokens=1000, total_tokens=251000, ) - with patch( - "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", - return_value=mock_model_info, - ): - prompt_cost, completion_cost = generic_cost_per_token( - model="test-model", usage=usage, custom_llm_provider="test-provider" - ) + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", + usage=usage, + custom_llm_provider="test-provider", + model_info=cast(ModelInfo, model_info), + ) # Expected costs using threshold pricing (string values converted to float) - expected_prompt_cost = 250000 * 5e-7 # threshold cost - expected_completion_cost = 1000 * 1e-6 # threshold cost + expected_prompt_cost: Final = 250000 * 5e-7 # threshold cost + expected_completion_cost: Final = 1000 * 1e-6 # threshold cost assert round(prompt_cost, 12) == round(expected_prompt_cost, 12) assert round(completion_cost, 12) == round(expected_completion_cost, 12) diff --git a/tests/unit/llms/reducto/test_cost.py b/tests/unit/llms/reducto/test_cost.py index 73340dc8729..6cffe8f04ec 100644 --- a/tests/unit/llms/reducto/test_cost.py +++ b/tests/unit/llms/reducto/test_cost.py @@ -1,26 +1,28 @@ -import litellm +from typing import Final + import pytest +import litellm + from litellm.cost_calculator import completion_cost from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo -def test_ocr_cost_prefers_credit_pricing_when_pages_processed_is_none(monkeypatch): - monkeypatch.setattr( - litellm, - "get_model_info", - lambda model, custom_llm_provider=None: {"ocr_cost_per_credit": 0.003}, +def test_ocr_cost_prefers_credit_pricing_when_pages_processed_is_none() -> None: + model: Final = "litellm-test-credit-ocr" + litellm.register_model( + {f"reducto/{model}": {"litellm_provider": "reducto", "mode": "ocr", "ocr_cost_per_credit": 0.003}} ) - response = OCRResponse( + response: Final = OCRResponse( pages=[OCRPage(index=0, markdown="credit priced")], - model="parse-v3", + model=model, usage_info=OCRUsageInfo(pages_processed=None, credits=10), ) - cost = completion_cost( + cost: Final = completion_cost( completion_response=response, - model="reducto/parse-v3", + model=f"reducto/{model}", custom_llm_provider="reducto", call_type="ocr", ) @@ -28,25 +30,28 @@ def test_ocr_cost_prefers_credit_pricing_when_pages_processed_is_none(monkeypatc assert cost == 0.03 -def test_ocr_cost_prefers_zero_credit_pricing_over_page_pricing(monkeypatch): - monkeypatch.setattr( - litellm, - "get_model_info", - lambda model, custom_llm_provider=None: { - "ocr_cost_per_credit": 0.0, - "ocr_cost_per_page": 0.5, - }, +def test_ocr_cost_prefers_zero_credit_pricing_over_page_pricing() -> None: + model: Final = "litellm-test-zero-credit-ocr" + litellm.register_model( + { + f"reducto/{model}": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.0, + "ocr_cost_per_page": 0.5, + } + } ) - response = OCRResponse( + response: Final = OCRResponse( pages=[OCRPage(index=0, markdown="free credit priced")], - model="parse-v3", + model=model, usage_info=OCRUsageInfo(pages_processed=2, credits=10), ) - cost = completion_cost( + cost: Final = completion_cost( completion_response=response, - model="reducto/parse-v3", + model=f"reducto/{model}", custom_llm_provider="reducto", call_type="ocr", ) @@ -54,22 +59,21 @@ def test_ocr_cost_prefers_zero_credit_pricing_over_page_pricing(monkeypatch): assert cost == 0.0 -def test_ocr_cost_falls_back_to_page_pricing(monkeypatch): - monkeypatch.setattr( - litellm, - "get_model_info", - lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5}, +def test_ocr_cost_falls_back_to_page_pricing() -> None: + model: Final = "litellm-test-page-ocr" + litellm.register_model( + {f"mistral/{model}": {"litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.5}} ) - response = OCRResponse( + response: Final = OCRResponse( pages=[OCRPage(index=0, markdown="page priced")], - model="mistral-ocr-latest", + model=model, usage_info=OCRUsageInfo(pages_processed=2), ) - cost = completion_cost( + cost: Final = completion_cost( completion_response=response, - model="mistral/mistral-ocr-latest", + model=f"mistral/{model}", custom_llm_provider="mistral", call_type="ocr", ) @@ -77,22 +81,19 @@ def test_ocr_cost_falls_back_to_page_pricing(monkeypatch): assert cost == 1.0 -def test_ocr_cost_returns_zero_when_no_pricing_and_no_pages(monkeypatch): - monkeypatch.setattr( - litellm, - "get_model_info", - lambda model, custom_llm_provider=None: {}, - ) +def test_ocr_cost_returns_zero_when_no_pricing_and_no_pages() -> None: + model: Final = "litellm-test-unpriced-ocr" + litellm.register_model({f"reducto/{model}": {"litellm_provider": "reducto", "mode": "ocr"}}) - response = OCRResponse( + response: Final = OCRResponse( pages=[OCRPage(index=0, markdown="unpriced")], - model="parse-v3", + model=model, usage_info=OCRUsageInfo(pages_processed=None, credits=5), ) - cost = completion_cost( + cost: Final = completion_cost( completion_response=response, - model="reducto/parse-v3", + model=f"reducto/{model}", custom_llm_provider="reducto", call_type="ocr", ) @@ -100,23 +101,22 @@ def test_ocr_cost_returns_zero_when_no_pricing_and_no_pages(monkeypatch): assert cost == 0.0 -def test_ocr_cost_raises_when_pages_processed_missing_for_page_pricing(monkeypatch): - monkeypatch.setattr( - litellm, - "get_model_info", - lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5}, +def test_ocr_cost_raises_when_pages_processed_missing_for_page_pricing() -> None: + model: Final = "litellm-test-missing-page-ocr" + litellm.register_model( + {f"mistral/{model}": {"litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.5}} ) - response = OCRResponse( + response: Final = OCRResponse( pages=[OCRPage(index=0, markdown="missing pages")], - model="mistral-ocr-latest", + model=model, usage_info=OCRUsageInfo(pages_processed=None), ) with pytest.raises(ValueError, match="OCR response pages_processed is None"): completion_cost( completion_response=response, - model="mistral/mistral-ocr-latest", + model=f"mistral/{model}", custom_llm_provider="mistral", call_type="ocr", )