scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge upstream main f553c80cd3 into litellm_forward_reasoning_content

Browse source
This commit is contained in:
jibanez-staticduo 2026-10-01 03:12:06 +02:00
commit 0be9a3d43f
No known key found for this signature in database
13 changed files with 897 additions and 227 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/anthropic_beta_headers_config.json
View file

@ -49,7 +49,7 @@
"computer-use-2025-11-24": "computer-use-2025-11-24",
"context-1m-2025-08-07": "context-1m-2025-08-07",
"context-management-2025-06-27": "context-management-2025-06-27",
"dangerous-tool-use-2026-09-03": null,
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": "files-api-2025-04-14",

13
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -76,6 +76,9 @@ _SERVICE_TIER_TO_COST_KEY_SUFFIX: Final[Mapping[str, str]] = MappingProxyType(
ServiceTier.ULTRAFAST.value: ServiceTier.ULTRAFAST.value,
}
)
SERVICE_TIER_COST_KEY_SUFFIXES: Final[tuple[str, ...]] = tuple(
sorted(frozenset(f"_{suffix}" for suffix in _SERVICE_TIER_TO_COST_KEY_SUFFIX.values()))
)
_INCLUSIVE_THRESHOLD_PROVIDERS: Final = frozenset({"xai"})
_BATCH_KEY_SUFFIX: Final = "_batches"
@ -663,13 +666,15 @@ def _get_token_base_cost(
## CHECK IF ABOVE THRESHOLD
# Optimization: collect threshold keys first to avoid sorting all model_info keys.
# Exclude service_tier-specific variants (e.g. input_cost_per_token_above_200k_tokens_priority)
# so that the threshold detection loop only processes standard keys. The
# service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key.
# Standard thresholds and thresholds suffixed for this request's service tier both count.
tier_key_suffix: Final = _get_service_tier_cost_key("", service_tier)
threshold_keys: Final = [
k
for k in model_info
if k.startswith("input_cost_per_token_above_") and not k.endswith(_NON_STANDARD_THRESHOLD_SUFFIXES)
if k.startswith("input_cost_per_token_above_")
and (
not k.endswith(_NON_STANDARD_THRESHOLD_SUFFIXES) or (tier_key_suffix != "" and k.endswith(tier_key_suffix))
)
]
# Only sort the threshold keys (typically 1-2 keys instead of 66+)

46
litellm/router.py
View file

@ -87,6 +87,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import (
is_registered_custom_provider,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES
from litellm.litellm_core_utils.ptu_pricing import (
PTU_COST_ATTRIBUTION_ENV_VAR,
declares_ptu,
@ -8812,6 +8813,41 @@ class Router:
if backend_value is not None:
model_info[field] = backend_value
@staticmethod
def _cost_map_backend_model(deployment: Deployment) -> str:
model_info_base_model: Final = deployment.model_info.base_model
if isinstance(model_info_base_model, str) and model_info_base_model:
return model_info_base_model
params_base_model: Final = deployment.litellm_params.get("base_model")
if isinstance(params_base_model, str) and params_base_model:
return params_base_model
return deployment.litellm_params.model
@staticmethod
def _inherit_builtin_service_tier_pricing(
model_info: dict, # mutable-ok: deployment cost-map entry filled in place
backend_model: str,
custom_llm_provider: str | None,
) -> None:
"""Inherit missing tier rates so a standalone entry does not fall back to custom standard rates."""
if ptu_terms(model_info) is not None and is_ptu_cost_attribution_enabled():
return
if all(model_info.get(field) is None for field in ("input_cost_per_token", "output_cost_per_token")):
return
try:
backend_info: Final = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider)
except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model
return
backend_entry: Final = litellm.model_cost.get(backend_info.get("key") or "")
if not isinstance(backend_entry, dict):
return
for field, backend_value in backend_entry.items():
if not field.endswith(SERVICE_TIER_COST_KEY_SUFFIXES):
continue
if model_info.get(field) is not None or backend_value is None:
continue
model_info[field] = copy.deepcopy(backend_value)
@staticmethod
def _inherit_builtin_base_rates_for_off_peak(
model_info: dict, # mutable-ok: cost-map entry filled in place
@ -8960,6 +8996,11 @@ class Router:
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_service_tier_pricing(
model_info=_model_info,
backend_model=Router._cost_map_backend_model(deployment),
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_tiered_output_rate(
model_info=_model_info,
backend_model=deployment.litellm_params.model,
@ -10000,6 +10041,11 @@ class Router:
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_service_tier_pricing(
model_info=model_info,
backend_model=Router._cost_map_backend_model(deployment),
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_tiered_output_rate(
model_info=model_info,
backend_model=deployment.litellm_params.model,

5
proxy_server_config.yaml
View file

@ -141,6 +141,11 @@ model_list:
- model_name: mistral-embed
litellm_params:
model: mistral/mistral-embed
- model_name: gpt-6-luna
litellm_params:
model: openai/gpt-6-luna
reasoning_effort: none
api_key: os.environ/OPENAI_API_KEY
- model_name: gpt-instruct # [PROD TEST] - tests if `/health` automatically infers this to be a text completion model
litellm_params:
model: text-completion-openai/gpt-3.5-turbo-instruct

140
tests/integration/pricing/test_service_tier_pricing.py
View file

@ -1,5 +1,6 @@
import json
import uuid
from pathlib import Path
from typing import Final, Literal
import httpx
@ -260,3 +261,142 @@ def test_ultrafast_long_context_prompt_bills_ultrafast_long_context_rates(
body: Final = object_value(observation["body"])
assert body.get("service_tier") == service_tier, body
assert not set(LONG_CONTEXT_PRICING).intersection(body), body
BUNDLED_COST_MAP: Final = (
Path(__file__).resolve().parents[3] / "litellm" / "model_prices_and_context_window_backup.json"
)
CUSTOM_STANDARD_INPUT_RATE: Final = 0.001
CUSTOM_STANDARD_OUTPUT_RATE: Final = 0.002
def _bundled_rate(model: str, field: str) -> float:
rate: Final = object_value(JSON_OBJECT.validate_json(BUNDLED_COST_MAP.read_bytes())[model])[field]
assert isinstance(rate, float) and rate > 0, f"{model}.{field} in {BUNDLED_COST_MAP.name}: {rate}"
return rate
@pytest.mark.parametrize(
("service_tier", "input_field", "output_field"),
(
("ultrafast", "input_cost_per_token_ultrafast", "output_cost_per_token_ultrafast"),
(None, None, None),
),
ids=("ultrafast", "standard"),
)
def test_custom_standard_rates_bill_served_ultrafast_tier_at_the_catalog_tier_rate(
gateway: Gateway, service_tier: str | None, input_field: str | None, output_field: str | None
) -> None:
input_rate: Final = CUSTOM_STANDARD_INPUT_RATE if input_field is None else _bundled_rate("gpt-6-astra", input_field)
output_rate: Final = (
CUSTOM_STANDARD_OUTPUT_RATE if output_field is None else _bundled_rate("gpt-6-astra", output_field)
)
with gateway.scenario() as scenario:
scenario_id: Final = f"custom-standard-ultrafast-{uuid.uuid4().hex}"
handle: Final = register_scenario(
scenario_id,
JsonResponse(
content_type="application/json",
body={
"id": "chatcmpl-$UNIQUE_ID",
"object": "chat.completion",
"created": 1,
"model": "gpt-6-astra",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100},
**({} if service_tier is None else {"service_tier": service_tier}),
},
),
)
scenario.cleanups.callback(delete_scenario, handle)
model: Final = scenario.model(
model="openai/gpt-6-astra",
api_key=scenario_id,
api_base=handle.api_base(),
input_cost_per_token=CUSTOM_STANDARD_INPUT_RATE,
output_cost_per_token=CUSTOM_STANDARD_OUTPUT_RATE,
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "OK"}],
**({} if service_tier is None else {"service_tier": service_tier}),
},
key=scenario.key(),
)
assert response.status_code == 200, response.text
expected: Final = 1000 * input_rate + 100 * output_rate
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6), response.text
request_id: Final = string_value(object_value(response.json())["id"])
rows: Final = eventually(
lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)),
lambda values: len(values) == 1,
seconds=70,
)
assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6), rows
def test_custom_standard_rates_bill_catalog_ultrafast_long_context_rates(gateway: Gateway) -> None:
input_rate: Final = _bundled_rate("gpt-6-astra", "input_cost_per_token_above_272k_tokens_ultrafast")
output_rate: Final = _bundled_rate("gpt-6-astra", "output_cost_per_token_above_272k_tokens_ultrafast")
with gateway.scenario() as scenario:
scenario_id: Final = f"custom-standard-ultrafast-long-context-{uuid.uuid4().hex}"
handle: Final = register_scenario(
scenario_id,
JsonResponse(
content_type="application/json",
body={
"id": "chatcmpl-$UNIQUE_ID",
"object": "chat.completion",
"created": 1,
"model": "gpt-6-astra",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}
],
"usage": {
"prompt_tokens": LONG_PROMPT_TOKENS,
"completion_tokens": 100,
"total_tokens": LONG_PROMPT_TOKENS + 100,
},
"service_tier": "ultrafast",
},
),
)
scenario.cleanups.callback(delete_scenario, handle)
model: Final = scenario.model(
model="openai/gpt-6-astra",
api_key=scenario_id,
api_base=handle.api_base(),
input_cost_per_token=CUSTOM_STANDARD_INPUT_RATE,
output_cost_per_token=CUSTOM_STANDARD_OUTPUT_RATE,
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "long context ultrafast pricing"}],
"service_tier": "ultrafast",
},
key=scenario.key(),
)
assert response.status_code == 200, response.text
expected: Final = LONG_PROMPT_TOKENS * input_rate + 100 * output_rate
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6), response.text
request_id: Final = string_value(object_value(response.json())["id"])
rows: Final = eventually(
lambda: read_rows(
'SELECT 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"] == LONG_PROMPT_TOKENS
assert rows[0]["completion_tokens"] == 100
assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6), rows

95
tests/local_testing/test_streaming.py
View file

@ -2,6 +2,7 @@
# This tests streaming for the completion endpoint
import asyncio
from typing import Final
import json
import os
import time
@ -1546,45 +1547,24 @@ async def test_openai_stream_options_call(model, sync):
)
def test_openai_stream_options_call_text_completion():
litellm.set_verbose = False
for idx in range(3):
try:
response = litellm.text_completion(
model="gpt-3.5-turbo-instruct",
prompt="say GM - we're going to make it ",
stream=True,
stream_options={"include_usage": True},
max_tokens=10,
)
usage = None
chunks = []
for chunk in response:
print("chunk: ", chunk)
chunks.append(chunk)
last_chunk = chunks[-1]
print("last chunk: ", last_chunk)
"""
Assert that:
- Last Chunk includes Usage
- All chunks prior to last chunk have usage=None
"""
assert last_chunk.usage is not None
assert last_chunk.usage.total_tokens > 0
assert last_chunk.usage.prompt_tokens > 0
assert last_chunk.usage.completion_tokens > 0
# assert all non last chunks have usage=None
assert all(chunk.usage is None for chunk in chunks[:-1])
break
except Exception as e:
if idx < 2:
pass
else:
raise e
def test_openai_stream_options_call_text_completion() -> None:
chunks: Final = tuple(
litellm.text_completion(
model="gpt-6-luna",
reasoning_effort="none",
prompt="say GM - we're going to make it ",
stream=True,
stream_options={"include_usage": True},
max_tokens=10,
)
)
assert chunks
assert chunks[-1].usage is not None
assert chunks[-1].usage.total_tokens > 0
assert chunks[-1].usage.prompt_tokens > 0
assert chunks[-1].usage.completion_tokens > 0
assert all(chunk.usage is None for chunk in chunks[:-1])
assert any(chunk.choices[0].text for chunk in chunks)
def test_openai_text_completion_call():
@ -1676,8 +1656,8 @@ def test_together_ai_completion_call_starcoder_bad_key():
#### Test Function calling + streaming ####
def test_completion_openai_with_functions():
function1 = [
def test_completion_openai_with_functions() -> None:
functions: Final = [
{
"name": "get_current_weather",
"description": "Get the current weather in a given location",
@ -1694,24 +1674,25 @@ def test_completion_openai_with_functions():
},
}
]
try:
litellm.set_verbose = False
response = completion(
model="gpt-3.5-turbo-1106",
messages=[{"role": "user", "content": "what's the weather in SF"}],
functions=function1,
messages: Final = [{"role": "user", "content": "what's the weather in SF"}]
chunks: Final = tuple(
completion(
model="gpt-6-luna",
reasoning_effort="none",
messages=messages,
functions=functions,
function_call={"name": "get_current_weather"},
stream=True,
max_tokens=128,
)
# Add any assertions here to check the response
print(response)
for chunk in response:
print(chunk)
if chunk["choices"][0]["finish_reason"] == "stop":
break
print(chunk["choices"][0]["finish_reason"])
print(chunk["choices"][0]["delta"]["content"])
except Exception as e:
pytest.fail(f"Error occurred: {e}")
)
response: Final = litellm.stream_chunk_builder(chunks, messages=messages)
assert response is not None
function_call: Final = response.choices[0].message.function_call
assert function_call is not None
assert function_call.name == "get_current_weather"
assert json.loads(function_call.arguments)["location"]
assert sum(chunk.choices[0].finish_reason is not None for chunk in chunks) == 1
#### Test Async streaming ####

91
tests/local_testing/test_text_completion.py
View file

@ -1,4 +1,5 @@
import asyncio
from typing import Final
import json
import traceback
@ -3790,42 +3791,30 @@ def test_completion_openai_prompt():
# test_completion_openai_prompt()
def test_completion_openai_engine_and_model():
try:
print("\n text 003 test\n")
litellm.set_verbose = True
response = text_completion(
model="gpt-3.5-turbo-instruct",
engine="anything",
prompt="What's the weather in SF?",
max_tokens=5,
)
print(response)
response_str = response["choices"][0]["text"]
# print(response.choices[0])
# print(response.choices[0].text)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_openai_engine_and_model() -> None:
response: Final = text_completion(
model="gpt-6-luna",
engine="anything",
reasoning_effort="none",
prompt="What's the weather in SF?",
max_tokens=5,
)
assert response.model == "gpt-6-luna"
assert response.choices[0].text
# test_completion_openai_engine_and_model()
def test_completion_openai_engine():
try:
print("\n text 003 test\n")
litellm.set_verbose = True
response = text_completion(
engine="gpt-3.5-turbo-instruct",
prompt="What's the weather in SF?",
max_tokens=5,
)
print(response)
response_str = response["choices"][0]["text"]
# print(response.choices[0])
# print(response.choices[0].text)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_openai_engine() -> None:
response: Final = text_completion(
engine="gpt-6-luna",
reasoning_effort="none",
prompt="What's the weather in SF?",
max_tokens=5,
)
assert response.model == "gpt-6-luna"
assert response.choices[0].text
# test_completion_openai_engine()
@ -4048,34 +4037,18 @@ def test_async_text_completion_together_ai():
# test_async_text_completion()
def test_async_text_completion_stream():
# tests atext_completion + streaming - assert only one finish reason sent
litellm.set_verbose = False
print("test_async_text_completion with stream")
async def test_get_response():
try:
response = await litellm.atext_completion(
model="gpt-3.5-turbo-instruct",
prompt="good morning",
stream=True,
)
print(f"response: {response}")
num_finish_reason = 0
async for chunk in response:
print(chunk)
if chunk["choices"][0].get("finish_reason") is not None:
num_finish_reason += 1
print("finish_reason", chunk["choices"][0].get("finish_reason"))
assert (
num_finish_reason == 1
), f"expected only one finish reason. Got {num_finish_reason}"
except Exception as e:
pytest.fail(f"GOT exception for gpt-3.5 instruct In streaming{e}")
asyncio.run(test_get_response())
@pytest.mark.asyncio
async def test_async_text_completion_stream() -> None:
response: Final = await litellm.atext_completion(
model="gpt-6-luna",
reasoning_effort="none",
prompt="good morning",
stream=True,
max_tokens=32,
)
chunks: Final = [chunk async for chunk in response]
assert sum(chunk.choices[0].finish_reason is not None for chunk in chunks) == 1
assert any(chunk.choices[0].text for chunk in chunks)
# test_async_text_completion_stream()

165
tests/test_fallbacks.py
View file

@ -1,3 +1,6 @@
import os
from typing import Final
# What is this?
## This tests if the proxy fallbacks work as expected
import pytest
@ -6,6 +9,9 @@ import aiohttp
from tests.large_text import text
import time
from typing import Optional
from openai import AsyncOpenAI, PermissionDeniedError
PROXY_BASE_URL: Final = os.environ.get("LITELLM_PROXY_BASE_URL", "http://0.0.0.0:4000")
async def generate_key(
@ -14,7 +20,7 @@ async def generate_key(
models: list,
calling_key="sk-1234",
):
url = "http://0.0.0.0:4000/key/generate"
url: Final = f"{PROXY_BASE_URL}/key/generate"
headers = {
"Authorization": f"Bearer {calling_key}",
"Content-Type": "application/json",
@ -48,7 +54,7 @@ async def chat_completion(
extra_headers: Optional[dict] = None,
**kwargs,
):
url = "http://0.0.0.0:4000/chat/completions"
url: Final = f"{PROXY_BASE_URL}/chat/completions"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
@ -94,42 +100,30 @@ async def test_chat_completion():
@pytest.mark.parametrize("has_access", [True, False])
@pytest.mark.asyncio
async def test_chat_completion_client_fallbacks(has_access):
"""
make chat completion call with prompt > context window. expect it to work with fallback
"""
async def test_chat_completion_client_fallbacks(has_access: bool) -> None:
models: Final = ["gpt-3.5-turbo", "gpt-6-luna"] if has_access else ["gpt-3.5-turbo"]
async with aiohttp.ClientSession() as session:
models = ["gpt-3.5-turbo"]
if has_access:
models.append("gpt-instruct")
## CREATE KEY WITH MODELS
generated_key = await generate_key(session=session, i=0, models=models)
calling_key = generated_key["key"]
model = "gpt-3.5-turbo"
messages = [
{"role": "user", "content": "Who was Alexander?"},
]
## CALL PROXY
try:
await chat_completion(
session=session,
key=calling_key,
model=model,
messages=messages,
mock_testing_fallbacks=True,
fallbacks=["gpt-instruct"],
)
if not has_access:
pytest.fail(
"Expected this to fail, submitted fallback model that key did not have access to"
)
except Exception as e:
if has_access:
pytest.fail("Expected this to work: {}".format(str(e)))
generated_key: Final = await generate_key(session=session, i=0, models=models)
async with AsyncOpenAI(api_key=generated_key["key"], base_url=PROXY_BASE_URL, max_retries=0) as client:
request: Final = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "Who was Alexander?"}],
"max_tokens": 32,
"temperature": 0,
"extra_body": {
"mock_testing_fallbacks": True,
"fallbacks": ["gpt-6-luna"],
},
}
if not has_access:
with pytest.raises(PermissionDeniedError) as denied:
await client.chat.completions.create(**request)
assert denied.value.status_code == 403
assert "gpt-6-luna" in str(denied.value)
return
response: Final = await client.chat.completions.create(**request)
assert response.model == "gpt-6-luna"
assert response.choices[0].message.content
@pytest.mark.asyncio
@ -241,55 +235,66 @@ async def test_chat_completion_with_timeout_from_request():
@pytest.mark.parametrize("has_access", [True, False])
@pytest.mark.asyncio
async def test_chat_completion_client_fallbacks_with_custom_message(has_access):
"""
make chat completion call with prompt > context window. expect it to work with fallback
"""
async def test_chat_completion_client_fallbacks_with_custom_message(has_access: bool) -> None:
original_messages: Final = [{"role": "user", "content": "Who was Alexander?"}]
custom_messages: Final = [
{
"role": "user",
"content": (
"Describe the weather in a coastal city during winter, including the usual temperature, rain, wind, "
"and the clothing a visitor should bring."
),
}
]
models: Final = ["gpt-3.5-turbo", "gpt-6-luna"] if has_access else ["gpt-3.5-turbo"]
async with aiohttp.ClientSession() as session:
models = ["gpt-3.5-turbo"]
if has_access:
models.append("gpt-instruct")
## CREATE KEY WITH MODELS
generated_key = await generate_key(session=session, i=0, models=models)
calling_key = generated_key["key"]
model = "gpt-3.5-turbo"
messages = [
{"role": "user", "content": "Who was Alexander?"},
]
## CALL PROXY
try:
await chat_completion(
session=session,
key=calling_key,
model=model,
messages=messages,
mock_testing_fallbacks=True,
fallbacks=[
generated_key: Final = await generate_key(session=session, i=0, models=models)
async with AsyncOpenAI(api_key=generated_key["key"], base_url=PROXY_BASE_URL, max_retries=0) as client:
request: Final = {
"model": "gpt-3.5-turbo",
"messages": original_messages,
"max_tokens": 32,
"temperature": 0,
"extra_body": {
"mock_testing_fallbacks": True,
"fallbacks": [
{
"model": "gpt-instruct",
"messages": [
{
"role": "assistant",
"content": "This is a custom message",
}
],
"model": "gpt-6-luna",
"messages": custom_messages,
}
],
)
if not has_access:
pytest.fail(
"Expected this to fail, submitted fallback model that key did not have access to"
)
except Exception as e:
if has_access:
pytest.fail("Expected this to work: {}".format(str(e)))
},
}
if not has_access:
with pytest.raises(PermissionDeniedError) as denied:
await client.chat.completions.create(**request)
assert denied.value.status_code == 403
assert "gpt-6-luna" in str(denied.value)
return
response: Final = await client.chat.completions.create(**request)
assert response.model == "gpt-6-luna"
assert response.choices[0].message.content
custom_control: Final = await client.chat.completions.create(
model="gpt-6-luna",
messages=custom_messages,
max_tokens=32,
temperature=0,
)
original_control: Final = await client.chat.completions.create(
model="gpt-6-luna",
messages=original_messages,
max_tokens=32,
temperature=0,
)
assert response.usage is not None
assert custom_control.usage is not None
assert original_control.usage is not None
assert custom_control.usage.completion_tokens > 0
assert original_control.usage.completion_tokens > 0
assert custom_control.usage.prompt_tokens != original_control.usage.prompt_tokens
assert response.usage.prompt_tokens == custom_control.usage.prompt_tokens
from openai import AsyncOpenAI
from typing import List

40
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -4503,7 +4503,7 @@ class TestPriceDataReloadAPI:
"""Test cases for price data reload API endpoints"""
@pytest.fixture
def client_with_auth(self):
def client_with_auth(self, monkeypatch):
"""Create a test client with authentication"""
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.proxy_server import cleanup_router_config_variables
@ -4516,7 +4516,7 @@ class TestPriceDataReloadAPI:
# Mock admin user authentication
mock_auth = MagicMock()
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
return TestClient(app)
@ -4557,12 +4557,12 @@ class TestPriceDataReloadAPI:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_reload_model_cost_map_non_admin_access(self, client_with_auth):
def test_reload_model_cost_map_non_admin_access(self, client_with_auth, monkeypatch):
"""Test that non-admin users cannot access the reload endpoint"""
# Mock non-admin user
mock_auth = MagicMock()
mock_auth.user_role = "user" # Non-admin role
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
response = client_with_auth.post("/reload/model_cost_map")
@ -4623,12 +4623,12 @@ class TestPriceDataReloadAPI:
assert set(create_payload.keys()) == {"param_name", "param_value"}
assert json.loads(create_payload["param_value"]) == {"interval_hours": 6}
def test_schedule_model_cost_map_reload_non_admin_access(self, client_with_auth):
def test_schedule_model_cost_map_reload_non_admin_access(self, client_with_auth, monkeypatch):
"""Test that non-admin users cannot schedule periodic reload"""
# Mock non-admin user
mock_auth = MagicMock()
mock_auth.user_role = "user" # Non-admin role
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
response = client_with_auth.post("/schedule/model_cost_map_reload?hours=6")
@ -4663,12 +4663,12 @@ class TestPriceDataReloadAPI:
}
mock_prisma.db.litellm_config.delete.assert_not_called()
def test_cancel_model_cost_map_reload_non_admin_access(self, client_with_auth):
def test_cancel_model_cost_map_reload_non_admin_access(self, client_with_auth, monkeypatch):
"""Test that non-admin users cannot cancel periodic reload"""
# Mock non-admin user
mock_auth = MagicMock()
mock_auth.user_role = "user" # Non-admin role
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
response = client_with_auth.delete("/schedule/model_cost_map_reload")
@ -4701,12 +4701,12 @@ class TestPriceDataReloadAPI:
assert data["last_run"] == "2024-01-01T06:00:00+00:00"
assert data["next_run"] == "2024-01-01T12:00:00+00:00"
def test_get_model_cost_map_reload_status_non_admin_access(self, client_with_auth):
def test_get_model_cost_map_reload_status_non_admin_access(self, client_with_auth, monkeypatch):
"""Test that non-admin users cannot get reload status"""
# Mock non-admin user
mock_auth = MagicMock()
mock_auth.user_role = "user" # Non-admin role
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
response = client_with_auth.get("/schedule/model_cost_map_reload/status")
@ -4769,7 +4769,7 @@ class TestPriceDataReloadIntegration:
"""Integration tests for the complete price data reload feature"""
@pytest.fixture
def client_with_auth(self):
def client_with_auth(self, monkeypatch):
"""Create a test client with authentication"""
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.proxy_server import cleanup_router_config_variables
@ -4782,7 +4782,7 @@ class TestPriceDataReloadIntegration:
# Mock admin user authentication
mock_auth = MagicMock()
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
return TestClient(app)
@ -5262,7 +5262,7 @@ class TestPriceDataReloadIntegration:
litellm_utils._runtime_registered_model_cost.update(original_registry)
_invalidate_model_cost_lowercase_map()
def test_manual_reload_preserves_interval_hours(self):
def test_manual_reload_preserves_interval_hours(self, monkeypatch):
"""
Regression: manual reload owns only the run columns, so it never reads or rewrites
param_value and cannot destroy an existing schedule
@ -5277,7 +5277,7 @@ class TestPriceDataReloadIntegration:
mock_auth = MagicMock()
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
client = TestClient(app)
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
@ -5358,7 +5358,7 @@ class TestPriceDataReloadIntegration:
"dropping it causes the schedule to self-destruct"
)
def test_anthropic_beta_headers_manual_reload_preserves_interval_hours(self):
def test_anthropic_beta_headers_manual_reload_preserves_interval_hours(self, monkeypatch):
"""Test that manual reload via /reload/anthropic_beta_headers preserves existing interval_hours.
Regression test: the manual reload endpoint was overwriting param_value with
@ -5374,7 +5374,7 @@ class TestPriceDataReloadIntegration:
mock_auth = MagicMock()
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
client = TestClient(app)
with patch("litellm.anthropic_beta_headers_manager.reload_beta_headers_config") as mock_reload:
@ -7155,7 +7155,7 @@ class TestInvitationEndpoints:
"""Tests for /invitation/new and /invitation/delete endpoints."""
@pytest.fixture
def client_with_auth(self):
def client_with_auth(self, monkeypatch):
"""Create a test client with admin authentication."""
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.proxy_server import cleanup_router_config_variables
@ -7169,7 +7169,7 @@ class TestInvitationEndpoints:
mock_auth.user_id = "admin-user-id"
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
mock_auth.api_key = "sk-test"
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
return TestClient(app)
@ -7238,7 +7238,7 @@ class TestInvitationEndpoints:
("/invitation/delete", {"invitation_id": "inv-456"}),
],
)
def test_invitation_endpoints_non_admin_denied(self, client_with_auth, endpoint, payload):
def test_invitation_endpoints_non_admin_denied(self, client_with_auth, endpoint, payload, monkeypatch):
"""Non-admin users cannot access invitation endpoints."""
from litellm.proxy._types import LitellmUserRoles
@ -7246,7 +7246,7 @@ class TestInvitationEndpoints:
mock_auth.user_id = "regular-user"
mock_auth.user_role = LitellmUserRoles.INTERNAL_USER
mock_auth.api_key = "sk-regular"
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: mock_auth)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_prisma.db.litellm_invitationlink = MagicMock()

12
tests/test_openai_endpoints.py
View file

@ -1,3 +1,5 @@
import os
from typing import Final
# What this tests ?
## Tests /chat/completions by generating a key and then making a chat completions-request
import pytest
@ -398,10 +400,12 @@ async def test_completion_streaming_usage_metrics():
"""
[PROD Test] Ensures usage metrics are returned correctly when `include_usage` is set to `True`
"""
client = AsyncOpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
client: Final = AsyncOpenAI(
api_key="sk-1234", base_url=os.environ.get("LITELLM_PROXY_BASE_URL", "http://0.0.0.0:4000")
)
response = await client.completions.create(
model="gpt-instruct",
model="gpt-6-luna",
prompt="hey",
stream=True,
stream_options={"include_usage": True},
@ -417,9 +421,7 @@ async def test_completion_streaming_usage_metrics():
assert last_chunk is not None, "No chunks were received"
assert last_chunk.usage is not None, "Usage information was not received"
assert last_chunk.usage.prompt_tokens > 0, "Prompt tokens should be greater than 0"
assert (
last_chunk.usage.completion_tokens > 0
), "Completion tokens should be greater than 0"
assert last_chunk.usage.completion_tokens > 0, "Completion tokens should be greater than 0"
assert last_chunk.usage.total_tokens > 0, "Total tokens should be greater than 0"

58
tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py
View file

@ -78,6 +78,64 @@ def test_completion_cost_bills_the_price_columns_of_the_service_tier(
assert cost == pytest.approx(_cost_at(TIER_ROW, column_suffix))
LONG_CONTEXT_TIER_MODEL: Final = "long-context-tier-priced-test-model"
LONG_CONTEXT_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType(
{
"input_cost_per_token": 4e-06,
"output_cost_per_token": 8e-06,
"input_cost_per_token_ultrafast": 1e-05,
"output_cost_per_token_ultrafast": 2e-05,
"input_cost_per_token_above_272k_tokens_ultrafast": 5e-05,
"output_cost_per_token_above_272k_tokens_ultrafast": 6e-05,
}
)
@pytest.mark.parametrize(
("service_tier", "prompt_tokens", "input_rate", "output_rate"),
(
pytest.param("ultrafast", 300_000, 5e-05, 6e-05, id="long-ultrafast"),
pytest.param(None, 300_000, 4e-06, 8e-06, id="long-standard"),
pytest.param("ultrafast", 1_000, 1e-05, 2e-05, id="short-ultrafast"),
pytest.param("priority", 300_000, 4e-06, 8e-06, id="long-priority-falls-back"),
),
)
def test_completion_cost_uses_only_the_request_tiers_long_context_rates(
local_model_cost_map: None,
service_tier: str | None,
prompt_tokens: int,
input_rate: float,
output_rate: float,
) -> None:
litellm.register_model(
{
LONG_CONTEXT_TIER_MODEL: {
"litellm_provider": "openai",
"mode": "chat",
**dict(LONG_CONTEXT_TIER_ROW),
}
}
)
completion_tokens: Final = 100
response: Final = ModelResponse(
model=LONG_CONTEXT_TIER_MODEL,
usage=Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
),
)
cost: Final = litellm.completion_cost(
completion_response=response,
model=LONG_CONTEXT_TIER_MODEL,
custom_llm_provider="openai",
service_tier=service_tier,
)
assert cost == pytest.approx(prompt_tokens * input_rate + completion_tokens * output_rate)
class _CostRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()

2
tests/unit/test_anthropic_beta_headers_filtering.py
View file

@ -444,7 +444,7 @@ class TestAnthropicBetaHeadersFiltering:
assert filtered == ["thinking-binding-controls-2026-08-01"]
@pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"])
@pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "bedrock", "bedrock_mantle", "vertex_ai"])
def test_dangerous_tool_use_forwarded(self, provider):
"""Claude Code's server-side auto-mode classifier sends `safeguards` together with
dangerous-tool-use-2026-09-03. Bedrock Invoke, Bedrock Mantle, and Vertex rawPredict

455
tests/unit/test_router_model_cost_isolation.py
View file

@ -23,6 +23,7 @@ from litellm import Router
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE
from litellm.litellm_core_utils.ptu_pricing import ptu_config_error
from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@ -862,6 +863,425 @@ def test_inherit_builtin_cache_pricing_noop_for_unknown_backend():
assert model_info == {"input_cost_per_token": 0.000003}
_TIER_BACKEND_MODEL: Final = "tier-priced-backend"
_TIER_BACKEND_KEY: Final = f"openai/{_TIER_BACKEND_MODEL}"
_CUSTOM_STANDARD_INPUT_RATE: Final = 0.00011
_CUSTOM_STANDARD_OUTPUT_RATE: Final = 0.00022
_TIER_BACKEND_ENTRY: Final = {
"key": _TIER_BACKEND_KEY,
"litellm_provider": "openai",
"mode": "chat",
"max_tokens": 123456,
"input_cost_per_token": 0.00021,
"output_cost_per_token": 0.00032,
"input_cost_per_token_ultrafast": 0.00031,
"output_cost_per_token_ultrafast": 0.00042,
"input_cost_per_token_priority": 0.00051,
"output_cost_per_token_priority": 0.00062,
"input_cost_per_token_flex": 0.00071,
"output_cost_per_token_flex": 0.00082,
"input_cost_per_token_balanced": 0.00091,
"output_cost_per_token_balanced": 0.00102,
"cache_read_input_token_cost_ultrafast": 0.00013,
"input_cost_per_token_above_272k_tokens_ultrafast": 0.00014,
"output_cost_per_token_above_272k_tokens_ultrafast": 0.00015,
"input_cost_per_token_batches": 0.00016,
"input_cost_per_token_above_272k_tokens": 0.00017,
}
_AZURE_TIER_BACKEND_KEY: Final = "azure/tier-priced-backend"
_AZURE_TIER_BACKEND_ENTRY: Final = {
**_TIER_BACKEND_ENTRY,
"key": _AZURE_TIER_BACKEND_KEY,
"litellm_provider": "azure",
}
def _register_tier_backend() -> None:
litellm.model_cost[_TIER_BACKEND_KEY] = copy.deepcopy(_TIER_BACKEND_ENTRY)
litellm.get_model_info.cache_clear()
_invalidate_model_cost_lowercase_map()
def _register_azure_tier_backend() -> None:
litellm.model_cost[_AZURE_TIER_BACKEND_KEY] = copy.deepcopy(_AZURE_TIER_BACKEND_ENTRY)
litellm.get_model_info.cache_clear()
_invalidate_model_cost_lowercase_map()
def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> None:
model_cost_entries: Final = {
key: copy.deepcopy(litellm.model_cost.get(key))
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
}
try:
_register_tier_backend()
model_info: Final = {
"id": "custom-priced-tier-deployment",
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
"output_cost_per_token_ultrafast": 0.00999,
}
Router._inherit_builtin_service_tier_pricing(
model_info=model_info,
backend_model=_TIER_BACKEND_MODEL,
custom_llm_provider="openai",
)
assert model_info == {
"id": "custom-priced-tier-deployment",
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
"input_cost_per_token_ultrafast": _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"],
"output_cost_per_token_ultrafast": 0.00999,
"input_cost_per_token_priority": _TIER_BACKEND_ENTRY["input_cost_per_token_priority"],
"output_cost_per_token_priority": _TIER_BACKEND_ENTRY["output_cost_per_token_priority"],
"input_cost_per_token_flex": _TIER_BACKEND_ENTRY["input_cost_per_token_flex"],
"output_cost_per_token_flex": _TIER_BACKEND_ENTRY["output_cost_per_token_flex"],
"input_cost_per_token_balanced": _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"],
"output_cost_per_token_balanced": _TIER_BACKEND_ENTRY["output_cost_per_token_balanced"],
"cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY[
"cache_read_input_token_cost_ultrafast"
],
"input_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[
"input_cost_per_token_above_272k_tokens_ultrafast"
],
"output_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[
"output_cost_per_token_above_272k_tokens_ultrafast"
],
}
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend() -> None:
model_cost_entries: Final = {
key: copy.deepcopy(litellm.model_cost.get(key))
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
}
try:
_register_tier_backend()
model_info_without_base_rate: Final = {
"id": "custom-priced-no-base-rate",
"input_cost_per_token_ultrafast": 0.00031,
}
expected_without_base_rate: Final = copy.deepcopy(model_info_without_base_rate)
Router._inherit_builtin_service_tier_pricing(
model_info=model_info_without_base_rate,
backend_model=_TIER_BACKEND_MODEL,
custom_llm_provider="openai",
)
model_info_with_unknown_backend: Final = {
"id": "custom-priced-unknown-backend",
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
}
expected_with_unknown_backend: Final = copy.deepcopy(model_info_with_unknown_backend)
Router._inherit_builtin_service_tier_pricing(
model_info=model_info_with_unknown_backend,
backend_model="tier-priced-backend-unknown",
custom_llm_provider="openai",
)
assert model_info_without_base_rate == expected_without_base_rate
assert model_info_with_unknown_backend == expected_with_unknown_backend
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None:
model_id: Final = "tier-priced-deployment"
model_cost_entries: Final = {
key: copy.deepcopy(litellm.model_cost.get(key))
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
}
try:
_register_tier_backend()
router: Final = Router(
model_list=[
{
"model_name": "tier-priced-router",
"litellm_params": {
"model": _TIER_BACKEND_MODEL,
"custom_llm_provider": "openai",
"api_key": "sk-tier-pricing-not-used",
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
},
"model_info": {
"id": model_id,
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
},
}
]
)
ultrafast_response: Final = router.completion(
model="tier-priced-router",
messages=[{"role": "user", "content": "tiered pricing"}],
service_tier="ultrafast",
mock_response=litellm.ModelResponse(
model=_TIER_BACKEND_MODEL,
service_tier="ultrafast",
usage=litellm.Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100),
),
)
standard_response: Final = router.completion(
model="tier-priced-router",
messages=[{"role": "user", "content": "standard pricing"}],
mock_response=litellm.ModelResponse(
model=_TIER_BACKEND_MODEL,
usage=litellm.Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100),
),
)
assert isinstance(ultrafast_response, litellm.ModelResponse)
assert ultrafast_response._hidden_params["response_cost"] == pytest.approx(
1000 * _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"]
+ 100 * _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"]
)
assert isinstance(standard_response, litellm.ModelResponse)
assert standard_response._hidden_params["response_cost"] == pytest.approx(
1000 * _CUSTOM_STANDARD_INPUT_RATE + 100 * _CUSTOM_STANDARD_OUTPUT_RATE
)
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
def test_router_completion_uses_backend_ultrafast_long_context_rates() -> None:
model_id: Final = "tier-priced-long-context-deployment"
model_cost_entries: Final = {
key: copy.deepcopy(litellm.model_cost.get(key))
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
}
try:
_register_tier_backend()
router: Final = Router(
model_list=[
{
"model_name": "tier-priced-long-context-router",
"litellm_params": {
"model": _TIER_BACKEND_MODEL,
"custom_llm_provider": "openai",
"api_key": "sk-tier-pricing-not-used",
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
},
"model_info": {
"id": model_id,
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
},
}
]
)
response: Final = router.completion(
model="tier-priced-long-context-router",
messages=[{"role": "user", "content": "long context tiered pricing"}],
service_tier="ultrafast",
mock_response=litellm.ModelResponse(
model=_TIER_BACKEND_MODEL,
service_tier="ultrafast",
usage=litellm.Usage(prompt_tokens=300_000, completion_tokens=100, total_tokens=300_100),
),
)
assert isinstance(response, litellm.ModelResponse)
assert response._hidden_params["response_cost"] == pytest.approx(
300_000 * _TIER_BACKEND_ENTRY["input_cost_per_token_above_272k_tokens_ultrafast"]
+ 100 * _TIER_BACKEND_ENTRY["output_cost_per_token_above_272k_tokens_ultrafast"]
)
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
@pytest.mark.parametrize("ptu_enabled", (True, False))
def test_ptu_service_tier_pricing_is_disabled_only_when_attribution_is_enabled(
monkeypatch: pytest.MonkeyPatch, ptu_enabled: bool
) -> None:
model_id: Final = f"ptu-tier-deployment-{ptu_enabled}"
model_cost_entries: Final = {
key: copy.deepcopy(litellm.model_cost.get(key))
for key in (_TIER_BACKEND_KEY, model_id)
}
try:
_register_tier_backend()
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True" if ptu_enabled else "")
router: Final = Router(
model_list=[
{
"model_name": f"ptu-tier-model-{ptu_enabled}",
"litellm_params": {
"model": _TIER_BACKEND_MODEL,
"custom_llm_provider": "openai",
"api_key": "sk-tier-pricing-not-used",
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
},
"model_info": {**_PTU_MODEL_INFO, "id": model_id},
}
]
)
registered: Final = litellm.model_cost[model_id]
tier_fields: Final = tuple(
field for field in _TIER_BACKEND_ENTRY if field.endswith(SERVICE_TIER_COST_KEY_SUFFIXES)
)
if ptu_enabled:
assert all(field not in registered for field in tier_fields)
else:
assert all(field in registered for field in tier_fields)
response: Final = router.completion(
model=f"ptu-tier-model-{ptu_enabled}",
messages=[{"role": "user", "content": "ptu service tier pricing"}],
service_tier="priority",
mock_response=litellm.ModelResponse(
model=_TIER_BACKEND_MODEL,
service_tier="priority",
usage=litellm.Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100),
),
)
assert isinstance(response, litellm.ModelResponse)
expected_cost: Final = (
0.0
if ptu_enabled
else 1000 * _TIER_BACKEND_ENTRY["input_cost_per_token_priority"]
+ 100 * _TIER_BACKEND_ENTRY["output_cost_per_token_priority"]
)
assert response._hidden_params["response_cost"] == pytest.approx(expected_cost)
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_payload() -> None:
model_id: Final = "azure-tier-priced-alias"
payload_id: Final = "azure-tier-priced-payload"
model_cost_entries: Final = {
key: copy.deepcopy(litellm.model_cost.get(key))
for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id)
}
try:
_register_azure_tier_backend()
router: Final = Router(
model_list=[
{
"model_name": "azure/tier-priced-alias",
"litellm_params": {
"model": "azure/tier-priced-alias",
"custom_llm_provider": "azure",
"api_key": "sk-tier-pricing-not-used",
"api_base": "https://tier-priced.azure.invalid",
},
"model_info": {
"id": model_id,
"base_model": _AZURE_TIER_BACKEND_KEY,
"input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE,
"output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE,
},
}
]
)
response: Final = router.completion(
model="azure/tier-priced-alias",
messages=[{"role": "user", "content": "azure base model pricing"}],
service_tier="priority",
allowed_openai_params=["service_tier"],
mock_response=litellm.ModelResponse(
model=_AZURE_TIER_BACKEND_KEY,
service_tier="priority",
usage=litellm.Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100),
),
)
assert isinstance(response, litellm.ModelResponse)
assert response._hidden_params["response_cost"] == pytest.approx(
1000 * _AZURE_TIER_BACKEND_ENTRY["input_cost_per_token_priority"]
+ 100 * _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"]
)
payload: Final = Router._deployment_model_cost_payload(
deployment=Deployment(
model_name="azure/tier-priced-alias-from-params",
litellm_params=LiteLLM_Params(
model="azure/tier-priced-alias",
custom_llm_provider="azure",
base_model=_AZURE_TIER_BACKEND_KEY,
input_cost_per_token=_CUSTOM_STANDARD_INPUT_RATE,
output_cost_per_token=_CUSTOM_STANDARD_OUTPUT_RATE,
),
model_info=ModelInfo(id=payload_id),
)
)
assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[
"input_cost_per_token_priority"
]
assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[
"output_cost_per_token_priority"
]
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
@pytest.mark.parametrize(
("model_info_base_model", "params_base_model", "model", "expected"),
(
pytest.param(
"azure/tier-priced-model-info-base",
"azure/tier-priced-params-base",
"azure/tier-priced-deployment-alias",
"azure/tier-priced-model-info-base",
id="model-info-base-model-wins",
),
pytest.param(
None,
"azure/tier-priced-params-base",
"azure/tier-priced-deployment-alias",
"azure/tier-priced-params-base",
id="params-base-model-fallback",
),
pytest.param(
None,
None,
"azure/tier-priced-deployment-alias",
"azure/tier-priced-deployment-alias",
id="model-fallback",
),
pytest.param(
"",
"azure/tier-priced-params-base",
"azure/tier-priced-deployment-alias",
"azure/tier-priced-params-base",
id="empty-model-info-base-model-falls-through",
),
),
)
def test_cost_map_backend_model_uses_canonical_model_precedence(
model_info_base_model: str | None,
params_base_model: str | None,
model: str,
expected: str,
) -> None:
deployment: Final = Deployment(
model_name="azure/tier-priced-cost-map-backend",
litellm_params=LiteLLM_Params(model=model, base_model=params_base_model),
model_info=ModelInfo(id="tier-priced-cost-map-backend", base_model=model_info_base_model),
)
assert Router._cost_map_backend_model(deployment) == expected
def test_inherit_builtin_base_rates_for_off_peak_fills_missing_rates():
"""Direct unit test of the helper: an entry carrying only an
off_peak_pricing block inherits the backend model's built-in base token
@ -1803,6 +2223,41 @@ def test_deployment_model_cost_payload_folds_in_litellm_params_pricing():
assert payload["cache_read_input_token_cost"] > 0
def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -> None:
model_id: Final = "tier-priced-payload"
model_cost_entries: Final = {
key: copy.deepcopy(litellm.model_cost.get(key))
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
}
try:
_register_tier_backend()
payload: Final = Router._deployment_model_cost_payload(
deployment=Deployment(
model_name="tier-priced-payload",
litellm_params=LiteLLM_Params(
model=_TIER_BACKEND_MODEL,
custom_llm_provider="openai",
input_cost_per_token=_CUSTOM_STANDARD_INPUT_RATE,
output_cost_per_token=_CUSTOM_STANDARD_OUTPUT_RATE,
),
model_info=ModelInfo(id=model_id),
)
)
assert (
payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"]
)
assert (
payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"]
)
assert payload["input_cost_per_token_balanced"] == _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"]
assert payload["input_cost_per_token"] == _CUSTOM_STANDARD_INPUT_RATE
assert payload["output_cost_per_token"] == _CUSTOM_STANDARD_OUTPUT_RATE
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
def test_register_deployment_in_model_cost_writes_both_key_families():
"""
A deployment contributes its full model_info under its unique id and the