diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index ab29b70bdd4..bf669ccce29 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -31,7 +31,45 @@ EXTRA_BOOLEAN_KEYS = frozenset( } ) +HOURS_UTC: JsonSchema = { + "description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.', + "oneOf": [STRING, {"type": "array", "items": STRING, "minItems": 1}], +} + +OFF_PEAK_WINDOW: 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}, STRING]}, + "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, + }, + "additionalProperties": False, + }, "search_context_cost_per_query": { "type": "object", "description": "USD cost per web search query, keyed by search context size.", diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index 2f9b11a16b7..65aaa2870ef 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -125,6 +125,50 @@ 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"]}], + "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}, + ], +) +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] = []