mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
* fix(pricing): add GitHub Copilot MAI Code Flash pricing (#30415) * fix(pricing): add GitHub Copilot MAI Code Flash pricing Add GitHub Copilot pricing entries for MAI-Code-1-Flash and the internal Copilot CLI model name so cost calculation can price input, cached input, and output tokens. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test(pricing): cover GitHub Copilot MAI Code Flash pricing Add regression coverage for both GitHub Copilot MAI-Code-1-Flash model names, including cached input pricing, chat endpoint metadata, and cost_per_token arithmetic. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(router/proxy): propagate completed_response through FallbackResponsesStreamWrapper for streaming /v1/responses container ownership (#30210) (#30213) * fix(router/proxy): propagate completed_response through FallbackResponsesStreamWrapper for streaming /v1/responses container ownership (#30210) #28990 added ownership recording for streaming /v1/responses via _wrap_responses_stream_for_container_ownership, which reads `getattr(stream_response, 'completed_response', None)` to extract the ResponsesAPIResponse. The unit test bypassed the Router, so it never exercised the production wrapping path. Through the Router (every proxy deployment), the stream is wrapped by FallbackResponsesStreamWrapper (router.py:2527). Its __init__ set `self.completed_response = None` and __anext__ only forwarded chunks — the inner source iterator's terminal event never bubbled up to the attribute the ownership hook reads, so the hook silently recorded nothing and every follow-up /v1/containers/<id>/files call returned 403 for non-admin keys. This commit: - router.py: pre-resolves the responses-API terminal event tuple (response.completed / .incomplete / .failed) once per _aresponses_streaming_iterator call, and has the wrapper's __anext__ sniff each forwarded chunk's .type. First terminal event hit gets stored on the wrapper's completed_response. Iterator-agnostic — works for source_iterator AND any future wrapper. - common_request_processing.py: when _extract_completed_responses_response returns None we now warn instead of silently skipping. Reporter on #30210 lost a day to this exact silent skip; the warning surfaces future regressions of the same shape directly in operator logs. Fixes #30210 * fix(router): type-ignore wrapper getattr-defaults; broaden ownership-skip warning CI lint (mypy) flagged the three pre-existing getattr(..., None) assignments in FallbackResponsesStreamWrapper.__init__: router.py:2564 self.response = getattr(source_iterator, 'response', None) router.py:2565 self.model = getattr(source_iterator, 'model', None) router.py:2566 self.logging_obj = getattr(..., None) Those lines also exist on litellm_internal_staging and pass mypy there. Adding the typed terminal-event tuple above the class made the function body more narrowable, which surfaced the pre-existing mismatch — base class declares non-Optional types but the bridge path (LiteLLMCompletionStreamingIterator) legitimately omits these. Keep the None fallback and silence with type: ignore[assignment]. Greptile 4/5 note: the ownership-skip warning hard-named code_interpreter which misleads operators when a non-code_interpreter stream aborts. Generalize to 'any tool container (e.g. code_interpreter)'. * fix(register_model): drop synthesized zero costs to preserve sparse entries (#30198) (#30201) * fix(register_model): drop synthesized zero costs to preserve sparse entries (#30198) get_model_info synthesizes input_cost_per_token / output_cost_per_token = 0 when they are absent from the raw entry (the price-unknown and free cases share the same representation). register_model then merges that result back into litellm.model_cost, which flips a sparse entry from 'no cost keys' (priced via model name) to 'cost keys = 0' (free). That defeats _is_cost_explicitly_configured (#24949) on re-registration: _is_model_cost_zero returns True, common_checks skips every tag / key / team / user / org budget check for the group, and over-budget traffic keeps returning 200. Spend keeps recording because cost calc still resolves by model name, so the symptom is silent and only triggers on the second register_model pass (router rebuild, /model/update, config sync). Mirror the existing litellm_provider-None guard one block above and pop the cost fields from the synthesized result when they are absent from the raw entry and not in the caller's value. Caller-provided zeros (genuinely free models, BYOK overrides) are preserved. Fixes #30198 * fix(register_model): switch _raw_entry to is-None checks + drop dead test assertion Greptile #30201 review notes: - the `or`-chain in the raw-entry lookup treated an empty dict (a key with no fields) as falsy and fell through to the second arm — replace with explicit `is None` checks so a present-but-empty entry is still taken at face value. - the first assertion in `test_router_double_init_keeps_db_model_entry_sparse` used `in (None, 0)` which passes under the bug condition (cost = 0 matches the tuple); the strong follow-up assertion already covers every shape, so drop the dead branch. * fix(bedrock mantle): use unique function-call id for responses->chat tool calls (#30426) * fix(bedrock mantle): use unique function-call id for responses->chat tool calls ... * fix(bedrock mantle): scope unique tool-call id fallback to degenerate call_id The previous revision preferred the Responses item id for every tool call, which broke providers (and existing tests) where call_id is a unique, canonical correlation key. Restrict the fallback to the degenerate index-based call_id that Bedrock Mantle returns (call_0, call_1, ... resetting per response) and keep call_id otherwise. Revert the change to the OUTPUT_ITEM_DONE streaming handler, whose tool_call_chunk is never emitted (dead code, per review). Extend the regression tests to assert a normal call_id is preserved. * fix(router): preserve azure_ad_token through CredentialLiteLLMParams for /v1/files + batches (#30235) (#30241) * fix(router): preserve azure_ad_token through CredentialLiteLLMParams for /v1/files + batches (#30235) Router.get_deployment_credentials_with_provider re-validates a deployment's litellm_params through CredentialLiteLLMParams before handing them to file/batch/passthrough callers: return CredentialLiteLLMParams( **deployment.litellm_params.model_dump(exclude_none=True) ).model_dump(exclude_none=True) Any field NOT declared on CredentialLiteLLMParams gets silently dropped on the way through. azure_ad_token was undeclared, so Azure deployments using OAuth/M2M (azure_ad_token instead of a static api_key) silently lost their token at the files endpoint and the proxy returned: Missing credentials. Please pass one of api_key, azure_ad_token, azure_ad_token_provider, ... Declare azure_ad_token on CredentialLiteLLMParams alongside api_key / api_base / api_version so it rides through the round-trip. Static-key deployments stay unaffected (Optional, default None, dropped by exclude_none=True). Provider-callable (azure_ad_token_provider) is a separate concern and out of scope here. Fixes #30235 * fix(ui-types): regenerate schema.d.ts for new azure_ad_token field CI's 'Verify schema.d.ts matches the proxy OpenAPI spec' check auto-detected the new field and emitted the exact diff to apply. Two schemas had `aws_secret_access_key` from CredentialLiteLLMParams, both get the new azure_ad_token marker next to it. * fix(proxy): org_admin with own user_id now sees all org teams on /v2/team/list (#30247) When the UI sends the callers own user_id (as it does for non-Admin global roles), _enforce_list_team_v2_access now nulls it out for org admins so _build_team_list_where_conditions scopes by organization_id only -- matching the legacy /team/list behavior and the documented intent. Fixes #30215 Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * test(vertex_ai): multi-region regression coverage for cachedContents host (#29571) (#29707) litellm_internal_staging already routes the cachedContents URL through get_vertex_base_url, fixing the multi-region 404 reported in #29571 — but carries no test coverage for the actual regression scenario (eu/us must resolve to the REP host aiplatform.{geo}.rep.googleapis.com). Add TestContextCachingMultiRegionUrls: parametrized eu/us REP-host assertions (including absence of the old broken {geo}-aiplatform host), plus regional (us-central1) and global no-regression checks. * fix(proxy): close upstream LLM stream when client disconnects mid-stream (#30245) * fix(proxy): close upstream LLM stream when client disconnects mid-stream When a streaming client disconnects, Starlette abandons the response body iterator without calling aclose(), so the proxy's connection to the upstream backend stays open until garbage collection, which may never come. The backend (e.g. vLLM) keeps generating into a dead pipe: small responses drain invisibly into TCP buffers while large ones block the backend on a full send buffer indefinitely (observed via lsof as an ESTABLISHED proxy->backend connection minutes after the client left) create_response now returns a StreamingResponse subclass that closes both its body iterator and the wrapped upstream-facing generator in a shielded finally. The upstream generator is closed directly rather than through a cascade because aclose() on a never-started generator skips its body, which would make the cascade a no-op when the client disconnects before the first chunk is sent. async_streaming_data_generator also gains the same shielded finally-aclose that async_data_generator in proxy_server.py already had, covering the Anthropic and Google SSE paths With this, killing a streaming client causes the backend to observe the abort within about a second and free its slot, while completed streams are unaffected. No flag is needed, unlike the non-streaming opt-in cancel in #30223: this only releases resources after the client is already gone and does not change any response a client can observe Fixes #30244 * fix(proxy): close upstream even when body iterator aclose raises BaseException Addresses the Greptile finding on #30245: the cleanup loop caught only Exception while the generator-level cleanup catches BaseException, so a CancelledError or GeneratorExit escaping body_iterator.aclose() would skip closing the upstream generator. Both sites now use the same scope and a regression test pins that the upstream is closed even when the body iterator explodes with a BaseException * fix(llms): expose aclose on BaseModelResponseIterator so stream close reaches the provider connection The response-level close added for #30244 only worked for SDK-based providers (e.g. openai), whose streams expose aclose all the way down. Providers served by base_llm_http_handler (hosted_vllm and most modern transformation-based providers) wrap a bare response.aiter_lines() generator in BaseModelResponseIterator, which had no aclose or close at all, and nothing retained the httpx response object; so CustomStreamWrapper.aclose() silently did nothing and the upstream connection stayed open. Verified with a vLLM-style mock: with hosted_vllm/ the backend streamed all 100 chunks to completion after the client disconnected, while openai/ aborted at chunk 6 BaseModelResponseIterator now carries an optional http_response and an aclose() that closes it; make_async_call_stream_helper attaches the response after building the iterator. With this, hosted_vllm aborts the backend within ~1.6s of the client dropping, and completed streams are unaffected --------- Co-authored-by: kursad <kursad.lacin@brado.net> * feat(anthropic): surface compaction usage iterations data (#27065) * feat(anthropic): surface compaction usage iterations data * style: apply black formatting to fix lint checks * fix(usage): correct calculate usage with cached tokens when use ChatCompletionUsageBlock (#30422) * fix(usage): correct calculate usage with cached tokens when use ChatCompletionUsageBlock * fix(usage): optimize test imports * feat: add fastCRW search provider (#30434) * feat(provider): add LibertAI as a JSON-configured OpenAI-compatible provider (#30203) * feat(provider): add LibertAI as a JSON-configured OpenAI-compatible provider * libertai: update served endpoints backup + add mode/matrix tests Addresses review feedback: - Add libertai to litellm/provider_endpoints_support_backup.json, the file actually served by GET /public/supported_endpoints (the root provider_endpoints_support.json already had it). - Add tests asserting bge-m3 normalizes to mode='embedding' and that the served matrix lists libertai. embeddings stays false: the JSON-configured provider path only wires chat routing (OpenAILike embedding handler is reached only for literal openai_like/llamafile/lm_studio), matching the llamagate precedent; bge-m3 remains in the cost map for metadata. --------- Co-authored-by: Moshe Malawach <moshemalawach@users.noreply.github.com> * feat(provider): add ModelScope as an OpenAI-compatible provider (#28460) * add ModelScope API support * add modelscope api support * update modelscope model list * add image-genetation support * update test and multimodal * fix: address PR review feedback for modelscope provider * update README * fix(customer_endpoints): restrict /customer/daily/activity to admin-only (#28849) * fix(customer_endpoints): restrict /customer/daily/activity to admin-only * fix(customer_endpoints): check role before prisma_client guard * fix(custom_guardrail): key disable_global_guardrails takes precedence over team guardrail list (#28563) * fix(fallbacks): preserve fallback model in SDK fallback responses (#28260) * fix(fallbacks): preserve fallback model in response when using SDK-level fallbacks * fix(fallbacks): gate x-litellm-* passthrough to trusted callers only The previous patch unconditionally let `x-litellm-*` keys bypass the `llm_provider-` prefix in `process_response_headers`. That function is also called on raw upstream-provider response headers (e.g. from `llm_http_handler.py`), so a malicious provider could return `x-litellm-attempted-fallbacks` and spoof a LiteLLM-internal marker, bypassing the proxy model-override guard. Add a `preserve_litellm_internal_headers` flag (default False). Only `response_metadata.py`, which re-processes the already-built `_hidden_params["additional_headers"]` dict (LiteLLM-owned), passes True. Raw provider header callsites keep the default False, so upstream `x-litellm-*` still gets the `llm_provider-` prefix. Adds a regression test for the spoofing case and renames the existing preserve test to make the trusted-path semantics explicit. * fix(fallbacks): ignore preserve_litellm_internal_headers for raw httpx.Headers inputs * style(core_helpers): apply black formatting * fix(lint): remove banned typing.List/Dict/Any imports and suppress PLR0913 on interface overrides Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): apply black formatting to modelscope chat transformation Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): replace noqa with proper fixes — use **kwargs and Awaitable instead of Any/List Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): remove unused AllMessageValues import Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * revert: restore base_model_iterator.py to original PR state Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): restore full method signatures for MyPy compatibility; bump PLR0913 budget for new provider files Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): use @override to suppress PLR0913 on inherited signatures instead of bumping budget The overrides keep their full base-class signatures for MyPy compatibility, but those signatures carry more than five parameters, which tripped PLR0913 on each subclass redeclaration. Since the arity is dictated by the base class and cannot be reduced, decorate the overrides with typing_extensions.override; ruff treats that as the intended signal that the parameter count is not under the author's control and skips PLR0913. This restores the PLR0913 baseline to 1813. * fix(lint): add @override to modelscope image generation overrides Apply the same typing_extensions.override treatment to the image generation config so its inherited-signature overrides do not count against PLR0913. --------- Co-authored-by: Joel Tony <github@jaytau.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: hcl <chenglunhu@gmail.com> Co-authored-by: ztko <96878659+koztkozt@users.noreply.github.com> Co-authored-by: Nahrin <nahrin@nahrinoda.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Humphrey <a739376838@gmail.com> Co-authored-by: kursadlacin <kursadlacin@gmail.com> Co-authored-by: kursad <kursad.lacin@brado.net> Co-authored-by: Dushyant Acharya <dushyantacharya873@gmail.com> Co-authored-by: Yuriy <yuriy.shuyskiy@gmail.com> Co-authored-by: Recep S <22618852+us@users.noreply.github.com> Co-authored-by: Moshe Malawach <moshe.malawach@protonmail.com> Co-authored-by: Moshe Malawach <moshemalawach@users.noreply.github.com> Co-authored-by: Rongkun Yan <2493404415@qq.com> Co-authored-by: Varshith <kvarshithgowda@gmail.com> Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
1449 lines
56 KiB
Python
1449 lines
56 KiB
Python
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
|
from litellm.types.utils import GuardrailTracingDetail
|
|
|
|
|
|
class TestCustomGuardrailDeploymentHook:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_pre_call_deployment_hook_no_guardrails(self):
|
|
"""Test that method returns kwargs unchanged when no guardrails are present"""
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Test with guardrails as None
|
|
kwargs = {
|
|
"messages": [{"role": "user", "content": "test message"}],
|
|
"model": "gpt-3.5-turbo",
|
|
"guardrails": None,
|
|
}
|
|
|
|
result = await custom_guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
assert result == kwargs
|
|
|
|
# Test with guardrails as non-list
|
|
kwargs["guardrails"] = "not_a_list"
|
|
|
|
result = await custom_guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
assert result == kwargs
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_pre_call_deployment_hook_with_guardrails_and_message_update(
|
|
self,
|
|
):
|
|
"""Test that method processes guardrails and updates messages when result contains messages"""
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Mock the async_pre_call_hook method
|
|
mock_result = {"messages": [{"role": "user", "content": "filtered message"}]}
|
|
custom_guardrail.async_pre_call_hook = AsyncMock(return_value=mock_result)
|
|
|
|
original_messages = [{"role": "user", "content": "original message"}]
|
|
kwargs = {
|
|
"messages": original_messages,
|
|
"model": "gpt-3.5-turbo",
|
|
"guardrails": ["some_guardrail"],
|
|
"user_api_key_user_id": "test_user",
|
|
"user_api_key_team_id": "test_team",
|
|
"user_api_key_end_user_id": "test_end_user",
|
|
"user_api_key_hash": "test_hash",
|
|
"user_api_key_request_route": "test_route",
|
|
}
|
|
|
|
result = await custom_guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
# Verify async_pre_call_hook was called with correct parameters
|
|
custom_guardrail.async_pre_call_hook.assert_called_once()
|
|
call_args = custom_guardrail.async_pre_call_hook.call_args
|
|
|
|
# Check that UserAPIKeyAuth was created properly
|
|
user_api_key_dict = call_args[1]["user_api_key_dict"]
|
|
assert isinstance(user_api_key_dict, UserAPIKeyAuth)
|
|
assert user_api_key_dict.user_id == "test_user"
|
|
assert user_api_key_dict.team_id == "test_team"
|
|
assert user_api_key_dict.end_user_id == "test_end_user"
|
|
assert user_api_key_dict.api_key == "test_hash"
|
|
assert user_api_key_dict.request_route == "test_route"
|
|
|
|
# Check other parameters
|
|
assert call_args[1]["data"] == kwargs
|
|
assert call_args[1]["call_type"] == "completion"
|
|
|
|
# Verify messages were updated in result
|
|
assert result["messages"] == mock_result["messages"]
|
|
assert result["messages"] != original_messages
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deployment_hook_skips_when_pre_call_already_ran(self):
|
|
"""The deployment hook must not re-run async_pre_call_hook once the proxy
|
|
pre-call loop has already run it for this request."""
|
|
|
|
class CountingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(guardrail_name="g1", default_on=True)
|
|
self.pre_call_count = 0
|
|
|
|
async def async_pre_call_hook(
|
|
self, user_api_key_dict, cache, data, call_type
|
|
):
|
|
self.pre_call_count += 1
|
|
return data
|
|
|
|
guardrail = CountingGuardrail()
|
|
kwargs = {
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"model": "gpt-3.5-turbo",
|
|
"guardrails": ["g1"],
|
|
"metadata": {},
|
|
}
|
|
|
|
guardrail.mark_pre_call_hook_ran(kwargs)
|
|
await guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
assert guardrail.pre_call_count == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deployment_hook_runs_when_not_marked(self):
|
|
"""Without the proxy marker (direct-SDK usage) the deployment hook is the
|
|
only execution path and must still run the guardrail exactly once."""
|
|
|
|
class CountingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(guardrail_name="g1", default_on=True)
|
|
self.pre_call_count = 0
|
|
|
|
async def async_pre_call_hook(
|
|
self, user_api_key_dict, cache, data, call_type
|
|
):
|
|
self.pre_call_count += 1
|
|
return data
|
|
|
|
guardrail = CountingGuardrail()
|
|
kwargs = {
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"model": "gpt-3.5-turbo",
|
|
"guardrails": ["g1"],
|
|
"metadata": {},
|
|
}
|
|
|
|
await guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
assert guardrail.pre_call_count == 1
|
|
|
|
def test_mark_pre_call_hook_ran_uses_litellm_metadata(self):
|
|
"""The marker is recorded in litellm_metadata when that is the metadata
|
|
bucket in use, and is then visible to the skip check."""
|
|
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
|
|
|
|
guardrail = CustomGuardrail(guardrail_name="g1")
|
|
kwargs = {"litellm_metadata": {}}
|
|
|
|
guardrail.mark_pre_call_hook_ran(kwargs)
|
|
|
|
assert kwargs["litellm_metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY]
|
|
assert guardrail._pre_call_hook_already_ran(kwargs) is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deployment_hook_ignores_forged_caller_marker(self):
|
|
"""A direct-SDK caller controls request metadata but cannot know the
|
|
per-process token, so a hand-crafted marker must not suppress a
|
|
requested guardrail in async_pre_call_deployment_hook."""
|
|
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
|
|
|
|
class CountingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(guardrail_name="g1", default_on=True)
|
|
self.pre_call_count = 0
|
|
|
|
async def async_pre_call_hook(
|
|
self, user_api_key_dict, cache, data, call_type
|
|
):
|
|
self.pre_call_count += 1
|
|
return data
|
|
|
|
guardrail = CountingGuardrail()
|
|
kwargs = {
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"model": "gpt-3.5-turbo",
|
|
"guardrails": ["g1"],
|
|
"metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]},
|
|
}
|
|
|
|
await guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
assert guardrail.pre_call_count == 1
|
|
|
|
|
|
class TestCustomGuardrailShouldRunGuardrail:
|
|
|
|
def test_should_run_guardrail_with_litellm_metadata(self):
|
|
"""Test that should_run_guardrail works with litellm_metadata pattern"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test with guardrails in litellm_metadata
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"litellm_metadata": {"guardrails": ["test_guardrail"]},
|
|
}
|
|
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
|
|
assert result is True
|
|
|
|
def test_should_run_guardrail_with_metadata(self):
|
|
"""Test that should_run_guardrail works with metadata pattern"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test with guardrails in metadata
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"metadata": {"guardrails": ["test_guardrail"]},
|
|
}
|
|
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
|
|
assert result is True
|
|
|
|
def test_should_run_guardrail_with_root_level_guardrails(self):
|
|
"""Test that should_run_guardrail works with root level guardrails"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test with guardrails at root level
|
|
data = {"model": "gpt-3.5-turbo", "guardrails": ["test_guardrail"]}
|
|
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
|
|
assert result is True
|
|
|
|
def test_should_run_guardrail_no_matching_guardrail(self):
|
|
"""Test that should_run_guardrail returns False when guardrail name doesn't match"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test with different guardrail name
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"litellm_metadata": {"guardrails": ["different_guardrail"]},
|
|
}
|
|
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
|
|
assert result is False
|
|
|
|
def test_should_run_guardrail_with_disable_global_guardrail(self):
|
|
"""Test that disable_global_guardrails only works from admin metadata"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="global_guardrail",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test 1: Global guardrail runs by default
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert result is True, "Global guardrail should run when default_on=True"
|
|
|
|
# Test 2: User-injected disable at root level is IGNORED
|
|
data_with_disable_root = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"disable_global_guardrails": True,
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert (
|
|
result is True
|
|
), "User-injected disable_global_guardrails should be ignored"
|
|
|
|
# Test 3: User-injected disable in metadata is IGNORED
|
|
data_with_disable_metadata = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"disable_global_guardrails": True},
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_with_disable_metadata, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert result is True, "User-injected metadata disable should be ignored"
|
|
|
|
# Test 4: Admin-configured disable via user_api_key_metadata IS respected
|
|
data_with_admin_disable = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_with_admin_disable, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert result is False, "Admin-configured disable should be respected"
|
|
|
|
# Test 5: Admin config in metadata isn't shadowed by user-supplied litellm_metadata
|
|
data_cross_key = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
|
|
"litellm_metadata": {"request_tags": ["user-supplied"]},
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_cross_key, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert (
|
|
result is False
|
|
), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata"
|
|
|
|
# Test 6: After the pre-call strip runs, user-injected
|
|
# user_api_key_metadata in the non-authoritative metadata key is gone.
|
|
# _get_admin_metadata must then surface admin config unchanged.
|
|
data_post_strip = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
|
|
"litellm_metadata": {}, # post-strip: attacker payload removed
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_post_strip, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert (
|
|
result is False
|
|
), "Admin config in metadata must be respected when other metadata key is empty"
|
|
|
|
def test_should_run_guardrail_key_disable_global_not_overruled_by_team_guardrail_list(
|
|
self,
|
|
):
|
|
"""Key disable_global_guardrails must take precedence over the guardrail
|
|
appearing in the team's explicit guardrails list."""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="global_guardrail",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Key disabled globals; team added the same guardrail to its explicit list
|
|
# (simulates what _add_guardrails_from_key_or_team_metadata produces).
|
|
data_key_disabled_team_listed = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {
|
|
"user_api_key_metadata": {"disable_global_guardrails": True},
|
|
"guardrails": ["global_guardrail"],
|
|
},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_key_disabled_team_listed,
|
|
event_type=GuardrailEventHooks.pre_call,
|
|
)
|
|
is False
|
|
), "Key disable_global_guardrails must win over team's explicit guardrail list"
|
|
|
|
# Complementary: key NOT disabled, team added guardrail → should run
|
|
data_key_enabled_team_listed = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {
|
|
"user_api_key_metadata": {},
|
|
"guardrails": ["global_guardrail"],
|
|
},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_key_enabled_team_listed,
|
|
event_type=GuardrailEventHooks.pre_call,
|
|
)
|
|
is True
|
|
), "Guardrail in team's explicit list should run when key has not disabled globals"
|
|
|
|
def test_should_run_guardrail_with_opted_out_global_guardrails(self):
|
|
"""Test that per-guardrail opt-out only works from admin metadata"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="global_guardrail",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test 1: User-injected opt-out at root level is IGNORED
|
|
data_root = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"opted_out_global_guardrails": ["global_guardrail"],
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_root, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
# Test 2: User-injected opt-out in metadata is IGNORED
|
|
data_metadata = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"opted_out_global_guardrails": ["global_guardrail"]},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_metadata, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
# Test 4: a different guardrail in the opt-out list → still runs
|
|
data_other = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"opted_out_global_guardrails": ["some_other_guardrail"]},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_other, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
# Test 5: empty opt-out list → still runs
|
|
data_empty = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"opted_out_global_guardrails": []},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_empty, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
# Test 6: malformed value (bool instead of list) → safely ignored, guardrail runs
|
|
data_malformed = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"opted_out_global_guardrails": True},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_malformed, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
def test_should_run_guardrail_opt_out_does_not_affect_non_global(self):
|
|
"""Opt-out list only matters for default_on=True guardrails"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
non_global = CustomGuardrail(
|
|
guardrail_name="opt_in_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# An opt-in guardrail named in opted_out_global_guardrails is still controlled
|
|
# by the explicit `guardrails` request list, not by the global opt-out list.
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {
|
|
"opted_out_global_guardrails": ["opt_in_guardrail"],
|
|
"guardrails": ["opt_in_guardrail"],
|
|
},
|
|
}
|
|
assert (
|
|
non_global.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
|
|
class TestApplyGuardrailCheck:
|
|
def test_apply_guardrail_check_only_on_direct_implementation(self):
|
|
"""
|
|
Test that "apply_guardrail" in type(callback).__dict__ only returns True
|
|
when the object's own class implements the method, not when it's inherited
|
|
from a parent class.
|
|
|
|
This is critical for properly routing guardrail handling to the unified
|
|
guardrail handler vs the guardrail's own implementation.
|
|
"""
|
|
|
|
# Parent class with apply_guardrail (CustomGuardrail already has it)
|
|
class ParentGuardrail(CustomGuardrail):
|
|
"""Parent that inherits apply_guardrail from CustomGuardrail"""
|
|
|
|
pass
|
|
|
|
# Child class that only inherits apply_guardrail (doesn't override)
|
|
class ChildGuardrailWithoutOverride(ParentGuardrail):
|
|
"""Child that only inherits apply_guardrail"""
|
|
|
|
pass
|
|
|
|
# Child class that overrides apply_guardrail
|
|
class ChildGuardrailWithOverride(ParentGuardrail):
|
|
"""Child that overrides apply_guardrail"""
|
|
|
|
async def apply_guardrail(self, text, language=None, entities=None):
|
|
return f"modified: {text}"
|
|
|
|
# Instantiate the classes
|
|
parent_instance = ParentGuardrail()
|
|
child_without_override = ChildGuardrailWithoutOverride()
|
|
child_with_override = ChildGuardrailWithOverride()
|
|
|
|
# Test: CustomGuardrail itself has apply_guardrail in its __dict__
|
|
assert (
|
|
"apply_guardrail" in type(CustomGuardrail()).__dict__
|
|
), "CustomGuardrail should have apply_guardrail in its own __dict__"
|
|
|
|
# Test: ParentGuardrail inherits but doesn't override, so it should NOT be in __dict__
|
|
assert (
|
|
"apply_guardrail" not in type(parent_instance).__dict__
|
|
), "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)"
|
|
|
|
# Test: ChildGuardrailWithoutOverride only inherits, should NOT be in __dict__
|
|
assert (
|
|
"apply_guardrail" not in type(child_without_override).__dict__
|
|
), "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)"
|
|
|
|
# Test: ChildGuardrailWithOverride overrides the method, SHOULD be in __dict__
|
|
assert (
|
|
"apply_guardrail" in type(child_with_override).__dict__
|
|
), "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)"
|
|
|
|
# Verify that all instances still have the method via inheritance (hasattr)
|
|
assert hasattr(
|
|
parent_instance, "apply_guardrail"
|
|
), "All instances should have apply_guardrail via inheritance"
|
|
assert hasattr(
|
|
child_without_override, "apply_guardrail"
|
|
), "All instances should have apply_guardrail via inheritance"
|
|
assert hasattr(
|
|
child_with_override, "apply_guardrail"
|
|
), "All instances should have apply_guardrail via inheritance"
|
|
|
|
|
|
class TestGuardrailLoggingAggregation:
|
|
def _make_guardrail(self):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
return CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
def _invoke_add_log(self, request_data: dict) -> None:
|
|
guardrail = self._make_guardrail()
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
start_time=1.0,
|
|
end_time=2.0,
|
|
duration=1.0,
|
|
masked_entity_count={"EMAIL": 1},
|
|
guardrail_provider="presidio",
|
|
)
|
|
|
|
def test_appends_to_existing_metadata_list(self):
|
|
request_data = {
|
|
"metadata": {
|
|
"standard_logging_guardrail_information": [
|
|
{"guardrail_name": "existing_guardrail"}
|
|
]
|
|
}
|
|
}
|
|
|
|
self._invoke_add_log(request_data)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert isinstance(info, list)
|
|
assert len(info) == 2
|
|
assert info[0]["guardrail_name"] == "existing_guardrail"
|
|
assert info[1]["guardrail_name"] == "test_guardrail"
|
|
|
|
def test_converts_existing_metadata_dict_to_list(self):
|
|
request_data = {
|
|
"metadata": {
|
|
"standard_logging_guardrail_information": {"guardrail_name": "legacy"}
|
|
}
|
|
}
|
|
|
|
self._invoke_add_log(request_data)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert isinstance(info, list)
|
|
assert len(info) == 2
|
|
assert info[0]["guardrail_name"] == "legacy"
|
|
assert info[1]["guardrail_name"] == "test_guardrail"
|
|
|
|
def test_appends_to_litellm_metadata(self):
|
|
request_data = {
|
|
"litellm_metadata": {
|
|
"standard_logging_guardrail_information": [
|
|
{"guardrail_name": "litellm_existing"}
|
|
]
|
|
}
|
|
}
|
|
|
|
self._invoke_add_log(request_data)
|
|
|
|
info = request_data["litellm_metadata"][
|
|
"standard_logging_guardrail_information"
|
|
]
|
|
assert isinstance(info, list)
|
|
assert len(info) == 2
|
|
assert info[1]["guardrail_name"] == "test_guardrail"
|
|
|
|
|
|
class TestGuardrailOtelSpanEmission:
|
|
"""Recording a guardrail emits its otel span inline, so every guardrail
|
|
execution produces a span — including the pass-through allow path that never
|
|
reaches a post-call hook."""
|
|
|
|
def _make_guardrail(self):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
return CustomGuardrail(
|
|
guardrail_name="emit_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
def _record(self, guardrail, request_data):
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
start_time=1.0,
|
|
end_time=2.0,
|
|
duration=1.0,
|
|
)
|
|
|
|
def test_emits_span_for_recorded_entry(self, monkeypatch):
|
|
captured = []
|
|
monkeypatch.setattr(
|
|
"litellm.integrations.otel.logger.emit_guardrail_span",
|
|
captured.append,
|
|
)
|
|
|
|
request_data = {"metadata": {}}
|
|
self._record(self._make_guardrail(), request_data)
|
|
|
|
assert len(captured) == 1
|
|
emitted = captured[0]
|
|
recorded = request_data["metadata"]["standard_logging_guardrail_information"][
|
|
-1
|
|
]
|
|
assert emitted is recorded
|
|
assert emitted["guardrail_name"] == "emit_guard"
|
|
assert emitted["start_time"] == 1.0
|
|
assert emitted["end_time"] == 2.0
|
|
|
|
def test_span_emission_failure_does_not_break_recording(self, monkeypatch):
|
|
def _boom(_entry):
|
|
raise RuntimeError("otel exporter down")
|
|
|
|
monkeypatch.setattr(
|
|
"litellm.integrations.otel.logger.emit_guardrail_span", _boom
|
|
)
|
|
|
|
request_data = {"metadata": {}}
|
|
self._record(self._make_guardrail(), request_data)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(info) == 1
|
|
assert info[0]["guardrail_name"] == "emit_guard"
|
|
|
|
|
|
class TestGuardrailSensitiveFieldStripping:
|
|
"""Tests that secret_fields is stripped from guardrail responses before logging.
|
|
|
|
Matches the pattern used by Langfuse and Arize integrations which also
|
|
pop("secret_fields") to prevent raw Authorization headers from being persisted.
|
|
"""
|
|
|
|
def _make_guardrail(self):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
return CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
def test_secret_fields_stripped_from_guardrail_response(self):
|
|
"""Ensure secret_fields (containing raw Authorization headers) is not persisted."""
|
|
guardrail = self._make_guardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
guardrail_response_with_secrets = {
|
|
"model": "gpt-4",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"secret_fields": {
|
|
"raw_headers": {
|
|
"authorization": "Bearer sk-live-secret-key-12345",
|
|
"content-type": "application/json",
|
|
}
|
|
},
|
|
"proxy_server_request": {"url": "http://localhost:4000/chat/completions"},
|
|
}
|
|
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=guardrail_response_with_secrets,
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
duration=1.0,
|
|
)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(info) == 1
|
|
logged_response = info[0]["guardrail_response"]
|
|
|
|
# secret_fields must be stripped
|
|
assert "secret_fields" not in logged_response
|
|
|
|
# Other fields should be preserved
|
|
assert "model" in logged_response
|
|
assert "messages" in logged_response
|
|
assert "proxy_server_request" in logged_response
|
|
|
|
def test_string_guardrail_response_not_affected(self):
|
|
"""String responses (e.g. 'allow', 'deny') should pass through unchanged."""
|
|
guardrail = self._make_guardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="allow",
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
duration=0.5,
|
|
)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert info[0]["guardrail_response"] == "allow"
|
|
|
|
def test_no_authorization_header_in_logged_response(self):
|
|
"""Verify no plaintext Authorization header ends up in the logged guardrail response."""
|
|
import json
|
|
|
|
guardrail = self._make_guardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={
|
|
"model": "gpt-4",
|
|
"secret_fields": {
|
|
"raw_headers": {
|
|
"authorization": "Bearer sk-live-SHOULD-NOT-APPEAR",
|
|
}
|
|
},
|
|
},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
duration=1.0,
|
|
)
|
|
|
|
logged_response = request_data["metadata"][
|
|
"standard_logging_guardrail_information"
|
|
][0]["guardrail_response"]
|
|
assert "secret_fields" not in logged_response
|
|
assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response)
|
|
|
|
def test_secret_fields_stripped_from_list_dict_response(self):
|
|
"""Ensure secret_fields is stripped from List[dict] guardrail responses too."""
|
|
guardrail = self._make_guardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=[
|
|
{
|
|
"result": "ok",
|
|
"secret_fields": {
|
|
"raw_headers": {"authorization": "Bearer sk-secret"}
|
|
},
|
|
},
|
|
{"result": "also_ok"},
|
|
],
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
duration=1.0,
|
|
)
|
|
|
|
import json
|
|
|
|
serialized = json.dumps(request_data)
|
|
assert "secret_fields" not in serialized
|
|
assert "sk-secret" not in serialized
|
|
|
|
|
|
class TestCustomGuardrailPassthroughSupport:
|
|
"""Tests for passthrough endpoint guardrail support - Issue fixes."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_post_call_success_deployment_hook_with_httpx_response(self):
|
|
"""
|
|
Test that async_post_call_success_deployment_hook handles raw httpx.Response objects
|
|
from passthrough endpoints without crashing with TypeError.
|
|
|
|
This tests Fix #3: TypeError: TypedDict does not support instance and class checks
|
|
"""
|
|
import httpx
|
|
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Mock the async_post_call_success_hook to return None (guardrail didn't modify response)
|
|
custom_guardrail.async_post_call_success_hook = AsyncMock(return_value=None)
|
|
|
|
# Create a mock httpx.Response object (typical passthrough response)
|
|
mock_response = AsyncMock(spec=httpx.Response)
|
|
mock_response.status_code = 200
|
|
mock_response.text = "Mock response"
|
|
|
|
request_data = {
|
|
"guardrails": ["test_guardrail"],
|
|
"user_api_key_user_id": "test_user",
|
|
"user_api_key_team_id": "test_team",
|
|
"user_api_key_end_user_id": "test_end_user",
|
|
"user_api_key_hash": "test_hash",
|
|
"user_api_key_request_route": "passthrough_route",
|
|
}
|
|
|
|
# This should not raise TypeError: TypedDict does not support instance and class checks
|
|
result = await custom_guardrail.async_post_call_success_deployment_hook(
|
|
request_data=request_data,
|
|
response=mock_response,
|
|
call_type=CallTypes.allm_passthrough_route,
|
|
)
|
|
|
|
# When result is None, should return the original response
|
|
assert result == mock_response
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_post_call_success_deployment_hook_with_none_call_type(self):
|
|
"""
|
|
Test that async_post_call_success_deployment_hook handles None call_type gracefully.
|
|
|
|
This ensures that even if call_type is None (before fix #1), the guardrail doesn't crash.
|
|
"""
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Mock the async_post_call_success_hook to return None
|
|
custom_guardrail.async_post_call_success_hook = AsyncMock(return_value=None)
|
|
|
|
mock_response = AsyncMock()
|
|
|
|
request_data = {
|
|
"guardrails": ["test_guardrail"],
|
|
"user_api_key_user_id": "test_user",
|
|
}
|
|
|
|
# Call with None call_type - should not crash
|
|
result = await custom_guardrail.async_post_call_success_deployment_hook(
|
|
request_data=request_data,
|
|
response=mock_response,
|
|
call_type=None,
|
|
)
|
|
|
|
# Should return the original response when result is None
|
|
assert result == mock_response
|
|
|
|
def test_is_valid_response_type_with_none(self):
|
|
"""
|
|
Test _is_valid_response_type helper method correctly identifies None as invalid.
|
|
|
|
This is part of Fix #3: Safely handling TypedDict types that don't support isinstance checks.
|
|
"""
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# None should be invalid
|
|
assert custom_guardrail._is_valid_response_type(None) is False
|
|
|
|
def test_is_valid_response_type_with_typeddict_error(self):
|
|
"""
|
|
Test _is_valid_response_type gracefully handles TypeError from TypedDict.
|
|
|
|
This tests Fix #3: When isinstance() is called with TypedDict types, it raises TypeError.
|
|
The method should catch this and allow the response through.
|
|
"""
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Create a valid LiteLLM response object
|
|
response = ModelResponse(
|
|
id="test-id",
|
|
choices=[],
|
|
created=0,
|
|
model="test-model",
|
|
object="chat.completion",
|
|
)
|
|
|
|
# This should return True (it's a valid response type or TypeError is caught)
|
|
result = custom_guardrail._is_valid_response_type(response)
|
|
assert result is True
|
|
|
|
|
|
class TestEventTypeLogging:
|
|
"""Tests for event_type logging in guardrail information."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_infers_event_type_from_async_pre_call_hook(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.pre_call
|
|
from async_pre_call_hook function name.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=[
|
|
GuardrailEventHooks.pre_call,
|
|
GuardrailEventHooks.post_call,
|
|
],
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_pre_call_hook(self, data: dict, **kwargs):
|
|
return {"result": "pre_call_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.async_pre_call_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode was set to pre_call (not the full list)
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_infers_event_type_from_async_post_call_success_hook(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
|
|
from async_post_call_success_hook function name.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=[
|
|
GuardrailEventHooks.pre_call,
|
|
GuardrailEventHooks.post_call,
|
|
],
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_post_call_success_hook(self, data: dict, **kwargs):
|
|
return {"result": "post_call_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.async_post_call_success_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode was set to post_call (not the full list)
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_infers_event_type_from_async_moderation_hook(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.during_call
|
|
from async_moderation_hook function name.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=[
|
|
GuardrailEventHooks.during_call,
|
|
GuardrailEventHooks.post_call,
|
|
],
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_moderation_hook(self, data: dict, **kwargs):
|
|
return {"result": "moderation_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.async_moderation_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode was set to during_call (not the full list)
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.during_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_infers_event_type_from_async_post_call_streaming_hook(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
|
|
from async_post_call_streaming_hook function name.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=[
|
|
GuardrailEventHooks.pre_call,
|
|
GuardrailEventHooks.post_call,
|
|
],
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_post_call_streaming_hook(self, data: dict, **kwargs):
|
|
return {"result": "streaming_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.async_post_call_streaming_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode was set to post_call (not the full list)
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_returns_none_for_unknown_function_name(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator returns None for event_type
|
|
when function name doesn't match known patterns, and falls back to self.event_hook.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def some_other_hook(self, data: dict, **kwargs):
|
|
return {"result": "other_hook_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.some_other_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode falls back to self.event_hook
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
|
|
|
|
def test_add_standard_logging_uses_event_type_over_event_hook(self):
|
|
"""
|
|
Test that add_standard_logging_guardrail_information_to_request_data
|
|
prioritizes event_type parameter over self.event_hook.
|
|
"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
|
|
)
|
|
|
|
request_data = {"metadata": {}}
|
|
|
|
# Call with explicit event_type
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
event_type=GuardrailEventHooks.post_call,
|
|
)
|
|
|
|
# Should use the provided event_type (post_call), not the full event_hook list
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_skips_auto_record_if_function_already_recorded(
|
|
self,
|
|
):
|
|
"""When a wrapped guardrail function records its own entry directly
|
|
(e.g. block_code_execution.apply_guardrail records a rich
|
|
``[detections...]`` payload), the decorator must NOT also append its
|
|
own ``"allow"``/raw-response entry — otherwise every backend
|
|
(OTEL spans, Datadog, Langfuse, spend logs) double-records one
|
|
logical guardrail invocation."""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="block-code",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def apply_guardrail(self, inputs, request_data, **kwargs):
|
|
self.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=[{"action_taken": "block"}],
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
event_type=GuardrailEventHooks.pre_call,
|
|
)
|
|
return inputs
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.apply_guardrail(
|
|
inputs={"texts": ["x"]}, request_data=request_data
|
|
)
|
|
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1, (
|
|
f"Decorator must not double-record when the wrapped function "
|
|
f"already appended its own entry; got {len(logged_info)} entries"
|
|
)
|
|
assert logged_info[0]["guardrail_response"] == [{"action_taken": "block"}]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_skips_auto_record_on_exception_if_function_already_recorded(
|
|
self,
|
|
):
|
|
"""Same as above on the failure path: if the wrapped function
|
|
appended an entry in its ``finally`` block before re-raising, the
|
|
decorator must just re-raise without auto-recording on top."""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="block-code",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def apply_guardrail(self, inputs, request_data, **kwargs):
|
|
try:
|
|
raise ValueError("blocked")
|
|
finally:
|
|
self.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=[{"action_taken": "block"}],
|
|
request_data=request_data,
|
|
guardrail_status="guardrail_intervened",
|
|
event_type=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
with pytest.raises(ValueError, match="blocked"):
|
|
await guardrail.apply_guardrail(
|
|
inputs={"texts": ["x"]}, request_data=request_data
|
|
)
|
|
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_status"] == "guardrail_intervened"
|
|
|
|
def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
|
|
self,
|
|
):
|
|
"""
|
|
Test that add_standard_logging_guardrail_information_to_request_data
|
|
falls back to self.event_hook when event_type is None.
|
|
"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
request_data = {"metadata": {}}
|
|
|
|
# Call with event_type=None
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
event_type=None,
|
|
)
|
|
|
|
# Should fall back to self.event_hook
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
|
|
|
|
|
|
class TestTracingFieldsPopulation:
|
|
"""Verify add_standard_logging_guardrail_information_to_request_data passes tracing_detail fields."""
|
|
|
|
def test_new_fields_set_on_slg(self):
|
|
cg = CustomGuardrail(guardrail_name="test-rail")
|
|
request_data = {"metadata": {}}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
tracing_detail=GuardrailTracingDetail(
|
|
guardrail_id="rail-123",
|
|
policy_template="EU AI Act Article 5",
|
|
detection_method="regex",
|
|
confidence_score=0.95,
|
|
match_details=[{"type": "pattern", "action_taken": "BLOCK"}],
|
|
patterns_checked=12,
|
|
alert_recipients=["admin@example.com"],
|
|
),
|
|
)
|
|
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(slg_list) == 1
|
|
slg = slg_list[0]
|
|
assert slg["guardrail_id"] == "rail-123"
|
|
assert slg["policy_template"] == "EU AI Act Article 5"
|
|
assert slg["detection_method"] == "regex"
|
|
assert slg["confidence_score"] == 0.95
|
|
assert slg["patterns_checked"] == 12
|
|
assert slg["alert_recipients"] == ["admin@example.com"]
|
|
assert len(slg["match_details"]) == 1
|
|
|
|
def test_new_fields_default_to_absent(self):
|
|
"""When tracing_detail is not passed, new fields are absent from the SLG dict."""
|
|
cg = CustomGuardrail(guardrail_name="test-rail")
|
|
request_data = {"metadata": {}}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="ok",
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
)
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert slg.get("guardrail_id") is None
|
|
assert slg.get("policy_template") is None
|
|
assert slg.get("confidence_score") is None
|
|
|
|
def test_multiple_guardrails_with_different_policies(self):
|
|
"""One request, multiple guardrails each with own policy_template."""
|
|
cg1 = CustomGuardrail(guardrail_name="rail-1")
|
|
cg2 = CustomGuardrail(guardrail_name="rail-2")
|
|
request_data = {"metadata": {}}
|
|
|
|
cg1.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="ok",
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
tracing_detail=GuardrailTracingDetail(policy_template="GDPR"),
|
|
)
|
|
cg2.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="blocked",
|
|
request_data=request_data,
|
|
guardrail_status="guardrail_intervened",
|
|
tracing_detail=GuardrailTracingDetail(
|
|
policy_template="EU AI Act Article 5"
|
|
),
|
|
)
|
|
|
|
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(slg_list) == 2
|
|
assert slg_list[0]["policy_template"] == "GDPR"
|
|
assert slg_list[1]["policy_template"] == "EU AI Act Article 5"
|
|
|
|
def test_classification_field_passed_through(self):
|
|
"""Classification dict for LLM-judge guardrails is passed through."""
|
|
cg = CustomGuardrail(guardrail_name="judge-rail")
|
|
request_data = {"metadata": {}}
|
|
classification = {
|
|
"flagged": True,
|
|
"category": "workplace_emotion_recognition",
|
|
"article_reference": "Article 5(1)(f)",
|
|
"confidence": 0.94,
|
|
"reason": "Request asks to analyze employee sentiment",
|
|
}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="blocked",
|
|
request_data=request_data,
|
|
guardrail_status="guardrail_intervened",
|
|
tracing_detail=GuardrailTracingDetail(
|
|
classification=classification,
|
|
detection_method="llm-judge",
|
|
confidence_score=0.94,
|
|
),
|
|
)
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert slg["classification"] == classification
|
|
assert slg["detection_method"] == "llm-judge"
|
|
assert slg["confidence_score"] == 0.94
|
|
|
|
|
|
class TestCustomGuardrailSpendLogMatchRedaction:
|
|
"""Guardrail JSON persisted via standard_logging must not contain raw match spans."""
|
|
|
|
def test_add_standard_logging_redacts_nested_match(self):
|
|
cg = CustomGuardrail(guardrail_name="test-rail")
|
|
raw = {
|
|
"assessments": [
|
|
{
|
|
"sensitiveInformationPolicy": {
|
|
"piiEntities": [
|
|
{"type": "NAME", "match": "GG", "action": "BLOCKED"}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
request_data: dict = {"metadata": {}}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=raw,
|
|
request_data=request_data,
|
|
guardrail_status="guardrail_intervened",
|
|
)
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert (
|
|
slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][
|
|
"piiEntities"
|
|
][0]["match"]
|
|
== "[REDACTED]"
|
|
)
|
|
assert (
|
|
raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
|
"match"
|
|
]
|
|
== "GG"
|
|
)
|
|
|
|
def test_add_standard_logging_redacts_regex_field(self):
|
|
cg = CustomGuardrail(guardrail_name="test-rail")
|
|
raw = {"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}
|
|
request_data: dict = {"metadata": {}}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=raw,
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
)
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert slg["guardrail_response"]["filters"][0]["regex"] == "[REDACTED]"
|
|
assert raw["filters"][0]["regex"] == r"\d{3}-\d{2}-\d{4}"
|
|
|
|
|
|
class TestGuardrailInterventionClassification:
|
|
"""A routing decision is a deliberate guardrail intervention, not a failure."""
|
|
|
|
def test_sensitive_data_route_exception_is_intervention(self):
|
|
from litellm.exceptions import SensitiveDataRouteException
|
|
|
|
exc = SensitiveDataRouteException(
|
|
route_to_model="on-prem-model",
|
|
session_id="sess-1",
|
|
guardrail_name="pii-rail",
|
|
)
|
|
assert CustomGuardrail._is_guardrail_intervention(exc) is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_routing_logged_as_intervened_not_failed(self):
|
|
from litellm.exceptions import SensitiveDataRouteException
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class RoutingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="pii-rail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_pre_call_hook(self, data, **kwargs):
|
|
raise SensitiveDataRouteException(
|
|
route_to_model="on-prem-model",
|
|
session_id="sess-1",
|
|
guardrail_name=self.guardrail_name,
|
|
)
|
|
|
|
guardrail = RoutingGuardrail()
|
|
request_data: dict = {"metadata": {}}
|
|
|
|
with pytest.raises(SensitiveDataRouteException):
|
|
await guardrail.async_pre_call_hook(data=request_data)
|
|
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert slg["guardrail_status"] == "guardrail_intervened"
|