mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
Merge branch 'litellm_lakera_v2_gemini_contents' into litellm_passthrough_applied_guardrails_header
This commit is contained in:
yucheng
commit
b61b4084a4
6 changed files with 134 additions and 32 deletions
|
|
@ -221,6 +221,9 @@ def _unmaskable_reason(
|
|||
|
||||
|
||||
def _gemini_items(value: object) -> Sequence[object]:
|
||||
"""A repeated Gemini field as a sequence: a JSON array, or the single object Google also accepts in its place."""
|
||||
if isinstance(value, dict):
|
||||
return (value,)
|
||||
return _OBJECT_LIST.validate_python(value) if isinstance(value, list) else ()
|
||||
|
||||
|
||||
|
|
@ -235,7 +238,7 @@ def _has_gemini_contents(data: object) -> bool:
|
|||
"""True if ``data`` is a Gemini ``generateContent`` body, which carries its prompt in
|
||||
``contents`` and ``systemInstruction`` instead of ``messages``."""
|
||||
fields: Final = as_str_mapping(data) or {}
|
||||
return isinstance(fields.get("contents"), list) or isinstance(
|
||||
return isinstance(fields.get("contents"), list | dict) or isinstance(
|
||||
fields.get("systemInstruction") or fields.get("system_instruction"), dict
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -960,7 +960,7 @@ async def _inherit_auth_guardrails(
|
|||
merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails")
|
||||
if not isinstance(merged, list) or not merged or not _named_guardrail_runs_post_call(parsed_body):
|
||||
return None
|
||||
return tuple(_GUARDRAIL_ENTRIES.validate_python(merged))
|
||||
return _GUARDRAIL_ENTRIES.validate_python(merged)
|
||||
|
||||
|
||||
def _named_guardrail_runs_post_call(data: Mapping[str, object]) -> bool:
|
||||
|
|
|
|||
|
|
@ -184,3 +184,13 @@ def test_monitor_mode_forwards_the_body_unchanged(rig: GeminiLakeraRig) -> None:
|
|||
assert response.status_code == 200, response.text
|
||||
assert len(_lakera_messages(rig)) == 1
|
||||
assert [_JSON.validate_json(request.body) for request in rig.upstream.drain()] == [body]
|
||||
|
||||
|
||||
def test_card_in_a_single_contents_object_is_blocked(rig: GeminiLakeraRig) -> None:
|
||||
body: Final[dict[str, JsonValue]] = {"contents": {"role": "user", "parts": {"text": PII_PROMPT}}}
|
||||
|
||||
response: Final = _send(rig, "block", body)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert _lakera_messages(rig) == ((("user", PII_PROMPT),),)
|
||||
assert rig.upstream.drain() == ()
|
||||
|
|
|
|||
|
|
@ -19,12 +19,15 @@ from pydantic import JsonValue, TypeAdapter
|
|||
|
||||
GUARDRAIL: Final = "passthrough-auth-guardrail"
|
||||
GLOBAL_POST_CALL: Final = "passthrough-global-post-call"
|
||||
KEY_POST_CALL: Final = "passthrough-key-post-call"
|
||||
GUARDED_MODEL: Final = "passthrough-guarded-model"
|
||||
BLOCK_MARKER: Final = "block-this-prompt"
|
||||
OPENAI_ROUTE: Final = "/openai/v1/chat/completions"
|
||||
ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages"
|
||||
GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent"
|
||||
_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
_JSON: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_JSON_LIST: Final = TypeAdapter(list[JsonValue])
|
||||
|
||||
Level = Literal["key", "team", "project"]
|
||||
|
||||
|
|
@ -53,8 +56,20 @@ def _openai_reply(_request: Request) -> Reply:
|
|||
return Reply(body=_OPENAI_BODY)
|
||||
|
||||
|
||||
def _vendor_down(_request: Request) -> Reply:
|
||||
return Reply(status=500, body=b'{"error": "guardrail vendor down"}')
|
||||
def _lakera_clean(_request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"flagged": False, "breakdown": [], "payload": []}).encode())
|
||||
|
||||
|
||||
def _lakera_flagged(_request: Request) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"flagged": True,
|
||||
"breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 0}],
|
||||
"payload": [],
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_reply(_request: Request) -> Reply:
|
||||
|
|
@ -77,6 +92,7 @@ class AuthGuardrailRig:
|
|||
owned: OwnedProxy
|
||||
vendor: Wire
|
||||
post_call_vendor: Wire
|
||||
key_post_call_vendor: Wire
|
||||
upstreams: Mapping[str, Wire]
|
||||
|
||||
|
||||
|
|
@ -85,7 +101,8 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
directory: Final = tmp_path_factory.mktemp("passthrough-auth-guardrails")
|
||||
with ExitStack() as stack:
|
||||
vendor: Final = stack.enter_context(wire_server(_verdict))
|
||||
post_call_vendor: Final = stack.enter_context(wire_server(_vendor_down))
|
||||
post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean))
|
||||
key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged))
|
||||
upstreams: Final = MappingProxyType(
|
||||
{
|
||||
OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)),
|
||||
|
|
@ -99,6 +116,17 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
yaml.safe_dump(
|
||||
{
|
||||
**config,
|
||||
"model_list": [
|
||||
*_JSON_LIST.validate_python(config.get("model_list") or []),
|
||||
{
|
||||
"model_name": GUARDED_MODEL,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.6",
|
||||
"api_key": "synthetic-openai-key",
|
||||
"guardrails": [GUARDRAIL],
|
||||
},
|
||||
},
|
||||
],
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": GUARDRAIL,
|
||||
|
|
@ -110,10 +138,20 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": KEY_POST_CALL,
|
||||
"litellm_params": {
|
||||
"guardrail": "lakera_v2",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"api_base": key_post_call_vendor.url,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": GLOBAL_POST_CALL,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"guardrail": "lakera_v2",
|
||||
"mode": "post_call",
|
||||
"default_on": True,
|
||||
"api_base": post_call_vendor.url,
|
||||
|
|
@ -133,7 +171,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
"GEMINI_API_KEY": "synthetic-gemini-key",
|
||||
}
|
||||
owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path))
|
||||
yield AuthGuardrailRig(owned, vendor, post_call_vendor, upstreams)
|
||||
yield AuthGuardrailRig(owned, vendor, post_call_vendor, key_post_call_vendor, upstreams)
|
||||
|
||||
|
||||
def _key_with_guardrail(scenario: Scenario, level: Level) -> str:
|
||||
|
|
@ -155,6 +193,7 @@ def _send(gateway: Gateway, route: str, body: Mapping[str, JsonValue], key: str)
|
|||
def _drain(rig: AuthGuardrailRig) -> None:
|
||||
rig.vendor.drain()
|
||||
rig.post_call_vendor.drain()
|
||||
rig.key_post_call_vendor.drain()
|
||||
for upstream in rig.upstreams.values():
|
||||
upstream.drain()
|
||||
|
||||
|
|
@ -228,3 +267,21 @@ def test_pre_call_guardrail_relays_the_provider_reply_untouched(rig: AuthGuardra
|
|||
assert response.content == _OPENAI_BODY
|
||||
assert len(rig.vendor.drain()) == 1
|
||||
assert rig.post_call_vendor.drain() == ()
|
||||
|
||||
|
||||
def test_key_post_call_guardrail_runs_on_a_model_with_its_own_guardrails(rig: AuthGuardrailRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(metadata={"guardrails": [KEY_POST_CALL]})
|
||||
_drain(rig)
|
||||
|
||||
response: Final = _send(
|
||||
gateway,
|
||||
OPENAI_ROUTE,
|
||||
{**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": GUARDED_MODEL},
|
||||
key,
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert len(rig.key_post_call_vendor.drain()) == 1
|
||||
assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1
|
||||
|
|
|
|||
|
|
@ -5,16 +5,18 @@ PR checklist requires at least one test in tests/test_litellm/.
|
|||
Additional tests live in tests/guardrails_tests/test_lakera_v2.py.
|
||||
"""
|
||||
|
||||
import importlib, logging, os
|
||||
from collections.abc import Mapping
|
||||
import importlib, json, logging, os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
filter_messages_by_skip_flags,
|
||||
)
|
||||
|
|
@ -1742,23 +1744,39 @@ PROMPT_ATTACK_ONLY_LAKERA_RESPONSE: Final = {
|
|||
|
||||
async def _run_gemini_pre_call(
|
||||
guardrail: LakeraAIGuardrail, lakera_response: Mapping[str, object], data: dict[str, object]
|
||||
) -> AsyncMock:
|
||||
with patch.object(guardrail, "call_v2_guard", new_callable=AsyncMock, return_value=(lakera_response, {})) as call:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
|
||||
cache=MagicMock(),
|
||||
data=data,
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
return call
|
||||
) -> Sequence[object]:
|
||||
"""Run the pre-call hook against a Lakera double and return the messages each guard request carried."""
|
||||
sent: Final = list[object]()
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
sent.append(json.loads(request.content)["messages"])
|
||||
return httpx.Response(200, json={"payload": [], **lakera_response})
|
||||
|
||||
guardrail.async_handler = AsyncHTTPHandler(transport=httpx.MockTransport(respond))
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
|
||||
cache=MagicMock(),
|
||||
data=data,
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
return sent
|
||||
|
||||
|
||||
def _without_metadata(data: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""The request body as forwarded, without the guardrail logging the hook records under ``metadata``."""
|
||||
return {key: value for key, value in data.items() if key != "metadata"}
|
||||
|
||||
|
||||
class TestGeminiContents:
|
||||
def test_reads_system_instruction_and_text_parts(self):
|
||||
assert _build_lakera_inspection_messages(_gemini_request()) == GEMINI_INSPECTION_MESSAGES
|
||||
|
||||
def test_reads_a_single_contents_object_with_a_single_parts_object(self):
|
||||
data: Final = {"contents": {"role": "user", "parts": {"text": "hello world"}}}
|
||||
assert _build_lakera_inspection_messages(data) == [{"role": "user", "content": "hello world"}]
|
||||
|
||||
def test_reads_snake_case_system_instruction_and_joins_text_parts(self):
|
||||
data = {
|
||||
data: Final = {
|
||||
"system_instruction": {"parts": [{"text": "a"}, {"text": "b"}]},
|
||||
"contents": [{"parts": [{"text": "c"}]}, {"role": "function", "parts": [{"text": "d"}]}],
|
||||
}
|
||||
|
|
@ -1770,41 +1788,41 @@ class TestGeminiContents:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clean_body_reaches_lakera_and_is_forwarded_unchanged(self):
|
||||
data = _gemini_request()
|
||||
call = await _run_gemini_pre_call(LakeraAIGuardrail(api_key="test_key"), {"flagged": False}, data)
|
||||
assert list(call.call_args.kwargs["messages"]) == GEMINI_INSPECTION_MESSAGES
|
||||
assert data == _gemini_request()
|
||||
data: Final = _gemini_request()
|
||||
sent: Final = await _run_gemini_pre_call(LakeraAIGuardrail(api_key="test_key"), {"flagged": False}, data)
|
||||
assert sent == [GEMINI_INSPECTION_MESSAGES]
|
||||
assert _without_metadata(data) == _gemini_request()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("on_flagged", ["block", "mask"])
|
||||
async def test_pii_is_blocked_since_contents_cannot_be_masked(self, on_flagged):
|
||||
data = _gemini_request()
|
||||
async def test_pii_is_blocked_since_contents_cannot_be_masked(self, on_flagged: str):
|
||||
data: Final = _gemini_request()
|
||||
with pytest.raises(HTTPException) as blocked:
|
||||
await _run_gemini_pre_call(
|
||||
LakeraAIGuardrail(api_key="test_key", on_flagged=on_flagged), PII_ONLY_LAKERA_RESPONSE, data
|
||||
)
|
||||
assert blocked.value.status_code == 400
|
||||
assert data == _gemini_request()
|
||||
assert _without_metadata(data) == _gemini_request()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("lakera_response", [PROMPT_ATTACK_ONLY_LAKERA_RESPONSE, PII_ONLY_LAKERA_RESPONSE])
|
||||
async def test_advisory_blocks_instead_of_adding_a_messages_field(self, lakera_response):
|
||||
data = _gemini_request()
|
||||
async def test_advisory_blocks_instead_of_adding_a_messages_field(self, lakera_response: Mapping[str, object]):
|
||||
data: Final = _gemini_request()
|
||||
with pytest.raises(HTTPException) as blocked:
|
||||
await _run_gemini_pre_call(
|
||||
LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message"), lakera_response, data
|
||||
)
|
||||
assert blocked.value.status_code == 400
|
||||
assert data == _gemini_request()
|
||||
assert _without_metadata(data) == _gemini_request()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_monitor_forwards_the_body_and_logs_why_it_was_not_masked(self):
|
||||
data = _gemini_request()
|
||||
data: Final = _gemini_request()
|
||||
with patch("litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2.verbose_proxy_logger") as logger:
|
||||
await _run_gemini_pre_call(
|
||||
LakeraAIGuardrail(api_key="test_key", on_flagged="monitor"), PII_ONLY_LAKERA_RESPONSE, data
|
||||
)
|
||||
assert data == _gemini_request()
|
||||
assert _without_metadata(data) == _gemini_request()
|
||||
assert logger.error.call_args.args[1] == (
|
||||
"Gemini contents, which have no write-back path for the redacted content"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import (
|
|||
)
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.proxy.utils import ProxyLogging, _merge_guardrails_with_existing
|
||||
from litellm.types import utils as types_utils
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import CallTypes
|
||||
|
|
@ -11840,6 +11840,20 @@ class TestInheritAuthGuardrails:
|
|||
"team-guard",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_guardrails_stay_named_after_a_model_guardrail_merge(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
_, post_call = await self._run(
|
||||
self._chat_body(),
|
||||
inherit=True,
|
||||
user_api_key_dict=auth,
|
||||
guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),),
|
||||
)
|
||||
|
||||
merged: Final = _merge_guardrails_with_existing(post_call.call_args.kwargs["data"], ["model-guard"])
|
||||
assert merged["metadata"]["guardrails"] == ["key-guard", "model-guard"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_only_guardrails_relay_the_response_untouched(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue