mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-22 18:25:39 +00:00
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 <yuneng-jiang@users.noreply.github.com>
This commit is contained in:
co-authored by
yuneng-jiang
parent
a5b86d3b2f
commit
a9e45e70ea
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user