diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index 9e96cda974e..c1b46ea4f7f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -1,7 +1,7 @@ import time import uuid from collections import OrderedDict -from typing import TYPE_CHECKING, Any, Dict, List, MutableMapping, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -61,7 +61,9 @@ class PurviewGuardrailBase: # Protection scope cache: user_id -> (etag, scope_response, fetched_at) # Capped at 1000 entries (LRU eviction) to avoid unbounded growth. - self._scope_cache: MutableMapping[str, Tuple[str, Dict, float]] = OrderedDict() + self._scope_cache: OrderedDict[str, Tuple[str, Dict[str, Any], float]] = ( + OrderedDict() + ) self._scope_cache_maxsize = 1000 # ------------------------------------------------------------------ @@ -145,6 +147,7 @@ class PurviewGuardrailBase: now = time.time() if cached and (now - cached[2]) < SCOPE_CACHE_TTL_SECONDS: + self._scope_cache.move_to_end(user_id) return cached[0], cached[1] url = ( @@ -165,9 +168,9 @@ class PurviewGuardrailBase: etag = response_headers.get("etag", response_headers.get("ETag", "")) self._scope_cache[user_id] = (etag, response_json, now) - # Evict oldest entry when cache exceeds max size. + # Evict least-recently-used entry when cache exceeds max size. while len(self._scope_cache) > self._scope_cache_maxsize: - self._scope_cache.popitem(last=False) # type: ignore[attr-defined] + self._scope_cache.popitem(last=False) return etag, response_json # ------------------------------------------------------------------ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index c4ed94b75f5..f10b93fa85a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -531,6 +531,42 @@ class TestScopeCaching: assert etag2 == "scope-etag" assert mock_post.call_count == 1 + @pytest.mark.asyncio + async def test_scope_cache_lru_keeps_hot_user_on_eviction(self): + """Frequently accessed users should not be evicted before cold entries.""" + guardrail = _make_guardrail() + guardrail._scope_cache_maxsize = 3 + + scope_payload = ( + { + "value": [ + {"activities": "uploadText", "executionMode": "evaluateInline"} + ] + }, + {"ETag": "scope-etag"}, + ) + + with patch.object( + guardrail, "_graph_post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = scope_payload + + await guardrail._compute_protection_scopes("user-a") + await guardrail._compute_protection_scopes("user-b") + await guardrail._compute_protection_scopes("user-c") + assert mock_post.call_count == 3 + + await guardrail._compute_protection_scopes("user-a") + assert mock_post.call_count == 3 + + await guardrail._compute_protection_scopes("user-d") + assert mock_post.call_count == 4 + + await guardrail._compute_protection_scopes("user-a") + assert mock_post.call_count == 4 + assert "user-a" in guardrail._scope_cache + assert "user-b" not in guardrail._scope_cache + @pytest.mark.asyncio async def test_scope_invalidated_on_modified(self): guardrail = _make_guardrail()