From 592e41801244b3c6715c93b874235da515a3ed53 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 2 Apr 2026 11:29:19 +0530 Subject: [PATCH] fix(guardrails/purview): raise_for_status on HTTP errors, cap scope cache, reuse executor --- .../guardrail_hooks/microsoft_purview/base.py | 12 ++++++++++-- .../guardrail_hooks/microsoft_purview/purview_dlp.py | 8 ++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index c0b2021728a..fd9cc4b86d5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -1,6 +1,7 @@ import time import uuid -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Dict, List, MutableMapping, Optional, Tuple from litellm._logging import verbose_proxy_logger from litellm.llms.custom_httpx.http_handler import ( @@ -56,7 +57,9 @@ class PurviewGuardrailBase: self._token_cache: Optional[Tuple[str, float]] = None # Protection scope cache: user_id -> (etag, scope_response, fetched_at) - self._scope_cache: Dict[str, Tuple[str, Dict, float]] = {} + # Capped at 1000 entries (LRU eviction) to avoid unbounded growth. + self._scope_cache: MutableMapping[str, Tuple[str, Dict, float]] = OrderedDict() + self._scope_cache_maxsize = 1000 # ------------------------------------------------------------------ # OAuth2 token management @@ -80,6 +83,7 @@ class PurviewGuardrailBase: data=data, headers={"Content-Type": "application/x-www-form-urlencoded"}, ) + response.raise_for_status() token_data = response.json() access_token = token_data["access_token"] expires_in = int(token_data.get("expires_in", 3599)) @@ -116,6 +120,7 @@ class PurviewGuardrailBase: response = await self.async_handler.post( url=url, headers=headers, json=json_body ) + response.raise_for_status() response_json: Dict[str, Any] = response.json() response_headers = dict(response.headers) verbose_proxy_logger.debug("Purview Graph response: %s", response_json) @@ -157,6 +162,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. + while len(self._scope_cache) > self._scope_cache_maxsize: + self._scope_cache.popitem(last=False) # type: ignore[attr-defined] return etag, response_json # ------------------------------------------------------------------ diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py index f6fcae77ce2..1bbbf531f55 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py @@ -8,6 +8,7 @@ Supports three modes: """ import asyncio +from concurrent.futures import ThreadPoolExecutor from datetime import datetime from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union @@ -75,6 +76,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): ) self._logging_only = logging_only self.guardrail_provider = "microsoft_purview" + self._executor = ThreadPoolExecutor(max_workers=1) verbose_proxy_logger.info( "Initialized Microsoft Purview DLP Guardrail: %s (logging_only=%s)", guardrail_name, @@ -229,7 +231,6 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): self, kwargs: dict, result: Any, call_type: str ) -> Tuple[dict, Any]: """Sync wrapper for async_logging_hook (follows Presidio pattern).""" - from concurrent.futures import ThreadPoolExecutor def run_in_new_loop() -> Tuple[dict, Any]: new_loop = asyncio.new_event_loop() @@ -246,9 +247,8 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): try: _ = asyncio.get_running_loop() - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(run_in_new_loop) - return future.result() + future = self._executor.submit(run_in_new_loop) + return future.result() except RuntimeError: return run_in_new_loop()