mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
4faeabc254
* fix(guardrails): run pre_call hook once for model-level guardrails A CustomGuardrail attached to a deployment via litellm_params.guardrails gets its async_pre_call_hook invoked twice per request: once by the proxy pre-call loop and again by async_pre_call_deployment_hook after the router spreads the model-level guardrails into the top-level request kwargs. Record in request metadata that the proxy pre-call loop already ran a given guardrail, and have the deployment hook skip it when the marker is present. Direct-SDK usage never runs the proxy loop, so the deployment hook stays the sole invocation there and still fires exactly once. The marker key is stripped from untrusted caller metadata so a request body cannot suppress a model-only guardrail by pre-seeding it. * fix(guardrails): mark pre_call dedup on the post-hook request data Record the exactly-once marker after async_pre_call_hook runs, on the data object that flows downstream, rather than before it. A guardrail whose hook returns a brand-new request dict (instead of mutating or spreading the one it received) would otherwise discard the marker, letting the deployment hook re-run the guardrail a second time.
577 lines
21 KiB
Python
577 lines
21 KiB
Python
"""
|
|
Unit tests for model-level guardrails in post_call paths.
|
|
|
|
Tests verify that guardrails configured via litellm_params.guardrails on a
|
|
deployment are merged into request metadata and trigger execution for both
|
|
streaming and non-streaming post_call hooks.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")))
|
|
|
|
from litellm.proxy.utils import (
|
|
_check_and_merge_model_level_guardrails,
|
|
_merge_guardrails_with_existing,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Unit tests for _check_and_merge_model_level_guardrails
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCheckAndMergeModelLevelGuardrails:
|
|
"""Tests for the _check_and_merge_model_level_guardrails function."""
|
|
|
|
def test_merge_adds_model_guardrails_to_metadata(self):
|
|
"""Model-level guardrails are added to metadata.guardrails."""
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"model_info": {"id": "model-uuid-123"}},
|
|
}
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = ["openai-moderation"]
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
result = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=mock_router
|
|
)
|
|
|
|
assert "openai-moderation" in result["metadata"]["guardrails"]
|
|
mock_router.get_deployment.assert_called_once_with(model_id="model-uuid-123")
|
|
|
|
def test_merge_combines_with_existing_guardrails(self):
|
|
"""Model-level guardrails merge with existing request guardrails."""
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {
|
|
"model_info": {"id": "model-uuid-123"},
|
|
"guardrails": ["existing-guardrail"],
|
|
},
|
|
}
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = ["model-guardrail"]
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
result = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=mock_router
|
|
)
|
|
|
|
assert "existing-guardrail" in result["metadata"]["guardrails"]
|
|
assert "model-guardrail" in result["metadata"]["guardrails"]
|
|
|
|
def test_no_duplicates_when_guardrail_already_in_metadata(self):
|
|
"""No duplicates when the same guardrail is in both model and request."""
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {
|
|
"model_info": {"id": "model-uuid-123"},
|
|
"guardrails": ["openai-moderation"],
|
|
},
|
|
}
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = ["openai-moderation"]
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
result = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=mock_router
|
|
)
|
|
|
|
assert result["metadata"]["guardrails"].count("openai-moderation") == 1
|
|
|
|
def test_returns_data_unchanged_when_no_router(self):
|
|
"""Returns data unchanged when llm_router is None."""
|
|
data = {"model": "gpt-4", "metadata": {}}
|
|
result = _check_and_merge_model_level_guardrails(data=data, llm_router=None)
|
|
assert result is data
|
|
|
|
def test_returns_data_unchanged_when_no_model_info(self):
|
|
"""Returns data unchanged when metadata has no model_info."""
|
|
data = {"model": "gpt-4", "metadata": {}}
|
|
mock_router = MagicMock()
|
|
result = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=mock_router
|
|
)
|
|
assert result is data
|
|
|
|
def test_returns_data_unchanged_when_deployment_has_no_guardrails(self):
|
|
"""Returns data unchanged when deployment has no guardrails configured."""
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"model_info": {"id": "model-uuid-123"}},
|
|
}
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = None
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
result = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=mock_router
|
|
)
|
|
|
|
assert result is data
|
|
|
|
def test_returns_data_unchanged_when_deployment_not_found(self):
|
|
"""Returns data unchanged when router can't find the deployment."""
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"model_info": {"id": "nonexistent-id"}},
|
|
}
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment.return_value = None
|
|
|
|
result = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=mock_router
|
|
)
|
|
|
|
assert result is data
|
|
|
|
def test_returns_new_data_dict(self):
|
|
"""Returns a new top-level dict (shallow copy), not the same object."""
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {
|
|
"model_info": {"id": "model-uuid-123"},
|
|
"guardrails": ["existing"],
|
|
},
|
|
}
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = ["new-guardrail"]
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
result = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=mock_router
|
|
)
|
|
|
|
# Result is a different top-level dict
|
|
assert result is not data
|
|
# Result should have the merged guardrail
|
|
assert "new-guardrail" in result["metadata"]["guardrails"]
|
|
assert "existing" in result["metadata"]["guardrails"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Regression test: pre_call hook must run exactly once with model-level guardrails
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_hook_runs_once_with_model_level_guardrails():
|
|
"""
|
|
A guardrail attached at the model level (litellm_params.guardrails) is
|
|
spread into the top-level request kwargs by the router. The proxy pre-call
|
|
loop (async_pre_call_hook) and the deployment-level hook
|
|
(async_pre_call_deployment_hook) must together invoke async_pre_call_hook
|
|
exactly once, not twice.
|
|
"""
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class CountingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="counting-guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
default_on=True,
|
|
)
|
|
self.pre_call_count = 0
|
|
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
self.pre_call_count += 1
|
|
return data
|
|
|
|
guardrail = CountingGuardrail()
|
|
|
|
with patch("litellm.callbacks", [guardrail]):
|
|
ProxyLogging._callback_capabilities_cache.clear()
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
data = {
|
|
"model": "gpt-4",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"metadata": {},
|
|
}
|
|
|
|
# Path A: proxy pre-call loop runs the guardrail and records that it ran
|
|
data = await proxy_logging.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
call_type="acompletion",
|
|
)
|
|
|
|
# Path B: the router spreads the deployment's model-level guardrails into
|
|
# the top-level kwargs, then litellm.acompletion fires the deployment hook
|
|
data["guardrails"] = ["counting-guardrail"]
|
|
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
|
|
|
|
assert guardrail.pre_call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_hook_runs_once_when_hook_returns_fresh_dict():
|
|
"""
|
|
async_pre_call_hook may return a brand-new request dict instead of mutating
|
|
or spreading the one it received. The exactly-once marker must live on the
|
|
data that flows downstream, so the deployment hook still skips the guardrail
|
|
even when the proxy loop swapped in a fresh dict that never carried it.
|
|
"""
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class FreshDictGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="counting-guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
default_on=True,
|
|
)
|
|
self.pre_call_count = 0
|
|
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
self.pre_call_count += 1
|
|
return {"model": data["model"], "messages": data["messages"]}
|
|
|
|
guardrail = FreshDictGuardrail()
|
|
|
|
with patch("litellm.callbacks", [guardrail]):
|
|
ProxyLogging._callback_capabilities_cache.clear()
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
data = {
|
|
"model": "gpt-4",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"metadata": {},
|
|
}
|
|
|
|
data = await proxy_logging.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
call_type="acompletion",
|
|
)
|
|
|
|
data["guardrails"] = ["counting-guardrail"]
|
|
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
|
|
|
|
assert guardrail.pre_call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deployment_hook_runs_pre_call_without_proxy_loop():
|
|
"""
|
|
Direct-SDK usage (litellm.acompletion(..., guardrails=[...]) without the
|
|
proxy) never runs the proxy pre-call loop, so the deployment hook is the
|
|
only place the guardrail executes and it must still run exactly once.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import CallTypes
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class CountingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="counting-guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
default_on=True,
|
|
)
|
|
self.pre_call_count = 0
|
|
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
self.pre_call_count += 1
|
|
return data
|
|
|
|
guardrail = CountingGuardrail()
|
|
|
|
data = {
|
|
"model": "gpt-4",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"guardrails": ["counting-guardrail"],
|
|
"metadata": {},
|
|
}
|
|
|
|
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
|
|
|
|
assert guardrail.pre_call_count == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration test: post_call_success_hook with model-level guardrails
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_success_hook_runs_model_level_guardrail():
|
|
"""
|
|
Model-level guardrails configured on a deployment should execute in
|
|
post_call_success_hook (non-streaming path).
|
|
"""
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
from litellm.types.utils import Choices, Message, ModelResponse, Usage
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test-model-guardrail",
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
)
|
|
self.was_called = False
|
|
|
|
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
|
self.was_called = True
|
|
return response
|
|
|
|
guardrail = TestGuardrail()
|
|
|
|
# Mock router that returns a deployment with guardrails configured
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = ["test-model-guardrail"]
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
with (
|
|
patch("litellm.callbacks", [guardrail]),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"model_info": {"id": "model-uuid-123"}},
|
|
}
|
|
response = ModelResponse(
|
|
id="resp-1",
|
|
choices=[
|
|
Choices(
|
|
message=Message(content="Hello", role="assistant"),
|
|
index=0,
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
model="gpt-4",
|
|
usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10),
|
|
)
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
await proxy_logging.post_call_success_hook(
|
|
data=data,
|
|
response=response,
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
assert guardrail.was_called is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_success_hook_skips_guardrail_not_on_model():
|
|
"""
|
|
Guardrails NOT configured on the model should not execute when
|
|
no other source (request body, key, team) enables them.
|
|
"""
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
from litellm.types.utils import Choices, Message, ModelResponse, Usage
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="unrelated-guardrail",
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
)
|
|
self.was_called = False
|
|
|
|
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
|
self.was_called = True
|
|
return response
|
|
|
|
guardrail = TestGuardrail()
|
|
|
|
# Deployment has a DIFFERENT guardrail configured
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = ["some-other-guardrail"]
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
with (
|
|
patch("litellm.callbacks", [guardrail]),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"model_info": {"id": "model-uuid-123"}},
|
|
}
|
|
response = ModelResponse(
|
|
id="resp-1",
|
|
choices=[
|
|
Choices(
|
|
message=Message(content="Hello", role="assistant"),
|
|
index=0,
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
model="gpt-4",
|
|
usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10),
|
|
)
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
await proxy_logging.post_call_success_hook(
|
|
data=data,
|
|
response=response,
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
assert guardrail.was_called is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration test: async_post_call_streaming_iterator_hook with model-level guardrails
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_iterator_hook_runs_model_level_guardrail():
|
|
"""
|
|
Model-level guardrails configured on a deployment should execute in
|
|
async_post_call_streaming_iterator_hook (streaming path) — even when
|
|
`default_on: false` and the guardrail is not in the request body.
|
|
"""
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestStreamingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test-model-guardrail",
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
)
|
|
self.was_called = False
|
|
|
|
async def async_post_call_streaming_iterator_hook(
|
|
self, user_api_key_dict, response, request_data
|
|
):
|
|
self.was_called = True
|
|
async for chunk in response:
|
|
yield chunk
|
|
|
|
guardrail = TestStreamingGuardrail()
|
|
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = ["test-model-guardrail"]
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
async def fake_response():
|
|
yield "chunk-1"
|
|
yield "chunk-2"
|
|
|
|
with (
|
|
patch("litellm.callbacks", [guardrail]),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
request_data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"model_info": {"id": "model-uuid-123"}},
|
|
}
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
chunks = []
|
|
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
|
response=fake_response(),
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data=request_data,
|
|
):
|
|
chunks.append(chunk)
|
|
|
|
assert guardrail.was_called is True
|
|
assert chunks == ["chunk-1", "chunk-2"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_iterator_hook_skips_guardrail_not_on_model():
|
|
"""
|
|
Streaming guardrails NOT configured on the model (and not in the request
|
|
body / key / team) should not execute, even after the dispatcher merge
|
|
runs. Confirms the gate stays closed for unrelated guardrails.
|
|
"""
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestStreamingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="unrelated-guardrail",
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
)
|
|
self.was_called = False
|
|
|
|
async def async_post_call_streaming_iterator_hook(
|
|
self, user_api_key_dict, response, request_data
|
|
):
|
|
self.was_called = True
|
|
async for chunk in response:
|
|
yield chunk
|
|
|
|
guardrail = TestStreamingGuardrail()
|
|
|
|
# Deployment has a DIFFERENT guardrail configured
|
|
mock_router = MagicMock()
|
|
mock_deployment = MagicMock()
|
|
mock_deployment.litellm_params.get.return_value = ["some-other-guardrail"]
|
|
mock_router.get_deployment.return_value = mock_deployment
|
|
|
|
async def fake_response():
|
|
yield "chunk-1"
|
|
|
|
with (
|
|
patch("litellm.callbacks", [guardrail]),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
request_data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"model_info": {"id": "model-uuid-123"}},
|
|
}
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
chunks = []
|
|
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
|
response=fake_response(),
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data=request_data,
|
|
):
|
|
chunks.append(chunk)
|
|
|
|
assert guardrail.was_called is False
|
|
assert chunks == ["chunk-1"]
|