mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
test(proxy): cover Bedrock spend logging_event_type and during_call routing
- Assert make_bedrock_api_request forwards logging_event_type to standard logging and legacy INPUT maps to pre_call when omitted. - Assert during_call_hook invokes Bedrock async_moderation_hook when native path is used. Made-with: Cursor
This commit is contained in:
parent
840fdfa915
commit
19f34b4b9d
1 changed files with 94 additions and 0 deletions
|
|
@ -6,16 +6,20 @@ import os
|
|||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../../.."))
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
||||
BedrockGuardrail,
|
||||
_redact_pii_matches,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1209,6 +1213,96 @@ def test_bedrock_guardrail_uses_native_during_call_hook():
|
|||
assert BedrockGuardrail.use_native_during_call_hook is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_make_bedrock_api_request_logging_event_type_for_spend_logs():
|
||||
"""
|
||||
Spend/UI use event_type from the proxy hook, not Bedrock's INPUT/OUTPUT alone.
|
||||
When logging_event_type is set, it must be forwarded to standard guardrail logging.
|
||||
When omitted, INPUT maps to pre_call (legacy).
|
||||
"""
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
|
||||
)
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "test-access-key"
|
||||
mock_credentials.secret_key = "test-secret-key"
|
||||
mock_credentials.token = None
|
||||
|
||||
mock_bedrock_response = MagicMock()
|
||||
mock_bedrock_response.status_code = 200
|
||||
mock_bedrock_response.json.return_value = {"action": "NONE", "assessments": []}
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler, "post", new_callable=AsyncMock
|
||||
) as mock_post, patch.object(
|
||||
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
|
||||
), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object(
|
||||
guardrail,
|
||||
"add_standard_logging_guardrail_information_to_request_data",
|
||||
) as mock_log:
|
||||
mock_post.return_value = mock_bedrock_response
|
||||
|
||||
await guardrail.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=request_data["messages"],
|
||||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.during_call,
|
||||
)
|
||||
assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call
|
||||
|
||||
mock_log.reset_mock()
|
||||
|
||||
await guardrail.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=request_data["messages"],
|
||||
request_data=request_data,
|
||||
)
|
||||
assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.pre_call
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_call_hook_invokes_bedrock_async_moderation_hook():
|
||||
"""
|
||||
Bedrock sets use_native_during_call_hook so ProxyLogging runs the real
|
||||
async_moderation_hook (unified apply_guardrail would log INPUT as pre_call).
|
||||
"""
|
||||
cache = DualCache()
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=cache)
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="bedrock-during-test",
|
||||
guardrailIdentifier="gid",
|
||||
guardrailVersion="1",
|
||||
event_hook=GuardrailEventHooks.during_call,
|
||||
default_on=True,
|
||||
)
|
||||
mock_mod = AsyncMock(return_value=None)
|
||||
guardrail.async_moderation_hook = mock_mod # type: ignore[method-assign]
|
||||
|
||||
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
|
||||
try:
|
||||
litellm.callbacks = [guardrail]
|
||||
await proxy_logging.during_call_hook(
|
||||
data={
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "test"}],
|
||||
},
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test_key", user_id="test_user"
|
||||
),
|
||||
call_type="completion",
|
||||
)
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
mock_mod.assert_awaited_once()
|
||||
|
||||
|
||||
def test_extract_blocked_assessments_pii_entity():
|
||||
"""L3: PII entity match (BLOCKED) is surfaced with category, type, and matched term."""
|
||||
g = _make_guardrail()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue