scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 134

fix(memory): share activation hints through coordination cache

Browse source
This commit is contained in:
moe-berri 2026-09-12 03:00:29 -07:00
parent 7297247aab
commit 1e7963abba
2 changed files with 59 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/proxy/memory/policy.py
View file

@ -8,6 +8,8 @@ from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.config_sync_pubsub import coordination_redis_cache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient
from litellm.repositories.table_repositories import MemoryPolicyRepository, MemoryPreferenceRepository
from litellm.types.memory_v2 import MemoryPolicy, MemoryScope, MemoryStatus
@ -25,24 +27,34 @@ async def gateway_memory_is_configured(prisma_client: object, cache: DualCache)
cached: Final = await cache.async_get_cache(key=_CONFIGURED_CACHE_KEY)
if cached is True:
return True
redis_cache: Final = cache.redis_cache or coordination_redis_cache()
# An empty local view reads shared Redis through DualCache's existing
# circuit-breaker/error handling, falling back to the primary on a miss.
shared_cache: Final = DualCache(redis_cache=redis_cache) if redis_cache is not None else None
if cached is False:
if cache.redis_cache is None:
if shared_cache is None:
return False
# A backend mutation evicts Redis, but another worker can still hold
# a negative local hint (Redis Cluster may not support pub/sub).
shared: Final = await cache.redis_cache.async_get_cache(key=_CONFIGURED_CACHE_KEY)
shared: Final = await shared_cache.async_get_cache(key=_CONFIGURED_CACHE_KEY)
if shared is False:
return False
rows: Final = await MemoryPolicyRepository(memory_primary_client(prisma_client)).table.find_many(take=1)
configured: Final = bool(rows)
await cache.async_set_cache(key=_CONFIGURED_CACHE_KEY, value=configured, ttl=30)
if shared_cache is not None and cache.redis_cache is None:
await shared_cache.async_set_cache(key=_CONFIGURED_CACHE_KEY, value=configured, ttl=30)
return configured
async def invalidate_memory_configuration() -> None:
from litellm.proxy.proxy_server import user_api_key_cache
await evict_and_broadcast(cache_keys=(_CONFIGURED_CACHE_KEY,), user_api_key_cache=user_api_key_cache)
cache: Final = UserApiKeyCache(
in_memory_cache=user_api_key_cache.in_memory_cache,
redis_cache=user_api_key_cache.redis_cache or coordination_redis_cache(),
)
await evict_and_broadcast(cache_keys=(_CONFIGURED_CACHE_KEY,), user_api_key_cache=cache)
def memory_primary_client(prisma_client: object) -> WriterPinnedClient:

57
tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py
View file

@ -351,8 +351,10 @@ async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client
)
auth = UserAPIKeyAuth(token="a" * 64, user_id="owner", team_id="team", project_id="project", org_id="org")
before = get_request_stash()
with patch.multiple( # test-quality-ok: Inject database/cache/ASGI provider edges; run real memory and limiter code.
proxy_server, app=provider, prisma_client=prisma_edge, user_api_key_cache=DualCache()
with (
patch.multiple( # test-quality-ok: Inject database/cache/ASGI provider edges; run real memory and limiter code.
proxy_server, app=provider, prisma_client=prisma_edge, user_api_key_cache=DualCache()
)
):
prepared = await prepare_gateway_memory(original, request, auth, "anthropic_messages")
release.set()
@ -368,7 +370,10 @@ async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client
@pytest.mark.asyncio
async def test_backend_activation_invalidates_a_gateway_negative_hint_without_pubsub(prisma_edge: MagicMock) -> None:
@pytest.mark.parametrize("share_auth_cache", [False, True])
async def test_backend_activation_invalidates_a_gateway_negative_hint_without_pubsub(
prisma_edge: MagicMock, share_auth_cache: bool
) -> None:
from unittest.mock import patch
from litellm.caching.caching import DualCache
@ -390,17 +395,43 @@ async def test_backend_activation_invalidates_a_gateway_negative_hint_without_pu
async_set_cache=AsyncMock(side_effect=set_value),
async_delete_cache=AsyncMock(side_effect=delete),
)
gateway_cache = DualCache(redis_cache=redis)
backend_cache = DualCache(redis_cache=redis)
gateway_cache = DualCache(redis_cache=redis if share_auth_cache else None)
backend_cache = DualCache(redis_cache=redis if share_auth_cache else None)
policies = prisma_edge.db.litellm_memorypolicy.find_many
policies.return_value = []
assert not await gateway_memory_is_configured(prisma_edge, gateway_cache)
assert not await gateway_memory_is_configured(prisma_edge, gateway_cache)
policies.assert_awaited_once()
policies.return_value = [_POLICY]
with patch.multiple( # test-quality-ok: Inject external worker caches and Redis; run real invalidation.
"litellm.proxy.proxy_server", user_api_key_cache=backend_cache, redis_usage_cache=None
"litellm.proxy.proxy_server", user_api_key_cache=backend_cache, redis_usage_cache=redis
):
policies.return_value = []
assert not await gateway_memory_is_configured(prisma_edge, gateway_cache)
assert not await gateway_memory_is_configured(prisma_edge, gateway_cache)
policies.assert_awaited_once()
policies.return_value = [_POLICY]
await invalidate_memory_configuration()
assert await gateway_memory_is_configured(prisma_edge, gateway_cache)
assert policies.await_count == 2
assert await gateway_memory_is_configured(prisma_edge, gateway_cache)
assert policies.await_count == 2
@pytest.mark.asyncio
async def test_redis_circuit_breaker_falls_back_to_primary_configuration(prisma_edge: MagicMock) -> None:
from unittest.mock import patch
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
from litellm.proxy.memory.policy import gateway_memory_is_configured, invalidate_memory_configuration
redis = MagicMock(
async_get_cache=AsyncMock(side_effect=RedisCircuitBreakerOpenError("open")),
async_set_cache=AsyncMock(side_effect=RedisCircuitBreakerOpenError("open")),
async_delete_cache=AsyncMock(side_effect=RedisCircuitBreakerOpenError("open")),
)
cache = DualCache()
prisma_edge.db.litellm_memorypolicy.find_many.return_value = []
with patch.multiple( # test-quality-ok: Inject external Redis failure and local worker cache; exercise real fallback.
"litellm.proxy.proxy_server", user_api_key_cache=cache, redis_usage_cache=redis
):
assert not await gateway_memory_is_configured(prisma_edge, cache)
prisma_edge.db.litellm_memorypolicy.find_many.return_value = [_POLICY]
assert await gateway_memory_is_configured(prisma_edge, cache)
await invalidate_memory_configuration()
assert prisma_edge.db.litellm_memorypolicy.find_many.await_count == 2
redis.async_get_cache.assert_awaited_once()