mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
d4a32771fd
* fix(proxy): scan batch records with the content hooks that are not guardrails Guardrails were made to run on batch uploads by scanning each record through the pre-call hook with the walk limited to guardrails. That limit exists because the same branch carries the rate limiters and budget accounting, which must count an upload once rather than once per line. It also excluded every enforcement hook written as a plain CustomLogger, so prompt-injection detection, Azure content safety, banned keywords and the blocked-user check never saw a batch record at all. Content that is a hard 400 online reached the provider verbatim through batch. A CustomLogger now declares whether its pre-call hook judges the payload or merely counts the request. The four that judge it opt in, the walk admits them, and both short-circuits learn about them, including the one that decides whether the file is streamed off disk in the first place: a proxy configured only with one of these hooks was skipping the scan entirely. Nothing that counts a request is marked, so an upload still costs one slot and one budget check. * refactor(proxy): drop the per-hook comment the attribute contract already states * test(proxy): make the classification a ledger, and pin the wiring with a real hook The classification test listed the two non-enterprise hooks by hand, so unmarking either enterprise one changed nothing and the mutation matrix passed with both surviving. It now walks the hook registries and fails on any pre-call CustomLogger that is on neither side, which also gives the flag the forcing function it lacked: an enforcement hook added later would otherwise default to off and silently skip batch records, which is the bug being fixed here. Nothing exercised the path the bug actually lived on either, since every test raised its own exception rather than a real hook's. One test now drives the shipped prompt-injection hook through the scan, which pins the part no synthetic exception reaches: a chained exception reads as a failure to judge, so refactoring any of these hooks to `raise ... from` would turn every per-record drop into an aborted upload. Also records why a hook that rewrites the payload for routing stays unmarked, and that only the leaf class is consulted. * test(proxy): set the callback list through monkeypatch rather than writing the global
116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
# +------------------------------+
|
|
#
|
|
# Banned Keywords
|
|
#
|
|
# +------------------------------+
|
|
# Thank you users! We ❤️ you! - Krrish & Ishaan
|
|
## Reject a call / response if it contains certain keywords
|
|
|
|
|
|
from typing import Literal
|
|
import litellm
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.guardrails._content_utils import (
|
|
is_text_content_call_type,
|
|
iter_message_text,
|
|
)
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm._logging import verbose_proxy_logger
|
|
from fastapi import HTTPException
|
|
|
|
|
|
class _ENTERPRISE_BannedKeywords(CustomLogger):
|
|
enforces_request_content: bool = True
|
|
# Class variables or attributes
|
|
def __init__(self):
|
|
banned_keywords_list = litellm.banned_keywords_list
|
|
|
|
if banned_keywords_list is None:
|
|
raise Exception(
|
|
"`banned_keywords_list` can either be a list or filepath. None set."
|
|
)
|
|
|
|
if isinstance(banned_keywords_list, list):
|
|
self.banned_keywords_list = banned_keywords_list
|
|
|
|
if isinstance(banned_keywords_list, str): # assume it's a filepath
|
|
try:
|
|
with open(banned_keywords_list, "r") as file:
|
|
data = file.read()
|
|
self.banned_keywords_list = data.split("\n")
|
|
except FileNotFoundError:
|
|
raise Exception(
|
|
f"File not found. banned_keywords_list={banned_keywords_list}"
|
|
)
|
|
except Exception as e:
|
|
raise Exception(
|
|
f"An error occurred: {str(e)}, banned_keywords_list={banned_keywords_list}"
|
|
)
|
|
|
|
def print_verbose(self, print_statement, level: Literal["INFO", "DEBUG"] = "DEBUG"):
|
|
if level == "INFO":
|
|
verbose_proxy_logger.info(print_statement)
|
|
elif level == "DEBUG":
|
|
verbose_proxy_logger.debug(print_statement)
|
|
|
|
if litellm.set_verbose is True:
|
|
print(print_statement) # noqa
|
|
|
|
def test_violation(self, test_str: str):
|
|
for word in self.banned_keywords_list:
|
|
if word in test_str.lower():
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={"error": f"Keyword banned. Keyword={word}"},
|
|
)
|
|
|
|
async def async_pre_call_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
cache: DualCache,
|
|
data: dict,
|
|
call_type: str, # "completion", "embeddings", "image_generation", "moderation"
|
|
):
|
|
try:
|
|
"""
|
|
- check if user id part of call
|
|
- check if user id part of blocked list
|
|
"""
|
|
self.print_verbose("Inside Banned Keyword List Pre-Call Hook")
|
|
if is_text_content_call_type(call_type):
|
|
for text in iter_message_text(data):
|
|
self.test_violation(test_str=text)
|
|
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
verbose_proxy_logger.exception(
|
|
"litellm.enterprise.enterprise_hooks.banned_keywords::async_pre_call_hook - Exception occurred - {}".format(
|
|
str(e)
|
|
)
|
|
)
|
|
|
|
async def async_post_call_success_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response,
|
|
):
|
|
if not isinstance(response, litellm.ModelResponse):
|
|
return
|
|
|
|
for choice in response.choices:
|
|
if not isinstance(choice, litellm.utils.Choices):
|
|
continue
|
|
message = getattr(choice, "message", None)
|
|
content = getattr(message, "content", None)
|
|
if isinstance(content, str):
|
|
self.test_violation(test_str=content)
|
|
|
|
async def async_post_call_streaming_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: str,
|
|
):
|
|
self.test_violation(test_str=response)
|