feat(proxy): add per-component response cost headers

- Extract input_cost, output_cost, cache_read_cost, cache_creation_cost, reasoning_cost, and tool_usage_cost from logging object cost breakdown
- Populate x-litellm-response-cost-* component headers in ProxyBaseLLMRequestProcessing.get_custom_headers
- Ensure headers are omitted when cost breakdown is absent or values are None
- Add comprehensive test suite covering component headers, math invariants, caching, reasoning, and discounts/margins
This commit is contained in:
abhinav 2026-08-14 23:43:04 +05:30
parent c9917cbf99
commit 1241bd5ce1
2 changed files with 216 additions and 11 deletions

View file

@ -930,26 +930,49 @@ def _override_openai_response_model(
def _get_cost_breakdown_from_logging_obj(
litellm_logging_obj: LiteLLMLoggingObj | None,
) -> tuple[float | None, float | None, float | None, float | None]:
"""
Extract discount and margin information from logging object's cost breakdown.
Returns:
Tuple of (original_cost, discount_amount, margin_total_amount, margin_percent)
"""
) -> tuple[
float | None,
float | None,
float | None,
float | None,
float | None,
float | None,
float | None,
float | None,
float | None,
float | None,
]:
"""Extract discount, margin, and per-component cost information from logging object's cost breakdown."""
if not litellm_logging_obj or not hasattr(litellm_logging_obj, "cost_breakdown"):
return None, None, None, None
return None, None, None, None, None, None, None, None, None, None
cost_breakdown: Final = litellm_logging_obj.cost_breakdown
if not cost_breakdown:
return None, None, None, None
return None, None, None, None, None, None, None, None, None, None
original_cost: Final = cost_breakdown.get("original_cost")
discount_amount: Final = cost_breakdown.get("discount_amount")
margin_total_amount: Final = cost_breakdown.get("margin_total_amount")
margin_percent: Final = cost_breakdown.get("margin_percent")
input_cost: Final = cost_breakdown.get("input_cost")
output_cost: Final = cost_breakdown.get("output_cost")
cache_read_cost: Final = cost_breakdown.get("cache_read_cost")
cache_creation_cost: Final = cost_breakdown.get("cache_creation_cost")
reasoning_cost: Final = cost_breakdown.get("reasoning_cost")
tool_usage_cost: Final = cost_breakdown.get("tool_usage_cost")
return original_cost, discount_amount, margin_total_amount, margin_percent
return (
original_cost,
discount_amount,
margin_total_amount,
margin_percent,
input_cost,
output_cost,
cache_read_cost,
cache_creation_cost,
reasoning_cost,
tool_usage_cost,
)
def _classifier_cost_from_request_data(request_data: Mapping[str, object] | None) -> float | None:
@ -1075,12 +1098,18 @@ class ProxyBaseLLMRequestProcessing:
exclude_values: Final = {"", None, "None"}
hidden_params = hidden_params or {}
# Extract discount and margin info from cost_breakdown if available
# Extract discount, margin, and per-component cost info from cost_breakdown if available
(
original_cost,
discount_amount,
margin_total_amount,
margin_percent,
input_cost,
output_cost,
cache_read_cost,
cache_creation_cost,
reasoning_cost,
tool_usage_cost,
) = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=litellm_logging_obj)
# Calculate updated spend for header (include current response_cost)
@ -1116,6 +1145,14 @@ class ProxyBaseLLMRequestProcessing:
str(margin_total_amount) if margin_total_amount is not None else None
),
"x-litellm-response-cost-margin-percent": (str(margin_percent) if margin_percent is not None else None),
"x-litellm-response-cost-input": (str(input_cost) if input_cost is not None else None),
"x-litellm-response-cost-output": (str(output_cost) if output_cost is not None else None),
"x-litellm-response-cost-cache-read": (str(cache_read_cost) if cache_read_cost is not None else None),
"x-litellm-response-cost-cache-creation": (
str(cache_creation_cost) if cache_creation_cost is not None else None
),
"x-litellm-response-cost-reasoning": (str(reasoning_cost) if reasoning_cost is not None else None),
"x-litellm-response-cost-tool-usage": (str(tool_usage_cost) if tool_usage_cost is not None else None),
"x-litellm-classifier-cost": (str(classifier_cost) if classifier_cost is not None else None),
"x-litellm-key-tpm-limit": str(user_api_key_dict.tpm_limit),
"x-litellm-key-rpm-limit": str(user_api_key_dict.rpm_limit),

View file

@ -834,6 +834,174 @@ class TestProxyBaseLLMRequestProcessing:
assert "x-litellm-response-cost-margin-amount" not in headers
assert "x-litellm-response-cost-margin-percent" not in headers
def test_get_custom_headers_per_component_cost_breakdown(self):
"""Test that per-component cost headers are included when component breakdown is available."""
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict.tpm_limit = None
mock_user_api_key_dict.rpm_limit = None
mock_user_api_key_dict.max_budget = None
mock_user_api_key_dict.spend = 0
logging_obj = LiteLLMLoggingObj(
model="gpt-5.4-nano",
messages=[{"role": "user", "content": "hello"}],
stream=False,
call_type="completion",
start_time=None,
litellm_call_id="test-call-id-components",
function_id="test-function",
)
input_cost: Final = 0.00002
output_cost: Final = 0.00004
cache_read_cost: Final = 0.000005
cache_creation_cost: Final = 0.00001
reasoning_cost: Final = 0.000015
tool_usage_cost: Final = 0.00003
total_cost: Final = (
input_cost + cache_read_cost + cache_creation_cost + output_cost + tool_usage_cost
)
logging_obj.set_cost_breakdown(
input_cost=input_cost,
output_cost=output_cost,
total_cost=total_cost,
cost_for_built_in_tools_cost_usd_dollar=tool_usage_cost,
cache_read_cost=cache_read_cost,
cache_creation_cost=cache_creation_cost,
reasoning_cost=reasoning_cost,
)
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
call_id="test-call-id-components",
response_cost=total_cost,
litellm_logging_obj=logging_obj,
)
assert "x-litellm-response-cost" in headers
assert float(headers["x-litellm-response-cost"]) == pytest.approx(total_cost)
assert "x-litellm-response-cost-input" in headers
assert float(headers["x-litellm-response-cost-input"]) == pytest.approx(input_cost)
assert "x-litellm-response-cost-output" in headers
assert float(headers["x-litellm-response-cost-output"]) == pytest.approx(output_cost)
assert "x-litellm-response-cost-cache-read" in headers
assert float(headers["x-litellm-response-cost-cache-read"]) == pytest.approx(cache_read_cost)
assert "x-litellm-response-cost-cache-creation" in headers
assert float(headers["x-litellm-response-cost-cache-creation"]) == pytest.approx(cache_creation_cost)
assert "x-litellm-response-cost-reasoning" in headers
assert float(headers["x-litellm-response-cost-reasoning"]) == pytest.approx(reasoning_cost)
assert "x-litellm-response-cost-tool-usage" in headers
assert float(headers["x-litellm-response-cost-tool-usage"]) == pytest.approx(tool_usage_cost)
component_sum: Final = (
float(headers["x-litellm-response-cost-input"])
+ float(headers["x-litellm-response-cost-cache-read"])
+ float(headers["x-litellm-response-cost-cache-creation"])
+ float(headers["x-litellm-response-cost-output"])
+ float(headers["x-litellm-response-cost-tool-usage"])
)
assert component_sum == pytest.approx(float(headers["x-litellm-response-cost"]))
assert float(headers["x-litellm-response-cost-reasoning"]) <= float(headers["x-litellm-response-cost-output"])
def test_get_custom_headers_without_cost_breakdown_omits_component_headers(self):
"""Test that when litellm_logging_obj has no cost_breakdown, component headers are omitted."""
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict.tpm_limit = None
mock_user_api_key_dict.rpm_limit = None
mock_user_api_key_dict.max_budget = None
mock_user_api_key_dict.spend = 0
logging_obj = LiteLLMLoggingObj(
model="gpt-4",
messages=[],
stream=False,
call_type="completion",
start_time=None,
litellm_call_id="test-call-id-no-breakdown",
function_id="test-function",
)
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
response_cost=0.0001,
litellm_logging_obj=logging_obj,
)
assert "x-litellm-response-cost" in headers
assert "x-litellm-response-cost-input" not in headers
assert "x-litellm-response-cost-output" not in headers
assert "x-litellm-response-cost-cache-read" not in headers
assert "x-litellm-response-cost-cache-creation" not in headers
assert "x-litellm-response-cost-reasoning" not in headers
assert "x-litellm-response-cost-tool-usage" not in headers
def test_get_custom_headers_per_component_with_discount_and_margin(self):
"""Test that component headers co-exist accurately with discount and margin headers."""
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict.tpm_limit = None
mock_user_api_key_dict.rpm_limit = None
mock_user_api_key_dict.max_budget = None
mock_user_api_key_dict.spend = 0
logging_obj = LiteLLMLoggingObj(
model="gpt-4",
messages=[],
stream=False,
call_type="completion",
start_time=None,
litellm_call_id="test-call-id-combined",
function_id="test-function",
)
logging_obj.set_cost_breakdown(
input_cost=0.00006,
output_cost=0.00004,
total_cost=0.000105,
cost_for_built_in_tools_cost_usd_dollar=0.0,
original_cost=0.0001,
discount_percent=0.05,
discount_amount=0.000005,
margin_percent=0.10,
margin_total_amount=0.00001,
)
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
response_cost=0.000105,
litellm_logging_obj=logging_obj,
)
assert float(headers["x-litellm-response-cost"]) == pytest.approx(0.000105)
assert float(headers["x-litellm-response-cost-original"]) == pytest.approx(0.0001)
assert float(headers["x-litellm-response-cost-discount-amount"]) == pytest.approx(0.000005)
assert float(headers["x-litellm-response-cost-margin-amount"]) == pytest.approx(0.00001)
assert float(headers["x-litellm-response-cost-margin-percent"]) == pytest.approx(0.10)
assert float(headers["x-litellm-response-cost-input"]) == pytest.approx(0.00006)
assert float(headers["x-litellm-response-cost-output"]) == pytest.approx(0.00004)
assert "x-litellm-response-cost-cache-read" not in headers
assert "x-litellm-response-cost-cache-creation" not in headers
assert "x-litellm-response-cost-reasoning" not in headers
assert float(headers["x-litellm-response-cost-tool-usage"]) == pytest.approx(0.0)
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
def test_get_custom_headers_classifier_cost_from_routing_decision(self, metadata_key):
"""The auto-router's LLM classifier cost must surface as its own header.