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
# Conflicts: # litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
This commit is contained in:
yucheng
commit
b654684c8c
10 changed files with 324 additions and 28 deletions
|
|
@ -3,6 +3,11 @@ from typing import List, Optional, Union
|
|||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
|
||||
|
||||
def _names_a_mode_tag(tag: object, mode: Mode) -> bool:
|
||||
"""Whether a request tag, which a client's body can make anything, is one of ``mode``'s tags"""
|
||||
return isinstance(tag, str) and tag in mode.tags
|
||||
|
||||
|
||||
class EnterpriseCustomGuardrailHelper:
|
||||
@staticmethod
|
||||
def _should_run_if_mode_by_tag(
|
||||
|
|
@ -55,7 +60,7 @@ class EnterpriseCustomGuardrailHelper:
|
|||
matched_mode = None
|
||||
if request_tags:
|
||||
for tag in request_tags:
|
||||
if tag in event_hook.tags:
|
||||
if _names_a_mode_tag(tag, event_hook):
|
||||
matched_mode = event_hook.tags[tag]
|
||||
break
|
||||
|
||||
|
|
|
|||
|
|
@ -243,7 +243,8 @@ class PassThroughEndpointHandler(BaseTranslation):
|
|||
guardrail_to_apply: The guardrail instance to apply
|
||||
litellm_logging_obj: Optional logging object
|
||||
user_api_key_dict: User API key metadata to pass to guardrails
|
||||
read_whole_response: Without targeted fields, read the response as JSON, ids and usage included
|
||||
read_whole_response: Without targeted fields, read the response as JSON, ids and usage included, as a call
|
||||
with no logging object does, such as a logging_only scan
|
||||
"""
|
||||
response_fields: Final = as_str_mapping(response)
|
||||
if not isinstance(response, dict):
|
||||
|
|
@ -263,7 +264,7 @@ class PassThroughEndpointHandler(BaseTranslation):
|
|||
# Extract text to check
|
||||
text_to_check: Final = (
|
||||
self._extract_text_for_guardrail(response, field_expressions)
|
||||
if field_expressions or read_whole_response
|
||||
if field_expressions or read_whole_response or litellm_logging_obj is None
|
||||
else _response_text(response_fields or {})
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -378,12 +378,13 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
# Fields in metadata are already prefixed with 'user_api_key_'
|
||||
for field_name in GenericGuardrailAPIMetadata.__annotations__:
|
||||
value = metadata_dict.get(field_name)
|
||||
if value is not None:
|
||||
if isinstance(value, str):
|
||||
result_metadata[field_name] = value
|
||||
|
||||
# handle user_api_key_token = user_api_key_hash
|
||||
if metadata_dict.get("user_api_key_token") is not None:
|
||||
result_metadata["user_api_key_hash"] = metadata_dict.get("user_api_key_token")
|
||||
token: Final = metadata_dict.get("user_api_key_token")
|
||||
if isinstance(token, str):
|
||||
result_metadata["user_api_key_hash"] = token
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Generic Guardrail API: Extracted user metadata: %s",
|
||||
|
|
@ -392,6 +393,14 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
|
||||
return result_metadata
|
||||
|
||||
def _dynamic_request_body_params(self, request_body: object) -> Mapping[str, object]:
|
||||
"""The ``extra_body`` a request's ``body`` asks this guardrail to send. A custom endpoint's client can send a
|
||||
``body`` field of any shape, which asks for nothing unless it names this guardrail's params in the usual form."""
|
||||
try:
|
||||
return as_str_mapping(self.get_guardrail_dynamic_request_body_params(request_body)) or {} # pyright: ignore[reportArgumentType, reportUnknownArgumentType] # a client's JSON, validated by the call
|
||||
except (TypeError, ValidationError):
|
||||
return {}
|
||||
|
||||
def _fail_open_passthrough(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -561,16 +570,16 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
tools: Final = inputs.get("tools")
|
||||
structured_messages: Final = inputs.get("structured_messages")
|
||||
tool_calls: Final = inputs.get("tool_calls")
|
||||
model: Final = inputs.get("model")
|
||||
model: Final[object] = inputs.get("model")
|
||||
|
||||
# Use provided request_data or create an empty dict
|
||||
if request_data is None:
|
||||
request_data = {}
|
||||
|
||||
request_body: Final = request_data.get("body") or {}
|
||||
request_body: Final[object] = request_data.get("body")
|
||||
|
||||
# Merge additional provider specific params from config and dynamic params
|
||||
dynamic_params: Final[Mapping[str, object]] = self.get_guardrail_dynamic_request_body_params(request_body) or {}
|
||||
dynamic_params: Final = self._dynamic_request_body_params(request_body)
|
||||
additional_params: Final = {**self.additional_provider_specific_params, **dynamic_params}
|
||||
|
||||
# Extract user API key metadata
|
||||
|
|
@ -597,7 +606,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
tool_calls=tool_calls,
|
||||
additional_provider_specific_params=additional_params,
|
||||
input_type=input_type,
|
||||
model=model,
|
||||
model=model if isinstance(model, str) else None,
|
||||
)
|
||||
|
||||
headers: Final = self._build_request_headers()
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
import copy
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from itertools import chain, product
|
||||
from itertools import chain
|
||||
from string import Formatter
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
|
@ -261,14 +261,20 @@ def _gemini_text(content: object) -> str:
|
|||
return "\n".join(text for text in texts if isinstance(text, str) and text)
|
||||
|
||||
|
||||
def _gemini_tool_part_values(part: object) -> Iterator[object]:
|
||||
fields: Final = as_str_mapping(part) or {}
|
||||
return (
|
||||
(as_str_mapping(fields[name]) or {}).get(_GEMINI_TOOL_PART_FIELDS[name])
|
||||
for name in fields
|
||||
if name in _GEMINI_TOOL_PART_FIELDS
|
||||
)
|
||||
|
||||
|
||||
def _gemini_tool_text(content: object) -> str:
|
||||
"""What one Gemini ``Content`` carries in its function call, function response and code execution parts, with
|
||||
structured values as JSON."""
|
||||
parts: Final = _gemini_items((as_str_mapping(content) or {}).get("parts"))
|
||||
values: Final = (
|
||||
(as_str_mapping((as_str_mapping(part) or {}).get(part_name)) or {}).get(field)
|
||||
for part, (part_name, field) in product(parts, _GEMINI_TOOL_PART_FIELDS.items())
|
||||
)
|
||||
values: Final = chain.from_iterable(map(_gemini_tool_part_values, parts))
|
||||
return "\n".join(
|
||||
value if isinstance(value, str) else json.dumps(value, default=str, ensure_ascii=False)
|
||||
for value in values
|
||||
|
|
|
|||
|
|
@ -210,6 +210,11 @@ _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata"))
|
|||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object])
|
||||
_RESPONSE_HEADERS: Final = TypeAdapter(Mapping[str, str])
|
||||
_FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = (
|
||||
("guardrails", (list, dict)),
|
||||
("guardrail_config", (dict,)),
|
||||
("proxy_server_request", (dict,)),
|
||||
)
|
||||
_POST_CALL_SELECTION_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails", "tags")
|
||||
_KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState))
|
||||
_PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",))
|
||||
|
|
@ -1069,12 +1074,11 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
def _drop_unmergeable_guardrail_fields(
|
||||
parsed_body: dict[str, object], # mutable-ok: strips the request body in place before the guardrail merge
|
||||
) -> None:
|
||||
"""Drop a ``guardrails`` that isn't a list or an object and a ``guardrail_config`` that isn't an object, which the
|
||||
merge can't combine with the metadata's own. The provider receives neither field."""
|
||||
if "guardrails" in parsed_body and not isinstance(parsed_body["guardrails"], list | dict):
|
||||
parsed_body.pop("guardrails")
|
||||
if "guardrail_config" in parsed_body and not isinstance(parsed_body["guardrail_config"], dict):
|
||||
parsed_body.pop("guardrail_config")
|
||||
"""Drop a ``guardrails`` that isn't a list or an object, and a ``guardrail_config`` or ``proxy_server_request`` that
|
||||
isn't an object, which the merge and the guardrails can't read. The provider receives none of these fields."""
|
||||
for field, shape in _FIELDS_GUARDRAILS_READ_BY_SHAPE:
|
||||
if field in parsed_body and not isinstance(parsed_body[field], shape):
|
||||
parsed_body.pop(field)
|
||||
|
||||
|
||||
async def _inherit_auth_guardrails(
|
||||
|
|
@ -1093,6 +1097,7 @@ async def _inherit_auth_guardrails(
|
|||
if configured_guardrails or not inherit or not parsed_body:
|
||||
return configured_guardrails
|
||||
_drop_unmergeable_guardrail_fields(parsed_body)
|
||||
client_sent_metadata: Final = "metadata" in parsed_body
|
||||
if not isinstance(parsed_body.get("metadata"), dict):
|
||||
parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata
|
||||
client_named: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails")
|
||||
|
|
@ -1102,6 +1107,9 @@ async def _inherit_auth_guardrails(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
metadata: Final = parsed_body["metadata"]
|
||||
if not metadata and not client_sent_metadata:
|
||||
parsed_body.pop("metadata")
|
||||
return None
|
||||
merged: Final = (as_str_mapping(metadata) or {}).get("guardrails")
|
||||
if not isinstance(merged, list):
|
||||
return None
|
||||
|
|
@ -1135,6 +1143,12 @@ def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str,
|
|||
)
|
||||
|
||||
|
||||
def _holds_resolved_pipelines(pipelines: object) -> bool:
|
||||
"""Whether a pipeline slot holds the policy engine's (name, pipeline) tuples, which a JSON body can't"""
|
||||
entries: Final = as_object_list(pipelines)
|
||||
return entries is not None and all(isinstance(entry, tuple) for entry in entries)
|
||||
|
||||
|
||||
def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]:
|
||||
"""The policy pipelines resolved for a request that inherits auth guardrails, and the tags a tag-scoped guardrail
|
||||
mode reads, from the logging metadata that request setup moved the body's ``metadata`` into"""
|
||||
|
|
@ -1142,7 +1156,10 @@ def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool)
|
|||
return MappingProxyType({})
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
|
||||
metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {}
|
||||
return {field: metadata[field] for field in _POST_CALL_SELECTION_FIELDS if field in metadata}
|
||||
selection: Final = {field: metadata[field] for field in _POST_CALL_SELECTION_FIELDS if field in metadata}
|
||||
if "_guardrail_pipelines" in selection and not _holds_resolved_pipelines(selection["_guardrail_pipelines"]):
|
||||
selection.pop("_guardrail_pipelines")
|
||||
return selection
|
||||
|
||||
|
||||
def _starts_post_call_pass(guardrails: Sequence[object] | None, policy_pipelines: Mapping[str, object]) -> bool:
|
||||
|
|
@ -1160,10 +1177,11 @@ def _post_call_hook_data(
|
|||
"""The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails,
|
||||
which ``should_run_guardrail`` checks, the policy pipelines to run, and the logging object, which unified
|
||||
(apply_guardrail) guardrails resolve a pass-through response's call type from when the route has no
|
||||
API_ROUTE_TO_CALL_TYPES entry"""
|
||||
API_ROUTE_TO_CALL_TYPES entry. A custom endpoint whose guardrails don't run post_call gets no logging object, so a
|
||||
global unified guardrail still can't resolve its response."""
|
||||
body: Final = parsed_body or {}
|
||||
admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) if inherited else MappingProxyType({})
|
||||
return {
|
||||
data: Final = {
|
||||
**body,
|
||||
"metadata": {
|
||||
**(as_str_mapping(body.get("metadata")) or {}),
|
||||
|
|
@ -1171,8 +1189,10 @@ def _post_call_hook_data(
|
|||
"guardrails": guardrails or (),
|
||||
**admin_metadata,
|
||||
},
|
||||
"litellm_logging_obj": logging_obj,
|
||||
}
|
||||
if not inherited and not _named_guardrail_runs_post_call(data, guardrails or ()):
|
||||
return data
|
||||
return {**data, "litellm_logging_obj": logging_obj}
|
||||
|
||||
|
||||
def _provider_charge_settlement(
|
||||
|
|
|
|||
|
|
@ -47,6 +47,8 @@ ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages"
|
|||
GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent"
|
||||
CUSTOM_ROUTE: Final = "/passthrough-auth-custom"
|
||||
CUSTOM_GUARDED_ROUTE: Final = "/passthrough-auth-custom-guarded"
|
||||
CUSTOM_PRE_CALL_ROUTE: Final = "/passthrough-auth-custom-pre-call"
|
||||
GLOBAL_UNIFIED_POST_CALL: Final = "passthrough-global-unified-post-call"
|
||||
UNPRICED_MODEL: Final = "passthrough-auth-unpriced-model"
|
||||
MONITOR_POST_CALL: Final = "passthrough-auth-monitor-post-call"
|
||||
MONITORED_PROMPT: Final = "reply-with-the-monitored-id"
|
||||
|
|
@ -175,6 +177,7 @@ class AuthGuardrailRig:
|
|||
post_call_vendor: Wire
|
||||
key_post_call_vendor: Wire
|
||||
generic_post_call_vendor: Wire
|
||||
global_unified_vendor: Wire
|
||||
model_post_call_vendor: Wire
|
||||
upstreams: Mapping[str, Wire]
|
||||
|
||||
|
|
@ -187,6 +190,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean))
|
||||
key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged))
|
||||
generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict))
|
||||
global_unified_vendor: Final = stack.enter_context(wire_server(_verdict))
|
||||
blocking_post_call_vendor: Final = stack.enter_context(wire_server(_block_every_reply))
|
||||
monitor_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_pii_only))
|
||||
model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict))
|
||||
|
|
@ -224,6 +228,12 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
"cost_per_request": 0.25,
|
||||
"guardrails": {BLOCKING_POST_CALL: None},
|
||||
},
|
||||
{
|
||||
"path": CUSTOM_PRE_CALL_ROUTE,
|
||||
"target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions",
|
||||
"auth": True,
|
||||
"guardrails": {GUARDRAIL: None},
|
||||
},
|
||||
],
|
||||
},
|
||||
"policies": {
|
||||
|
|
@ -317,6 +327,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": GLOBAL_UNIFIED_POST_CALL,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": "post_call",
|
||||
"default_on": True,
|
||||
"api_base": global_unified_vendor.url,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": GENERIC_POST_CALL,
|
||||
"litellm_params": {
|
||||
|
|
@ -367,6 +387,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
post_call_vendor,
|
||||
key_post_call_vendor,
|
||||
generic_post_call_vendor,
|
||||
global_unified_vendor,
|
||||
model_post_call_vendor,
|
||||
upstreams,
|
||||
)
|
||||
|
|
@ -866,6 +887,78 @@ def test_custom_endpoint_writes_no_charge_for_a_response_its_post_call_guardrail
|
|||
assert [float(str(row["spend"])) for row in _row_fields(gateway, key)] == [0], response.text
|
||||
|
||||
|
||||
def test_a_custom_endpoint_with_only_a_pre_call_guardrail_never_calls_a_global_unified_guardrail(
|
||||
rig: AuthGuardrailRig,
|
||||
) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
custom_key: Final = scenario.key(metadata={"allowed_passthrough_routes": [CUSTOM_PRE_CALL_ROUTE]})
|
||||
built_in_key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]})
|
||||
_drain(rig)
|
||||
rig.global_unified_vendor.drain()
|
||||
|
||||
custom_response: Final = gateway.request(
|
||||
"POST", CUSTOM_PRE_CALL_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=custom_key
|
||||
)
|
||||
|
||||
assert custom_response.status_code == 200, custom_response.text
|
||||
assert len(rig.vendor.drain()) == 1
|
||||
assert rig.global_unified_vendor.drain() == ()
|
||||
|
||||
built_in_response: Final = _send(gateway, OPENAI_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), built_in_key)
|
||||
|
||||
assert built_in_response.status_code == 200, built_in_response.text
|
||||
assert len(rig.global_unified_vendor.drain()) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field", ["Hello there", 7, ["a", "b"], {"guardrails": 7}], ids=["text", "int", "list", "dict"]
|
||||
)
|
||||
def test_a_custom_endpoints_own_body_field_still_gets_its_reply_scanned(
|
||||
rig: AuthGuardrailRig, field: JsonValue
|
||||
) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(metadata={"allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE]})
|
||||
|
||||
response: Final = gateway.request(
|
||||
"POST", CUSTOM_GUARDED_ROUTE, {**_route_body(OPENAI_ROUTE, "Hello"), "body": field}, key=key
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "synthetic block" in response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("guardrail", "fields"),
|
||||
[
|
||||
(GUARDRAIL, {"proxy_server_request": 5}),
|
||||
(GUARDRAIL, {"metadata": {"user_api_key_alias": 5}}),
|
||||
(TAG_SCOPED_POST_CALL, {"metadata": {"tags": [1, {"a": 1}]}}),
|
||||
(
|
||||
None,
|
||||
{
|
||||
"metadata": {"_pipeline_managed_guardrails": []},
|
||||
"litellm_metadata": {"_guardrail_pipelines": [["p", "q"]]},
|
||||
},
|
||||
),
|
||||
],
|
||||
ids=["proxy-server-request", "key-alias", "tags", "pipeline-slots"],
|
||||
)
|
||||
def test_a_malformed_proxy_field_in_a_body_still_gets_the_provider_reply(
|
||||
rig: AuthGuardrailRig, guardrail: str | None, fields: Mapping[str, JsonValue]
|
||||
) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(metadata={"guardrails": [guardrail]} if guardrail else {})
|
||||
_drain(rig)
|
||||
|
||||
response: Final = _send(gateway, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE, "Hello"), **fields}, key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1
|
||||
|
||||
|
||||
def test_a_monitored_lakera_flag_keeps_the_provider_id_on_the_row(rig: AuthGuardrailRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
|
|
|
|||
|
|
@ -330,3 +330,29 @@ async def test_baggage_never_changes_the_guardrail_mode_key_team_and_deployment_
|
|||
|
||||
assert await runs_post_call_guardrail(()) is runs_post_call
|
||||
assert await runs_post_call_guardrail(("usecase:dev",)) is runs_post_call
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"])
|
||||
def test_a_request_tag_that_cant_match_falls_back_to_the_default_mode(monkeypatch, bucket):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
cg = CustomGuardrail(
|
||||
guardrail_name="test_guardrail",
|
||||
supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
|
||||
event_hook=Mode(tags={"test_tag": "post_call"}, default="pre_call"),
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
assert (
|
||||
cg.should_run_guardrail(
|
||||
data={bucket: {"tags": [1, {"a": 1}, ["b"]]}},
|
||||
event_type=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
cg.should_run_guardrail(
|
||||
data={bucket: {"tags": [{"a": 1}, "test_tag"]}},
|
||||
event_type=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
|
|
|||
|
|
@ -231,10 +231,22 @@ def test_passthrough_route_block_before_any_invoke_frame_is_a_whole_message_in_c
|
|||
|
||||
async def _scanned_response_texts(response: dict[str, object]) -> tuple[str, ...]:
|
||||
guardrail: Final = MagicMock(guardrail_name="post-call", apply_guardrail=AsyncMock())
|
||||
await PassThroughEndpointHandler().process_output_response(response, guardrail, request_data={})
|
||||
await PassThroughEndpointHandler().process_output_response(
|
||||
response, guardrail, MagicMock(passthrough_guardrails_config=None), request_data={}
|
||||
)
|
||||
return tuple(text for call in guardrail.apply_guardrail.await_args_list for text in call.kwargs["inputs"]["texts"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_scan_with_no_logging_object_reads_the_response_as_json():
|
||||
response: Final = {"response": {"id": "chatcmpl-1", "object": "chat.completion", "choices": [{"text": "Paris."}]}}
|
||||
guardrail: Final = MagicMock(guardrail_name="logging-only", apply_guardrail=AsyncMock())
|
||||
|
||||
await PassThroughEndpointHandler().process_output_response(response, guardrail, request_data={})
|
||||
|
||||
assert json.loads(guardrail.apply_guardrail.await_args.kwargs["inputs"]["texts"][0]) == response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_scan_reads_the_reply_text_and_not_the_numbers_around_it():
|
||||
response: Final = {
|
||||
|
|
|
|||
|
|
@ -2703,3 +2703,52 @@ def test_extract_user_api_key_metadata_reads_past_a_litellm_metadata_that_is_not
|
|||
)
|
||||
|
||||
assert metadata["user_api_key_hash"] == "hash"
|
||||
|
||||
|
||||
def test_extract_user_api_key_metadata_leaves_out_a_value_that_isnt_a_string() -> None:
|
||||
guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic")
|
||||
|
||||
metadata: Final = guardrail._extract_user_api_key_metadata(
|
||||
{"metadata": {"user_api_key_alias": 5, "user_api_key_token": 7, "user_api_key_team_id": "team"}}
|
||||
)
|
||||
|
||||
assert metadata == {"user_api_key_team_id": "team"}
|
||||
|
||||
|
||||
async def _vendor_payload(request_data: dict[str, object], inputs: dict[str, object]) -> dict[str, object]:
|
||||
guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic")
|
||||
vendor_reply: Final = MagicMock(json=MagicMock(return_value={"action": "NONE"}))
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", return_value=vendor_reply) as post:
|
||||
await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="response")
|
||||
|
||||
return post.call_args.kwargs["json"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
(
|
||||
"Hello there",
|
||||
7,
|
||||
True,
|
||||
["a", "b"],
|
||||
{"guardrails": 7},
|
||||
{"metadata": {"guardrails": 7}},
|
||||
{"guardrails": [{"generic": 7}]},
|
||||
),
|
||||
)
|
||||
async def test_a_body_field_a_custom_endpoints_client_sent_asks_the_vendor_for_nothing(body: object) -> None:
|
||||
payload: Final = await _vendor_payload({"body": body}, {"texts": ["reply"]})
|
||||
|
||||
assert payload["texts"] == ["reply"]
|
||||
assert not payload.get("additional_provider_specific_params")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", ({"id": 1}, ["gpt"], 7))
|
||||
async def test_a_reply_model_that_isnt_a_string_is_left_out_of_the_vendor_request(model: object) -> None:
|
||||
payload: Final = await _vendor_payload({}, {"texts": ["reply"], "model": model})
|
||||
|
||||
assert payload["texts"] == ["reply"]
|
||||
assert payload.get("model") is None
|
||||
|
|
|
|||
|
|
@ -59,6 +59,9 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
SafeRouteAdder,
|
||||
_anthropic_user_id_to_forward,
|
||||
_inherit_auth_guardrails,
|
||||
_post_call_hook_data,
|
||||
_post_call_selection_metadata,
|
||||
_starts_post_call_pass,
|
||||
WebsocketFrameModelAllowed,
|
||||
WebsocketFrameModelGate,
|
||||
WebsocketFrameModelRefused,
|
||||
|
|
@ -13416,11 +13419,83 @@ class TestInheritAuthGuardrails:
|
|||
await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test", metadata=key_metadata), None, True)
|
||||
|
||||
assert "guardrails" not in body
|
||||
assert (body["metadata"].get("guardrails") or []) == [
|
||||
assert (body.get("metadata", {}).get("guardrails") or []) == [
|
||||
*key_metadata.get("guardrails", ()),
|
||||
*(body_metadata or {}).get("guardrails", ()),
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_body_with_nothing_to_merge_gets_no_metadata(self):
|
||||
body: Final[dict[str, object]] = self._chat_body()
|
||||
|
||||
assert await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) is None
|
||||
|
||||
assert body == self._chat_body()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_body_with_nothing_to_merge_keeps_the_metadata_it_sent(self):
|
||||
body: Final[dict[str, object]] = self._chat_body(metadata={})
|
||||
|
||||
await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True)
|
||||
|
||||
assert body["metadata"] == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", [5, "text", ["list"], None])
|
||||
async def test_a_proxy_server_request_that_isnt_an_object_never_reaches_the_guardrails(self, value):
|
||||
body: Final[dict[str, object]] = self._chat_body(proxy_server_request=value)
|
||||
|
||||
await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True)
|
||||
|
||||
assert "proxy_server_request" not in body
|
||||
|
||||
@pytest.mark.parametrize("slot", [[["p", "q"]], [7], 7, "text", [{"mode": "post_call"}]])
|
||||
def test_a_pipeline_slot_a_body_sent_does_not_start_the_post_call_pass(self, slot):
|
||||
kwargs: Final = {
|
||||
"litellm_params": {
|
||||
"metadata": {"_guardrail_pipelines": slot, "_pipeline_managed_guardrails": [], "tags": ["t"]}
|
||||
}
|
||||
}
|
||||
|
||||
selection: Final = _post_call_selection_metadata(kwargs, inherited=True)
|
||||
|
||||
assert selection == {"_pipeline_managed_guardrails": [], "tags": ["t"]}
|
||||
assert _starts_post_call_pass(None, selection) is False
|
||||
|
||||
def test_a_pipeline_the_policy_engine_resolved_starts_the_post_call_pass(self):
|
||||
pipeline: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="guard")])
|
||||
kwargs: Final = {"litellm_params": {"metadata": {"_guardrail_pipelines": [("p", pipeline)]}}}
|
||||
|
||||
selection: Final = _post_call_selection_metadata(kwargs, inherited=True)
|
||||
|
||||
assert selection == {"_guardrail_pipelines": [("p", pipeline)]}
|
||||
assert _starts_post_call_pass(None, selection) is True
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("event_hook", "inherited", "gets_the_logging_object"),
|
||||
[
|
||||
(GuardrailEventHooks.pre_call, False, False),
|
||||
(GuardrailEventHooks.post_call, False, True),
|
||||
(GuardrailEventHooks.pre_call, True, True),
|
||||
],
|
||||
ids=["custom-pre-call", "custom-post-call", "built-in-route"],
|
||||
)
|
||||
def test_a_custom_endpoint_gets_the_logging_object_only_for_a_post_call_guardrail(
|
||||
self, event_hook, inherited, gets_the_logging_object
|
||||
):
|
||||
logging_obj: Final = MagicMock()
|
||||
|
||||
with patch.object(
|
||||
litellm.logging_callback_manager,
|
||||
"get_custom_loggers_for_type",
|
||||
return_value=[self._guardrail("guard", event_hook)],
|
||||
):
|
||||
data: Final = _post_call_hook_data(
|
||||
self._chat_body(), ["guard"], logging_obj, UserAPIKeyAuth(api_key="sk-test"), {}, inherited
|
||||
)
|
||||
|
||||
assert data.get("litellm_logging_obj") is (logging_obj if gets_the_logging_object else None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", [None, False, 7, "config", ["config"]])
|
||||
async def test_a_guardrail_config_the_merge_cant_add_to_is_dropped(self, value):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue