refactor: hoist shared model_info extraction to reduce duplication

Address Greptile review: deduplicate the litellm_params -> metadata ->
model_info navigation used by both per-token and per-second custom
pricing extraction blocks.
This commit is contained in:
Dor Amir 2026-03-09 08:08:16 -04:00
parent af8b79eb5b
commit 0b32789312

View file

@ -1108,47 +1108,33 @@ def completion_cost( # noqa: PLR0915
elif isinstance(cost_per_token_usage_object, dict):
service_tier = cost_per_token_usage_object.get("service_tier")
# Extract custom_cost_per_token from litellm_logging_obj when custom_pricing=True
# Extract custom pricing from litellm_logging_obj model_info when custom_pricing=True
# This enables cost calculation for custom/HuggingFace providers with per-deployment pricing
if (
custom_pricing is True
and custom_cost_per_token is None
and litellm_logging_obj is not None
):
_litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
if _litellm_params is not None:
_metadata = _litellm_params.get("metadata", {}) or {}
_model_info = _metadata.get("model_info", {}) or {}
_input_cost = _model_info.get("input_cost_per_token")
_output_cost = _model_info.get("output_cost_per_token")
if _input_cost is not None or _output_cost is not None:
custom_cost_per_token = {
"input_cost_per_token": _input_cost if _input_cost is not None else 0.0,
"output_cost_per_token": _output_cost if _output_cost is not None else 0.0,
}
# Extract custom_cost_per_second from litellm_logging_obj when custom_pricing=True
# This is independent of the per-token extraction above
if (
custom_pricing is True
and custom_cost_per_second is None
and litellm_logging_obj is not None
):
_litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
if _litellm_params is not None:
_metadata = _litellm_params.get("metadata", {}) or {}
_model_info = _metadata.get("model_info", {}) or {}
# Prefer input_cost_per_second; fall back to output_cost_per_second
# Use `is not None` guards to correctly handle explicit 0.0 costs
_input_cost_per_second = _model_info.get("input_cost_per_second")
_output_cost_per_second = _model_info.get("output_cost_per_second")
_cost_per_second = (
_input_cost_per_second
if _input_cost_per_second is not None
else _output_cost_per_second
_custom_pricing_model_info: dict = {}
if custom_pricing is True and litellm_logging_obj is not None:
_lp = getattr(litellm_logging_obj, "litellm_params", None)
if _lp is not None:
_custom_pricing_model_info = (
(_lp.get("metadata", {}) or {}).get("model_info", {}) or {}
)
if _cost_per_second is not None:
custom_cost_per_second = _cost_per_second
if custom_cost_per_token is None and _custom_pricing_model_info:
_input_cost = _custom_pricing_model_info.get("input_cost_per_token")
_output_cost = _custom_pricing_model_info.get("output_cost_per_token")
if _input_cost is not None or _output_cost is not None:
custom_cost_per_token = {
"input_cost_per_token": _input_cost if _input_cost is not None else 0.0,
"output_cost_per_token": _output_cost if _output_cost is not None else 0.0,
}
if custom_cost_per_second is None and _custom_pricing_model_info:
# Prefer input_cost_per_second; fall back to output_cost_per_second
# Use `is not None` guards to correctly handle explicit 0.0 costs
_input_cps = _custom_pricing_model_info.get("input_cost_per_second")
_output_cps = _custom_pricing_model_info.get("output_cost_per_second")
_cost_per_second = _input_cps if _input_cps is not None else _output_cps
if _cost_per_second is not None:
custom_cost_per_second = _cost_per_second
selected_model = _select_model_name_for_cost_calc(
model=model,