feat(guardrails.py): return specific litellm params in /guardrails/list endpoint

support returning mode, default_on and guardrail name on `/guardrails/list` endpoint
This commit is contained in:
Krrish Dholakia 2025-02-10 17:18:51 -08:00
parent 93595826fb
commit e9a861ec32
4 changed files with 54 additions and 5 deletions

View file

@ -49,4 +49,12 @@ general_settings:
router_settings: router_settings:
redis_host: os.environ/REDIS_HOST redis_host: os.environ/REDIS_HOST
redis_password: os.environ/REDIS_PASSWORD redis_password: os.environ/REDIS_PASSWORD
redis_port: os.environ/REDIS_PORT redis_port: os.environ/REDIS_PORT
guardrails:
- guardrail_name: "aporia-pre-guard"
litellm_params:
guardrail: aporia # supported values: "aporia", "lakera"
mode: "during_call"
api_key: os.environ/APORIO_API_KEY
api_base: os.environ/APORIO_API_BASE

View file

@ -25,6 +25,7 @@ def _get_guardrails_list_response(
guardrail_configs.append( guardrail_configs.append(
GuardrailInfoResponse( GuardrailInfoResponse(
guardrail_name=guardrail.get("guardrail_name"), guardrail_name=guardrail.get("guardrail_name"),
litellm_params=guardrail.get("litellm_params"),
guardrail_info=guardrail.get("guardrail_info"), guardrail_info=guardrail.get("guardrail_info"),
) )
) )

View file

@ -1,7 +1,7 @@
from enum import Enum from enum import Enum
from typing import Any, Dict, List, Literal, Optional, TypedDict from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, ConfigDict, Field, SecretStr
from typing_extensions import Required, TypedDict from typing_extensions import Required, TypedDict
""" """
@ -83,7 +83,7 @@ class LakeraCategoryThresholds(TypedDict, total=False):
class LitellmParams(TypedDict): class LitellmParams(TypedDict):
guardrail: str guardrail: str
mode: str mode: str
api_key: str api_key: Optional[str]
api_base: Optional[str] api_base: Optional[str]
# Lakera specific params # Lakera specific params
@ -140,9 +140,28 @@ class DynamicGuardrailParams(TypedDict):
extra_body: Dict[str, Any] extra_body: Dict[str, Any]
class GuardrailLiteLLMParamsResponse(BaseModel):
"""The returned LiteLLM Params object for /guardrails/list"""
guardrail: str
mode: Union[str, List[str]]
default_on: bool = Field(default=False)
def __init__(self, **kwargs):
default_on = kwargs.get("default_on")
if default_on is None:
default_on = False
super().__init__(**kwargs)
class GuardrailInfoResponse(BaseModel): class GuardrailInfoResponse(BaseModel):
guardrail_name: Optional[str] guardrail_name: Optional[str]
guardrail_info: Optional[Dict] # This will contain all other fields litellm_params: GuardrailLiteLLMParamsResponse
guardrail_info: Optional[Dict]
def __init__(self, **kwargs):
super().__init__(**kwargs)
class ListGuardrailsResponse(BaseModel): class ListGuardrailsResponse(BaseModel):

View file

@ -90,3 +90,24 @@ def test_guardrail_list_of_event_hooks():
assert not cg.should_run_guardrail( assert not cg.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.during_call data=data, event_type=GuardrailEventHooks.during_call
) )
def test_guardrail_info_response():
from litellm.types.guardrails import GuardrailInfoResponse, LitellmParams
guardrail_info = GuardrailInfoResponse(
guardrail_name="aporia-pre-guard",
litellm_params=LitellmParams(
guardrail="aporia",
mode="pre_call",
),
guardrail_info={
"guardrail_name": "aporia-pre-guard",
"litellm_params": {
"guardrail": "aporia",
"mode": "always_on",
},
},
)
assert guardrail_info.litellm_params.default_on == False