Replace hardcoded azure_model_router_flat_cost with flexible additional_costs field

- Changed CostBreakdown.azure_model_router_flat_cost to additional_costs (Dict[str, float])
- Updated cost calculator to use additional_costs dict for extensibility
- Azure Model Router flat cost now stored as {'Azure Model Router Flat Cost': 0.00014}
- Updated UI CostBreakdownViewer to render all additional_costs dynamically
- Added comprehensive test for additional_costs in cost breakdown
- This flexible approach allows any future cost components to be added without schema changes

Co-authored-by: ishaan <ishaan@berri.ai>
This commit is contained in:
Cursor Agent 2026-01-30 20:32:24 +00:00
parent 4ba7bbc145
commit 941162c342
5 changed files with 92 additions and 13 deletions

View file

@ -808,7 +808,7 @@ def _store_cost_breakdown_in_logging_obj(
completion_tokens_cost_usd_dollar: float,
cost_for_built_in_tools_cost_usd_dollar: float,
total_cost_usd_dollar: float,
azure_model_router_flat_cost: Optional[float] = None,
additional_costs: Optional[dict] = None,
original_cost: Optional[float] = None,
discount_percent: Optional[float] = None,
discount_amount: Optional[float] = None,
@ -825,7 +825,7 @@ def _store_cost_breakdown_in_logging_obj(
completion_tokens_cost_usd_dollar: Cost of completion tokens (includes reasoning if applicable)
cost_for_built_in_tools_cost_usd_dollar: Cost of built-in tools
total_cost_usd_dollar: Total cost of request
azure_model_router_flat_cost: Azure Model Router flat infrastructure cost
additional_costs: Free-form additional costs dict (e.g., {"azure_model_router_flat_cost": 0.00014})
original_cost: Cost before discount
discount_percent: Discount percentage applied (0.05 = 5%)
discount_amount: Discount amount in USD
@ -843,7 +843,7 @@ def _store_cost_breakdown_in_logging_obj(
output_cost=completion_tokens_cost_usd_dollar,
total_cost=total_cost_usd_dollar,
cost_for_built_in_tools_cost_usd_dollar=cost_for_built_in_tools_cost_usd_dollar,
azure_model_router_flat_cost=azure_model_router_flat_cost,
additional_costs=additional_costs,
original_cost=original_cost,
discount_percent=discount_percent,
discount_amount=discount_amount,
@ -1342,13 +1342,17 @@ def completion_cost( # noqa: PLR0915
response=completion_response,
)
# Get Azure Model Router flat cost if available (for azure_ai provider)
azure_model_router_flat_cost: Optional[float] = None
# Get additional costs (e.g., Azure Model Router flat cost for azure_ai provider)
additional_costs: Optional[dict] = None
if custom_llm_provider == "azure_ai":
from litellm.llms.azure_ai.cost_calculator import (
get_azure_model_router_flat_cost,
)
azure_model_router_flat_cost = get_azure_model_router_flat_cost()
azure_router_flat_cost = get_azure_model_router_flat_cost()
if azure_router_flat_cost is not None and azure_router_flat_cost > 0:
additional_costs = {
"Azure Model Router Flat Cost": azure_router_flat_cost
}
_final_cost = (
prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar
@ -1389,7 +1393,7 @@ def completion_cost( # noqa: PLR0915
completion_tokens_cost_usd_dollar=completion_tokens_cost_usd_dollar,
cost_for_built_in_tools_cost_usd_dollar=cost_for_built_in_tools,
total_cost_usd_dollar=_final_cost,
azure_model_router_flat_cost=azure_model_router_flat_cost,
additional_costs=additional_costs,
original_cost=original_cost,
discount_percent=discount_percent,
discount_amount=discount_amount,

View file

@ -1297,7 +1297,7 @@ class Logging(LiteLLMLoggingBaseClass):
output_cost: float,
total_cost: float,
cost_for_built_in_tools_cost_usd_dollar: float,
azure_model_router_flat_cost: Optional[float] = None,
additional_costs: Optional[dict] = None,
original_cost: Optional[float] = None,
discount_percent: Optional[float] = None,
discount_amount: Optional[float] = None,
@ -1313,7 +1313,7 @@ class Logging(LiteLLMLoggingBaseClass):
output_cost: Cost of output/completion tokens
cost_for_built_in_tools_cost_usd_dollar: Cost of built-in tools
total_cost: Total cost of request
azure_model_router_flat_cost: Azure Model Router flat infrastructure cost
additional_costs: Free-form additional costs dict (e.g., {"azure_model_router_flat_cost": 0.00014})
original_cost: Cost before discount
discount_percent: Discount percentage (0.05 = 5%)
discount_amount: Discount amount in USD
@ -1329,9 +1329,9 @@ class Logging(LiteLLMLoggingBaseClass):
tool_usage_cost=cost_for_built_in_tools_cost_usd_dollar,
)
# Store Azure Model Router flat cost if provided
if azure_model_router_flat_cost is not None and azure_model_router_flat_cost > 0:
self.cost_breakdown["azure_model_router_flat_cost"] = azure_model_router_flat_cost
# Store additional costs if provided (free-form dict for extensibility)
if additional_costs and isinstance(additional_costs, dict) and len(additional_costs) > 0:
self.cost_breakdown["additional_costs"] = additional_costs
# Store discount information if provided
if original_cost is not None:

View file

@ -2635,7 +2635,7 @@ class CostBreakdown(TypedDict, total=False):
)
total_cost: float # Total cost (input + output + tool usage)
tool_usage_cost: float # Cost of usage of built-in tools
azure_model_router_flat_cost: float # Azure AI Foundry Model Router flat cost ($0.14 per M input tokens)
additional_costs: Dict[str, float] # Free-form additional costs (e.g., {"azure_model_router_flat_cost": 0.00014})
original_cost: float # Cost before discount (optional)
discount_percent: float # Discount percentage applied (e.g., 0.05 = 5%) (optional)
discount_amount: float # Discount amount in USD (optional)

View file

@ -233,3 +233,66 @@ class TestAzureModelRouterCostBreakdown:
assert cost > expected_flat_cost
print(f"Total cost with flat fee: ${cost:.6f}")
print(f"Expected minimum flat cost: ${expected_flat_cost:.6f}")
def test_additional_costs_in_cost_breakdown(self):
"""Test that Azure Model Router flat cost appears in additional_costs dict."""
from litellm.cost_calculator import completion_cost
from litellm.litellm_core_utils.litellm_logging import LitellmLoggingObject
from litellm.types.utils import Usage, ModelResponse, Choices, Message
# Create logging object
logging_obj = LitellmLoggingObject()
# Create a mock response for azure_ai model router
response = ModelResponse(
id="test-123",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
role="assistant",
content="Test response",
),
)
],
created=1234567890,
model="azure-model-router",
object="chat.completion",
usage=Usage(
prompt_tokens=5000,
completion_tokens=2000,
total_tokens=7000,
),
)
# Set hidden params for provider
response._hidden_params = {"custom_llm_provider": "azure_ai"}
# Calculate cost with logging object
cost = completion_cost(
completion_response=response,
model="azure-model-router",
custom_llm_provider="azure_ai",
litellm_logging_obj=logging_obj,
)
# Check that cost breakdown contains additional_costs
assert hasattr(logging_obj, "cost_breakdown")
assert logging_obj.cost_breakdown is not None
assert "additional_costs" in logging_obj.cost_breakdown
assert isinstance(logging_obj.cost_breakdown["additional_costs"], dict)
# Check that the Azure Model Router flat cost is in additional_costs
additional_costs = logging_obj.cost_breakdown["additional_costs"]
assert "Azure Model Router Flat Cost" in additional_costs
# Verify the flat cost value
expected_flat_cost = (
5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000
)
actual_flat_cost = additional_costs["Azure Model Router Flat Cost"]
assert actual_flat_cost == pytest.approx(expected_flat_cost, rel=1e-9)
print(f"Additional costs in breakdown: {additional_costs}")
print(f"Azure Model Router Flat Cost: ${actual_flat_cost:.6f}")

View file

@ -7,6 +7,7 @@ export interface CostBreakdown {
output_cost?: number;
total_cost?: number;
tool_usage_cost?: number;
additional_costs?: Record<string, number>;
original_cost?: number;
discount_percent?: number;
discount_amount?: number;
@ -88,6 +89,17 @@ export const CostBreakdownViewer: React.FC<CostBreakdownViewerProps> = ({
<span className="text-gray-900">{formatCost(costBreakdown.tool_usage_cost)}</span>
</div>
)}
{/* Additional Costs (free-form) */}
{costBreakdown.additional_costs && Object.keys(costBreakdown.additional_costs).length > 0 && (
<>
{Object.entries(costBreakdown.additional_costs).map(([key, value]) => (
<div key={key} className="flex text-sm">
<span className="text-gray-600 font-medium w-1/3">{key}:</span>
<span className="text-gray-900">{formatCost(value)}</span>
</div>
))}
</>
)}
</div>
{/* Subtotal / Original Cost */}