mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
a76b59db9f
Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
1062 lines
42 KiB
Python
1062 lines
42 KiB
Python
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
import litellm
|
|
from litellm.caching import DualCache
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
from litellm.types.utils import CallTypesLiteral
|
|
|
|
|
|
def test_has_post_call_response_headers_callbacks_ignores_empty_callbacks(
|
|
monkeypatch,
|
|
):
|
|
monkeypatch.setattr(litellm, "callbacks", [])
|
|
|
|
assert ProxyLogging.has_post_call_response_headers_callbacks() is False
|
|
|
|
|
|
def test_has_post_call_response_headers_callbacks_requires_override(
|
|
monkeypatch,
|
|
):
|
|
"""A vanilla ``CustomLogger`` inherits the no-op response-headers hook;
|
|
the capability flag must stay False so the proxy can skip the headers
|
|
loop entirely. Only callbacks that *override* the hook should flip it."""
|
|
monkeypatch.setattr(litellm, "callbacks", [CustomLogger()])
|
|
assert ProxyLogging.has_post_call_response_headers_callbacks() is False
|
|
|
|
class _AddsHeaders(CustomLogger):
|
|
async def async_post_call_response_headers_hook(self, **kwargs):
|
|
return {"x-custom": "1"}
|
|
|
|
monkeypatch.setattr(litellm, "callbacks", [_AddsHeaders()])
|
|
assert ProxyLogging.has_post_call_response_headers_callbacks() is True
|
|
|
|
|
|
def test_has_streaming_callbacks_uses_custom_logger_detection(monkeypatch):
|
|
monkeypatch.setattr(litellm, "callbacks", [])
|
|
assert ProxyLogging.has_streaming_callbacks() is False
|
|
|
|
monkeypatch.setattr(litellm, "callbacks", [CustomLogger()])
|
|
assert ProxyLogging.has_streaming_callbacks() is False
|
|
|
|
class StreamingLogger(CustomLogger):
|
|
async def async_post_call_streaming_hook(self, **kwargs):
|
|
return kwargs.get("response")
|
|
|
|
monkeypatch.setattr(litellm, "callbacks", [StreamingLogger()])
|
|
assert ProxyLogging.has_streaming_callbacks() is True
|
|
|
|
|
|
def test_has_streaming_callbacks_detects_guardrails(monkeypatch):
|
|
monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail()])
|
|
assert ProxyLogging.has_streaming_callbacks() is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_response_headers_hook_returns_early_without_callbacks(
|
|
monkeypatch,
|
|
):
|
|
monkeypatch.setattr(litellm, "callbacks", [])
|
|
proxy_logging_obj = ProxyLogging(user_api_key_cache={}) # type: ignore[arg-type]
|
|
|
|
result = await proxy_logging_obj.post_call_response_headers_hook(
|
|
data={},
|
|
user_api_key_dict=None, # type: ignore[arg-type]
|
|
response=None,
|
|
request_headers={},
|
|
)
|
|
|
|
assert result == {}
|
|
|
|
|
|
def test_callback_capabilities_skips_default_custom_logger(monkeypatch):
|
|
"""
|
|
Internal proxy hooks (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit
|
|
the default ``async_post_call_streaming_iterator_hook`` body. The
|
|
capability scanner must NOT report them as iterator overrides — wrapping
|
|
the chunk stream through every no-op layer was responsible for ~10x
|
|
streaming overhead on default deployments.
|
|
"""
|
|
|
|
class _InternalNoopHook(CustomLogger):
|
|
pass
|
|
|
|
monkeypatch.setattr(litellm, "callbacks", [_InternalNoopHook()])
|
|
|
|
caps = ProxyLogging._callback_capabilities()
|
|
# Subclass inherits the base no-op for every hook — every capability flag
|
|
# must stay False so the proxy short-circuits the corresponding loops.
|
|
assert caps.has_post_call_response_headers is False
|
|
assert caps.iterator_overrides == ()
|
|
assert caps.has_iterator_override is False
|
|
assert caps.has_streaming_chunk_override is False
|
|
assert caps.has_guardrail is False
|
|
|
|
|
|
def test_callback_capabilities_captures_iterator_override(monkeypatch):
|
|
class _OverridesIterator(CustomLogger):
|
|
async def async_post_call_streaming_iterator_hook( # type: ignore[override]
|
|
self, user_api_key_dict, response, request_data
|
|
):
|
|
async for item in response:
|
|
yield item
|
|
|
|
override = _OverridesIterator()
|
|
monkeypatch.setattr(litellm, "callbacks", [override])
|
|
|
|
caps = ProxyLogging._callback_capabilities()
|
|
assert caps.has_iterator_override is True
|
|
assert len(caps.iterator_overrides) == 1
|
|
resolved, kind = caps.iterator_overrides[0]
|
|
assert resolved is override
|
|
assert kind == "override"
|
|
|
|
|
|
def test_callback_capabilities_detects_inherited_streaming_chunk_override(monkeypatch):
|
|
"""
|
|
``async_post_call_streaming_hook`` must be detected even when the override
|
|
lives on an intermediate parent class — a vendor base class can carry the
|
|
override and the registered class can add nothing else. Before this PR the
|
|
hook was unconditionally invoked, so a leaf-class ``__dict__`` miss here
|
|
would silently drop the inherited hook.
|
|
"""
|
|
ProxyLogging._callback_capabilities_cache.clear()
|
|
|
|
class _StreamingBase(CustomLogger):
|
|
async def async_post_call_streaming_hook(self, *args, **kwargs): # type: ignore[override]
|
|
return kwargs.get("response")
|
|
|
|
class _LeafWithoutOverride(_StreamingBase):
|
|
pass
|
|
|
|
monkeypatch.setattr(litellm, "callbacks", [_LeafWithoutOverride()])
|
|
caps = ProxyLogging._callback_capabilities()
|
|
assert caps.has_streaming_chunk_override is True
|
|
|
|
|
|
def test_callback_capabilities_cache_invalidates_on_list_change(monkeypatch):
|
|
"""The cache key includes (length, id-of-each-callback). Mutating the
|
|
callback list must produce a fresh capability snapshot."""
|
|
monkeypatch.setattr(litellm, "callbacks", [])
|
|
assert ProxyLogging._callback_capabilities().resolved_callbacks == ()
|
|
|
|
class _OverridesPreCall(CustomLogger):
|
|
async def async_pre_call_hook(self, *args, **kwargs):
|
|
return kwargs.get("data")
|
|
|
|
pre = _OverridesPreCall()
|
|
monkeypatch.setattr(litellm, "callbacks", [pre])
|
|
caps = ProxyLogging._callback_capabilities()
|
|
assert caps.has_pre_call_override is True
|
|
assert pre in caps.resolved_callbacks
|
|
|
|
|
|
def _sse_bytes(event: str, payload: dict) -> bytes:
|
|
import json
|
|
|
|
return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode()
|
|
|
|
|
|
def _anthropic_stream_chunks(text_parts):
|
|
chunks = [
|
|
_sse_bytes(
|
|
"message_start",
|
|
{
|
|
"type": "message_start",
|
|
"message": {
|
|
"model": "claude-sonnet-5",
|
|
"id": "msg_1",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [],
|
|
"stop_reason": None,
|
|
"usage": {"input_tokens": 20, "output_tokens": 1},
|
|
},
|
|
},
|
|
),
|
|
_sse_bytes(
|
|
"content_block_start",
|
|
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
|
),
|
|
]
|
|
for part in text_parts:
|
|
chunks.append(
|
|
_sse_bytes(
|
|
"content_block_delta",
|
|
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": part}},
|
|
)
|
|
)
|
|
chunks.append(_sse_bytes("content_block_stop", {"type": "content_block_stop", "index": 0}))
|
|
chunks.append(
|
|
_sse_bytes(
|
|
"message_delta",
|
|
{
|
|
"type": "message_delta",
|
|
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
|
"usage": {"input_tokens": 20, "output_tokens": 8},
|
|
},
|
|
)
|
|
)
|
|
chunks.append(_sse_bytes("message_stop", {"type": "message_stop"}))
|
|
return chunks
|
|
|
|
|
|
def _content_filter_guardrail(action: str, guardrail_cls=None, **guardrail_kwargs):
|
|
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
|
|
ContentFilterGuardrail,
|
|
)
|
|
from litellm.types.guardrails import BlockedWord, ContentFilterAction
|
|
|
|
cls = guardrail_cls or ContentFilterGuardrail
|
|
return cls(
|
|
guardrail_name="output-filter",
|
|
blocked_words=[BlockedWord(keyword="zebra", action=ContentFilterAction(action))],
|
|
event_hook="post_call",
|
|
default_on=True,
|
|
**guardrail_kwargs,
|
|
)
|
|
|
|
|
|
def _streaming_logging_obj():
|
|
import datetime
|
|
import uuid
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
|
|
return Logging(
|
|
model="claude-sonnet-5",
|
|
messages=[{"role": "user", "content": "Reply with exactly: the zebra runs"}],
|
|
stream=True,
|
|
call_type="anthropic_messages",
|
|
start_time=datetime.datetime.now(),
|
|
litellm_call_id=str(uuid.uuid4()),
|
|
function_id="test",
|
|
)
|
|
|
|
|
|
def test_stream_requires_guardrail_translation_route_detection():
|
|
assert (
|
|
ProxyLogging._stream_requires_guardrail_translation(
|
|
UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages")
|
|
)
|
|
is True
|
|
)
|
|
assert (
|
|
ProxyLogging._stream_requires_guardrail_translation(
|
|
UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions")
|
|
)
|
|
is False
|
|
)
|
|
assert ProxyLogging._stream_requires_guardrail_translation(UserAPIKeyAuth(api_key="sk-1234")) is False
|
|
assert (
|
|
ProxyLogging._stream_requires_guardrail_translation(
|
|
UserAPIKeyAuth(api_key="sk-1234", request_route="/route/without/call/types")
|
|
)
|
|
is False
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_stream_guardrail_blocks_anthropic_messages_stream(monkeypatch):
|
|
"""
|
|
Regression test for https://github.com/BerriAI/litellm/issues/35257.
|
|
|
|
/v1/messages streams raw Anthropic SSE bytes. A guardrail whose custom
|
|
iterator hook only understands OpenAI ModelResponseStream chunks used to
|
|
receive those bytes directly and silently pass every chunk through
|
|
unscanned. The dispatch must route apply_guardrail-capable guardrails
|
|
through unified_guardrail's anthropic translation so blocked output
|
|
raises instead of streaming to the client. Because the guardrail's own
|
|
iterator hook withheld content until scanned, the rerouted invocation
|
|
defaults to buffer_until_moderated, so nothing may reach the client
|
|
before the block fires.
|
|
"""
|
|
from fastapi import HTTPException
|
|
|
|
from litellm.caching.caching import DualCache
|
|
guardrail = _content_filter_guardrail("BLOCK")
|
|
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
|
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
request_data = {
|
|
"model": "claude-sonnet-5",
|
|
"litellm_logging_obj": _streaming_logging_obj(),
|
|
"metadata": {},
|
|
}
|
|
|
|
async def fake_stream():
|
|
for chunk in _anthropic_stream_chunks(["the", " zebra runs"]):
|
|
yield chunk
|
|
|
|
delivered = []
|
|
async def _drain():
|
|
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
|
response=fake_stream(),
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
|
|
request_data=request_data,
|
|
):
|
|
delivered.append(chunk)
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await _drain()
|
|
|
|
detail = exc_info.value.detail
|
|
assert detail["guardrail_name"] == "output-filter"
|
|
assert detail["keyword"] == "zebra"
|
|
assert delivered == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions(monkeypatch):
|
|
"""
|
|
On /chat/completions the guardrail's own iterator hook must keep running:
|
|
it masks incrementally inside ModelResponseStream chunks, which the
|
|
unified block_only path never does. Masked output proves the own-hook
|
|
path was used.
|
|
"""
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
|
|
|
guardrail = _content_filter_guardrail("MASK")
|
|
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
|
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
async def fake_stream():
|
|
yield ModelResponseStream(
|
|
choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))]
|
|
)
|
|
yield ModelResponseStream(
|
|
choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")]
|
|
)
|
|
|
|
delivered_text = ""
|
|
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
|
response=fake_stream(),
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"),
|
|
request_data={"model": "gpt-4o-mini", "metadata": {}},
|
|
):
|
|
for choice in chunk.choices:
|
|
delivered_text += choice.delta.content or ""
|
|
|
|
assert "zebra" not in delivered_text
|
|
assert delivered_text != ""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch):
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
|
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
|
|
|
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail("MASK")])
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
logging_obj = _streaming_logging_obj()
|
|
|
|
async def fake_stream():
|
|
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))])
|
|
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")])
|
|
|
|
delivered_text = ""
|
|
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
|
response=fake_stream(),
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"),
|
|
request_data={"model": "gpt-4o-mini", "metadata": {}, "litellm_logging_obj": logging_obj},
|
|
):
|
|
for choice in chunk.choices:
|
|
delivered_text += choice.delta.content or ""
|
|
|
|
assert "zebra" not in delivered_text
|
|
assert logging_obj.model_call_details[SERVED_OUTPUT_TEXTS_KEY] == (delivered_text,)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_stream_records_the_served_text_when_the_client_disconnects(monkeypatch):
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
|
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
|
|
|
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail("MASK")])
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
logging_obj = _streaming_logging_obj()
|
|
|
|
async def fake_stream():
|
|
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))])
|
|
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=" far"))])
|
|
|
|
stream = proxy_logging.async_post_call_streaming_iterator_hook(
|
|
response=fake_stream(),
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"),
|
|
request_data={"model": "gpt-4o-mini", "metadata": {}, "litellm_logging_obj": logging_obj},
|
|
)
|
|
first = await stream.__anext__()
|
|
await stream.aclose()
|
|
|
|
delivered_text = "".join(choice.delta.content or "" for choice in first.choices)
|
|
assert "zebra" not in delivered_text
|
|
assert logging_obj.model_call_details[SERVED_OUTPUT_TEXTS_KEY] == (delivered_text,)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unified_guardrail_iterator_accepts_explicit_guardrail():
|
|
"""
|
|
The dispatch passes each guardrail explicitly instead of through a shared
|
|
request_data key, so chaining two unified-routed guardrails cannot drop
|
|
all but the last one. The block fires after the deltas were already
|
|
flushed to the client, so it surfaces as a trailing in-stream error frame
|
|
rather than a raised HTTPException.
|
|
"""
|
|
from litellm.proxy.utils import unified_guardrail
|
|
|
|
guardrail = _content_filter_guardrail("BLOCK")
|
|
request_data = {
|
|
"model": "claude-sonnet-5",
|
|
"litellm_logging_obj": _streaming_logging_obj(),
|
|
"metadata": {},
|
|
}
|
|
|
|
async def fake_stream():
|
|
for chunk in _anthropic_stream_chunks(["the", " zebra runs"]):
|
|
yield chunk
|
|
|
|
delivered = []
|
|
async for item in unified_guardrail.async_post_call_streaming_iterator_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
|
|
response=fake_stream(),
|
|
request_data=request_data,
|
|
guardrail_to_apply=guardrail,
|
|
):
|
|
delivered.append(item)
|
|
|
|
raw = b"".join(c for c in delivered if isinstance(c, bytes)).decode()
|
|
assert "event: error" in raw
|
|
assert "guardrail_error" in raw
|
|
assert raw.index("guardrail_error") > raw.index(" zebra runs")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_stream_guardrail_reroutes_inherited_apply_guardrail(monkeypatch):
|
|
"""
|
|
The reroute predicate must recognize apply_guardrail implementations
|
|
inherited from a parent class, not only ones defined on the registered
|
|
leaf class. A vendor base class can carry apply_guardrail while the leaf
|
|
only overrides the streaming iterator; a leaf-class ``__dict__`` check
|
|
would leave that guardrail on the raw Anthropic SSE path unscanned.
|
|
"""
|
|
from fastapi import HTTPException
|
|
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
|
|
ContentFilterGuardrail,
|
|
)
|
|
|
|
class _InheritsApplyGuardrail(ContentFilterGuardrail):
|
|
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
|
async for item in response:
|
|
yield item
|
|
|
|
guardrail = _content_filter_guardrail("BLOCK", guardrail_cls=_InheritsApplyGuardrail)
|
|
assert "apply_guardrail" not in type(guardrail).__dict__
|
|
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
|
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
request_data = {
|
|
"model": "claude-sonnet-5",
|
|
"litellm_logging_obj": _streaming_logging_obj(),
|
|
"metadata": {},
|
|
}
|
|
|
|
async def fake_stream():
|
|
for chunk in _anthropic_stream_chunks(["the", " zebra runs"]):
|
|
yield chunk
|
|
|
|
delivered = []
|
|
async def _drain():
|
|
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
|
response=fake_stream(),
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
|
|
request_data=request_data,
|
|
):
|
|
delivered.append(chunk)
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await _drain()
|
|
|
|
assert exc_info.value.detail["keyword"] == "zebra"
|
|
assert delivered == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_stream_masking_guardrail_keeps_own_iterator_on_anthropic(monkeypatch):
|
|
"""
|
|
A guardrail with mask_response_content=True must stay on its own iterator
|
|
hook on /v1/messages. The unified streaming path cannot re-emit rewritten
|
|
text on raw Anthropic SSE (block_only drops rewrites and buffered replay
|
|
releases the unredacted originals), so rerouting such a guardrail would
|
|
deliver content it decided to mask. PANW Prisma AIRS is the concrete
|
|
case: its own hook parses the raw bytes and blocks instead of masking.
|
|
"""
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
|
|
ContentFilterGuardrail,
|
|
)
|
|
|
|
own_hook_streams = []
|
|
|
|
class _MasksViaOwnRawStreamHook(ContentFilterGuardrail):
|
|
apply_guardrail = ContentFilterGuardrail.apply_guardrail
|
|
|
|
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
|
own_hook_streams.append(request_data.get("model"))
|
|
async for item in response:
|
|
yield item
|
|
|
|
guardrail = _content_filter_guardrail(
|
|
"BLOCK", guardrail_cls=_MasksViaOwnRawStreamHook, mask_response_content=True
|
|
)
|
|
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
|
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
chunks = _anthropic_stream_chunks(["the", " zebra runs"])
|
|
|
|
async def fake_stream():
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
delivered = []
|
|
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
|
response=fake_stream(),
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
|
|
request_data={
|
|
"model": "claude-sonnet-5",
|
|
"litellm_logging_obj": _streaming_logging_obj(),
|
|
"metadata": {},
|
|
},
|
|
):
|
|
delivered.append(chunk)
|
|
|
|
assert own_hook_streams == ["claude-sonnet-5"]
|
|
assert delivered == chunks
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_stream_presidio_output_masking_masks_anthropic_messages_stream(monkeypatch):
|
|
"""Regression: the presidio output-masking callback built by initialize_presidio
|
|
was rerouted onto the unified scan-only path on /v1/messages, so a card number
|
|
the analyzer flagged still streamed to the caller unmasked."""
|
|
import json
|
|
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
|
|
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
|
|
|
handler = InMemoryGuardrailHandler()
|
|
result = handler.initialize_guardrail(
|
|
guardrail={
|
|
"guardrail_name": "presidio-card-mask",
|
|
"litellm_params": {
|
|
"guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
|
|
"mode": ["pre_call", "post_call"],
|
|
"default_on": True,
|
|
"presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
|
|
"presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
|
|
"pii_entities_config": {"CREDIT_CARD": "MASK"},
|
|
"mock_redacted_text": {"text": "<CREDIT_CARD>", "items": []},
|
|
},
|
|
}
|
|
)
|
|
guardrail_id = result["guardrail_id"]
|
|
callbacks = [
|
|
handler.guardrail_id_to_custom_guardrail[guardrail_id],
|
|
*handler.guardrail_id_to_sibling_callbacks[guardrail_id],
|
|
]
|
|
monkeypatch.setattr(litellm, "callbacks", callbacks)
|
|
|
|
chunks = _anthropic_stream_chunks(["4111", " 1111 1111 1111"])
|
|
|
|
async def fake_stream():
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
delivered = []
|
|
async for chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook(
|
|
response=fake_stream(),
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
|
|
request_data={
|
|
"model": "claude-sonnet-5",
|
|
"litellm_logging_obj": _streaming_logging_obj(),
|
|
"metadata": {},
|
|
},
|
|
):
|
|
delivered.append(chunk)
|
|
|
|
wire = b"".join(delivered).decode()
|
|
text_deltas = [
|
|
json.loads(line[6:])["delta"]["text"]
|
|
for line in wire.split("\n")
|
|
if line.startswith("data: ") and json.loads(line[6:]).get("delta", {}).get("type") == "text_delta"
|
|
]
|
|
assert "4111" not in wire, wire
|
|
assert "".join(text_deltas) == "<CREDIT_CARD>", wire
|
|
assert wire.count("event: message_stop") == 1, wire
|
|
|
|
|
|
class _AppliesGuardrail(CustomGuardrail):
|
|
"""Implements the unified interface only, so the proxy routes it to unified_guardrail."""
|
|
|
|
def __init__(self, **kwargs):
|
|
super().__init__(guardrail_name="applies", **kwargs)
|
|
self.native_hooks_ran = []
|
|
|
|
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
|
return inputs
|
|
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
self.native_hooks_ran.append("pre_call")
|
|
|
|
async def async_moderation_hook(self, data, user_api_key_dict, call_type):
|
|
self.native_hooks_ran.append("during_call")
|
|
|
|
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
|
self.native_hooks_ran.append("post_call")
|
|
return response
|
|
|
|
|
|
class _KeepsNativeHooks(CustomGuardrail):
|
|
"""Same, plus the opt-out that keeps request traffic on its own hooks.
|
|
|
|
apply_guardrail is redefined here rather than inherited because the proxy's
|
|
dispatch check reads the leaf class __dict__, so an inherited override would
|
|
take the native path for the wrong reason and the flag would go untested."""
|
|
|
|
use_native_lifecycle_hooks = True
|
|
|
|
def __init__(self, **kwargs):
|
|
super().__init__(guardrail_name="keeps_native", **kwargs)
|
|
self.native_hooks_ran = []
|
|
|
|
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
|
return inputs
|
|
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
self.native_hooks_ran.append("pre_call")
|
|
|
|
async def async_moderation_hook(self, data, user_api_key_dict, call_type):
|
|
self.native_hooks_ran.append("during_call")
|
|
|
|
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
|
self.native_hooks_ran.append("post_call")
|
|
return response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("hook_type", ["pre_call", "post_call"])
|
|
async def test_execute_guardrail_hook_routes_apply_guardrail_implementers_to_unified(hook_type):
|
|
guardrail = _AppliesGuardrail()
|
|
data = {"messages": [{"role": "user", "content": "hi"}]}
|
|
|
|
await ProxyLogging(user_api_key_cache=DualCache())._execute_guardrail_hook(
|
|
callback=guardrail,
|
|
hook_type=hook_type,
|
|
data=data,
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
call_type="completion",
|
|
response=None,
|
|
)
|
|
|
|
assert guardrail.native_hooks_ran == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("hook_type", ["pre_call", "post_call"])
|
|
async def test_execute_guardrail_hook_keeps_native_hooks_when_opted_out(hook_type):
|
|
"""A guardrail that implements apply_guardrail purely to serve
|
|
/guardrails/apply_guardrail must not have its request traffic rerouted."""
|
|
guardrail = _KeepsNativeHooks()
|
|
data = {"messages": [{"role": "user", "content": "hi"}]}
|
|
|
|
await ProxyLogging(user_api_key_cache=DualCache())._execute_guardrail_hook(
|
|
callback=guardrail,
|
|
hook_type=hook_type,
|
|
data=data,
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
call_type="completion",
|
|
response=None,
|
|
)
|
|
|
|
assert guardrail.native_hooks_ran == [hook_type]
|
|
assert "guardrail_to_apply" not in data
|
|
|
|
|
|
def test_azure_content_safety_guardrails_keep_their_native_hooks():
|
|
from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import (
|
|
AzureContentSafetyPromptShieldGuardrail,
|
|
)
|
|
from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import (
|
|
AzureContentSafetyTextModerationGuardrail,
|
|
)
|
|
|
|
assert CustomGuardrail.use_native_lifecycle_hooks is False
|
|
assert AzureContentSafetyPromptShieldGuardrail.use_native_lifecycle_hooks is True
|
|
assert AzureContentSafetyTextModerationGuardrail.use_native_lifecycle_hooks is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_during_call_hook_keeps_native_moderation_hook_when_opted_out(monkeypatch):
|
|
opted_out = _KeepsNativeHooks(event_hook=GuardrailEventHooks.during_call, default_on=True)
|
|
routed = _AppliesGuardrail(event_hook=GuardrailEventHooks.during_call, default_on=True)
|
|
monkeypatch.setattr(litellm, "callbacks", [opted_out, routed])
|
|
|
|
await ProxyLogging(user_api_key_cache=DualCache()).during_call_hook(
|
|
data={"messages": [{"role": "user", "content": "hi"}]},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
call_type="completion",
|
|
)
|
|
|
|
assert opted_out.native_hooks_ran == ["during_call"]
|
|
assert routed.native_hooks_ran == []
|
|
|
|
|
|
class _RejectsInModeration(CustomLogger):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.moderated: list[str] = []
|
|
|
|
async def async_moderation_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
call_type: CallTypesLiteral,
|
|
) -> None:
|
|
self.moderated.append(call_type)
|
|
raise HTTPException(status_code=400, detail={"error": "rejected"})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_during_call_hook_runs_custom_logger_moderation_override(monkeypatch):
|
|
moderator = _RejectsInModeration()
|
|
monkeypatch.setattr(litellm, "callbacks", [CustomLogger(), moderator])
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await ProxyLogging(user_api_key_cache=DualCache()).during_call_hook(
|
|
data={"messages": [{"role": "user", "content": "hi"}]},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
call_type="acompletion",
|
|
)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
assert moderator.moderated == ["acompletion"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_during_call_hook_skips_custom_logger_moderation_without_auth(monkeypatch):
|
|
moderator = _RejectsInModeration()
|
|
monkeypatch.setattr(litellm, "callbacks", [moderator])
|
|
data = {"messages": [{"role": "user", "content": "hi"}]}
|
|
|
|
result = await ProxyLogging(user_api_key_cache=DualCache()).during_call_hook(
|
|
data=data,
|
|
user_api_key_dict=None,
|
|
call_type="acompletion",
|
|
)
|
|
|
|
assert result == data
|
|
assert moderator.moderated == []
|
|
|
|
|
|
class _InheritsModerationOverride(_RejectsInModeration):
|
|
pass
|
|
|
|
|
|
class _V1PreCallGuardrail(CustomGuardrail):
|
|
def __init__(self) -> None:
|
|
super().__init__(guardrail_name="v1-pre-call")
|
|
self.moderation_check = "pre_call"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.filterwarnings("error::RuntimeWarning")
|
|
async def test_during_call_hook_runs_moderation_override_after_v1_pre_call_guardrail(monkeypatch):
|
|
moderator = _RejectsInModeration()
|
|
monkeypatch.setattr(litellm, "callbacks", [_V1PreCallGuardrail(), moderator])
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await ProxyLogging(user_api_key_cache=DualCache()).during_call_hook(
|
|
data={"messages": [{"role": "user", "content": "hi"}]},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
call_type="acompletion",
|
|
)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
assert moderator.moderated == ["acompletion"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_during_call_hook_runs_moderation_override_inherited_from_parent(monkeypatch):
|
|
moderator = _InheritsModerationOverride()
|
|
monkeypatch.setattr(litellm, "callbacks", [moderator])
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await ProxyLogging(user_api_key_cache=DualCache()).during_call_hook(
|
|
data={"messages": [{"role": "user", "content": "hi"}]},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
call_type="acompletion",
|
|
)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
assert moderator.moderated == ["acompletion"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_success_hook_keeps_native_hook_when_opted_out(monkeypatch):
|
|
from litellm.types.utils import Choices, Message, ModelResponse
|
|
|
|
opted_out = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
|
routed = _AppliesGuardrail(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
|
monkeypatch.setattr(litellm, "callbacks", [opted_out, routed])
|
|
response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))])
|
|
|
|
await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook(
|
|
data={"messages": [{"role": "user", "content": "hi"}]},
|
|
response=response,
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
)
|
|
|
|
assert opted_out.native_hooks_ran == ["post_call"]
|
|
assert routed.native_hooks_ran == []
|
|
|
|
|
|
def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overrides(monkeypatch):
|
|
"""An opted-out guardrail must not be registered as an apply_guardrail iterator
|
|
override, or its streamed responses run through the unified pipeline instead of
|
|
its own hooks."""
|
|
ProxyLogging._callback_capabilities_cache.clear()
|
|
opted_out = _KeepsNativeHooks()
|
|
routed = _AppliesGuardrail()
|
|
monkeypatch.setattr(litellm, "callbacks", [opted_out, routed])
|
|
|
|
caps = ProxyLogging._callback_capabilities()
|
|
|
|
assert [(cb, kind) for cb, kind in caps.iterator_overrides if cb is routed] == [(routed, "apply_guardrail")]
|
|
assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == []
|
|
|
|
|
|
def test_deployment_hook_target_stays_native_when_opted_out():
|
|
"""Model-level guardrails resolve their target here rather than through ProxyLogging."""
|
|
assert _KeepsNativeHooks()._deployment_hook_target() is not None
|
|
opted_out = _KeepsNativeHooks()
|
|
assert opted_out._deployment_hook_target() is opted_out
|
|
assert _AppliesGuardrail()._deployment_hook_target() is not None
|
|
routed = _AppliesGuardrail()
|
|
assert routed._deployment_hook_target() is not routed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_stream_guardrails_run_native_hook_when_opted_out(monkeypatch):
|
|
"""The deferred path skips unified-routed guardrails because the streaming iterator
|
|
already scanned. An opted-out guardrail never reached that iterator, so its own
|
|
post-call hook has to run here."""
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
from litellm.types.utils import Choices, Message, ModelResponse
|
|
|
|
opted_out = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
|
routed = _AppliesGuardrail(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
|
monkeypatch.setattr(litellm, "callbacks", [opted_out, routed])
|
|
|
|
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
|
|
captured_data={"messages": [{"role": "user", "content": "hi"}]},
|
|
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
captured_logging_obj=_streaming_logging_obj(),
|
|
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
|
|
cache_hit=False,
|
|
)
|
|
|
|
assert opted_out.native_hooks_ran == ["post_call"]
|
|
assert routed.native_hooks_ran == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monkeypatch):
|
|
"""A post_call pipeline step already ran the opted-out guardrail's own hook against
|
|
the buffered stream, so the deferred audit must not run it a second time."""
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
|
|
from litellm.types.utils import Choices, Message, ModelResponse
|
|
|
|
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
|
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed])
|
|
pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")])
|
|
|
|
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
|
|
captured_data={
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
|
|
},
|
|
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"),
|
|
captured_logging_obj=_streaming_logging_obj(),
|
|
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
|
|
cache_hit=False,
|
|
)
|
|
|
|
assert pipeline_managed.native_hooks_ran == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_stream_guardrails_run_native_hook_whose_pipeline_could_not_stream(monkeypatch):
|
|
"""A pipeline step with neither streaming interface keeps the whole pipeline off the
|
|
stream, so the deferred audit is the only place the opted-out guardrail's own hook
|
|
still runs, the way it did before pipelines ran on streams."""
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
|
|
from litellm.types.utils import Choices, Message, ModelResponse
|
|
|
|
class NeitherHookGuardrail(CustomGuardrail):
|
|
pass
|
|
|
|
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
|
neither = NeitherHookGuardrail(guardrail_name="gr-neither", event_hook=GuardrailEventHooks.post_call)
|
|
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed, neither])
|
|
pipeline = GuardrailPipeline(
|
|
mode="post_call",
|
|
steps=[
|
|
PipelineStep(guardrail="keeps_native", on_fail="next"),
|
|
PipelineStep(guardrail="gr-neither", on_fail="block"),
|
|
],
|
|
)
|
|
|
|
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
|
|
captured_data={
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
|
|
},
|
|
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"),
|
|
captured_logging_obj=_streaming_logging_obj(),
|
|
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
|
|
cache_hit=False,
|
|
)
|
|
|
|
assert pipeline_managed.native_hooks_ran == ["post_call"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_stream_guardrails_run_native_hook_on_route_without_translation(monkeypatch):
|
|
"""A route with no endpoint guardrail translation cannot gate the stream through its
|
|
pipelines, so the deferred audit still owes the opted-out guardrail its own hook."""
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
|
|
from litellm.types.utils import Choices, Message, ModelResponse
|
|
|
|
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
|
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed])
|
|
pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")])
|
|
|
|
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
|
|
captured_data={
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
|
|
},
|
|
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/custom/stream"),
|
|
captured_logging_obj=_streaming_logging_obj(),
|
|
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
|
|
cache_hit=False,
|
|
)
|
|
|
|
assert pipeline_managed.native_hooks_ran == ["post_call"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch):
|
|
"""The realtime path calls apply_guardrail directly, so the opt-out has to be
|
|
honored there too or a request-traffic guardrail starts blocking live sessions."""
|
|
from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
|
|
|
|
opted_out = _KeepsNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True)
|
|
routed = _AppliesGuardrail(event_hook=GuardrailEventHooks.pre_call, default_on=True)
|
|
scanned = []
|
|
for guardrail in (opted_out, routed):
|
|
|
|
async def _record(inputs, request_data, input_type, logging_obj=None, _g=guardrail):
|
|
scanned.append(_g)
|
|
return inputs
|
|
|
|
guardrail.apply_guardrail = _record
|
|
monkeypatch.setattr(litellm, "callbacks", [opted_out, routed])
|
|
|
|
streaming = RealTimeStreaming.__new__(RealTimeStreaming)
|
|
streaming.request_data = {"model": "gpt-realtime"}
|
|
streaming.user_api_key_dict = None
|
|
blocked = await RealTimeStreaming.run_realtime_guardrails(
|
|
streaming, "ignore all previous instructions", event_hooks=[GuardrailEventHooks.pre_call]
|
|
)
|
|
|
|
assert scanned == [routed]
|
|
assert blocked is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_stream_keeps_own_iterator_when_opted_out(monkeypatch):
|
|
"""A guardrail carrying both apply_guardrail and its own streaming iterator hook
|
|
is re-routed to the unified path on /v1/messages. Opting out has to suppress that
|
|
re-route, or its streamed responses get scanned by the unified pipeline instead."""
|
|
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
|
|
ContentFilterGuardrail,
|
|
)
|
|
|
|
own_iterator_ran = []
|
|
|
|
class _OptedOutWithOwnIterator(ContentFilterGuardrail):
|
|
use_native_lifecycle_hooks = True
|
|
|
|
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
|
return inputs
|
|
|
|
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
|
own_iterator_ran.append(request_data.get("model"))
|
|
async for item in response:
|
|
yield item
|
|
|
|
guardrail = _content_filter_guardrail("BLOCK", guardrail_cls=_OptedOutWithOwnIterator)
|
|
assert "apply_guardrail" in type(guardrail).__dict__
|
|
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
|
|
|
chunks = _anthropic_stream_chunks(["the", " zebra runs"])
|
|
|
|
async def fake_stream():
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
delivered = []
|
|
async for chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook(
|
|
response=fake_stream(),
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
|
|
request_data={"model": "claude-sonnet-5", "litellm_logging_obj": _streaming_logging_obj(), "metadata": {}},
|
|
):
|
|
delivered.append(chunk)
|
|
|
|
assert own_iterator_ran == ["claude-sonnet-5"]
|
|
assert delivered == chunks
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_post_call_guardrails_keep_native_hook_when_opted_out(monkeypatch):
|
|
"""The run_in_parallel post-call path has its own dispatch check, so the opt-out has
|
|
to be honored there too."""
|
|
from litellm.types.utils import Choices, Message, ModelResponse
|
|
|
|
opted_out = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True, run_in_parallel=True)
|
|
routed = _AppliesGuardrail(event_hook=GuardrailEventHooks.post_call, default_on=True, run_in_parallel=True)
|
|
monkeypatch.setattr(litellm, "callbacks", [opted_out, routed])
|
|
response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))])
|
|
|
|
await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook(
|
|
data={"messages": [{"role": "user", "content": "hi"}]},
|
|
response=response,
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"),
|
|
)
|
|
|
|
assert opted_out.native_hooks_ran == ["post_call"]
|
|
assert routed.native_hooks_ran == []
|