scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(cost): use deployment model for router aliases

Browse source 
Signed-off-by: YiMing Zhang <zzz777110@qq.com>
This commit is contained in:
YiMing Zhang 2026-08-24 17:42:41 +08:00
parent f005afa146
commit 28267bf72b
No known key found for this signature in database
GPG key ID: 8C32C9D2C3609583
2 changed files with 53 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/litellm_core_utils/litellm_logging.py
View file

@ -1609,9 +1609,10 @@ class Logging(LiteLLMLoggingBaseClass):
cache_hit = self.model_call_details.get("cache_hit", False)
try:
model_for_cost: Final = litellm_model_name or self.get_deployment_model_for_cost()
response_cost_calculator_kwargs: Final = {
"response_object": result,
"model": litellm_model_name or self.model,
"model": model_for_cost,
"cache_hit": cache_hit,
"custom_llm_provider": self.model_call_details.get("custom_llm_provider", None),
"base_model": _get_base_model_from_metadata(model_call_details=self.model_call_details),
@ -1632,7 +1633,7 @@ class Logging(LiteLLMLoggingBaseClass):
custom_llm_provider=self.model_call_details.get("custom_llm_provider", None),
litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None),
optional_params=self.optional_params,
model=litellm_model_name or self.model,
model=model_for_cost,
),
}
except Exception as e: # error creating kwargs for cost calculation

50
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -281,6 +281,56 @@ def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata():
litellm.model_cost.pop(custom_model_id, None)
def test_response_cost_calculator_uses_deployment_model_for_router_alias():
"""A streamed router alias must not hide the provider-qualified deployment model."""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
alias_model = "vertex/claude-alias-cost-test"
deployment_model = "vertex_ai/claude-alias-cost-test"
input_cost = 2.0
output_cost = 3.0
litellm.register_model(
model_cost={
deployment_model: {
"input_cost_per_token": input_cost,
"output_cost_per_token": output_cost,
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"litellm_provider": "vertex_ai",
}
}
)
try:
logging_obj = LiteLLMLoggingObj(
model=alias_model,
messages=[{"role": "user", "content": "Hi"}],
stream=True,
call_type="completion",
start_time=time.time(),
litellm_call_id="test-router-alias-cost",
function_id="test-fn",
)
logging_obj.optional_params = {}
logging_obj.model_call_details["model"] = deployment_model
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
response_obj = ModelResponse(
id="chatcmpl-router-alias-cost",
model=alias_model,
choices=[],
usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
)
cost = logging_obj._response_cost_calculator(result=response_obj)
assert cost == pytest.approx((10 * input_cost) + (5 * output_cost))
finally:
litellm.model_cost.pop(deployment_model, None)
class TestGetRouterModelId:
"""Tests for the get_router_model_id helper method."""