From 654bd3d0d6dcf2033bb6afd51132717733da23d8 Mon Sep 17 00:00:00 2001 From: zhengguangzhuo <100548634+zhengguangzhuo@users.noreply.github.com> Date: Thu, 8 Oct 2026 17:12:33 +0800 Subject: [PATCH 1/4] fix(guardrails): extract text from embedding inputs --- litellm/integrations/custom_guardrail.py | 40 ++++++++++++++---- .../integrations/test_custom_guardrail.py | 41 +++++++++++++++++++ 2 files changed, 72 insertions(+), 9 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 94f1207b9ae..b6b88ea7cc0 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -9,6 +9,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args import httpx +from pydantic import TypeAdapter, ValidationError from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -64,6 +65,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions" # field to suppress a guardrail on the direct-SDK path that never reaches the # proxy's metadata sanitizer. _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) +_EMBEDDING_INPUT_ADAPTER: Final = TypeAdapter[list[str]](list[str]) _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) @@ -1537,21 +1539,41 @@ class CustomGuardrail(CustomLogger): ): return data.get("messages") - ######################################################### - # /responses - # User/System messages are stored in the "input" key, use litellm transformation to get the messages - ######################################################### - if call_type == CallTypes.responses.value or call_type == CallTypes.aresponses.value: + if ( + call_type == CallTypes.embedding.value + or call_type == CallTypes.aembedding.value + or call_type == CallTypes.responses.value + or call_type == CallTypes.aresponses.value + ): + input_data: Final = data.get("input") + + if call_type == CallTypes.embedding.value or call_type == CallTypes.aembedding.value: + if isinstance(input_data, str): + return [{"role": "user", "content": input_data}] + if isinstance(input_data, list): + try: + embedding_input: Final = _EMBEDDING_INPUT_ADAPTER.validate_python( + input_data, # pyright: ignore[reportUnknownArgumentType] # 请求体字段由上游动态传入 + strict=True, + ) + except ValidationError: + return None + return [{"role": "user", "content": item} for item in embedding_input] + return None + + if input_data is None: + return None + + ######################################################### + # /responses + # User/System messages are stored in the "input" key, use litellm transformation to get the messages + ######################################################### from typing import cast from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) - input_data: Final = data.get("input") - if input_data is None: - return None - messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=input_data, responses_api_request=data, diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 72c362425bb..f949575b017 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -25,6 +25,47 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +class TestGetGuardrailsMessagesForCallType: + @pytest.mark.parametrize("call_type", [CallTypes.embedding, CallTypes.aembedding]) + @pytest.mark.parametrize( + ("embedding_input", "expected_messages"), + [ + ("hello", [{"role": "user", "content": "hello"}]), + ( + ["first", "second"], + [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ], + ), + ], + ids=["single-text", "batch-text"], + ) + def test_embedding_text_input_is_extracted_as_user_messages( + self, + call_type: CallTypes, + embedding_input: str | list[str], + expected_messages: list[dict[str, str]], + ) -> None: + guardrail: Final = CustomGuardrail(guardrail_name="embedding-extractor") + data: Final = {"input": embedding_input} + + messages: Final = guardrail.get_guardrails_messages_for_call_type(call_type=call_type, data=data) + + assert messages == expected_messages + 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: + guardrail: Final = CustomGuardrail(guardrail_name="embedding-extractor") + data: Final = {"input": [1, 2]} + + messages: Final = guardrail.get_guardrails_messages_for_call_type(call_type=call_type, data=data) + + assert messages is None + assert data == {"input": [1, 2]} + + class TestCustomGuardrailDeploymentHook: @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_no_guardrails(self): From e013a40b57e6fc1da253ee27aa575c730607c77f Mon Sep 17 00:00:00 2001 From: zhengguangzhuo <100548634+zhengguangzhuo@users.noreply.github.com> Date: Thu, 8 Oct 2026 17:28:45 +0800 Subject: [PATCH 2/4] fix(guardrails): satisfy strict comparison lint --- litellm/integrations/custom_guardrail.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index b6b88ea7cc0..4ca2bf7b305 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1539,15 +1539,15 @@ class CustomGuardrail(CustomLogger): ): return data.get("messages") - if ( - call_type == CallTypes.embedding.value - or call_type == CallTypes.aembedding.value - or call_type == CallTypes.responses.value - or call_type == CallTypes.aresponses.value + if call_type in ( + CallTypes.embedding.value, + CallTypes.aembedding.value, + CallTypes.responses.value, + CallTypes.aresponses.value, ): input_data: Final = data.get("input") - if call_type == CallTypes.embedding.value or call_type == CallTypes.aembedding.value: + if call_type in (CallTypes.embedding.value, CallTypes.aembedding.value): if isinstance(input_data, str): return [{"role": "user", "content": input_data}] if isinstance(input_data, list): From 0757f964315b3a0ccf6c6186547bee126db08a3c Mon Sep 17 00:00:00 2001 From: zhengguangzhuo <100548634+zhengguangzhuo@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:10:00 +0800 Subject: [PATCH 3/4] [AI] test(guardrails): cover unsupported embedding inputs --- .../integrations/test_custom_guardrail.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index f949575b017..8a8917fb59e 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -65,6 +65,27 @@ class TestGetGuardrailsMessagesForCallType: assert messages is None assert data == {"input": [1, 2]} + @pytest.mark.parametrize("call_type", [CallTypes.embedding, CallTypes.aembedding]) + @pytest.mark.parametrize("embedding_input", [None, 123, {"text": "hello"}], ids=["none", "number", "object"]) + def test_embedding_unsupported_input_is_not_extracted(self, call_type: CallTypes, embedding_input: object) -> None: + guardrail: Final = CustomGuardrail(guardrail_name="embedding-extractor") + 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": embedding_input} + + @pytest.mark.parametrize("call_type", [CallTypes.responses, CallTypes.aresponses]) + def test_responses_missing_input_is_not_extracted(self, call_type: CallTypes) -> None: + guardrail: Final = CustomGuardrail(guardrail_name="responses-extractor") + data: Final = {"input": None} + + messages: Final = guardrail.get_guardrails_messages_for_call_type(call_type=call_type, data=data) + + assert messages is None + assert data == {"input": None} + class TestCustomGuardrailDeploymentHook: @pytest.mark.asyncio From 8420aed84f6cd8ff4efa9493beb7306628509b86 Mon Sep 17 00:00:00 2001 From: zhengguangzhuo <100548634+zhengguangzhuo@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:11:04 +0800 Subject: [PATCH 4/4] test(guardrails): cover mixed text and bytes embedding batches --- tests/unit/integrations/test_custom_guardrail.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 8a8917fb59e..5b8946dc2a9 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -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"])