scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge pull request #30680 from BerriAI/litellm_backport_1_88_x_bp-30543-30542-30274-0617

Browse source 
chore(release): backport #30543, #30542 to stable/1.88.x and cut 1.88.3
This commit is contained in:
yuneng-jiang 2026-06-17 12:44:42 -07:00 • committed by GitHub
commit 9c135abbf0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 582 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/constants.py
View file

@ -190,6 +190,10 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int(
# Override with LITELLM_MAX_CALLBACKS env var for large deployments (e.g., many teams with guardrails)
MAX_CALLBACKS = get_env_int("LITELLM_MAX_CALLBACKS", 100)
# Metadata key recording which pre_call guardrails the proxy loop already ran,
# so the deployment-level hook does not re-run them for the same request
PRE_CALL_EXECUTED_GUARDRAILS_KEY = "_pre_call_executed_guardrails"
# Generic fallback for unknown models
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int(
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)

54
litellm/integrations/custom_guardrail.py
View file

@ -1,3 +1,4 @@
import secrets
from datetime import datetime
from typing import (
TYPE_CHECKING,
@ -43,12 +44,19 @@ if TYPE_CHECKING:
dc = DualCache()
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.exceptions import (
BlockedPiiEntityError,
GuardrailRaisedException,
ModifyResponseException,
)
# Per-process secret tagging each recorded marker. The deployment hook only
# honors markers carrying this token, so a caller cannot forge the metadata
# field to suppress a guardrail on the direct-SDK path that never reaches the
# proxy's metadata sanitizer.
_PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16)
class CustomGuardrail(CustomLogger):
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
@ -325,6 +333,49 @@ class CustomGuardrail(CustomLogger):
return False
def _pre_call_marker(self) -> Optional[str]:
name = self.guardrail_name
if not name:
return None
return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
def mark_pre_call_hook_ran(self, data: Dict[str, Any]) -> None:
"""
Record that this guardrail's ``async_pre_call_hook`` already ran for this
request, so the deployment-level hook does not run it a second time.
The proxy runs pre-call guardrails in ``ProxyLogging.pre_call_hook``. The
router later spreads a deployment's model-level ``guardrails`` into the
top-level request kwargs, which would otherwise re-trigger the same hook
from ``async_pre_call_deployment_hook``.
"""
marker = self._pre_call_marker()
if marker is None:
return
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key)
if isinstance(meta, dict):
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
if isinstance(executed, list):
if marker not in executed:
executed.append(marker)
else:
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
return
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
def _pre_call_hook_already_ran(self, data: Dict[str, Any]) -> bool:
marker = self._pre_call_marker()
if marker is None:
return False
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key)
if isinstance(meta, dict):
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
if isinstance(executed, list) and marker in executed:
return True
return False
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
) -> Optional[dict]:
@ -335,6 +386,9 @@ class CustomGuardrail(CustomLogger):
if litellm_guardrails is None or not isinstance(litellm_guardrails, list):
return kwargs
if self._pre_call_hook_already_ran(kwargs):
return kwargs
if (
self.should_run_guardrail(
data=kwargs, event_type=GuardrailEventHooks.pre_call

16
litellm/litellm_core_utils/logging_callback_manager.py
View file

@ -394,6 +394,22 @@ class LoggingCallbackManager:
+ litellm._async_failure_callback
)
def remove_callback_from_all_lists(self, obj, require_self=False) -> None:
"""
Remove a callback object from every callback list it may have been
promoted into, so a re-initialized callback leaves no stale instance behind.
"""
for callback_list in (
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
self.remove_callback_from_list_by_object(
callback_list, obj, require_self=require_self
)
def get_active_additional_logging_utils_from_custom_logger(
self,
) -> Set[AdditionalLoggingUtils]:

2
litellm/proxy/common_utils/callback_utils.py
View file

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal,
import litellm
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams
@ -490,6 +491,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS = frozenset(
"guardrail_config",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
"disable_global_guardrails",
"disable_global_guardrail",
"opted_out_global_guardrails",

64
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -5,6 +5,8 @@ import os
from datetime import datetime, timezone
from typing import Any, Dict, List, Literal, Optional, Set, Type, cast
from pydantic import ValidationError
import litellm
from litellm import Router
from litellm._logging import verbose_proxy_logger
@ -598,21 +600,25 @@ class InMemoryGuardrailHandler:
def delete_in_memory_guardrail(self, guardrail_id: str) -> None:
"""
Delete a guardrail in memory and remove from litellm callbacks.
The callback is purged from every callback list, not just
litellm.callbacks: request handling promotes guardrail callbacks into the
success/failure/async lists, so removing it from only litellm.callbacks
leaves the old instance stranded in those lists on every re-initialization.
"""
# Remove from in-memory storage
self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None)
self._sources.pop(guardrail_id, None)
# Remove the callback from litellm.callbacks
custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop(
guardrail_id, None
)
if custom_guardrail_callback:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
callback_list=litellm.callbacks,
obj=custom_guardrail_callback,
require_self=False,
)
if custom_guardrail_callback is None:
return
litellm.logging_callback_manager.remove_callback_from_all_lists(
custom_guardrail_callback
)
def list_in_memory_guardrails(self) -> List[Guardrail]:
"""
@ -654,6 +660,34 @@ class InMemoryGuardrailHandler:
self.delete_in_memory_guardrail(guardrail_id)
return stale_ids
@staticmethod
def _normalize_litellm_params_for_comparison(
params: Optional[Any],
) -> Optional[Dict[str, Any]]:
"""
Render litellm_params to a canonical dict so an in-memory LitellmParams and
the raw dict loaded from the DB compare equal when they describe the same
config. The in-memory side is a LitellmParams whose model_dump() carries
every field default and coerces enums, while the DB side is the raw stored
dict holding only the keys originally provided. Comparing those two shapes
directly never matches, so each DB poll would re-initialize the guardrail
forever; normalizing both through LitellmParams keeps the diff meaningful.
"""
if params is None:
return None
if isinstance(params, LitellmParams):
return params.model_dump()
if isinstance(params, dict):
try:
return LitellmParams(**params).model_dump()
except ValidationError as e:
verbose_proxy_logger.warning(
f"Could not normalize guardrail litellm_params for comparison; "
f"treating the guardrail as changed. Error: {e}"
)
return params
return params
def _has_guardrail_params_changed(
self, guardrail_id: str, new_guardrail: Guardrail
) -> bool:
@ -670,19 +704,11 @@ class InMemoryGuardrailHandler:
return True
# Compare litellm_params
existing_params = existing.get("litellm_params")
new_params = new_guardrail.get("litellm_params")
# Convert to dicts for comparison
existing_dict = (
existing_params.model_dump()
if isinstance(existing_params, LitellmParams)
else existing_params
existing_dict = self._normalize_litellm_params_for_comparison(
existing.get("litellm_params")
)
new_dict = (
new_params.model_dump()
if isinstance(new_params, LitellmParams)
else new_params
new_dict = self._normalize_litellm_params_for_comparison(
new_guardrail.get("litellm_params")
)
# Compare and identify specific differences

2
litellm/proxy/litellm_pre_call_utils.py
View file

@ -13,6 +13,7 @@ from starlette.datastructures import Headers
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host
@ -161,6 +162,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS = (
"secret_fields",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
)
_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS = frozenset(

4
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -171,6 +171,10 @@ class PipelineExecutor:
data=data,
call_type=call_type, # type: ignore
)
if isinstance(callback, CustomGuardrail):
callback.mark_pre_call_hook_ran(data)
if isinstance(response, dict):
callback.mark_pre_call_hook_ran(response)
elif mode == "post_call":
response = await target.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,

2
litellm/proxy/utils.py
View file

@ -1150,6 +1150,8 @@ class ProxyLogging:
response=response, data=data, call_type=call_type
)
callback.mark_pre_call_hook_ran(data)
except Exception as e:
status = "error"
error_type = type(e).__name__

4
pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.88.2"
version = "1.88.3"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.14"
@ -264,7 +264,7 @@ source-exclude = [
profile = "black"
[tool.commitizen]
version = "1.88.2"
version = "1.88.3"
version_files = [
"pyproject.toml:^version",
]

23
tests/litellm_utils_tests/test_logging_callback_manager.py
View file

@ -192,6 +192,29 @@ def test_remove_callback_from_list_by_object():
assert len(litellm._async_failure_callback) == 0
def test_remove_callback_from_all_lists():
manager = LoggingCallbackManager()
manager._reset_all_callbacks()
class TestLogger(CustomLogger):
pass
obj = TestLogger()
manager.add_litellm_callback(obj)
manager.add_litellm_success_callback(obj)
manager.add_litellm_failure_callback(obj)
manager.add_litellm_async_success_callback(obj)
manager.add_litellm_async_failure_callback(obj)
manager.remove_callback_from_all_lists(obj)
assert obj not in litellm.callbacks
assert obj not in litellm.success_callback
assert obj not in litellm.failure_callback
assert obj not in litellm._async_success_callback
assert obj not in litellm._async_failure_callback
def test_reset_callbacks(callback_manager):
# Add various callbacks
callback_manager.add_litellm_callback("test")

106
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -84,6 +84,112 @@ class TestCustomGuardrailDeploymentHook:
assert result["messages"] == mock_result["messages"]
assert result["messages"] != original_messages
@pytest.mark.asyncio
async def test_deployment_hook_skips_when_pre_call_already_ran(self):
"""The deployment hook must not re-run async_pre_call_hook once the proxy
pre-call loop has already run it for this request."""
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", 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()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {},
}
guardrail.mark_pre_call_hook_ran(kwargs)
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 0
@pytest.mark.asyncio
async def test_deployment_hook_runs_when_not_marked(self):
"""Without the proxy marker (direct-SDK usage) the deployment hook is the
only execution path and must still run the guardrail exactly once."""
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", 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()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {},
}
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 1
def test_mark_pre_call_hook_ran_uses_litellm_metadata(self):
"""The marker is recorded in litellm_metadata when that is the metadata
bucket in use, and is then visible to the skip check."""
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
guardrail = CustomGuardrail(guardrail_name="g1")
kwargs = {"litellm_metadata": {}}
guardrail.mark_pre_call_hook_ran(kwargs)
assert kwargs["litellm_metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY]
assert guardrail._pre_call_hook_already_ran(kwargs) is True
@pytest.mark.asyncio
async def test_deployment_hook_ignores_forged_caller_marker(self):
"""A direct-SDK caller controls request metadata but cannot know the
per-process token, so a hand-crafted marker must not suppress a
requested guardrail in async_pre_call_deployment_hook."""
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", 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()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]},
}
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 1
class TestCustomGuardrailShouldRunGuardrail:

169
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -180,3 +180,172 @@ def test_sync_guardrail_from_db_marks_source_db_when_unchanged():
handler.sync_guardrail_from_db(g)
assert handler.get_source("collide") == "db"
def _db_litellm_params() -> dict:
"""
Shape produced by GuardrailRegistry.get_all_guardrails_from_db: litellm_params
is a raw dict (not a LitellmParams), holding only the keys originally stored,
a non-schema extra key, and plain-string enum values.
"""
return {
"guardrail": "litellm_content_filter",
"mode": "pre_call",
"default_on": True,
"version": 2,
"blocked_words": [{"keyword": "secret", "action": "BLOCK"}],
}
def test_unchanged_db_params_do_not_register_as_changed():
"""
A DB poll returns litellm_params as a raw dict while the in-memory copy is a
LitellmParams whose model_dump() fills every field default and coerces enums.
The two shapes must compare equal when the config is identical; otherwise
every poll cycle re-initializes the guardrail indefinitely.
"""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "11111111-1111-1111-1111-111111111111"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=dict(raw))
assert handler._has_guardrail_params_changed(gid, new) is False
def test_changed_db_params_register_as_changed():
"""Normalizing both sides must still surface a genuine config change."""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "22222222-2222-2222-2222-222222222222"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
changed = {**raw, "blocked_words": [{"keyword": "different", "action": "BLOCK"}]}
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=changed)
assert handler._has_guardrail_params_changed(gid, new) is True
def test_unnormalizable_db_params_register_as_changed_without_raising():
"""
A DB row whose litellm_params fail LitellmParams validation must not crash the
poll loop. The comparison falls back to treating the guardrail as changed so it
re-initializes (and surfaces the bad row in logs) rather than propagating the
validation error up through the polling cycle.
"""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "55555555-5555-5555-5555-555555555555"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
malformed = {**raw, "default_on": "not-a-bool-xyz"}
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=malformed)
assert handler._has_guardrail_params_changed(gid, new) is True
def _all_callback_lists():
import litellm
return [
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
]
def test_delete_in_memory_guardrail_removes_callback_from_all_lists():
"""
Request handling promotes guardrail callbacks from litellm.callbacks into the
success/failure/async lists. delete_in_memory_guardrail must purge the callback
from every list, otherwise a re-initialized guardrail leaves its old instance
stranded in those lists and instances accumulate.
"""
handler = InMemoryGuardrailHandler()
callback = CustomGuardrail(
guardrail_name="cf-delete",
default_on=True,
event_hook=GuardrailEventHooks.pre_call,
)
gid = "33333333-3333-3333-3333-333333333333"
handler.IN_MEMORY_GUARDRAILS[gid] = _make_guardrail(gid, "cf-delete")
handler._sources[gid] = "db"
handler.guardrail_id_to_custom_guardrail[gid] = callback
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
for cb_list in lists:
cb_list.append(callback)
handler.delete_in_memory_guardrail(gid)
for cb_list in lists:
assert callback not in cb_list
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def test_repeated_db_sync_does_not_accumulate_runner_instances():
"""
End-to-end regression for the OOM: across repeated DB polls (with the config
genuinely changing each cycle to force re-initialization), exactly one live
guardrail instance must exist across all callback lists. On the unfixed code
the stale instance lingers in the success/failure lists and the distinct count
climbs above one.
"""
import litellm
handler = InMemoryGuardrailHandler()
gid = "44444444-4444-4444-4444-444444444444"
name = "cf-accum"
def db_guardrail(word: str) -> Guardrail:
params = {
**_db_litellm_params(),
"blocked_words": [{"keyword": word, "action": "BLOCK"}],
}
return Guardrail(guardrail_id=gid, guardrail_name=name, litellm_params=params)
def promote_into_request_lists() -> None:
manager = litellm.logging_callback_manager
for callback in list(litellm.callbacks):
manager.add_litellm_success_callback(callback)
manager.add_litellm_failure_callback(callback)
manager.add_litellm_async_success_callback(callback)
manager.add_litellm_async_failure_callback(callback)
def distinct_runner_instances() -> int:
seen = set()
for callback in litellm.logging_callback_manager._get_all_callbacks():
if (
isinstance(callback, CustomGuardrail)
and getattr(callback, "guardrail_name", None) == name
):
seen.add(id(callback))
return len(seen)
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
for cycle in range(5):
handler.sync_guardrail_from_db(db_guardrail(f"word-{cycle}"))
promote_into_request_lists()
assert distinct_runner_instances() == 1
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot

152
tests/test_litellm/proxy/test_model_level_guardrails.py
View file

@ -19,7 +19,6 @@ from litellm.proxy.utils import (
_merge_guardrails_with_existing,
)
# ---------------------------------------------------------------------------
# Unit tests for _check_and_merge_model_level_guardrails
# ---------------------------------------------------------------------------
@ -159,6 +158,157 @@ class TestCheckAndMergeModelLevelGuardrails:
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
# ---------------------------------------------------------------------------

4
uv.lock generated
View file

@ -9,7 +9,7 @@ resolution-markers = [
]
[options]
exclude-newer = "2026-06-11T00:01:45.852753Z"
exclude-newer = "2026-06-14T19:22:16.739045Z"
exclude-newer-span = "P3D"
[manifest]
@ -3280,7 +3280,7 @@ wheels = [
[[package]]
name = "litellm"
version = "1.88.2"
version = "1.88.3"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },