From 87715dce1ec51ea4ded0a3658e30cf0d2d26c166 Mon Sep 17 00:00:00 2001 From: Mark Chu Date: Thu, 26 Mar 2026 00:52:49 +0800 Subject: [PATCH] fix(proxy/presidio): stabilize output_parse_pii and user_config routing --- .../guardrails/guardrail_hooks/presidio.py | 604 ++++++-- .../unified_guardrail/unified_guardrail.py | 29 +- .../guardrails/guardrail_initializers.py | 36 +- litellm/proxy/route_llm_request.py | 260 +++- litellm/proxy/utils.py | 65 +- litellm/types/guardrails.py | 8 + .../guardrail_hooks/test_presidio.py | 1334 ++++++++++++++++- .../test_unified_guardrail.py | 152 +- tests/test_litellm/proxy/test_proxy_utils.py | 166 ++ .../proxy/test_route_llm_request.py | 662 +++++++- 10 files changed, 3099 insertions(+), 217 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 0f4ebbd4880..b2d4ed28117 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -11,8 +11,9 @@ import asyncio import json import threading -from contextlib import asynccontextmanager +import re from datetime import datetime +from contextlib import asynccontextmanager from typing import ( TYPE_CHECKING, Any, @@ -29,6 +30,7 @@ from typing import ( import aiohttp import litellm # noqa: E401 +from litellm._uuid import uuid from litellm import get_secret from litellm._logging import verbose_proxy_logger from litellm.types.utils import GenericGuardrailAPIInputs @@ -46,6 +48,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import ( GuardrailEventHooks, LitellmParams, + Mode, PiiAction, PiiEntityType, PresidioPerRequestConfig, @@ -63,6 +66,204 @@ from litellm.utils import ( ) +# Max trailing alphabetic chars to allow when matching corrupted uuid-style placeholders +# (e.g. LLM outputs "...fa9den" instead of "...fa9d"). Tune down (e.g. 3–5) +# if LLM rarely adds more than a few chars to reduce false matches. +_MAX_TRAILING_CHARS_CORRUPTED_PLACEHOLDER = 15 +# Matches placeholders generated in anonymize_text as "" + str(uuid.uuid4()), +# i.e. no underscore and a full UUID4 suffix. str(uuid.uuid4()) currently produces +# lowercase hex only; uppercase [A-F] is accepted here for robustness if the UUID +# generator changes or a caller normalizes case differently. Keep this in sync with +# the placeholder generator in anonymize_text. +_UUID_SUFFIX_PLACEHOLDER_RE = re.compile( + r"^<[^>]+>[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$" +) + + +def _ensure_event_hook_includes_post_call( + event_hook: Optional[Union[GuardrailEventHooks, List[Any], Mode, str]], + include_pre_call_on_none: bool = False, +) -> Optional[Union[Mode, List[str], str]]: + post_call = GuardrailEventHooks.post_call.value + + def _hook_value(hook: Any) -> str: + if isinstance(hook, GuardrailEventHooks): + return hook.value + return str(hook) + + def _normalize_hook_list(hooks: List[Any]) -> List[str]: + normalized: List[str] = [] + for hook in hooks: + hook_value = _hook_value(hook) + if hook_value not in normalized: + normalized.append(hook_value) + if post_call not in normalized: + normalized.append(post_call) + return normalized + + if event_hook is None: + if include_pre_call_on_none: + return [GuardrailEventHooks.pre_call.value, post_call] + return post_call + if isinstance(event_hook, Mode): + mode_copy = event_hook.model_copy(deep=True) + mode_copy.tags = { + tag: _normalize_hook_list(values if isinstance(values, list) else [values]) + for tag, values in mode_copy.tags.items() + } + # Preserve tag-only Mode semantics: if default is None, untagged requests do + # not run the guardrail. We only expand post_call for explicitly configured + # tag/default hooks, rather than turning an absent default into a new one. + if mode_copy.default is not None: + mode_copy.default = _normalize_hook_list( + mode_copy.default + if isinstance(mode_copy.default, list) + else [mode_copy.default] + ) + return mode_copy + if isinstance(event_hook, list): + return _normalize_hook_list(event_hook) + + hook_value = _hook_value(event_hook) + if hook_value == post_call: + return hook_value + return [hook_value, post_call] + + +def _get_corrupted_placeholder_pattern(key: str) -> re.Pattern: + return re.compile( + re.escape(key) + + rf"[a-zA-Z0-9]{{1,{_MAX_TRAILING_CHARS_CORRUPTED_PLACEHOLDER}}}" + + r"(?![a-zA-Z0-9])" + ) + + +def _replace_pii_tokens_in_text(text: str, pii_tokens: Dict[str, str]) -> str: + """ + Replace PII placeholders in text with original values. Handles LLM corruption + of uuid-style placeholders (e.g. uuid becomes uiden or + uuid2) by matching key + trailing alphanumeric chars. + """ + if not pii_tokens: + return text + corrupted_placeholder_patterns = { + key: _get_corrupted_placeholder_pattern(key) + for key in pii_tokens + if _UUID_SUFFIX_PLACEHOLDER_RE.match(key) is not None + } + consumed_tokens = set() + # Do regex pass first for uuid-style keys so "key+trailing" is replaced in one go. + # If we did exact replace first, we'd replace the key and leave trailing chars (e.g. "Jane Doeen"). + for key, pattern in corrupted_placeholder_patterns.items(): + text, replacement_count = pattern.subn(pii_tokens[key], text) + if replacement_count > 0: + consumed_tokens.add(key) + # Replace longer keys first so "uuid" is replaced before "" + for key, value in sorted(pii_tokens.items(), key=lambda x: -len(x[0])): + if key in text: + consumed_tokens.add(key) + text = text.replace(key, value) + # Fallback for tokens truncated by max_tokens: if the end of the text is a + # sufficiently long prefix of a placeholder, replace that suffix with the + # original value so standard and Anthropic response paths behave the same. + for token, original_text in pii_tokens.items(): + if token in consumed_tokens: + continue + if token in text: + continue + min_overlap = max(1, min(20, len(token) // 2)) + # Re-capture the latest text on each outer iteration. If a prior token + # replacement shortened or lengthened the response, the next token's + # suffix scan must use the updated string length and tail positions. + current_text = text + current_len = len(current_text) + scan_start = max(0, current_len - len(token)) + for i in range(scan_start, current_len): + sub = current_text[i:] + if token.startswith(sub) and len(sub) >= min_overlap: + # Safe to break: at most one suffix of `current_text` can be the + # active truncated match for this token, and the next outer-loop + # iteration will re-snapshot `text` after this replacement. + text = current_text[:i] + original_text + break + return text + + +def _score_analyze_span_for_unmask_pairing( + span: Dict[str, Any] +) -> Tuple[int, float, int]: + """Score span quality for replace-token pairing: prefer longer, then higher-score, then earlier span.""" + start = cast(int, span["start"]) + end = cast(int, span["end"]) + raw_score = span.get("score") + score = ( + float(raw_score) + if isinstance(raw_score, (int, float)) and not isinstance(raw_score, bool) + else -1.0 + ) + return (end - start, score, -start) + + +def _to_int_offset(value: Any) -> Optional[int]: + if isinstance(value, int) and not isinstance(value, bool): + return value + if isinstance(value, float) and value.is_integer(): + return int(value) + return None + + +def _dedupe_overlapping_analyze_spans_for_unmask( + analyze_spans_sorted: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """ + Collapse exact-duplicate and overlapping analyze spans for output_parse_pii mapping. + Overlap clusters keep the most representative span to reduce mismatched pairing + when analyzer returns nested/overlapping entities for the same text region. + """ + if not analyze_spans_sorted: + return [] + + deduped_exact_spans: List[Dict[str, Any]] = [] + seen_spans = set() + for span in analyze_spans_sorted: + span_key = (cast(int, span["start"]), cast(int, span["end"])) + if span_key in seen_spans: + continue + seen_spans.add(span_key) + deduped_exact_spans.append(span) + + collapsed_spans: List[Dict[str, Any]] = [] + overlap_cluster: List[Dict[str, Any]] = [] + cluster_end = -1 + + for span in deduped_exact_spans: + span_start = cast(int, span["start"]) + span_end = cast(int, span["end"]) + if not overlap_cluster: + overlap_cluster = [span] + cluster_end = span_end + continue + + # Overlap only when current start is inside the previous cluster. + if span_start < cluster_end: + overlap_cluster.append(span) + cluster_end = max(cluster_end, span_end) + continue + + collapsed_spans.append( + max(overlap_cluster, key=_score_analyze_span_for_unmask_pairing) + ) + overlap_cluster = [span] + cluster_end = span_end + + if overlap_cluster: + collapsed_spans.append( + max(overlap_cluster, key=_score_analyze_span_for_unmask_pairing) + ) + + return collapsed_spans + + class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_cache = None ad_hoc_recognizers = None @@ -93,9 +294,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): kwargs["event_hook"] = GuardrailEventHooks.logging_only super().__init__(**kwargs) self.guardrail_provider = "presidio" - self.pii_tokens: dict = ( - {} - ) # mapping of PII token to original text - only used with Presidio `replace` operation + # Deprecated request state. Keep attribute for backward compatibility with + # tests/instrumentation, but request processing uses request-local mappings. + self.pii_tokens: dict = {} self.mock_redacted_text = mock_redacted_text self.output_parse_pii = output_parse_pii or False self.apply_to_output = apply_to_output @@ -104,15 +305,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # also run on post_call to unmask/mask the response. Expand the event_hook # so should_run_guardrail returns True for both pre_call and post_call. if (self.output_parse_pii or self.apply_to_output) and not logging_only: - current_hook = self.event_hook - if isinstance(current_hook, str) and current_hook != "post_call": - self.event_hook = cast( - List[GuardrailEventHooks], [current_hook, "post_call"] - ) - elif isinstance(current_hook, list) and "post_call" not in current_hook: - self.event_hook = cast( - List[GuardrailEventHooks], current_hook + ["post_call"] - ) + self.event_hook = _ensure_event_hook_includes_post_call( + self.event_hook, + include_pre_call_on_none=self.output_parse_pii, + ) self.pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = ( pii_entities_config or {} ) @@ -439,7 +635,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): analyze_results: Any, output_parse_pii: bool, masked_entity_count: Dict[str, int], - request_data: Optional[Dict] = None, + pii_tokens: Optional[Dict[str, str]] = None, ) -> str: """ Send analysis results to the Presidio anonymizer endpoint to get redacted text @@ -485,62 +681,200 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): redacted_text = await response.json() - new_text = text + if output_parse_pii and pii_tokens is None: + verbose_proxy_logger.warning( + "Presidio output_parse_pii enabled but pii_tokens is None; " + "token mappings will be discarded and response unmasking may fail." + ) + token_store = pii_tokens if pii_tokens is not None else {} if redacted_text is not None: verbose_proxy_logger.debug("redacted_text: %s", redacted_text) - # Process items in reverse order by start position so that - # replacing later spans first does not shift earlier coordinates. - for item in sorted( - redacted_text["items"], key=lambda x: x["start"], reverse=True - ): - start = item["start"] - end = item["end"] - replacement = item["text"] # replacement token - if item["operator"] == "replace" and output_parse_pii is True: - if request_data is None: - verbose_proxy_logger.warning( - "Presidio anonymize_text called without request_data — " - "PII tokens cannot be stored per-request. " - "This may indicate a missing caller update." - ) - request_data = {} - # Store pii_tokens in metadata to avoid leaking to LLM providers. - # Providers like Anthropic reject unknown top-level fields. - if not request_data.get("metadata"): - request_data["metadata"] = {} - if "pii_tokens" not in request_data["metadata"]: - request_data["metadata"]["pii_tokens"] = {} - pii_tokens = request_data["metadata"]["pii_tokens"] - - # Append a sequential number to make each token unique - # per request, so unmasking maps back to the correct - # original value. Format: , - # This is LLM-friendly and degrades gracefully if the - # LLM doesn't echo the token verbatim. - seq = len(pii_tokens) + 1 - if replacement.endswith(">"): - replacement = f"{replacement[:-1]}_{seq}>" - else: - replacement = f"{replacement}_{seq}" - - # Use ORIGINAL text (not new_text) since start/end - # reference the original text's coordinates. - pii_tokens[replacement] = text[start:end] - - new_text = new_text[:start] + replacement + new_text[end:] + items = redacted_text["items"] + for item in items: entity_type = item.get("entity_type", None) if entity_type is not None: masked_entity_count[entity_type] = ( masked_entity_count.get(entity_type, 0) + 1 ) - # When output_parse_pii is True, new_text contains sequentially - # numbered tokens (e.g. ) that match the keys - # in pii_tokens. Returning redacted_text["text"] (Presidio's - # original output) would send un-numbered tokens to the LLM, - # making unmasking impossible. - # When output_parse_pii is False, new_text == redacted_text["text"] - # because no suffix is appended. - return new_text + + # output_parse_pii is designed for replace-and-unmask flow. + # If Presidio returns non-replace operators, use Presidio's text as-is + # to avoid incorrect offset arithmetic on mixed operator outputs. + if output_parse_pii: + replace_items = [i for i in items if i.get("operator") == "replace"] + non_replace_items = [ + i for i in items if i.get("operator") != "replace" + ] + if non_replace_items: + replace_entity_types = [ + str(entity_type) + for entity_type in ( + i.get("entity_type") for i in replace_items + ) + if entity_type is not None + ] + verbose_proxy_logger.warning( + "Presidio output_parse_pii fallback: detected non-replace " + "operators (%s); %d replace-operator entities (%s) will also " + "NOT be unmasked; returning redacted_text without unmask mapping.", + sorted({str(i.get("operator")) for i in non_replace_items}), + len(replace_items), + replace_entity_types, + ) + return redacted_text["text"] + + if not isinstance(analyze_results, list): + verbose_proxy_logger.warning( + "Presidio output_parse_pii fallback: analyze_results is not a list; " + "returning redacted_text without unmask mapping." + ) + return redacted_text["text"] + + analyze_spans = [] + for r in analyze_results: + if not isinstance(r, dict): + continue + start_v = _to_int_offset(r.get("start")) + end_v = _to_int_offset(r.get("end")) + if start_v is None or end_v is None: + continue + analyze_spans.append( + { + "start": cast(int, start_v), + "end": cast(int, end_v), + "entity_type": r.get("entity_type"), + "score": r.get("score"), + } + ) + # `replace_items` starts are in anonymized-output coordinates, while + # `analyze_spans` starts are in original-input coordinates. We still + # sort both lists left-to-right and pair them positionally, relying on + # Presidio preserving entity order across analyze/anonymize results. + # If that ordering ever changes, this logic would need stronger + # placeholder-to-span matching than simple positional pairing. + replace_items_sorted = sorted( + replace_items, key=lambda i: i["start"] + ) + analyze_spans_sorted = sorted( + analyze_spans, key=lambda r: r["start"] + ) + analyze_spans_sorted = _dedupe_overlapping_analyze_spans_for_unmask( + analyze_spans_sorted + ) + + if len(replace_items_sorted) != len(analyze_spans_sorted): + # Best-effort mapping: preserve left-to-right order and prefer + # matching entity_type before falling back to next available span. + if len(analyze_spans_sorted) >= len(replace_items_sorted): + best_effort_spans: List[Dict[str, Any]] = [] + used_span_indices = set() + last_used_index = -1 + for replace_item in replace_items_sorted: + replace_entity = replace_item.get("entity_type") + replace_entity_str = ( + str( + getattr(replace_entity, "value", replace_entity) + ) + if replace_entity is not None + else None + ) + selected_index: Optional[int] = None + + for idx, span in enumerate(analyze_spans_sorted): + if ( + idx in used_span_indices + or idx <= last_used_index + ): + continue + span_entity = span.get("entity_type") + span_entity_str = ( + str(getattr(span_entity, "value", span_entity)) + if span_entity is not None + else None + ) + if span_entity_str == replace_entity_str: + selected_index = idx + break + + if selected_index is None: + for idx, _ in enumerate(analyze_spans_sorted): + if ( + idx in used_span_indices + or idx <= last_used_index + ): + continue + selected_index = idx + break + + if selected_index is None: + break + + used_span_indices.add(selected_index) + last_used_index = selected_index + best_effort_spans.append( + analyze_spans_sorted[selected_index] + ) + + if len(best_effort_spans) == len(replace_items_sorted): + verbose_proxy_logger.warning( + "Presidio output_parse_pii best-effort mapping: replace item count (%s) " + "does not match analyze span count (%s); using ordered span pairing.", + len(replace_items_sorted), + len(analyze_spans_sorted), + ) + analyze_spans_sorted = best_effort_spans + else: + verbose_proxy_logger.warning( + "Presidio output_parse_pii fallback: replace item count (%s) " + "does not match analyze span count (%s); returning redacted_text.", + len(replace_items_sorted), + len(analyze_spans_sorted), + ) + return redacted_text["text"] + else: + verbose_proxy_logger.warning( + "Presidio output_parse_pii fallback: replace item count (%s) " + "does not match analyze span count (%s); returning redacted_text.", + len(replace_items_sorted), + len(analyze_spans_sorted), + ) + return redacted_text["text"] + + # Build unique placeholders in original-text coordinates (from + # analyze results) so replacement offsets stay stable. + new_text = text + # Both lists are already sorted left-to-right in their own + # coordinate systems (analyze: original input, anonymize: + # anonymized output). Positional zipping is still valid + # because Presidio preserves entity order across both + # endpoints even when anonymized-output offsets shift. + for analyze_span, replace_item in zip( + reversed(analyze_spans_sorted), reversed(replace_items_sorted) + ): + start = cast(int, analyze_span["start"]) + end = cast(int, analyze_span["end"]) + if not (isinstance(start, int) and isinstance(end, int)): + raise RuntimeError( + "Presidio output_parse_pii: unexpected non-int span offsets " + f"start={start!r} end={end!r}" + ) + replacement = cast(str, replace_item["text"]) + replacement = f"{replacement}{str(uuid.uuid4())}" + token_store[replacement] = text[start:end] + verbose_proxy_logger.debug( + "Presidio output_parse_pii pair: placeholder=%s entity=%s original=%r", + replacement, + analyze_span.get("entity_type"), + text[start:end], + ) + # Safe to splice against `new_text` using original-text offsets: + # upstream span dedupe removes overlaps, and right-to-left + # replacement keeps positions to the left of `start` stable. + new_text = new_text[:start] + replacement + new_text[end:] + return new_text + + # output_parse_pii disabled: return Presidio's redacted text directly. + return redacted_text["text"] else: raise Exception("Invalid anonymizer response: received None") except Exception as e: @@ -634,6 +968,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): output_parse_pii: bool, presidio_config: Optional[PresidioPerRequestConfig], request_data: dict, + pii_tokens: Optional[Dict[str, str]] = None, ) -> str: """ Calls Presidio Analyze + Anonymize endpoints for PII Analysis + Masking @@ -674,7 +1009,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): analyze_results=analyze_results, output_parse_pii=output_parse_pii, masked_entity_count=masked_entity_count, - request_data=request_data, + pii_tokens=pii_tokens, ) return anonymized_text return redacted_text["text"] @@ -722,6 +1057,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """ try: + request_pii_tokens: Dict[str, str] = {} content_safety = data.get("content_safety", None) verbose_proxy_logger.debug("content_safety: %s", content_safety) presidio_config = self.get_presidio_settings_from_request_data(data) @@ -744,6 +1080,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): output_parse_pii=self.output_parse_pii, presidio_config=presidio_config, request_data=data, + pii_tokens=request_pii_tokens, ) ) task_mappings.append( @@ -760,6 +1097,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): output_parse_pii=self.output_parse_pii, presidio_config=presidio_config, request_data=data, + pii_tokens=request_pii_tokens, ) ) task_mappings.append((msg_idx, int(content_idx))) @@ -785,6 +1123,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): f"Presidio PII Masking: Redacted pii message: {data['messages']}" ) data["messages"] = messages + # Store pii_tokens in request data so post_call can unmask using the same + # request's mappings (guardrail instance is shared across requests). + if self.output_parse_pii and request_pii_tokens: + data.setdefault("_presidio_pii_tokens", {})[self.guardrail_name] = dict( + request_pii_tokens + ) return data except Exception as e: raise e @@ -921,18 +1265,36 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if self.output_parse_pii is False and litellm.output_parse_pii is False: return response - if isinstance(response, ModelResponse) and not isinstance( - response.choices[0], StreamingChoices - ): # /chat/completions requests - await self._process_response_for_pii( - response=response, - request_data=data, - mode="unmask", - ) - elif self._is_anthropic_message_response(response): - await self._process_anthropic_response_for_pii( - response=cast(dict, response), request_data=data, mode="unmask" - ) + # Use only request-scoped pii_tokens; do not fall back to self.pii_tokens + # or we may use stale mappings from a previous request (e.g. wrong name). + pii_tokens = data.get("_presidio_pii_tokens", {}).get(self.guardrail_name, {}) + if not pii_tokens: + return response + + attempted_non_streaming_unmask = False + defer_cleanup_to_streaming_hook = isinstance(response, ModelResponseStream) or ( + isinstance(response, ModelResponse) + and len(response.choices) > 0 + and isinstance(response.choices[0], StreamingChoices) + ) + try: + if isinstance(response, ModelResponse) and not isinstance( + response.choices[0], StreamingChoices + ): # /chat/completions requests + attempted_non_streaming_unmask = True + await self._process_response_for_pii( + response=response, + request_data=data, + mode="unmask", + ) + elif self._is_anthropic_message_response(response): + attempted_non_streaming_unmask = True + await self._process_anthropic_response_for_pii( + response=cast(dict, response), request_data=data, mode="unmask" + ) + finally: + if attempted_non_streaming_unmask or not defer_cleanup_to_streaming_hook: + self._clear_request_scoped_pii_tokens(data) return response @staticmethod @@ -947,19 +1309,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ``min(20, len(token) // 2)`` to reduce the risk of false positives when multiple tokens share a common prefix. """ - for token, original_text in pii_tokens.items(): - if token in text: - text = text.replace(token, original_text) - else: - # FALLBACK: Handle truncated tokens (token cut off by max_tokens) - # Only check at the very end of the text. - min_overlap = min(20, len(token) // 2) - for i in range(max(0, len(text) - len(token)), len(text)): - sub = text[i:] - if token.startswith(sub) and len(sub) >= min_overlap: - text = text[:i] + original_text - break - return text + return _replace_pii_tokens_in_text(text, pii_tokens) @staticmethod def _is_anthropic_message_response(response: Any) -> bool: @@ -980,11 +1330,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Process an Anthropic native message dict for PII masking/unmasking. Handles content blocks with type == "text". """ - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) + pii_tokens = ( + request_data.get("_presidio_pii_tokens", {}).get(self.guardrail_name, {}) + if request_data + else {} + ) if not pii_tokens and mode == "unmask": verbose_proxy_logger.debug( - "No pii_tokens in metadata for Anthropic response unmask" + "No pii_tokens found in request_data — nothing to unmask (anthropic response)" ) presidio_config = self.get_presidio_settings_from_request_data( request_data or {} @@ -1022,11 +1375,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Helper to recursively process a ModelResponse for PII. Handles all choices and tool calls. """ - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) + pii_tokens = ( + request_data.get("_presidio_pii_tokens", {}).get(self.guardrail_name, {}) + if request_data + else {} + ) if not pii_tokens and mode == "unmask": verbose_proxy_logger.debug( - "No pii_tokens found in request_data['metadata'] — nothing to unmask" + "No pii_tokens found in request_data — nothing to unmask" ) presidio_config = self.get_presidio_settings_from_request_data( request_data or {} @@ -1041,7 +1397,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): content = getattr(message, "content", None) if isinstance(content, str): if mode == "unmask": - message.content = self._unmask_pii_text(content, pii_tokens) + message.content = _replace_pii_tokens_in_text(content, pii_tokens) elif mode == "mask": message.content = await self.check_pii( text=content, @@ -1057,7 +1413,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if text_value is None: continue if mode == "unmask": - item["text"] = self._unmask_pii_text(text_value, pii_tokens) + item["text"] = _replace_pii_tokens_in_text( + text_value, pii_tokens + ) elif mode == "mask": item["text"] = await self.check_pii( text=text_value, @@ -1075,7 +1433,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): args = function.arguments if isinstance(args, str): if mode == "unmask": - function.arguments = self._unmask_pii_text( + function.arguments = _replace_pii_tokens_in_text( args, pii_tokens ) elif mode == "mask": @@ -1092,7 +1450,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): args = function_call.arguments if isinstance(args, str): if mode == "unmask": - function_call.arguments = self._unmask_pii_text( + function_call.arguments = _replace_pii_tokens_in_text( args, pii_tokens ) elif mode == "mask": @@ -1232,6 +1590,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): verbose_proxy_logger.error(f"Error in PII streaming processing: {str(e)}") for chunk in remaining_chunks: yield chunk + finally: + self._clear_request_scoped_pii_tokens(request_data) async def async_post_call_streaming_iterator_hook( # type: ignore[override] self, @@ -1253,11 +1613,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk return - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) + pii_tokens = ( + request_data.get("_presidio_pii_tokens", {}).get(self.guardrail_name, {}) + if request_data + else {} + ) if not pii_tokens and request_data: verbose_proxy_logger.debug( - "No pii_tokens in request_data['metadata'] for streaming unmask path" + "No pii_tokens found in request_data for streaming unmask path" ) if not (self.output_parse_pii and pii_tokens): async for chunk in response: @@ -1300,6 +1663,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except Exception: pass + def _clear_request_scoped_pii_tokens(self, data: dict) -> None: + pii_token_map = data.get("_presidio_pii_tokens") + if not isinstance(pii_token_map, dict): + return + pii_token_map.pop(self.guardrail_name, None) + if not pii_token_map: + data.pop("_presidio_pii_tokens", None) + @log_guardrail_information async def apply_guardrail( self, @@ -1315,15 +1686,18 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """ texts = inputs.get("texts", []) - # When input_type is "response" and pii_tokens are available, - # unmask the text instead of masking it. - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) - + request_pii_tokens: Dict[str, str] = {} + response_pii_tokens = ( + request_data.get("_presidio_pii_tokens", {}).get(self.guardrail_name, {}) + if request_data + else {} + ) new_texts = [] - if input_type == "response" and pii_tokens: + if input_type == "response" and response_pii_tokens: for text in texts: - new_texts.append(self._unmask_pii_text(text, pii_tokens)) + new_texts.append(_replace_pii_tokens_in_text(text, response_pii_tokens)) + if request_data is not None: + self._clear_request_scoped_pii_tokens(request_data) else: for text in texts: modified_text = await self.check_pii( @@ -1331,9 +1705,21 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): output_parse_pii=self.output_parse_pii, presidio_config=None, request_data=request_data or {}, + pii_tokens=request_pii_tokens, ) new_texts.append(modified_text) inputs["texts"] = new_texts + # When using unified guardrail path, pre_call uses apply_guardrail instead of + # async_pre_call_hook; store pii_tokens in request_data so post_call can unmask. + if ( + input_type == "request" + and self.output_parse_pii + and request_pii_tokens + and request_data is not None + ): + request_data.setdefault("_presidio_pii_tokens", {})[ + self.guardrail_name + ] = dict(request_pii_tokens) return inputs def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index a1623121da5..d468839aecb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -202,7 +202,9 @@ class UnifiedLLMGuardrails(CustomLogger): ) from litellm.types.guardrails import GuardrailEventHooks - guardrail_to_apply: CustomGuardrail = data.pop("guardrail_to_apply", None) + guardrail_to_apply: Optional[CustomGuardrail] = data.pop( + "guardrail_to_apply", None + ) if guardrail_to_apply is None: return @@ -219,6 +221,31 @@ class UnifiedLLMGuardrails(CustomLogger): "async_post_call_success_hook response: %s", response ) + # Presidio (and similar) with output_parse_pii: unmask tokens in response + # using request-scoped pii_tokens. Must call the guardrail's hook directly + # with full data; process_output_response would mask output again. + if getattr(guardrail_to_apply, "output_parse_pii", False): + if hasattr(guardrail_to_apply, "async_post_call_success_hook"): + hook_result = await guardrail_to_apply.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + ) + if hook_result is not None: + response = hook_result + add_guardrail_to_applied_guardrails_header( + request_data=data, + guardrail_name=guardrail_to_apply.guardrail_name, + ) + return response + verbose_proxy_logger.warning( + "Guardrail %s has output_parse_pii=True but no " + "async_post_call_success_hook; preserving original response " + "without falling back to process_output_response.", + getattr(guardrail_to_apply, "guardrail_name", repr(guardrail_to_apply)), + ) + return response + call_type: Optional[CallTypesLiteral] = None if user_api_key_dict.request_route is not None: call_types = get_call_types_for_route(user_api_key_dict.request_route) diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 109f2237165..b25c6975933 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -79,6 +79,9 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): filter_scope = getattr(litellm_params, "presidio_filter_scope", None) or "both" run_input = filter_scope in ("input", "both") run_output = filter_scope in ("output", "both") + mask_residual_output_pii = bool( + getattr(litellm_params, "presidio_mask_residual_output_pii", False) + ) def _make_presidio_callback(**overrides): params = dict( @@ -104,22 +107,29 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): primary_callback = None if run_input: + # output_parse_pii-specific hook expansion is handled inside + # _OPTIONAL_PresidioPIIMasking.__init__. primary_callback = _make_presidio_callback() - if litellm_params.output_parse_pii: - _make_presidio_callback( - output_parse_pii=True, - event_hook=GuardrailEventHooks.post_call.value, - ) - - if run_output: - output_callback = _make_presidio_callback( - apply_to_output=True, - event_hook=GuardrailEventHooks.post_call.value, - output_parse_pii=False, - ) + # By default, output_parse_pii uses post_call for token unmasking only. A + # separate output-masking callback is created only when output masking is + # requested without input token unmasking, or when users explicitly opt in to + # masking residual model-generated PII after output_parse_pii flows. + if run_output and ( + not (litellm_params.output_parse_pii and run_input) or mask_residual_output_pii + ): if primary_callback is None: - primary_callback = output_callback + primary_callback = _make_presidio_callback( + apply_to_output=True, + event_hook=GuardrailEventHooks.post_call.value, + output_parse_pii=False, + ) + else: + _make_presidio_callback( + apply_to_output=True, + event_hook=GuardrailEventHooks.post_call.value, + output_parse_pii=False, + ) return primary_callback diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index f1590b16c24..4dddaf3dbb5 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,9 +1,14 @@ import asyncio -from typing import TYPE_CHECKING, Any, Literal, Optional +import os +import threading +import time +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple from fastapi import HTTPException, status import litellm +from litellm._logging import verbose_proxy_logger if TYPE_CHECKING: from litellm.router import Router as _Router @@ -13,19 +18,171 @@ else: LitellmRouter = Any -def _route_user_config_request(data: dict, route_type: str): - """Route a request using the user-provided router config.""" - router_config = data.pop("user_config") +_USER_CONFIG_ROUTER_CACHE_MAX_SIZE = max( + 1, int(os.getenv("LITELLM_USER_CONFIG_ROUTER_CACHE_MAX_SIZE", "64")) +) +_USER_CONFIG_ROUTER_CACHE_TTL_SECONDS = max( + 1, int(os.getenv("LITELLM_USER_CONFIG_ROUTER_CACHE_TTL_SECONDS", "300")) +) +_USER_CONFIG_ROUTER_CACHE: "OrderedDict[Tuple[Any, ...], Tuple[LitellmRouter, float]]" = ( + OrderedDict() +) +_USER_CONFIG_ROUTER_CACHE_LOCK = threading.Lock() + +def _discard_router_safely(router: LitellmRouter) -> None: + try: + router.discard() + except Exception: + pass + + +def _freeze_user_config_cache_key(value: Any) -> Any: + if isinstance(value, dict): + return tuple( + (str(key), _freeze_user_config_cache_key(val)) + for key, val in sorted(value.items(), key=lambda item: str(item[0])) + ) + if isinstance(value, list): + return tuple(_freeze_user_config_cache_key(item) for item in value) + if isinstance(value, tuple): + return tuple(_freeze_user_config_cache_key(item) for item in value) + if isinstance(value, set): + frozen_items = [_freeze_user_config_cache_key(item) for item in value] + return tuple( + sorted( + frozen_items, + key=lambda item: (type(item).__name__, repr(item)), + ) + ) + return value + + +def _get_user_config_router_cache_key(filtered_config: dict) -> Tuple[Any, ...]: + frozen_config = _freeze_user_config_cache_key(filtered_config) + if isinstance(frozen_config, tuple): + return frozen_config + return (frozen_config,) + + +def _prune_expired_user_config_routers(now: float) -> None: + expired_keys = [ + key + for key, (_, expires_at) in _USER_CONFIG_ROUTER_CACHE.items() + if expires_at <= now + ] + for key in expired_keys: + _USER_CONFIG_ROUTER_CACHE.pop(key) + + +def _clear_user_config_router_cache() -> None: + with _USER_CONFIG_ROUTER_CACHE_LOCK: + while _USER_CONFIG_ROUTER_CACHE: + _USER_CONFIG_ROUTER_CACHE.popitem(last=False) + + +def _get_or_create_user_config_router(filtered_config: dict) -> LitellmRouter: + cache_key = _get_user_config_router_cache_key(filtered_config) + with _USER_CONFIG_ROUTER_CACHE_LOCK: + now = time.monotonic() + _prune_expired_user_config_routers(now) + cached_entry = _USER_CONFIG_ROUTER_CACHE.get(cache_key) + if cached_entry is not None: + router, expires_at = cached_entry + if expires_at > now: + _USER_CONFIG_ROUTER_CACHE[cache_key] = ( + router, + now + _USER_CONFIG_ROUTER_CACHE_TTL_SECONDS, + ) + _USER_CONFIG_ROUTER_CACHE.move_to_end(cache_key) + return router + _USER_CONFIG_ROUTER_CACHE.pop(cache_key, None) + + new_router = litellm.Router(**filtered_config) + inserted_into_cache = False + evicted_count = 0 + try: + with _USER_CONFIG_ROUTER_CACHE_LOCK: + now = time.monotonic() + _prune_expired_user_config_routers(now) + cached_entry = _USER_CONFIG_ROUTER_CACHE.get(cache_key) + if cached_entry is not None: + router, expires_at = cached_entry + if expires_at > now: + _USER_CONFIG_ROUTER_CACHE[cache_key] = ( + router, + now + _USER_CONFIG_ROUTER_CACHE_TTL_SECONDS, + ) + _USER_CONFIG_ROUTER_CACHE.move_to_end(cache_key) + _discard_router_safely(new_router) + return router + _USER_CONFIG_ROUTER_CACHE.pop(cache_key, None) + + _USER_CONFIG_ROUTER_CACHE[cache_key] = ( + new_router, + now + _USER_CONFIG_ROUTER_CACHE_TTL_SECONDS, + ) + inserted_into_cache = True + _USER_CONFIG_ROUTER_CACHE.move_to_end(cache_key) + + while len(_USER_CONFIG_ROUTER_CACHE) > _USER_CONFIG_ROUTER_CACHE_MAX_SIZE: + _USER_CONFIG_ROUTER_CACHE.popitem(last=False) + evicted_count += 1 + except Exception: + with _USER_CONFIG_ROUTER_CACHE_LOCK: + if inserted_into_cache: + cached_entry = _USER_CONFIG_ROUTER_CACHE.get(cache_key) + if cached_entry is not None and cached_entry[0] is new_router: + _USER_CONFIG_ROUTER_CACHE.pop(cache_key, None) + _discard_router_safely(new_router) + raise + if evicted_count > 0: + verbose_proxy_logger.warning( + "user_config Router cache full (evicted %d entries). " + "Increase LITELLM_USER_CONFIG_ROUTER_CACHE_MAX_SIZE if this is frequent.", + evicted_count, + ) + return new_router + + +async def _route_user_config_request(data: dict, user_config: dict, route_type: str): + """ + Route a request using the user-provided router config. + + This is `async def` and returns the final response when awaited. + `route_request()` deliberately returns this coroutine without awaiting it, + preserving the existing double-await call-site pattern: + `llm_call = await route_request(...); response = await llm_call`. + """ # Filter router_config to only include valid Router.__init__ arguments # This prevents TypeError when invalid parameters are stored in the database valid_args = litellm.Router.get_valid_args() - filtered_config = {k: v for k, v in router_config.items() if k in valid_args} + filtered_config = {k: v for k, v in user_config.items() if k in valid_args} - user_router = litellm.Router(**filtered_config) - ret_val = getattr(user_router, f"{route_type}")(**data) - user_router.discard() - return ret_val + user_router = await asyncio.to_thread( + _get_or_create_user_config_router, filtered_config + ) + + # Handle batch completions with comma-separated models for user-provided routers + if ( + route_type == "acompletion" + and data.get("model", "") is not None + and "," in data.get("model", "") + ): + if data.get("fastest_response", False): + models = [m.strip() for m in str(data.get("model", "")).split(",")] + kwargs = dict(_kwargs_for_llm(data)) + kwargs.pop("model", None) + return await user_router.abatch_completion_fastest_response( + models=models, **kwargs + ) + + models = [m.strip() for m in str(data.get("model", "")).split(",")] + kwargs = dict(_kwargs_for_llm(data)) + kwargs.pop("model", None) + return await user_router.abatch_completion(models=models, **kwargs) + + return await getattr(user_router, f"{route_type}")(**_kwargs_for_llm(data)) def _is_a2a_agent_model(model_name: Any) -> bool: @@ -140,6 +297,25 @@ def _get_shared_session_lock() -> asyncio.Lock: return _shared_session_lock +# Keys that are only for proxy/guardrail use and must not be sent to the LLM API +_INTERNAL_REQUEST_KEYS = frozenset( + {"_presidio_pii_tokens", "fastest_response", "user_config"} +) + + +def _kwargs_for_llm(data: dict) -> dict: + """ + Strip internal proxy keys so they are not sent to the LLM provider. + + NOTE: Returns `data` by reference on the fast path when no internal keys are + present. Callers that need to mutate the result (for example `pop("model")`) + must wrap it in `dict(...)` first. + """ + if _INTERNAL_REQUEST_KEYS.isdisjoint(data): + return data + return {k: v for k, v in data.items() if k not in _INTERNAL_REQUEST_KEYS} + + async def add_shared_session_to_data(data: dict) -> None: """ Add shared aiohttp session for connection reuse (prevents cold starts). @@ -329,10 +505,18 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin if "generationConfig" in data and "config" not in data: data["config"] = data.pop("generationConfig") if "api_key" in data or "api_base" in data: + kwargs = _kwargs_for_llm(data) if llm_router is not None: - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**kwargs) else: - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**kwargs) + + elif "user_config" in data: + # user_config stays authoritative for routing, including comma-separated + # batch model requests. Keep this check ahead of the main-router batch + # branch so user-scoped router resolution is not bypassed. + user_config = data.pop("user_config") + return _route_user_config_request(data, user_config, route_type) elif ( route_type == "acompletion" @@ -340,16 +524,18 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin and "," in data.get("model", "") and llm_router is not None ): - # Handle batch completions with comma-separated models BEFORE user_config check - # This ensures batch completion logic is applied even when user_config is set if data.get("fastest_response", False): - return llm_router.abatch_completion_fastest_response(**data) + models = [model.strip() for model in str(data.get("model", "")).split(",")] + kwargs = dict(_kwargs_for_llm(data)) + kwargs.pop("model", None) + return llm_router.abatch_completion_fastest_response( + models=models, **kwargs + ) else: - models = [model.strip() for model in data.pop("model").split(",")] - return llm_router.abatch_completion(models=models, **data) - - elif "user_config" in data: - return _route_user_config_request(data, route_type) + models = [model.strip() for model in str(data.get("model", "")).split(",")] + kwargs = dict(_kwargs_for_llm(data)) + kwargs.pop("model", None) + return llm_router.abatch_completion(models=models, **kwargs) elif "router_settings_override" in data: # Apply per-request router settings overrides from key/team config @@ -374,10 +560,12 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin data[key] = override_settings[key] # Use main router with overridden kwargs + kwargs = _kwargs_for_llm(data) if llm_router is not None: - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**kwargs) else: - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**kwargs) + elif llm_router is not None: # Evals API: always route to litellm directly (not through router) # But extract model credentials if a model is provided @@ -421,7 +609,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin # If we can't get deployment creds, continue without them pass - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**_kwargs_for_llm(data)) # Skip model-based routing for container operations if route_type in [ "acreate_container", @@ -434,7 +622,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin "adelete_container_file", "aretrieve_container_file_content", ]: - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**_kwargs_for_llm(data)) # Interactions API: create with agent, get/delete/cancel don't need model routing if route_type in [ "acreate_interaction", @@ -442,7 +630,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin "adelete_interaction", "acancel_interaction", ]: - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**_kwargs_for_llm(data)) if route_type in [ "avideo_list", "avideo_status", @@ -463,7 +651,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin "aingest", ] and (data.get("model") is None or data.get("model") == ""): # These endpoints don't need a model, use custom_llm_provider directly - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**_kwargs_for_llm(data)) team_model_name = ( llm_router.map_team_model(data["model"], team_id) @@ -472,33 +660,33 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin ) if team_model_name is not None: data["model"] = team_model_name - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**_kwargs_for_llm(data)) elif data["model"] in router_model_names or llm_router.has_model_id( data["model"] ): - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**_kwargs_for_llm(data)) elif ( llm_router.model_group_alias is not None and data["model"] in llm_router.model_group_alias ): - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**_kwargs_for_llm(data)) elif data["model"] not in router_model_names: # Check wildcards before checking deployment_names # Priority: 1. Exact model_name match, 2. Wildcard match, 3. deployment_names match if llm_router.router_general_settings.pass_through_all_models: - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**_kwargs_for_llm(data)) elif ( llm_router.default_deployment is not None or len(llm_router.pattern_router.patterns) > 0 ): - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**_kwargs_for_llm(data)) elif data["model"] in llm_router.deployment_names: # Only match deployment_names if no wildcard matched return getattr(llm_router, f"{route_type}")( - **data, specific_deployment=True + **_kwargs_for_llm(data), specific_deployment=True ) elif route_type in [ "amoderation", @@ -526,7 +714,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin "aretrieve_container_file_content", ]: # These endpoints can work with or without model parameter - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**_kwargs_for_llm(data)) elif route_type in [ "avideo_status", "avideo_content", @@ -539,10 +727,10 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin # Video endpoints: If model is provided (e.g., from decoded video_id or target_model_names), # try router first to allow for multi-deployment load balancing try: - return getattr(llm_router, f"{route_type}")(**data) + return getattr(llm_router, f"{route_type}")(**_kwargs_for_llm(data)) except Exception: # If router fails (e.g., model not found in router), fall back to direct call - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**_kwargs_for_llm(data)) elif _is_a2a_agent_model(data.get("model", "")): from litellm.proxy.agent_endpoints.a2a_routing import ( route_a2a_agent_request, @@ -554,9 +742,9 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin # Fall through to raise exception below if result is None elif user_model is not None: - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**_kwargs_for_llm(data)) elif route_type == "allm_passthrough_route": - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**_kwargs_for_llm(data)) # if no route found then it's a bad request route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ec98cfd4d1e..bfa460a2106 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -141,6 +141,48 @@ else: unified_guardrail = UnifiedLLMGuardrails() +def _should_run_output_parse_pii_last() -> bool: + """ + Opt-in ordering change for output_parse_pii post-call guardrails. + + By default we preserve callback registration order for backwards + compatibility. Deployments that want output_parse_pii unmasking to run + after all other post-call guardrails can set + LITELLM_RUN_OUTPUT_PARSE_PII_LAST=true. + """ + return os.getenv("LITELLM_RUN_OUTPUT_PARSE_PII_LAST", "false").lower() == "true" + + +def _should_auto_run_output_parse_pii_last( + guardrail_callbacks: List[CustomGuardrail], +) -> bool: + """ + Auto-enable output_parse_pii-last ordering for the explicit Presidio + residual-output-masking flow. + + When the same guardrail registers both: + - an output_parse_pii post-call callback (unmask original request tokens), and + - an apply_to_output post-call callback (mask residual model-generated PII), + the residual-masking callback must run before the unmask callback. Otherwise + the unmask step is immediately overwritten and output_parse_pii produces + incorrect results. + """ + output_parse_guardrail_names = { + getattr(callback, "guardrail_name", None) + for callback in guardrail_callbacks + if getattr(callback, "output_parse_pii", False) + and getattr(callback, "guardrail_name", None) is not None + } + if not output_parse_guardrail_names: + return False + + return any( + getattr(callback, "apply_to_output", False) + and getattr(callback, "guardrail_name", None) in output_parse_guardrail_names + for callback in guardrail_callbacks + ) + + def print_verbose(print_statement): """ Prints the given `print_statement` to the console if `litellm.set_verbose` is True. @@ -1968,8 +2010,27 @@ class ProxyLogging: guardrail_data = _check_and_merge_model_level_guardrails( data=data, llm_router=llm_router ) - - for callback in guardrail_callbacks: + ordered_guardrail_callbacks = guardrail_callbacks + if ( + _should_run_output_parse_pii_last() + or _should_auto_run_output_parse_pii_last(guardrail_callbacks) + ): + # Opt-in ordering guarantee: all non-output_parse_pii guardrails + # run first, then output_parse_pii guardrails run last. This + # prevents request-token unmasking from being overwritten by + # later post-call guardrails (for example apply_to_output + # masking the response again). + output_parse_callbacks = [] + other_guardrail_callbacks = [] + for callback in guardrail_callbacks: + if getattr(callback, "output_parse_pii", False): + output_parse_callbacks.append(callback) + else: + other_guardrail_callbacks.append(callback) + ordered_guardrail_callbacks = ( + other_guardrail_callbacks + output_parse_callbacks + ) + for callback in ordered_guardrail_callbacks: # Main - V2 Guardrails implementation if ( diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index cfec0398c81..0387b795801 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -298,6 +298,14 @@ class PresidioPresidioConfigModelUserInterface(BaseModel): # extra param to let the ui know this is a boolean json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL}, ) + presidio_mask_residual_output_pii: Optional[bool] = Field( + default=False, + description=( + "When True and output_parse_pii is enabled, also run a separate " + "post-call output-masking pass for new PII introduced by the model." + ), + json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL}, + ) presidio_language: Optional[str] = Field( default="en", description="Language code for Presidio PII analysis (e.g., 'en', 'de', 'es', 'fr')", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 32a8c1b1070..fc4cdb66212 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -4,7 +4,9 @@ Tests PII detection and masking for different message formats """ import asyncio +import logging import os +import re import sys from contextlib import asynccontextmanager from unittest.mock import MagicMock, patch @@ -15,13 +17,27 @@ sys.path.insert(0, os.path.abspath("../../../../../..")) import litellm from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.presidio import ( _OPTIONAL_PresidioPIIMasking, + _replace_pii_tokens_in_text, ) from litellm.exceptions import GuardrailRaisedException -from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType -from litellm.types.utils import Choices, Message, ModelResponse +from litellm.proxy.route_llm_request import ( + _clear_user_config_router_cache, + route_request, +) +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import LitellmParams, Mode, PiiAction, PiiEntityType +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) def _make_mock_session_iterator( @@ -72,6 +88,13 @@ def _make_mock_session_iterator( return mock_iterator +@pytest.fixture(autouse=True) +def clear_user_config_router_cache(): + _clear_user_config_router_cache() + yield + _clear_user_config_router_cache() + + @pytest.fixture def presidio_guardrail(): """Create a Presidio guardrail instance for testing""" @@ -137,7 +160,9 @@ async def test_multimodal_message_format_completion_call_type( } # Mock the check_pii method to return redacted text - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): # Simulate PII detection and masking redacted_text = text redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]") @@ -206,7 +231,9 @@ async def test_multimodal_message_format_anthropic_messages_call_type( } # Mock the check_pii method to return redacted text - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): # Simulate PII detection and masking redacted_text = text redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]") @@ -276,7 +303,9 @@ async def test_multimodal_message_multiple_content_items( } # Mock the check_pii method - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): redacted_text = text redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]") redacted_text = redacted_text.replace("test@example.com", "[EMAIL]") @@ -336,7 +365,9 @@ async def test_mixed_string_and_list_content( } # Mock the check_pii method - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): redacted_text = text redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]") redacted_text = redacted_text.replace("test@example.com", "[EMAIL]") @@ -400,7 +431,9 @@ async def test_content_list_without_text_field( } # Mock the check_pii method - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): redacted_text = text.replace("test@example.com", "[EMAIL]") return redacted_text @@ -476,6 +509,615 @@ async def test_no_messages_field(presidio_guardrail, mock_user_api_key, mock_cac print("✓ No messages field test passed") +# --- _replace_pii_tokens_in_text (output_parse_pii unmasking) --- + + +def test_replace_pii_tokens_exact_replace(): + """Exact placeholder keys are replaced with original values.""" + text = "Hello and ." + pii_tokens = {"": "Jane", "": "jane@example.com"} + result = _replace_pii_tokens_in_text(text, pii_tokens) + assert result == "Hello Jane and jane@example.com." + + +def test_replace_pii_tokens_longest_key_first(): + """Longer keys are replaced first so uuid is replaced before .""" + text = "User a1b2c3 and ." + pii_tokens = { + "": "Unknown", + "a1b2c3": "Alice", + } + result = _replace_pii_tokens_in_text(text, pii_tokens) + assert result == "User Alice and Unknown." + + +def test_replace_pii_tokens_corrupted_uuid_placeholder(): + """LLM-corrupted uuid-style placeholders (key + trailing letters) are unmasked.""" + # e.g. ...fa9d -> ...fa9den + text = "Contact 123e4567-e89b-12d3-a456-426614174000en for details." + pii_tokens = { + "123e4567-e89b-12d3-a456-426614174000": "Jane Doe", + } + result = _replace_pii_tokens_in_text(text, pii_tokens) + assert result == "Contact Jane Doe for details." + + +def test_replace_pii_tokens_corrupted_uuid_placeholder_with_digit_suffix(): + """UUID-style placeholders with trailing digits are unmasked in one pass.""" + text = "Contact 123e4567-e89b-12d3-a456-4266141740002 for details." + pii_tokens = { + "123e4567-e89b-12d3-a456-426614174000": "Jane Doe", + } + result = _replace_pii_tokens_in_text(text, pii_tokens) + assert result == "Contact Jane Doe for details." + + +def test_replace_pii_tokens_truncated_placeholder_suffix(): + """A placeholder truncated at the end of the response should still unmask.""" + token = "123e4567-e89b-12d3-a456-426614174000" + text = f"Contact {token[:24]}" + pii_tokens = {token: "Jane Doe"} + result = _replace_pii_tokens_in_text(text, pii_tokens) + assert result == "Contact Jane Doe" + + +def test_replace_pii_tokens_empty_dict_returns_unchanged(): + """Empty pii_tokens returns text unchanged.""" + text = "Hello ." + result = _replace_pii_tokens_in_text(text, {}) + assert result == text + + +@pytest.mark.asyncio +async def test_post_call_uses_request_scoped_pii_tokens_not_instance(mock_user_api_key): + """Post_call uses only data['_presidio_pii_tokens'] and does not fall back to self.pii_tokens.""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_guardrail", + mock_testing=True, + output_parse_pii=True, + ) + # Intentionally set instance pii_tokens to a different (stale) value + presidio.pii_tokens = {"": "WrongPerson"} + # Request-scoped mapping (what pre_call would have stored) + data = { + "_presidio_pii_tokens": { + "test_guardrail": {"": "Jane Doe"}, + }, + } + response = ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content="Hello !"), + index=0, + finish_reason="stop", + ) + ], + ) + result = await presidio.async_post_call_success_hook( + data=data, + user_api_key_dict=mock_user_api_key, + response=response, + ) + assert result.choices[0].message.content == "Hello Jane Doe!" + # Proves we did not use self.pii_tokens ("WrongPerson") + assert "WrongPerson" not in result.choices[0].message.content + assert "_presidio_pii_tokens" not in data + + +@pytest.mark.asyncio +async def test_post_call_clears_request_scoped_pii_tokens_for_unrecognized_response_type( + mock_user_api_key, +): + """Unknown non-streaming response types should still clear request-scoped pii tokens.""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_guardrail", + mock_testing=True, + output_parse_pii=True, + ) + data = { + "_presidio_pii_tokens": { + "test_guardrail": {"": "Jane Doe"}, + }, + } + response = {"not": "a-known-response-type"} + + result = await presidio.async_post_call_success_hook( + data=data, + user_api_key_dict=mock_user_api_key, + response=response, # type: ignore[arg-type] + ) + + assert result == response + assert "_presidio_pii_tokens" not in data + + +@pytest.mark.asyncio +async def test_post_call_clears_request_scoped_pii_tokens_on_unmask_exception( + mock_user_api_key, +): + """Post_call should clear request-scoped pii tokens even if unmasking raises.""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_guardrail", + mock_testing=True, + output_parse_pii=True, + ) + data = { + "_presidio_pii_tokens": { + "test_guardrail": {"": "Jane Doe"}, + }, + } + response = ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content="Hello !"), + index=0, + finish_reason="stop", + ) + ], + ) + + async def raising_process_response_for_pii(*args, **kwargs): + raise RuntimeError("unmask failed") + + presidio._process_response_for_pii = raising_process_response_for_pii # type: ignore[method-assign] + + with pytest.raises(RuntimeError, match="unmask failed"): + await presidio.async_post_call_success_hook( + data=data, + user_api_key_dict=mock_user_api_key, + response=response, + ) + + assert "_presidio_pii_tokens" not in data + + +@pytest.mark.asyncio +async def test_pre_call_stores_pii_tokens_in_data_when_output_parse_pii( + mock_user_api_key, mock_cache +): + """Pre_call stores _presidio_pii_tokens in data when output_parse_pii is True.""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="pre_call_test", + mock_testing=True, + output_parse_pii=True, + ) + + # Simulate check_pii having recorded placeholders + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): + assert pii_tokens is not None + pii_tokens[""] = "Alice" + return text + + presidio.check_pii = mock_check_pii + data = { + "messages": [{"role": "user", "content": "Hello Alice"}], + "model": "gpt-4", + } + result = await presidio.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=data, + call_type="completion", + ) + assert result is not None + assert "_presidio_pii_tokens" in result + assert "pre_call_test" in result["_presidio_pii_tokens"] + assert result["_presidio_pii_tokens"]["pre_call_test"] == {"": "Alice"} + + +@pytest.mark.asyncio +async def test_pre_call_parallel_requests_use_request_local_pii_tokens( + mock_user_api_key, mock_cache +): + """Parallel pre_call requests keep independent token mappings.""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="race_test_guard", + mock_testing=True, + output_parse_pii=True, + ) + + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): + assert pii_tokens is not None + await asyncio.sleep(0.01) + pii_tokens[f"{text}"] = f"orig-{text}" + return f"masked-{text}" + + presidio.check_pii = mock_check_pii + + data_a = {"messages": [{"role": "user", "content": "A"}], "model": "gpt-4"} + data_b = {"messages": [{"role": "user", "content": "B"}], "model": "gpt-4"} + + result_a, result_b = await asyncio.gather( + presidio.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=data_a, + call_type="completion", + ), + presidio.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=data_b, + call_type="completion", + ), + ) + + assert result_a["_presidio_pii_tokens"]["race_test_guard"] == { + "A": "orig-A" + } + assert result_b["_presidio_pii_tokens"]["race_test_guard"] == { + "B": "orig-B" + } + + +@pytest.mark.asyncio +async def test_post_call_multiple_choices_unmask(mock_user_api_key): + """Post_call unmasks content in every choice (response.choices loop).""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="multi_choice_guard", + mock_testing=True, + output_parse_pii=True, + ) + data = { + "_presidio_pii_tokens": { + "multi_choice_guard": {"": "Alice", "": "alice@test.com"}, + }, + } + response = ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content="Hello !"), + index=0, + finish_reason="stop", + ), + Choices( + message=Message(role="assistant", content="Email: "), + index=1, + finish_reason="stop", + ), + ], + ) + result = await presidio.async_post_call_success_hook( + data=data, + user_api_key_dict=mock_user_api_key, + response=response, + ) + assert result.choices[0].message.content == "Hello Alice!" + assert result.choices[1].message.content == "Email: alice@test.com" + + +@pytest.mark.asyncio +async def test_post_call_unmasks_truncated_placeholder_suffix(mock_user_api_key): + """Standard ModelResponse path should recover placeholders truncated by max_tokens.""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="truncated_guard", + mock_testing=True, + output_parse_pii=True, + ) + token = "123e4567-e89b-12d3-a456-426614174000" + data = { + "_presidio_pii_tokens": { + "truncated_guard": {token: "Jane Doe"}, + }, + } + response = ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content=f"Contact {token[:24]}"), + index=0, + finish_reason="length", + ) + ], + ) + result = await presidio.async_post_call_success_hook( + data=data, + user_api_key_dict=mock_user_api_key, + response=response, + ) + assert result.choices[0].message.content == "Contact Jane Doe" + + +@pytest.mark.asyncio +async def test_post_call_pii_tokens_missing_guardrail_name_returns_unchanged( + mock_user_api_key, +): + """When _presidio_pii_tokens exists but has no key for this guardrail, response is returned unchanged.""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="my_guard", + mock_testing=True, + output_parse_pii=True, + ) + # data has _presidio_pii_tokens but no "my_guard" key (or empty dict for my_guard) + data = {"_presidio_pii_tokens": {}} + response = ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content="Hello !"), + index=0, + finish_reason="stop", + ) + ], + ) + result = await presidio.async_post_call_success_hook( + data=data, + user_api_key_dict=mock_user_api_key, + response=response, + ) + assert result.choices[0].message.content == "Hello !" + # Same when another guardrail has tokens but not ours (my_guard not in dict) + data2 = {"_presidio_pii_tokens": {"other_guardrail": {"": "Bob"}}} + result2 = await presidio.async_post_call_success_hook( + data=data2, + user_api_key_dict=mock_user_api_key, + response=ModelResponse( + id="2", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content="Hi "), + index=0, + finish_reason="stop", + ) + ], + ), + ) + assert result2.choices[0].message.content == "Hi " + + +@pytest.mark.asyncio +async def test_streaming_post_call_corrupted_uuid_placeholder_unmask(mock_user_api_key): + """Streaming response with LLM-corrupted uuid placeholder is unmasked correctly.""" + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="stream_guard", + mock_testing=True, + output_parse_pii=True, + ) + # Corrupted placeholder: key + "en" -> should unmask to "Jane Doe" + request_data = { + "_presidio_pii_tokens": { + "stream_guard": { + "123e4567-e89b-12d3-a456-426614174000": "Jane Doe", + }, + }, + } + + async def mock_stream(): + # Simulate chunks that together form: + # "Contact 123e4567-e89b-12d3-a456-426614174000en for details." + yield ModelResponseStream( + id="chunk1", + created=0, + model="gpt-test", + choices=[ + StreamingChoices( + delta=Delta(content="Contact "), + index=0, + ) + ], + ) + yield ModelResponseStream( + id="chunk2", + created=0, + model="gpt-test", + choices=[ + StreamingChoices( + delta=Delta( + content="123e4567-e89b-12d3-a456-426614174000en" + ), + index=0, + ) + ], + ) + yield ModelResponseStream( + id="chunk3", + created=0, + model="gpt-test", + choices=[ + StreamingChoices( + delta=Delta(content=" for details."), + index=0, + ) + ], + ) + + collected = [] + async for chunk in presidio.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data=request_data, + ): + if ( + chunk.choices + and hasattr(chunk.choices[0], "delta") + and chunk.choices[0].delta + and getattr(chunk.choices[0].delta, "content", None) + ): + collected.append(chunk.choices[0].delta.content or "") + elif ( + chunk.choices + and hasattr(chunk.choices[0], "message") + and getattr(chunk.choices[0].message, "content", None) + ): + collected.append(chunk.choices[0].message.content or "") + + full_content = "".join(collected) + assert full_content == "Contact Jane Doe for details." + assert "fa9den" not in full_content + assert "_presidio_pii_tokens" not in request_data + + +@pytest.mark.asyncio +async def test_integration_user_config_output_parse_pii_batch_streaming_no_remask_or_stale( + mock_user_api_key, mock_cache, monkeypatch +): + """ + Integration test for: + - route_request user_config branch + batch model routing + - output_parse_pii streaming unmask + - no stale pii token carryover between requests + - no re-mask in final output stream + """ + + class _FakeRouter: + calls = [] + + @staticmethod + def get_valid_args(): + return ["model_list"] + + def __init__(self, **kwargs): + self.init_kwargs = kwargs + + async def abatch_completion(self, models, **kwargs): + _FakeRouter.calls.append({"models": models, "kwargs": kwargs}) + token = "" + masked_content = kwargs.get("messages", [{}])[0].get("content", "") + for part in str(masked_content).split(): + if part.startswith(""): + token = part + break + + async def _stream(): + yield ModelResponseStream( + id="chunk-1", + created=0, + model="gpt-test", + choices=[StreamingChoices(delta=Delta(content="Hello "), index=0)], + ) + yield ModelResponseStream( + id="chunk-2", + created=0, + model="gpt-test", + choices=[StreamingChoices(delta=Delta(content=token), index=0)], + ) + + return _stream() + + def discard(self): + pass + + class _ReMaskLogger(CustomLogger): + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict, response, request_data + ): + async for chunk in response: + if chunk.choices and chunk.choices[0].delta is not None: + content = chunk.choices[0].delta.content + if isinstance(content, str): + chunk.choices[0].delta.content = content.replace( + "Alice", "" + ).replace("Bob", "") + yield chunk + + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="integration_presidio", + mock_testing=True, + output_parse_pii=True, + default_on=True, + ) + + token_map = {"Alice": "alice-token", "Bob": "bob-token"} + + async def _mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): + assert pii_tokens is not None + for name, token in token_map.items(): + if name in text: + pii_tokens[token] = name + return text.replace(name, token) + return text + + presidio.check_pii = _mock_check_pii + monkeypatch.setattr(litellm, "Router", _FakeRouter) + + async def _run_one_request(person_name: str) -> str: + data = { + "model": "gpt-4o-mini, claude-3-haiku", + "stream": True, + "messages": [{"role": "user", "content": f"My name is {person_name}"}], + "user_config": { + "model_list": [{"model_name": "x", "litellm_params": {"model": "x"}}], + "invalid_router_key": "ignore-me", + }, + } + + pre_call_data = await presidio.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=data, + call_type="completion", + ) + + llm_call = await route_request( + data=pre_call_data, + llm_router=None, + user_model=None, + route_type="acompletion", + ) + raw_stream = await llm_call + + proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + remask_logger = _ReMaskLogger() + with patch.object(litellm, "callbacks", [remask_logger, presidio]): + processed_stream = proxy_logging.async_post_call_streaming_iterator_hook( + response=raw_stream, + user_api_key_dict=mock_user_api_key, + request_data=pre_call_data, + ) + chunks = [] + async for chunk in processed_stream: + if ( + chunk.choices + and chunk.choices[0].delta is not None + and isinstance(chunk.choices[0].delta.content, str) + ): + chunks.append(chunk.choices[0].delta.content) + + return "".join(chunks) + + first_response = await _run_one_request("Alice") + second_response = await _run_one_request("Bob") + + assert first_response == "Hello Alice" + assert second_response == "Hello Bob" + assert "Alice" not in second_response # stale token carryover guard + assert "" not in first_response + assert "" not in second_response + + assert len(_FakeRouter.calls) == 2 + for call in _FakeRouter.calls: + assert call["models"] == ["gpt-4o-mini", "claude-3-haiku"] + assert "model" not in call["kwargs"] # batch should not forward model kwarg + assert "_presidio_pii_tokens" not in call["kwargs"] # internal key stripped + + @pytest.mark.asyncio async def test_logging_hook_multimodal_message_format(presidio_guardrail): """ @@ -502,7 +1144,9 @@ async def test_logging_hook_multimodal_message_format(presidio_guardrail): mock_result = {"choices": [{"message": {"content": "Response"}}]} # Mock the check_pii method - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): redacted_text = text redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]") redacted_text = redacted_text.replace("test@example.com", "[EMAIL]") @@ -558,7 +1202,9 @@ async def test_logging_hook_multiple_content_items(presidio_guardrail): mock_result = {"choices": [{"message": {"content": "Response"}}]} # Mock the check_pii method - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): redacted_text = text redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]") redacted_text = redacted_text.replace("test@example.com", "[EMAIL]") @@ -604,7 +1250,9 @@ async def test_presidio_sets_guardrail_information_in_request_data(): "metadata": {}, } - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): assert request_data is not None presidio.add_standard_logging_guardrail_information_to_request_data( @@ -664,7 +1312,9 @@ async def test_request_data_flows_to_apply_guardrail(): "metadata": {}, } - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): assert request_data is not None, "request_data should be passed to check_pii" assert "metadata" in request_data, "request_data should have metadata" @@ -698,7 +1348,9 @@ async def test_output_masking_apply_to_output_only(mock_user_api_key): pii_entities_config={PiiEntityType.CREDIT_CARD: PiiAction.MASK}, ) - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): return text.replace("4111-1111-1111-1111", "[CREDIT_CARD]") presidio.check_pii = mock_check_pii @@ -735,13 +1387,26 @@ async def test_presidio_filter_scope_initializer(monkeypatch): """ Ensure initializer respects presidio_filter_scope for input/output/both. """ + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _ensure_event_hook_includes_post_call, + ) created = [] class DummyGuardrail: - def __init__(self, apply_to_output: bool = False, event_hook=None, **kwargs): + def __init__( + self, + apply_to_output: bool = False, + event_hook=None, + output_parse_pii: bool = False, + **kwargs, + ): self.apply_to_output = apply_to_output - self.event_hook = event_hook + self.event_hook = ( + _ensure_event_hook_includes_post_call(event_hook) + if (apply_to_output or output_parse_pii) + else event_hook + ) created.append(self) def update_in_memory_litellm_params(self, litellm_params): @@ -787,7 +1452,7 @@ async def test_presidio_filter_scope_initializer(monkeypatch): assert len(created) == 1 assert created[0].apply_to_output is True - # both -> expect two callbacks (input + output) + # both without output_parse_pii -> expect two callbacks (input + output) created.clear() params_both = LitellmParams( guardrail="presidio", mode="pre_call", presidio_filter_scope="both" @@ -797,6 +1462,47 @@ async def test_presidio_filter_scope_initializer(monkeypatch): assert any(not c.apply_to_output for c in created) assert any(c.apply_to_output for c in created) + # both with output_parse_pii -> single callback (pre_call + post_call), no output-mask + created.clear() + params_both_output_parse = LitellmParams( + guardrail="presidio", + mode="pre_call", + presidio_filter_scope="both", + output_parse_pii=True, + ) + cb = initialize_presidio(params_both_output_parse, guardrail_dict) + assert len(created) == 1 + assert created[0].apply_to_output is False + assert created[0].event_hook == ["pre_call", "post_call"] + + created.clear() + params_both_output_parse_list_mode = LitellmParams( + guardrail="presidio", + mode=["pre_call", "during_call"], + presidio_filter_scope="both", + output_parse_pii=True, + ) + cb = initialize_presidio(params_both_output_parse_list_mode, guardrail_dict) + assert len(created) == 1 + assert created[0].apply_to_output is False + assert created[0].event_hook == ["pre_call", "during_call", "post_call"] + + # output_parse_pii + explicit residual output masking opt-in -> keep both callbacks + created.clear() + params_both_output_parse_with_residual_mask = LitellmParams( + guardrail="presidio", + mode="pre_call", + presidio_filter_scope="both", + output_parse_pii=True, + presidio_mask_residual_output_pii=True, + ) + cb = initialize_presidio( + params_both_output_parse_with_residual_mask, guardrail_dict + ) + assert len(created) == 2 + assert any(not c.apply_to_output for c in created) + assert any(c.apply_to_output for c in created) + @pytest.mark.asyncio async def test_empty_content_handling( @@ -834,7 +1540,9 @@ async def test_empty_content_handling( } # Mock check_pii to simulate PII processing without needing Presidio API - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): # Empty text returns as-is (this is what our fix ensures) return text @@ -875,7 +1583,9 @@ async def test_whitespace_only_content( } # Mock check_pii to simulate PII processing - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): return text presidio_guardrail.check_pii = mock_check_pii @@ -1104,7 +1814,9 @@ async def test_tool_calling_complete_scenario( } # Mock check_pii to simulate PII masking - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): if "john.doe@example.com" in text: return text.replace("john.doe@example.com", "[EMAIL]") return text @@ -1281,6 +1993,471 @@ async def test_anonymize_skips_when_no_detections_after_filter(): assert masked_entity_count == {} +@pytest.mark.asyncio +async def test_anonymize_counts_masked_entities_when_output_parse_disabled(): + """ + masked_entity_count should still be updated when output_parse_pii=False. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": "Call me at ", + "items": [ + { + "start": 11, + "end": 23, + "text": "", + "operator": "replace", + "entity_type": "PHONE_NUMBER", + } + ], + } + ) + masked_entity_count = {} + result = await guardrail.anonymize_text( + text="Call me at 555-123-4567", + analyze_results=[ + {"entity_type": "PHONE_NUMBER", "score": 0.99, "start": 11, "end": 23} + ], + output_parse_pii=False, + masked_entity_count=masked_entity_count, + ) + + assert result == "Call me at " + assert masked_entity_count == {"PHONE_NUMBER": 1} + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_applies_non_replace_operators(): + """ + output_parse_pii=True should still apply non-replace operators to the text + sent to the LLM, while only replace operators are tracked for unmasking. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": "Call me ", + "items": [ + { + "start": 8, + "end": 20, + "text": "", + "operator": "mask", + "entity_type": "PHONE_NUMBER", + } + ], + } + ) + pii_tokens = {} + masked_entity_count = {} + + result = await guardrail.anonymize_text( + text="Call me 555-123-4567", + analyze_results=[ + {"entity_type": "PHONE_NUMBER", "score": 0.99, "start": 8, "end": 20} + ], + output_parse_pii=True, + masked_entity_count=masked_entity_count, + pii_tokens=pii_tokens, + ) + + assert result == "Call me " + assert pii_tokens == {} + assert masked_entity_count == {"PHONE_NUMBER": 1} + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_mixed_operators_falls_back_to_redacted_text( + caplog, +): + """ + With mixed operators, output_parse_pii should return Presidio output directly + to avoid offset corruption from mixed-length placeholder rewrites. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": "Hello and ", + "items": [ + { + "start": 6, + "end": 14, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + { + "start": 19, + "end": 33, + "text": "", + "operator": "mask", + "entity_type": "PHONE_NUMBER", + }, + ], + } + ) + pii_tokens = {} + masked_entity_count = {} + + with caplog.at_level(logging.WARNING): + result = await guardrail.anonymize_text( + text="Hello Jane Doe and 555-123-4567", + analyze_results=[ + {"entity_type": "PERSON", "score": 0.99, "start": 6, "end": 14}, + {"entity_type": "PHONE_NUMBER", "score": 0.99, "start": 19, "end": 31}, + ], + output_parse_pii=True, + masked_entity_count=masked_entity_count, + pii_tokens=pii_tokens, + ) + + assert result == "Hello and " + assert pii_tokens == {} + assert masked_entity_count == {"PERSON": 1, "PHONE_NUMBER": 1} + assert ( + "1 replace-operator entities (['PERSON']) will also NOT be unmasked" + in caplog.text + ) + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_best_effort_when_analyze_has_extra_spans(): + """ + If analyzer returns more spans than anonymizer replace items (e.g. overlap/duplicate + detections), output_parse_pii should still do best-effort mapping instead of + dropping unmask metadata entirely. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": "Hello ", + "items": [ + { + "start": 6, + "end": 14, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + } + ], + } + ) + pii_tokens = {} + masked_entity_count = {} + + result = await guardrail.anonymize_text( + text="Hello John Smith", + analyze_results=[ + {"entity_type": "PERSON", "score": 0.99, "start": 6, "end": 16}, + # overlapping extra detection + {"entity_type": "PERSON", "score": 0.98, "start": 11, "end": 16}, + ], + output_parse_pii=True, + masked_entity_count=masked_entity_count, + pii_tokens=pii_tokens, + ) + + assert "John Smith" not in result + assert re.match(r"^Hello [0-9a-fA-F]{8}-", result) + assert set(pii_tokens.values()) == {"John Smith"} + assert masked_entity_count == {"PERSON": 1} + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_overlapping_spans_prefers_representative_span(): + """ + Overlapping analyze spans should collapse to a representative span so mapping + does not pair a short nested span (e.g. "John") instead of full entity text. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": " and ", + "items": [ + { + "start": 0, + "end": 8, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + { + "start": 13, + "end": 21, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + ], + } + ) + pii_tokens = {} + masked_entity_count = {} + + result = await guardrail.anonymize_text( + text="John Smith and Alice", + analyze_results=[ + {"entity_type": "PERSON", "score": 0.90, "start": 0, "end": 4}, # John + { + "entity_type": "PERSON", + "score": 0.95, + "start": 0, + "end": 10, + }, # John Smith + {"entity_type": "PERSON", "score": 0.99, "start": 15, "end": 20}, # Alice + ], + output_parse_pii=True, + masked_entity_count=masked_entity_count, + pii_tokens=pii_tokens, + ) + + assert "John Smith" not in result + assert "Alice" not in result + assert set(pii_tokens.values()) == {"John Smith", "Alice"} + assert "John" not in set(pii_tokens.values()) + assert masked_entity_count == {"PERSON": 2} + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_falls_back_when_replace_exceeds_analyze_spans(): + """ + If there are fewer analyze spans than replace items, fallback to redacted text + (cannot build a complete unmask mapping safely). + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": " and ", + "items": [ + { + "start": 0, + "end": 8, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + { + "start": 13, + "end": 21, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + ], + } + ) + pii_tokens = {} + masked_entity_count = {} + + result = await guardrail.anonymize_text( + text="Alice and Bob", + analyze_results=[ + {"entity_type": "PERSON", "score": 0.99, "start": 0, "end": 5}, + ], + output_parse_pii=True, + masked_entity_count=masked_entity_count, + pii_tokens=pii_tokens, + ) + + assert result == " and " + assert pii_tokens == {} + assert masked_entity_count == {"PERSON": 2} + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_multiple_replace_items_offset_safe(): + """ + Multiple replacements should be offset-safe when output_parse_pii=True. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": " and ", + # Deliberately ascending by start to catch offset-shift bugs. + "items": [ + { + "start": 0, + "end": 5, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + { + "start": 10, + "end": 13, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + ], + } + ) + pii_tokens = {} + masked_entity_count = {} + + result = await guardrail.anonymize_text( + text="Alice and Bob", + analyze_results=[ + {"entity_type": "PERSON", "score": 0.99, "start": 0, "end": 5}, + {"entity_type": "PERSON", "score": 0.99, "start": 10, "end": 13}, + ], + output_parse_pii=True, + masked_entity_count=masked_entity_count, + pii_tokens=pii_tokens, + ) + + assert "Alice" not in result + assert "Bob" not in result + assert " and " in result + assert len(pii_tokens) == 2 + assert set(pii_tokens.values()) == {"Alice", "Bob"} + assert masked_entity_count == {"PERSON": 2} + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_accepts_integral_float_offsets(): + """ + Analyze spans with integral float offsets (e.g. 0.0 from JSON decoding) should + still participate in unmask mapping instead of being silently dropped. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": "", + "items": [ + { + "start": 0, + "end": 8, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + } + ], + } + ) + pii_tokens = {} + + result = await guardrail.anonymize_text( + text="Alice", + analyze_results=[ + {"entity_type": "PERSON", "score": 0.99, "start": 0.0, "end": 5.0}, + ], + output_parse_pii=True, + masked_entity_count={}, + pii_tokens=pii_tokens, + ) + + assert "Alice" not in result + assert len(pii_tokens) == 1 + assert next(iter(pii_tokens.values())) == "Alice" + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_uses_analyze_spans_for_original_mapping(): + """ + Presidio items offsets are in anonymized-output coordinates. Ensure original + value mapping uses analyze spans (original coordinates), not item offsets. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": "Hello and ", + "items": [ + { + "start": 6, + "end": 14, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + { + "start": 19, + "end": 27, + "text": "", + "operator": "replace", + "entity_type": "PERSON", + }, + ], + } + ) + pii_tokens = {} + masked_entity_count = {} + + result = await guardrail.anonymize_text( + text="Hello Jane and Christopher", + analyze_results=[ + {"entity_type": "PERSON", "score": 0.99, "start": 6, "end": 10}, + {"entity_type": "PERSON", "score": 0.99, "start": 15, "end": 26}, + ], + output_parse_pii=True, + masked_entity_count=masked_entity_count, + pii_tokens=pii_tokens, + ) + + assert "Jane" not in result + assert "Christopher" not in result + assert set(pii_tokens.values()) == {"Jane", "Christopher"} + assert masked_entity_count == {"PERSON": 2} + + +@pytest.mark.asyncio +async def test_anonymize_output_parse_pii_without_pii_tokens_logs_warning(): + """ + output_parse_pii=True without pii_tokens should emit a warning, since token + mappings cannot be retained for post-call unmasking. + """ + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + guardrail.presidio_anonymizer_api_base = "http://mock-presidio/" + guardrail._get_session_iterator = _make_mock_session_iterator( + { + "text": "Call me at ", + "items": [ + { + "start": 11, + "end": 23, + "text": "", + "operator": "replace", + "entity_type": "PHONE_NUMBER", + } + ], + } + ) + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger.warning" + ) as mock_warning: + result = await guardrail.anonymize_text( + text="Call me at 555-123-4567", + analyze_results=[ + { + "entity_type": "PHONE_NUMBER", + "score": 0.99, + "start": 11, + "end": 23, + } + ], + output_parse_pii=True, + masked_entity_count={}, + pii_tokens=None, + ) + + assert re.match( + r"^Call me at [0-9a-fA-F]{8}-[0-9a-fA-F]{4}-", + result, + ) + assert mock_warning.called is True + assert "pii_tokens is None" in mock_warning.call_args[0][0] + + def test_blocking_respects_threshold_filter(): """ Entities filtered out by score should not trigger blocking, but high-score detections should. @@ -1611,12 +2788,11 @@ async def test_anonymize_text_http_error_status(): @pytest.mark.asyncio -async def test_pii_tokens_stored_in_metadata_not_top_level(presidio_guardrail): +async def test_pii_tokens_stored_in_request_scoped_internal_key(presidio_guardrail): """ - Regression test: pii_tokens must be stored in data['metadata']['pii_tokens'], - NOT in data['pii_tokens']. Storing at the top level leaks the field to LLM - providers like Anthropic, which reject unknown fields with - 'pii_tokens: Extra inputs are not permitted'. + Regression test: pii_tokens must be stored in request-scoped internal + data['_presidio_pii_tokens'][guardrail_name], not in legacy data['metadata']['pii_tokens'] + or data['pii_tokens']. """ guardrail = _OPTIONAL_PresidioPIIMasking( mock_testing=True, @@ -1638,16 +2814,11 @@ async def test_pii_tokens_stored_in_metadata_not_top_level(presidio_guardrail): "metadata": {}, } - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): - # Simulate PII masking with token storage (mimics real anonymize_text behavior) - if request_data is not None and output_parse_pii: - if "metadata" not in request_data: - request_data["metadata"] = {} - if "pii_tokens" not in request_data["metadata"]: - request_data["metadata"]["pii_tokens"] = {} - pii_tokens = request_data["metadata"]["pii_tokens"] - seq = len(pii_tokens) + 1 - token = f"" + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): + if pii_tokens is not None and output_parse_pii: + token = "123e4567-e89b-12d3-a456-426614174000" pii_tokens[token] = "John" text = text.replace("John", token) return text @@ -1667,19 +2838,20 @@ async def test_pii_tokens_stored_in_metadata_not_top_level(presidio_guardrail): "it would leak to LLM providers and cause 'Extra inputs are not permitted' errors" ) - # pii_tokens must be inside metadata (safe from provider leakage) - assert "metadata" in result - assert "pii_tokens" in result["metadata"] - assert len(result["metadata"]["pii_tokens"]) > 0 + assert "pii_tokens" not in (result.get("metadata") or {}) + assert "_presidio_pii_tokens" in result + assert guardrail.guardrail_name in result["_presidio_pii_tokens"] + assert len(result["_presidio_pii_tokens"][guardrail.guardrail_name]) > 0 @pytest.mark.asyncio -async def test_pii_tokens_in_metadata_used_for_unmasking(): +async def test_request_scoped_pii_tokens_used_for_unmasking(): """ - Regression test: _process_response_for_pii must read pii_tokens from - data['metadata']['pii_tokens'] and correctly unmask the response. + Regression test: _process_response_for_pii must read request-scoped + data['_presidio_pii_tokens'][guardrail_name] and correctly unmask the response. """ guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", mock_testing=True, output_parse_pii=True, ) @@ -1687,7 +2859,7 @@ async def test_pii_tokens_in_metadata_used_for_unmasking(): token_key = "" request_data = { "model": "claude-haiku-4-5-20251001", - "metadata": {"pii_tokens": {token_key: "John"}}, + "_presidio_pii_tokens": {guardrail.guardrail_name: {token_key: "John"}}, } response = ModelResponse( @@ -1743,6 +2915,57 @@ def test_event_hook_no_expansion_when_already_post_call(): assert guardrail.event_hook == "post_call" +def test_event_hook_none_with_output_parse_pii_defaults_to_pre_and_post_call(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook=None, + ) + assert guardrail.event_hook == ["pre_call", "post_call"] + + +def test_event_hook_mode_expands_to_include_post_call(): + mode = Mode( + tags={"team-a": ["pre_call", "during_call"], "team-b": "pre_call"}, + default="pre_call", + ) + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook=mode, + ) + + assert isinstance(guardrail.event_hook, Mode) + assert guardrail.event_hook.tags["team-a"] == [ + "pre_call", + "during_call", + "post_call", + ] + assert guardrail.event_hook.tags["team-b"] == ["pre_call", "post_call"] + assert guardrail.event_hook.default == ["pre_call", "post_call"] + + +def test_event_hook_mode_without_default_preserves_tag_only_behavior(): + mode = Mode( + tags={"team-a": ["pre_call", "during_call"], "team-b": "pre_call"}, + default=None, + ) + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook=mode, + ) + + assert isinstance(guardrail.event_hook, Mode) + assert guardrail.event_hook.tags["team-a"] == [ + "pre_call", + "during_call", + "post_call", + ] + assert guardrail.event_hook.tags["team-b"] == ["pre_call", "post_call"] + assert guardrail.event_hook.default is None + + @pytest.mark.asyncio async def test_metadata_none_does_not_crash(): """ @@ -1886,7 +3109,7 @@ async def test_anonymize_text_multiple_items_position_correctness(): mock_iterator = _make_mock_session_iterator(anonymizer_response) - request_data = {"metadata": {}} + pii_tokens = {} with patch.object(guardrail, "_get_session_iterator", mock_iterator): result = await guardrail.anonymize_text( text="Call John at 555-123-4567", @@ -1896,11 +3119,9 @@ async def test_anonymize_text_multiple_items_position_correctness(): ], output_parse_pii=True, masked_entity_count={}, - request_data=request_data, + pii_tokens=pii_tokens, ) - pii_tokens = request_data["metadata"]["pii_tokens"] - # Verify tokens captured the correct ORIGINAL text values person_token = [k for k in pii_tokens if "PERSON" in k][0] phone_token = [k for k in pii_tokens if "PHONE" in k][0] @@ -1924,14 +3145,15 @@ async def test_anthropic_native_response_unmasking(): when output_parse_pii is enabled. """ guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", mock_testing=True, output_parse_pii=True, ) request_data = { "model": "claude-3-haiku", - "metadata": { - "pii_tokens": { + "_presidio_pii_tokens": { + guardrail.guardrail_name: { "": "John Smith", "": "555-123-4567", } @@ -1977,7 +3199,9 @@ async def test_anthropic_native_response_masking(): apply_to_output=True, ) - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): return text.replace("John Smith", "[PERSON]").replace("555-123-4567", "[PHONE]") guardrail.check_pii = mock_check_pii @@ -2012,13 +3236,14 @@ async def test_anthropic_native_response_non_text_blocks_untouched(): should be left untouched during unmasking. """ guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", mock_testing=True, output_parse_pii=True, ) request_data = { "model": "claude-3-haiku", - "metadata": {"pii_tokens": {"": "John"}}, + "_presidio_pii_tokens": {guardrail.guardrail_name: {"": "John"}}, } anthropic_response = { @@ -2100,7 +3325,7 @@ async def test_streaming_unmask_path_bytes_passthrough(): byte_chunk = b'data: {"type":"content_block_delta"}\n\n' request_data = { - "metadata": {"pii_tokens": {"": "John"}}, + "_presidio_pii_tokens": {guardrail.guardrail_name: {"": "John"}}, } async def mock_stream(): @@ -2138,8 +3363,8 @@ async def test_apply_guardrail_unmask_on_response(): request_data = { "model": "gpt-4o", - "metadata": { - "pii_tokens": { + "_presidio_pii_tokens": { + guardrail.guardrail_name: { "": "John Smith", "": "555-123-4567", } @@ -2159,6 +3384,7 @@ async def test_apply_guardrail_unmask_on_response(): ) assert result["texts"][0] == "Hello John Smith, your number is 555-123-4567." + assert "_presidio_pii_tokens" not in request_data @pytest.mark.asyncio @@ -2172,7 +3398,9 @@ async def test_apply_guardrail_masks_on_request(): mock_testing=True, ) - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, pii_tokens=None + ): return text.replace("John Smith", "") guardrail.check_pii = mock_check_pii diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index bbba8e4d03d..7585c9eca26 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -18,7 +18,15 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai UnifiedLLMGuardrails, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices +from litellm.types.utils import ( + CallTypes, + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) class RecordingGuardrail(CustomGuardrail): @@ -38,6 +46,29 @@ class RecordingGuardrail(CustomGuardrail): return {"texts": inputs.get("texts", [])} +class OutputParseReturnsNoneGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="output-parse-none") + self.output_parse_pii = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def async_post_call_success_hook( # type: ignore[override] + self, data, user_api_key_dict, response + ): + return None + + +class OutputParseNoPostHookGuardrail: + def __init__(self): + self.guardrail_name = "output-parse-no-hook" + self.output_parse_pii = True + + def should_run_guardrail(self, data, event_type): + return True + + class _NoopTranslation(BaseTranslation): """Test translation handler that simply echoes input/output.""" @@ -404,3 +435,122 @@ class TestUnifiedLLMGuardrails: # Response returned with pages intact assert result.pages[0].markdown == "Some text" + + @pytest.mark.asyncio + async def test_output_parse_none_hook_preserves_original_response(self): + handler = UnifiedLLMGuardrails() + guardrail = OutputParseReturnsNoneGuardrail() + data = { + "guardrail_to_apply": guardrail, + "model": "gpt-test", + "metadata": {}, + } + + class _ExplodingTranslation(_NoopTranslation): + async def process_output_response( # type: ignore[override] + self, + response, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + ): + raise AssertionError( + "process_output_response should not run for output_parse_pii" + ) + + response = ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content="Hello"), + index=0, + finish_reason="stop", + ) + ], + ) + + original_endpoint_mappings = ( + unified_module.endpoint_guardrail_translation_mappings + ) + unified_module.endpoint_guardrail_translation_mappings = { + CallTypes.anthropic_messages: _ExplodingTranslation + } + original_infer_call_type = unified_module._infer_call_type + unified_module._infer_call_type = ( + lambda call_type=None, completion_response=None: CallTypes.anthropic_messages + ) # type: ignore[assignment] + + try: + result = await handler.async_post_call_success_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=response, + ) + finally: + unified_module._infer_call_type = original_infer_call_type + unified_module.endpoint_guardrail_translation_mappings = ( + original_endpoint_mappings + ) + + assert result is response + assert result.choices[0].message.content == "Hello" + assert data["metadata"]["applied_guardrails"] == [guardrail.guardrail_name] + + @pytest.mark.asyncio + async def test_output_parse_missing_post_hook_preserves_original_response( + self, caplog + ): + handler = UnifiedLLMGuardrails() + guardrail = OutputParseNoPostHookGuardrail() + response = ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content="Hello"), + index=0, + finish_reason="stop", + ) + ], + ) + + with caplog.at_level("WARNING"): + result = await handler.async_post_call_success_hook( + data={"guardrail_to_apply": guardrail, "model": "gpt-test"}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=response, + ) + + assert result is response + assert "no async_post_call_success_hook" in caplog.text + + @pytest.mark.asyncio + async def test_no_guardrail_to_apply_returns_none_and_removes_key(self): + handler = UnifiedLLMGuardrails() + data = {"guardrail_to_apply": None, "model": "gpt-test"} + + result = await handler.async_post_call_success_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message(role="assistant", content="Hello"), + index=0, + finish_reason="stop", + ) + ], + ), + ) + + assert result is None + assert "guardrail_to_apply" not in data diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 4b50e9a4d31..eb1a8e7cb1a 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -6,9 +6,14 @@ import sys import pytest from fastapi import HTTPException +import litellm from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import ProxyErrorTypes +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import Choices, Message, ModelResponse, Usage sys.path.insert( 0, os.path.abspath("../../..") @@ -190,3 +195,164 @@ def test_get_projected_spend_over_limit_includes_current_spend(monkeypatch): projected_spend, projected_exceeded_date = result assert projected_spend == 290.0 assert projected_exceeded_date == real_datetime.date(2026, 4, 21) + + +class _TrackingPostCallGuardrail(CustomGuardrail): + def __init__( + self, + guardrail_name: str, + label: str, + seen: list[str], + output_parse_pii: bool, + apply_to_output: bool = False, + ): + super().__init__(guardrail_name=guardrail_name, event_hook="post_call") + self.label = label + self.seen = seen + self.output_parse_pii = output_parse_pii + self.apply_to_output = apply_to_output + + def should_run_guardrail(self, data, event_type) -> bool: + return event_type == GuardrailEventHooks.post_call + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.seen.append(self.label) + return response + + +@pytest.mark.asyncio +async def test_post_call_guardrails_preserve_registration_order_by_default(monkeypatch): + monkeypatch.delenv("LITELLM_RUN_OUTPUT_PARSE_PII_LAST", raising=False) + seen: list[str] = [] + callbacks = [ + _TrackingPostCallGuardrail( + guardrail_name="output-parse", + label="output-parse", + seen=seen, + output_parse_pii=True, + ), + _TrackingPostCallGuardrail( + guardrail_name="audit", + label="audit", + seen=seen, + output_parse_pii=False, + ), + ] + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + response = ModelResponse( + id="resp", + choices=[ + Choices( + message=Message(content="ok", role="assistant"), + index=0, + finish_reason="stop", + ) + ], + model="test-model", + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + monkeypatch.setattr(litellm, "callbacks", callbacks) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + await proxy_logging.post_call_success_hook( + data={"model": "test-model"}, + response=response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + + assert seen == ["output-parse", "audit"] + + +@pytest.mark.asyncio +async def test_post_call_guardrails_can_opt_in_to_run_output_parse_last(monkeypatch): + monkeypatch.setenv("LITELLM_RUN_OUTPUT_PARSE_PII_LAST", "true") + seen: list[str] = [] + callbacks = [ + _TrackingPostCallGuardrail( + guardrail_name="output-parse", + label="output-parse", + seen=seen, + output_parse_pii=True, + ), + _TrackingPostCallGuardrail( + guardrail_name="audit", + label="audit", + seen=seen, + output_parse_pii=False, + ), + ] + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + response = ModelResponse( + id="resp", + choices=[ + Choices( + message=Message(content="ok", role="assistant"), + index=0, + finish_reason="stop", + ) + ], + model="test-model", + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + monkeypatch.setattr(litellm, "callbacks", callbacks) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + await proxy_logging.post_call_success_hook( + data={"model": "test-model"}, + response=response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + + assert seen == ["audit", "output-parse"] + + +@pytest.mark.asyncio +async def test_post_call_guardrails_auto_reorder_presidio_residual_masking( + monkeypatch, +): + monkeypatch.delenv("LITELLM_RUN_OUTPUT_PARSE_PII_LAST", raising=False) + seen: list[str] = [] + callbacks = [ + _TrackingPostCallGuardrail( + guardrail_name="presidio", + label="unmask", + seen=seen, + output_parse_pii=True, + ), + _TrackingPostCallGuardrail( + guardrail_name="presidio", + label="residual-mask", + seen=seen, + output_parse_pii=False, + apply_to_output=True, + ), + ] + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + response = ModelResponse( + id="resp", + choices=[ + Choices( + message=Message(content="ok", role="assistant"), + index=0, + finish_reason="stop", + ) + ], + model="test-model", + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + monkeypatch.setattr(litellm, "callbacks", callbacks) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + await proxy_logging.post_call_success_hook( + data={"model": "test-model"}, + response=response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + + assert seen == ["residual-mask", "unmask"] diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 1283d2ccbe7..dbaadc40d60 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1,5 +1,6 @@ import os import sys +from collections import OrderedDict import pytest @@ -10,7 +11,117 @@ sys.path.insert( from unittest.mock import MagicMock -from litellm.proxy.route_llm_request import route_request +import litellm.proxy.route_llm_request as route_llm_request_module +from litellm.proxy.route_llm_request import ( + _clear_user_config_router_cache, + _kwargs_for_llm, + route_request, +) + + +@pytest.fixture(autouse=True) +def clear_user_config_router_cache(): + _clear_user_config_router_cache() + yield + _clear_user_config_router_cache() + + +def test_kwargs_for_llm_strips_presidio_pii_tokens(): + """_kwargs_for_llm removes _presidio_pii_tokens so it is not sent to the LLM provider.""" + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "_presidio_pii_tokens": {"guardrail_1": {"": "Jane"}}, + } + result = _kwargs_for_llm(data) + assert "model" in result + assert "messages" in result + assert "_presidio_pii_tokens" not in result + assert result.get("_presidio_pii_tokens") is None + + +def test_kwargs_for_llm_preserves_other_keys(): + """_kwargs_for_llm leaves all other keys unchanged.""" + data = {"model": "gpt-4", "temperature": 0.7, "api_key": "sk-xxx"} + result = _kwargs_for_llm(data) + assert result == data + assert result is data + + +def test_kwargs_for_llm_strips_fastest_response(): + data = {"model": "gpt-4", "fastest_response": True, "messages": []} + result = _kwargs_for_llm(data) + assert "fastest_response" not in result + assert result["model"] == "gpt-4" + + +def test_kwargs_for_llm_strips_user_config(): + data = { + "model": "gpt-4", + "api_key": "sk-xxx", + "user_config": {"model_list": [{"model_name": "x"}]}, + } + result = _kwargs_for_llm(data) + assert "user_config" not in result + assert result["model"] == "gpt-4" + + +def test_discard_router_safely_swallows_discard_errors(): + class _BadRouter: + def discard(self): + raise RuntimeError("discard failed") + + route_llm_request_module._discard_router_safely(_BadRouter()) + + +def test_freeze_user_config_cache_key_stabilizes_nested_values(): + frozen = route_llm_request_module._freeze_user_config_cache_key( + { + "b": {3, 1}, + "a": [ + {"y": ("two", "one"), "x": 1}, + {"z": ["nested"]}, + ], + } + ) + + assert frozen == ( + ( + "a", + ( + (("x", 1), ("y", ("two", "one"))), + (("z", ("nested",)),), + ), + ), + ("b", (1, 3)), + ) + + +def test_get_user_config_router_cache_key_wraps_non_tuple(monkeypatch): + monkeypatch.setattr( + route_llm_request_module, + "_freeze_user_config_cache_key", + lambda value: "scalar-key", + ) + + assert route_llm_request_module._get_user_config_router_cache_key({}) == ( + "scalar-key", + ) + + +def test_prune_expired_user_config_routers_removes_only_expired(): + route_llm_request_module._USER_CONFIG_ROUTER_CACHE = OrderedDict( + { + ("expired",): (object(), 10.0), + ("active",): (object(), 20.0), + } + ) + + route_llm_request_module._prune_expired_user_config_routers(15.0) + + assert list(route_llm_request_module._USER_CONFIG_ROUTER_CACHE.keys()) == [ + ("active",) + ] @pytest.mark.parametrize( @@ -168,7 +279,9 @@ async def test_route_request_with_router_settings_override(): assert call_kwargs["fallbacks"] == [{"gpt-3.5-turbo": ["gpt-4"]}] assert call_kwargs["num_retries"] == 5 assert call_kwargs["timeout"] == 30 - assert call_kwargs["model_group_retry_policy"] == {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}} + assert call_kwargs["model_group_retry_policy"] == { + "gpt-3.5-turbo": {"RateLimitErrorRetries": 3} + } # Verify unsupported settings were NOT merged assert "routing_strategy" not in call_kwargs assert "model_group_alias" not in call_kwargs @@ -237,3 +350,548 @@ async def test_route_request_with_router_settings_override_preserves_existing(): assert call_kwargs["num_retries"] == 10 # Key/team timeout should be applied since not in request assert call_kwargs["timeout"] == 30 + + +@pytest.mark.asyncio +async def test_route_request_user_config_filters_router_args_and_reuses_cached_router( + monkeypatch, +): + import litellm + + state = {"init_kwargs": None, "init_count": 0, "discard_called": False} + + class _FakeRouter: + @staticmethod + def get_valid_args(): + return ["model_list", "routing_strategy"] + + def __init__(self, **kwargs): + state["init_count"] += 1 + state["init_kwargs"] = kwargs + + async def acompletion(self, **kwargs): + return {"ok": True, "kwargs": kwargs} + + def discard(self): + state["discard_called"] = True + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + + data_1 = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "_presidio_pii_tokens": {"guardrail": {"": "Jane"}}, + "user_config": { + "model_list": [ + {"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4"}} + ], + "routing_strategy": "least-busy", + "invalid_key": "should_be_ignored", + }, + } + + data_2 = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello again"}], + "fastest_response": True, + "_presidio_pii_tokens": {"guardrail": {"": "Bob"}}, + "user_config": { + "model_list": [ + {"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4"}} + ], + "routing_strategy": "least-busy", + "invalid_key": "should_be_ignored", + }, + } + + llm_call_1 = await route_request(data_1, None, None, "acompletion") + result_1 = await llm_call_1 + llm_call_2 = await route_request(data_2, None, None, "acompletion") + result_2 = await llm_call_2 + + assert result_1["ok"] is True + assert result_2["ok"] is True + assert "_presidio_pii_tokens" not in result_1["kwargs"] + assert "_presidio_pii_tokens" not in result_2["kwargs"] + assert "fastest_response" not in result_2["kwargs"] + assert state["init_kwargs"] == { + "model_list": [ + {"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4"}} + ], + "routing_strategy": "least-busy", + } + assert state["init_count"] == 1 + assert state["discard_called"] is False + + +@pytest.mark.asyncio +async def test_route_request_user_config_batch_does_not_forward_model_kwarg( + monkeypatch, +): + import litellm + + state = {"models": None, "kwargs": None} + + class _FakeRouter: + @staticmethod + def get_valid_args(): + return [] + + def __init__(self, **kwargs): + pass + + async def abatch_completion(self, models, **kwargs): + state["models"] = models + state["kwargs"] = kwargs + return "ok" + + def discard(self): + pass + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + + data = { + "model": "gpt-4o-mini, claude-3-haiku", + "messages": [{"role": "user", "content": "hi"}], + "user_config": {}, + } + + llm_call = await route_request(data, None, None, "acompletion") + result = await llm_call + + assert result == "ok" + assert state["models"] == ["gpt-4o-mini", "claude-3-haiku"] + assert "model" not in (state["kwargs"] or {}) + assert "fastest_response" not in (state["kwargs"] or {}) + assert data["model"] == "gpt-4o-mini, claude-3-haiku" + + +@pytest.mark.asyncio +async def test_route_request_user_config_batch_fastest_response_uses_models_list( + monkeypatch, +): + import litellm + + state = {"models": None, "kwargs": None} + + class _FakeRouter: + @staticmethod + def get_valid_args(): + return [] + + def __init__(self, **kwargs): + pass + + async def abatch_completion_fastest_response(self, models, **kwargs): + state["models"] = models + state["kwargs"] = kwargs + return "ok" + + def discard(self): + pass + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + + data = { + "model": "gpt-4o-mini, claude-3-haiku", + "messages": [{"role": "user", "content": "hi"}], + "fastest_response": True, + "user_config": {}, + } + + llm_call = await route_request(data, None, None, "acompletion") + result = await llm_call + + assert result == "ok" + assert state["models"] == ["gpt-4o-mini", "claude-3-haiku"] + assert "model" not in (state["kwargs"] or {}) + assert data["model"] == "gpt-4o-mini, claude-3-haiku" + + +@pytest.mark.asyncio +async def test_route_request_pops_user_config_before_returning_coroutine(monkeypatch): + import litellm + + class _FakeRouter: + @staticmethod + def get_valid_args(): + return ["model_list"] + + def __init__(self, **kwargs): + pass + + async def acompletion(self, **kwargs): + return {"ok": True} + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "user_config": {"model_list": [{"model_name": "x"}]}, + } + + llm_call = await route_request(data, None, None, "acompletion") + + assert "user_config" not in data + assert llm_call is not None + await llm_call + + +@pytest.mark.asyncio +async def test_route_request_user_config_builds_router_via_to_thread(monkeypatch): + import litellm + + captured = {"func": None, "args": None} + + class _FakeRouter: + @staticmethod + def get_valid_args(): + return ["model_list"] + + async def acompletion(self, **kwargs): + return {"ok": True} + + async def _fake_to_thread(func, *args, **kwargs): + captured["func"] = func + captured["args"] = args + return _FakeRouter() + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + monkeypatch.setattr(route_llm_request_module.asyncio, "to_thread", _fake_to_thread) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "user_config": {"model_list": [{"model_name": "x"}]}, + } + + llm_call = await route_request(data, None, None, "acompletion") + result = await llm_call + + assert result == {"ok": True} + assert ( + captured["func"] is route_llm_request_module._get_or_create_user_config_router + ) + assert captured["args"] == ({"model_list": [{"model_name": "x"}]},) + + +@pytest.mark.asyncio +async def test_route_request_user_config_router_cache_evicts_lru(monkeypatch): + import litellm + + state = {"init_count": 0, "discard_count": 0} + + class _FakeRouter: + @staticmethod + def get_valid_args(): + return ["model_list"] + + def __init__(self, **kwargs): + state["init_count"] += 1 + + async def acompletion(self, **kwargs): + return {"ok": True} + + def discard(self): + state["discard_count"] += 1 + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + monkeypatch.setattr( + route_llm_request_module, "_USER_CONFIG_ROUTER_CACHE_MAX_SIZE", 1 + ) + monkeypatch.setattr( + route_llm_request_module, "_USER_CONFIG_ROUTER_CACHE_TTL_SECONDS", 3600 + ) + + data_a_1 = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "a1"}], + "user_config": {"model_list": [{"model_name": "a"}]}, + } + data_b = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "b"}], + "user_config": {"model_list": [{"model_name": "b"}]}, + } + data_a_2 = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "a2"}], + "user_config": {"model_list": [{"model_name": "a"}]}, + } + + llm_call = await route_request(data_a_1, None, None, "acompletion") + await llm_call + llm_call = await route_request(data_b, None, None, "acompletion") + await llm_call + llm_call = await route_request(data_a_2, None, None, "acompletion") + await llm_call + + assert state["init_count"] == 3 + assert state["discard_count"] == 0 + + +@pytest.mark.asyncio +async def test_route_request_user_config_router_cache_expires_by_ttl(monkeypatch): + import litellm + + state = {"init_count": 0, "discard_count": 0} + clock = {"t": 1000.0} + + class _FakeRouter: + @staticmethod + def get_valid_args(): + return ["model_list"] + + def __init__(self, **kwargs): + state["init_count"] += 1 + + async def acompletion(self, **kwargs): + return {"ok": True} + + def discard(self): + state["discard_count"] += 1 + + def _fake_monotonic(): + return clock["t"] + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + monkeypatch.setattr( + route_llm_request_module, "_USER_CONFIG_ROUTER_CACHE_MAX_SIZE", 16 + ) + monkeypatch.setattr( + route_llm_request_module, "_USER_CONFIG_ROUTER_CACHE_TTL_SECONDS", 1 + ) + monkeypatch.setattr(route_llm_request_module.time, "monotonic", _fake_monotonic) + + data_1 = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "first"}], + "user_config": {"model_list": [{"model_name": "same"}]}, + } + data_2 = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "second"}], + "user_config": {"model_list": [{"model_name": "same"}]}, + } + + llm_call = await route_request(data_1, None, None, "acompletion") + await llm_call + + clock["t"] = 1002.0 # beyond ttl + llm_call = await route_request(data_2, None, None, "acompletion") + await llm_call + + assert state["init_count"] == 2 + assert state["discard_count"] == 0 + + +def test_user_config_router_cache_discards_new_router_on_cache_block_exception( + monkeypatch, +): + import litellm + + state = {"init_count": 0, "discard_count": 0} + + class _FakeRouter: + def __init__(self, **kwargs): + state["init_count"] += 1 + + def discard(self): + state["discard_count"] += 1 + + class _FailingCache(OrderedDict): + def move_to_end(self, key, last=True): + raise RuntimeError("cache move_to_end failed") + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + monkeypatch.setattr( + route_llm_request_module, + "_USER_CONFIG_ROUTER_CACHE", + _FailingCache(), + ) + + with pytest.raises(RuntimeError, match="cache move_to_end failed"): + route_llm_request_module._get_or_create_user_config_router( + {"model_list": [{"model_name": "x"}]} + ) + + assert state["init_count"] == 1 + assert state["discard_count"] == 1 + assert len(route_llm_request_module._USER_CONFIG_ROUTER_CACHE) == 0 + + +def test_user_config_router_cache_discards_duplicate_new_router_after_double_check( + monkeypatch, +): + import litellm + + state = {"init_count": 0, "discard_count": 0} + created_routers = [] + + class _FakeRouter: + def __init__(self, **kwargs): + state["init_count"] += 1 + self.name = f"router-{state['init_count']}" + created_routers.append(self) + if state["init_count"] == 1: + cache_key = route_llm_request_module._get_user_config_router_cache_key( + kwargs + ) + route_llm_request_module._USER_CONFIG_ROUTER_CACHE[cache_key] = ( + _FakeRouter.__new__(_FakeRouter), + float("inf"), + ) + route_llm_request_module._USER_CONFIG_ROUTER_CACHE[cache_key][ + 0 + ].name = "existing-router" + + def discard(self): + state["discard_count"] += 1 + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + + router = route_llm_request_module._get_or_create_user_config_router( + {"model_list": [{"model_name": "x"}]} + ) + + assert router.name == "existing-router" + assert state["init_count"] == 1 + assert state["discard_count"] == 1 + assert len(created_routers) == 1 + + +def test_user_config_router_cache_reuses_cached_entry_and_refreshes_ttl(monkeypatch): + import litellm + + state = {"init_count": 0} + clock = {"t": 1000.0} + + class _FakeRouter: + def __init__(self, **kwargs): + state["init_count"] += 1 + + def _fake_monotonic(): + return clock["t"] + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + monkeypatch.setattr(route_llm_request_module.time, "monotonic", _fake_monotonic) + monkeypatch.setattr( + route_llm_request_module, "_USER_CONFIG_ROUTER_CACHE_TTL_SECONDS", 5 + ) + + router = object() + cache_key = route_llm_request_module._get_user_config_router_cache_key( + {"model_list": [{"model_name": "x"}]} + ) + route_llm_request_module._USER_CONFIG_ROUTER_CACHE = OrderedDict( + {cache_key: (router, 1001.0)} + ) + + result = route_llm_request_module._get_or_create_user_config_router( + {"model_list": [{"model_name": "x"}]} + ) + + assert result is router + assert state["init_count"] == 0 + assert route_llm_request_module._USER_CONFIG_ROUTER_CACHE[cache_key] == ( + router, + 1005.0, + ) + + +def test_user_config_router_cache_logs_aggregated_eviction_warning(monkeypatch): + import litellm + + class _FakeRouter: + def __init__(self, **kwargs): + self.kwargs = kwargs + + warnings = [] + + monkeypatch.setattr(litellm, "Router", _FakeRouter) + monkeypatch.setattr( + route_llm_request_module, "_USER_CONFIG_ROUTER_CACHE_MAX_SIZE", 1 + ) + monkeypatch.setattr( + route_llm_request_module.verbose_proxy_logger, + "warning", + lambda msg, count: warnings.append((msg, count)), + ) + + route_llm_request_module._get_or_create_user_config_router( + {"model_list": [{"model_name": "a"}]} + ) + route_llm_request_module._get_or_create_user_config_router( + {"model_list": [{"model_name": "b"}]} + ) + + assert warnings == [ + ( + "user_config Router cache full (evicted %d entries). " + "Increase LITELLM_USER_CONFIG_ROUTER_CACHE_MAX_SIZE if this is frequent.", + 1, + ) + ] + + +@pytest.mark.asyncio +async def test_route_request_batch_with_router_does_not_forward_model_kwarg(): + data = { + "model": "gpt-4o-mini, claude-3-haiku", + "messages": [{"role": "user", "content": "hi"}], + } + llm_router = MagicMock() + llm_router.model_names = [] + llm_router.abatch_completion.return_value = "ok" + + response = await route_request(data, llm_router, None, "acompletion") + + assert response == "ok" + call_kwargs = llm_router.abatch_completion.call_args[1] + assert call_kwargs["models"] == ["gpt-4o-mini", "claude-3-haiku"] + assert "model" not in call_kwargs + + +@pytest.mark.asyncio +async def test_route_request_batch_fastest_response_with_router_uses_models_list(): + data = { + "model": "gpt-4o-mini, claude-3-haiku", + "messages": [{"role": "user", "content": "hi"}], + "fastest_response": True, + } + llm_router = MagicMock() + llm_router.model_names = [] + llm_router.abatch_completion_fastest_response.return_value = "ok" + + response = await route_request(data, llm_router, None, "acompletion") + + assert response == "ok" + call_kwargs = llm_router.abatch_completion_fastest_response.call_args[1] + assert call_kwargs["models"] == ["gpt-4o-mini", "claude-3-haiku"] + assert "model" not in call_kwargs + assert "fastest_response" not in call_kwargs + assert data["model"] == "gpt-4o-mini, claude-3-haiku" + + +@pytest.mark.asyncio +async def test_route_request_evals_path_strips_internal_keys(): + import litellm + + data = { + "name": "eval-test", + "_presidio_pii_tokens": {"guardrail": {"": "Alice"}}, + } + llm_router = MagicMock() + + original_func = litellm.acreate_eval + mock_create_eval = MagicMock(return_value={"ok": True}) + litellm.acreate_eval = mock_create_eval + try: + response = await route_request(data, llm_router, None, "acreate_eval") + assert response == {"ok": True} + call_kwargs = mock_create_eval.call_args[1] + assert "_presidio_pii_tokens" not in call_kwargs + assert call_kwargs["name"] == "eval-test" + finally: + litellm.acreate_eval = original_func