build(model-catalog): ship the cost map inside the native extension (#48907)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

.github/scripts/smoke_test_native_wheel.py

@@ -7,6 +7,8 @@ import zipfile
 from pathlib import Path
 from typing import Final
 
+COST_MAP: Final = Path(__file__).resolve().parents[2] / "model_prices_and_context_window.json"
+
 CHILD_SCRIPT: Final = """
 from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
@@ -34,6 +36,9 @@ else:
 after = module.gil_stats()
 if not isinstance(after.get("releases"), int):
     raise AssertionError(f"native module unusable after panic: {after!r}")
+
+if module.bundled_model_cost_map() != Path(sys.argv[2]).read_bytes():
+    raise AssertionError("the wheel embeds a different model_prices_and_context_window.json than this checkout")
 """
 
 
@@ -55,7 +60,9 @@ def main() -> int:
 
         native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name
         native_path.write_bytes(archive.read(native_members[0]))
-        result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False)
+        result: Final = subprocess.run(
+            (sys.executable, "-c", CHILD_SCRIPT, str(native_path), str(COST_MAP)), check=False
+        )
 
     if result.returncode != 0:
         sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n")

backend/Dockerfile

@@ -24,6 +24,7 @@ ENV UV_PYTHON_DOWNLOADS=0
 COPY pyproject.toml README.md rust-toolchain.toml ./
 COPY .cargo ./.cargo
 COPY litellm-rust ./litellm-rust
+COPY model_prices_and_context_window.json ./
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     mkdir -p litellm/rust_bridge /wheels && \

gateway/Dockerfile

@@ -40,6 +40,7 @@ ENV UV_PYTHON_DOWNLOADS=0
 COPY pyproject.toml README.md rust-toolchain.toml ./
 COPY .cargo ./.cargo
 COPY litellm-rust ./litellm-rust
+COPY model_prices_and_context_window.json ./
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     mkdir -p litellm/rust_bridge /wheels && \

litellm-rust/Cargo.lock

@@ -4345,6 +4345,7 @@ dependencies = [
  "litellm-inference-transcription",
  "litellm-llms",
  "litellm-llms-types",
+ "litellm-model-catalog",
  "litellm-router-types",
  "litellm-secrets",
  "litellm-secrets-aws",

litellm-rust/crates/python-bridge/AGENTS.md

@@ -106,6 +106,18 @@ GIL handling to `litellm-host-python`.
 
 ## Data Handling
 
+- Reuse large immutable bundled data and parsed catalogs across calls, including in
+  replacement APIs. Share Rust storage and cache repeated Python projections instead
+  of reparsing, serializing or copying the full payload on every access
+- Cache owned Python objects with `PyOnceLock<Py<T>>`, then bind and clone references
+  while attached. Do not use blocking `OnceLock::get_or_init` while attached
+  ([PyO3 initialization](https://pyo3.rs/v0.29.2/faq.html),
+  [reference ownership](https://pyo3.rs/v0.29.2/types.html))
+- Never reuse Python objects across interpreters or interpreter reinitialization.
+  Once-only caches are for immutable data. Reloadable/custom data needs versioned
+  entries or invalidation, and caller-mutable Python containers must not be shared
+- Test unchanged contents and allocation reuse, not just equal values. Keep this
+  coverage when replacing a bundled-data accessor with a catalog API
 - OCR payloads can contain personal data and large base64 images. Do not log
   payloads or provider responses.
 - Avoid copying large payloads more than needed. The current JSON round-trip is

litellm-rust/crates/python-bridge/Cargo.toml

@@ -49,6 +49,7 @@ litellm-secrets-types.workspace = true
 litellm-llms-types.workspace = true
 litellm-router-types.workspace = true
 litellm-host-python.workspace = true
+litellm-model-catalog.workspace = true
 litellm-token-counter = { path = "../token-counter", default-features = false }
 pyo3.workspace = true
 prost.workspace = true

litellm-rust/crates/python-bridge/src/lib.rs

@@ -8,6 +8,7 @@ mod execution;
 mod http;
 mod logger;
 mod marshal;
+mod model_registry;
 mod preflight;
 mod python_settings;
 mod routes;
@@ -30,6 +31,8 @@ mod _native {
     #[pymodule_export]
     use crate::logger::NativeDiagnosticProcessor;
     #[pymodule_export]
+    use crate::model_registry::bundled_model_cost_map;
+    #[pymodule_export]
     use crate::routes::audio_transcription::{atranscription, transcription};
     #[pymodule_export]
     use crate::routes::chat_completions::{acompletion, completion};
@@ -110,6 +113,7 @@ mod tests {
                 "NativeClickHouseSpendConfig",
                 "NativeClickHouseSpendStorage",
                 "trace_encode_error",
+                "bundled_model_cost_map",
                 "TokenCounter",
                 "Tokenizer",
                 "gil_stats",

litellm-rust/crates/python-bridge/src/model_registry.rs

@@ -0,0 +1,30 @@
+use pyo3::prelude::*;
+use pyo3::sync::PyOnceLock;
+use pyo3::types::PyBytes;
+
+#[pyfunction]
+pub fn bundled_model_cost_map(py: Python<'_>) -> Bound<'_, PyBytes> {
+    static BYTES: PyOnceLock<Py<PyBytes>> = PyOnceLock::new();
+    BYTES
+        .get_or_init(py, || {
+            PyBytes::new(py, litellm_model_catalog::BUNDLED_REGISTRY).unbind()
+        })
+        .bind(py)
+        .clone()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[rstest::rstest]
+    fn repeated_calls_reuse_the_bundled_registry_bytes() {
+        Python::initialize();
+        let first = Python::attach(|py| bundled_model_cost_map(py).unbind());
+        Python::attach(|py| {
+            let second = bundled_model_cost_map(py);
+            assert!(first.bind(py).is(&second));
+            assert_eq!(second.as_bytes(), litellm_model_catalog::BUNDLED_REGISTRY);
+        });
+    }
+}

litellm/rust_bridge/_native.pyi

@@ -14,6 +14,7 @@ class ForkedAfterNativeRuntimeStarted(RuntimeError): ...
 class ProcessReservedForForking(RuntimeError): ...
 
 def trace_encode_error(message: str) -> bytes: ...
+def bundled_model_cost_map() -> bytes: ...
 
 @final
 class NativeClickHouseSpendConfig:
@@ -220,6 +221,7 @@ __all__ = [
     "aocr",
     "aresponses",
     "atranscription",
+    "bundled_model_cost_map",
     "completion",
     "embedding",
     "gil_stats",

migrations/Dockerfile

@@ -43,6 +43,7 @@ ENV UV_PYTHON_DOWNLOADS=0
 COPY pyproject.toml README.md rust-toolchain.toml ./
 COPY .cargo ./.cargo
 COPY litellm-rust ./litellm-rust
+COPY model_prices_and_context_window.json ./
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     mkdir -p litellm/rust_bridge /wheels && \

packaging/litellm-core/pyproject.toml

@@ -56,6 +56,7 @@ editable-profile = "dev"
 include = [
     { path = "rust-toolchain.toml", format = "sdist" },
     { path = ".cargo/config.toml", format = "sdist" },
+    { path = "model_prices_and_context_window.json", format = "sdist" },
     "litellm/router_strategy/complexity_router/artifacts/*.json",
     "litellm/router_strategy/complexity_router/fuse_presets.json",
     "litellm/proxy/model_insights_tasks.json",

pyproject.toml

@@ -335,6 +335,7 @@ include = [
     "litellm/proxy/public_endpoints/whats_new.json",
     "litellm/proxy/common_utils/codex_base_instructions.md",
     "litellm/proxy/common_utils/codex_bundled_models_0.159.3.json",
+    { path = "model_prices_and_context_window.json", format = "sdist" },
 ]
 exclude = [
     "litellm/proxy/enterprise",
@@ -357,6 +358,7 @@ cache-keys = [
     { file = "litellm-rust/Cargo.lock" },
     { file = "litellm-rust/Cargo.toml" },
     { file = "litellm-rust/crates/**/*" },
+    { file = "model_prices_and_context_window.json" },
 ]
 constraint-dependencies = [
     "tornado>=6.5.8",

scripts/build_core_distribution.py

@@ -9,7 +9,15 @@ from pathlib import Path
 from typing import Final
 
 ROOT: Final = Path(__file__).resolve().parents[1]
-SOURCES: Final = ("litellm", "litellm-rust", ".cargo", "rust-toolchain.toml", "README.md", "LICENSE")
+SOURCES: Final = (
+    "litellm",
+    "litellm-rust",
+    ".cargo",
+    "rust-toolchain.toml",
+    "model_prices_and_context_window.json",
+    "README.md",
+    "LICENSE",
+)
 
 
 def stage_core_distribution(source: Path, destination: Path) -> None:

tests/base_sdk_tests/test_core_distribution.py

@@ -132,7 +132,9 @@ def test_core_artifacts_exclude_local_configuration(distributions: tuple[Path, P
             assert not any(member.name.endswith(f"/{name}") for member in archive.getmembers())
 
 
-@pytest.mark.parametrize("relative_path", ["rust-toolchain.toml", ".cargo/config.toml"])
+@pytest.mark.parametrize(
+    "relative_path", ["rust-toolchain.toml", ".cargo/config.toml", "model_prices_and_context_window.json"]
+)
 def test_core_sdist_preserves_native_build_configuration(distribution_directory: Path, relative_path: str) -> None:
     sdists: Final = tuple(distribution_directory.glob("litellm_core-*.tar.gz"))
     assert len(sdists) == 1