From 93fdd70275ec4e542d66184b762cbac697259789 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 5 Feb 2026 13:17:02 -0800 Subject: [PATCH] init OpenAIEmbeddingsHandler --- .../guardrail_translation/handler.py | 179 ++++++++++++++++++ 1 file changed, 179 insertions(+) create mode 100644 litellm/llms/openai/embeddings/guardrail_translation/handler.py diff --git a/litellm/llms/openai/embeddings/guardrail_translation/handler.py b/litellm/llms/openai/embeddings/guardrail_translation/handler.py new file mode 100644 index 00000000000..7458020e109 --- /dev/null +++ b/litellm/llms/openai/embeddings/guardrail_translation/handler.py @@ -0,0 +1,179 @@ +""" +OpenAI Embeddings Handler for Unified Guardrails + +This module provides guardrail translation support for OpenAI's embeddings endpoint. +The handler processes the 'input' parameter for guardrails. +""" + +from typing import TYPE_CHECKING, Any, List, Optional, Union + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.utils import EmbeddingResponse + + +class OpenAIEmbeddingsHandler(BaseTranslation): + """ + Handler for processing OpenAI embeddings requests with guardrails. + + This class provides methods to: + 1. Process input text (pre-call hook) + 2. Process output response (post-call hook) - embeddings don't typically need output guardrails + + The handler specifically processes the 'input' parameter which can be: + - A single string + - A list of strings (for batch embeddings) + - A list of integers (token IDs - not processed by guardrails) + - A list of lists of integers (batch token IDs - not processed by guardrails) + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + ) -> Any: + """ + Process input text by applying guardrails to text content. + + Args: + data: Request data dictionary containing 'input' parameter + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + + Returns: + Modified data with guardrails applied to input + """ + input_data = data.get("input") + if input_data is None: + verbose_proxy_logger.debug( + "OpenAI Embeddings: No input found in request data" + ) + return data + + if isinstance(input_data, str): + data = await self._process_string_input( + data, input_data, guardrail_to_apply, litellm_logging_obj + ) + elif isinstance(input_data, list): + data = await self._process_list_input( + data, input_data, guardrail_to_apply, litellm_logging_obj + ) + else: + verbose_proxy_logger.warning( + "OpenAI Embeddings: Unexpected input type: %s. Expected string or list.", + type(input_data), + ) + + return data + + async def _process_string_input( + self, + data: dict, + input_data: str, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any], + ) -> dict: + """Process a single string input through the guardrail.""" + inputs = GenericGuardrailAPIInputs(texts=[input_data]) + if model := data.get("model"): + inputs["model"] = model + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + + if guardrailed_texts := guardrailed_inputs.get("texts"): + data["input"] = guardrailed_texts[0] + verbose_proxy_logger.debug( + "OpenAI Embeddings: Applied guardrail to string input. " + "Original length: %d, New length: %d", + len(input_data), + len(data["input"]), + ) + + return data + + async def _process_list_input( + self, + data: dict, + input_data: List[Union[str, int, List[int]]], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any], + ) -> dict: + """Process a list input through the guardrail (if it contains strings).""" + if len(input_data) == 0: + return data + + first_item = input_data[0] + + # Skip non-text inputs (token IDs) + if isinstance(first_item, (int, list)): + verbose_proxy_logger.debug( + "OpenAI Embeddings: Input is token IDs, skipping guardrail processing" + ) + return data + + if not isinstance(first_item, str): + verbose_proxy_logger.warning( + "OpenAI Embeddings: Unexpected input list item type: %s", + type(first_item), + ) + return data + + # List of strings - apply guardrail + inputs = GenericGuardrailAPIInputs(texts=input_data) # type: ignore + if model := data.get("model"): + inputs["model"] = model + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + + if guardrailed_texts := guardrailed_inputs.get("texts"): + data["input"] = guardrailed_texts + verbose_proxy_logger.debug( + "OpenAI Embeddings: Applied guardrail to %d inputs", + len(guardrailed_texts), + ) + + return data + + async def process_output_response( + self, + response: "EmbeddingResponse", + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + user_api_key_dict: Optional[Any] = None, + ) -> Any: + """ + Process output response - embeddings responses contain vectors, not text. + + For embeddings, the output is numerical vectors, so there's typically + no text content to apply guardrails to. This method is a no-op but + is included for interface consistency. + + Args: + response: Embedding response object + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + user_api_key_dict: User API key metadata + + Returns: + Unmodified response (embeddings don't have text output to guard) + """ + verbose_proxy_logger.debug( + "OpenAI Embeddings: Output response processing skipped - " + "embeddings contain vectors, not text" + ) + return response