mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(proxy)!: run key, team and project guardrails on built-in provider pass-through routes
The /openai, /azure, /anthropic, /vertex_ai and /gemini pass-through routes now merge the caller's key, team, project and policy guardrails with any named in the request body, the same way the native routes and /bedrock do. Requests without a JSON body are not scanned, and custom pass-through endpoints keep their opt-in guardrails config
This commit is contained in:
yucheng
parent
6ad4f9cfdf
commit
5011b1263f
5 changed files with 361 additions and 2 deletions
|
|
@ -414,6 +414,7 @@ async def gemini_proxy_route(
|
|||
custom_llm_provider="gemini",
|
||||
is_streaming_request=is_streaming_request,
|
||||
query_params=merged_params,
|
||||
inherit_auth_guardrails=True,
|
||||
) # dynamically construct pass-through endpoint based on incoming path
|
||||
received_value: Final = await endpoint_func(
|
||||
request,
|
||||
|
|
@ -1079,6 +1080,7 @@ async def anthropic_proxy_route(
|
|||
target=str(updated_url),
|
||||
custom_headers=_upstream_headers_for_anthropic_route(request, user_api_key_dict, auth_header),
|
||||
is_streaming_request=is_streaming_request,
|
||||
inherit_auth_guardrails=True,
|
||||
) # dynamically construct pass-through endpoint based on incoming path
|
||||
received_value: Final = await endpoint_func(
|
||||
request,
|
||||
|
|
@ -3192,6 +3194,7 @@ async def _base_vertex_proxy_route(
|
|||
upstream_headers_filter=partial(
|
||||
_without_dropped_anthropic_betas, _vertex_dropped_anthropic_betas(general_settings, request.headers)
|
||||
),
|
||||
inherit_auth_guardrails=True,
|
||||
) # dynamically construct pass-through endpoint based on incoming path
|
||||
|
||||
try:
|
||||
|
|
@ -3670,6 +3673,7 @@ class BaseOpenAIPassThroughHandler:
|
|||
if custom_llm_provider
|
||||
else None
|
||||
),
|
||||
inherit_auth_guardrails=True,
|
||||
) # dynamically construct pass-through endpoint based on incoming path
|
||||
received_value: Final = await endpoint_func(
|
||||
request,
|
||||
|
|
|
|||
|
|
@ -60,6 +60,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
as_str_mapping,
|
||||
bind_budget_reservation_to_callbacks,
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
get_or_create_metadata_bucket,
|
||||
|
|
@ -126,6 +127,7 @@ from litellm.proxy.litellm_pre_call_utils import ( # noqa: F401 # legacy modul
|
|||
_key_or_team_allows_client_pricing_override, # pyright: ignore[reportPrivateUsage] # reuse the proxy's pricing trust policy
|
||||
_strip_client_pricing_overrides, # pyright: ignore[reportPrivateUsage] # sanitize before trusted hooks add guardrail costs
|
||||
get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above
|
||||
move_guardrails_to_metadata, # pyright: ignore[reportUnknownVariableType] # untyped data param; the same merge the native routes run
|
||||
redact_credential_headers,
|
||||
strip_client_guardrail_information,
|
||||
)
|
||||
|
|
@ -173,6 +175,7 @@ pass_through_endpoint_logging: Final = PassThroughEndpointLogging()
|
|||
|
||||
_METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata"))
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object])
|
||||
_KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState))
|
||||
_PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",))
|
||||
|
||||
|
|
@ -854,6 +857,31 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
_update_stream_param_based_on_request_body = update_stream_param_based_on_request_body
|
||||
|
||||
|
||||
async def _inherit_auth_guardrails(
|
||||
parsed_body: dict[str, object] | None, # mutable-ok: merges the guardrails into the live request body
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
configured_guardrails: Mapping[str, bool] | None,
|
||||
inherit: bool,
|
||||
) -> Mapping[str, bool] | Sequence[object] | None:
|
||||
"""Return the guardrails a pass-through request runs.
|
||||
|
||||
The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team,
|
||||
project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` and returned, or None
|
||||
when there are none.
|
||||
"""
|
||||
if configured_guardrails or not inherit or not parsed_body:
|
||||
return configured_guardrails
|
||||
if not isinstance(parsed_body.get("metadata"), dict):
|
||||
parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata
|
||||
await move_guardrails_to_metadata(
|
||||
data=parsed_body,
|
||||
_metadata_variable_name="metadata",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails")
|
||||
return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) if isinstance(merged, list) and merged else None
|
||||
|
||||
|
||||
def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | None) -> None:
|
||||
"""Copy guardrail logging entries from ``guardrail_data`` onto ``request_data``.
|
||||
|
||||
|
|
@ -1212,6 +1240,7 @@ async def pass_through_request(
|
|||
timeout: float | None = None,
|
||||
upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers,
|
||||
operation: GenAIOperation | None = None,
|
||||
inherit_auth_guardrails: bool = False,
|
||||
):
|
||||
"""
|
||||
Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called
|
||||
|
|
@ -1234,6 +1263,8 @@ async def pass_through_request(
|
|||
general_settings.pass_through_request_timeout, then 600s.
|
||||
upstream_headers_filter: Optional last rewrite of the headers sent upstream, after x-pass- expansion
|
||||
operation: Optional GenAI operation the route reports, unless the upstream names its own
|
||||
inherit_auth_guardrails: Run the key, team, project and policy guardrails on a JSON body, as the
|
||||
native routes do, when the pass-through guardrails config resolves to none
|
||||
"""
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
|
@ -1342,6 +1373,13 @@ async def pass_through_request(
|
|||
_parsed_body["metadata"]["guardrails"] = guardrails_to_run
|
||||
verbose_proxy_logger.debug("Added guardrails to passthrough request metadata: %s", guardrails_to_run)
|
||||
|
||||
post_call_guardrails: Final = await _inherit_auth_guardrails(
|
||||
parsed_body=cast(dict[str, object] | None, _parsed_body), # cast-ok: json
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
configured_guardrails=guardrails_to_run,
|
||||
inherit=inherit_auth_guardrails,
|
||||
)
|
||||
|
||||
## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it
|
||||
# Surface the requested model (when the body carries one) so logging/spans
|
||||
# read e.g. ``chat gpt-4o`` instead of ``chat unknown``.
|
||||
|
|
@ -1876,7 +1914,7 @@ async def pass_through_request(
|
|||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
if response.status_code < 400 and response_body is not None and guardrails_to_run:
|
||||
if response.status_code < 400 and response_body is not None and post_call_guardrails:
|
||||
# Build an enriched data dict: _parsed_body has been stripped of
|
||||
# `metadata` by both pre_call_hook and _init_kwargs_for_pass_through_endpoint,
|
||||
# so we re-attach the configured guardrails here so should_run_guardrail
|
||||
|
|
@ -1887,7 +1925,7 @@ async def pass_through_request(
|
|||
existing_metadata = {}
|
||||
hook_data["metadata"] = {
|
||||
**existing_metadata,
|
||||
"guardrails": guardrails_to_run,
|
||||
"guardrails": post_call_guardrails,
|
||||
}
|
||||
post_call_guardrail_data = hook_data
|
||||
response_body = await proxy_logging_obj.post_call_success_hook(
|
||||
|
|
@ -2271,6 +2309,7 @@ def create_pass_through_route(
|
|||
timeout: float | None = None,
|
||||
upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers,
|
||||
operation: GenAIOperation | None = None,
|
||||
inherit_auth_guardrails: bool = False,
|
||||
):
|
||||
# check if target is an adapter.py or a url
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -2420,6 +2459,7 @@ def create_pass_through_route(
|
|||
timeout=cast(float | None, param_timeout),
|
||||
upstream_headers_filter=upstream_headers_filter,
|
||||
operation=operation,
|
||||
inherit_auth_guardrails=inherit_auth_guardrails,
|
||||
)
|
||||
finally:
|
||||
if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,193 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.anthropic_sse import message_json
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.openai_wire import chat_reply
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
GUARDRAIL: Final = "passthrough-auth-guardrail"
|
||||
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])
|
||||
|
||||
Level = Literal["key", "team", "project"]
|
||||
|
||||
|
||||
def _route_body(route: str, prompt: str) -> dict[str, JsonValue]:
|
||||
match route:
|
||||
case "/anthropic/v1/messages":
|
||||
return {"model": "claude-sonnet-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}
|
||||
case "/gemini/v1beta/models/gemini-2.5-flash:generateContent":
|
||||
return {"contents": [{"role": "user", "parts": [{"text": prompt}]}]}
|
||||
case _:
|
||||
return {"model": "gpt-5.6", "messages": [{"role": "user", "content": prompt}]}
|
||||
|
||||
|
||||
def _verdict(request: Request) -> Reply:
|
||||
action: Final = "BLOCKED" if BLOCK_MARKER.encode() in request.body else "NONE"
|
||||
return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode())
|
||||
|
||||
|
||||
def _openai_reply(_request: Request) -> Reply:
|
||||
return chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False)
|
||||
|
||||
|
||||
def _anthropic_reply(_request: Request) -> Reply:
|
||||
return Reply(body=message_json("msg_passthrough_auth", "claude-sonnet-5-5", "Hi"))
|
||||
|
||||
|
||||
def _gemini_reply(_request: Request) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"candidates": [{"content": {"role": "model", "parts": [{"text": "Hi"}]}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 1, "totalTokenCount": 5},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AuthGuardrailRig:
|
||||
owned: OwnedProxy
|
||||
vendor: Wire
|
||||
upstreams: Mapping[str, Wire]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[AuthGuardrailRig]:
|
||||
directory: Final = tmp_path_factory.mktemp("passthrough-auth-guardrails")
|
||||
with ExitStack() as stack:
|
||||
vendor: Final = stack.enter_context(wire_server(_verdict))
|
||||
upstreams: Final = MappingProxyType(
|
||||
{
|
||||
OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)),
|
||||
ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)),
|
||||
GEMINI_ROUTE: stack.enter_context(wire_server(_gemini_reply)),
|
||||
}
|
||||
)
|
||||
config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text()))
|
||||
path: Final = directory / "passthrough-auth-guardrails.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**config,
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": GUARDRAIL,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": "pre_call",
|
||||
"default_on": False,
|
||||
"api_base": vendor.url,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
environment: Final = {
|
||||
"OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url,
|
||||
"OPENAI_API_KEY": "synthetic-openai-key",
|
||||
"ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url,
|
||||
"ANTHROPIC_API_KEY": "synthetic-anthropic-key",
|
||||
"GEMINI_API_BASE": upstreams[GEMINI_ROUTE].url,
|
||||
"GEMINI_API_KEY": "synthetic-gemini-key",
|
||||
}
|
||||
owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path))
|
||||
yield AuthGuardrailRig(owned, vendor, upstreams)
|
||||
|
||||
|
||||
def _key_with_guardrail(scenario: Scenario, level: Level) -> str:
|
||||
attached: Final[dict[str, JsonValue]] = {"metadata": {"guardrails": [GUARDRAIL]}}
|
||||
team: Final = scenario.team(**(attached if level == "team" else {}))
|
||||
match level:
|
||||
case "project":
|
||||
return scenario.key(team_id=team, project_id=scenario.project(team, **attached))
|
||||
case "key":
|
||||
return scenario.key(team_id=team, **attached)
|
||||
case _:
|
||||
return scenario.key(team_id=team)
|
||||
|
||||
|
||||
def _send(gateway: Gateway, route: str, body: Mapping[str, JsonValue], key: str) -> httpx.Response:
|
||||
return gateway.request("POST", route, body, key=key, headers={"x-goog-api-key": key})
|
||||
|
||||
|
||||
def _drain(rig: AuthGuardrailRig) -> None:
|
||||
rig.vendor.drain()
|
||||
for upstream in rig.upstreams.values():
|
||||
upstream.drain()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("level", ["key", "team", "project"])
|
||||
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE])
|
||||
def test_auth_guardrail_blocks_the_prompt_before_the_provider(rig: AuthGuardrailRig, route: str, level: Level) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = _key_with_guardrail(scenario, level)
|
||||
_drain(rig)
|
||||
|
||||
response: Final = _send(gateway, route, _route_body(route, BLOCK_MARKER), key)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True]
|
||||
assert rig.upstreams[route].drain() == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE])
|
||||
def test_auth_guardrail_lets_a_clean_prompt_reach_the_provider(rig: AuthGuardrailRig, route: str) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = _key_with_guardrail(scenario, "key")
|
||||
marker: Final = f"clean-{uuid.uuid4().hex}"
|
||||
_drain(rig)
|
||||
|
||||
response: Final = _send(gateway, route, _route_body(route, marker), key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [marker.encode() in request.body for request in rig.vendor.drain()] == [True]
|
||||
assert [marker.encode() in request.body for request in rig.upstreams[route].drain()] == [True]
|
||||
|
||||
|
||||
def test_empty_request_guardrails_list_does_not_drop_the_key_guardrail(rig: AuthGuardrailRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = _key_with_guardrail(scenario, "key")
|
||||
_drain(rig)
|
||||
|
||||
response: Final = gateway.request(
|
||||
"POST", OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE, BLOCK_MARKER), "guardrails": []}, key=key
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert rig.upstreams[OPENAI_ROUTE].drain() == ()
|
||||
|
||||
|
||||
def test_request_without_a_body_is_not_sent_to_the_guardrail(rig: AuthGuardrailRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = _key_with_guardrail(scenario, "key")
|
||||
_drain(rig)
|
||||
|
||||
response: Final = gateway.request("GET", "/openai/v1/models", key=key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert rig.vendor.drain() == ()
|
||||
assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1
|
||||
|
|
@ -344,6 +344,7 @@ class TestVertexAIPassThroughHandler:
|
|||
custom_headers={"Authorization": f"Bearer {test_token}"},
|
||||
is_streaming_request=False,
|
||||
upstream_headers_filter=mock.ANY,
|
||||
inherit_auth_guardrails=True,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -444,6 +445,7 @@ class TestVertexAIPassThroughHandler:
|
|||
custom_headers={"Authorization": f"Bearer {test_token}"},
|
||||
is_streaming_request=False,
|
||||
upstream_headers_filter=mock.ANY,
|
||||
inherit_auth_guardrails=True,
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -540,6 +542,7 @@ class TestVertexAIPassThroughHandler:
|
|||
custom_headers={"Authorization": f"Bearer {default_credentials}"},
|
||||
is_streaming_request=False,
|
||||
upstream_headers_filter=mock.ANY,
|
||||
inherit_auth_guardrails=True,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -11353,3 +11353,122 @@ async def test_passthrough_header_and_callback_share_finalized_cost(
|
|||
finally:
|
||||
cleanup()
|
||||
await client.aclose()
|
||||
|
||||
|
||||
class TestInheritAuthGuardrails:
|
||||
"""Built-in provider routes run the caller's key, team and project guardrails on pass-through."""
|
||||
|
||||
@staticmethod
|
||||
async def _run(
|
||||
body: dict[str, object], *, inherit: bool, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> tuple[list[dict[str, object]], AsyncMock]:
|
||||
pre_call_data: Final[list[dict[str, object]]] = []
|
||||
|
||||
async def _record_pre_call(**kwargs: object) -> object:
|
||||
pre_call_data.append(json.loads(json.dumps(kwargs["data"], default=str)))
|
||||
return kwargs["data"]
|
||||
|
||||
proxy_logging: Final = MagicMock()
|
||||
proxy_logging.pre_call_hook = AsyncMock(side_effect=_record_pre_call)
|
||||
proxy_logging.post_call_success_hook = AsyncMock(return_value={"id": "chatcmpl-1"})
|
||||
proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
upstream: Final = httpx.Response(
|
||||
status_code=200,
|
||||
headers={"content-type": "application/json"},
|
||||
content=b'{"id": "chatcmpl-1"}',
|
||||
request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions"),
|
||||
)
|
||||
request: Final = MagicMock()
|
||||
request.method = "POST"
|
||||
request.query_params = {}
|
||||
request.headers = MagicMock()
|
||||
request.headers.copy.return_value = {}
|
||||
module: Final = "litellm.proxy.pass_through_endpoints.pass_through_endpoints"
|
||||
with ExitStack() as stack:
|
||||
stack.enter_context(
|
||||
patch(
|
||||
f"{module}.HttpPassThroughEndpointHelpers.non_streaming_http_request_handler",
|
||||
new_callable=AsyncMock,
|
||||
return_value=upstream,
|
||||
)
|
||||
)
|
||||
stack.enter_context(patch(f"{module}._is_streaming_response", return_value=False))
|
||||
stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging))
|
||||
stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True))
|
||||
stack.enter_context(
|
||||
patch(
|
||||
f"{module}.pass_through_endpoint_logging",
|
||||
MagicMock(pass_through_async_success_handler=AsyncMock()),
|
||||
)
|
||||
)
|
||||
stack.enter_context(patch(f"{module}.read_request_body", new_callable=AsyncMock, return_value=body))
|
||||
stack.enter_context(patch(f"{module}.safe_get_request_headers", return_value={}))
|
||||
await pass_through_request(
|
||||
request=request,
|
||||
target="https://api.openai.com/v1/chat/completions",
|
||||
custom_headers={},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
stream=False,
|
||||
inherit_auth_guardrails=inherit,
|
||||
)
|
||||
return pre_call_data, proxy_logging.post_call_success_hook
|
||||
|
||||
@staticmethod
|
||||
def _chat_body(**extra: object) -> dict[str, object]:
|
||||
return {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], **extra}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_team_and_project_guardrails_reach_pre_and_post_call(self):
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
api_key="sk-test",
|
||||
metadata={"guardrails": ["key-guard"]},
|
||||
team_metadata={"guardrails": ["team-guard"]},
|
||||
project_metadata={"guardrails": ["project-guard"]},
|
||||
)
|
||||
|
||||
pre_call, post_call = await self._run(self._chat_body(), inherit=True, user_api_key_dict=auth)
|
||||
|
||||
assert sorted(pre_call[0]["metadata"]["guardrails"]) == ["key-guard", "project-guard", "team-guard"]
|
||||
post_call.assert_awaited_once()
|
||||
assert sorted(post_call.call_args.kwargs["data"]["metadata"]["guardrails"]) == [
|
||||
"key-guard",
|
||||
"project-guard",
|
||||
"team-guard",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_body_guardrails_are_added_to_the_key_guardrails(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
pre_call, _ = await self._run(
|
||||
self._chat_body(guardrails=["body-guard"]), inherit=True, user_api_key_dict=auth
|
||||
)
|
||||
|
||||
assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "body-guard"]
|
||||
assert "guardrails" not in pre_call[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_request_body_guardrails_keep_the_key_guardrails(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
pre_call, _ = await self._run(self._chat_body(guardrails=[]), inherit=True, user_api_key_dict=auth)
|
||||
|
||||
assert pre_call[0]["metadata"]["guardrails"] == ["key-guard"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routes_without_inheritance_keep_key_guardrails_off(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
pre_call, post_call = await self._run(self._chat_body(), inherit=False, user_api_key_dict=auth)
|
||||
|
||||
assert "guardrails" not in pre_call[0].get("metadata", {})
|
||||
post_call.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_without_a_body_gets_no_guardrails(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
pre_call, post_call = await self._run({}, inherit=True, user_api_key_dict=auth)
|
||||
|
||||
assert "guardrails" not in pre_call[0].get("metadata", {})
|
||||
post_call.assert_not_awaited()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue