scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-23 05:34:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 25

test(guardrails): cover mixed text and bytes embedding batches

Browse source
This commit is contained in:
zhengguangzhuo 2026-10-15 11:11:04 +08:00
parent 0757f96431
commit 8420aed84f
1 changed files with 10 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
tests/unit/integrations/test_custom_guardrail.py
View file

@ -56,14 +56,21 @@ class TestGetGuardrailsMessagesForCallType:
assert data == {"input": embedding_input}
@pytest.mark.parametrize("call_type", [CallTypes.embedding, CallTypes.aembedding])
def test_embedding_batch_with_non_text_values_is_not_extracted(self, call_type: CallTypes) -> None:
@pytest.mark.parametrize(
"embedding_input",
[[1, 2], ["first", b"second"]],
ids=["numbers", "mixed-text-bytes"],
)
def test_embedding_batch_with_non_text_values_is_not_extracted(
self, call_type: CallTypes, embedding_input: list[object]
) -> None:
guardrail: Final = CustomGuardrail(guardrail_name="embedding-extractor")
data: Final = {"input": [1, 2]}
data: Final = {"input": embedding_input}
messages: Final = guardrail.get_guardrails_messages_for_call_type(call_type=call_type, data=data)
assert messages is None
assert data == {"input": [1, 2]}
assert data == {"input": embedding_input}
@pytest.mark.parametrize("call_type", [CallTypes.embedding, CallTypes.aembedding])
@pytest.mark.parametrize("embedding_input", [None, 123, {"text": "hello"}], ids=["none", "number", "object"])