mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
799a559871
The /guardrails/usage/{overview,detail,logs} endpoints resolved guardrails only
from the litellm_guardrailstable Prisma table, so guardrails defined in
config.yaml (which live only in IN_MEMORY_GUARDRAIL_HANDLER) were invisible:
detail 404'd, overview omitted them or rendered them as Custom/Guardrail
orphans, and logs missed their logical-name alias.
Add config-owned accessors (list_config_guardrails, get_config_guardrail_by_id)
to the in-memory handler and use them in the usage endpoints, mirroring the
union/fallback already used by list_guardrails_v2 and get_guardrail_info. Also
preserve guardrail_info when storing a config guardrail (type/description were
dropped at initialize time) and read the Prisma-row / dict / LitellmParams
shapes uniformly.
Resolves LIT-2529
64 lines
2.5 KiB
Python
64 lines
2.5 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
|
|
|
|
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
|
|
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
|
|
|
|
|
def test_initialize_presidio_guardrail():
|
|
"""
|
|
Test that initialize_guardrail correctly uses registered initializers
|
|
for presidio guardrail
|
|
"""
|
|
# Setup test data for a non-custom guardrail (using Presidio as an example)
|
|
test_guardrail = {
|
|
"guardrail_name": "test_presidio_guardrail",
|
|
"litellm_params": {
|
|
"guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
|
|
"mode": "pre_call",
|
|
"presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
|
|
"presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
|
|
},
|
|
}
|
|
|
|
# Call the initialize_guardrail method
|
|
guardrail_handler = InMemoryGuardrailHandler()
|
|
result = guardrail_handler.initialize_guardrail(
|
|
guardrail=test_guardrail,
|
|
)
|
|
|
|
assert result["guardrail_name"] == "test_presidio_guardrail"
|
|
assert result["litellm_params"].guardrail == SupportedGuardrailIntegrations.PRESIDIO.value
|
|
assert result["litellm_params"].mode == "pre_call"
|
|
|
|
|
|
def test_initialize_guardrail_preserves_guardrail_info():
|
|
"""
|
|
Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the
|
|
stored in-memory Guardrail. Dropping it left the Guardrail Monitor's usage
|
|
endpoints unable to render type/description for YAML-defined guardrails.
|
|
"""
|
|
test_guardrail = {
|
|
"guardrail_name": "test_presidio_with_info",
|
|
"litellm_params": {
|
|
"guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
|
|
"mode": "pre_call",
|
|
"presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
|
|
"presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
|
|
},
|
|
"guardrail_info": {"type": "PII", "description": "masks PII"},
|
|
}
|
|
|
|
guardrail_handler = InMemoryGuardrailHandler()
|
|
result = guardrail_handler.initialize_guardrail(guardrail=test_guardrail)
|
|
|
|
assert result is not None
|
|
assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"}
|
|
stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]]
|
|
assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"}
|