diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index f8d1710dfdc..a09a2afe26e 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -701,8 +701,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the signed headers signed_headers = dict(aws_request.headers.items()) + request_url = prepped.url or url response = await self.async_httpx_client.get( - prepped.url, headers=signed_headers + request_url, headers=signed_headers ) if response.status_code != 200: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index e048ca21cba..67cb281029c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -433,6 +433,109 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # contain API keys or other secrets) in error responses. raise Exception(f"Presidio PII analysis failed: {type(e).__name__}") from e + async def _post_presidio_anonymize( + self, text: str, analyze_results: Any + ) -> Any: + """POST to Presidio anonymize; returns parsed JSON body.""" + # Use shared session to prevent memory leak (issue #14540) + async with self._get_session_iterator() as session: + anonymize_url = f"{self.presidio_anonymizer_api_base}anonymize" + verbose_proxy_logger.debug("Making request to: %s", anonymize_url) + anonymize_payload = { + "text": text, + "analyzer_results": analyze_results, + } + async with session.post( + anonymize_url, + json=anonymize_payload, + headers={"Accept": "application/json"}, + ) as response: + if response.status >= 400: + error_body = await response.text() + raise Exception( + f"Presidio anonymizer returned HTTP {response.status}: {error_body[:200]}" + ) + content_type = getattr( + response, + "content_type", + response.headers.get("Content-Type", ""), + ) + if "application/json" not in content_type: + error_body = await response.text() + raise Exception( + f"Presidio anonymizer returned non-JSON Content-Type '{content_type}'; body: '{error_body[:200]}'" + ) + return await response.json() + + def _finalize_presidio_anonymize_simple( + self, + redacted_text: Dict[str, Any], + masked_entity_count: Dict[str, int], + ) -> str: + # No need to build numbered tokens — just use Presidio's + # already-anonymized text directly. The old code incorrectly + # applied anonymizer item positions (which reference the + # *output* text) to the *original* text, causing offset errors. + for item in redacted_text.get("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 + ) + return redacted_text["text"] + + def _finalize_presidio_anonymize_numbered_tokens( + self, + text: str, + analyze_results: Any, + request_data: Optional[Dict], + masked_entity_count: Dict[str, int], + ) -> str: + # output_parse_pii is True — we need sequentially numbered + # tokens and a pii_tokens mapping for later unmasking. + # Use analyze_results positions (which reference the ORIGINAL + # text) instead of anonymizer items (which reference the output). + new_text = text + 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 = {} + 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"] + + # Assign sequence numbers in forward (left-to-right) order so + # that is the first entity in the text, etc. + sorted_forward = sorted(analyze_results, key=lambda x: x["start"]) + seq_map = {} + for idx, ar in enumerate(sorted_forward, start=1): + seq_map[(ar["start"], ar["end"])] = idx + + # Apply replacements in reverse order by start position so + # that replacing later spans first does not shift earlier + # coordinates in the original text. + for ar in reversed(sorted_forward): + start = ar["start"] + end = ar["end"] + entity_type = ar["entity_type"] + replacement = f"<{entity_type}>" + seq = seq_map[(start, end)] + if replacement.endswith(">"): + replacement = f"{replacement[:-1]}_{seq}>" + else: + replacement = f"{replacement}_{seq}" + pii_tokens[replacement] = text[start:end] + new_text = new_text[:start] + replacement + new_text[end:] + masked_entity_count[entity_type] = ( + masked_entity_count.get(entity_type, 0) + 1 + ) + return new_text + async def anonymize_text( self, text: str, @@ -449,110 +552,20 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if isinstance(analyze_results, list) and len(analyze_results) == 0: return text - # Use shared session to prevent memory leak (issue #14540) - async with self._get_session_iterator() as session: - # Make the request to /anonymize - anonymize_url = f"{self.presidio_anonymizer_api_base}anonymize" - verbose_proxy_logger.debug("Making request to: %s", anonymize_url) - anonymize_payload = { - "text": text, - "analyzer_results": analyze_results, - } - - async with session.post( - anonymize_url, - json=anonymize_payload, - headers={"Accept": "application/json"}, - ) as response: - # Validate HTTP status - if response.status >= 400: - error_body = await response.text() - raise Exception( - f"Presidio anonymizer returned HTTP {response.status}: {error_body[:200]}" - ) - - # Validate Content-Type is JSON - content_type = getattr( - response, - "content_type", - response.headers.get("Content-Type", ""), - ) - if "application/json" not in content_type: - error_body = await response.text() - raise Exception( - f"Presidio anonymizer returned non-JSON Content-Type '{content_type}'; body: '{error_body[:200]}'" - ) - - redacted_text = await response.json() - - if redacted_text is not None: - verbose_proxy_logger.debug("redacted_text: %s", redacted_text) - - if not output_parse_pii: - # No need to build numbered tokens — just use Presidio's - # already-anonymized text directly. The old code incorrectly - # applied anonymizer item positions (which reference the - # *output* text) to the *original* text, causing offset errors. - for item in redacted_text.get("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 - ) - return redacted_text["text"] - - # output_parse_pii is True — we need sequentially numbered - # tokens and a pii_tokens mapping for later unmasking. - # Use analyze_results positions (which reference the ORIGINAL - # text) instead of anonymizer items (which reference the output). - new_text = text - 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 = {} - 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"] - - # Assign sequence numbers in forward (left-to-right) order so - # that is the first entity in the text, etc. - sorted_forward = sorted( - analyze_results, key=lambda x: x["start"] - ) - seq_map = {} - for idx, ar in enumerate(sorted_forward, start=1): - seq_map[(ar["start"], ar["end"])] = idx - - # Apply replacements in reverse order by start position so - # that replacing later spans first does not shift earlier - # coordinates in the original text. - for ar in reversed(sorted_forward): - start = ar["start"] - end = ar["end"] - entity_type = ar["entity_type"] - replacement = f"<{entity_type}>" - - seq = seq_map[(start, end)] - if replacement.endswith(">"): - replacement = f"{replacement[:-1]}_{seq}>" - else: - replacement = f"{replacement}_{seq}" - - pii_tokens[replacement] = text[start:end] - new_text = new_text[:start] + replacement + new_text[end:] - - masked_entity_count[entity_type] = ( - masked_entity_count.get(entity_type, 0) + 1 - ) - - return new_text - else: + redacted_text = await self._post_presidio_anonymize(text, analyze_results) + if redacted_text is None: raise Exception("Invalid anonymizer response: received None") + + verbose_proxy_logger.debug("redacted_text: %s", redacted_text) + + if not output_parse_pii: + return self._finalize_presidio_anonymize_simple( + redacted_text, masked_entity_count + ) + + return self._finalize_presidio_anonymize_numbered_tokens( + text, analyze_results, request_data, masked_entity_count + ) except Exception as e: # Sanitize exception to avoid leaking the original text (which may # contain API keys or other secrets) in error responses. diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index daabed0def1..c32a1bdd463 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -104,7 +104,10 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): ) assert response.status_code == 200 - assert response.json() == {"redirect_url": "http://testserver/ui/?login=success"} + assert response.json() == { + "redirect_url": "http://testserver/ui/?login=success", + "token": "signed-token", + } assert response.cookies.get("token") == "signed-token" mock_authenticate_user.assert_awaited_once_with(