forked from phoenix/litellm-mirror
57 lines
1.5 KiB
Python
57 lines
1.5 KiB
Python
from enum import Enum
|
|
from typing import List, Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict
|
|
from typing_extensions import Required, TypedDict
|
|
|
|
"""
|
|
Pydantic object defining how to set guardrails on litellm proxy
|
|
|
|
litellm_settings:
|
|
guardrails:
|
|
- prompt_injection:
|
|
callbacks: [lakera_prompt_injection, prompt_injection_api_2]
|
|
default_on: true
|
|
enabled_roles: [system, user]
|
|
- detect_secrets:
|
|
callbacks: [hide_secrets]
|
|
default_on: true
|
|
"""
|
|
|
|
class Role(Enum):
|
|
SYSTEM = "system"
|
|
ASSISTANT = "assistant"
|
|
USER = "user"
|
|
|
|
default_roles = [Role.SYSTEM, Role.ASSISTANT, Role.USER];
|
|
|
|
class GuardrailItemSpec(TypedDict, total=False):
|
|
callbacks: Required[List[str]]
|
|
default_on: bool
|
|
logging_only: Optional[bool]
|
|
enabled_roles: Optional[List[Role]]
|
|
|
|
|
|
class GuardrailItem(BaseModel):
|
|
callbacks: List[str]
|
|
default_on: bool
|
|
logging_only: Optional[bool]
|
|
guardrail_name: str
|
|
enabled_roles: List[Role]
|
|
model_config = ConfigDict(use_enum_values=True)
|
|
|
|
def __init__(
|
|
self,
|
|
callbacks: List[str],
|
|
guardrail_name: str,
|
|
default_on: bool = False,
|
|
logging_only: Optional[bool] = None,
|
|
enabled_roles: List[Role] = default_roles,
|
|
):
|
|
super().__init__(
|
|
callbacks=callbacks,
|
|
default_on=default_on,
|
|
logging_only=logging_only,
|
|
guardrail_name=guardrail_name,
|
|
enabled_roles=enabled_roles,
|
|
)
|