from __future__ import annotations import argparse import json import sys from collections.abc import Iterator, Mapping, Sequence from itertools import chain from pathlib import Path from types import MappingProxyType from typing import Final from pydantic import JsonValue, TypeAdapter ROOT: Final = Path(__file__).resolve().parents[1] COST_MAP_PATH: Final = ROOT / "model_prices_and_context_window.json" SOURCE_PATH: Final = ROOT / "provider_endpoints_support.json" OUTPUT_PATHS: Final = ( SOURCE_PATH, ROOT / "litellm/provider_endpoints_support_backup.json", ) JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) MODE_TO_FLAG: Final[Mapping[str, str]] = MappingProxyType( { "chat": "chat_completions", "responses": "responses", "completion": "text_completion", "embedding": "embeddings", "image_generation": "image_generations", "image_edit": "image_edits", "audio_transcription": "audio_transcriptions", "audio_speech": "audio_speech", "rerank": "rerank", "ocr": "ocr", "search": "search", "moderation": "moderations", "realtime": "realtime", "video_generation": "video_generations", } ) SUPPORTED_ENDPOINT_TO_FLAG: Final[Mapping[str, str]] = MappingProxyType( { "/v1/chat/completions": "chat_completions", "/v1/responses": "responses", "/v1/completions": "text_completion", "/v1/embeddings": "embeddings", "/v1/images/generations": "image_generations", "/v1/images/edits": "image_edits", "/v1/audio/transcriptions": "audio_transcriptions", "/v1/audio/speech": "audio_speech", "/v1/ocr": "ocr", "/v1/realtime": "realtime", "/v1/videos": "video_generations", "/v1/batch": "batches", "/v1/messages": "messages", "/v1beta/interactions": "interactions", "/v1/systemone": "systemone", } ) def _json_object(value: JsonValue, name: str) -> Mapping[str, JsonValue]: if not isinstance(value, dict): raise ValueError(f"{name} must be a JSON object") return value def _provider_slug(provider: str) -> str: if provider.startswith("vertex_ai-"): return "vertex_ai" if provider == "fireworks_ai-embedding-models": return "fireworks_ai" if provider == "bedrock_converse": return "bedrock" if provider == "text-completion-inception": return "inception" return provider def _model_provider_slug(model: JsonValue) -> str | None: if not isinstance(model, dict): return None provider: Final = model.get("litellm_provider") if not isinstance(provider, str) or not provider: return None return _provider_slug(provider) def _model_flags(model: Mapping[str, JsonValue]) -> tuple[str, ...]: mode: Final = model.get("mode") mode_flag: Final = MODE_TO_FLAG.get(mode) if isinstance(mode, str) else None supported_endpoints: Final = model.get("supported_endpoints") endpoint_flags: Final = ( tuple( SUPPORTED_ENDPOINT_TO_FLAG[endpoint] for endpoint in supported_endpoints if isinstance(endpoint, str) and endpoint in SUPPORTED_ENDPOINT_TO_FLAG ) if isinstance(supported_endpoints, list) else () ) return ((mode_flag,) if mode_flag is not None else ()) + endpoint_flags def _bridge_field(endpoint: JsonValue) -> str | None: if not isinstance(endpoint, dict) or endpoint.get("bridges_to_chat_completion") is not True: return None provider_json_field: Final = endpoint.get("provider_json_field") return provider_json_field if isinstance(provider_json_field, str) else None def _bridged_fields(endpoints: Mapping[str, JsonValue]) -> tuple[str, ...]: return tuple( dict.fromkeys( provider_json_field for endpoint in endpoints.values() if (provider_json_field := _bridge_field(endpoint)) is not None ) ) def _derived_keys(bridged_fields: tuple[str, ...]) -> tuple[str, ...]: return tuple(dict.fromkeys((*MODE_TO_FLAG.values(), *SUPPORTED_ENDPOINT_TO_FLAG.values(), *bridged_fields))) def _model_evidence( model: JsonValue, bridged_fields: tuple[str, ...] ) -> Iterator[tuple[str, str]]: provider: Final = _model_provider_slug(model) if provider is None or not isinstance(model, dict): return flags: Final = _model_flags(model) yield from ((provider, flag) for flag in flags) if "chat_completions" not in flags: return yield from ((provider, flag) for flag in bridged_fields) def _provider_evidence( cost_map: Mapping[str, JsonValue], bridged_fields: tuple[str, ...] ) -> Iterator[tuple[str, str]]: return chain.from_iterable( _model_evidence(model, bridged_fields) for model_name, model in cost_map.items() if model_name != "sample_spec" ) def _provider_slugs(cost_map: Mapping[str, JsonValue]) -> frozenset[str]: return frozenset( provider_slug for model_name, model in cost_map.items() if model_name != "sample_spec" and (provider_slug := _model_provider_slug(model)) is not None ) def _generated_provider( slug: str, value: JsonValue, derived_keys: tuple[str, ...], derived_key_set: frozenset[str], evidence: frozenset[tuple[str, str]], ) -> dict[str, JsonValue]: provider: Final = _json_object(value, f"provider {slug}") current_flags_value: Final = provider.get("endpoints") current_flags: Final = ( _json_object(current_flags_value, f"provider {slug} endpoints") if current_flags_value is not None else {} ) derived_flags: Final = {flag: (slug, flag) in evidence for flag in derived_keys} handwritten_flags: Final = { flag: value for flag, value in current_flags.items() if flag not in derived_key_set } return { **provider, "endpoints": {**derived_flags, **handwritten_flags}, } def generate_provider_endpoints_support( cost_map: Mapping[str, JsonValue], current_json: Mapping[str, JsonValue], ) -> dict[str, JsonValue]: current_providers: Final = _json_object( current_json.get("providers"), "provider_endpoints_support providers" ) endpoints: Final = _json_object(current_json.get("endpoints"), "provider_endpoints_support endpoints") bridged_fields: Final = _bridged_fields(endpoints) derived_keys: Final = _derived_keys(bridged_fields) derived_key_set: Final = frozenset(derived_keys) provider_slugs: Final = _provider_slugs(cost_map) missing_slugs: Final = tuple(sorted(provider_slugs - frozenset(current_providers))) if missing_slugs: raise ValueError(f"Cost map providers missing from provider_endpoints_support.json: {', '.join(missing_slugs)}") evidence: Final = frozenset(_provider_evidence(cost_map, bridged_fields)) providers: Final = { slug: _generated_provider(slug, value, derived_keys, derived_key_set, evidence) for slug, value in current_providers.items() } return {**current_json, "providers": providers} def generate_files( cost_map: Mapping[str, JsonValue], current_json: Mapping[str, JsonValue], output_paths: tuple[Path, ...], check: bool, ) -> int: generated: Final = generate_provider_endpoints_support(cost_map, current_json) content: Final = json.dumps(generated, indent=2, ensure_ascii=False) + "\n" stale_paths: Final = tuple( path for path in output_paths if not path.exists() or path.read_text(encoding="utf-8") != content ) if check: if not stale_paths: return 0 sys.stderr.write( "provider endpoint support is stale in " + ", ".join(str(path) for path in stale_paths) + "\nRun `uv run --no-sync python scripts/generate_provider_endpoints_support.py` to regenerate it\n" ) return 1 for path in output_paths: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(content, encoding="utf-8") return 0 def main(argv: Sequence[str] | None = None, *, repo_root: Path = ROOT) -> int: parser: Final = argparse.ArgumentParser(description="Generate provider endpoint support from the model cost map") parser.add_argument("--check", action="store_true", help="fail if either generated file is stale") args: Final = parser.parse_args(argv) cost_map: Final = JSON_OBJECT.validate_json((repo_root / COST_MAP_PATH.name).read_bytes()) current_json: Final = JSON_OBJECT.validate_json((repo_root / SOURCE_PATH.name).read_bytes()) return generate_files( cost_map, current_json, ( repo_root / SOURCE_PATH.name, repo_root / OUTPUT_PATHS[1].relative_to(ROOT), ), args.check, ) if __name__ == "__main__": sys.exit(main())