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(huggingface): use typed registered-model lookup

Browse source
This commit is contained in:
er-s-an 2026-10-12 17:16:10 +08:00
parent aa0b28df25
commit eb3bbd6fa8
1 changed files with 16 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
tests/unit/test_cost_calculator.py
View file

@ -1907,7 +1907,18 @@ def test_tiered_pricing_only_deployment_completion_cost_is_nonzero():
assert cost > 0
def test_huggingface_rerank_uses_registered_input_cost_per_query():
@pytest.mark.parametrize(
("registered_model", "model"),
[
("hf-tei-rerank-deployment", "hf-tei-rerank-deployment"),
("huggingface/hf-tei-rerank-deployment", "huggingface/hf-tei-rerank-deployment"),
("hf-tei-rerank-deployment", "HF-TEI-RERANK-DEPLOYMENT"),
],
)
@pytest.mark.usefixtures("local_model_cost_map")
def test_huggingface_rerank_uses_registered_input_cost_per_query(
registered_model: str, model: str, request: pytest.FixtureRequest
):
"""Registered per-query price must apply for huggingface/TEI rerank (#45975).
``_get_model_info_helper`` used to return a synthetic zero-cost huggingface
@ -1918,11 +1929,12 @@ def test_huggingface_rerank_uses_registered_input_cost_per_query():
from litellm.utils import cached_get_model_info_helper
cached_get_model_info_helper.cache_clear()
litellm.model_cost.pop("hf-tei-rerank-deployment", None)
litellm.register_model({"hf-tei-rerank-deployment": {"input_cost_per_query": 0.001}})
request.addfinalizer(cached_get_model_info_helper.cache_clear)
litellm.model_cost.pop(registered_model, None)
litellm.register_model({registered_model: {"input_cost_per_query": 0.001}})
cost = rerank_cost(
model="hf-tei-rerank-deployment",
model=model,
custom_llm_provider="huggingface",
billed_units={"search_units": 1},
)