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/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f59cd966261..edb4d56a5b7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -316,6 +316,58 @@ _DEPLOYMENT_PRICING_KEYS: Final = ( ) +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, \ @@ -604,59 +656,11 @@ class Logging(LiteLLMLoggingBaseClass): return next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None) def get_router_deployment_model_info(self) -> 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. - """ - model_id: Final = self.get_router_model_id() - 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 = self._published_model_info() - 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 + """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(), ) - 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_model_info(self) -> ModelInfo | None: - """The cost map's own entry for this deployment's model, when it resolves.""" - deployment_model: Final = self.get_deployment_model_for_cost() - 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 def update_environment_variables( self, 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