Merge pull request #41847 from BerriAI/litellm_lit_8128_off_peak_pricing_schema

fix(schema): classify off_peak_pricing as a structured object in the model prices schema generator
This commit is contained in:
kerry-berri 2026-09-18 17:53:13 -07:00 committed by GitHub
commit f71129f65b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 99 additions and 4 deletions

View file

@ -3,7 +3,7 @@ from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Optional
from typing import Final, Optional
import jsonschema
@ -19,6 +19,10 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
BOOLEAN: JsonSchema = {"type": "boolean"}
STRING: JsonSchema = {"type": "string"}
TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
WEEKDAY_PATTERN: Final = (
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
)
EXTRA_BOOLEAN_KEYS = frozenset(
{
@ -31,7 +35,51 @@ EXTRA_BOOLEAN_KEYS = frozenset(
}
)
HOURS_UTC: Final[JsonSchema] = {
"description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
"oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
}
OFF_PEAK_WINDOW: Final[JsonSchema] = {
"type": "object",
"properties": {
"hours_utc": HOURS_UTC,
"weekdays": {
"type": "array",
"description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
"items": {
"oneOf": [
{"type": "integer", "minimum": 1, "maximum": 7},
{"type": "string", "pattern": WEEKDAY_PATTERN},
]
},
"minItems": 1,
},
},
"required": ["hours_utc"],
"additionalProperties": False,
}
OBJECT_KEYS: dict[str, JsonSchema] = {
"off_peak_pricing": {
"type": "object",
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
"properties": {
"hours_utc": HOURS_UTC,
"windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
"weekday_timezone": {
"type": "string",
"description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
},
"anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
"additionalProperties": False,
},
"search_context_cost_per_query": {
"type": "object",
"description": "USD cost per web search query, keyed by search context size.",
@ -327,9 +375,7 @@ def render(schema: JsonSchema) -> str:
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
validator = jsonschema.Draft202012Validator(
schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER
)
validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
return tuple(
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)

View file

@ -125,6 +125,55 @@ def test_schema_accepts_cache_creation_cost_inside_a_pricing_tier(committed_sche
assert validator.is_valid({"some-model": entry})
OFF_PEAK_ENTRY: Final = MappingProxyType(
{
"litellm_provider": "openrouter",
"mode": "chat",
"input_cost_per_token": 2e-6,
"output_cost_per_token": 8e-6,
"off_peak_pricing": {
"hours_utc": "16:30-00:30",
"windows": [{"hours_utc": ["00:30-02:00"], "weekdays": [6, "Sunday", "mon", "THURS"]}],
"weekday_timezone": "Asia/Shanghai",
"input_cost_per_token": 1e-6,
"output_cost_per_token": 4e-6,
"cache_read_input_token_cost": 1e-7,
},
}
)
def test_generator_classifies_off_peak_pricing_as_a_windowed_rate_block():
generator = load_generator()
schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)})))
validator = build_validator(schema)
assert validator.is_valid({"some-model": dict(OFF_PEAK_ENTRY)})
@pytest.mark.parametrize(
"block",
[
{"hours_utc": "16:30-00:30", "input_cost_per_token": "1e-6"},
{"hours_utc": "16:30-00:30", "input_cost_per_token": -1e-6},
{"hours_utc": 1630, "input_cost_per_token": 1e-6},
{"hours_utc": "16:30-00:30", "discount": 0.5},
{"windows": [{"weekdays": [6]}], "input_cost_per_token": 1e-6},
{"windows": [{"hours_utc": "00:30-02:00", "weekdays": [0]}], "input_cost_per_token": 1e-6},
{"windows": [], "input_cost_per_token": 1e-6},
{"input_cost_per_token": 1e-6},
{"hours_utc": "16:30", "input_cost_per_token": 1e-6},
{"hours_utc": "25:00-01:00", "input_cost_per_token": 1e-6},
{"hours_utc": ["16:30-00:30", "4pm-midnight"], "input_cost_per_token": 1e-6},
{"windows": [{"hours_utc": "00:30-02:00", "weekdays": ["Funday"]}], "input_cost_per_token": 1e-6},
],
)
def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict):
generator = load_generator()
schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)})))
validator = build_validator(schema)
assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}})
def find_duplicate_keys(path: Path) -> list[str]:
duplicates: list[str] = []