mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
Merge branch 'litellm_lakera_v2_gemini_contents' into litellm_passthrough_applied_guardrails_header
This commit is contained in:
yucheng
commit
fdb6e5f65d
9 changed files with 394 additions and 9 deletions
|
|
@ -365,7 +365,10 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
top_level_metadata: Final = request_data.get("metadata", {})
|
||||
|
||||
# Merge both sources, preferring litellm_metadata if both exist
|
||||
metadata_dict: Final = {**top_level_metadata, **litellm_metadata}
|
||||
metadata_dict: Final = {
|
||||
**(top_level_metadata if isinstance(top_level_metadata, dict) else {}),
|
||||
**(litellm_metadata if isinstance(litellm_metadata, dict) else {}),
|
||||
}
|
||||
|
||||
if not metadata_dict:
|
||||
return result_metadata
|
||||
|
|
|
|||
|
|
@ -219,7 +219,10 @@ class lakeraAI_Moderation(CustomGuardrail):
|
|||
}
|
||||
)
|
||||
elif "input" in data and isinstance(data["input"], list):
|
||||
text = "\n".join(data["input"])
|
||||
text = "\n".join(item for item in data["input"] if isinstance(item, str))
|
||||
if not text:
|
||||
verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text")
|
||||
return
|
||||
_json_data = json.dumps(
|
||||
{
|
||||
"input": text,
|
||||
|
|
|
|||
|
|
@ -279,11 +279,14 @@ def _gemini_tool_text(content: object) -> str:
|
|||
def _has_gemini_contents(data: object, call_type: str | None) -> bool:
|
||||
"""True if ``data`` is a Gemini ``generateContent`` body on a pass-through or Google ``generateContent`` call,
|
||||
which carries its prompt in ``contents`` and ``systemInstruction``. A pass-through body with chat ``messages`` or
|
||||
Responses ``input`` is read as chat, as before."""
|
||||
Responses ``input`` or ``instructions`` is read as chat, as before."""
|
||||
fields: Final = as_str_mapping(data) or {}
|
||||
return (
|
||||
call_type in _GEMINI_BODY_CALL_TYPES
|
||||
and not (call_type == CallTypes.pass_through.value and (fields.get("messages") or fields.get("input")))
|
||||
and not (
|
||||
call_type == CallTypes.pass_through.value
|
||||
and (fields.get("messages") or fields.get("input") or fields.get("instructions"))
|
||||
)
|
||||
and (
|
||||
isinstance(fields.get("contents"), list | dict)
|
||||
or isinstance(fields.get("systemInstruction") or fields.get("system_instruction"), dict)
|
||||
|
|
@ -291,6 +294,25 @@ def _has_gemini_contents(data: object, call_type: str | None) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _opted_out_of_global_guardrail(guardrail: "LakeraAIGuardrail", user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
"""True if the key's or team's admin metadata opts out of this ``default_on`` guardrail. A pass-through body has no
|
||||
admin metadata for ``should_run_guardrail`` to read, so it is read from the key here."""
|
||||
admin_fields: Final = user_api_key_dict.model_dump(include={"metadata", "team_metadata"})
|
||||
return (
|
||||
guardrail.default_on is True
|
||||
and guardrail.should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # takes an untyped data dict
|
||||
data={
|
||||
"metadata": {
|
||||
"user_api_key_metadata": admin_fields.get("metadata"),
|
||||
"user_api_key_team_metadata": admin_fields.get("team_metadata"),
|
||||
}
|
||||
},
|
||||
event_type=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
is not True
|
||||
)
|
||||
|
||||
|
||||
def _gemini_inspection_messages(data: Mapping[str, object], call_type: str | None) -> Sequence[Mapping[str, str]]:
|
||||
"""Chat-style messages for a Gemini ``generateContent`` body on a pass-through or Google ``generateContent``
|
||||
call: ``systemInstruction`` first as a system message, then each ``contents`` turn with its text parts joined and
|
||||
|
|
@ -647,7 +669,10 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
verbose_proxy_logger.debug("Lakera AI: pre_call_hook")
|
||||
|
||||
event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.pre_call
|
||||
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
|
||||
is_gemini_body: Final = _has_gemini_contents(data, call_type) # pyright: ignore[reportUnknownArgumentType] # hook data is an untyped dict
|
||||
if self.should_run_guardrail(data=data, event_type=event_type) is not True or (
|
||||
is_gemini_body and _opted_out_of_global_guardrail(self, user_api_key_dict)
|
||||
):
|
||||
verbose_proxy_logger.debug("Lakera AI: not running guardrail. Guardrail is disabled.")
|
||||
return data
|
||||
|
||||
|
|
@ -679,7 +704,6 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
# handled safely by _apply_redacted_messages_back_preserving_fields's
|
||||
# scope-index merge, which never touches a message outside the scope
|
||||
# it actually redacted instead of reconstructing the list from scratch.
|
||||
is_gemini_body: Final = _has_gemini_contents(data, call_type) # pyright: ignore[reportUnknownArgumentType] # hook data is an untyped dict
|
||||
is_multimodal_input: Final = (
|
||||
has_non_string_content(data)
|
||||
or _has_combined_messages_and_input(data)
|
||||
|
|
@ -891,6 +915,9 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
|
||||
# Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"]
|
||||
post_call_messages: Final = list(copy.deepcopy(original_messages)) + response_messages
|
||||
if not post_call_messages:
|
||||
verbose_proxy_logger.debug("Lakera AI: not running guardrail. No inspectable text in data")
|
||||
return response
|
||||
|
||||
# Call Lakera guardrail
|
||||
lakera_guardrail_response, _ = await self.call_v2_guard(
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -63,6 +64,11 @@ class _Action(str, enum.Enum):
|
|||
GUARDRAIL_INTERVENED = "GUARDRAIL_INTERVENED"
|
||||
|
||||
|
||||
def _user_api_key_alias(request_data: Mapping[str, object], metadata_key: str) -> object:
|
||||
"""The key alias recorded under ``metadata_key``, which a pass-through client can set to anything."""
|
||||
return (as_str_mapping(request_data.get(metadata_key)) or {}).get("user_api_key_alias")
|
||||
|
||||
|
||||
class NomaV2Guardrail(CustomGuardrail):
|
||||
reads_stream_marks: ClassVar[bool] = True
|
||||
"""Noma keeps one record per streamed reply when every scan carries ``stream`` and only the last one carries the
|
||||
|
|
@ -336,8 +342,8 @@ class NomaV2Guardrail(CustomGuardrail):
|
|||
# (ports v1 fallback from PR #16832).
|
||||
if application_id is None:
|
||||
application_id = self._get_non_empty_str(
|
||||
request_data.get("litellm_metadata", {}).get("user_api_key_alias")
|
||||
) or self._get_non_empty_str(request_data.get("metadata", {}).get("user_api_key_alias"))
|
||||
_user_api_key_alias(request_data, "litellm_metadata")
|
||||
) or self._get_non_empty_str(_user_api_key_alias(request_data, "metadata"))
|
||||
|
||||
try:
|
||||
payload: Final = self._build_scan_payload(
|
||||
|
|
|
|||
|
|
@ -4377,11 +4377,12 @@ async def add_guardrails_from_policy_engine(
|
|||
all_tags: Final = _policy_match_tags(data, metadata_variable_name)
|
||||
_team_alias: Final = user_api_key_dict.team_alias
|
||||
_key_alias: Final = user_api_key_dict.key_alias
|
||||
_model: Final = data.get("model")
|
||||
context: Final = PolicyMatchContext.model_validate(
|
||||
{
|
||||
"team_alias": _team_alias if isinstance(_team_alias, str) else None,
|
||||
"key_alias": _key_alias if isinstance(_key_alias, str) else None,
|
||||
"model": data.get("model"),
|
||||
"model": _model if isinstance(_model, str) else None,
|
||||
"tags": all_tags,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,121 @@
|
|||
"""A key or team opted out of a global Lakera guardrail keeps skipping it on a Gemini pass-through body."""
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Iterator
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent"
|
||||
CARD: Final = "4111 1111 1111 1111"
|
||||
GLOBAL_LAKERA: Final = "lakera-global"
|
||||
_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
_JSON: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": f"My card is {CARD}"}]}]}
|
||||
|
||||
|
||||
def _lakera_verdict(request: Request) -> Reply:
|
||||
flagged: Final = CARD.encode() in request.body
|
||||
breakdown: Final = [{"detector_type": "pii/credit_card", "detected": flagged, "message_id": 0}]
|
||||
return Reply(body=json.dumps({"flagged": flagged, "breakdown": breakdown, "payload": []}).encode())
|
||||
|
||||
|
||||
def _gemini_reply(_request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"candidates": [{"content": {"role": "model", "parts": [{"text": "Hi"}]}}]}).encode())
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GlobalLakeraRig:
|
||||
owned: OwnedProxy
|
||||
lakera: Wire
|
||||
upstream: Wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[GlobalLakeraRig]:
|
||||
directory: Final = tmp_path_factory.mktemp("lakera-gemini-global-opt-out")
|
||||
with ExitStack() as stack:
|
||||
lakera: Final = stack.enter_context(wire_server(_lakera_verdict))
|
||||
upstream: Final = stack.enter_context(wire_server(_gemini_reply))
|
||||
config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text()))
|
||||
path: Final = directory / "lakera-gemini-global-opt-out.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**config,
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": GLOBAL_LAKERA,
|
||||
"litellm_params": {
|
||||
"guardrail": "lakera_v2",
|
||||
"mode": "pre_call",
|
||||
"default_on": True,
|
||||
"on_flagged": "block",
|
||||
"api_base": lakera.url,
|
||||
"api_key": "synthetic-lakera-key",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
environment: Final = {"GEMINI_API_BASE": upstream.url, "GEMINI_API_KEY": "synthetic-gemini-key"}
|
||||
owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path))
|
||||
yield GlobalLakeraRig(owned, lakera, upstream)
|
||||
|
||||
|
||||
def _key_disabling_global(scenario: Scenario) -> str:
|
||||
return scenario.key(disable_global_guardrails=True)
|
||||
|
||||
|
||||
def _key_naming_it(scenario: Scenario) -> str:
|
||||
return scenario.key(metadata={"opted_out_global_guardrails": [GLOBAL_LAKERA]})
|
||||
|
||||
|
||||
def _team_disabling_global(scenario: Scenario) -> str:
|
||||
return scenario.key(team_id=scenario.team(disable_global_guardrails=True))
|
||||
|
||||
|
||||
def _team_naming_it(scenario: Scenario) -> str:
|
||||
return scenario.key(team_id=scenario.team(metadata={"opted_out_global_guardrails": [GLOBAL_LAKERA]}))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"opted_out_key", [_key_disabling_global, _key_naming_it, _team_disabling_global, _team_naming_it]
|
||||
)
|
||||
def test_opted_out_key_skips_the_global_guardrail_on_a_gemini_body(
|
||||
rig: GlobalLakeraRig, opted_out_key: Callable[[Scenario], str]
|
||||
) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = opted_out_key(scenario)
|
||||
rig.lakera.drain()
|
||||
rig.upstream.drain()
|
||||
|
||||
response: Final = gateway.request("POST", GEMINI_ROUTE, _BODY, key=key, headers={"x-goog-api-key": key})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert rig.lakera.drain() == ()
|
||||
assert len(rig.upstream.drain()) == 1
|
||||
|
||||
|
||||
def test_key_that_did_not_opt_out_is_inspected_by_the_global_guardrail(rig: GlobalLakeraRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
rig.lakera.drain()
|
||||
rig.upstream.drain()
|
||||
|
||||
response: Final = gateway.request("POST", GEMINI_ROUTE, _BODY, key=key, headers={"x-goog-api-key": key})
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert len(rig.lakera.drain()) == 1
|
||||
assert rig.upstream.drain() == ()
|
||||
|
|
@ -0,0 +1,83 @@
|
|||
"""Guardrails on built-in pass-through routes read client JSON that request setup hasn't normalized."""
|
||||
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
|
||||
GenericGuardrailAPI,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation
|
||||
from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail
|
||||
|
||||
NOT_AN_OBJECT: Final = (None, 0, False, "", [], "text", 7)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", NOT_AN_OBJECT)
|
||||
def test_generic_guardrail_reads_metadata_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None:
|
||||
guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic")
|
||||
|
||||
metadata: Final = guardrail._extract_user_api_key_metadata(
|
||||
{"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value}
|
||||
)
|
||||
|
||||
assert metadata["user_api_key_hash"] == "hash"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", NOT_AN_OBJECT)
|
||||
async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None:
|
||||
guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma")
|
||||
verdict: Final = httpx.Response(
|
||||
200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid")
|
||||
)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)):
|
||||
result: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["hello"]},
|
||||
request_data={"metadata": {}, "litellm_metadata": value},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert result["texts"] == ["hello"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]])
|
||||
async def test_legacy_lakera_has_no_text_to_scan_in_token_array_embeddings(token_arrays: list[object]) -> None:
|
||||
guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid")
|
||||
vendor_error: Final = httpx.HTTPStatusError(
|
||||
"vendor called",
|
||||
request=httpx.Request("POST", "http://vendor.invalid"),
|
||||
response=httpx.Response(500, text="vendor called"),
|
||||
)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)):
|
||||
await guardrail._check(
|
||||
data={"input": token_arrays},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="key"),
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_lakera_still_scans_string_inputs() -> None:
|
||||
guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid")
|
||||
vendor_error: Final = httpx.HTTPStatusError(
|
||||
"vendor called",
|
||||
request=httpx.Request("POST", "http://vendor.invalid"),
|
||||
response=httpx.Response(500, text="vendor called"),
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)),
|
||||
pytest.raises(Exception, match="vendor called"),
|
||||
):
|
||||
await guardrail._check(
|
||||
data={"input": ["hello"]},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="key"),
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
|
|
@ -1775,6 +1775,12 @@ def _without_metadata(data: Mapping[str, object]) -> Mapping[str, object]:
|
|||
return {key: value for key, value in data.items() if key != "metadata"}
|
||||
|
||||
|
||||
GEMINI_FLAGGED: Final = {
|
||||
"flagged": True,
|
||||
"breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 0}],
|
||||
}
|
||||
|
||||
|
||||
class TestGeminiContents:
|
||||
def test_reads_system_instruction_and_text_parts(self):
|
||||
assert (
|
||||
|
|
@ -1855,6 +1861,85 @@ class TestGeminiContents:
|
|||
{"role": "user", "content": "card"}
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"gemini_fields",
|
||||
[
|
||||
pytest.param({"contents": []}, id="contents-empty"),
|
||||
pytest.param({"contents": [{"parts": [{"text": "hello"}]}]}, id="contents-benign"),
|
||||
pytest.param({"systemInstruction": {"parts": [{"text": "hello"}]}}, id="system-instruction"),
|
||||
],
|
||||
)
|
||||
async def test_monitor_masks_instructions_on_a_pass_through_body_with_gemini_fields(
|
||||
self, gemini_fields: dict[str, object]
|
||||
):
|
||||
guardrail: Final = LakeraAIGuardrail(api_key="test_key", on_flagged="monitor")
|
||||
data: Final[dict[str, object]] = { # mutable-ok: the hook masks the body in place
|
||||
"instructions": "a@b.com is the contact",
|
||||
"input": "",
|
||||
**gemini_fields,
|
||||
}
|
||||
flagged: Final = {
|
||||
"flagged": True,
|
||||
"breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}],
|
||||
"payload": [{"detector_type": "pii/email", "start": 0, "end": 7, "message_id": 0}],
|
||||
}
|
||||
with patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(flagged, {}))):
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
|
||||
cache=DualCache(),
|
||||
data=data,
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
assert data["instructions"] == "[MASKED EMAIL] is the contact"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"key_fields",
|
||||
[
|
||||
pytest.param({"metadata": {"disable_global_guardrails": True}}, id="key-disables-global"),
|
||||
pytest.param({"metadata": {"opted_out_global_guardrails": ["lakera-global"]}}, id="key-names-it"),
|
||||
pytest.param({"team_metadata": {"disable_global_guardrails": True}}, id="team-disables-global"),
|
||||
pytest.param({"team_metadata": {"opted_out_global_guardrails": ["lakera-global"]}}, id="team-names-it"),
|
||||
],
|
||||
)
|
||||
async def test_key_opted_out_of_a_global_guardrail_skips_it_on_a_gemini_pass_through_body(
|
||||
self, key_fields: dict[str, dict[str, object]]
|
||||
):
|
||||
guardrail: Final = LakeraAIGuardrail(api_key="test_key", guardrail_name="lakera-global", default_on=True)
|
||||
data: Final = _gemini_request()
|
||||
with patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(GEMINI_FLAGGED, {}))):
|
||||
result: Final = await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", **key_fields),
|
||||
cache=DualCache(),
|
||||
data=data,
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
assert result == data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"key_fields",
|
||||
[
|
||||
pytest.param({}, id="no-opt-out"),
|
||||
pytest.param({"metadata": {"opted_out_global_guardrails": ["another-guardrail"]}}, id="another-named"),
|
||||
],
|
||||
)
|
||||
async def test_key_not_opted_out_of_a_global_guardrail_is_blocked_on_a_gemini_pass_through_body(
|
||||
self, key_fields: dict[str, dict[str, object]]
|
||||
):
|
||||
guardrail: Final = LakeraAIGuardrail(api_key="test_key", guardrail_name="lakera-global", default_on=True)
|
||||
with (
|
||||
patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(GEMINI_FLAGGED, {}))),
|
||||
pytest.raises(HTTPException),
|
||||
):
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", **key_fields),
|
||||
cache=DualCache(),
|
||||
data=_gemini_request(),
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("chat_field", ["messages", "input"])
|
||||
def test_a_pass_through_chat_body_with_stray_contents_is_read_as_chat(self, chat_field: str):
|
||||
data: Final = {chat_field: [USER_MSG], "contents": [{"parts": [{"text": "stray"}]}]}
|
||||
|
|
@ -2658,6 +2743,35 @@ async def test_lakera_post_call_allows_clean_content():
|
|||
|
||||
assert result is llm_response
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_with_no_messages_and_no_reply_text_returns_the_response():
|
||||
guardrail: Final = LakeraAIGuardrail(api_key="test_key")
|
||||
response: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}}]}
|
||||
rejected: Final = HTTPException(status_code=400, detail="either content, messages, or tools must be provided")
|
||||
with patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)):
|
||||
result: Final = await guardrail.async_post_call_success_hook(
|
||||
data={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
|
||||
response=response,
|
||||
)
|
||||
assert result == response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_with_messages_still_reaches_lakera():
|
||||
guardrail: Final = LakeraAIGuardrail(api_key="test_key")
|
||||
rejected: Final = HTTPException(status_code=400, detail="rejected")
|
||||
with (
|
||||
patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)),
|
||||
pytest.raises(HTTPException),
|
||||
):
|
||||
await guardrail.async_post_call_success_hook(
|
||||
data={"messages": [{"role": "user", "content": "hello"}]},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
|
||||
response={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
|
||||
@pytest.mark.asyncio
|
||||
async def test_lakera_post_call_masks_pii_and_allows():
|
||||
|
|
|
|||
|
|
@ -5327,6 +5327,33 @@ async def test_embedding_header_forwarding_without_model_group_config():
|
|||
litellm.model_group_settings = original_model_group_settings
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", [0, False, [], {}, 7, True])
|
||||
async def test_policy_engine_leaves_a_model_that_is_not_a_string_alone(
|
||||
model: object, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails
|
||||
|
||||
policy_registry = get_policy_registry()
|
||||
monkeypatch.setattr(
|
||||
policy_registry, "_policies", {"baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"]))}
|
||||
)
|
||||
monkeypatch.setattr(policy_registry, "_initialized", True)
|
||||
attachment_registry = get_attachment_registry()
|
||||
monkeypatch.setattr(attachment_registry, "_attachments", [PolicyAttachment(policy="baseline", scope="*")])
|
||||
monkeypatch.setattr(attachment_registry, "_initialized", True)
|
||||
data = {"model": model, "metadata": {}}
|
||||
|
||||
await add_guardrails_from_policy_engine(
|
||||
data=data, metadata_variable_name="metadata", user_api_key_dict=UserAPIKeyAuth(api_key="test-key")
|
||||
)
|
||||
|
||||
assert data["model"] == model
|
||||
assert data["metadata"]["guardrails"] == ["pii_blocker"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_guardrails_from_policy_engine():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue