diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index a05cbefd52e..a8e46349917 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -583,6 +583,7 @@ class CheckBatchCost: from litellm.files.main import afile_content from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, ) @@ -703,15 +704,20 @@ class CheckBatchCost: f"{_file_attr}={_raw_file_id!r}: {_e}" ) - # Pass deployment model_info so custom batch pricing - # (input_cost_per_token_batches etc.) is used for cost calc - deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {} + # Pass the deployment's router-registered pricing (litellm_params custom + # rates merged with the model's published rates) so custom batch pricing + # (input_cost_per_token_batches etc.) is used for cost calc, exactly as + # the inline retrieve path does. + deployment_model_info = deployment_pricing_model_info( + model_id=model_id, + deployment_model=litellm_model_name, + ) batch_cost, batch_usage, batch_models = ( await calculate_batch_cost_and_usage( file_content_dictionary=file_content_as_dict, custom_llm_provider=llm_provider, # type: ignore model_name=model_name, - model_info=deployment_model_info, # type: ignore[arg-type] + model_info=deployment_model_info, ) ) logging_obj = LiteLLMLogging( diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index be23972de14..c2cbb9604e5 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -48,6 +48,7 @@ async def _handle_completed_batch( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: str | None = None, litellm_params: dict | None = None, + model_info: ModelInfo | None = None, ) -> tuple[float, Usage, list[str]]: """Fetch a completed batch's output file and aggregate its cost, usage, and models in a single pass over the JSONL lines, so the parsed file content is @@ -58,6 +59,9 @@ async def _handle_completed_batch( custom_llm_provider: The LLM provider model_name: Optional model name litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) + model_info: Optional deployment-level model info with custom pricing, + threaded through so a deployment's configured rates win over the + global cost map. """ # A completed batch whose request lines all failed has no output file - the # results are written to a separate error_file_id and output_file_id is None. @@ -86,6 +90,7 @@ async def _handle_completed_batch( entries=_iter_batch_input_entries(file_content), custom_llm_provider=custom_llm_provider, model_name=model_name, + model_info=model_info, ) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b37ff865c65..8369bc3a6a2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2160,7 +2160,7 @@ def batch_cost_calculator( output_cost_per_token: Final = model_info.get("output_cost_per_token") total_prompt_cost = 0.0 total_completion_cost = 0.0 - if input_cost_per_token_batches: + if input_cost_per_token_batches is not None: total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches elif input_cost_per_token: details: Final = parse_prompt_tokens_details(usage) @@ -2180,7 +2180,7 @@ def batch_cost_calculator( cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") or input_cost_per_token total_prompt_cost += cache_creation_tokens * cache_creation_cost / 2 - if output_cost_per_token_batches: + if output_cost_per_token_batches is not None: total_completion_cost = usage.completion_tokens * output_cost_per_token_batches elif output_cost_per_token: total_completion_cost = ( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4a4a97b1d85..edb4d56a5b7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -108,6 +108,7 @@ from litellm.types.utils import ( LiteLLMBatch, LiteLLMLoggingBaseClass, LiteLLMRealtimeStreamLoggingObject, + ModelInfo, ModelResponse, ModelResponseStream, RawRequestTypedDict, @@ -307,6 +308,66 @@ def _get_cached_prometheus_logger(): return _PrometheusLogger +_DEPLOYMENT_PRICING_KEYS: Final = ( + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_token_batches", + "output_cost_per_token_batches", +) + + +def deployment_pricing_model_info(model_id: str | None, deployment_model: str | None) -> ModelInfo | None: + """Pricing the router registered under this deployment's model_info.id. + + Returns None when the deployment declares no pricing of its own, so the + caller falls back to the global cost map. The raw registration is what + decides that: the router registers an entry for every deployment, and + get_model_info fills absent costs with 0, so asking it directly cannot + tell "configured as free" apart from "no pricing configured". A deployment + may declare only one side of its pricing, so the side it leaves out keeps + the model's published rates instead of billing as zero. Ownership is per + token direction: declaring either rate for a direction takes that whole + direction, so a published batch rate can never displace a standard rate + the deployment configured itself. + """ + if model_id is None: + return None + registered: Final = litellm.model_cost.get(model_id) + if not isinstance(registered, dict) or not any(registered.get(key) is not None for key in _DEPLOYMENT_PRICING_KEYS): + return None + try: + merged: Final = litellm.get_model_info(model=model_id).copy() + except Exception: # noqa: BLE001 # get_model_info raises for ids it cannot resolve a provider for + return None + published: Final = _published_pricing(deployment_model) + if published is None: + return merged + declares_input: Final = ( + registered.get("input_cost_per_token") is not None or registered.get("input_cost_per_token_batches") is not None + ) + declares_output: Final = ( + registered.get("output_cost_per_token") is not None + or registered.get("output_cost_per_token_batches") is not None + ) + if not declares_input: + merged["input_cost_per_token"] = published.get("input_cost_per_token") + merged["input_cost_per_token_batches"] = published.get("input_cost_per_token_batches") + if not declares_output: + merged["output_cost_per_token"] = published.get("output_cost_per_token") + merged["output_cost_per_token_batches"] = published.get("output_cost_per_token_batches") + return merged + + +def _published_pricing(deployment_model: str | None) -> ModelInfo | None: + """The cost map's own entry for the deployment's model, when it resolves.""" + if deployment_model is None: + return None + try: + return litellm.get_model_info(model=deployment_model) + except Exception: # noqa: BLE001 # no published entry to layer the declared rates over + return None + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -579,6 +640,28 @@ class Logging(LiteLLMLoggingBaseClass): return model_id return None + def get_deployment_model_for_cost(self) -> str | None: + """The provider-qualified model to price against. + + On a batch retrieve both self.model and litellm_params["model"] can be + unset, and self.model can otherwise carry the router's model_group alias, + which no cost map resolves. model_call_details holds the deployment's own + provider-qualified model, so it is preferred. + """ + candidates: Final = ( + (self.model_call_details or {}).get("model") if hasattr(self, "model_call_details") else None, + self.litellm_params.get("model") if hasattr(self, "litellm_params") else None, + self.model, + ) + return next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None) + + def get_router_deployment_model_info(self) -> ModelInfo | None: + """See deployment_pricing_model_info; None means fall back to the global cost map.""" + return deployment_pricing_model_info( + model_id=self.get_router_model_id(), + deployment_model=self.get_deployment_model_for_cost(), + ) + def update_environment_variables( self, litellm_params: dict, @@ -2600,7 +2683,9 @@ class Logging(LiteLLMLoggingBaseClass): ) = await _handle_completed_batch( batch=result, custom_llm_provider=self.custom_llm_provider, + model_name=self.get_deployment_model_for_cost(), litellm_params=self.litellm_params, + model_info=self.get_router_deployment_model_info(), ) result._hidden_params["response_cost"] = response_cost diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 2ac15502840..1dbbbfc43a0 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -449,6 +449,108 @@ class TestCheckBatchCost: ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" assert snapshot["s3_bucket_name"] == "configured-batch-bucket" + @pytest.mark.asyncio + async def test_poller_prices_with_deployment_registered_batch_rates( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """The cost poller must price with the rates the router registered for the deployment. + + The deployment's raw model_info dict carries no litellm_params pricing, so passing + its model_dump() made the poller bill custom-rate batches at the public cost-map + price while the inline retrieve path billed the declared rate. + """ + from unittest.mock import patch + + import litellm + + deployment_id = "deploy-poller-registered-rates-1" + litellm.model_cost[deployment_id] = { + "id": deployment_id, + "input_cost_per_token_batches": 2e-06, + "output_cost_per_token_batches": 4e-06, + "litellm_provider": "bedrock", + "mode": "chat", + } + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_job = MagicMock() + mock_job.id = "job-poller-rates-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"custom_llm_provider": "bedrock", "aws_region_name": "us-east-1"} + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "bedrock" + mock_deployment.litellm_params.model = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"recordId":"req-1"}' + + decoded_id = f"llm_model_id,{deployment_id};llm_batch_id,batch-456;" + + try: + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + side_effect=[decoded_id, None], + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value=deployment_id, + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ), + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"recordId": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=(0.0052, {"prompt_tokens": 1400, "completion_tokens": 600}, ["claude-haiku-4-5"]), + ) as mock_calculate, + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("us.anthropic.claude-haiku-4-5-20251001-v1:0", "bedrock", None, None), + ), + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await check_batch_cost_instance.check_batch_cost() + finally: + litellm.model_cost.pop(deployment_id, None) + + mock_calculate.assert_awaited_once() + passed_model_info = mock_calculate.await_args.kwargs["model_info"] + assert passed_model_info is not None, "poller must pass the deployment's registered pricing" + assert passed_model_info["input_cost_per_token_batches"] == 2e-06 + assert passed_model_info["output_cost_per_token_batches"] == 4e-06 + @pytest.mark.asyncio async def test_primary_path_completion_update_includes_batch_processed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 68e40ac2c81..573882ebfca 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1323,6 +1323,94 @@ async def test_output_file_content_bedrock_reads_with_deployment_aws_credentials assert "model" not in captured +# =========================================================================== # +# _handle_completed_batch threads the deployment's model identity + pricing +# =========================================================================== # + + +def _bedrock_row(model: str, input_tokens: int, output_tokens: int) -> dict[str, object]: + return { + "modelInput": {"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]}, + "modelOutput": { + "model": model, + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, + }, + "recordId": "r", + } + + +@pytest.mark.asyncio +async def test_handle_completed_bedrock_batch_prices_from_deployment_model(monkeypatch) -> None: + """A bedrock batch must price from the deployment model, not the response model.""" + rows = [_bedrock_row("claude-sonnet-4-6", 18, 10)] * 100 + + async def fake_fetch(batch: object, custom_llm_provider: str, litellm_params: dict | None = None) -> bytes: + return _vertex_jsonl(rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + + cost, usage, _ = await bu._handle_completed_batch( + _batch("of"), + custom_llm_provider="bedrock", + model_name="bedrock/global.anthropic.claude-sonnet-4-6", + ) + + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (1800, 1000, 2800) + # 3e-06 / 1.5e-05 on-demand, halved for batch. + assert cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2) + + # The response model alone cannot price a bedrock batch: this is the $0 bug. + zero_cost, zero_usage, _ = await bu._handle_completed_batch( + _batch("of"), + custom_llm_provider="bedrock", + model_name=None, + ) + assert zero_cost == 0.0 + assert zero_usage.total_tokens == 2800 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> None: + """A deployment's configured rates must win over the global cost map.""" + rows = [_success_row(model="gemini-2.5-flash", usage=_usage(60, 75))] + + async def fake_fetch(batch: object, custom_llm_provider: str, litellm_params: dict | None = None) -> bytes: + return _vertex_jsonl(rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + + free_cost, _, _ = await bu._handle_completed_batch( + _batch("of"), + custom_llm_provider="vertex_ai", + model_name="vertex_ai/gemini-2.5-flash", + model_info={ + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "input_cost_per_token_batches": 0.0, + "output_cost_per_token_batches": 0.0, + }, + ) + assert free_cost == 0.0 + + billed_cost, _, _ = await bu._handle_completed_batch( + _batch("of"), + custom_llm_provider="vertex_ai", + model_name="vertex_ai/gemini-2.5-flash", + model_info=None, + ) + assert billed_cost > 0.0 + + # =========================================================================== # # _get_batch_job_usage_from_response_body: bedrock usage shapes # =========================================================================== # diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 28a6c8dd18d..946f19b7658 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,3 +1,4 @@ +import contextlib import os import sys import asyncio @@ -340,6 +341,292 @@ class TestGetRouterModelId: assert obj.get_router_model_id() is None +class TestGetRouterDeploymentModelInfo: + """Pricing a deployment registered under its own model_info.id.""" + + def test_returns_registered_deployment_pricing(self, logging_obj) -> None: + deployment_id = "deploy-zero-cost-1" + litellm.model_cost[deployment_id] = { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "input_cost_per_token_batches": 0.0, + "output_cost_per_token_batches": 0.0, + "litellm_provider": "vertex_ai", + "mode": "chat", + } + logging_obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}} + try: + info = logging_obj.get_router_deployment_model_info() + assert info is not None + assert info["input_cost_per_token"] == 0.0 + assert info["output_cost_per_token_batches"] == 0.0 + finally: + litellm.model_cost.pop(deployment_id, None) + + def test_returns_none_for_unregistered_deployment(self, logging_obj) -> None: + logging_obj.litellm_params = {"litellm_metadata": {"model_info": {"id": "deploy-never-registered"}}} + assert logging_obj.get_router_deployment_model_info() is None + + def test_returns_none_when_deployment_registered_without_pricing(self, logging_obj) -> None: + """The router registers an entry for EVERY deployment, priced or not. + + get_model_info fills absent costs with 0, so consulting it directly would + hand back free pricing for an ordinary deployment and bill its batches $0. + """ + deployment_id = "deploy-no-pricing-1" + litellm.register_model( + model_cost={deployment_id: {"id": deployment_id, "access_groups": ["x"]}}, + persist_across_reloads=False, + ) + logging_obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}} + try: + assert litellm.get_model_info(model=deployment_id)["input_cost_per_token"] == 0 + assert logging_obj.get_router_deployment_model_info() is None + finally: + litellm.model_cost.pop(deployment_id, None) + + def test_returns_none_without_a_deployment_id(self, logging_obj) -> None: + logging_obj.litellm_params = {"api_base": ""} + assert logging_obj.get_router_deployment_model_info() is None + + @pytest.mark.parametrize( + "declared,expected_input,expected_output", + [ + ({"input_cost_per_token": 1e-06}, 1e-06, 1.5e-05), + ({"output_cost_per_token": 5e-06}, 3e-06, 5e-06), + ({"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}, 0.0, 0.0), + ], + ids=["input-only", "output-only", "both-zero"], + ) + def test_one_sided_override_keeps_the_published_rate_for_the_other_side( + self, + declared: dict[str, float], + expected_input: float, + expected_output: float, + ) -> None: + """A deployment may configure one direction only. + + Substituting its pricing wholesale billed the direction it left unset at + zero, because get_model_info fills an absent cost with 0 and that + suppressed the global fallback. + """ + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + model = "bedrock/global.anthropic.claude-sonnet-4-6" + published = litellm.get_model_info(model=model) + assert (published["input_cost_per_token"], published["output_cost_per_token"]) == (3e-06, 1.5e-05) + + deployment_id = f"deploy-one-sided-{'-'.join(sorted(declared))}" + litellm.model_cost[deployment_id] = {"id": deployment_id, **declared} + obj = LiteLLMLoggingObj( + model=model, + messages=[], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="one-sided", + function_id="f", + ) + obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}, "model": model} + obj.model_call_details["model"] = model + try: + info = obj.get_router_deployment_model_info() + assert info is not None + assert info["input_cost_per_token"] == expected_input + assert info["output_cost_per_token"] == expected_output + finally: + litellm.model_cost.pop(deployment_id, None) + + def test_a_published_batch_rate_never_displaces_a_declared_standard_rate(self) -> None: + """Ownership is per token direction, not per field. + + Filling the batch field from the published entry let that rate win, so a + deployment configuring only its standard rate had batches billed at the + published batch price instead of half the rate it configured. + """ + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + model = "ft:gpt-3.5-turbo" + published = litellm.get_model_info(model=model) + assert published["input_cost_per_token_batches"] is not None + + deployment_id = "deploy-standard-input-only-1" + litellm.model_cost[deployment_id] = { + "id": deployment_id, + "input_cost_per_token": 1e-06, + "litellm_provider": "openai", + "mode": "chat", + } + obj = LiteLLMLoggingObj( + model=model, + messages=[], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="direction-ownership", + function_id="f", + ) + obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}, "model": model} + obj.model_call_details["model"] = model + try: + info = obj.get_router_deployment_model_info() + assert info is not None + assert info["input_cost_per_token"] == 1e-06 + assert info["input_cost_per_token_batches"] is None + assert info["output_cost_per_token"] == published["output_cost_per_token"] + assert info["output_cost_per_token_batches"] == published["output_cost_per_token_batches"] + finally: + litellm.model_cost.pop(deployment_id, None) + + def test_merging_does_not_mutate_the_cached_model_info(self) -> None: + """The published-rate merge must not write into get_model_info's lru-cached dict. + + get_model_info returns the same cached object on every call, so writing + the published rates into it poisoned every later lookup of the + deployment id for the life of the process. + """ + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + model = "bedrock/global.anthropic.claude-sonnet-4-6" + deployment_id = "deploy-cache-not-poisoned-1" + litellm.model_cost[deployment_id] = {"id": deployment_id, "input_cost_per_token": 1e-06} + obj = LiteLLMLoggingObj( + model=model, + messages=[], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="cache-not-poisoned", + function_id="f", + ) + obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}, "model": model} + obj.model_call_details["model"] = model + try: + cached_before = dict(litellm.get_model_info(model=deployment_id)) + info = obj.get_router_deployment_model_info() + assert info is not None + assert info["output_cost_per_token"] == 1.5e-05 + assert dict(litellm.get_model_info(model=deployment_id)) == cached_before + finally: + litellm.model_cost.pop(deployment_id, None) + + def test_keeps_declared_rates_when_no_model_is_resolvable(self, logging_obj) -> None: + """With no model to look a published entry up by, the declared rates stand alone.""" + deployment_id = "deploy-no-model-at-all-1" + litellm.model_cost[deployment_id] = { + "id": deployment_id, + "input_cost_per_token": 9e-06, + "output_cost_per_token": 2e-05, + "litellm_provider": "bedrock", + "mode": "chat", + } + logging_obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}} + logging_obj.model_call_details["model"] = None + logging_obj.model = None + try: + assert logging_obj.get_deployment_model_for_cost() is None + info = logging_obj.get_router_deployment_model_info() + assert info is not None + assert info["input_cost_per_token"] == 9e-06 + assert info["output_cost_per_token"] == 2e-05 + finally: + litellm.model_cost.pop(deployment_id, None) + + def test_returns_none_when_the_deployment_id_resolves_no_provider(self, logging_obj) -> None: + """A registration whose id get_model_info cannot resolve yields no pricing.""" + deployment_id = "deploy-unresolvable-provider-1" + litellm.model_cost[deployment_id] = {"id": deployment_id, "input_cost_per_token": 4e-06} + logging_obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}} + logging_obj.model_call_details["model"] = None + logging_obj.model = None + try: + with patch.object(litellm, "get_model_info", side_effect=Exception("unresolvable")): + assert logging_obj.get_router_deployment_model_info() is None + finally: + litellm.model_cost.pop(deployment_id, None) + + def test_falls_back_to_declared_rates_when_the_model_has_no_published_entry(self, logging_obj) -> None: + """With no published entry to layer under, the declared rates still apply.""" + deployment_id = "deploy-unpublished-model-1" + litellm.model_cost[deployment_id] = {"id": deployment_id, "input_cost_per_token": 7e-06} + logging_obj.litellm_params = { + "litellm_metadata": {"model_info": {"id": deployment_id}}, + "model": "not-a-real-provider/not-a-real-model-xyz", + } + logging_obj.model_call_details["model"] = "not-a-real-provider/not-a-real-model-xyz" + try: + info = logging_obj.get_router_deployment_model_info() + assert info is not None + assert info["input_cost_per_token"] == 7e-06 + finally: + litellm.model_cost.pop(deployment_id, None) + + +class TestRetrieveBatchCostPassesModelIdentity: + """Regression: retrieving a batch priced it with no model identity at all. + + _handle_completed_batch was called without model_name or model_info, so a + bedrock batch fell back to the provider's own response model (unresolvable + under custom_llm_provider="bedrock") and silently cost $0, and a deployment's + configured rates were ignored entirely. + """ + + @pytest.mark.asyncio + async def test_forwards_deployment_model_and_pricing(self, monkeypatch) -> None: + from litellm.litellm_core_utils import litellm_logging as logging_module + from litellm.types.utils import LiteLLMBatch, Usage + + deployment_id = "deploy-batch-pricing-1" + litellm.model_cost[deployment_id] = { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "bedrock", + "mode": "chat", + } + + captured: dict[str, object] = {} + + async def fake_handle_completed_batch(**kwargs: object) -> tuple[float, Usage, list[str]]: + captured.update(kwargs) + return 1.25, Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), ["m"] + + monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch) + + obj = LitellmLogging( + model="bedrock/global.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="batch-call-1", + function_id="f", + ) + obj.custom_llm_provider = "bedrock" + obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}} + + batch = LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="completed", + output_file_id="file-out", + ) + + try: + with contextlib.suppress(Exception): + await obj._async_success_handler_body(result=batch, start_time=None, end_time=None) + finally: + litellm.model_cost.pop(deployment_id, None) + + assert captured, "_handle_completed_batch was never called" + assert captured["model_name"] == "bedrock/global.anthropic.claude-sonnet-4-6" + assert captured["model_info"] is not None + assert captured["model_info"]["input_cost_per_token"] == 0.0 + + class TestAnthropicPassthroughCustomPricing: """Verify the Anthropic pass-through handler forwards custom pricing.""" diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index a51f4e733b6..75c90d793fe 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -20,7 +20,7 @@ from litellm.cost_calculator import ( response_cost_calculator, ) from litellm.types.llms.openai import OpenAIRealtimeStreamList -from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +from litellm.types.utils import ModelInfo, ModelResponse, PromptTokensDetailsWrapper, Usage from litellm.utils import TranscriptionResponse @@ -3562,16 +3562,18 @@ def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate( """ from litellm.cost_calculator import batch_cost_calculator + model_info: ModelInfo = { + "supported_openai_params": [], + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + } prompt_cost, completion_cost_value = batch_cost_calculator( usage=_batch_cache_usage(), model="claude-sonnet-4-5-20250929", custom_llm_provider="anthropic", - model_info={ # type: ignore[arg-type] - "input_cost_per_token": 3e-6, - "output_cost_per_token": 15e-6, - "cache_read_input_token_cost": 3e-7, - "cache_creation_input_token_cost": 3.75e-6, - }, + model_info=model_info, ) assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3.75e-6) / 2) @@ -3581,20 +3583,69 @@ def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate( def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate(): from litellm.cost_calculator import batch_cost_calculator + model_info: ModelInfo = { + "supported_openai_params": [], + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + } prompt_cost, _ = batch_cost_calculator( usage=_batch_cache_usage(), model="claude-sonnet-4-5-20250929", custom_llm_provider="anthropic", - model_info={ # type: ignore[arg-type] - "input_cost_per_token": 3e-6, - "output_cost_per_token": 15e-6, - "cache_read_input_token_cost": 3e-7, - }, + model_info=model_info, ) assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2) +@pytest.mark.parametrize( + "batch_rate,expected_prompt,expected_completion", + [ + (0.0, 0.0, 0.0), + (1e-6, 1000 * 1e-6, 500 * 1e-6), + (None, 1000 * 3e-6 / 2, 500 * 15e-6 / 2), + ], + ids=["explicit-zero", "explicit-nonzero", "unset"], +) +def test_batch_cost_calculator_honors_an_explicitly_zero_batch_rate( + batch_rate: float | None, + expected_prompt: float, + expected_completion: float, +) -> None: + """A batch rate configured as 0.0 means free, not unset. + + Gating the batch fields on truthiness read an explicit 0.0 as absent and + charged half the standard rate for that token direction instead. + """ + from litellm.cost_calculator import batch_cost_calculator + + base_model_info: ModelInfo = { + "supported_openai_params": [], + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + } + model_info: ModelInfo = ( + base_model_info + if batch_rate is None + else { + **base_model_info, + "input_cost_per_token_batches": batch_rate, + "output_cost_per_token_batches": batch_rate, + } + ) + + prompt_cost, completion_cost_value = batch_cost_calculator( + usage=Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), + model="claude-sonnet-4-5-20250929", + custom_llm_provider="anthropic", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(expected_prompt) + assert completion_cost_value == pytest.approx(expected_completion) + + def test_combine_usage_objects_sums_mirrored_cache_write_fields_once(): """ cache_write_tokens and cache_creation_tokens mirror each other on