scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(batches): price poller-tracked batches from the deployment's registered rates

Browse source
This commit is contained in:
mateo-berri 2026-08-17 15:36:39 -07:00
parent 678a636846
commit 5a11fe141e
3 changed files with 168 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -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(

108
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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,

102
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -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