fix: preserve moved Rust configuration and support relocated cost maps

litellm/scripts/build_core_distribution.py

@@ -12,8 +12,6 @@ ROOT: Final = Path(__file__).resolve().parents[2]
 SOURCES: Final = (
     "litellm",
     "litellm/rust",
-    ".cargo",
-    "rust-toolchain.toml",
     "model_prices_and_context_window.json",
     "README.md",
     "LICENSE",

litellm/scripts/ci_cd/cost_map_guard.py

@@ -26,6 +26,7 @@ from types import MappingProxyType
 from typing import Final
 
 COST_MAP_PATH: Final = "model_prices_and_context_window.json"
+COST_MAP_PATHS: Final = ("litellm/data/model_prices_and_context_window.json", COST_MAP_PATH)
 BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
 PINS_PATH: Final = "litellm/scripts/ci_cd/cost_map_pins.json"
 PINS_PATHS: Final = (
@@ -33,7 +34,7 @@ PINS_PATHS: Final = (
     "scripts/ci_cd/cost_map_pins.json",
     "ci_cd/cost_map_pins.json",
 )
-GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
+GUARDED_PATHS: Final = (*COST_MAP_PATHS, BACKUP_PATH)
 CHECKED_PATHS: Final = (*GUARDED_PATHS, *PINS_PATHS)
 BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
 PRICE_SYNC_BRANCH_PREFIX: Final = "litellm-providers/price-sync"
@@ -258,7 +259,8 @@ def _git(*args: str) -> str | None:
 
 def snapshot(revision: str) -> Snapshot:
     return Snapshot(
-        *(_git("show", f"{revision}:{path}") or "" for path in GUARDED_PATHS),
+        next((text for path in COST_MAP_PATHS if (text := _git("show", f"{revision}:{path}")) is not None), ""),
+        _git("show", f"{revision}:{BACKUP_PATH}") or "",
         next((text for path in PINS_PATHS if (text := _git("show", f"{revision}:{path}")) is not None), ""),
     )
 

tests/base_sdk_tests/test_core_distribution.py

@@ -78,6 +78,10 @@ def test_staging_stamps_release_version_without_modifying_sources(
         elif name == "litellm/rust":
             path.mkdir()
             (path / "Cargo.toml").write_text("shared rust")
+            for configuration in ("rust-toolchain.toml", ".cargo/config.toml"):
+                target: Final = path / configuration
+                target.parent.mkdir(parents=True, exist_ok=True)
+                target.write_bytes((ROOT / "litellm" / "rust" / configuration).read_bytes())
         else:
             path.parent.mkdir(parents=True, exist_ok=True)
             path.write_text(f"shared {name}")
@@ -96,7 +100,10 @@ def test_staging_stamps_release_version_without_modifying_sources(
         assert not (stage / "litellm" / "__init__.py").exists()
         assert all(
             (stage / name).exists()
-            for name in ("litellm/rust", ".cargo", "rust-toolchain.toml", "model_prices_and_context_window.json")
+            for name in (
+                "litellm/rust", "litellm/rust/.cargo/config.toml", "litellm/rust/rust-toolchain.toml",
+                "model_prices_and_context_window.json",
+            )
         )
     assert manifest.read_bytes() == original
     assert (source / "pyproject.toml").read_text() == '[project]\nname = "litellm"\nversion = "9.8.7rc1"\n'
@@ -198,7 +205,8 @@ def test_core_artifacts_exclude_local_configuration(distributions: tuple[Path, P
 
 
 @pytest.mark.parametrize(
-    "relative_path", ["rust-toolchain.toml", ".cargo/config.toml", "model_prices_and_context_window.json"]
+    "relative_path",
+    ["litellm/rust/rust-toolchain.toml", "litellm/rust/.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"))

tests/unit/test_cost_map_guard.py

@@ -374,12 +374,31 @@ def test_relocated_pins_cannot_be_changed_by_bots(pins_path: str) -> None:
     assert guard.pin_author_failures("litellm_fix_pricing", "Bot", (pins_path,)) == (BOT_AUTHOR_PIN_FAILURE,)
 
 
+@pytest.mark.parametrize("cost_map_path", guard.COST_MAP_PATHS)
+@pytest.mark.parametrize("content", (_serialize(BASE_MAP) + "\n", "", "not json"))
+def test_main_reads_cost_maps_in_both_repository_layouts(
+    tmp_path: Path, cost_map_path: str, content: str
+) -> None:
+    subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
+    base: Final = _commit(tmp_path, BASE_MAP, "base")
+    destination: Final = tmp_path / cost_map_path
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    if destination != tmp_path / guard.COST_MAP_PATH:
+        (tmp_path / guard.COST_MAP_PATH).rename(destination)
+    destination.write_text(content)
+    (tmp_path / guard.BACKUP_PATH).write_text(content)
+    head: Final = _git_commit(tmp_path, "update cost map layout")
+    result: Final = _run_guard(tmp_path, base, head, "litellm_fix_pricing")
+    assert result.returncode == (0 if content.strip().startswith("{") else 1), result.stdout + result.stderr
+
+
 def _commit(repo: Path, cost_map: dict[str, object], message: str, pins: str | None = None) -> str:
     text = _serialize(cost_map)
     pins_file = repo / guard.PINS_PATH
     if pins is not None or not pins_file.exists():
         pins_file.parent.mkdir(parents=True, exist_ok=True)
         pins_file.write_text(NO_PINS if pins is None else pins)
+    (repo / guard.COST_MAP_PATH).parent.mkdir(parents=True, exist_ok=True)
     (repo / guard.COST_MAP_PATH).write_text(text)
     (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True)
     (repo / guard.BACKUP_PATH).write_text(text)