mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
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:
parent
93595826fb
commit
e9a861ec32
4 changed files with 54 additions and 5 deletions
|
@ -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
|
|
@ -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"),
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
|
@ -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):
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue