From a9e45e70ea634bbe06d96cb3e145b30072b5f83b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 13 Mar 2026 05:31:25 +0000 Subject: [PATCH] fix: revert presidio streaming type changes (unsafe cast) Revert the return type narrowing and cast() calls in async_post_call_streaming_iterator_hook. The internal generators _stream_apply_output_masking and _stream_pii_unmasking genuinely yield bytes objects for Anthropic native SSE chunks. Casting them to ModelResponseStream masks a real design issue. Restore the original Union[ModelResponseStream, bytes] return type and accept the known mypy override error for now. Co-authored-by: yuneng-jiang --- litellm/proxy/guardrails/guardrail_hooks/presidio.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index ce2419e97c..5701f7dd9e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1234,7 +1234,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, response: Any, request_data: dict, - ) -> AsyncGenerator[ModelResponseStream, None]: + ) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]: """ Process streaming response chunks to unmask PII tokens when needed. """ @@ -1242,7 +1242,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async for chunk in self._stream_apply_output_masking( response, request_data ): - yield cast(ModelResponseStream, chunk) + yield chunk return metadata = (request_data.get("metadata") or {}) if request_data else {} @@ -1257,7 +1257,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return async for chunk in self._stream_pii_unmasking(response, request_data): - yield cast(ModelResponseStream, chunk) + yield chunk @staticmethod def _preserve_usage_from_last_chunk(