fix(azure_ai): keep the request_model keyword on cost_per_token

Restores the public keyword removed at 415bdbfd8f. A direct caller that
names the Model Router as the request model gets the routing fee folded
into the prompt cost once; completion_cost never passes it and charges
the fee through the additional-costs hook as before

litellm/cost_calculator.py

@@ -341,6 +341,8 @@ def cost_per_token(
     ### VERTEX LOCATION ###
     vertex_location: str | None = None,  # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global")
     response: Any | None = None,
+    ### REQUEST MODEL ###
+    request_model: str | None = None,  # original request model for router detection
 ) -> tuple[float, float]:
     """
     Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -662,6 +664,7 @@ def cost_per_token(
             model=model,
             usage=usage_block,
             response_time_ms=response_time_ms,
+            request_model=request_model,
             service_tier=service_tier,
         )
     else:

litellm/llms/azure_ai/cost_calculator.py

@@ -63,32 +63,7 @@ def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> fl
     return 0.0
 
 
-def cost_per_token(
-    model: str,
-    usage: Usage,
-    response_time_ms: float | None = 0.0,
-    service_tier: str | None = None,
-) -> tuple[float, float]:
-    """
-    Price the response model's own tokens for Azure AI.
-
-    The Azure AI Foundry Model Router fee is not part of this: completion_cost charges it once through
-    AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost
-    breakdown, and a response priced as the router entry itself already carries it. A router deployment name
-    that is missing from the cost map prices at zero here so that line item is the whole cost.
-
-    Args:
-        model: str, the model name without provider prefix (from response)
-        usage: LiteLLM Usage block
-        response_time_ms: Optional response time in milliseconds
-        service_tier: Optional service tier the request was priced on
-
-    Returns:
-        Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
-
-    Raises:
-        ValueError: If a model that is not a Model Router name is missing from the cost map
-    """
+def _response_model_cost(model: str, usage: Usage, service_tier: str | None) -> tuple[float, float]:
     try:
         return generic_cost_per_token(
             model=model, usage=usage, custom_llm_provider="azure_ai", service_tier=service_tier
@@ -100,3 +75,38 @@ def cost_per_token(
             "Azure AI Model Router: model '%s' not in cost map, only the routing fee applies. Error: %s", model, e
         )
         return 0.0, 0.0
+
+
+def cost_per_token(
+    model: str,
+    usage: Usage,
+    response_time_ms: float | None = 0.0,
+    request_model: str | None = None,
+    service_tier: str | None = None,
+) -> tuple[float, float]:
+    """
+    Price the response model's own tokens for Azure AI, plus the Model Router fee when the caller names the
+    router as the request model.
+
+    completion_cost never passes request_model: it charges the fee once through
+    AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost
+    breakdown. A response priced as the router entry itself already carries the fee, so request_model adds
+    nothing on top of it, and a router deployment name that is missing from the cost map prices at zero here.
+
+    Args:
+        model: str, the model name without provider prefix (from response)
+        usage: LiteLLM Usage block
+        response_time_ms: Optional response time in milliseconds
+        request_model: Optional[str], the original request model name; a Model Router name adds the routing fee
+        service_tier: Optional service tier the request was priced on
+
+    Returns:
+        Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
+
+    Raises:
+        ValueError: If a model that is not a Model Router name is missing from the cost map
+    """
+    prompt_cost, completion_cost = _response_model_cost(model=model, usage=usage, service_tier=service_tier)
+    if request_model is None or not _is_azure_model_router(request_model) or is_router_fee_entry(model):
+        return prompt_cost, completion_cost
+    return prompt_cost + calculate_azure_model_router_flat_cost(request_model, usage.prompt_tokens), completion_cost

tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py

@@ -155,6 +155,39 @@ class TestAzureModelRouterFlatCost:
         with pytest.raises(Exception, match="no-such-azure-ai-model"):
             cost_per_token(model="no-such-azure-ai-model", usage=usage)
 
+    def test_request_model_through_the_router_adds_the_fee_once(self) -> None:
+        routed_prompt_cost, routed_completion_cost = _routed_model_cost()
+        prompt_cost, completion_cost_usd = cost_per_token(
+            model=ROUTED_MODEL, usage=ROUTED_USAGE, request_model="azure_ai/model-router"
+        )
+        assert prompt_cost == pytest.approx(routed_prompt_cost + ROUTED_FEE, rel=1e-9)
+        assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9)
+
+    def test_request_model_that_is_not_the_router_adds_nothing(self) -> None:
+        routed_prompt_cost, routed_completion_cost = _routed_model_cost()
+        assert cost_per_token(
+            model=ROUTED_MODEL, usage=ROUTED_USAGE, request_model=f"azure_ai/{ROUTED_MODEL}"
+        ) == pytest.approx((routed_prompt_cost, routed_completion_cost), rel=1e-9)
+
+    @pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"])
+    def test_request_model_does_not_double_the_router_entry(self, router_entry_name: str) -> None:
+        prompt_cost, completion_cost_usd = cost_per_token(
+            model=router_entry_name, usage=ROUTED_USAGE, request_model=f"azure_ai/{router_entry_name}"
+        )
+        assert prompt_cost == pytest.approx(ROUTED_FEE, rel=1e-9)
+        assert completion_cost_usd == 0.0
+
+    def test_public_cost_per_token_keeps_the_request_model_keyword(self) -> None:
+        routed_prompt_cost, routed_completion_cost = _routed_model_cost()
+        prompt_cost, completion_cost_usd = litellm.cost_per_token(
+            model=ROUTED_MODEL,
+            custom_llm_provider="azure_ai",
+            usage_object=ROUTED_USAGE,
+            request_model="azure_ai/model-router",
+        )
+        assert prompt_cost == pytest.approx(routed_prompt_cost + ROUTED_FEE, rel=1e-9)
+        assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9)
+
     def test_flat_cost_helper(self) -> None:
         assert calculate_azure_model_router_flat_cost(
             model="azure-model-router", prompt_tokens=10_000