scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-25 05:54:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 125

test(greenpt): cover public rerank cost path

Browse source
This commit is contained in:
Robert Keus 2026-07-27 19:12:10 +02:00
parent 49e568a978
commit 7d1ea76130
1 changed files with 28 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
tests/test_litellm/llms/greenpt/rerank/test_greenpt_rerank_transformation.py
View file

@ -1,8 +1,26 @@
import json
from pathlib import Path
import pytest
import litellm
from litellm.cost_calculator import rerank_cost
from litellm.llms.greenpt.rerank.transformation import GreenPTRerankConfig
@pytest.fixture
def local_model_cost_map():
original_model_cost = litellm.model_cost
model_cost_path = Path(__file__).resolve().parents[5] / "model_prices_and_context_window.json"
litellm.model_cost = json.loads(model_cost_path.read_text(encoding="utf-8"))
litellm.get_model_info.cache_clear()
try:
yield
finally:
litellm.model_cost = original_model_cost
litellm.get_model_info.cache_clear()
def test_greenpt_supported_params():
assert GreenPTRerankConfig().get_supported_cohere_rerank_params("green-rerank") == [
"query",
@ -76,6 +94,16 @@ def test_calculate_rerank_cost():
assert completion_cost == 0.0
def test_public_rerank_cost_uses_token_pricing(local_model_cost_map):
prompt_cost, completion_cost = rerank_cost(
model="greenpt/green-rerank",
custom_llm_provider=None,
billed_units={"total_tokens": 1000},
)
assert prompt_cost == pytest.approx(0.000136524)
assert completion_cost == 0.0
@pytest.mark.parametrize(
("billed_units", "model_info"),
[