scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(lint): comply with LIT009 type discipline and TQ003 test quality gates

Browse source
This commit is contained in:
Ilya Sibiryakov 2026-09-30 14:09:36 +03:00
parent d44d99c895
commit cf105fd8b2
2 changed files with 16 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/proxy/guardrails/guardrail_hooks/ztds.py
View file

@ -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)

10
tests/guardrails_tests/test_ztds_guardrail.py
View file

@ -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: