fix(responses): presidio PII masking for Azure WebSocket and streaming

litellm/llms/azure/responses/transformation.py

@@ -185,6 +185,40 @@ def get_complete_url(
             default_api_version=AZURE_DEFAULT_RESPONSES_API_VERSION,
         )
 
+    def supports_native_websocket(self) -> bool:
+        return True
+
+    def get_websocket_url(
+        self,
+        api_base: Optional[str],
+        litellm_params: dict,
+    ) -> str:
+        """
+        Azure Responses WebSocket endpoint is at /openai/v1/responses with no
+        api-version query param. Auth is via Authorization header, model is sent
+        in the response.create body — not the URL.
+        """
+        if api_base is None:
+            raise ValueError("api_base is required for Azure WebSocket")
+
+        parsed_url = httpx.URL(api_base)
+        path = parsed_url.path.rstrip("/")
+        # Strip existing /openai/responses path if the api_base already contains it
+        for suffix in ("/openai/v1/responses", "/openai/responses"):
+            if path.endswith(suffix):
+                path = path[: -len(suffix)]
+                break
+        scheme = "wss" if parsed_url.scheme == "https" else "ws"
+        return str(
+            parsed_url.copy_with(
+                scheme=scheme, path=f"{path}/openai/v1/responses", query=None
+            )
+        )
+
+    def model_in_websocket_url(self) -> bool:
+        # Azure sends the model in the response.create body, not the URL
+        return False
+
     #########################################################
     ########## DELETE RESPONSE API TRANSFORMATION ##############
     #########################################################

litellm/llms/base_llm/responses/transformation.py

@@ -258,6 +258,31 @@ def supports_native_websocket(self) -> bool:
         """
         return False
 
+    def get_websocket_url(
+        self,
+        api_base: Optional[str],
+        litellm_params: dict,
+    ) -> str:
+        """
+        Return the wss:// URL for the provider's native Responses WebSocket endpoint.
+
+        Defaults to converting the HTTP URL from get_complete_url. Providers whose
+        WebSocket path differs from their HTTP path (e.g. Azure uses
+        /openai/v1/responses without api-version) should override this.
+        """
+        http_url = self.get_complete_url(
+            api_base=api_base, litellm_params=litellm_params
+        )
+        return http_url.replace("https://", "wss://").replace("http://", "ws://")
+
+    def model_in_websocket_url(self) -> bool:
+        """
+        Return True if the model should be appended as a ?model= query param to
+        the WebSocket URL. Providers that identify the model via the request body
+        (e.g. Azure Responses API) should override this to return False.
+        """
+        return True
+
     #########################################################
     ########## CANCEL RESPONSE API TRANSFORMATION ##########
     #########################################################

litellm/llms/custom_httpx/llm_http_handler.py

@@ -5575,7 +5575,7 @@ async def async_realtime_calls_handler(
                 )
             raise
 
-    async def async_responses_websocket(
+    async def async_responses_websocket(  # noqa: PLR0915
         self,
         model: str,
         websocket: Any,
@@ -5628,7 +5628,11 @@ async def async_responses_websocket(
         import websockets
         from websockets.asyncio.client import ClientConnection
 
-        litellm_params = GenericLiteLLMParams()
+        litellm_params = GenericLiteLLMParams(
+            api_base=api_base,
+            api_key=api_key,
+            **kwargs,
+        )
         headers = responses_api_provider_config.validate_environment(
             headers={},
             model=model,
@@ -5637,21 +5641,21 @@ async def async_responses_websocket(
         if api_key:
             headers["Authorization"] = f"Bearer {api_key}"
 
-        http_url = responses_api_provider_config.get_complete_url(
+        ws_url = responses_api_provider_config.get_websocket_url(
             api_base=api_base,
-            litellm_params={},
+            litellm_params=dict(litellm_params),
         )
-        ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://")
-        # OpenAI's WebSocket responses endpoint requires ?model= in the URL,
-        # matching the Realtime API convention (wss://.../v1/realtime?model=...).
-        # Use urllib.parse so existing query params (e.g. api-version) are preserved.
-        _parsed = urlparse(ws_url)
-        _qs = parse_qs(_parsed.query)
-        if "model" not in _qs:
-            _qs["model"] = [model]
-            ws_url = urlunparse(
-                _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()}))
-            )
+        # Some providers (e.g. OpenAI) require ?model= in the WebSocket URL.
+        # Providers that send the model in the request body (e.g. Azure) set
+        # model_in_websocket_url() to False to suppress this append.
+        if responses_api_provider_config.model_in_websocket_url():
+            _parsed = urlparse(ws_url)
+            _qs = parse_qs(_parsed.query)
+            if "model" not in _qs:
+                _qs["model"] = [model]
+                ws_url = urlunparse(
+                    _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()}))
+                )
 
         try:
             ssl_context = get_shared_realtime_ssl_context()
@@ -5679,13 +5683,51 @@ async def async_responses_websocket(
                 _request_data: Dict[str, Any] = {}
                 if litellm_metadata:
                     _request_data["litellm_metadata"] = litellm_metadata
+
+                _ws_guardrail_callbacks: list = []
+                _ws_output_guardrail_callbacks: list = []
+                try:
+                    import litellm as _litellm
+
+                    # Use duck-typing so any guardrail that exposes the PII
+                    # masking interface works, not just _OPTIONAL_PresidioPIIMasking.
+                    # This avoids a layering violation (SDK importing from proxy).
+                    _ws_guardrail_callbacks = [
+                        cb
+                        for cb in _litellm.callbacks
+                        if callable(getattr(cb, "check_pii", None))
+                        and callable(
+                            getattr(cb, "get_presidio_settings_from_request_data", None)
+                        )
+                        and callable(getattr(cb, "_unmask_pii_text", None))
+                        and getattr(cb, "output_parse_pii", False)
+                    ]
+                    _ws_output_guardrail_callbacks = [
+                        cb
+                        for cb in _litellm.callbacks
+                        if callable(getattr(cb, "check_pii", None))
+                        and callable(
+                            getattr(cb, "get_presidio_settings_from_request_data", None)
+                        )
+                        and getattr(cb, "apply_to_output", False)
+                    ]
+                except Exception as _guardrail_exc:
+                    verbose_logger.warning(
+                        "Responses WebSocket: failed to collect guardrail "
+                        "callbacks — PII masking will be skipped. Error: %s",
+                        _guardrail_exc,
+                    )
+
                 streaming = ResponsesWebSocketStreaming(
                     websocket=websocket,
                     backend_ws=cast(ClientConnection, backend_ws),
                     logging_obj=logging_obj,
                     user_api_key_dict=user_api_key_dict,
                     request_data=_request_data,
                     first_message=first_message,
+                    guardrail_callbacks=_ws_guardrail_callbacks,
+                    output_guardrail_callbacks=_ws_output_guardrail_callbacks,
+                    authorized_model=model,
                 )
                 await streaming.bidirectional_forward()
 

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -35,7 +35,6 @@
 
 from litellm._logging import verbose_proxy_logger
 from litellm.completion_extras.litellm_responses_transformation.transformation import (
-    LiteLLMResponsesTransformationHandler,
     OpenAiResponsesToChatCompletionStreamIterator,
 )
 from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
@@ -479,90 +478,137 @@ async def process_output_streaming_response(
     ) -> List[Any]:
         """
         Process output streaming response by applying guardrails to text content.
+
+        Mirrors the Chat Completions handler pattern: extract text from the final
+        chunk, apply the guardrail, then write the result back in-place so the
+        caller sees the modified content (e.g. PII tokens replaced).
+
+        For ``response.completed`` events (the normal end-of-stream signal) we
+        use the same per-item extraction + task-mapping approach as
+        ``process_output_response`` so that unmasking / blocking works correctly
+        for every output item.
         """
+        if not responses_so_far:
+            return responses_so_far
 
         final_chunk = responses_so_far[-1]
+        # Accept both plain dicts and Pydantic models (BaseLiteLLMOpenAIResponseObject
+        # exposes a .get() shim, so all the .get() calls below work for both).
+        if not (isinstance(final_chunk, dict) or hasattr(final_chunk, "get")):
+            return responses_so_far
+
+        # ------------------------------------------------------------------ #
+        # Case 1: response.completed — full response is available in the      #
+        # final chunk; iterate output items, apply guardrail, write back.     #
+        # ------------------------------------------------------------------ #
+        if final_chunk.get("type") == "response.completed":
+            response_obj = final_chunk.get("response") or {}
+            if not hasattr(response_obj, "get"):
+                return responses_so_far
+            outputs: List[Any] = response_obj.get("output") or []
+
+            texts_to_check: List[str] = []
+            tool_calls_to_check: List[ChatCompletionToolCallChunk] = []
+            task_mappings: List[Tuple[int, int]] = []
+
+            for output_idx, output_item in enumerate(outputs):
+                self._extract_output_text_and_images(
+                    output_item=output_item,
+                    output_idx=output_idx,
+                    texts_to_check=texts_to_check,
+                    images_to_check=[],
+                    task_mappings=task_mappings,
+                    tool_calls_to_check=tool_calls_to_check,
+                )
+
+            if texts_to_check or tool_calls_to_check:
+                if request_data is None:
+                    request_data = {}
+                if "response" not in request_data:
+                    request_data["response"] = response_obj
+                if "litellm_metadata" not in request_data:
+                    user_metadata = self.transform_user_api_key_dict_to_metadata(
+                        user_api_key_dict
+                    )
+                    if user_metadata:
+                        request_data["litellm_metadata"] = user_metadata
+
+                inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
+                if tool_calls_to_check:
+                    inputs["tool_calls"] = cast(
+                        List[ChatCompletionToolCallChunk], tool_calls_to_check
+                    )
+                response_model = response_obj.get("model")
+                if response_model:
+                    inputs["model"] = response_model
+
+                guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
+                    inputs=inputs,
+                    request_data=request_data,
+                    input_type="response",
+                    logging_obj=litellm_logging_obj,
+                )
+
+                guardrailed_texts = guardrailed_inputs.get("texts", [])
+
+                # Write guardrailed texts back into the output items in-place.
+                # final_chunk is a reference into responses_so_far so this
+                # mutates the list that the caller holds.
+                await self._apply_guardrail_responses_to_output(
+                    response=response_obj,
+                    responses=guardrailed_texts,
+                    task_mappings=task_mappings,
+                )
+
+            return responses_so_far
 
+        # ------------------------------------------------------------------ #
+        # Case 2: response.output_item.done — extract tool calls only.        #
+        # ------------------------------------------------------------------ #
         if final_chunk.get("type") == "response.output_item.done":
-            # convert openai response to model response
             model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
                 final_chunk
             )
-
             tool_calls = model_response_stream.choices[0].delta.tool_calls
             if tool_calls:
                 inputs = GenericGuardrailAPIInputs()
                 inputs["tool_calls"] = cast(
                     List[ChatCompletionToolCallChunk], tool_calls
                 )
-                # Include model information if available
                 if (
                     hasattr(model_response_stream, "model")
                     and model_response_stream.model
                 ):
                     inputs["model"] = model_response_stream.model
-                _guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
+                await guardrail_to_apply.apply_guardrail(
                     inputs=inputs,
                     request_data=request_data if request_data is not None else {},
                     input_type="response",
                     logging_obj=litellm_logging_obj,
                 )
-                return responses_so_far
-        elif final_chunk.get("type") == "response.completed":
-            # convert openai response to model response
-            outputs = final_chunk.get("response", {}).get("output", [])
+            return responses_so_far
 
-            model_response_choices = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices(
-                output_items=outputs,
-                handle_raw_dict_callback=None,
-            )
-
-            if model_response_choices:
-                tool_calls = model_response_choices[0].message.tool_calls
-                text = model_response_choices[0].message.content
-                guardrail_inputs = GenericGuardrailAPIInputs()
-                if text:
-                    guardrail_inputs["texts"] = [text]
-                if tool_calls:
-                    guardrail_inputs["tool_calls"] = cast(
-                        List[ChatCompletionToolCallChunk], tool_calls
-                    )
-                # Include model information from the response if available
-                response_model = final_chunk.get("response", {}).get("model")
-                if response_model:
-                    guardrail_inputs["model"] = response_model
-                if tool_calls or text:
-                    _guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
-                        inputs=guardrail_inputs,
-                        request_data=request_data if request_data is not None else {},
-                        input_type="response",
-                        logging_obj=litellm_logging_obj,
-                    )
-                    return responses_so_far
-            else:
-                verbose_proxy_logger.debug(
-                    "Skipping output guardrail - model response has no choices"
-                )
-        # model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk)
-        # tool_calls = model_response_stream.choices[0].tool_calls
-        # convert openai response to model response
+        # ------------------------------------------------------------------ #
+        # Fallback: apply guardrail to the accumulated text string.           #
+        # No structured write-back is possible here; guardrails that only     #
+        # need to block/flag (not rewrite) still work correctly.             #
+        # ------------------------------------------------------------------ #
         string_so_far = self.get_streaming_string_so_far(responses_so_far)
-        inputs = GenericGuardrailAPIInputs(texts=[string_so_far])
-        # Try to get model from the final chunk if available
-        if isinstance(final_chunk, dict):
+        if string_so_far:
+            fallback_inputs = GenericGuardrailAPIInputs(texts=[string_so_far])
             response_model = (
                 final_chunk.get("response", {}).get("model")
                 if isinstance(final_chunk.get("response"), dict)
                 else None
             )
             if response_model:
-                inputs["model"] = response_model
-        _guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
-            inputs=inputs,
-            request_data=request_data if request_data is not None else {},
-            input_type="response",
-            logging_obj=litellm_logging_obj,
-        )
+                fallback_inputs["model"] = response_model
+            await guardrail_to_apply.apply_guardrail(
+                inputs=fallback_inputs,
+                request_data=request_data if request_data is not None else {},
+                input_type="response",
+                logging_obj=litellm_logging_obj,
+            )
         return responses_so_far
 
     def _check_streaming_has_ended(self, responses_so_far: List[Any]) -> bool:
@@ -721,7 +767,7 @@ def _extract_output_text_and_images(
 
     async def _apply_guardrail_responses_to_output(
         self,
-        response: "ResponsesAPIResponse",
+        response: Union["ResponsesAPIResponse", Dict[Any, Any]],
         responses: List[str],
         task_mappings: List[Tuple[int, int]],
     ) -> None: