mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
808 lines
30 KiB
Python
808 lines
30 KiB
Python
import os
|
|
import io, asyncio
|
|
import json
|
|
import pytest
|
|
import time
|
|
from litellm import mock_completion
|
|
from unittest.mock import MagicMock, AsyncMock, patch
|
|
|
|
import litellm
|
|
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
|
|
_OPTIONAL_PresidioPIIMasking,
|
|
PresidioPerRequestConfig,
|
|
)
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.types.utils import (
|
|
StandardLoggingPayload,
|
|
StandardLoggingGuardrailInformation,
|
|
)
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.caching.caching import DualCache
|
|
from typing import Optional
|
|
|
|
|
|
class CustomLoggerForTesting(CustomLogger):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.standard_logging_payload: Optional[StandardLoggingPayload] = None
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
self.standard_logging_payload = kwargs.get("standard_logging_object")
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_standard_logging_payload_includes_guardrail_information():
|
|
"""
|
|
Test that the standard logging payload includes the guardrail information when a guardrail is applied
|
|
"""
|
|
test_custom_logger = CustomLoggerForTesting()
|
|
litellm.callbacks = [test_custom_logger]
|
|
presidio_guard = _OPTIONAL_PresidioPIIMasking(
|
|
guardrail_name="presidio_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
presidio_analyzer_api_base="https://mock-presidio-analyzer.com/",
|
|
presidio_anonymizer_api_base="https://mock-presidio-anonymizer.com/",
|
|
)
|
|
|
|
# Mock the Presidio API responses
|
|
mock_analyze_response = [
|
|
{
|
|
"analysis_explanation": {
|
|
"recognizer": "PhoneRecognizer",
|
|
"pattern": "phone",
|
|
},
|
|
"start": 26,
|
|
"end": 40,
|
|
"score": 0.75,
|
|
"entity_type": "PHONE_NUMBER",
|
|
}
|
|
]
|
|
|
|
mock_anonymize_response = {
|
|
"text": "Hello, my phone number is <PHONE_NUMBER>",
|
|
"items": [
|
|
{
|
|
"start": 26,
|
|
"end": 40,
|
|
"entity_type": "PHONE_NUMBER",
|
|
"text": "<PHONE_NUMBER>",
|
|
"operator": "replace",
|
|
}
|
|
],
|
|
}
|
|
|
|
# Create mock response objects
|
|
mock_analyze_resp = MagicMock()
|
|
mock_analyze_resp.status = 200
|
|
mock_analyze_resp.content_type = "application/json"
|
|
mock_analyze_resp.json = AsyncMock(return_value=mock_analyze_response)
|
|
|
|
mock_anonymize_resp = MagicMock()
|
|
mock_anonymize_resp.status = 200
|
|
mock_anonymize_resp.content_type = "application/json"
|
|
mock_anonymize_resp.json = AsyncMock(return_value=mock_anonymize_response)
|
|
|
|
# Mock the aiohttp ClientSession with global call tracking
|
|
call_counter = {"count": 0}
|
|
|
|
class MockClientSession:
|
|
def __init__(self):
|
|
self.closed = False
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
async def close(self):
|
|
self.closed = True
|
|
|
|
def post(self, url, json=None, **kwargs):
|
|
class MockResponse:
|
|
def __init__(self, response_obj):
|
|
self.response_obj = response_obj
|
|
|
|
async def __aenter__(self):
|
|
return self.response_obj
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
# Return analyze response first, then anonymize response
|
|
call_counter["count"] += 1
|
|
if "analyze" in url:
|
|
return MockResponse(mock_analyze_resp)
|
|
else:
|
|
return MockResponse(mock_anonymize_resp)
|
|
|
|
# 1. call the pre call hook with guardrail
|
|
request_data = {
|
|
"model": "gpt-5.5",
|
|
"messages": [
|
|
{"role": "user", "content": "Hello, my phone number is +1 412 555 1212"},
|
|
],
|
|
"mock_response": "Hello",
|
|
"guardrails": ["presidio_guard"],
|
|
"metadata": {},
|
|
}
|
|
|
|
with patch("aiohttp.ClientSession", MockClientSession):
|
|
await presidio_guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(),
|
|
cache=DualCache(),
|
|
data=request_data,
|
|
call_type="acompletion",
|
|
)
|
|
|
|
# 2. call litellm.acompletion
|
|
response = await litellm.acompletion(**request_data)
|
|
|
|
# 3. assert that the standard logging payload includes the guardrail information
|
|
await asyncio.sleep(1)
|
|
print(
|
|
"got standard logging payload=",
|
|
json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str),
|
|
)
|
|
assert test_custom_logger.standard_logging_payload is not None
|
|
assert (
|
|
test_custom_logger.standard_logging_payload["guardrail_information"] is not None
|
|
)
|
|
|
|
# guardrail_information is now a list
|
|
assert isinstance(
|
|
test_custom_logger.standard_logging_payload["guardrail_information"], list
|
|
)
|
|
assert len(test_custom_logger.standard_logging_payload["guardrail_information"]) > 0
|
|
|
|
guardrail_info = test_custom_logger.standard_logging_payload[
|
|
"guardrail_information"
|
|
][0]
|
|
assert guardrail_info.get("guardrail_name") == "presidio_guard"
|
|
assert guardrail_info.get("guardrail_mode") == GuardrailEventHooks.pre_call
|
|
|
|
# assert that the guardrail_response is a response from presidio analyze
|
|
presidio_response = guardrail_info.get("guardrail_response")
|
|
assert isinstance(presidio_response, list)
|
|
for response_item in presidio_response:
|
|
assert "analysis_explanation" in response_item
|
|
assert "start" in response_item
|
|
assert "end" in response_item
|
|
assert "score" in response_item
|
|
assert "entity_type" in response_item
|
|
|
|
# assert that the duration is not None
|
|
duration = guardrail_info.get("duration")
|
|
assert duration is not None
|
|
assert duration > 0
|
|
|
|
# assert that we get the count of masked entities
|
|
masked_entity_count = guardrail_info.get("masked_entity_count")
|
|
assert masked_entity_count is not None
|
|
assert masked_entity_count["PHONE_NUMBER"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skip(reason="Local only test")
|
|
async def test_langfuse_trace_includes_guardrail_information():
|
|
"""
|
|
Test that the langfuse trace includes the guardrail information when a guardrail is applied
|
|
"""
|
|
import httpx
|
|
from unittest.mock import AsyncMock, patch
|
|
from litellm.integrations.langfuse.langfuse_prompt_management import (
|
|
LangfusePromptManagement,
|
|
)
|
|
|
|
callback = LangfusePromptManagement(flush_interval=3)
|
|
import json
|
|
|
|
# Create a mock Response object
|
|
mock_response = AsyncMock(spec=httpx.Response)
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {"status": "success"}
|
|
|
|
# Create mock for httpx.Client.post
|
|
mock_post = AsyncMock()
|
|
mock_post.return_value = mock_response
|
|
|
|
with patch("httpx.Client.post", mock_post):
|
|
litellm._turn_on_debug()
|
|
litellm.callbacks = [callback]
|
|
presidio_guard = _OPTIONAL_PresidioPIIMasking(
|
|
guardrail_name="presidio_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
presidio_analyzer_api_base=os.getenv("PRESIDIO_ANALYZER_API_BASE"),
|
|
presidio_anonymizer_api_base=os.getenv("PRESIDIO_ANONYMIZER_API_BASE"),
|
|
)
|
|
# 1. call the pre call hook with guardrail
|
|
request_data = {
|
|
"model": "gpt-5.5",
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "Hello, my phone number is +1 412 555 1212",
|
|
},
|
|
],
|
|
"mock_response": "Hello",
|
|
"guardrails": ["presidio_guard"],
|
|
"metadata": {},
|
|
}
|
|
await presidio_guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(),
|
|
cache=DualCache(),
|
|
data=request_data,
|
|
call_type="acompletion",
|
|
)
|
|
|
|
# 2. call litellm.acompletion
|
|
response = await litellm.acompletion(**request_data)
|
|
|
|
# 3. Wait for async logging operations to complete
|
|
await asyncio.sleep(5)
|
|
|
|
# 4. Verify the Langfuse payload
|
|
assert mock_post.call_count >= 1
|
|
url = mock_post.call_args[0][0]
|
|
request_body = mock_post.call_args[1].get("content")
|
|
|
|
# Parse the JSON body
|
|
actual_payload = json.loads(request_body)
|
|
print("\nLangfuse payload:", json.dumps(actual_payload, indent=2))
|
|
|
|
# Look for the guardrail span in the payload
|
|
guardrail_span = None
|
|
for item in actual_payload["batch"]:
|
|
if (
|
|
item["type"] == "span-create"
|
|
and item["body"].get("name") == "guardrail"
|
|
):
|
|
guardrail_span = item
|
|
break
|
|
|
|
# Assert that the guardrail span exists
|
|
assert guardrail_span is not None, "No guardrail span found in Langfuse payload"
|
|
|
|
# Validate the structure of the guardrail span
|
|
assert guardrail_span["body"]["name"] == "guardrail"
|
|
assert "metadata" in guardrail_span["body"]
|
|
assert guardrail_span["body"]["metadata"]["guardrail_name"] == "presidio_guard"
|
|
assert (
|
|
guardrail_span["body"]["metadata"]["guardrail_mode"]
|
|
== GuardrailEventHooks.pre_call
|
|
)
|
|
assert "guardrail_masked_entity_count" in guardrail_span["body"]["metadata"]
|
|
assert (
|
|
guardrail_span["body"]["metadata"]["guardrail_masked_entity_count"][
|
|
"PHONE_NUMBER"
|
|
]
|
|
== 1
|
|
)
|
|
|
|
# Validate the output format matches the expected structure
|
|
assert "output" in guardrail_span["body"]
|
|
assert isinstance(guardrail_span["body"]["output"], list)
|
|
assert len(guardrail_span["body"]["output"]) > 0
|
|
|
|
# Validate the first output item has the expected structure
|
|
output_item = guardrail_span["body"]["output"][0]
|
|
assert "entity_type" in output_item
|
|
assert output_item["entity_type"] == "PHONE_NUMBER"
|
|
assert "score" in output_item
|
|
assert "start" in output_item
|
|
assert "end" in output_item
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bedrock_guardrail_status_blocked():
|
|
"""
|
|
Test that Bedrock guardrail sets correct status fields when blocking content.
|
|
|
|
This test verifies that when Bedrock guardrail blocks content:
|
|
1. The guardrail_information contains guardrail_status="blocked"
|
|
2. The status_fields.guardrail_status is set to "guardrail_intervened"
|
|
3. The status_fields.llm_api_status remains "success" (mock LLM call succeeds)
|
|
"""
|
|
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
|
BedrockGuardrail,
|
|
)
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
litellm._turn_on_debug()
|
|
|
|
# Setup custom logger to capture standard logging payload
|
|
test_custom_logger = CustomLoggerForTesting()
|
|
litellm.callbacks = [test_custom_logger]
|
|
|
|
# Create Bedrock guardrail with mock AWS credentials
|
|
bedrock_guard = BedrockGuardrail(
|
|
guardrail_name="bedrock_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
guardrailIdentifier="test-id",
|
|
guardrailVersion="1",
|
|
aws_access_key_id="test-key",
|
|
aws_secret_access_key="test-secret",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"action": "GUARDRAIL_INTERVENED",
|
|
"outputs": [{"text": "Blocked"}],
|
|
"assessments": [
|
|
{"topicPolicy": {"topics": [{"name": "harmful", "action": "BLOCKED"}]}}
|
|
],
|
|
}
|
|
with patch.object(
|
|
bedrock_guard.async_handler, "post", AsyncMock(return_value=mock_response)
|
|
):
|
|
request_data = {
|
|
"model": "gpt-5.5",
|
|
"messages": [{"role": "user", "content": "harmful content"}],
|
|
"mock_response": "Hello",
|
|
"metadata": {},
|
|
}
|
|
|
|
# Mock should_run_guardrail to ensure guardrail logic executes
|
|
with patch.object(bedrock_guard, "should_run_guardrail", return_value=True):
|
|
# Call guardrail pre_call hook - this will raise an exception when content is blocked
|
|
try:
|
|
await bedrock_guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(),
|
|
cache=DualCache(),
|
|
data=request_data,
|
|
call_type="completion",
|
|
)
|
|
except Exception:
|
|
# Expected exception when guardrail blocks content
|
|
pass
|
|
|
|
# Call litellm.acompletion to trigger logging callbacks
|
|
# This populates the standard_logging_payload in our custom logger
|
|
response = await litellm.acompletion(**request_data)
|
|
await asyncio.sleep(1)
|
|
|
|
# Verify the standard logging payload was captured
|
|
assert test_custom_logger.standard_logging_payload is not None
|
|
assert (
|
|
test_custom_logger.standard_logging_payload["guardrail_information"] is not None
|
|
)
|
|
assert isinstance(
|
|
test_custom_logger.standard_logging_payload["guardrail_information"], list
|
|
)
|
|
assert len(test_custom_logger.standard_logging_payload["guardrail_information"]) > 0
|
|
|
|
# Verify guardrail information fields (guardrail_information is now a list)
|
|
guardrail_info = test_custom_logger.standard_logging_payload[
|
|
"guardrail_information"
|
|
][0]
|
|
assert guardrail_info.get("guardrail_status") == "guardrail_intervened"
|
|
assert guardrail_info.get("guardrail_provider") == "bedrock"
|
|
|
|
# Verify the new typed status fields
|
|
# guardrail_status should be "guardrail_intervened" when content is blocked
|
|
# llm_api_status should be "success" since the mock LLM call itself succeeded
|
|
status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {})
|
|
assert status_fields.get("llm_api_status") == "success"
|
|
assert status_fields.get("guardrail_status") == "guardrail_intervened"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bedrock_guardrail_status_success():
|
|
"""
|
|
Test that Bedrock guardrail sets correct status fields when allowing content.
|
|
|
|
This test verifies that when Bedrock guardrail allows content through:
|
|
1. The guardrail_information contains guardrail_status="success"
|
|
2. The status_fields.guardrail_status is set to "success"
|
|
3. The status_fields.llm_api_status is "success"
|
|
"""
|
|
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
|
BedrockGuardrail,
|
|
)
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
# Reset callbacks completely to avoid event loop conflicts
|
|
litellm.callbacks = []
|
|
await asyncio.sleep(0.1) # Let previous callbacks finish
|
|
|
|
# Setup custom logger to capture standard logging payload
|
|
test_custom_logger = CustomLoggerForTesting()
|
|
litellm.callbacks = [test_custom_logger]
|
|
|
|
# Create Bedrock guardrail
|
|
bedrock_guard = BedrockGuardrail(
|
|
guardrail_name="bedrock_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
guardrailIdentifier="test-id",
|
|
guardrailVersion="1",
|
|
aws_access_key_id="test-key",
|
|
aws_secret_access_key="test-secret",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
# Mock success response
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"action": "NONE",
|
|
"outputs": [{"text": "Safe content"}],
|
|
"assessments": [],
|
|
}
|
|
with patch.object(
|
|
bedrock_guard.async_handler, "post", AsyncMock(return_value=mock_response)
|
|
):
|
|
request_data = {
|
|
"model": "gpt-5.5",
|
|
"messages": [{"role": "user", "content": "safe content"}],
|
|
"mock_response": "Hello",
|
|
"metadata": {},
|
|
}
|
|
|
|
# Mock should_run_guardrail to return True
|
|
with patch.object(bedrock_guard, "should_run_guardrail", return_value=True):
|
|
await bedrock_guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(),
|
|
cache=DualCache(),
|
|
data=request_data,
|
|
call_type="completion",
|
|
)
|
|
|
|
# Call litellm.acompletion to trigger logging
|
|
response = await litellm.acompletion(**request_data)
|
|
await asyncio.sleep(1)
|
|
|
|
# Check standard logging payload status fields
|
|
assert test_custom_logger.standard_logging_payload is not None
|
|
assert (
|
|
test_custom_logger.standard_logging_payload["guardrail_information"] is not None
|
|
)
|
|
assert isinstance(
|
|
test_custom_logger.standard_logging_payload["guardrail_information"], list
|
|
)
|
|
assert len(test_custom_logger.standard_logging_payload["guardrail_information"]) > 0
|
|
|
|
guardrail_info = test_custom_logger.standard_logging_payload[
|
|
"guardrail_information"
|
|
][0]
|
|
assert guardrail_info.get("guardrail_status") == "success"
|
|
assert guardrail_info.get("guardrail_provider") == "bedrock"
|
|
|
|
# Check status fields
|
|
status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {})
|
|
assert status_fields.get("llm_api_status") == "success"
|
|
assert status_fields.get("guardrail_status") == "success"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bedrock_guardrail_status_failure():
|
|
"""
|
|
Test that Bedrock guardrail sets correct status fields when the API endpoint fails.
|
|
|
|
This test verifies that when Bedrock guardrail API is down/fails:
|
|
1. The guardrail_information contains guardrail_status="failure"
|
|
2. The status_fields.guardrail_status is set to "guardrail_failed_to_respond"
|
|
3. The exception is still raised (maintaining existing behavior)
|
|
"""
|
|
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
|
BedrockGuardrail,
|
|
)
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
import httpx
|
|
|
|
# Reset callbacks completely to avoid event loop conflicts
|
|
litellm.callbacks = []
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Setup custom logger to capture standard logging payload
|
|
test_custom_logger = CustomLoggerForTesting()
|
|
litellm.callbacks = [test_custom_logger]
|
|
|
|
# Create Bedrock guardrail
|
|
bedrock_guard = BedrockGuardrail(
|
|
guardrail_name="bedrock_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
guardrailIdentifier="test-id",
|
|
guardrailVersion="1",
|
|
aws_access_key_id="test-key",
|
|
aws_secret_access_key="test-secret",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
# Mock network failure (endpoint down)
|
|
with patch.object(
|
|
bedrock_guard.async_handler,
|
|
"post",
|
|
AsyncMock(side_effect=httpx.ConnectError("Connection failed")),
|
|
):
|
|
request_data = {
|
|
"model": "gpt-5.5",
|
|
"messages": [{"role": "user", "content": "test content"}],
|
|
"mock_response": "Hello",
|
|
"metadata": {},
|
|
}
|
|
|
|
# Mock should_run_guardrail to return True
|
|
with patch.object(bedrock_guard, "should_run_guardrail", return_value=True):
|
|
# Call guardrail (will raise exception on network failure)
|
|
try:
|
|
await bedrock_guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(),
|
|
cache=DualCache(),
|
|
data=request_data,
|
|
call_type="completion",
|
|
)
|
|
except Exception:
|
|
# Expected exception when endpoint is down
|
|
pass
|
|
|
|
# Call litellm.acompletion to trigger logging
|
|
response = await litellm.acompletion(**request_data)
|
|
await asyncio.sleep(1)
|
|
|
|
# Check standard logging payload status fields
|
|
assert test_custom_logger.standard_logging_payload is not None
|
|
assert (
|
|
test_custom_logger.standard_logging_payload["guardrail_information"] is not None
|
|
)
|
|
assert isinstance(
|
|
test_custom_logger.standard_logging_payload["guardrail_information"], list
|
|
)
|
|
assert len(test_custom_logger.standard_logging_payload["guardrail_information"]) > 0
|
|
|
|
guardrail_info = test_custom_logger.standard_logging_payload[
|
|
"guardrail_information"
|
|
][0]
|
|
assert guardrail_info.get("guardrail_status") == "guardrail_failed_to_respond"
|
|
assert guardrail_info.get("guardrail_provider") == "bedrock"
|
|
|
|
# Check status fields
|
|
status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {})
|
|
assert status_fields.get("llm_api_status") == "success"
|
|
assert status_fields.get("guardrail_status") == "guardrail_failed_to_respond"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_noma_guardrail_status_blocked():
|
|
"""
|
|
Test that Noma guardrail sets correct status fields when blocking content.
|
|
|
|
This test verifies that when Noma guardrail blocks content (verdict=False):
|
|
1. The guardrail_information contains guardrail_status="blocked"
|
|
2. The status_fields.guardrail_status is set to "guardrail_intervened"
|
|
3. The status_fields.llm_api_status remains "success"
|
|
"""
|
|
from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaGuardrail
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
# Reset callbacks completely to avoid event loop conflicts
|
|
litellm.callbacks = []
|
|
await asyncio.sleep(0.1) # Let previous callbacks finish
|
|
|
|
# Setup custom logger to capture standard logging payload
|
|
test_custom_logger = CustomLoggerForTesting()
|
|
litellm.callbacks = [test_custom_logger]
|
|
|
|
# Create Noma guardrail
|
|
noma_guard = NomaGuardrail(
|
|
guardrail_name="noma_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
api_key="test-key",
|
|
monitor_mode=False,
|
|
)
|
|
|
|
# Mock blocked response
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"verdict": False,
|
|
"aggregatedScanResult": True,
|
|
"originalResponse": {
|
|
"prompt": {"topicDetector": {"harmful": {"result": True}}}
|
|
},
|
|
}
|
|
mock_response.raise_for_status = MagicMock()
|
|
with patch.object(
|
|
noma_guard.async_handler, "post", AsyncMock(return_value=mock_response)
|
|
):
|
|
request_data = {
|
|
"model": "gpt-5.5",
|
|
"messages": [{"role": "user", "content": "harmful content"}],
|
|
"mock_response": "Hello",
|
|
"metadata": {},
|
|
}
|
|
|
|
# Mock should_run_guardrail to return True
|
|
with patch.object(noma_guard, "should_run_guardrail", return_value=True):
|
|
# Call guardrail (will raise exception on block)
|
|
try:
|
|
await noma_guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(),
|
|
cache=DualCache(),
|
|
data=request_data,
|
|
call_type="completion",
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
# Call litellm.acompletion to trigger logging
|
|
response = await litellm.acompletion(**request_data)
|
|
await asyncio.sleep(1)
|
|
|
|
# Check standard logging payload status fields
|
|
assert test_custom_logger.standard_logging_payload is not None
|
|
assert (
|
|
test_custom_logger.standard_logging_payload["guardrail_information"] is not None
|
|
)
|
|
assert isinstance(
|
|
test_custom_logger.standard_logging_payload["guardrail_information"], list
|
|
)
|
|
assert len(test_custom_logger.standard_logging_payload["guardrail_information"]) > 0
|
|
|
|
guardrail_info = test_custom_logger.standard_logging_payload[
|
|
"guardrail_information"
|
|
][0]
|
|
assert guardrail_info.get("guardrail_status") == "guardrail_intervened"
|
|
assert guardrail_info.get("guardrail_provider") == "noma"
|
|
|
|
# Check status fields
|
|
status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {})
|
|
assert status_fields.get("llm_api_status") == "success"
|
|
assert status_fields.get("guardrail_status") == "guardrail_intervened"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_noma_guardrail_status_success():
|
|
"""
|
|
Test that Noma guardrail sets correct status fields when allowing content.
|
|
|
|
This test verifies that when Noma guardrail allows content (verdict=True):
|
|
1. The guardrail_information contains guardrail_status="success"
|
|
2. The status_fields.guardrail_status is set to "success"
|
|
3. The status_fields.llm_api_status is "success"
|
|
"""
|
|
from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaGuardrail
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
# Reset callbacks completely to avoid event loop conflicts
|
|
litellm.callbacks = []
|
|
await asyncio.sleep(0.1) # Let previous callbacks finish
|
|
|
|
# Setup custom logger to capture standard logging payload
|
|
test_custom_logger = CustomLoggerForTesting()
|
|
litellm.callbacks = [test_custom_logger]
|
|
|
|
# Create Noma guardrail
|
|
noma_guard = NomaGuardrail(
|
|
guardrail_name="noma_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
api_key="test-key",
|
|
monitor_mode=False,
|
|
)
|
|
|
|
# Mock success response
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"verdict": True,
|
|
"aggregatedScanResult": False,
|
|
"originalResponse": {"prompt": {}},
|
|
}
|
|
mock_response.raise_for_status = MagicMock()
|
|
with patch.object(
|
|
noma_guard.async_handler, "post", AsyncMock(return_value=mock_response)
|
|
):
|
|
request_data = {
|
|
"model": "gpt-5.5",
|
|
"messages": [{"role": "user", "content": "safe content"}],
|
|
"mock_response": "Hello",
|
|
"metadata": {},
|
|
}
|
|
|
|
# Mock should_run_guardrail to return True
|
|
with patch.object(noma_guard, "should_run_guardrail", return_value=True):
|
|
await noma_guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(),
|
|
cache=DualCache(),
|
|
data=request_data,
|
|
call_type="completion",
|
|
)
|
|
|
|
# Call litellm.acompletion to trigger logging
|
|
response = await litellm.acompletion(**request_data)
|
|
await asyncio.sleep(1)
|
|
|
|
# Check standard logging payload status fields
|
|
assert test_custom_logger.standard_logging_payload is not None
|
|
assert (
|
|
test_custom_logger.standard_logging_payload["guardrail_information"] is not None
|
|
)
|
|
assert isinstance(
|
|
test_custom_logger.standard_logging_payload["guardrail_information"], list
|
|
)
|
|
assert len(test_custom_logger.standard_logging_payload["guardrail_information"]) > 0
|
|
|
|
guardrail_info = test_custom_logger.standard_logging_payload[
|
|
"guardrail_information"
|
|
][0]
|
|
assert guardrail_info.get("guardrail_status") == "success"
|
|
assert guardrail_info.get("guardrail_provider") == "noma"
|
|
|
|
# Check status fields
|
|
status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {})
|
|
assert status_fields.get("llm_api_status") == "success"
|
|
assert status_fields.get("guardrail_status") == "success"
|
|
|
|
|
|
def test_guardrail_status_fields_computation():
|
|
"""
|
|
Test that status fields are computed correctly from guardrail information.
|
|
|
|
This unit test verifies the _get_status_fields function correctly maps:
|
|
- guardrail_status="blocked" -> status_fields.guardrail_status="guardrail_intervened" (legacy)
|
|
- guardrail_status="guardrail_intervened" -> status_fields.guardrail_status="guardrail_intervened"
|
|
- guardrail_status="success" -> status_fields.guardrail_status="success"
|
|
- guardrail_status="failure" -> status_fields.guardrail_status="guardrail_failed_to_respond" (legacy)
|
|
- guardrail_status="guardrail_failed_to_respond" -> status_fields.guardrail_status="guardrail_failed_to_respond"
|
|
- no guardrail -> status_fields.guardrail_status="not_run"
|
|
"""
|
|
from litellm.litellm_core_utils.litellm_logging import _get_status_fields
|
|
|
|
# Test guardrail_intervened status (content was blocked by guardrail)
|
|
# guardrail_information is now a list
|
|
intervened_info = [{"guardrail_status": "guardrail_intervened"}]
|
|
status_fields_intervened = _get_status_fields(
|
|
status="success", guardrail_information=intervened_info, error_str=None
|
|
)
|
|
assert status_fields_intervened.get("llm_api_status") == "success"
|
|
assert status_fields_intervened.get("guardrail_status") == "guardrail_intervened"
|
|
|
|
# Test legacy blocked status (for backward compatibility)
|
|
blocked_info = [{"guardrail_status": "blocked"}]
|
|
status_fields_blocked = _get_status_fields(
|
|
status="success", guardrail_information=blocked_info, error_str=None
|
|
)
|
|
assert status_fields_blocked.get("llm_api_status") == "success"
|
|
assert status_fields_blocked.get("guardrail_status") == "guardrail_intervened"
|
|
|
|
# Test success status
|
|
success_info = [{"guardrail_status": "success"}]
|
|
status_fields_success = _get_status_fields(
|
|
status="success", guardrail_information=success_info, error_str=None
|
|
)
|
|
assert status_fields_success.get("llm_api_status") == "success"
|
|
assert status_fields_success.get("guardrail_status") == "success"
|
|
|
|
# Test guardrail_failed_to_respond status
|
|
failed_info = [{"guardrail_status": "guardrail_failed_to_respond"}]
|
|
status_fields_failed = _get_status_fields(
|
|
status="failure", guardrail_information=failed_info, error_str=None
|
|
)
|
|
assert status_fields_failed.get("llm_api_status") == "failure"
|
|
assert status_fields_failed.get("guardrail_status") == "guardrail_failed_to_respond"
|
|
|
|
# Test legacy failure status (for backward compatibility)
|
|
failure_info = [{"guardrail_status": "failure"}]
|
|
status_fields_failure = _get_status_fields(
|
|
status="failure", guardrail_information=failure_info, error_str=None
|
|
)
|
|
assert status_fields_failure.get("llm_api_status") == "failure"
|
|
assert (
|
|
status_fields_failure.get("guardrail_status") == "guardrail_failed_to_respond"
|
|
)
|
|
|
|
# Test no guardrail run
|
|
no_guardrail = None
|
|
status_fields_no_guardrail = _get_status_fields(
|
|
status="success", guardrail_information=no_guardrail, error_str=None
|
|
)
|
|
assert status_fields_no_guardrail.get("llm_api_status") == "success"
|
|
assert status_fields_no_guardrail.get("guardrail_status") == "not_run"
|