mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
222 lines
11 KiB
Python
222 lines
11 KiB
Python
from collections.abc import Iterator, Mapping
|
|
from typing import Final
|
|
from pathlib import Path
|
|
import json
|
|
import uuid
|
|
|
|
import pytest
|
|
import yaml
|
|
from pydantic import JsonValue
|
|
|
|
from tests.integration._support.client import Gateway, eventually, object_value, string_value
|
|
from tests.integration._support.database import read_rows
|
|
|
|
|
|
@pytest.mark.covers("quota_management.spend_tracking.custom_price.matches_input_rates")
|
|
def test_custom_price_is_reported_and_charged(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
response: Final = gateway.request(
|
|
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "price control"}]}
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(20 * 0.001 + 20 * 0.002)
|
|
entries: Final = gateway.get("/model/info")["data"]
|
|
assert isinstance(entries, list)
|
|
matching: Final = tuple(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model)
|
|
assert len(matching) == 1
|
|
params: Final = object_value(matching[0]["litellm_params"])
|
|
assert params["input_cost_per_token"] == 0.001
|
|
assert params["output_cost_per_token"] == 0.002
|
|
|
|
|
|
@pytest.mark.covers("quota_management.cost_estimate.configured_price.reported_for_model_absent_from_cost_map")
|
|
def test_cost_estimate_reports_configured_prices_for_model_absent_from_cost_map(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model(
|
|
model=f"openai/integration-on-prem-{uuid.uuid4().hex}",
|
|
input_cost_per_token=0.003,
|
|
output_cost_per_token=0.007,
|
|
)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/cost/estimate",
|
|
{"model": model, "input_tokens": 1000, "output_tokens": 500, "num_requests_per_day": 10},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
body: Final = object_value(response.json())
|
|
assert body["input_cost_per_token"] == pytest.approx(0.003), response.text
|
|
assert body["output_cost_per_token"] == pytest.approx(0.007), response.text
|
|
assert body["input_cost_per_request"] == pytest.approx(1000 * 0.003), response.text
|
|
assert body["output_cost_per_request"] == pytest.approx(500 * 0.007), response.text
|
|
margin: Final = body["margin_cost_per_request"]
|
|
assert isinstance(margin, float), response.text
|
|
assert body["cost_per_request"] == pytest.approx(1000 * 0.003 + 500 * 0.007 + margin), response.text
|
|
assert body["daily_cost"] == pytest.approx(10 * (1000 * 0.003 + 500 * 0.007 + margin)), response.text
|
|
|
|
|
|
@pytest.mark.covers("quota_management.spend_tracking.default_prices.survive_nullable_sibling_reload")
|
|
def test_default_prices_survive_nullable_sibling_and_reload(gateway: Gateway) -> None:
|
|
for registration_order in (("custom", "omitted", "nullable"), ("nullable", "omitted", "custom")):
|
|
with gateway.scenario() as scenario:
|
|
configured: Final = {
|
|
"custom": {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002},
|
|
"omitted": {},
|
|
"nullable": {"input_cost_per_token": None, "output_cost_per_token": None},
|
|
}
|
|
rates: Final = {
|
|
"custom": (0.001, 0.002),
|
|
"omitted": (0.00000015, 0.0000006),
|
|
"nullable": (0.00000015, 0.0000006),
|
|
}
|
|
models: Final = {kind: scenario.model(**configured[kind]) for kind in registration_order}
|
|
|
|
def observe_requests(
|
|
registration_order: tuple[str, ...],
|
|
models: Mapping[str, str],
|
|
rates: Mapping[str, tuple[float, float]],
|
|
) -> Iterator[tuple[str, float]]:
|
|
for generation in range(2):
|
|
entries: Final = gateway.get("/model/info")["data"]
|
|
assert isinstance(entries, list)
|
|
kinds: Final = tuple(reversed(registration_order)) if generation else registration_order
|
|
for index, kind in enumerate(kinds):
|
|
model: Final = models[kind]
|
|
target: Final = next(
|
|
object_value(entry) for entry in entries if object_value(entry)["model_name"] == model
|
|
)
|
|
info: Final = object_value(target["model_info"])
|
|
assert info["input_cost_per_token"] == rates[kind][0]
|
|
assert info["output_cost_per_token"] == rates[kind][1]
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "user", "content": f"price {generation * len(registration_order) + index}"}
|
|
],
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
expected: Final = 20 * rates[kind][0] + 20 * rates[kind][1]
|
|
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6)
|
|
request_id: Final = string_value(object_value(response.json())["id"])
|
|
yield request_id, expected
|
|
target: Final = next(
|
|
object_value(entry)
|
|
for entry in entries
|
|
if object_value(entry)["model_name"] == models["nullable"]
|
|
)
|
|
identity: Final = string_value(object_value(target["model_info"])["id"])
|
|
updated: Final = gateway.request(
|
|
"PATCH", f"/model/{identity}/update", {"model_info": {"description": "reload price contract"}}
|
|
)
|
|
assert updated.status_code == 200, updated.text
|
|
|
|
observations: Final = tuple(observe_requests(registration_order, models, rates))
|
|
for request_id, expected in observations:
|
|
rows: Final = eventually(
|
|
lambda request_id=request_id: read_rows(
|
|
'SELECT request_id, spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" '
|
|
"WHERE request_id = %s",
|
|
(request_id,),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert rows[0]["prompt_tokens"] == 20
|
|
assert rows[0]["completion_tokens"] == 20
|
|
assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6)
|
|
|
|
|
|
COST_MAP_DISPLAY_PRICING_KEYS: Final = frozenset(
|
|
{
|
|
"input_cost_per_token",
|
|
"output_cost_per_token",
|
|
"cache_read_input_token_cost",
|
|
"cache_creation_input_token_cost",
|
|
}
|
|
)
|
|
|
|
|
|
def persisted_model_info(identity: str) -> dict[str, JsonValue]:
|
|
rows: Final = read_rows('SELECT model_info FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,))
|
|
assert len(rows) == 1, f"Deployment {identity} has {len(rows)} rows"
|
|
stored: Final = rows[0]["model_info"]
|
|
return object_value(json.loads(stored) if isinstance(stored, str) else stored)
|
|
|
|
|
|
@pytest.mark.covers("pricing.model_update.echoed_cost_map_price_is_not_persisted_as_override")
|
|
def test_saving_echoed_model_info_does_not_freeze_cost_map_price_into_deployment(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
entries: Final = gateway.get("/model/info")["data"]
|
|
assert isinstance(entries, list)
|
|
target: Final = next(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model)
|
|
displayed: Final = object_value(target["model_info"])
|
|
identity: Final = string_value(displayed["id"])
|
|
assert isinstance(displayed["input_cost_per_token"], float), displayed
|
|
assert isinstance(displayed["output_cost_per_token"], float), displayed
|
|
fresh: Final = persisted_model_info(identity)
|
|
assert {key: value for key, value in fresh.items() if key in COST_MAP_DISPLAY_PRICING_KEYS} == {}, fresh
|
|
saved: Final = gateway.request(
|
|
"PATCH", f"/model/{identity}/update", {"model_info": {**displayed, "description": "echoed ui save"}}
|
|
)
|
|
assert saved.status_code == 200, saved.text
|
|
stored: Final = persisted_model_info(identity)
|
|
assert stored["description"] == "echoed ui save", stored
|
|
assert {key: value for key, value in stored.items() if key in COST_MAP_DISPLAY_PRICING_KEYS} == {}, stored
|
|
|
|
|
|
@pytest.mark.covers("quota_management.spend_tracking.default_prices.loaded_router_preserves_cached_defaults")
|
|
def test_loaded_router_preserves_cached_defaults_during_real_requests(gateway: Gateway, tmp_path: Path) -> None:
|
|
from litellm import Router
|
|
|
|
aliases: Final = tuple(f"pricing-{uuid.uuid4().hex}" for _ in range(3))
|
|
path: Final = tmp_path / "models.yaml"
|
|
path.write_text(
|
|
yaml.safe_dump(
|
|
{
|
|
"model_list": [
|
|
{
|
|
"model_name": alias,
|
|
"litellm_params": {
|
|
"model": "openai/gpt-4o-mini",
|
|
"api_key": "integration-provider-key",
|
|
"api_base": f"{gateway.upstream_url}/v1",
|
|
},
|
|
"model_info": {"id": alias, **pricing},
|
|
}
|
|
for alias, pricing in zip(
|
|
aliases,
|
|
(
|
|
{},
|
|
{"input_cost_per_token": None, "output_cost_per_token": None},
|
|
{"input_cost_per_token": 0.0, "output_cost_per_token": 0.0},
|
|
),
|
|
strict=True,
|
|
)
|
|
]
|
|
}
|
|
)
|
|
)
|
|
for reverse in (False, True):
|
|
configured: Final = yaml.safe_load(path.read_text())["model_list"]
|
|
router: Final = Router(model_list=list(reversed(configured)) if reverse else configured, num_retries=0)
|
|
try:
|
|
for alias in (*aliases, *reversed(aliases)):
|
|
result: Final = router.completion(
|
|
model=alias, messages=[{"role": "user", "content": "router price control"}]
|
|
)
|
|
assert result.usage.prompt_tokens == 20
|
|
assert result.usage.completion_tokens == 20
|
|
expected_cost: Final = 0.0 if alias == aliases[2] else 20 * 0.00000015 + 20 * 0.0000006
|
|
assert result._hidden_params["response_cost"] == pytest.approx(expected_cost, rel=1e-6)
|
|
deployment: Final = router.get_deployment(model_id=alias)
|
|
assert deployment is not None
|
|
info: Final = router.get_router_model_info(deployment=deployment, received_model_name=alias)
|
|
assert info["input_cost_per_token"] == (0.0 if alias == aliases[2] else 0.00000015)
|
|
assert info["output_cost_per_token"] == (0.0 if alias == aliases[2] else 0.0000006)
|
|
finally:
|
|
router.reset()
|