mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
8cf2e2eb43
Key and team `router_settings.model_group_alias` was accepted, persisted and echoed back by `/key/info`, but never applied at request time, so the request ran on the group the caller asked for. `route_request` forwards only the settings the Router accepts as per-request kwargs, and `model_group_alias` is not one of them: the Router resolves aliases from its own instance attribute, which holds the global config map and is shared across requests. Resolve the alias in the proxy instead, alongside the existing model-alias rewrites and ahead of the pre-call hooks, so per-model limits and guardrails key off the group that actually serves the request. Authorize the alias target before the rewrite; model access was checked against the requested group, so a key whose alias points at a group it cannot call gets the usual 403 rather than being quietly served it. Resolves LIT-4879
664 lines
25 KiB
Python
664 lines
25 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 AsyncMock, 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 AND the
|
|
model alias does not resolve to a deployment."""
|
|
data = {"model": "gpt-4", "metadata": {}}
|
|
mock_router = MagicMock()
|
|
# Neither the model_id lookup nor the alias-fallback lookup
|
|
# finds a deployment.
|
|
mock_router.get_deployment.return_value = None
|
|
mock_router.get_deployment_by_model_group_name.return_value = None
|
|
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"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Regression: pre_call ordering — _check_and_merge_model_level_guardrails
|
|
# must run BEFORE pre_call_hook so DB/UI-configured guardrails fire on
|
|
# pre_call paths (#29652; #23774 only covered post_call).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_merges_model_level_guardrails_before_pre_call_hook():
|
|
"""
|
|
common_processing_pre_call_logic must merge model-level guardrails into
|
|
data BEFORE proxy_logging_obj.pre_call_hook is invoked. Otherwise
|
|
pre_call guardrails (e.g. apply_guardrail event) never see the
|
|
UI/DB-assigned guardrail name.
|
|
"""
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
|
|
# Stub router that reports one deployment in the group with one
|
|
# model-level guardrail. Mirrors the real proxy: at pre_call_hook time
|
|
# model_info has been stripped by add_litellm_data_to_request (see
|
|
# veria-ai review on PR #29654) and route_request hasn't yet populated
|
|
# model_info.id — so the resolver has to fall back to the model alias
|
|
# and union guardrails across all deployments in the group.
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment.return_value = None
|
|
mock_router.get_model_list.return_value = [{"litellm_params": {"guardrails": ["my-pre-call-guardrail"]}}]
|
|
|
|
processing = ProxyBaseLLMRequestProcessing(
|
|
data={
|
|
"model": "my-model",
|
|
"metadata": {}, # model_info already stripped
|
|
}
|
|
)
|
|
|
|
captured_pre_call_guardrails: list = []
|
|
|
|
async def fake_pre_call_hook(*, user_api_key_dict, data, call_type):
|
|
# Snapshot the list rather than the dict: metadata is shared by
|
|
# reference, so a merge that happens after this point would otherwise
|
|
# show up here retroactively and the assertion would pass either way.
|
|
captured_pre_call_guardrails.extend(
|
|
(data.get("metadata") or {}).get("guardrails") or data.get("guardrails") or []
|
|
)
|
|
return data
|
|
|
|
proxy_logging = MagicMock()
|
|
proxy_logging.pre_call_hook = fake_pre_call_hook
|
|
|
|
# Minimal stubs for the surrounding setup steps in
|
|
# common_processing_pre_call_logic. We only care about the ordering
|
|
# between _check_and_merge_model_level_guardrails and pre_call_hook.
|
|
async def passthrough_add_litellm_data(*, data, **kwargs):
|
|
return data
|
|
|
|
proxy_config = MagicMock()
|
|
proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None)
|
|
|
|
# Assert on what pre_call_hook was handed rather than short-circuiting the
|
|
# function part way through: a sentinel keyed to one particular later call
|
|
# silently stops testing the ordering as soon as that call moves.
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.common_request_processing.add_litellm_data_to_request",
|
|
side_effect=passthrough_add_litellm_data,
|
|
),
|
|
patch(
|
|
"litellm.proxy.common_request_processing.litellm.utils.function_setup",
|
|
return_value=(MagicMock(), processing.data),
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client",
|
|
None,
|
|
),
|
|
):
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
|
|
await processing.common_processing_pre_call_logic(
|
|
request=MagicMock(headers={}, url=MagicMock(path="/v1/chat/completions")),
|
|
general_settings={},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
|
|
proxy_logging_obj=proxy_logging,
|
|
proxy_config=proxy_config,
|
|
route_type="acompletion",
|
|
version=None,
|
|
user_model=None,
|
|
user_temperature=None,
|
|
user_request_timeout=None,
|
|
user_max_tokens=None,
|
|
user_api_base=None,
|
|
model=None,
|
|
llm_router=mock_router,
|
|
)
|
|
|
|
# The pre_call_hook must have received data with the model-level
|
|
# guardrail already merged in. Before the fix, this assertion fails
|
|
# because pre_call_hook saw the original data without merge.
|
|
assert "my-pre-call-guardrail" in captured_pre_call_guardrails
|