From cf105fd8b2377b11e69ded50572db4fae51fde99 Mon Sep 17 00:00:00 2001 From: Ilya Sibiryakov Date: Wed, 30 Sep 2026 14:09:36 +0300 Subject: [PATCH] fix(lint): comply with LIT009 type discipline and TQ003 test quality gates --- .../proxy/guardrails/guardrail_hooks/ztds.py | 31 +++++++++---------- tests/guardrails_tests/test_ztds_guardrail.py | 10 +----- 2 files changed, 16 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ztds.py b/litellm/proxy/guardrails/guardrail_hooks/ztds.py index 69b1f29d486..691003727d7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ztds.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ztds.py @@ -16,7 +16,7 @@ from __future__ import annotations import re import uuid -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterable from typing import ClassVar try: @@ -273,7 +273,7 @@ class ZTDSGuardrail(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: object, - response: object, + response: AsyncIterable[object], request_data: dict[str, object], ) -> AsyncGenerator[object, None]: """ @@ -282,20 +282,19 @@ class ZTDSGuardrail(CustomGuardrail): """ session_id = request_data.get("_ztds_session_id") if isinstance(request_data, dict) else None try: - if hasattr(response, "__aiter__"): - async for chunk in response: # type: ignore[union-attr] - if session_id and isinstance(session_id, str) and self.reverse_on_output: - if hasattr(chunk, "choices") and chunk.choices: - for choice in chunk.choices: - delta = getattr(choice, "delta", None) - if delta and hasattr(delta, "content") and isinstance(delta.content, str): - delta.content = self.restore_text(delta.content, session_id) - elif isinstance(chunk, dict) and "choices" in chunk: - for choice in chunk["choices"]: - delta = choice.get("delta") if isinstance(choice, dict) else None - if delta and isinstance(delta, dict) and isinstance(delta.get("content"), str): - delta["content"] = self.restore_text(delta["content"], session_id) - yield chunk + async for chunk in response: + if session_id and isinstance(session_id, str) and self.reverse_on_output: + if hasattr(chunk, "choices") and chunk.choices: + for choice in chunk.choices: + delta = getattr(choice, "delta", None) + if delta and hasattr(delta, "content") and isinstance(delta.content, str): + delta.content = self.restore_text(delta.content, session_id) + elif isinstance(chunk, dict) and "choices" in chunk: + for choice in chunk["choices"]: + delta = choice.get("delta") if isinstance(choice, dict) else None + if delta and isinstance(delta, dict) and isinstance(delta.get("content"), str): + delta["content"] = self.restore_text(delta["content"], session_id) + yield chunk finally: if session_id and isinstance(session_id, str): self.zeroize_session(session_id) diff --git a/tests/guardrails_tests/test_ztds_guardrail.py b/tests/guardrails_tests/test_ztds_guardrail.py index 40e6969cbce..2961807a665 100644 --- a/tests/guardrails_tests/test_ztds_guardrail.py +++ b/tests/guardrails_tests/test_ztds_guardrail.py @@ -6,17 +6,9 @@ https://datatracker.ietf.org/doc/draft-sibiryakov-ztds-protocol/ from __future__ import annotations -import sys import unittest -from pathlib import Path -try: - from litellm.proxy.guardrails.guardrail_hooks.ztds import ZTDSGuardrail -except (ImportError, ModuleNotFoundError): - hook_dir = Path(__file__).resolve().parents[2] / "litellm" / "proxy" / "guardrails" / "guardrail_hooks" - if str(hook_dir) not in sys.path: - sys.path.insert(0, str(hook_dir)) - from ztds import ZTDSGuardrail +from litellm.proxy.guardrails.guardrail_hooks.ztds import ZTDSGuardrail class MockMessage: