mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
test: add tracing fields tests for custom guardrail base class
This commit is contained in:
parent
69834a1fb8
commit
3d109bc63a
1 changed files with 97 additions and 0 deletions
|
|
@ -4,6 +4,7 @@ import pytest
|
|||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
||||
from litellm.types.utils import GuardrailTracingDetail
|
||||
|
||||
|
||||
class TestCustomGuardrailDeploymentHook:
|
||||
|
|
@ -723,3 +724,99 @@ class TestEventTypeLogging:
|
|||
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(logged_info) == 1
|
||||
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
|
||||
|
||||
|
||||
class TestTracingFieldsPopulation:
|
||||
"""Verify add_standard_logging_guardrail_information_to_request_data passes tracing_detail fields."""
|
||||
|
||||
def test_new_fields_set_on_slg(self):
|
||||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
request_data = {"metadata": {}}
|
||||
cg.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response={"result": "ok"},
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
tracing_detail=GuardrailTracingDetail(
|
||||
guardrail_id="rail-123",
|
||||
policy_template="EU AI Act Article 5",
|
||||
detection_method="regex",
|
||||
confidence_score=0.95,
|
||||
match_details=[{"type": "pattern", "action_taken": "BLOCK"}],
|
||||
patterns_checked=12,
|
||||
alert_recipients=["admin@example.com"],
|
||||
),
|
||||
)
|
||||
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(slg_list) == 1
|
||||
slg = slg_list[0]
|
||||
assert slg["guardrail_id"] == "rail-123"
|
||||
assert slg["policy_template"] == "EU AI Act Article 5"
|
||||
assert slg["detection_method"] == "regex"
|
||||
assert slg["confidence_score"] == 0.95
|
||||
assert slg["patterns_checked"] == 12
|
||||
assert slg["alert_recipients"] == ["admin@example.com"]
|
||||
assert len(slg["match_details"]) == 1
|
||||
|
||||
def test_new_fields_default_to_absent(self):
|
||||
"""When tracing_detail is not passed, new fields are absent from the SLG dict."""
|
||||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
request_data = {"metadata": {}}
|
||||
cg.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response="ok",
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
)
|
||||
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert slg.get("guardrail_id") is None
|
||||
assert slg.get("policy_template") is None
|
||||
assert slg.get("confidence_score") is None
|
||||
|
||||
def test_multiple_guardrails_with_different_policies(self):
|
||||
"""One request, multiple guardrails each with own policy_template."""
|
||||
cg1 = CustomGuardrail(guardrail_name="rail-1")
|
||||
cg2 = CustomGuardrail(guardrail_name="rail-2")
|
||||
request_data = {"metadata": {}}
|
||||
|
||||
cg1.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response="ok",
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
tracing_detail=GuardrailTracingDetail(policy_template="GDPR"),
|
||||
)
|
||||
cg2.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response="blocked",
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_intervened",
|
||||
tracing_detail=GuardrailTracingDetail(policy_template="EU AI Act Article 5"),
|
||||
)
|
||||
|
||||
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(slg_list) == 2
|
||||
assert slg_list[0]["policy_template"] == "GDPR"
|
||||
assert slg_list[1]["policy_template"] == "EU AI Act Article 5"
|
||||
|
||||
def test_classification_field_passed_through(self):
|
||||
"""Classification dict for LLM-judge guardrails is passed through."""
|
||||
cg = CustomGuardrail(guardrail_name="judge-rail")
|
||||
request_data = {"metadata": {}}
|
||||
classification = {
|
||||
"flagged": True,
|
||||
"category": "workplace_emotion_recognition",
|
||||
"article_reference": "Article 5(1)(f)",
|
||||
"confidence": 0.94,
|
||||
"reason": "Request asks to analyze employee sentiment",
|
||||
}
|
||||
cg.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response="blocked",
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_intervened",
|
||||
tracing_detail=GuardrailTracingDetail(
|
||||
classification=classification,
|
||||
detection_method="llm-judge",
|
||||
confidence_score=0.94,
|
||||
),
|
||||
)
|
||||
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert slg["classification"] == classification
|
||||
assert slg["detection_method"] == "llm-judge"
|
||||
assert slg["confidence_score"] == 0.94
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue