mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
test(guardrails): run bedrock assistant filter in CI
This commit is contained in:
parent
4ce773feee
commit
ec5931df5b
2 changed files with 44 additions and 22 deletions
|
|
@ -0,0 +1,44 @@
|
|||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.bedrock.passthrough.guardrail_translation.handler import (
|
||||
BedrockPassthroughGuardrailHandler,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_assistant_preserves_converse_history_and_masks_user() -> None:
|
||||
assistant = {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"text": "old reply"},
|
||||
{
|
||||
"toolUse": {
|
||||
"toolUseId": "t1",
|
||||
"name": "search",
|
||||
"input": {"q": "old"},
|
||||
}
|
||||
},
|
||||
],
|
||||
}
|
||||
data = {
|
||||
"endpoint": "model/test/converse",
|
||||
"data": {
|
||||
"messages": [
|
||||
assistant,
|
||||
{"role": "user", "content": [{"text": "private"}]},
|
||||
]
|
||||
},
|
||||
}
|
||||
guardrail = MagicMock()
|
||||
guardrail.apply_guardrail = AsyncMock(return_value={"texts": ["[MASKED]"]})
|
||||
guardrail.skip_system_message_in_guardrail = False
|
||||
guardrail.skip_tool_message_in_guardrail = False
|
||||
guardrail.skip_assistant_message_in_guardrail = True
|
||||
|
||||
await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["private"]
|
||||
assert data["data"]["messages"][0] == assistant
|
||||
assert data["data"]["messages"][1]["content"][0]["text"] == "[MASKED]"
|
||||
|
|
@ -10,7 +10,6 @@ Validates that:
|
|||
"""
|
||||
|
||||
import copy
|
||||
from typing import Final
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -36,27 +35,6 @@ def _make_guardrail(apply_result: dict) -> MagicMock:
|
|||
return g
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_assistant_preserves_converse_history_and_masks_user() -> None:
|
||||
assistant: Final = {
|
||||
"role": "assistant",
|
||||
"content": [{"text": "old reply"}, {"toolUse": {"toolUseId": "t1", "name": "search", "input": {"q": "old"}}}],
|
||||
}
|
||||
data: Final = {
|
||||
"endpoint": "model/test/converse",
|
||||
"data": {"messages": [assistant, {"role": "user", "content": [{"text": "private"}]}]},
|
||||
}
|
||||
original_assistant: Final = copy.deepcopy(assistant)
|
||||
guardrail: Final = _make_guardrail({"texts": ["[MASKED]"]})
|
||||
guardrail.skip_assistant_message_in_guardrail = True
|
||||
|
||||
await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["private"]
|
||||
assert data["data"]["messages"][0] == original_assistant
|
||||
assert data["data"]["messages"][1]["content"][0]["text"] == "[MASKED]"
|
||||
|
||||
|
||||
def _converse_data(endpoint: str = "model/anthropic.claude-3-sonnet/converse") -> dict:
|
||||
return {
|
||||
"endpoint": endpoint,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue