Merge 538d72cb9e into bc18e10f6b

litellm/caching/in_memory_cache.py

@@ -47,7 +47,7 @@ class InMemoryCache(BaseCache):
         self.max_size_per_item = max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB  # 1MB = 1024KB
 
         # in-memory cache
-        self.cache_dict: dict = {}
+        self.cache_dict: dict[object, Any] = {}
         self.ttl_dict: dict = {}
         self.expiration_heap: list[tuple[float, str]] = []
         self._increment_lock = threading.Lock()
@@ -100,17 +100,17 @@ class InMemoryCache(BaseCache):
         except KeyError:
             return False
 
-    def _cache_keys(self) -> tuple[str, ...]:
-        return tuple(self.cache_dict)  # pyright: ignore[reportUnknownVariableType]  # cache_dict is untyped legacy state
+    def _cache_keys(self) -> tuple[object, ...]:
+        return tuple(self.cache_dict)
 
-    def _remove_key(self, key: str) -> None:
+    def _remove_key(self, key: object) -> None:
         """
         Remove a key from both cache_dict and ttl_dict
         """
         self.cache_dict.pop(key, None)
         self.ttl_dict.pop(key, None)
 
-    def evict_cache(self):
+    def evict_cache(self, *, reserve_space: bool = True) -> None:
         """
         Eviction policy:
         1. First, remove expired items from ttl_dict and cache_dict
@@ -145,7 +145,7 @@ class InMemoryCache(BaseCache):
             heapq.heapify(self.expiration_heap)
 
         # Step 2: Evict if cache is still full
-        while self.expiration_heap and len(self.cache_dict) >= self.max_size_in_memory:
+        while self.expiration_heap and len(self.cache_dict) > self.max_size_in_memory - int(reserve_space):
             expiration_time, key = heapq.heappop(self.expiration_heap)
             # Skip if key was removed or updated
             if self.ttl_dict.get(key) == expiration_time:
@@ -153,7 +153,7 @@ class InMemoryCache(BaseCache):
 
         if not self.expiration_heap:
             untracked_keys: Final = self._cache_keys()
-            overflow: Final = len(untracked_keys) - self.max_size_in_memory + 1
+            overflow: Final = len(untracked_keys) - self.max_size_in_memory + int(reserve_space)
             for untracked_key in untracked_keys[: max(overflow, 0)]:
                 self._remove_key(untracked_key)
 
@@ -181,10 +181,13 @@ class InMemoryCache(BaseCache):
             # Always prune expired/outdated heap roots before inserting.
             # This keeps expiration_heap bounded even when the live cache stays
             # below max_size_in_memory and keys are reinserted after TTL expiry.
-            self.evict_cache()
+            self.evict_cache(reserve_space=key not in self.cache_dict)
             if not self.check_value_size(value):
                 return
 
+            if key not in self.cache_dict and len(self.cache_dict) >= self.max_size_in_memory:
+                self.evict_cache()
+
             self.cache_dict[key] = value
             if self.allow_ttl_override(key):  # if ttl is not set, set it to default ttl
                 if "ttl" in kwargs and kwargs["ttl"] is not None:

litellm/caching/llm_caching_handler.py

@@ -37,7 +37,7 @@ class LLMClientCache(InMemoryCache):
         )
         self.evicted_client_closer = evicted_client_closer or default_evicted_client_closer
 
-    def _remove_key(self, key: str) -> None:
+    def _remove_key(self, key: object) -> None:
         evicted: Final[object] = self.cache_dict.get(key)
         super()._remove_key(key)
         self.evicted_client_closer.schedule(evicted)

litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py

@@ -27,7 +27,7 @@ class LangfuseInMemoryCache(InMemoryCache):
     (https://github.com/BerriAI/litellm/issues/11169).
     """
 
-    def _remove_key(self, key: str) -> None:
+    def _remove_key(self, key: object) -> None:
         from litellm.integrations.langfuse.langfuse import LangFuseLogger
 
         evicted: Final = self.cache_dict.pop(key, None)

litellm/proxy/agent_endpoints/databricks_oauth.py

@@ -212,12 +212,12 @@ class DatabricksAppOAuthTokenCache(InMemoryCache):
 
     def __init__(self) -> None:
         super().__init__(default_ttl=_DEFAULT_TTL_SECONDS)
-        self._locks: dict[str, asyncio.Lock] = {}
+        self._locks: dict[object, asyncio.Lock] = {}
 
     def _get_lock(self, cache_key: str) -> asyncio.Lock:
         return self._locks.setdefault(cache_key, asyncio.Lock())
 
-    def _remove_key(self, key: str) -> None:
+    def _remove_key(self, key: object) -> None:
         # Drop the per-key lock alongside the cached token so ``_locks`` stays
         # bounded by the live key set rather than growing for every key ever seen.
         super()._remove_key(key)

litellm/proxy/common_utils/debug_utils.py

@@ -130,7 +130,7 @@ if os.environ.get("LITELLM_PROFILE", "false").lower() == "true":
         return {"top_50_memory_usage": result}
 
 
-def _prompt_caching_pins(llm_router: "Router | None") -> tuple[Mapping[str, object], Mapping[str, float]]:
+def _prompt_caching_pins(llm_router: "Router | None") -> tuple[Mapping[object, object], Mapping[str, float]]:
     if llm_router is None:
         return {}, {}
     pin_cache: Final = llm_router.prompt_caching_pin_cache.in_memory_cache

tests/unit/caching/test_in_memory_cache.py

@@ -1,21 +1,14 @@
-import asyncio
-import json
 import random
 import sys
 import threading
 import time
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
-from typing import Final
-from unittest.mock import MagicMock, patch
+from typing import Final, Literal
 
-import httpx
 import pytest
-import respx
-from fastapi.testclient import TestClient
-
-from unittest.mock import AsyncMock
 
+from litellm.caching.dual_cache import DualCache
 from litellm.caching.in_memory_cache import InMemoryCache
 
 
@@ -25,6 +18,17 @@ class _SlowInt(int):
         return _SlowInt(int(self) + value)
 
 
+class _FakeClock:
+    def __init__(self) -> None:
+        self.now = 0.0
+
+    def __call__(self) -> float:
+        return self.now
+
+    def advance(self, seconds: float) -> None:
+        self.now += seconds
+
+
 def test_increment_cache_is_atomic_under_thread_concurrency():
     cache = InMemoryCache()
     seed = 1000
@@ -139,6 +143,7 @@ def test_in_memory_cache_ttl():
     assert new_ttl_time is not None
     time.sleep(1)
     cached_obj = in_memory_cache.get_cache(key="new-fake-key")
+    assert cached_obj is None
     new_ttl_time = in_memory_cache.ttl_dict.get("new-fake-key")
     assert new_ttl_time is None
 
@@ -195,6 +200,48 @@ def test_in_memory_cache_max_size_with_ttl():
     assert "key_3" in in_memory_cache.cache_dict
 
 
+@pytest.mark.parametrize("updated_key", ["first", "second"])
+def test_updating_full_cache_preserves_entries_and_expiration(updated_key: Literal["first", "second"]) -> None:
+    clock: Final = _FakeClock()
+    cache: Final = InMemoryCache(max_size_in_memory=2, clock=clock)
+    cache.set_cache("first", "original-a", ttl=10)
+    cache.set_cache("second", "original-b", ttl=30)
+
+    cache.set_cache(updated_key, "updated", ttl=60)
+
+    assert cache.get_cache("first") == ("updated" if updated_key == "first" else "original-a")
+    assert cache.get_cache("second") == ("updated" if updated_key == "second" else "original-b")
+
+    clock.advance(11.0)
+    assert cache.get_cache("first") is None
+    assert cache.get_cache("second") == ("updated" if updated_key == "second" else "original-b")
+
+    clock.advance(20.0)
+    assert cache.get_cache("second") is None
+
+
+@pytest.mark.parametrize(
+    ("updated_key", "expected_values"),
+    [
+        ("first", ("updated", None, "original-c")),
+        ("third", (None, "original-b", "updated")),
+    ],
+)
+def test_updating_cache_after_capacity_reduction_respects_limit(
+    updated_key: Literal["first", "third"], expected_values: tuple[str | None, str | None, str | None]
+) -> None:
+    cache: Final = InMemoryCache(max_size_in_memory=3, clock=_FakeClock())
+    cache.set_cache("first", "original-a", ttl=10)
+    cache.set_cache("second", "original-b", ttl=20)
+    cache.set_cache("third", "original-c", ttl=30)
+    DualCache(in_memory_cache=cache).update_in_memory_max_size(2)
+
+    cache.set_cache(updated_key, "updated", ttl=60)
+
+    assert len(cache.cache_dict) == 2
+    assert tuple(cache.get_cache(key) for key in ("first", "second", "third")) == expected_values
+
+
 def test_in_memory_cache_expired_items_evicted_first():
     """
     Test that expired items are evicted before non-expired items when cache is full.
@@ -234,14 +281,9 @@ def test_in_memory_cache_eviction_order():
     in_memory_cache = InMemoryCache(max_size_in_memory=2)
 
     # Add items with different TTLs
-    now = time.time()
-    in_memory_cache.set_cache(
-        key="early_expire", value="value_1", ttl=100
-    )  # expires in 100 seconds
+    in_memory_cache.set_cache(key="early_expire", value="value_1", ttl=100)  # expires in 100 seconds
     time.sleep(0.01)
-    in_memory_cache.set_cache(
-        key="late_expire", value="value_2", ttl=200
-    )  # expires in 200 seconds
+    in_memory_cache.set_cache(key="late_expire", value="value_2", ttl=200)  # expires in 200 seconds
 
     # Verify TTL order
     early_ttl = in_memory_cache.ttl_dict["early_expire"]
@@ -337,6 +379,17 @@ def test_in_memory_cache_stays_bounded_when_entries_have_no_usable_expiry() -> N
     assert cache.get_cache("fresh") == "response"
 
 
+def test_updating_full_cache_with_untracked_expiry_preserves_other_entries() -> None:
+    cache: Final = InMemoryCache(max_size_in_memory=2, clock=_FakeClock())
+    cache.set_cache("first", "original-a", ttl="NaN")
+    cache.set_cache("second", "original-b", ttl="NaN")
+
+    cache.set_cache("second", "updated", ttl=60)
+
+    assert cache.get_cache("first") == "original-a"
+    assert cache.get_cache("second") == "updated"
+
+
 def test_in_memory_cache_prunes_expired_heap_entries_below_capacity():
     """
     Re-inserting expired keys below capacity should not grow expiration_heap