mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
2c733c00f5
* test: modernize models used in CircleCI e2e test suites
Replaces obsolete models (gpt-4o, gpt-4o-mini, gpt-3.5-turbo,
claude-3-5-sonnet-20240620, claude-sonnet-4-20250514) with current
equivalents across the e2e_openai_endpoints and
proxy_e2e_anthropic_messages_tests CircleCI jobs.
- gpt-4o -> gpt-5.5 (responses API e2e tests)
- gpt-4o-mini -> gpt-5-mini (websocket responses, oai_misc_config)
- gpt-4o-mini-2024-07-18 -> gpt-4.1-mini-2025-04-14 (fine-tuning,
still actively fine-tunable)
- gpt-4 / gpt-3.5-turbo target_model_names example -> gpt-5.5 /
gpt-5-mini
- bedrock claude-3-5-sonnet-20240620 batch entry -> haiku-4-5-20251001
(also aligning oai_misc_config model_name with what
test_bedrock_batches_api.py actually requests)
- bedrock claude-sonnet-4-20250514 (deprecated, retires 2026-06-15)
-> claude-sonnet-4-5-20250929
* test: point bedrock-claude-sonnet-4 alias at Sonnet 4.6, not 4.5
Greptile/Cursor flagged that after the previous commit, the
bedrock-claude-sonnet-4 alias collided with bedrock-claude-sonnet-4.5
(both pointed to claude-sonnet-4-5-20250929). Rename to
bedrock-claude-sonnet-4.6 and point it at the Sonnet 4.6 Bedrock ID
(us.anthropic.claude-sonnet-4-6, already in the litellm model
registry) so the alias name matches the underlying model version.
* test: modernize models across remaining CI-mounted configs & tests
Expands the modernization sweep to all CircleCI-mounted proxy configs
and to test directories where the model literal is a fixture/route key
(not the test's subject).
Config changes:
- proxy_server_config.yaml: bump gpt-3.5-turbo / gpt-3.5-turbo-1106 /
gpt-4o / gemini-1.5-flash / dall-e-3 underlying models; rename
gpt-3.5-turbo-end-user-test alias to gpt-5-mini-end-user-test; bump
text-embedding-ada-002 underlying to text-embedding-3-small. User-
facing aliases (gpt-3.5-turbo, gpt-4, text-embedding-ada-002, etc.)
preserved for backward compatibility with tests.
- simple_config.yaml, otel_test_config.yaml, spend_tracking_config.yaml:
bump gpt-3.5-turbo underlying to gpt-5-mini.
- pass_through_config.yaml: claude-3-5-sonnet / claude-3-7-sonnet /
claude-3-haiku entries replaced with claude-sonnet-4-5 / claude-
haiku-4-5 / claude-opus-4-7.
- oai_misc_config.yaml: align alias name with the gpt-5-mini rename.
Test changes (proactive: claude-sonnet-4-20250514 / claude-opus-4-
20250514 retire 2026-06-15):
- tests/llm_translation/test_anthropic_completion.py: bump 3 references
+ paired Vertex AI ID to claude-sonnet-4-5.
- tests/llm_translation/test_optional_params.py: bump 2 references.
- tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py
and test_bedrock_anthropic_messages_test.py: bump router fixtures
using the deprecated model IDs.
- tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py:
modernize docstring examples.
- tests/test_end_users.py: update references to renamed alias.
* test: modernize placeholder model literals in router_unit_tests
Mass replace_all on fixture/placeholder model literals across the
router_unit_tests/ suite (model name is a routing key / label, not the
test subject). Sub-agent sweep so far — additional commits will follow
for logging_callback_tests/, enterprise/, top-level tests/test_*.py,
and other CI-mounted dirs.
Mappings applied:
- gpt-3.5-turbo -> gpt-5-mini
- gpt-4 (bare) -> gpt-5.5
- gpt-4o (bare) -> gpt-5
- text-embedding-ada-002 -> text-embedding-3-small
- claude-3-sonnet-20240229 / claude-3-opus-20240229 /
claude-3-haiku-20240307 / claude-3-5-sonnet-20240620 ->
claude-sonnet-4-5-20250929 / claude-opus-4-7 /
claude-haiku-4-5-20251001 as appropriate
Explicitly preserved:
- gpt-4o-mini-* variants (transcribe, tts, etc.) where they're current
- gpt-4-turbo / gpt-4-vision-preview / gpt-4-0613 (subject literals)
- JSONL batch body literals
- Mock LLM response model fields (must match upstream)
- Fake/mock identifiers
* test: modernize placeholder model literals across remaining CI suites
Sub-agent sweep across logging_callback_tests/, guardrails_tests/,
enterprise/, pass_through_unit_tests/, otel_tests/,
llm_responses_api_testing/, batches_tests/, spend_tracking_tests/,
litellm_utils_tests/, unified_google_tests/, and a few top-level
tests/test_*.py files where the model literal is a fixture or
placeholder (router model_list, mock standard logging payload, mock
callback data) rather than the test's subject.
Mappings applied (see scope notes below):
- gpt-3.5-turbo -> gpt-5-mini
- gpt-4 (bare) -> gpt-5.5
- gpt-4o (bare) -> gpt-5.5 (corrected from initial gpt-5 — bare gpt-5
is not a valid OpenAI alias; only gpt-5.5 / gpt-5.4 / gpt-5.2-codex
/ gpt-5-mini exist)
- gpt-4o-mini (bare) -> gpt-5-mini
- text-embedding-ada-002 -> text-embedding-3-small
- claude-3-sonnet-20240229 -> claude-sonnet-4-5-20250929
- claude-3-opus-20240229 -> claude-opus-4-7
- claude-3-haiku-20240307 -> claude-haiku-4-5-20251001
- claude-3-5-sonnet-20240620/20241022 -> claude-sonnet-4-5-20250929
- claude-3-7-sonnet-20250219 -> claude-sonnet-4-6
- gemini-1.5-flash -> gemini-2.5-flash
- gemini-1.5-pro -> gemini-2.5-pro
Explicitly preserved (not modernized):
- llm_translation/ tests where model is the SUBJECT (provider-specific
translation/transformation logic). Only the deprecated 20250514
references were already bumped in a prior commit.
- Cost-calc / tokenizer subject tests in test_utils.py (skip-ranges
documented by the sub-agent).
- Bedrock model IDs in test_health_check.py path-stripping tests.
- JSONL batch request bodies and mock LLM response bodies (must match
upstream literal).
- Langfuse expected-request-body JSON fixtures (cost values are exact-
match-asserted; changing the model would shift response_cost).
- gpt-3.5-turbo-instruct (text-completion endpoint; no modern OpenAI
equivalent).
- Top-level tests calling the proxy through user-facing aliases
(gpt-3.5-turbo, gpt-4, text-embedding-ada-002, dall-e-3) — aliases
in proxy_server_config.yaml stay; only the underlying model was
bumped.
- tests/test_gpt5_azure_temperature_support.py (the test's whole point
is model-name handling).
- Fake / mock / openai/fake identifiers.
Notable side fixes:
- test_spend_accuracy_tests.py: UPSTREAM_MODEL now matches what
spend_tracking_config.yaml's proxy actually routes to (gpt-5-mini),
resolving a latent inconsistency.
- proxy_server_config.yaml: bare `gpt-5` alias renamed to `gpt-5.5`
(bare gpt-5 is not a valid OpenAI alias).
- test_batches_logging_unit_tests.py: explicit_models list entries
kept distinct (gpt-5-mini + gpt-5.5) after bulk rename.
* test: fix CI failures from model modernization sweep
CI surfaced 4 categories of regression from the bulk modernization:
1. Azure deployment names are customer-specific. Reverted:
- tests/litellm_utils_tests/test_health_check.py: azure/text-
embedding-3-small -> azure/text-embedding-ada-002 (the CI Azure
account does not have a text-embedding-3-small deployment).
- tests/logging_callback_tests/test_custom_callback_router.py:
same revert for two router fixtures driving aembedding.
2. gpt-5 family does not accept temperature != 1. Tests that pass a
custom temperature swapped from gpt-5-mini to gpt-4.1-mini (modern
non-reasoning OpenAI mini that still accepts temperature/logprobs):
- tests/logging_callback_tests/test_datadog.py
- tests/logging_callback_tests/test_langsmith_unit_test.py
- tests/logging_callback_tests/test_otel_logging.py
3. proxy_server_config.yaml's gpt-3.5-turbo-large alias was routing to
gpt-5.5 (a reasoning model that rejects logprobs). The proxy test
tests/test_openai_endpoints.py::test_chat_completion_streaming
exercises logprobs/top_logprobs through that alias. Bumped the
underlying model to gpt-4.1 (non-reasoning, still modern).
4. tests/logging_callback_tests/test_gcs_pub_sub.py asserts against a
pinned JSON fixture (gcs_pub_sub_body/spend_logs_payload.json) with
hardcoded model="gpt-4o" and a model-specific spend value. Reverted
the litellm.acompletion calls in the test to model="gpt-4o" so the
fixture's exact-match assertions still hold.
5. tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py:
anthropic.messages.create routing to openai/gpt-5-mini returned an
empty content[0] with max_tokens=100 (reasoning-token consumption).
Swapped to openai/gpt-4.1-mini.
* test: fix Assistants API model + 2 cursor[bot] review nits
1. pass_through_unit_tests/test_custom_logger_passthrough.py: gpt-5.5
isn't accepted by the /v1/assistants endpoint
("unsupported_model"). Switch to gpt-4.1-mini (modern, Assistants-
API-supported, non-reasoning).
2. example_config_yaml/pass_through_config.yaml: the previous sweep
bumped the claude-3-7-sonnet alias to claude-opus-4-7, which is a
tier change (Sonnet -> Opus). Map to claude-sonnet-4-6 to keep the
Sonnet tier intact. (Cursor bugbot review.)
3. example_config_yaml/simple_config.yaml: model_name was left as
gpt-3.5-turbo while the underlying was bumped to gpt-5-mini, which
muddles the "simple" example. Make both sides gpt-5-mini so the
most basic example is a straight 1:1 mapping again. (Cursor bugbot
review.)
* fix: revert gpt-4/gpt-3.5-turbo alias underlying to non-reasoning models
tests/test_openai_endpoints.py::test_completion calls the proxy alias
"gpt-4" with temperature=0, and other tests call gpt-3.5-turbo with
custom temperature / logprobs / the legacy /v1/completions endpoint.
The earlier modernization mapped both aliases to gpt-5.5 / gpt-5-mini,
which are reasoning models that reject temperature != 1 and don't
expose /v1/completions. Map the aliases to gpt-4.1 / gpt-4.1-mini
(modern non-reasoning OpenAI models) instead — keeps user-facing
aliases preserved while picking a current underlying that still
supports the parameters/endpoints the tests exercise.
736 lines
24 KiB
Python
736 lines
24 KiB
Python
import sys
|
|
import os
|
|
import io, asyncio
|
|
import pytest
|
|
import time
|
|
from litellm import mock_completion
|
|
from unittest.mock import MagicMock, AsyncMock, patch
|
|
|
|
sys.path.insert(0, os.path.abspath("../.."))
|
|
import litellm
|
|
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import LakeraAIGuardrail
|
|
from litellm.types.guardrails import PiiEntityType, PiiAction
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
|
from fastapi import HTTPException
|
|
from litellm.types.utils import CallTypes as LitellmCallTypes, ModelResponse
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_pre_call_hook_for_pii_masking():
|
|
"""Test for Lakera guardrail pre-call hook for PII masking"""
|
|
# Setup the guardrail with specific entities config
|
|
litellm._turn_on_debug()
|
|
lakera_guardrail = LakeraAIGuardrail(
|
|
api_key="test_key",
|
|
)
|
|
|
|
# Mock response with PII detections in payload (with start/end positions for masking)
|
|
mock_response = {
|
|
"payload": [
|
|
{
|
|
"detector_type": "pii/credit_card",
|
|
"start": 18,
|
|
"end": 37,
|
|
"message_id": 1,
|
|
}, # "4111-1111-1111-1111"
|
|
{
|
|
"detector_type": "pii/email",
|
|
"start": 54,
|
|
"end": 70,
|
|
"message_id": 1,
|
|
}, # "test@example.com"
|
|
],
|
|
"flagged": True,
|
|
"breakdown": [
|
|
{"detector_type": "pii/credit_card", "detected": True, "message_id": 1},
|
|
{"detector_type": "pii/email", "detected": True, "message_id": 1},
|
|
],
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
# Create a sample request with PII data
|
|
data = {
|
|
"messages": [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{
|
|
"role": "user",
|
|
"content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567",
|
|
},
|
|
],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
# Mock objects needed for the pre-call hook
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
cache = DualCache()
|
|
|
|
# Call the pre-call hook with the specified call type
|
|
modified_data = await lakera_guardrail.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=cache,
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
print(modified_data)
|
|
|
|
# Verify the messages have been modified to mask PII
|
|
assert (
|
|
modified_data["messages"][0]["content"] == "You are a helpful assistant."
|
|
) # System prompt should be unchanged
|
|
|
|
user_message = modified_data["messages"][1]["content"]
|
|
# Verify both credit card and email are masked
|
|
assert "4111-1111-1111-1111" not in user_message
|
|
assert "test@example.com" not in user_message
|
|
# Verify masking placeholders are present
|
|
assert "[MASKED CREDIT_CARD]" in user_message
|
|
assert "[MASKED EMAIL]" in user_message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_blocks_non_pii_violations():
|
|
"""Test that Lakera guardrail blocks requests with non-PII violations like hate speech, violence, etc."""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(
|
|
api_key="test_key",
|
|
)
|
|
|
|
# Mock the call_v2_guard method to return a response similar to the user's example
|
|
mock_response = {
|
|
"payload": [],
|
|
"flagged": True,
|
|
"dev_info": {
|
|
"git_revision": "f0bc093a",
|
|
"git_timestamp": "2025-09-23T15:28:06+00:00",
|
|
"model_version": "lakera-guard-1",
|
|
"version": "2.0.281",
|
|
},
|
|
"metadata": {"request_uuid": "b7cd4c8a-28aa-4285-a245-2befee514dbf"},
|
|
"breakdown": [
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/crime",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/hate",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/violence",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-prompt-attack",
|
|
"detector_type": "prompt_attack",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/email",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
],
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
# Create a sample request that would trigger violations
|
|
data = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "Some harmful content that triggers violations",
|
|
}
|
|
],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
# Mock objects needed for the pre-call hook
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
cache = DualCache()
|
|
|
|
# The guardrail should raise an HTTPException for non-PII violations
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await lakera_guardrail.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=cache,
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
# Verify the exception details include the Lakera response
|
|
assert exc_info.value.status_code == 400
|
|
assert "Violated guardrail policy" in str(exc_info.value.detail)
|
|
assert "lakera_guardrail_response" in exc_info.value.detail
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_only_pii_violations_are_masked():
|
|
"""Test that Lakera guardrail only masks PII violations and doesn't block the request."""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(
|
|
api_key="test_key",
|
|
)
|
|
|
|
# Mock response with only PII violations
|
|
mock_response = {
|
|
"payload": [
|
|
{"detector_type": "pii/email", "start": 10, "end": 25, "message_id": 0}
|
|
],
|
|
"flagged": True,
|
|
"breakdown": [
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"detector_type": "pii/email",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"detector_type": "moderated_content/hate",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"detector_type": "prompt_attack",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
],
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
data = {
|
|
"messages": [{"role": "user", "content": "My email test@example.com here"}],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
cache = DualCache()
|
|
|
|
# Should not raise an exception, just mask the PII
|
|
result = await lakera_guardrail.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=cache,
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
# Verify the request was not blocked
|
|
assert result is not None
|
|
assert "messages" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_blocks_flagged_content_with_user_scenario():
|
|
"""
|
|
Test the exact user scenario where Lakera flagged content but request went through.
|
|
This should now be blocked with the fix to check breakdown field instead of payload.
|
|
"""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(
|
|
api_key="test_key",
|
|
)
|
|
|
|
# Mock response matching the exact user scenario
|
|
mock_response = {
|
|
"payload": [], # Empty payload like in user's case
|
|
"flagged": True,
|
|
"dev_info": {
|
|
"git_revision": "f0bc093a",
|
|
"git_timestamp": "2025-09-23T15:28:06+00:00",
|
|
"model_version": "lakera-guard-1",
|
|
"version": "2.0.281",
|
|
},
|
|
"metadata": {"request_uuid": "b7cd4c8a-28aa-4285-a245-2befee514dbf"},
|
|
"breakdown": [
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/crime",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/hate",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/profanity",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/sexual",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/violence",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-moderated-content",
|
|
"detector_type": "moderated_content/weapons",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/address",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/credit_card",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/email",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/iban_code",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/ip_address",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/name",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/phone_number",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-pii",
|
|
"detector_type": "pii/us_social_security_number",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-prompt-attack",
|
|
"detector_type": "prompt_attack",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{
|
|
"project_id": "project-9770817088",
|
|
"policy_id": "policy-lakera-default",
|
|
"detector_id": "detector-lakera-default-unknown-links",
|
|
"detector_type": "unknown_links",
|
|
"detected": False,
|
|
"message_id": 0,
|
|
},
|
|
],
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
# Create a sample request that would trigger violations
|
|
data = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "Some harmful content that should be blocked",
|
|
}
|
|
],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
# Mock objects needed for the pre-call hook
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
cache = DualCache()
|
|
|
|
# With the fix, this should now raise an HTTPException instead of letting the request through
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await lakera_guardrail.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=cache,
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
# Verify the exception details
|
|
assert exc_info.value.status_code == 400
|
|
assert "Violated guardrail policy" in str(exc_info.value.detail)
|
|
assert "lakera_guardrail_response" in exc_info.value.detail
|
|
|
|
# Verify the full response is included in the exception
|
|
lakera_response = exc_info.value.detail["lakera_guardrail_response"]
|
|
assert lakera_response["flagged"] is True
|
|
assert (
|
|
lakera_response["metadata"]["request_uuid"]
|
|
== "b7cd4c8a-28aa-4285-a245-2befee514dbf"
|
|
)
|
|
assert (
|
|
len(lakera_response["breakdown"]) == 16
|
|
) # All the breakdown items from the user's scenario
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_monitor_mode_allows_flagged_content():
|
|
"""Test that monitor mode logs violations but allows requests to proceed."""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(
|
|
api_key="test_key",
|
|
on_flagged="monitor", # Monitor mode
|
|
)
|
|
|
|
# Mock response with violations
|
|
mock_response = {
|
|
"payload": [],
|
|
"flagged": True,
|
|
"breakdown": [
|
|
{
|
|
"detector_type": "moderated_content/violence",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
{"detector_type": "prompt_attack", "detected": True, "message_id": 0},
|
|
],
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
data = {
|
|
"messages": [{"role": "user", "content": "Some harmful content"}],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
cache = DualCache()
|
|
|
|
# Should NOT raise an exception in monitor mode
|
|
result = await lakera_guardrail.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=cache,
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
# Verify request was allowed through
|
|
assert result is not None
|
|
assert "messages" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_block_mode_raises_exception():
|
|
"""Test that block mode (default) raises HTTPException for violations."""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(
|
|
api_key="test_key",
|
|
on_flagged="block", # Block mode (default)
|
|
)
|
|
|
|
mock_response = {
|
|
"payload": [],
|
|
"flagged": True,
|
|
"breakdown": [
|
|
{
|
|
"detector_type": "moderated_content/violence",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
],
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
data = {
|
|
"messages": [{"role": "user", "content": "Harmful content"}],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
cache = DualCache()
|
|
|
|
# Should raise HTTPException in block mode
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await lakera_guardrail.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=cache,
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_monitor_mode_during_call():
|
|
"""Test monitor mode works with during_call (moderation_hook)."""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(
|
|
api_key="test_key",
|
|
on_flagged="monitor",
|
|
)
|
|
|
|
mock_response = {
|
|
"payload": [],
|
|
"flagged": True,
|
|
"breakdown": [
|
|
{"detector_type": "prompt_attack", "detected": True, "message_id": 0},
|
|
],
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
data = {
|
|
"messages": [{"role": "user", "content": "Test content"}],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
|
|
# Should NOT raise exception in monitor mode
|
|
result = await lakera_guardrail.async_moderation_hook(
|
|
data=data, user_api_key_dict=user_api_key_dict, call_type="completion"
|
|
)
|
|
|
|
assert result is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_post_call_blocks_flagged_content():
|
|
"""Post-call hook should block when violations are flagged."""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(api_key="test_key")
|
|
|
|
mock_response = {
|
|
"payload": [],
|
|
"flagged": True,
|
|
"breakdown": [
|
|
{
|
|
"detector_type": "moderated_content/violence",
|
|
"detected": True,
|
|
"message_id": 0,
|
|
},
|
|
],
|
|
}
|
|
|
|
# Mock LLM response object
|
|
llm_response = MagicMock()
|
|
llm_response.model_dump.return_value = {
|
|
"choices": [{"message": {"role": "assistant", "content": "some response"}}]
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
data = {
|
|
"messages": [{"role": "user", "content": "Harmful content"}],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await lakera_guardrail.async_post_call_success_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=llm_response,
|
|
)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_post_call_allows_clean_content():
|
|
"""Post-call hook should allow when not flagged."""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(api_key="test_key")
|
|
|
|
mock_response = {
|
|
"payload": [],
|
|
"flagged": False,
|
|
"breakdown": [],
|
|
}
|
|
|
|
llm_response = MagicMock()
|
|
llm_response.model_dump.return_value = {
|
|
"choices": [{"message": {"role": "assistant", "content": "clean response"}}]
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
data = {
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
|
|
result = await lakera_guardrail.async_post_call_success_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=llm_response,
|
|
)
|
|
|
|
assert result is llm_response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_post_call_masks_pii_and_allows():
|
|
"""Post-call hook should mask PII-only violations and allow response."""
|
|
|
|
lakera_guardrail = LakeraAIGuardrail(api_key="test_key")
|
|
|
|
mock_response = {
|
|
"payload": [
|
|
{"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 1}
|
|
],
|
|
"flagged": True,
|
|
"breakdown": [
|
|
{"detector_type": "pii/email", "detected": True, "message_id": 1},
|
|
],
|
|
}
|
|
|
|
llm_response = MagicMock()
|
|
llm_response.model_dump.return_value = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "Your email is test@example.com",
|
|
}
|
|
},
|
|
]
|
|
}
|
|
|
|
with patch.object(
|
|
lakera_guardrail, "call_v2_guard", new_callable=AsyncMock
|
|
) as mock_call:
|
|
mock_call.return_value = (mock_response, {})
|
|
|
|
data = {
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"model": "gpt-5-mini",
|
|
"metadata": {},
|
|
}
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
|
|
|
result = await lakera_guardrail.async_post_call_success_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=llm_response,
|
|
)
|
|
|
|
assert isinstance(
|
|
result, ModelResponse
|
|
), "PII masking path must return ModelResponse"
|
|
result_dict = result.model_dump()
|
|
assert (
|
|
result_dict["choices"][0]["message"]["content"]
|
|
!= "Your email is test@example.com"
|
|
)
|
|
assert "[MASKED" in result_dict["choices"][0]["message"]["content"]
|