Merge 8420aed84f into cd20ecace1

litellm/integrations/custom_guardrail.py

@@ -9,6 +9,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
 
 import httpx
+from pydantic import TypeAdapter, ValidationError
 
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_logger
@@ -69,6 +70,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})
 
@@ -1614,21 +1616,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 in (
+            CallTypes.embedding.value,
+            CallTypes.aembedding.value,
+            CallTypes.responses.value,
+            CallTypes.aresponses.value,
+        ):
+            input_data: Final = data.get("input")
+
+            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):
+                    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,

tests/unit/integrations/test_custom_guardrail.py

@@ -32,6 +32,75 @@ 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])
+    @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": 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.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
     async def test_async_pre_call_deployment_hook_no_guardrails(self):