mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
feat(proxy): opt-in include_guardrail_response returns guardrail_information in the response (#42327)
* feat(proxy): opt-in include_guardrail_response returns guardrail_information in the response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): read include_guardrail_response from the request metadata bucket the router did not reseed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): format common request processing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): redact matched content in guardrail_information and stop mutating cached responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): traverse guardrail diagnostics iteratively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): annotate guardrail traversal cast Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): reuse core redaction helper for guardrail_information Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(proxy): justify response rebind when attaching guardrail information Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
7516327898
commit
cb2f22533c
10 changed files with 379 additions and 8 deletions
|
|
@ -3,7 +3,7 @@
|
|||
import copy
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Iterable, Mapping
|
||||
from collections.abc import Collection, Iterable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
|
||||
|
||||
|
|
@ -709,10 +709,11 @@ def filter_internal_params(data: dict, additional_internal_params: set | None =
|
|||
|
||||
def redact_nested_match_and_regex_keys(
|
||||
payload: dict | list[Any] | str | None,
|
||||
keys: Collection[str] = ("match", "regex"),
|
||||
) -> dict | list[Any] | str | None:
|
||||
"""
|
||||
Deep-copy `payload` and replace every `match` / `regex` string field with
|
||||
"[REDACTED]" anywhere in nested dict/list structures.
|
||||
Deep-copy `payload` and replace every configured string field with "[REDACTED]"
|
||||
anywhere in nested dict/list structures.
|
||||
|
||||
Used for guardrail spend/compliance logging so raw spans are not persisted.
|
||||
"""
|
||||
|
|
@ -734,10 +735,9 @@ def redact_nested_match_and_regex_keys(
|
|||
continue
|
||||
seen.add(node_id)
|
||||
if isinstance(node, dict):
|
||||
if "match" in node:
|
||||
node["match"] = "[REDACTED]"
|
||||
if "regex" in node:
|
||||
node["regex"] = "[REDACTED]"
|
||||
for key in keys:
|
||||
if key in node:
|
||||
node[key] = "[REDACTED]"
|
||||
stack.extend(node.values())
|
||||
elif isinstance(node, list):
|
||||
stack.extend(node)
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ import httpx
|
|||
import orjson
|
||||
from fastapi import HTTPException, Request, status
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
import litellm
|
||||
|
|
@ -56,6 +56,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
get_or_create_metadata_bucket,
|
||||
independent_snapshot,
|
||||
is_expected_client_error,
|
||||
redact_nested_match_and_regex_keys,
|
||||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
|
||||
from litellm.litellm_core_utils.get_supported_openai_params import (
|
||||
|
|
@ -1401,6 +1402,42 @@ def _override_openai_response_model(
|
|||
)
|
||||
|
||||
|
||||
_METADATA_BUCKET_KEYS: Final = ("metadata", "litellm_metadata")
|
||||
_RESPONSE_REDACTED_KEYS: Final = ("keyword", "snippet", "match", "regex")
|
||||
|
||||
|
||||
def _request_metadata_buckets(request_data: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
return tuple(bucket for key in _METADATA_BUCKET_KEYS if isinstance(bucket := request_data.get(key), Mapping))
|
||||
|
||||
|
||||
def include_guardrail_response_requested(request_data: Mapping[str, object]) -> bool:
|
||||
return any(bucket.get("include_guardrail_response") is True for bucket in _request_metadata_buckets(request_data))
|
||||
|
||||
|
||||
def attach_guardrail_information(response: object, request_data: Mapping[str, object]) -> object:
|
||||
recorded: Final[Sequence[object]] = next(
|
||||
(
|
||||
entries
|
||||
for bucket in _request_metadata_buckets(request_data)
|
||||
if isinstance(
|
||||
entries := bucket.get("standard_logging_guardrail_information"),
|
||||
list,
|
||||
)
|
||||
),
|
||||
(),
|
||||
)
|
||||
guardrail_information: Final = [ # mutable-ok: response list contract
|
||||
redact_nested_match_and_regex_keys(entry, keys=_RESPONSE_REDACTED_KEYS)
|
||||
for entry in recorded
|
||||
if isinstance(entry, dict)
|
||||
]
|
||||
if isinstance(response, dict):
|
||||
return response | MappingProxyType({"guardrail_information": guardrail_information})
|
||||
if isinstance(response, BaseModel) and response.model_config.get("extra") == "allow":
|
||||
return response.model_copy(update=MappingProxyType({"guardrail_information": guardrail_information}))
|
||||
return response
|
||||
|
||||
|
||||
class CostBreakdownHeaderValues(NamedTuple):
|
||||
original_cost: float | None = None
|
||||
discount_amount: float | None = None
|
||||
|
|
@ -2898,6 +2935,11 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if isinstance(response, dict):
|
||||
response.pop("_hidden_params", None)
|
||||
|
||||
if include_guardrail_response_requested(self.data):
|
||||
response = attach_guardrail_information( # rebind-ok: response tail rebinds the copied response
|
||||
response=response, request_data=self.data
|
||||
)
|
||||
|
||||
# Call response headers hook for non-streaming success
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=self.data,
|
||||
|
|
|
|||
|
|
@ -3116,7 +3116,15 @@ async def move_guardrails_to_metadata(
|
|||
- If guardrails not set on API key, then checks request metadata
|
||||
- Adds guardrails from policies attached to key/team metadata
|
||||
- Adds guardrails from policy engine based on team/key/model context
|
||||
- Moves include_guardrail_response into request metadata before provider dispatch
|
||||
"""
|
||||
if "include_guardrail_response" in data:
|
||||
data[_metadata_variable_name][
|
||||
"include_guardrail_response"
|
||||
] = ( # rebind-ok: pre-call hooks mutate the shared request dict in place
|
||||
data.pop("include_guardrail_response") is True
|
||||
)
|
||||
|
||||
# Early-out: skip all guardrails processing when nothing is configured
|
||||
key_metadata: Final = user_api_key_dict.metadata
|
||||
team_metadata: Final = user_api_key_dict.team_metadata
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@
|
|||
- {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"}
|
||||
- {id: guardrail.litellm_content_filter.pre_call.blocks_video, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [videos], source: "test_key_guardrail_video_e2e.py", fail_before_fix: proven, rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/videos before the provider is called; before the fix the route's call type was unknown to the unified guardrail hook and the prompt went to the provider unscanned (LIT-6685)"}
|
||||
- {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"}
|
||||
- {id: guardrail.litellm_content_filter.pre_call.returns_guardrail_information, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "guardrails/test_guardrail_information_response_e2e.py", rationale: "Opt-in chat responses expose successful guardrail execution details"}
|
||||
- {id: guardrail.litellm_content_filter.apply_endpoint.blocks, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail blocks banned content for customers that call the apply surface directly"}
|
||||
- {id: guardrail.litellm_content_filter.apply_endpoint.allows, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [allows], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail returns clean text for allowed input"}
|
||||
- {id: guardrail.bedrock.during.blocks, module: guardrail, tier: P0, hook_point: during, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "During-call moderation for streaming"}
|
||||
|
|
|
|||
|
|
@ -291,6 +291,7 @@ class GuardrailsClient:
|
|||
text: str,
|
||||
*,
|
||||
guardrails: list[str] | None = None,
|
||||
include_guardrail_response: bool | None = None,
|
||||
max_tokens: int = 16,
|
||||
tools: list[ChatTool] | None = None,
|
||||
) -> Result[ChatResponse]:
|
||||
|
|
@ -306,6 +307,7 @@ class GuardrailsClient:
|
|||
messages=[ChatMessage(role="user", content=text)],
|
||||
max_tokens=max_tokens,
|
||||
guardrails=guardrails,
|
||||
include_guardrail_response=include_guardrail_response,
|
||||
tools=tools,
|
||||
),
|
||||
)
|
||||
|
|
|
|||
118
tests/e2e/guardrails/test_guardrail_information_response_e2e.py
Normal file
118
tests/e2e/guardrails/test_guardrail_information_response_e2e.py
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
"""Live e2e: an opted-in chat response includes the guardrail execution details."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import unwrap
|
||||
from guardrails_client import (
|
||||
BlockedWordBody,
|
||||
ContentFilterParamsBody,
|
||||
GuardrailsClient,
|
||||
)
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatResponse, GuardrailInformationEntry
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
GUARDRAIL_PROPAGATION_DEADLINE_SECONDS: Final = 40.0
|
||||
GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS: Final = 5.0
|
||||
|
||||
|
||||
def _register_content_filter(client: GuardrailsClient, resources: ResourceManager, *, name: str) -> None:
|
||||
guardrail_id = client.register(
|
||||
name,
|
||||
ContentFilterParamsBody(
|
||||
mode="pre_call",
|
||||
default_on=False,
|
||||
blocked_words=[BlockedWordBody(keyword=f"never-match-{unique_marker()}", action="MASK")],
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
|
||||
def _opted_in_entries(
|
||||
client: GuardrailsClient,
|
||||
key: str,
|
||||
model: str,
|
||||
name: str,
|
||||
) -> tuple[ChatResponse, tuple[GuardrailInformationEntry, ...]]:
|
||||
response = unwrap(
|
||||
client.chat(
|
||||
key,
|
||||
model,
|
||||
"Reply with the single word OK.",
|
||||
guardrails=[name],
|
||||
include_guardrail_response=True,
|
||||
max_tokens=16,
|
||||
)
|
||||
)
|
||||
entries = tuple(entry for entry in response.guardrail_information or () if entry.guardrail_name == name)
|
||||
return response, entries
|
||||
|
||||
|
||||
class TestGuardrailInformationResponse:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.litellm_content_filter.pre_call.returns_guardrail_information",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_flag_returns_guardrail_information_for_the_guardrail_that_ran(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name = f"e2e-guardrail-information-{unique_marker()}"
|
||||
_register_content_filter(client, resources, name=name)
|
||||
model = client.create_backend_model(
|
||||
resources,
|
||||
prefix="e2e-guardrail-info-backend",
|
||||
backend="openai/gpt-4.1-mini",
|
||||
api_key="os.environ/OPENAI_API_KEY",
|
||||
)
|
||||
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
|
||||
|
||||
while True:
|
||||
response, entries = _opted_in_entries(client, scoped_key, model, name)
|
||||
if len(entries) == 1:
|
||||
entry = entries[0]
|
||||
assert entry.guardrail_status == "success", (
|
||||
f"guardrail information should report a successful run, got {entry!r}; response: {response}"
|
||||
)
|
||||
assert entry.duration is not None and entry.duration >= 0, (
|
||||
f"guardrail information should report a non-negative duration; response: {response}"
|
||||
)
|
||||
return
|
||||
if time.monotonic() >= deadline:
|
||||
pytest.fail(
|
||||
f"guardrail information did not report exactly one successful {name!r} entry within "
|
||||
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; response: {response}"
|
||||
)
|
||||
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
|
||||
|
||||
def test_without_flag_response_has_no_guardrail_information(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name = f"e2e-guardrail-information-default-{unique_marker()}"
|
||||
_register_content_filter(client, resources, name=name)
|
||||
model = client.create_backend_model(
|
||||
resources,
|
||||
prefix="e2e-guardrail-info-backend",
|
||||
backend="openai/gpt-4.1-mini",
|
||||
api_key="os.environ/OPENAI_API_KEY",
|
||||
)
|
||||
|
||||
response = unwrap(
|
||||
client.chat(
|
||||
scoped_key,
|
||||
model,
|
||||
"Reply with the single word OK.",
|
||||
guardrails=[name],
|
||||
max_tokens=16,
|
||||
)
|
||||
)
|
||||
|
||||
assert "guardrail_information" not in response.model_fields_set, (
|
||||
f"guardrail information must remain absent without include_guardrail_response, got {response}"
|
||||
)
|
||||
|
|
@ -334,6 +334,7 @@ class ChatBody(BaseModel):
|
|||
tools: Sequence[ChatTool | McpChatTool] | None = None
|
||||
tool_choice: str | None = None
|
||||
guardrails: list[str] | None = None
|
||||
include_guardrail_response: bool | None = None
|
||||
response_format: dict[str, object] | None = None
|
||||
chat_template_kwargs: dict[str, bool] | None = None
|
||||
cache: dict[str, bool] | None = {"no-cache": True}
|
||||
|
|
@ -448,6 +449,14 @@ class Usage(BaseModel):
|
|||
completion_tokens_details: CompletionTokensDetails | None = None
|
||||
|
||||
|
||||
class GuardrailInformationEntry(BaseModel):
|
||||
guardrail_name: str
|
||||
guardrail_status: str
|
||||
guardrail_mode: object | None = None
|
||||
guardrail_response: object | None = None
|
||||
duration: float | None = None
|
||||
|
||||
|
||||
class ChatResponse(BaseModel):
|
||||
id: str | None = None
|
||||
object: str | None = None
|
||||
|
|
@ -455,6 +464,7 @@ class ChatResponse(BaseModel):
|
|||
choices: list[ChatChoice] = []
|
||||
usage: Usage | None = None
|
||||
service_tier: str | None = None
|
||||
guardrail_information: list[GuardrailInformationEntry] | None = None
|
||||
|
||||
|
||||
# ---------- anthropic /v1/messages + count_tokens ----------
|
||||
|
|
|
|||
|
|
@ -322,6 +322,28 @@ class TestRedactNestedMatchAndRegexKeys:
|
|||
assert redact_nested_match_and_regex_keys(None) is None
|
||||
assert redact_nested_match_and_regex_keys("plain") == "plain"
|
||||
|
||||
def test_redacts_custom_keys_without_changing_default_keys(self):
|
||||
payload = {
|
||||
"keyword": "secret-keyword",
|
||||
"snippet": "secret-snippet",
|
||||
"match": "secret-match",
|
||||
"regex": "secret-regex",
|
||||
"nested": [{"keyword": "nested-keyword", "match": "nested-match"}],
|
||||
}
|
||||
|
||||
custom_keys = redact_nested_match_and_regex_keys(payload, keys=("keyword", "snippet"))
|
||||
default_keys = redact_nested_match_and_regex_keys(payload)
|
||||
|
||||
assert custom_keys["keyword"] == "[REDACTED]"
|
||||
assert custom_keys["snippet"] == "[REDACTED]"
|
||||
assert custom_keys["nested"][0]["keyword"] == "[REDACTED]"
|
||||
assert custom_keys["match"] == "secret-match"
|
||||
assert custom_keys["regex"] == "secret-regex"
|
||||
assert default_keys["match"] == "[REDACTED]"
|
||||
assert default_keys["regex"] == "[REDACTED]"
|
||||
assert default_keys["keyword"] == "secret-keyword"
|
||||
assert default_keys["snippet"] == "secret-snippet"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, expected",
|
||||
|
|
|
|||
|
|
@ -35,10 +35,12 @@ from litellm.proxy.common_request_processing import (
|
|||
_buffer_first_chunk_honoring_disconnect,
|
||||
_cancel_llm_call_on_client_disconnect,
|
||||
_ClientDisconnectedBeforeFirstChunk,
|
||||
attach_guardrail_information,
|
||||
_extract_error_from_sse_chunk,
|
||||
_get_cost_breakdown_from_logging_obj,
|
||||
CostBreakdownHeaderValues,
|
||||
_has_attribute_error_in_chain,
|
||||
include_guardrail_response_requested,
|
||||
_is_azure_model_router_request,
|
||||
open_sse_before_first_byte,
|
||||
resolve_litellm_call_id,
|
||||
|
|
@ -61,6 +63,132 @@ from litellm.proxy.utils import ProxyLogging
|
|||
from litellm.router import Router
|
||||
|
||||
|
||||
def test_attach_guardrail_information_copies_recorded_entries_onto_model_response():
|
||||
recorded = [
|
||||
{"guardrail_name": "first", "guardrail_status": "success"},
|
||||
{"guardrail_name": "second", "guardrail_status": "success"},
|
||||
]
|
||||
response = litellm.ModelResponse()
|
||||
|
||||
result = attach_guardrail_information(
|
||||
response=response,
|
||||
request_data={"metadata": {"standard_logging_guardrail_information": recorded}},
|
||||
)
|
||||
|
||||
assert isinstance(result, litellm.ModelResponse)
|
||||
assert result.model_dump()["guardrail_information"] == recorded
|
||||
assert "guardrail_information" not in response.model_dump()
|
||||
|
||||
|
||||
def test_attach_guardrail_information_reports_empty_list_when_nothing_ran():
|
||||
response = litellm.ModelResponse()
|
||||
|
||||
result = attach_guardrail_information(response=response, request_data={})
|
||||
|
||||
assert isinstance(result, litellm.ModelResponse)
|
||||
assert result.model_dump()["guardrail_information"] == []
|
||||
assert "guardrail_information" not in response.model_dump()
|
||||
|
||||
|
||||
def test_attach_guardrail_information_sets_key_on_dict_response():
|
||||
recorded = [{"guardrail_name": "first", "guardrail_status": "success"}]
|
||||
response = {"id": "x"}
|
||||
|
||||
result = attach_guardrail_information(
|
||||
response=response,
|
||||
request_data={"metadata": {"standard_logging_guardrail_information": recorded}},
|
||||
)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert result == {"id": "x", "guardrail_information": recorded}
|
||||
assert response == {"id": "x"}
|
||||
|
||||
|
||||
def test_attach_guardrail_information_redacts_matched_content():
|
||||
recorded = [
|
||||
{
|
||||
"guardrail_name": "cf",
|
||||
"guardrail_status": "success",
|
||||
"guardrail_response": [
|
||||
{"type": "blocked_word", "keyword": "secret-word", "action": "MASK"}
|
||||
],
|
||||
"match_details": [{"snippet": "secret-word", "detection_method": "keyword"}],
|
||||
}
|
||||
]
|
||||
|
||||
result = attach_guardrail_information(
|
||||
response={"id": "x"},
|
||||
request_data={"metadata": {"standard_logging_guardrail_information": recorded}},
|
||||
)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
guardrail_information = result["guardrail_information"]
|
||||
assert isinstance(guardrail_information, list)
|
||||
assert guardrail_information[0]["guardrail_response"][0]["keyword"] == "[REDACTED]"
|
||||
assert guardrail_information[0]["match_details"][0]["snippet"] == "[REDACTED]"
|
||||
assert guardrail_information[0]["match_details"][0]["detection_method"] == "keyword"
|
||||
assert "secret-word" not in json.dumps(result)
|
||||
|
||||
|
||||
def test_attach_guardrail_information_leaves_cached_dict_response_untouched():
|
||||
recorded = [{"guardrail_name": "cf", "guardrail_status": "success"}]
|
||||
cached = {"id": "x", "content": []}
|
||||
|
||||
result = attach_guardrail_information(
|
||||
response=cached,
|
||||
request_data={
|
||||
"metadata": {
|
||||
"include_guardrail_response": True,
|
||||
"standard_logging_guardrail_information": recorded,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert "guardrail_information" not in cached
|
||||
assert result is not cached
|
||||
assert isinstance(result, dict)
|
||||
assert result["guardrail_information"] == recorded
|
||||
|
||||
original = litellm.ModelResponse()
|
||||
copied = attach_guardrail_information(
|
||||
response=original,
|
||||
request_data={"metadata": {"standard_logging_guardrail_information": recorded}},
|
||||
)
|
||||
|
||||
assert "guardrail_information" not in original.model_dump()
|
||||
assert isinstance(copied, litellm.ModelResponse)
|
||||
assert copied.model_dump()["guardrail_information"] == recorded
|
||||
|
||||
|
||||
def test_include_guardrail_response_requested_reads_flag_from_metadata_when_router_seeded_litellm_metadata():
|
||||
recorded = [
|
||||
{"guardrail_name": "first", "guardrail_status": "success"},
|
||||
{"guardrail_name": "second", "guardrail_status": "success"},
|
||||
]
|
||||
request_data = {
|
||||
"metadata": {
|
||||
"include_guardrail_response": True,
|
||||
"standard_logging_guardrail_information": recorded,
|
||||
},
|
||||
"litellm_metadata": {},
|
||||
}
|
||||
|
||||
assert include_guardrail_response_requested(request_data) is True
|
||||
|
||||
response = litellm.ModelResponse()
|
||||
result = attach_guardrail_information(response=response, request_data=request_data)
|
||||
|
||||
assert isinstance(result, litellm.ModelResponse)
|
||||
assert result.model_dump()["guardrail_information"] == recorded
|
||||
|
||||
|
||||
def test_include_guardrail_response_requested_is_false_without_exact_true():
|
||||
assert include_guardrail_response_requested(
|
||||
{"metadata": {"include_guardrail_response": "true"}, "litellm_metadata": {}}
|
||||
) is False
|
||||
assert include_guardrail_response_requested({}) is False
|
||||
|
||||
|
||||
class TestProxyBaseLLMRequestProcessing:
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_passthrough_process_llm_request_preserves_litellm_headers_for_non_streaming_response(
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
add_provider_specific_headers_to_request,
|
||||
check_if_token_is_service_account,
|
||||
clean_headers,
|
||||
move_guardrails_to_metadata,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
|
||||
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
|
||||
|
|
@ -5187,6 +5188,45 @@ def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_move_guardrails_to_metadata_moves_include_guardrail_response_before_the_no_guardrail_early_out():
|
||||
policy_registry = MagicMock()
|
||||
policy_registry.is_initialized.return_value = False
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
true_data = {
|
||||
"model": "gpt-4.1-mini",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"metadata": {},
|
||||
"include_guardrail_response": True,
|
||||
}
|
||||
with patch("litellm.proxy.policy_engine.policy_registry.get_policy_registry", return_value=policy_registry):
|
||||
await move_guardrails_to_metadata(
|
||||
data=true_data,
|
||||
_metadata_variable_name="metadata",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert "include_guardrail_response" not in true_data
|
||||
assert true_data["metadata"]["include_guardrail_response"] is True
|
||||
|
||||
string_data = {
|
||||
"model": "gpt-4.1-mini",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"metadata": {},
|
||||
"include_guardrail_response": "true",
|
||||
}
|
||||
with patch("litellm.proxy.policy_engine.policy_registry.get_policy_registry", return_value=policy_registry):
|
||||
await move_guardrails_to_metadata(
|
||||
data=string_data,
|
||||
_metadata_variable_name="metadata",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert "include_guardrail_response" not in string_data
|
||||
assert string_data["metadata"]["include_guardrail_response"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_guardrail_merges_with_global_policy():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue