mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(zscaler_ai_guard)!: scan only the last user turn on requests (#50139)
* fix(zscaler_ai_guard): scan only the last user turn on requests * fix(zscaler_ai_guard): keep scanning A2A requests and refuse scan_only_tool_results * fix(zscaler_ai_guard): join the last user turn's text parts with spaces like before * fix(zscaler_ai_guard): keep scanning tool results when scan_only_tool_results is set * fix(zscaler_ai_guard): scan the latest user turn when assistant or tool messages follow it * fix(zscaler_ai_guard): skip the scan when the latest user turn is only whitespace A blank latest user turn was sent to Zscaler on its own, which answers Content is empty, so the proxy returned 500 where main returned 200 --------- Co-authored-by: Arnav Maiti <amaiti@zscaler.com>
This commit is contained in:
yucheng-berri
GitHub
parent
7cdaa73fc0
commit
ff2709f5c3
2 changed files with 252 additions and 5 deletions
|
|
@ -4,7 +4,8 @@
|
|||
#
|
||||
# +-------------------------------------------------------------+
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from itertools import chain, dropwhile, takewhile
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -14,6 +15,10 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
effective_scan_only_tool_results_for_guardrail,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
|
|
@ -218,15 +223,14 @@ class ZscalerAIGuard(CustomGuardrail):
|
|||
zscaler_ai_guard_result = None
|
||||
direction: Final = "OUT" if input_type == "response" else "IN"
|
||||
verbose_proxy_logger.debug("direction: %s", direction)
|
||||
# Concatenate all texts and send to Zscaler AI Guard
|
||||
if texts:
|
||||
concatenated_text: Final = " ".join(texts)
|
||||
content: Final = self._content_to_scan(inputs, input_type)
|
||||
if content:
|
||||
zscaler_ai_guard_result = await self.make_zscaler_ai_guard_api_call(
|
||||
zscaler_ai_guard_url=self.zscaler_ai_guard_url,
|
||||
api_key=self.api_key,
|
||||
policy_id=policy_id,
|
||||
direction=direction,
|
||||
content=concatenated_text,
|
||||
content=content,
|
||||
user=user_for_payload,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
@ -244,6 +248,37 @@ class ZscalerAIGuard(CustomGuardrail):
|
|||
verbose_proxy_logger.debug("ZscalerAIGuard: Successfully applied guardrail.")
|
||||
return inputs
|
||||
|
||||
def _content_to_scan(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
input_type: Literal["request", "response"],
|
||||
) -> str | None:
|
||||
structured_messages: Final = inputs.get("structured_messages") or []
|
||||
scans_last_user_turn: Final = (
|
||||
input_type == "request"
|
||||
and not effective_scan_only_tool_results_for_guardrail(self)
|
||||
and any("content" in message for message in structured_messages)
|
||||
)
|
||||
texts: Final = (
|
||||
self._last_user_turn_texts(structured_messages) if scans_last_user_turn else inputs.get("texts", [])
|
||||
)
|
||||
return " ".join(texts) if texts else None
|
||||
|
||||
@staticmethod
|
||||
def _last_user_turn_texts(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
|
||||
from_last_user_message: Final = dropwhile(lambda message: message.get("role") != "user", reversed(messages))
|
||||
last_user_turn: Final = tuple(takewhile(lambda message: message.get("role") == "user", from_last_user_message))
|
||||
texts: Final = tuple(chain.from_iterable(map(ZscalerAIGuard._content_texts, reversed(last_user_turn))))
|
||||
return texts if any(text.strip() for text in texts) else ()
|
||||
|
||||
@staticmethod
|
||||
def _content_texts(message: Mapping[str, object]) -> tuple[str, ...]:
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return (content,)
|
||||
parts: Final = (as_str_mapping(part) for part in as_object_list(content) or ())
|
||||
return tuple(text for part in parts if part is not None and isinstance(text := part.get("text"), str))
|
||||
|
||||
def extract_blocking_info(self, response):
|
||||
"""
|
||||
Extracts transaction ID and blocking detector details from a response.
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ from fastapi import HTTPException
|
|||
|
||||
import litellm
|
||||
from litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import ZscalerAIGuard
|
||||
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
|
||||
|
||||
|
|
@ -141,6 +143,216 @@ async def test_apply_guardrail_text_concatenation(mock_api_call):
|
|||
assert call_args.kwargs["content"] == "Hello world"
|
||||
|
||||
|
||||
_CONVERSATION: Final = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is two plus two?"},
|
||||
{"role": "assistant", "content": "Four."},
|
||||
]
|
||||
|
||||
|
||||
async def _execute_policy_bodies(
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
input_type: Literal["request", "response"],
|
||||
scan_only_tool_results: bool = False,
|
||||
) -> tuple[object, ...]:
|
||||
received: Final = asyncio.Queue[object]()
|
||||
|
||||
async def execute_policy(request: web.Request) -> web.Response:
|
||||
await received.put(await request.json())
|
||||
return web.json_response({"statusCode": 200, "action": "ALLOW"})
|
||||
|
||||
app: Final = web.Application()
|
||||
app.router.add_post("/v1/detection/execute-policy", execute_policy)
|
||||
handler: Final = InMemoryGuardrailHandler()
|
||||
async with TestServer(app) as server:
|
||||
guardrail_id: Final = handler.initialize_guardrail(
|
||||
guardrail={
|
||||
"guardrail_name": "zscaler-last-user-turn",
|
||||
"litellm_params": {
|
||||
"guardrail": "zscaler_ai_guard",
|
||||
"mode": "pre_call",
|
||||
"api_key": "test_key",
|
||||
"api_base": str(server.make_url("/v1/detection/execute-policy")),
|
||||
"policy_id": 1,
|
||||
"scan_only_tool_results": scan_only_tool_results,
|
||||
},
|
||||
},
|
||||
)["guardrail_id"]
|
||||
try:
|
||||
guardrail: Final = handler.guardrail_id_to_custom_guardrail[guardrail_id]
|
||||
assert guardrail is not None
|
||||
returned: Final = await guardrail.apply_guardrail(inputs, {}, input_type)
|
||||
finally:
|
||||
handler.delete_in_memory_guardrail(guardrail_id)
|
||||
|
||||
assert returned is inputs
|
||||
return tuple(received.get_nowait() for _ in range(received.qsize()))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("inputs", "expected_content"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"texts": ["every", "message", "flattened"],
|
||||
"structured_messages": [*_CONVERSATION, {"role": "user", "content": "And three plus three?"}],
|
||||
},
|
||||
"And three plus three?",
|
||||
),
|
||||
(
|
||||
{
|
||||
"texts": ["every", "message", "flattened"],
|
||||
"structured_messages": [
|
||||
*_CONVERSATION,
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "And three"}, {"type": "text", "text": "plus three?"}],
|
||||
},
|
||||
],
|
||||
},
|
||||
"And three plus three?",
|
||||
),
|
||||
(
|
||||
{
|
||||
"texts": ["every", "message", "flattened"],
|
||||
"structured_messages": [
|
||||
*_CONVERSATION,
|
||||
{"role": "user", "content": "And three plus three?"},
|
||||
{"role": "user", "content": "Answer in words."},
|
||||
],
|
||||
},
|
||||
"And three plus three? Answer in words.",
|
||||
),
|
||||
(
|
||||
{
|
||||
"texts": ["Summarize", "this report"],
|
||||
"structured_messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "Summarize"}, {"kind": "text", "text": "this report"}],
|
||||
"messageId": "msg-1",
|
||||
}
|
||||
],
|
||||
},
|
||||
"Summarize this report",
|
||||
),
|
||||
({"texts": ["first input", "second input"]}, "first input second input"),
|
||||
(
|
||||
{
|
||||
"texts": ["every", "message"],
|
||||
"structured_messages": [*_CONVERSATION, {"role": "assistant", "content": "Prefilled answer:"}],
|
||||
},
|
||||
"What is two plus two?",
|
||||
),
|
||||
(
|
||||
{
|
||||
"texts": ["every", "message"],
|
||||
"structured_messages": [
|
||||
*_CONVERSATION,
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "tool output"},
|
||||
],
|
||||
},
|
||||
"What is two plus two?",
|
||||
),
|
||||
(
|
||||
{
|
||||
"texts": ["every", "message"],
|
||||
"structured_messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "assistant", "content": "Prefilled answer:"},
|
||||
],
|
||||
},
|
||||
None,
|
||||
),
|
||||
(
|
||||
{
|
||||
"texts": ["You are a helpful assistant.", " "],
|
||||
"structured_messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": " "},
|
||||
],
|
||||
},
|
||||
None,
|
||||
),
|
||||
(
|
||||
{
|
||||
"texts": ["every", "message", "\n"],
|
||||
"structured_messages": [
|
||||
*_CONVERSATION,
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "\n"},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
None,
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"string-content",
|
||||
"content-parts",
|
||||
"consecutive-user-messages",
|
||||
"a2a-message-parts",
|
||||
"texts-only-route",
|
||||
"ends-on-assistant",
|
||||
"ends-on-tool",
|
||||
"no-user-message",
|
||||
"whitespace-only-latest-turn",
|
||||
"whitespace-text-beside-image",
|
||||
],
|
||||
)
|
||||
async def test_request_scan_sends_only_the_last_user_turn(
|
||||
inputs: GenericGuardrailAPIInputs, expected_content: str | None
|
||||
) -> None:
|
||||
bodies: Final = await _execute_policy_bodies(inputs, "request")
|
||||
|
||||
assert bodies == (
|
||||
() if expected_content is None else ({"direction": "IN", "content": expected_content, "policyId": 1},)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_scan_sends_every_text() -> None:
|
||||
inputs: Final = GenericGuardrailAPIInputs(
|
||||
texts=["First choice.", "Second choice."],
|
||||
structured_messages=[*_CONVERSATION, {"role": "user", "content": "And three plus three?"}],
|
||||
)
|
||||
|
||||
bodies: Final = await _execute_policy_bodies(inputs, "response")
|
||||
|
||||
assert bodies == ({"direction": "OUT", "content": "First choice. Second choice.", "policyId": 1},)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scan_only_tool_results_sends_the_tool_results() -> None:
|
||||
inputs: Final = GenericGuardrailAPIInputs(
|
||||
texts=["Sunny, 22C"],
|
||||
structured_messages=[
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "Sunny, 22C"},
|
||||
],
|
||||
)
|
||||
|
||||
bodies: Final = await _execute_policy_bodies(inputs, "request", scan_only_tool_results=True)
|
||||
|
||||
assert bodies == ({"direction": "IN", "content": "Sunny, 22C", "policyId": 1},)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue