Merge 81c5ab480a into b6f91a221b

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
 
@@ -34,6 +34,10 @@ class StreamTransformSink:
     mutated_text_per_choice: dict[int, str] = field(default_factory=dict)
     holdback_per_choice: dict[int, int] = field(default_factory=dict)
 
+    def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None:
+        self.mutated_text_per_choice = dict(texts)
+        self.holdback_per_choice = dict(holdback)
+
 
 @dataclass(frozen=True, slots=True)
 class StreamingScanKey:

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

@@ -18,9 +18,11 @@ import json
 import time
 import uuid
 from collections.abc import Iterator, Mapping, Sequence
+from itertools import chain
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Union, cast
 
+from pydantic import TypeAdapter
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 import litellm
@@ -50,7 +52,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     stream_item_items,
     unappliable_request_rewrite,
 )
-from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam
 from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
     coerce_stream_holdback_value,
 )
@@ -172,8 +174,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 logging_obj=litellm_logging_obj,
             )
 
-            guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
-            guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", [])
+            guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
+            guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", ())
             guardrailed_tools: Final = guardrailed_inputs.get("tools")
             if guardrailed_tools is not None:
                 data["tools"] = (
@@ -392,7 +394,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
         user_api_key_dict: "UserAPIKeyAuth | None" = None,
-        request_data: dict | None = None,
+        request_data: dict[str, object] | None = None,
     ) -> ModelResponse:
         """
         Process output response by applying guardrails to text content.
@@ -443,9 +445,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             # create a standalone dict (SDK / direct-call path).
             if request_data is None:
                 request_data = {"response": response}
-            else:
-                if "response" not in request_data:
-                    request_data["response"] = response
+            request_data.setdefault("response", response)
 
             self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
 
@@ -465,7 +465,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 logging_obj=litellm_logging_obj,
             )
 
-            guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
+            guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
             returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
             guardrailed_tool_calls: Final[list[dict[str, object]]] = (
                 cast(list[dict[str, object]], returned_tool_calls)
@@ -499,7 +499,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
         user_api_key_dict: "UserAPIKeyAuth | None" = None,
-        request_data: dict | None = None,
+        request_data: dict[str, object] | None = None,
         stream_transform_sink: StreamTransformSink | None = None,
         deliver_ended_stream_rewrites: bool = False,
     ) -> list["ModelResponseStream"]:
@@ -555,14 +555,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None",
         user_api_key_dict: "UserAPIKeyAuth | None",
-        request_data: dict | None,
+        request_data: dict[str, object] | None,
         deliver_ended_stream_rewrites: bool = False,
     ) -> list["ModelResponseStream"]:
         """Block-only streaming path: run the guardrail so an in-flight BLOCK can
         terminate the stream. Text rewrites are not propagated to the client here
         (see ``_process_streaming_transform`` for the incremental_diff path) unless
         ``deliver_ended_stream_rewrites`` opts the ended-stream branch in."""
-        has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
+        has_stream_ended: Final = deliver_ended_stream_rewrites or self._first_choice_has_finished(responses_so_far)
 
         if has_stream_ended:
             await self._process_ended_stream(
@@ -628,7 +628,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 logging_obj=litellm_logging_obj,
             )
 
-            guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
+            guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
 
             # Step 4: Apply guardrailed text back to all streaming chunks
             # For each choice, replace the combined text across all chunks
@@ -661,13 +661,36 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
         pre_guardrail_texts: Final = self._string_choice_contents(model_response)
         pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response)
-        await self.process_output_response(
-            response=model_response,
-            guardrail_to_apply=guardrail_to_apply,
-            litellm_logging_obj=litellm_logging_obj,
-            user_api_key_dict=user_api_key_dict,
-            request_data=request_data,
+        inspection_responses: Final = (
+            tuple(
+                model_response.model_copy(update=MappingProxyType({"choices": [choice]}))
+                for choice in model_response.choices
+            )
+            if pre_guardrail_tool_calls and len(model_response.choices) > 1
+            else (model_response,)
         )
+        inspection_request_data: Final = request_data if request_data is not None else {}
+        try:
+            for inspection_response in inspection_responses:
+                inspection_request_data["response"] = inspection_response
+                inspection_request_data["responses"] = (
+                    [
+                        self._narrowed_to_choice(chunk, inspection_response.choices[0].index)
+                        for chunk in responses_so_far
+                    ]
+                    if len(inspection_response.choices) == 1
+                    else responses_so_far
+                )
+                await self.process_output_response(
+                    response=inspection_response,
+                    guardrail_to_apply=guardrail_to_apply,
+                    litellm_logging_obj=litellm_logging_obj,
+                    user_api_key_dict=user_api_key_dict,
+                    request_data=inspection_request_data,
+                )
+        finally:
+            inspection_request_data["response"] = model_response
+            inspection_request_data["responses"] = responses_so_far
         if not deliver_ended_stream_rewrites:
             return
         await self._write_ended_stream_text_rewrites(
@@ -768,7 +791,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None",
         user_api_key_dict: "UserAPIKeyAuth | None",
-        request_data: dict | None,
+        request_data: dict[str, object] | None,
         sink: StreamTransformSink,
     ) -> None:
         """Run the guardrail over the raw accumulated text and report the
@@ -780,8 +803,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         """
         raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far)
         if not raw_by_index:
-            sink.mutated_text_per_choice = {}
-            sink.holdback_per_choice = {}
+            sink.record(MappingProxyType({}), MappingProxyType({}))
             return
 
         # Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits
@@ -793,18 +815,72 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         indices: Final = sorted(raw_by_index.keys())
         texts_to_check: Final = [raw_by_index[i] for i in indices]
 
-        if request_data is None:
-            request_data = {"responses": responses_so_far}
-        elif "responses" not in request_data:
-            request_data["responses"] = responses_so_far
-        self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
+        resolved_request_data: Final[dict[str, object]] = (
+            request_data if request_data is not None else {"responses": responses_so_far}
+        )
+        if "responses" not in resolved_request_data:
+            resolved_request_data["responses"] = responses_so_far
+        self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict)
 
         inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
+        if self._streamed_tool_call_fingerprints(responses_so_far):
+            assembled: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
+            resolved_request_data["response"] = assembled
+            if len(assembled.choices) > 1:
+                choice_rounds: Final = tuple(
+                    (
+                        choice,
+                        StreamTransformSink(),
+                        [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far],
+                    )
+                    for choice in assembled.choices
+                )
+                try:
+                    for choice, choice_sink, choice_chunks in choice_rounds:
+                        resolved_request_data["response"] = assembled.model_copy(
+                            update=MappingProxyType({"choices": [choice]})
+                        )
+                        resolved_request_data["responses"] = choice_chunks
+                        await self._process_streaming_transform(
+                            responses_so_far=choice_chunks,
+                            guardrail_to_apply=guardrail_to_apply,
+                            litellm_logging_obj=litellm_logging_obj,
+                            user_api_key_dict=user_api_key_dict,
+                            request_data=resolved_request_data,
+                            sink=choice_sink,
+                        )
+                finally:
+                    resolved_request_data["response"] = assembled
+                    resolved_request_data["responses"] = responses_so_far
+                merged_texts: Final = MappingProxyType(
+                    dict(
+                        chain.from_iterable(
+                            choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds
+                        )
+                    )
+                )
+                merged_holdback: Final = MappingProxyType(
+                    dict(
+                        chain.from_iterable(
+                            choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds
+                        )
+                    )
+                )
+                sink.record(merged_texts, merged_holdback)
+                return
+            tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices)
+            inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python(
+                tuple(
+                    MappingProxyType({"index": index, **converted})
+                    for index, tool_call in enumerate(tool_calls)
+                    if (converted := self._convert_tool_call_to_dict(tool_call)) is not None
+                )
+            )
         if responses_so_far and getattr(responses_so_far[0], "model", None):
             inputs["model"] = responses_so_far[0].model
         guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
             inputs=inputs,
-            request_data=request_data,
+            request_data=resolved_request_data,
             input_type="response",
             logging_obj=litellm_logging_obj,
         )
@@ -824,13 +900,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 len(texts_to_check),
             )
 
-        holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or []
-        sink.mutated_text_per_choice = {
-            idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)
-        }
-        sink.holdback_per_choice = {
-            indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)
-        }
+        holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or ()
+        transformed_texts: Final = MappingProxyType(
+            {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)}
+        )
+        transformed_holdback: Final = MappingProxyType(
+            {indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)}
+        )
+        sink.record(transformed_texts, transformed_holdback)
 
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
@@ -1147,20 +1224,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     def _function_tool_call_fragments(
         responses_so_far: Sequence["ModelResponseStream"],
     ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]:
-        """Group the stream's function tool-call fragments by their tool-call index, in
-        the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping
-        only the indices the builder keeps (an id and a name somewhere in the stream)."""
         fragments: Final = tuple(
-            tool_call
+            (choice.index, tool_call)
             for response in responses_so_far
             for choice in response.choices
             for tool_call in choice.delta.tool_calls or ()
             if isinstance(tool_call, ChatCompletionDeltaToolCall)
         )
-        identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id)
-        named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name)
+        identified: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.id)
+        named: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.function.name)
         return tuple(
-            tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named)
+            tuple(fragment for choice, fragment in fragments if (choice, fragment.index) == key)
+            for key in sorted(identified & named)
         )
 
     def _write_ended_stream_tool_call_rewrites(
@@ -1170,28 +1245,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...],
         guardrail_name: str,
     ) -> None:
-        """Write ended-stream guardrail tool-call rewrites back across the buffered
-        chunks: the rewritten name and full arguments land in the tool call's first
-        fragment and the arguments of its later fragments are blanked, mirroring the
-        text write-back. A rewrite on a stream carrying more than one distinct choice
-        index, or whose fragments do not line up with the rebuilt tool calls, is
-        reported as undeliverable, so the pipeline executor discards it and releases
-        the original chunks."""
         post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response)
         if post_guardrail_tool_calls == pre_guardrail_tool_calls:
             return
-        stream_choice_indices: Final = frozenset(
-            choice.index for response in responses_so_far for choice in response.choices
-        )
         fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far)
-        if len(stream_choice_indices) != 1:
-            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
-
-            raise UndeliverableStreamRewrite(
-                guardrail_name,
-                f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written "
-                "back on single-choice streams",
-            )
         if len(fragments_by_tool_call) != len(post_guardrail_tool_calls):
             from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
@@ -1214,7 +1271,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     async def _apply_guardrail_responses_to_output_streaming(
         self,
         responses: list["ModelResponseStream"],
-        guardrailed_texts: list[str],
+        guardrailed_texts: Sequence[str],
         task_mappings: list[tuple[int, int | None]],
     ) -> None:
         """
@@ -1232,12 +1289,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         Override this method to customize how responses are applied to streaming responses.
         """
         # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx)
-        guardrail_map: Final[dict[tuple[int, int | None], str]] = {}
-        for task_idx, guardrail_response in enumerate(guardrailed_texts):
-            mapping = task_mappings[task_idx]
-            choice_idx = cast(int, mapping[0])
-            content_idx_optional = cast(int | None, mapping[1])
-            guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response
+        guardrail_map: Final = MappingProxyType(
+            {task_mappings[index]: text for index, text in enumerate(guardrailed_texts)}
+        )
 
         # Track which choices we've already set the guardrailed text for
         # Key: (choice_idx, content_idx), Value: boolean (True if already set)
@@ -1247,34 +1301,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         # choice by its index field: on n>1 streams a chunk usually carries one
         # choice at list position 0 whose index names the logical choice.
         for response in responses:
-            for choice in response.choices:
-                if isinstance(choice, litellm.StreamingChoices):
-                    content = choice.delta.content
-                elif isinstance(choice, litellm.Choices):
-                    content = choice.message.content
-                else:
-                    continue
-
+            for choice, message in (
+                (choice, choice.delta if isinstance(choice, StreamingChoices) else choice.message)
+                for choice in response.choices
+                if isinstance(choice, (StreamingChoices, Choices))
+            ):
+                content = message.content
                 if content is None:
                     continue
 
                 if isinstance(content, str):
-                    # String content
                     str_key: tuple[int, int | None] = (choice.index, None)
                     if str_key in guardrail_map:
-                        if str_key not in already_set:
-                            # First chunk - set the complete guardrailed text
-                            if isinstance(choice, litellm.StreamingChoices):
-                                choice.delta.content = guardrail_map[str_key]
-                            elif isinstance(choice, litellm.Choices):
-                                choice.message.content = guardrail_map[str_key]
-                            already_set[str_key] = True
-                        else:
-                            # Subsequent chunks - clear the content
-                            if isinstance(choice, litellm.StreamingChoices):
-                                choice.delta.content = ""
-                            elif isinstance(choice, litellm.Choices):
-                                choice.message.content = ""
+                        message.content = guardrail_map[str_key] if str_key not in already_set else ""
+                        already_set[str_key] = True
 
                 elif isinstance(content, list):
                     # List content - handle each content item

litellm/proxy/_lazy_openapi_snapshot.json

@@ -15724,6 +15724,22 @@
               "description": "Whether this guardrail runs on streaming requests, non-streaming requests, or both. A string applies to every configured mode. A map overrides named modes (pre_call, during_call, post_call, ...); omitted keys default to both. Unset means both, matching historical behavior.",
               "title": "Stream Scope"
             },
+            "streaming_transform_mode": {
+              "anyOf": [
+                {
+                  "enum": [
+                    "block_only",
+                    "incremental_diff"
+                  ],
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Whether a guardrail's text rewrite reaches a streaming client. Implemented by guardrail='prompt_security' and 'neuraltrust'; generic_guardrail_api takes the same setting under optional_params. 'block_only' (default) streams the raw model chunks, so a block still ends the stream but rewrites are dropped. 'incremental_diff' withholds those chunks and streams the guardrail's rewritten text instead. OpenAI chat completions streaming only.",
+              "title": "Streaming Transform Mode"
+            },
             "template_id": {
               "anyOf": [
                 {
@@ -15750,7 +15766,7 @@
             },
             "unreachable_fallback": {
               "default": "fail_closed",
-              "description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. Also implemented by guardrail='decision_model'.",
+              "description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', 'typesafe', and 'neuraltrust'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. Also implemented by guardrail='decision_model'.",
               "enum": [
                 "fail_closed",
                 "fail_open"
@@ -17372,6 +17388,18 @@
               "description": "Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable.",
               "title": "Client Secret"
             },
+            "collector_key": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "TrustGuard collector key (tgcol_...). Optional when the API key is bound to a collector. Env: TRUSTGUARD_COLLECTOR_KEY.",
+              "title": "Collector Key"
+            },
             "confidence_threshold": {
               "default": 0.5,
               "default_value": 0.5,
@@ -18804,6 +18832,22 @@
               "description": "Whether this guardrail runs on streaming requests, non-streaming requests, or both. A string applies to every configured mode. A map overrides named modes (pre_call, during_call, post_call, ...); omitted keys default to both. Unset means both, matching historical behavior.",
               "title": "Stream Scope"
             },
+            "streaming_transform_mode": {
+              "anyOf": [
+                {
+                  "enum": [
+                    "block_only",
+                    "incremental_diff"
+                  ],
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Whether a guardrail's text rewrite reaches a streaming client. Implemented by guardrail='prompt_security' and 'neuraltrust'; generic_guardrail_api takes the same setting under optional_params. 'block_only' (default) streams the raw model chunks, so a block still ends the stream but rewrites are dropped. 'incremental_diff' withholds those chunks and streams the guardrail's rewritten text instead. OpenAI chat completions streaming only.",
+              "title": "Streaming Transform Mode"
+            },
             "template_id": {
               "anyOf": [
                 {

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -2331,6 +2331,26 @@ async def test_custom_code_guardrail(
         )
 
 
+_GUARDRAIL_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
+
+
+def _metadata_fields(value: object) -> Mapping[str, object]:
+    try:
+        return _GUARDRAIL_METADATA_ADAPTER.validate_python(value)
+    except ValidationError:
+        return {}
+
+
+def _guardrail_request_metadata(caller: object, proxy: object) -> Mapping[str, object]:
+    caller_fields: Final = tuple(
+        (key, value) for key, value in _metadata_fields(caller).items() if not key.startswith("user_api_key_")
+    )
+    key_identity: Final = tuple(
+        (key, value) for key, value in _metadata_fields(proxy).items() if key.startswith("user_api_key_")
+    )
+    return dict((*caller_fields, *key_identity))
+
+
 def _execution_timeout_response(timeout: float) -> TestCustomCodeGuardrailResponse:
     return TestCustomCodeGuardrailResponse(
         success=False,
@@ -2493,9 +2513,16 @@ async def apply_guardrail(
         if litellm_logging_obj is not None:
             _patch_logging_obj_for_guardrail(litellm_logging_obj, request)
 
+        metadata: Final = _guardrail_request_metadata(
+            request.metadata, _GUARDRAIL_METADATA_ADAPTER.validate_python(data).get("metadata")
+        )
         request_data: Final[dict] = {
-            **({"messages": request.messages} if request.messages is not None else {}),
-            **({"metadata": request.metadata} if request.metadata is not None else {}),
+            key: value
+            for key, value in (
+                ("messages", request.messages),
+                ("metadata", metadata if request.metadata is not None or metadata else None),
+            )
+            if value is not None
         }
         _input_type: Final = _resolve_guardrail_input_type(active_guardrail, request.input_type)
         guardrailed_inputs: Final = await active_guardrail.apply_guardrail(

litellm/proxy/guardrails/guardrail_hooks/neuraltrust/README.md

@@ -0,0 +1,107 @@
+# NeuralTrust TrustGuard
+
+Native LiteLLM guardrail. Sends chat input and output to TrustGuard `POST /v1/evaluate`.
+
+Setup guide, verdict mapping, and the streaming caveat:
+[docs.neuraltrust.ai/integrations/litellm](https://docs.neuraltrust.ai/integrations/litellm).
+
+## Config
+
+```yaml
+guardrails:
+  - guardrail_name: neuraltrust-trustguard
+    litellm_params:
+      guardrail: neuraltrust
+      mode: [pre_call, post_call]
+      api_key: os.environ/TRUSTGUARD_API_KEY
+      api_base: os.environ/TRUSTGUARD_API_BASE  # default https://trustguard.neuraltrust.ai
+      collector_key: os.environ/TRUSTGUARD_COLLECTOR_KEY  # tgcol_… ; optional if the API key is bound
+      unreachable_fallback: fail_closed
+      timeout: 5
+      streaming_transform_mode: incremental_diff
+      default_on: true
+```
+
+## Auth
+
+Bearer `tgk_…` API key. Address the collector with `collector_key`, or omit it when the key is already bound to one.
+
+## Identity
+
+Each evaluate call carries `session_id` from the LiteLLM session and `consumer_id` from the virtual key: the key alias, else the key's user email, user id, or team alias. TrustGuard Activity and per-consumer policies group by that value.
+
+## Verdicts
+
+| TrustGuard `status` | LiteLLM |
+| --- | --- |
+| `block` | HTTP 400 (trace_id / request_id only; findings are not echoed) |
+| `ask` | HTTP 400 like `block`: a proxy has no approval flow, so the response names `verdict: ask` |
+| `transform` | rewrite the last user message / last text from `transformed_payload` |
+| `report` / `allow` | pass through (`report` is logged by trace_id) |
+
+Unknown verdicts, malformed bodies, and `transform` without a usable payload fail closed.
+
+## Fail-open vs fail-closed
+
+`unreachable_fallback` applies only to transport failures: connect errors, timeouts, HTTP 502/504. TrustGuard
+does not emit 502 or 504 itself, so those mean the request never reached a verdict.
+
+HTTP 503 entitlements, 429 rate limits, 401/403, other 4xx/5xx, and unusable TrustGuard verdicts always fail closed.
+
+503 and 429 stay fail-closed on purpose, unlike `generic_guardrail_api`, which treats 503 as unreachable. A
+TrustGuard 503 means its rate-limit entitlements could not be loaded, so TrustGuard itself declined to evaluate
+rather than evaluate blind; 429 means the plan's burst or monthly quota is exhausted. Both are verdicts about
+whether the traffic can be inspected, not signs that the service is down, so honouring `fail_open` for them would
+pass uninspected traffic the guardrail deliberately refused. 429 logs `X-RateLimit-Reason` so a transient burst is
+distinguishable from an exhausted monthly quota.
+
+`fail_open` means the request bypasses TrustGuard entirely when the endpoint is unreachable. It is off by default.
+
+## Streaming
+
+`block` and `ask` end the stream in either mode. `streaming_transform_mode` decides whether a `transform` verdict reaches the client.
+
+| Mode | What the client receives |
+| --- | --- |
+| `block_only` (default) | the raw model tokens, so the redaction is dropped |
+| `incremental_diff` | TrustGuard's rewritten reply |
+
+Under `incremental_diff` the reply is held until the end-of-stream evaluate returns, so the first token arrives with the last. TrustGuard re-reads the whole reply on every scan and its redaction spans move as the reply grows, so releasing tokens early would let a later scan try to rewrite text already on the wire. Holding them also means a blocking verdict ends the stream with nothing sent at all.
+
+### Evaluate volume
+
+Under `incremental_diff` the guardrail defaults to a single end-of-stream evaluate rather than the framework's
+sampled rounds. Because the whole reply is withheld until the end anyway, a sampled round cannot stop anything
+reaching the client; it only re-sends the reply for another evaluate, so a 400-chunk reply at the default
+`streaming_sampling_rate` of 5 would cost 81 calls instead of 1.
+
+The trade-off is that a blocking verdict is reached only once the model has finished generating, so a blocked
+stream is no longer aborted early. Nothing is delivered either way. Set `streaming_end_of_stream_only: false`
+under `guardrail_config` to restore sampled rounds where cutting off a blocked generation matters more than the
+evaluate volume.
+
+### Surface coverage
+
+`incremental_diff` is implemented for OpenAI chat completions streaming only. Every other streamed surface —
+including `/v1/messages` and `/v1/responses` — silently falls back to `block_only`, and the proxy logs a warning
+naming the guardrail when it does.
+
+**What that means on those surfaces:** `block` and `ask` still end the stream, so dangerous content is still
+stopped. What is lost is masking: a `transform` verdict is dropped and the raw model tokens reach the client
+unredacted. Treat `incremental_diff` as a chat-completions feature, and prefer non-streaming requests on
+`/v1/messages` and `/v1/responses` when redaction must be enforced rather than merely attempted.
+
+Extending the diff protocol to those surfaces means implementing the streamed text and tool-call rewrite
+contract in their translation handlers, whose event shapes differ from chat completions. That work belongs in
+the shared streaming layer, not in this guardrail.
+
+Under `incremental_diff`, LiteLLM buffers tool-call deltas until TrustGuard inspects the assembled response. A blocking verdict releases no tool-call arguments or finish signal. Tool calls retain their IDs and order, with transformed arguments written into the buffered deltas before delivery
+
+The default `block_only` mode still forwards tool-call deltas before inspection. Use `incremental_diff` or non-streaming requests when tool calls must be checked before delivery. The example selects `incremental_diff` for inspected text and tool arguments
+
+## References
+
+- [NeuralTrust TrustGuard on LiteLLM](https://docs.neuraltrust.ai/integrations/litellm)
+- [TrustGuard Evaluate API](https://docs.neuraltrust.ai/trustguard/api/evaluate)
+- [TrustGuard collectors](https://docs.neuraltrust.ai/trustguard/concepts/collectors)
+- [LiteLLM Guardrails Documentation](https://docs.litellm.ai/docs/proxy/guardrails/quick_start)

litellm/proxy/guardrails/guardrail_hooks/neuraltrust/__init__.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+from litellm.types.guardrails import SupportedGuardrailIntegrations
+
+from .neuraltrust import NeuralTrustGuardrail
+
+if TYPE_CHECKING:
+    from litellm.types.guardrails import Guardrail, LitellmParams
+
+
+def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> NeuralTrustGuardrail:
+    import litellm
+
+    _callback: Final = NeuralTrustGuardrail(
+        api_base=litellm_params.api_base,
+        api_key=litellm_params.api_key,
+        collector_key=litellm_params.collector_key,
+        unreachable_fallback=litellm_params.unreachable_fallback,
+        timeout=litellm_params.timeout,
+        streaming_transform_mode=litellm_params.streaming_transform_mode,
+        guardrail_name=guardrail["guardrail_name"],
+        event_hook=litellm_params.mode,
+        default_on=litellm_params.default_on,
+    )
+    litellm.logging_callback_manager.add_litellm_callback(_callback)
+    return _callback
+
+
+guardrail_initializer_registry: Final = {
+    SupportedGuardrailIntegrations.NEURALTRUST.value: initialize_guardrail,
+}
+
+guardrail_class_registry: Final = {
+    SupportedGuardrailIntegrations.NEURALTRUST.value: NeuralTrustGuardrail,
+}

litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py

@@ -0,0 +1,514 @@
+"""NeuralTrust TrustGuard native LiteLLM guardrail.
+
+Calls TrustGuard POST /v1/evaluate on pre_call (input) and post_call (output).
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable, Mapping, Sequence
+from itertools import chain, product
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final, Literal, ParamSpec, Protocol, TypeAlias, TypeVar, runtime_checkable
+
+import httpx
+from fastapi import HTTPException
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
+
+from litellm._logging import verbose_proxy_logger
+from litellm.exceptions import Timeout
+from litellm.integrations import custom_guardrail
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    get_session_id_from_request_data,
+)
+from litellm.llms.custom_httpx import http_handler
+from litellm.llms.custom_httpx.http_handler import httpxSpecialProvider
+from litellm.types.guardrails import GuardrailEventHooks, Mode
+from litellm.types.llms.openai import AllMessageValues
+from litellm.types.proxy.guardrails.guardrail_hooks.neuraltrust import (
+    DEFAULT_API_BASE,
+    DEFAULT_TIMEOUT,
+    NeuralTrustGuardrailConfigModel,
+)
+from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+
+EVALUATE_PATH: Final = "/v1/evaluate"
+CONSUMER_ID_KEYS: Final = (
+    ("user_api_key_alias", "user_api_key_key_alias"),
+    ("user_api_key_user_email",),
+    ("user_api_key_user_id",),
+    ("user_api_key_team_alias",),
+)
+METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
+EMPTY_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
+STATUS_BLOCK: Final = "block"
+STATUS_ASK: Final = "ask"
+STATUS_TRANSFORM: Final = "transform"
+STATUS_REPORT: Final = "report"
+STATUS_ALLOW: Final = "allow"
+BLOCKING_STATUSES: Final = frozenset({STATUS_BLOCK, STATUS_ASK})
+KNOWN_STATUSES: Final = frozenset({STATUS_ALLOW, STATUS_TRANSFORM, STATUS_REPORT, *BLOCKING_STATUSES})
+UNREACHABLE_HTTP_STATUSES: Final = frozenset({502, 504})
+TRANSFORM_MISSING: Final = "TrustGuard transform missing payload"
+
+
+_Parameters: Final = ParamSpec("_Parameters")
+_Result: Final = TypeVar("_Result")
+
+
+@runtime_checkable
+class _LoggingDecorator(Protocol):
+    def __call__(self, function: Callable[_Parameters, _Result]) -> Callable[_Parameters, _Result]: ...
+
+
+_DECORATOR_ADAPTER: Final[TypeAdapter[_LoggingDecorator]] = TypeAdapter(
+    _LoggingDecorator, config=ConfigDict(arbitrary_types_allowed=True)
+)
+log_guardrail_information: Final = _DECORATOR_ADAPTER.validate_python(
+    vars(custom_guardrail)["log_guardrail_information"]
+)
+
+
+_GuardrailInitializer: TypeAlias = Callable[
+    [
+        CustomGuardrail,
+        str | None,
+        list[GuardrailEventHooks],
+        GuardrailEventHooks | Mode | str | Sequence[str] | None,
+        bool,
+    ],
+    None,
+]
+
+
+class _TrustGuardClient(Protocol):
+    async def post(
+        self,
+        url: str,
+        *,
+        json: dict[str, object],
+        headers: dict[str, str],
+        timeout: float | httpx.Timeout | None,
+    ) -> httpx.Response: ...
+
+
+class _TrustGuardUnreachable(Exception):
+    """Transport or availability failure; eligible for unreachable_fallback."""
+
+
+def _metadata(block: object) -> Mapping[str, object]:
+    try:
+        return METADATA_ADAPTER.validate_python(block)
+    except ValidationError:
+        return EMPTY_METADATA
+
+
+def _consumer_id(request_data: Mapping[str, object]) -> str | None:
+    blocks: Final = tuple(_metadata(request_data.get(source)) for source in ("litellm_metadata", "metadata"))
+    candidates: Final = (block.get(name) for name, block in product(chain.from_iterable(CONSUMER_ID_KEYS), blocks))
+    return next((value for value in candidates if isinstance(value, str) and value), None)
+
+
+def _message_text(message: Mapping[str, object]) -> str:
+    content: Final = message.get("content")
+    return content if isinstance(content, str) else ""
+
+
+def _copy_message(value: object) -> Mapping[str, object] | None:
+    if not isinstance(value, Mapping):
+        return None
+    return {str(key): item for key, item in TypeAdapter(Mapping[object, object]).validate_python(value).items()}
+
+
+def _copy_messages(messages: Sequence[object]) -> tuple[Mapping[str, object], ...] | None:
+    copied: Final = tuple(copy for message in messages if (copy := _copy_message(message)) is not None)
+    return copied if len(copied) == len(messages) else None
+
+
+def _texts_from_messages(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
+    return tuple(_message_text(message) for message in messages)
+
+
+def _tool_calls_in_message(message: Mapping[str, object]) -> tuple[object, ...] | None:
+    raw: Final = message.get("tool_calls")
+    if raw is None:
+        return None
+    if not isinstance(raw, list):
+        raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
+    return TypeAdapter(tuple[object, ...]).validate_python(raw)
+
+
+def _tool_calls_from_messages(messages: Sequence[Mapping[str, object]]) -> tuple[object, ...] | None:
+    groups: Final = tuple(_tool_calls_in_message(message) for message in messages)
+    if all(group is None for group in groups):
+        return None
+    return tuple(chain.from_iterable(group for group in groups if group is not None))
+
+
+def _rewrite_last_user_message(
+    messages: Sequence[Mapping[str, object]],
+    redacted: str,
+) -> tuple[Mapping[str, object], ...]:
+    user_indices: Final = tuple(index for index, message in enumerate(messages) if message.get("role") == "user")
+    target: Final = user_indices[-1] if user_indices else len(messages) - 1
+    return tuple(
+        {**message, "content": redacted} if index == target else dict(message) for index, message in enumerate(messages)
+    )
+
+
+def _model_name(
+    inputs: GenericGuardrailAPIInputs,
+    logging_obj: LiteLLMLoggingObj | None,
+) -> str:
+    if logging_obj is not None and logging_obj.model:
+        return str(logging_obj.model)
+    return str(inputs.get("model") or "")
+
+
+def _assistant_message(text: str | None, tool_calls: object) -> Mapping[str, object]:
+    if tool_calls:
+        return {"role": "assistant", "content": text, "tool_calls": tool_calls}
+    return {"role": "assistant", "content": text}
+
+
+def _assistant_messages(texts: Sequence[str], tool_calls: object) -> tuple[Mapping[str, object], ...]:
+    if not texts:
+        return (_assistant_message(None if tool_calls else "", tool_calls),)
+    last: Final = len(texts) - 1
+    return tuple(_assistant_message(text, tool_calls if index == last else None) for index, text in enumerate(texts))
+
+
+def _sent_messages(
+    inputs: GenericGuardrailAPIInputs,
+    input_type: Literal["request", "response"],
+) -> Sequence[Mapping[str, object]]:
+    if input_type == "response":
+        return _assistant_messages(tuple(inputs.get("texts") or ()), inputs.get("tool_calls"))
+    structured: Final = inputs.get("structured_messages")
+    if structured:
+        return structured
+    return tuple({"role": "user", "content": text} for text in (inputs.get("texts") or ()))
+
+
+def _inputs_with_messages(
+    inputs: GenericGuardrailAPIInputs,
+    messages: Sequence[Mapping[str, object]],
+    *,
+    replace_tool_calls: bool,
+) -> GenericGuardrailAPIInputs:
+    extracted: Final = _tool_calls_from_messages(messages) if replace_tool_calls else None
+    original_tool_calls: Final = inputs.get("tool_calls")
+    if extracted is not None and original_tool_calls is not None and len(extracted) != len(original_tool_calls):
+        raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
+    try:
+        normalized_messages: Final = tuple(
+            {**message, "content": ""} if message.get("role") == "user" and message.get("content") is None else message
+            for message in messages
+        )
+        structured_messages: Final = TypeAdapter(list[AllMessageValues]).validate_python(normalized_messages)
+        rewritten_tool_calls: Final = (
+            TypeAdapter(list[ChatCompletionMessageToolCall]).validate_python(extracted)
+            if extracted is not None
+            else None
+        )
+    except ValidationError as exc:
+        raise HTTPException(status_code=400, detail=TRANSFORM_MISSING) from exc
+    merged: Final[GenericGuardrailAPIInputs] = {
+        **inputs,
+        "structured_messages": structured_messages,
+    }
+    rebuilt: Final[GenericGuardrailAPIInputs] = (
+        {**merged, "texts": list(_texts_from_messages(messages))} if inputs.get("texts") else merged
+    )
+    if extracted is None:
+        return rebuilt
+    return {**rebuilt, "tool_calls": rewritten_tool_calls} if rewritten_tool_calls is not None else rebuilt
+
+
+class NeuralTrustGuardrail(CustomGuardrail):
+    """LiteLLM hook that evaluates prompts and completions with TrustGuard."""
+
+    @staticmethod
+    def get_config_model() -> type[NeuralTrustGuardrailConfigModel[BaseModel]]:
+        return NeuralTrustGuardrailConfigModel
+
+    @classmethod
+    def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:  # mutable-ok: CustomGuardrail contract
+        return [
+            GuardrailEventHooks.pre_call,
+            GuardrailEventHooks.post_call,
+        ]
+
+    def __init__(
+        self,
+        api_base: str | None = None,
+        api_key: str | None = None,
+        collector_key: str | None = None,
+        unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
+        timeout: float | None = None,
+        streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None,
+        guardrail_name: str | None = None,
+        event_hook: GuardrailEventHooks | Mode | str | Sequence[str] | None = None,
+        default_on: bool | None = None,
+    ) -> None:
+        factory_adapter: Final[TypeAdapter[Callable[[httpxSpecialProvider], _TrustGuardClient]]] = TypeAdapter(
+            Callable[[httpxSpecialProvider], _TrustGuardClient]
+        )
+        client_factory: Final = factory_adapter.validate_python(vars(http_handler)["get_async_httpx_client"])
+        self.async_handler = client_factory(httpxSpecialProvider.GuardrailCallback)
+        self.api_base = (api_base or os.environ.get("TRUSTGUARD_API_BASE") or DEFAULT_API_BASE).rstrip("/")
+        self.api_key = api_key or os.environ.get("TRUSTGUARD_API_KEY") or ""
+        if not self.api_key:
+            raise ValueError(
+                "TrustGuard API key is required. Set TRUSTGUARD_API_KEY or pass api_key in litellm_params."
+            )
+        self.collector_key = collector_key or os.environ.get("TRUSTGUARD_COLLECTOR_KEY") or ""
+        self.unreachable_fallback: Literal["fail_closed", "fail_open"] = unreachable_fallback
+        resolved_timeout: Final = DEFAULT_TIMEOUT if timeout is None else float(timeout)
+        if resolved_timeout <= 0:
+            raise ValueError("TrustGuard timeout must be a positive number of seconds.")
+        self.timeout = resolved_timeout
+        # Read off the instance by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook.
+        self.streaming_transform_mode: Literal["block_only", "incremental_diff"] = (
+            streaming_transform_mode or "block_only"
+        )
+        # Read off the instance by UnifiedLLMGuardrails.resolve_streaming_flag, which lets
+        # guardrail_config and optional_params override it.
+        #
+        # Under incremental_diff _with_holdback withholds the whole reply on every round, so
+        # nothing reaches the client until the end-of-stream round flushes it. Sampled rounds
+        # therefore cannot reduce leakage: each one re-sends the entire reply for another
+        # evaluate (a 400-chunk reply at the default sampling_rate of 5 costs 81 calls) and buys
+        # only an earlier stream abort on a blocking verdict, which saves upstream generation
+        # rather than protecting the client. One end-of-stream scan is the honest cost here.
+        # Set streaming_end_of_stream_only to False in guardrail_config to restore sampled
+        # rounds when ending a blocked stream early matters more than the evaluate volume.
+        self.streaming_end_of_stream_only: bool = self.streaming_transform_mode == "incremental_diff"
+        initializer_adapter: Final[TypeAdapter[_GuardrailInitializer]] = TypeAdapter(_GuardrailInitializer)
+        initialize: Final = initializer_adapter.validate_python(vars(CustomGuardrail)["__init__"])
+        initialize(self, guardrail_name, self.get_supported_event_hooks(), event_hook, bool(default_on))
+
+    @log_guardrail_information
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict[str, object],  # mutable-ok: CustomGuardrail.apply_guardrail contract
+        input_type: Literal["request", "response"],
+        logging_obj: LiteLLMLoggingObj | None = None,
+    ) -> GenericGuardrailAPIInputs:
+        return self._with_holdback(
+            await self._evaluate(inputs, request_data, input_type, logging_obj),
+            input_type,
+        )
+
+    async def _evaluate(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict[str, object],  # mutable-ok: CustomGuardrail.apply_guardrail contract
+        input_type: Literal["request", "response"],
+        logging_obj: LiteLLMLoggingObj | None,
+    ) -> GenericGuardrailAPIInputs:
+        body: Final = self._evaluate_body(inputs, request_data, input_type, logging_obj)
+        try:
+            result: Final = await self._call_evaluate(body)
+        except HTTPException:
+            raise
+        except _TrustGuardUnreachable as exc:
+            return self._handle_unreachable(inputs, exc)
+
+        status: Final = result["status"]
+        if status in BLOCKING_STATUSES:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": "Violated guardrail policy",
+                    "neuraltrust_guardrail_response": "Blocked by NeuralTrust TrustGuard.",
+                    "verdict": status,
+                    "trace_id": result.get("trace_id"),
+                    "request_id": result.get("request_id"),
+                },
+            )
+        if status == STATUS_TRANSFORM:
+            return self._apply_transform(inputs, result.get("transformed_payload"), input_type=input_type)
+        if status == STATUS_REPORT:
+            verbose_proxy_logger.info("TrustGuard report-only findings trace_id=%s", result.get("trace_id"))
+        return inputs
+
+    def _with_holdback(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        input_type: Literal["request", "response"],
+    ) -> GenericGuardrailAPIInputs:
+        """Withhold the whole reply from each streaming round under incremental_diff.
+
+        TrustGuard re-reads the full reply every round and its redaction spans move as the reply grows, so a
+        later round can rewrite text an earlier one already streamed. The engine cannot retract streamed bytes
+        and answers that with stream_transform_underflow, so nothing is released until the end-of-stream round
+        forces the holdback to zero and flushes the final redacted reply.
+        """
+        texts: Final = inputs.get("texts")
+        if input_type == "request" or self.streaming_transform_mode != "incremental_diff" or not texts:
+            return inputs
+        held: Final[GenericGuardrailAPIInputs] = {
+            **inputs,
+            "stream_holdback_chars": [len(text) for text in texts],
+        }
+        return held
+
+    def _evaluate_body(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict[str, object],  # mutable-ok: CustomGuardrail.apply_guardrail contract
+        input_type: Literal["request", "response"],
+        logging_obj: LiteLLMLoggingObj | None,
+    ) -> dict[str, object]:  # mutable-ok: outbound JSON
+        session_id: Final = get_session_id_from_request_data(request_data)
+        consumer_id: Final = _consumer_id(request_data)
+        return {
+            "payload": self._payload(inputs, input_type),
+            "direction": "input" if input_type == "request" else "output",
+            "protocol": "llm",
+            "attributes": {
+                "content_type": "application/json",
+                "model": {"name": _model_name(inputs, logging_obj)},
+            },
+            **({"collector_key": self.collector_key} if self.collector_key else {}),
+            **({"session_id": session_id} if session_id else {}),
+            **({"consumer_id": consumer_id} if consumer_id is not None else {}),
+        }
+
+    @staticmethod
+    def _payload(
+        inputs: GenericGuardrailAPIInputs,
+        input_type: Literal["request", "response"],
+    ) -> Mapping[str, object]:
+        messages: Final = _sent_messages(inputs, input_type)
+        tools: Final = inputs.get("tools") if input_type == "request" else None
+        if tools:
+            return {"messages": messages, "tools": tools}
+        return {"messages": messages}
+
+    async def _call_evaluate(self, body: dict[str, object]) -> dict[str, object]:  # mutable-ok: TrustGuard JSON
+        url: Final = f"{self.api_base}{EVALUATE_PATH}"
+        headers: Final = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        try:
+            response: Final = await self.async_handler.post(
+                url,
+                json=body,
+                headers=headers,
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+        except Timeout as exc:
+            raise _TrustGuardUnreachable(exc) from exc
+        except httpx.HTTPStatusError as exc:
+            status_code: Final = exc.response.status_code
+            if status_code == 503:
+                raise HTTPException(
+                    status_code=503,
+                    detail="TrustGuard entitlements unavailable",
+                ) from exc
+            if status_code == 429:
+                rate_limit_reason: Final = exc.response.headers.get("X-RateLimit-Reason")
+                verbose_proxy_logger.warning(
+                    "TrustGuard rate limit reached; failing closed. reason=%s",
+                    rate_limit_reason or "unspecified",
+                )
+                raise HTTPException(
+                    status_code=503,
+                    detail="TrustGuard rate limit or quota exceeded",
+                ) from exc
+            if status_code in (401, 403):
+                raise HTTPException(
+                    status_code=status_code,
+                    detail="TrustGuard authentication failed",
+                ) from exc
+            if status_code in UNREACHABLE_HTTP_STATUSES:
+                raise _TrustGuardUnreachable(exc) from exc
+            raise HTTPException(
+                status_code=503,
+                detail="TrustGuard request failed",
+            ) from exc
+        except httpx.RequestError as exc:
+            raise _TrustGuardUnreachable(exc) from exc
+
+        try:
+            parsed: Final[object] = response.json()
+        except ValueError as exc:
+            raise _TrustGuardUnreachable("TrustGuard returned non-JSON body") from exc
+        if not isinstance(parsed, dict):
+            raise HTTPException(status_code=503, detail="TrustGuard returned an invalid response")
+        typed_response: Final = METADATA_ADAPTER.validate_python(parsed)
+        status: Final = typed_response.get("status")
+        if not isinstance(status, str) or status.lower() not in KNOWN_STATUSES:
+            raise HTTPException(status_code=503, detail="TrustGuard returned an unknown verdict")
+        return {**typed_response, "status": status.lower()}
+
+    def _handle_unreachable(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        error: Exception,
+    ) -> GenericGuardrailAPIInputs:
+        if self.unreachable_fallback == "fail_open":
+            verbose_proxy_logger.critical(
+                "TrustGuard unreachable (fail-open): %s",
+                error,
+                exc_info=error,
+            )
+            return inputs
+        verbose_proxy_logger.error("TrustGuard unreachable (fail-closed): %s", error)
+        raise HTTPException(
+            status_code=503,
+            detail="TrustGuard guardrail service unreachable",
+        ) from error
+
+    @staticmethod
+    def _apply_transform(
+        inputs: GenericGuardrailAPIInputs,
+        transformed: object,
+        *,
+        input_type: Literal["request", "response"],
+    ) -> GenericGuardrailAPIInputs:
+        if not isinstance(transformed, Mapping):
+            raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
+
+        try:
+            typed_transform: Final = METADATA_ADAPTER.validate_python(transformed)
+        except ValidationError as exc:
+            raise HTTPException(status_code=400, detail=TRANSFORM_MISSING) from exc
+        raw_messages: Final = typed_transform.get("messages")
+        if isinstance(raw_messages, list) and raw_messages:
+            rewritten_messages: Final = _copy_messages(TypeAdapter(tuple[object, ...]).validate_python(raw_messages))
+            if rewritten_messages is None or len(rewritten_messages) != len(_sent_messages(inputs, input_type)):
+                raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
+            return _inputs_with_messages(inputs, rewritten_messages, replace_tool_calls=True)
+
+        raw_input: Final = typed_transform.get("input")
+        if not isinstance(raw_input, str) or not raw_input:
+            raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
+
+        # structured_messages on a reply scan is the request conversation the framework attached for context,
+        # so a scalar payload rewrites the scanned reply text instead.
+        original_messages: Final = inputs.get("structured_messages") if input_type == "request" else None
+        if isinstance(original_messages, list) and original_messages:
+            copied: Final = _copy_messages(original_messages)
+            if copied is None:
+                raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
+            return _inputs_with_messages(
+                inputs,
+                _rewrite_last_user_message(copied, raw_input),
+                replace_tool_calls=False,
+            )
+
+        original_texts: Final = tuple(inputs.get("texts") or ())
+        if not original_texts:
+            raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
+        rewritten_texts: Final = (*original_texts[:-1], raw_input)
+        return {**inputs, "texts": list(rewritten_texts)}

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime
 
 import anyio
 from fastapi import HTTPException
+from pydantic import InstanceOf, TypeAdapter
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
@@ -49,6 +50,7 @@ if TYPE_CHECKING:
 A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message)
 
 GUARDRAIL_NAME: Final = "unified_llm_guardrails"
+GUARDRAIL_ADAPTER: Final[TypeAdapter[CustomGuardrail | None]] = TypeAdapter(InstanceOf[CustomGuardrail] | None)
 
 _RequestData: TypeAlias = dict[str, object]
 
@@ -531,9 +533,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         self,
         *,
         reference_chunk: object,
-        mutated_text_per_choice: dict[int, str],
+        mutated_text_per_choice: Mapping[int, str],
         emitted_text_per_choice: dict[int, str],
-        holdback_per_choice: dict[int, int],
+        holdback_per_choice: Mapping[int, int],
         finish_reason_per_choice: dict[int, str | None],
         held_chars_per_choice: dict[int, int],
         is_final: bool,
@@ -769,46 +771,13 @@ class UnifiedLLMGuardrails(CustomLogger):
 
         saw_tool_calls = False
         saw_text_content = False
-        tool_calls_released = False  # rebind-ok: set once a raw tool call reaches the client unscanned
-        end_of_stream_inspection_started = False  # rebind-ok: set once the end-of-stream inspection owns the verdict
 
         try:
             async for item in response:
-                # v1 transforms only text. A chunk carrying tool_calls is passed
-                # through raw so function-calling turns are not dropped, but ONLY
-                # its tool-call fields are forwarded: content is stripped so any
-                # response text (in the same delta, or in another choice of an n>1
-                # chunk) can never bypass the transform. The original chunk is kept
-                # in responses_so_far so its text is still accumulated + redacted +
-                # emitted as synthetic deltas, and so the guardrail inspects the
-                # assembled tool calls at end of stream (see the block inspection
-                # below), matching block_only. finish_reason rides on the raw
-                # tool-only chunk, so it is not recorded for the text flush.
                 if self._chunk_has_tool_calls(item):
                     saw_tool_calls = True
                     responses_so_far.append(item)
                     last_chunk = item
-                    # Fix #3 — flush accumulated text BEFORE the tool-call
-                    # passthrough. Without this, a stream of text chunks that
-                    # hasn't yet hit a sampled round can be trailed by a
-                    # tool-call chunk carrying finish_reason="tool_calls"; an
-                    # SSE-compliant client stops reading at that finish_reason
-                    # and drops the end-of-stream text flush that would follow.
-                    if saw_text_content:
-                        async for out in _round(item, is_final=False):
-                            yield out
-                    # Fix #1 — pass finish_reason_per_choice into the
-                    # passthrough so a mixed content+tool_call chunk defers its
-                    # finish_reason to the final text terminator (see the
-                    # _tool_call_passthrough_chunk docstring).
-                    tool_only = self._tool_call_passthrough_chunk(
-                        item,
-                        finish_reason_per_choice=finish_reason_per_choice,
-                        held_choices=_held_choices(held_chars_per_choice),
-                    )
-                    responses_yielded.append(tool_only)
-                    tool_calls_released = True
-                    yield tool_only
                     continue
 
                 if self._is_trailing_metadata_chunk(item):
@@ -826,39 +795,60 @@ class UnifiedLLMGuardrails(CustomLogger):
                 # sampled round here would guardrail the same content twice.
                 if (
                     not end_of_stream_only
+                    and not saw_tool_calls
                     and not self._chunk_has_finish_reason(item)
                     and chunk_counter % sampling_rate == 0
                 ):
                     async for out in _round(item, is_final=False):
                         yield out
 
-            # v1 does not transform streamed tool calls, but they must still go
-            # through the guardrail's block decision. Run the block_only inspection
-            # over the full assembled response so tool calls cannot bypass it.
-            #
-            # Pass a deep copy of responses_so_far — the block path routes through
-            # ``_process_streaming_block_only`` which mutates ``delta.content``
-            # in-place on the chunk objects it receives. For an n>1 chunk carrying
-            # text on one choice and tool_calls (with finish_reason) on another,
-            # ``has_stream_ended`` reads ``choices[0]`` alone and can miss the
-            # terminal signal, letting the block path rewrite the raw accumulator.
-            # The subsequent final ``_round`` would then re-read the already-mutated
-            # text, producing double-application for a non-idempotent guardrail or a
-            # ``stream_transform_underflow`` 400 from mismatched prefixes. A shallow
-            # list copy wouldn't help — the mutation is on the chunk objects
-            # themselves — so we deepcopy.
-            end_of_stream_inspection_started = True
             if saw_tool_calls:
-                async for out in self._inspect_full_response_for_block(
+                inspected_responses: Final = copy.deepcopy(responses_so_far)
+                async for out in self._inspect_full_response(
                     endpoint_translation=endpoint_translation,
                     guardrail_to_apply=guardrail_to_apply,
                     request_data=request_data,
                     user_api_key_dict=user_api_key_dict,
-                    responses_so_far=copy.deepcopy(responses_so_far),
+                    responses_so_far=inspected_responses,
                     responses_yielded=responses_yielded,
                 ):
                     yield out
 
+                if saw_text_content:
+                    async for out in _round(last_chunk, is_final=False):
+                        yield out
+                tool_chunks: Final = tuple(
+                    self._tool_call_passthrough_chunk(
+                        buffered_item,
+                        finish_reason_per_choice=finish_reason_per_choice,
+                        held_choices=_held_choices(held_chars_per_choice),
+                    )
+                    for buffered_item in inspected_responses
+                    if self._chunk_has_tool_calls(buffered_item)
+                )
+
+                tail_chunks: Final = tuple(
+                    [
+                        chunk
+                        async for chunk in self._checked_stream_tail(
+                            last_chunk=last_chunk,
+                            final_round=_round,
+                            responses_so_far=responses_so_far,
+                            responses_yielded=responses_yielded,
+                        )
+                    ]
+                )
+                if tail_chunks and isinstance(tail_chunks[-1], _StreamTerminated):
+                    for error_chunk in tail_chunks[:-1]:
+                        yield error_chunk
+                    return
+                for tool_only in tool_chunks:
+                    responses_yielded.append(tool_only)
+                    yield tool_only
+                for tail_chunk in tail_chunks:
+                    yield tail_chunk
+                return
+
             async for out in self._emit_stream_tail(
                 last_chunk=last_chunk,
                 final_round=_round,
@@ -868,37 +858,25 @@ class UnifiedLLMGuardrails(CustomLogger):
                 yield out
         except _StreamTerminated:
             return
-        except (GeneratorExit, asyncio.CancelledError):
-            await self._scan_uninspected_tool_calls_after_disconnect(
-                uninspected=tool_calls_released and not end_of_stream_inspection_started and not terminated.is_set(),
-                endpoint_translation=endpoint_translation,
-                responses_released=responses_yielded,
-                guardrail_to_apply=guardrail_to_apply,
-                user_api_key_dict=user_api_key_dict,
-                request_data=request_data,
-            )
-            raise
 
-    @staticmethod
-    async def _scan_uninspected_tool_calls_after_disconnect(
+    async def _checked_stream_tail(
+        self,
         *,
-        uninspected: bool,
-        endpoint_translation: _EndpointTranslation,
-        responses_released: Sequence[object],
-        guardrail_to_apply: CustomGuardrail,
-        user_api_key_dict: UserAPIKeyAuth,
-        request_data: _RequestData,
-    ) -> None:
-        if not uninspected:
-            return
-        await UnifiedLLMGuardrails._scan_released_stream_after_disconnect(
-            endpoint_translation=endpoint_translation,
-            responses_released=responses_released,
-            last_scan_key=None,
-            guardrail_to_apply=guardrail_to_apply,
-            user_api_key_dict=user_api_key_dict,
-            request_data=request_data,
-        )
+        last_chunk: object | None,
+        final_round: Callable[[object, bool], AsyncGenerator[object, None]],
+        responses_so_far: Sequence[object],
+        responses_yielded: list[object],
+    ) -> AsyncGenerator[object, None]:
+        try:
+            async for chunk in self._emit_stream_tail(
+                last_chunk=last_chunk,
+                final_round=final_round,
+                responses_so_far=responses_so_far,
+                responses_yielded=responses_yielded,
+            ):
+                yield chunk
+        except _StreamTerminated as exc:
+            yield exc
 
     async def _emit_stream_tail(
         self,
@@ -917,29 +895,21 @@ class UnifiedLLMGuardrails(CustomLogger):
             responses_yielded.append(trailing)
             yield trailing
 
-    async def _inspect_full_response_for_block(
+    async def _inspect_full_response(
         self,
         *,
         endpoint_translation: _EndpointTranslation,
         guardrail_to_apply: CustomGuardrail,
-        request_data: dict,
+        request_data: _RequestData,
         user_api_key_dict: UserAPIKeyAuth,
         responses_so_far: Sequence[object],
         responses_yielded: Sequence[object],
     ) -> AsyncGenerator[object, None]:
-        """Run the block-only guardrail inspection over the full assembled
-        response (text + tool calls) so nothing bypasses the block decision.
-
-        The guardrail's returned transforms are discarded here (v1 does not
-        transform tool calls); only its block decision matters. A block is
-        surfaced the same way as elsewhere: ModifyResponseException terminates the
-        stream via the shared block handler; a GenericGuardrailAPI block raises and
-        propagates, matching block_only.
-        """
         from litellm.integrations.custom_guardrail import ModifyResponseException
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
         try:
-            with anyio.CancelScope(shield=bool(responses_yielded)):
+            with anyio.CancelScope(shield=True):
                 await endpoint_translation.process_output_streaming_response(
                     responses_so_far=responses_so_far,
                     guardrail_to_apply=guardrail_to_apply,
@@ -947,7 +917,10 @@ class UnifiedLLMGuardrails(CustomLogger):
                     user_api_key_dict=user_api_key_dict,
                     request_data=request_data,
                     stream_transform_sink=None,
+                    deliver_ended_stream_rewrites=True,
                 )
+        except UndeliverableStreamRewrite as exc:
+            raise HTTPException(status_code=400, detail="Guardrail stream rewrite could not be applied") from exc
         except ModifyResponseException as e:
             if e.original_response is None:
                 e.original_response = responses_so_far
@@ -1131,12 +1104,15 @@ class UnifiedLLMGuardrails(CustomLogger):
         # litellm.integrations.custom_guardrail.
         from litellm.integrations.custom_guardrail import ModifyResponseException
 
-        if guardrail_to_apply is None:
-            guardrail_to_apply = request_data.pop("guardrail_to_apply", None)
+        resolved_guardrail: Final[CustomGuardrail | None] = (
+            guardrail_to_apply
+            if guardrail_to_apply is not None
+            else GUARDRAIL_ADAPTER.validate_python(request_data.pop("guardrail_to_apply", None))
+        )
         typed_request_data: Final[_RequestData] = request_data
 
         def _streaming_flag(name: str, default: object) -> Any:
-            return self.resolve_streaming_flag(guardrail_to_apply, name, default)
+            return self.resolve_streaming_flag(resolved_guardrail, name, default)
 
         sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5)
         # Only apply the guardrail at end of stream (not per chunk).
@@ -1159,30 +1135,30 @@ class UnifiedLLMGuardrails(CustomLogger):
 
         if (
             buffer_until_moderated
-            and guardrail_to_apply is not None
-            and getattr(guardrail_to_apply, "mask_response_content", False)
+            and resolved_guardrail is not None
+            and getattr(resolved_guardrail, "mask_response_content", False)
         ):
             verbose_proxy_logger.warning(
                 "UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s "
                 "because mask_response_content=True -- buffered replay would release "
                 "unredacted original chunks instead of the moderated output.",
-                guardrail_to_apply.guardrail_name,
+                resolved_guardrail.guardrail_name,
             )
             buffer_until_moderated = False
 
         if buffer_until_moderated and not release_on_scan:
             end_of_stream_only = True
 
-        if guardrail_to_apply is None:
+        if resolved_guardrail is None:
             async for item in response:
                 yield item
             return
 
         event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.post_call
-        if guardrail_to_apply.should_run_guardrail(data=request_data, event_type=event_type) is not True:
+        if resolved_guardrail.should_run_guardrail(data=request_data, event_type=event_type) is not True:
             verbose_proxy_logger.debug(
                 "UnifiedLLMGuardrails: Post-call streaming scanning disabled for %s",
-                guardrail_to_apply.guardrail_name,
+                resolved_guardrail.guardrail_name,
             )
             async for item in response:
                 yield item
@@ -1202,7 +1178,7 @@ class UnifiedLLMGuardrails(CustomLogger):
             if transform_call_type is not None:
                 async with contextlib.aclosing(
                     self._run_incremental_transform_stream(
-                        guardrail_to_apply=guardrail_to_apply,
+                        guardrail_to_apply=resolved_guardrail,
                         response=response,
                         request_data=typed_request_data,
                         user_api_key_dict=user_api_key_dict,
@@ -1219,7 +1195,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                 "UnifiedLLMGuardrails: streaming_transform_mode=incremental_diff is only supported "
                 "for the OpenAI chat completions streaming path with a resolvable request route; "
                 "falling back to block_only for %s",
-                getattr(guardrail_to_apply, "guardrail_name", None),
+                getattr(resolved_guardrail, "guardrail_name", None),
             )
 
         # Infer call type from first chunk
@@ -1298,7 +1274,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                         verbose_proxy_logger.debug(
                             "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
                             chunk_counter,
-                            guardrail_to_apply.guardrail_name,
+                            resolved_guardrail.guardrail_name,
                         )
                         if buffer_until_moderated:
                             if hold_window:
@@ -1318,7 +1294,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                         "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
                         chunk_counter,
                         sampling_rate,
-                        guardrail_to_apply.guardrail_name,
+                        resolved_guardrail.guardrail_name,
                     )
 
                     original_items = (
@@ -1328,7 +1304,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                     try:
                         await endpoint_translation.process_output_streaming_response(
                             responses_so_far=responses_so_far,
-                            guardrail_to_apply=guardrail_to_apply,
+                            guardrail_to_apply=resolved_guardrail,
                             litellm_logging_obj=request_data.get("litellm_logging_obj"),
                             user_api_key_dict=user_api_key_dict,
                             request_data=request_data,
@@ -1374,7 +1350,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                         verbose_proxy_logger.debug(
                             "Holding %s buffered chunks for guardrail %s: this round could not scan the whole window",
                             len(withheld_items),
-                            guardrail_to_apply.guardrail_name,
+                            resolved_guardrail.guardrail_name,
                         )
                         withheld_items[:] = original_items
                         continue
@@ -1394,7 +1370,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                 verbose_proxy_logger.debug(
                     "Processing final streaming response with all %s chunks for guardrail %s",
                     len(responses_so_far),
-                    guardrail_to_apply.guardrail_name,
+                    resolved_guardrail.guardrail_name,
                 )
 
                 endpoint_translation = mappings[CallTypes(call_type)]()
@@ -1409,11 +1385,11 @@ class UnifiedLLMGuardrails(CustomLogger):
                 end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
                 verdict_settled = True
                 if _is_redundant_scan(end_scan_key, last_scan_key) and not _owes_ended_scan(
-                    endpoint_translation, guardrail_to_apply, responses_so_far, last_scan_key
+                    endpoint_translation, resolved_guardrail, responses_so_far, last_scan_key
                 ):
                     verbose_proxy_logger.debug(
                         "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
-                        guardrail_to_apply.guardrail_name,
+                        resolved_guardrail.guardrail_name,
                     )
                     for buffered_item in buffered_items or ():
                         yield buffered_item
@@ -1426,7 +1402,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                     with anyio.CancelScope(shield=chunks_yielded):
                         await endpoint_translation.process_output_streaming_response(
                             responses_so_far=responses_so_far,
-                            guardrail_to_apply=guardrail_to_apply,
+                            guardrail_to_apply=resolved_guardrail,
                             litellm_logging_obj=request_data.get("litellm_logging_obj"),
                             user_api_key_dict=user_api_key_dict,
                             request_data=request_data,
@@ -1466,17 +1442,12 @@ class UnifiedLLMGuardrails(CustomLogger):
                         yield error_item
         except (GeneratorExit, asyncio.CancelledError):
             translation_class: Final = None if call_type is None else mappings.get(CallTypes(call_type))
-            if (
-                chunks_yielded
-                and not verdict_settled
-                and translation_class is not None
-                and isinstance(guardrail_to_apply, CustomGuardrail)
-            ):
+            if chunks_yielded and not verdict_settled and translation_class is not None:
                 await self._scan_released_stream_after_disconnect(
                     endpoint_translation=translation_class(),
                     responses_released=responses_yielded,
                     last_scan_key=last_scan_key,
-                    guardrail_to_apply=guardrail_to_apply,
+                    guardrail_to_apply=resolved_guardrail,
                     user_api_key_dict=user_api_key_dict,
                     request_data=typed_request_data,
                 )

litellm/types/guardrails.py

@@ -4,7 +4,7 @@ from enum import Enum
 from types import MappingProxyType
 from typing import Annotated, Final, Literal, TypeAlias, cast
 
-from pydantic import ConfigDict, Field, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from pydantic.fields import FieldInfo
 from typing_extensions import ReadOnly, Required, TypedDict
 
@@ -47,6 +47,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ibm import (
 from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import (
     ContentFilterCategoryConfig,
 )
+from litellm.types.proxy.guardrails.guardrail_hooks.neuraltrust import (
+    NeuralTrustGuardrailConfigModel,
+)
 from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import (
     OvalixGuardrailConfigModel,
 )
@@ -87,7 +90,7 @@ Pydantic object defining how to set guardrails on litellm proxy
 guardrails:
   - guardrail_name: "bedrock-pre-guard"
     litellm_params:
-      guardrail: bedrock  # supported values: "akto", "aporia", "bedrock", "lakera", "zscaler_ai_guard"
+      guardrail: bedrock  # supported values: "akto", "aporia", "bedrock", "lakera", "neuraltrust", "zscaler_ai_guard"
       mode: "during_call"
       guardrailIdentifier: ff6ujrregl1q
       guardrailVersion: "DRAFT"
@@ -105,6 +108,7 @@ class SupportedGuardrailIntegrations(Enum):
     PRESIDIO = "presidio"
     HIDE_SECRETS = "hide-secrets"
     HIDDENLAYER = "hiddenlayer"
+    NEURALTRUST = "neuraltrust"
     AIM = "aim"
     CATO_NETWORKS = "cato_networks"
     PANGEA = "pangea"
@@ -1193,7 +1197,7 @@ class BaseLitellmParams(ContentFilterConfigModel):  # works for new and patch up
         default="fail_closed",
         description=(
             "Behavior when a guardrail endpoint is unreachable due to network errors. "
-            "Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. "
+            "Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', 'typesafe', and 'neuraltrust'. "
             "'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. "
             "Also implemented by guardrail='decision_model'."
         ),
@@ -1211,6 +1215,17 @@ class BaseLitellmParams(ContentFilterConfigModel):  # works for new and patch up
         description="Maximum decisions calls in flight at once on the guardrail in each proxy worker. Implemented by guardrail='decision_model'.",
     )
 
+    streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = Field(
+        default=None,
+        description=(
+            "Whether a guardrail's text rewrite reaches a streaming client. "
+            "Implemented by guardrail='prompt_security' and 'neuraltrust'; generic_guardrail_api takes the "
+            "same setting under optional_params. 'block_only' (default) streams the raw model chunks, so a "
+            "block still ends the stream but rewrites are dropped. 'incremental_diff' withholds those chunks "
+            "and streams the guardrail's rewritten text instead. OpenAI chat completions streaming only."
+        ),
+    )
+
     forward_messages_as_sent: bool | None = Field(
         default=False,
         description=(
@@ -1395,6 +1410,7 @@ class LitellmParams(  # pyright: ignore[reportIncompatibleVariableOverride]  # o
     QualifireGuardrailConfigModel,
     BlockCodeExecutionGuardrailConfigModel,
     HiddenlayerGuardrailConfigModel,
+    NeuralTrustGuardrailConfigModel[BaseModel],
     QostodianNexusConfigModel,
     VigilGuardGuardrailConfigModel,
     SingulrGuardrailConfigModel,

litellm/types/proxy/guardrails/guardrail_hooks/neuraltrust.py

@@ -0,0 +1,68 @@
+from typing import Final, Literal, TypeVar
+
+from pydantic import BaseModel, Field
+
+from .base import GuardrailConfigModel
+
+DEFAULT_API_BASE: Final = "https://trustguard.neuraltrust.ai"
+DEFAULT_TIMEOUT: Final = 5.0
+
+
+_Config: Final = TypeVar("_Config", bound=BaseModel)
+
+
+class NeuralTrustGuardrailConfigModel(GuardrailConfigModel[_Config]):
+    """Config for the NeuralTrust TrustGuard native LiteLLM hook."""
+
+    api_key: str | None = Field(
+        default=None,
+        description=("TrustGuard API key (tgk_...). If not provided, TRUSTGUARD_API_KEY is checked."),
+    )
+
+    api_base: str | None = Field(
+        default=None,
+        description=("TrustGuard API base URL. Default https://trustguard.neuraltrust.ai. Env: TRUSTGUARD_API_BASE."),
+    )
+
+    collector_key: str | None = Field(
+        default=None,
+        description=(
+            "TrustGuard collector key (tgcol_...). Optional when the API key is bound to a "
+            "collector. Env: TRUSTGUARD_COLLECTOR_KEY."
+        ),
+    )
+
+    unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
+        default="fail_closed",
+        description=(
+            "What to do on transport failures (connect errors, timeouts, HTTP 502/504). "
+            "'fail_closed' blocks the request; 'fail_open' allows it. "
+            "HTTP 503 entitlements, 401/403, other 4xx/5xx, unknown verdicts, and "
+            "unusable transform payloads always fail closed. "
+            "'fail_open' means the request bypasses TrustGuard entirely."
+        ),
+    )
+
+    timeout: float | None = Field(
+        default=DEFAULT_TIMEOUT,
+        gt=0.0,
+        description=(
+            "Seconds to wait for each TrustGuard evaluate call before it counts as a "
+            "transport failure and unreachable_fallback applies. Default 5."
+        ),
+    )
+
+    streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = Field(
+        default=None,
+        description=(
+            "How a `transform` verdict reaches a streaming client. `block_only` (default) streams the raw "
+            "model chunks, so `block` and `ask` still end the stream but the redacted text is dropped. "
+            "`incremental_diff` withholds the model chunks and streams TrustGuard's rewritten text and tool arguments: "
+            "the reply arrives once the end-of-stream evaluate returns, and a blocking verdict ends the "
+            "stream with nothing already sent. OpenAI chat completions streaming only."
+        ),
+    )
+
+    @staticmethod
+    def ui_friendly_name() -> str:
+        return "NeuralTrust"

tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py

@@ -7,7 +7,7 @@ with guardrail transformations, including tool calls.
 
 import json
 from collections.abc import Mapping
-from typing import Any, Literal, Optional
+from typing import Any, Final, Literal, Optional
 
 import pytest
 
@@ -1395,29 +1395,65 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
         )
 
         assert [
-            (tool_call["id"], tool_call["function"]["arguments"])
-            for tool_call in guardrail.seen_inputs[-1]["tool_calls"]
-        ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')]
+            [(tool_call["id"], tool_call["function"]["arguments"]) for tool_call in inputs["tool_calls"]]
+            for inputs in guardrail.seen_inputs
+        ] == [[("call_1", '{"fruit": "persimmon"}')], [("call_2", '{"fruit": "durian"}')]]
 
     @pytest.mark.asyncio
-    async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self):
-        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+    @pytest.mark.parametrize("transform", [False, True])
+    async def test_each_choice_refreshes_shared_guardrail_response_context(self, transform: bool) -> None:
+        from fastapi import HTTPException
+        from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
+        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
 
-        handler = OpenAIChatCompletionsHandler()
-        chunks = self._two_choice_tool_call_stream_chunks()
+        class ContextGuardrail(CustomGuardrail):
+            async def apply_guardrail(
+                self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
+                input_type: Literal["request", "response"], logging_obj: object = None,
+            ) -> GenericGuardrailAPIInputs:
+                response: Final = request_data["response"]
+                assert isinstance(response, ModelResponse)
+                request_data["visited_choices"] = (*request_data.get("visited_choices", ()), response.choices[0].index)
+                if response.choices[0].message.content == "forbidden":
+                    raise HTTPException(status_code=400, detail="later choice blocked")
+                return inputs
 
-        with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised:
-            await handler.process_output_streaming_response(
-                responses_so_far=chunks,
-                guardrail_to_apply=MockGuardrail(guardrail_name="test"),
-                litellm_logging_obj=None,
-                deliver_ended_stream_rewrites=True,
+        chunks: Final = [ModelResponseStream(choices=[StreamingChoices(
+            index=index, delta=Delta(content=text, tool_calls=[{
+                "index": 0, "id": f"call_{index}", "type": "function",
+                "function": {"name": "contact", "arguments": "{}"},
+            }]), finish_reason="tool_calls",
+        ) for index, text in enumerate(("allowed", "forbidden"))])]
+        request_data: Final = {"metadata": {"trace": "retained"}}
+        with pytest.raises(HTTPException, match="later choice blocked"):
+            await OpenAIChatCompletionsHandler().process_output_streaming_response(
+                responses_so_far=chunks, guardrail_to_apply=ContextGuardrail(guardrail_name="context"),
+                request_data=request_data, deliver_ended_stream_rewrites=True,
+                stream_transform_sink=StreamTransformSink() if transform else None,
             )
+        assert request_data["visited_choices"] == (0, 1)
+        assert request_data["metadata"]["trace"] == "retained"
+        assert request_data["responses"] is chunks
+        restored: Final = request_data["response"]
+        assert isinstance(restored, ModelResponse)
+        assert tuple(choice.index for choice in restored.choices) == (0, 1)
 
-        assert raised.value.guardrail_name == "test"
-        assert raised.value.reason == (
-            "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams"
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None:
+        handler: Final = OpenAIChatCompletionsHandler()
+        chunks: Final = self._two_choice_tool_call_stream_chunks()
+        await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=MockGuardrail(guardrail_name="test"),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
         )
+        arguments: Final = tuple(
+            "".join(call.function.arguments for chunk in chunks for choice in chunk.choices
+                    if choice.index == index for call in choice.delta.tool_calls or ())
+            for index in range(2)
+        )
+        assert arguments == ('{"fruit": "PERSIMMON"}', '{"fruit": "DURIAN"}')
 
     @pytest.mark.asyncio
     async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):

tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py

@@ -8,6 +8,8 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator
 from types import SimpleNamespace
 from typing import TYPE_CHECKING, Final, Literal
 
+from pydantic import TypeAdapter
+
 import anyio
 import pytest
 
@@ -50,6 +52,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse
 from litellm.types.utils import (
     CallTypes,
     CallTypesLiteral,
+    ChatCompletionMessageToolCall,
     Delta,
     GenericGuardrailAPIInputs,
     ModelResponse,
@@ -1008,6 +1011,42 @@ def _delta_text(item):
     return item.choices[0].delta.content or ""
 
 
+class _ToolRedactingGuardrail(CustomGuardrail):
+    def __init__(self, require_tool_context: bool = False) -> None:
+        super().__init__(guardrail_name="tool-redactor", event_hook=GuardrailEventHooks.post_call, default_on=True)
+        self.streaming_transform_mode = "incremental_diff"
+        self.streaming_sampling_rate = 1
+        self.require_tool_context = require_tool_context
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict[str, object],
+        input_type: Literal["request", "response"],
+        logging_obj: object | None = None,
+    ) -> GenericGuardrailAPIInputs:
+        calls: Final = TypeAdapter(tuple[ChatCompletionMessageToolCall, ...]).validate_python(
+            inputs.get("tool_calls", ())
+        )
+        input_texts: Final = inputs.get("texts", ())
+        texts: Final = tuple(
+            "checked:" + text.replace("SECRET", "MASKED")
+            if not self.require_tool_context or (calls and index == len(input_texts) - 1) else text
+            for index, text in enumerate(input_texts)
+        )
+        return {
+            **inputs,
+            "texts": list(texts),
+            "stream_holdback_chars": [len(text) for text in texts],
+            "tool_calls": [
+                call.model_copy(update={"function": call.function.model_copy(update={
+                    "arguments": call.function.arguments.replace("SECRET", "MASKED"),
+                })})
+                for call in calls
+            ],
+        }
+
+
 class TestStreamingTransform:
     """Streaming text-transformation (incremental_diff) path on the OpenAI chat
     completions streaming surface."""
@@ -1016,6 +1055,103 @@ class TestStreamingTransform:
     def _use_openai_handler_mapping(self, monkeypatch):
         _patch_translation_mappings(monkeypatch, {CallTypes.acompletion: OpenAIChatCompletionsHandler})
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("require_tool_context", [False, True])
+    @pytest.mark.parametrize("include_text", [False, True])
+    @pytest.mark.parametrize("tool_count", [1, 2])
+    async def test_buffered_tool_arguments_are_rewritten_before_delivery(
+        self, include_text: bool, tool_count: int, require_tool_context: bool
+    ) -> None:
+        chunks: Final = (
+            *([_stream_chunk("hello SECRET")] if include_text else []),
+            ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[
+                {"index": index, "id": f"call_{index}", "type": "function",
+                 "function": {"name": "contact", "arguments": '{"contact":"SEC'}}
+                for index in range(tool_count)
+            ]))]),
+            ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[
+                {"index": index, "function": {"arguments": 'RET"}'}} for index in range(tool_count)
+            ]), finish_reason="tool_calls")]),
+            ModelResponseStream(choices=[], usage={"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}),
+        )
+        out: Final = await _drive_stream(
+            UnifiedLLMGuardrails(), _ToolRedactingGuardrail(require_tool_context=require_tool_context), chunks
+        )
+        calls: Final = tuple(
+            call for chunk in out for choice in chunk.choices for call in choice.delta.tool_calls or ()
+        )
+        for index in range(tool_count):
+            arguments: Final = "".join(call.function.arguments or "" for call in calls if call.index == index)
+            assert arguments == '{"contact":"MASKED"}'
+            assert next(call.id for call in calls if call.index == index and call.id) == f"call_{index}"
+        assert "".join(_delta_text(chunk) for chunk in out) == ("checked:hello MASKED" if include_text else "")
+        assert any(choice.finish_reason == "tool_calls" for chunk in out for choice in chunk.choices)
+        assert out[-1].usage.total_tokens == 18
+        assert all("SECRET" not in chunk.model_dump_json() for chunk in out)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("tool_choice_index", [0, 1])
+    async def test_tool_dependent_text_rewrites_keep_choice_context(self, tool_choice_index: int) -> None:
+        chunks: Final = (
+            ModelResponseStream(choices=[StreamingChoices(
+                index=index, delta=Delta(content="SECRET" if index == tool_choice_index else "plain"),
+            ) for index in (1, 0)]),
+            ModelResponseStream(choices=[StreamingChoices(
+                index=tool_choice_index,
+                delta=Delta(tool_calls=[{
+                    "index": 0, "id": "call_context", "type": "function",
+                    "function": {"name": "contact", "arguments": '{"contact":"SECRET"}'},
+                }]), finish_reason="tool_calls",
+            )]),
+        )
+        out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(True), chunks)
+        for index in (0, 1):
+            text: Final = "".join(
+                choice.delta.content or "" for chunk in out for choice in chunk.choices if choice.index == index
+            )
+            assert text == ("checked:MASKED" if index == tool_choice_index else "plain")
+        assert all("SECRET" not in chunk.model_dump_json() for chunk in out)
+
+    @pytest.mark.asyncio
+    async def test_tool_rewrites_keep_completion_choices_separate(self) -> None:
+        async def response() -> AsyncIterator[ModelResponseStream]:
+            yield ModelResponseStream(choices=[StreamingChoices(
+                index=index,
+                delta=Delta(tool_calls=[{"index": 0, "id": f"call_{index}", "type": "function",
+                                        "function": {"name": "contact", "arguments": '{"contact":"SECRET"}'}}]),
+                finish_reason="tool_calls",
+            ) for index in range(2)])
+
+        iterator: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
+            user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"),
+            response=response(), request_data={"guardrail_to_apply": _ToolRedactingGuardrail()},
+        )
+        chunks: Final = tuple([chunk async for chunk in iterator])
+        calls: Final = tuple(
+            (choice.index, call.id, call.function.arguments)
+            for chunk in chunks for choice in chunk.choices for call in choice.delta.tool_calls or ()
+        )
+        assert calls == ((0, "call_0", '{"contact":"MASKED"}'), (1, "call_1", '{"contact":"MASKED"}'))
+
+    @pytest.mark.asyncio
+    async def test_undeliverable_rewrite_is_a_closed_failure(self) -> None:
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        class UndeliverableTranslation(OpenAIChatCompletionsHandler):
+            async def process_output_streaming_response(
+                self, *args: object, **kwargs: object
+            ) -> list[ModelResponseStream]:
+                raise UndeliverableStreamRewrite("tool-redactor", "unmappable tool fragments")
+
+        iterator: Final = UnifiedLLMGuardrails()._inspect_full_response(
+            endpoint_translation=UndeliverableTranslation(), guardrail_to_apply=_ToolRedactingGuardrail(),
+            request_data={}, user_api_key_dict=UserAPIKeyAuth(), responses_so_far=(), responses_yielded=(),
+        )
+        with pytest.raises(unified_module.HTTPException) as error:
+            await anext(iterator)
+        assert error.value.status_code == 400
+        assert error.value.detail == "Guardrail stream rewrite could not be applied"
+
     @pytest.mark.asyncio
     async def test_block_only_drops_text_rewrites(self):
         """Default block_only: the guardrail's uppercasing never reaches the
@@ -1436,14 +1572,15 @@ class TestStreamingTransform:
         assert out[-2].choices[0].finish_reason == "stop"
 
     @pytest.mark.asyncio
-    async def test_tool_call_blocking_guardrail_is_enforced(self):
+    @pytest.mark.parametrize(("content", "allowed_scans"), [(None, 0), ("proposal", 0), ("proposal", 1)])
+    async def test_tool_call_blocking_guardrail_is_enforced(self, content: str | None, allowed_scans: int):
         """A guardrail that blocks on tool calls must terminate the incremental_diff
         stream: tool calls go through the block decision, not bypass it."""
         from litellm.exceptions import GuardrailRaisedException
 
         class _ToolCallBlocker(_StreamingTextGuardrail):
             async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
-                if input_type == "response" and inputs.get("tool_calls"):
+                if input_type == "response" and self.response_calls >= allowed_scans:
                     raise GuardrailRaisedException(
                         guardrail_name="tc-block",
                         message="blocked tool call",
@@ -1456,7 +1593,7 @@ class TestStreamingTransform:
                 StreamingChoices(
                     index=0,
                     delta=Delta(
-                        content=None,
+                        content=content,
                         tool_calls=[
                             {
                                 "index": 0,
@@ -1471,8 +1608,21 @@ class TestStreamingTransform:
             ],
         )
 
+        async def upstream():
+            yield tool_chunk
+
+        stream: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
+            user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"),
+            response=upstream(),
+            request_data={"guardrail_to_apply": _ToolCallBlocker(), "model": "gpt-4"},
+        )
+        async def consume_checked_stream() -> None:
+            async for chunk in stream:
+                assert isinstance(chunk, ModelResponseStream)
+                assert all(not choice.delta.tool_calls and choice.finish_reason is None for choice in chunk.choices)
+
         with pytest.raises(GuardrailRaisedException):
-            await _drive_stream(UnifiedLLMGuardrails(), _ToolCallBlocker(), [tool_chunk])
+            await consume_checked_stream()
 
     @pytest.mark.asyncio
     async def test_mixed_content_and_tool_call_chunk_does_not_leak_text(self):
@@ -2690,7 +2840,7 @@ class TestStreamingClientDisconnectScan:
         yield _stream_chunk(" tail", finish_reason="stop")
 
     @pytest.mark.asyncio
-    async def test_closing_while_an_incremental_diff_block_is_delivered_does_not_inspect_the_tool_call_again(self):
+    async def test_incremental_diff_block_withholds_tool_calls_and_does_not_inspect_again(self):
         guardrail = _MarkerBlockingStreamingTextGuardrail()
 
         async def upstream() -> AsyncIterator[ModelResponseStream]:
@@ -2704,12 +2854,16 @@ class TestStreamingClientDisconnectScan:
         released = [await stream.__anext__(), await stream.__anext__()]
         await stream.aclose()
 
-        assert [call.function.name for call in released[0].choices[0].delta.tool_calls] == ["get_weather"]
-        assert guardrail.received_texts == [["BLOCKME"]], guardrail.received_texts
-        assert guardrail.received_tool_calls == [], guardrail.received_tool_calls
+        assert all(isinstance(item, bytes) for item in released)
+        assert b"get_weather" not in b"".join(released)
+        assert b"Paris" not in b"".join(released)
+        assert guardrail.received_texts == [["BLOCKME tail"]], guardrail.received_texts
+        assert [[call["function"]["name"] for call in calls] for calls in guardrail.received_tool_calls] == [
+            ["get_weather"]
+        ]
 
     @pytest.mark.asyncio
-    async def test_closing_after_a_released_tool_call_under_incremental_diff_scans_no_held_back_text(self):
+    async def test_closing_after_a_buffered_tool_call_inspects_held_text_before_delivery(self):
         guardrail = _StreamingTextGuardrail(holdback_schedule=[len("held secret")] * 2)
 
         async def upstream() -> AsyncIterator[ModelResponseStream]:
@@ -2718,12 +2872,14 @@ class TestStreamingClientDisconnectScan:
                 yield tool_chunk
 
         stream = self._guarded_stream(guardrail, upstream())
-        received = await stream.__anext__()
+        received: Final = tuple([await stream.__anext__() for _ in range(2)])
+        scans_before_close: Final = len(guardrail.received_texts)
         await stream.aclose()
 
-        assert [call.function.name for call in received.choices[0].delta.tool_calls] == ["get_weather"]
+        assert [call.function.name for call in received[-1].choices[0].delta.tool_calls] == ["get_weather"]
         assert guardrail.received_tool_calls, guardrail.received_texts
-        assert all("held secret" not in text for text in guardrail.received_texts[-1]), guardrail.received_texts
+        assert guardrail.received_texts[0] == ["held secret"]
+        assert len(guardrail.received_texts) == scans_before_close
 
 
 def _responses_delta(sequence_number, text):

tests/unit/proxy/guardrails/test_guardrail_endpoints.py

@@ -2167,7 +2167,7 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker):
     }
 
 
-def _patch_apply_guardrail_env(mocker, guardrail_result):
+def _patch_apply_guardrail_env(mocker, guardrail_result, processed_data=None):
     mock_guardrail = mocker.Mock()
     mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result)
 
@@ -2180,7 +2180,7 @@ def _patch_apply_guardrail_env(mocker, guardrail_result):
     mock_logging_obj.model_call_details = {}
     mock_processor = mocker.Mock()
     mock_processor.common_processing_pre_call_logic = AsyncMock(
-        return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj)
+        return_value=(processed_data or {"guardrail_name": "test-guardrail"}, mock_logging_obj)
     )
     mocker.patch(
         "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
@@ -2251,6 +2251,68 @@ async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker):
     )
 
 
+@pytest.mark.asyncio
+async def test_apply_guardrail_replaces_caller_identity_with_the_authenticated_key(mocker):
+    """Identity fields come from the proxy's own sanitized metadata, never from
+    the caller, so a request cannot impersonate another key or probe its policy."""
+    mock_guardrail = _patch_apply_guardrail_env(
+        mocker,
+        {"texts": ["ok"]},
+        processed_data={
+            "guardrail_name": "test-guardrail",
+            "metadata": {
+                "route": "/apply_guardrail",
+                "user_api_key_alias": "billing-app",
+                "user_api_key_user_id": "u-1",
+            },
+        },
+    )
+
+    request = ApplyGuardrailRequest(
+        guardrail_name="test-guardrail",
+        text="hello",
+        metadata={
+            "forbidden_topics": ["tax"],
+            "user_api_key_alias": "someone-else",
+            "user_api_key_user_email": "victim@example.com",
+        },
+    )
+    await apply_guardrail(
+        fastapi_request=mocker.Mock(),
+        request=request,
+        user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"),
+    )
+
+    forwarded = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"]
+    assert forwarded == {
+        "forbidden_topics": ["tax"],
+        "user_api_key_alias": "billing-app",
+        "user_api_key_user_id": "u-1",
+    }
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_attaches_key_identity_when_caller_sends_no_metadata(mocker):
+    """A caller that sends no metadata still gets the authenticated identity
+    forwarded, the same shape every LLM route gives a guardrail."""
+    mock_guardrail = _patch_apply_guardrail_env(
+        mocker,
+        {"texts": ["ok"]},
+        processed_data={"guardrail_name": "test-guardrail", "metadata": {"user_api_key_alias": "billing-app"}},
+    )
+
+    request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello")
+    await apply_guardrail(
+        fastapi_request=mocker.Mock(),
+        request=request,
+        user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app"),
+    )
+
+    assert mock_guardrail.apply_guardrail.await_args.kwargs["request_data"] == {
+        "metadata": {"user_api_key_alias": "billing-app"}
+    }
+
+
 @pytest.mark.asyncio
 async def test_apply_guardrail_omits_metadata_when_not_sent(mocker):
     """Without metadata, request_data stays empty (backward-compatible)."""

ui/litellm-dashboard/public/assets/logos/neuraltrust.svg

@@ -0,0 +1,22 @@
+<svg xmlns="http://www.w3.org/2000/svg" width="32" height="32" viewBox="0 0 32 32" fill="none">
+  <g clip-path="url(#neuraltrustClip)">
+    <path fill="url(#neuraltrustGrad)" d="M32 0H0v32h32z" />
+    <path
+      fill="#fff"
+      d="M18.092 20.06a.67.67 0 0 1-.55.3.7.7 0 0 1-.565-.286l-2.704-3.814-1.45 2.103 2.197 3.098a3.08 3.08 0 0 0 2.51 1.297h.038a3.06 3.06 0 0 0 2.502-1.342l8.02-11.477h-2.926z"
+    />
+    <path
+      fill="#fff"
+      d="M14.292 11.518a.63.63 0 0 1 .552.286l2.652 3.74 1.449-2.103-2.145-3.024a3.08 3.08 0 0 0-2.509-1.297h-.039a3.06 3.06 0 0 0-2.506 1.35L3.925 21.85l-.085.123h2.91l6.98-10.155a.68.68 0 0 1 .562-.3"
+    />
+  </g>
+  <defs>
+    <linearGradient id="neuraltrustGrad" x1="30.667" x2="6.667" y1="0" y2="32" gradientUnits="userSpaceOnUse">
+      <stop stop-color="#03AFFF" />
+      <stop offset="1" stop-color="#9B29FF" />
+    </linearGradient>
+    <clipPath id="neuraltrustClip">
+      <path fill="#fff" d="M0 0h32v32H0z" />
+    </clipPath>
+  </defs>
+</svg>

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts

@@ -256,6 +256,14 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
     mode: "pre_call",
     defaultOn: false,
   },
+  neuraltrust: {
+    provider: "Neuraltrust",
+    guardrailNameSuggestion: "NeuralTrust TrustGuard",
+    // TrustGuard rewrites the request and masks the reply, so seeding pre_call alone
+    // would leave model output uninspected. Matches the hook README's example config.
+    mode: ["pre_call", "post_call"],
+    defaultOn: false,
+  },
   noma: {
     provider: "Noma",
     guardrailNameSuggestion: "Noma Security",

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts

@@ -18,6 +18,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record<string, string> = {
   panw: "palo_alto_networks.jpeg",
   cisco_ai_defense: "cisco.png",
   noma: "noma_security.png",
+  neuraltrust: "neuraltrust.svg",
   aporia: "aporia.png",
   aim: "aim_security.jpeg",
   cato_networks: "cato_networks.svg",
@@ -77,6 +78,16 @@ describe("guardrail_garden_data logos", () => {
       expect(card.logo, `card ${card.id}`).not.toContain("/ui/assets/logos/");
     }
   });
+
+  it("does not publish unsourced NeuralTrust eval numbers", () => {
+    const card = PARTNER_GUARDRAIL_CARDS.find((c) => c.id === "neuraltrust");
+    expect(card).toBeDefined();
+    expect(card?.eval).toBeUndefined();
+  });
+
+  it("seeds the NeuralTrust card with both pre_call and post_call so replies are inspected", () => {
+    expect(GUARDRAIL_PRESETS.neuraltrust.mode).toEqual(["pre_call", "post_call"]);
+  });
 });
 
 describe("decision model presets", () => {

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts

@@ -372,6 +372,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [
     tags: ["Enterprise", "Security", "Prompt Injection", "PII"],
     providerKey: "CiscoAiDefense",
   },
+  {
+    id: "neuraltrust",
+    name: "NeuralTrust",
+    description:
+      "TrustGuard runtime guardrails: prompt injection, toxicity, DLP, and policy enforcement on LLM input and output.",
+    category: "partner",
+    logo: guardrailLogoMap["NeuralTrust"],
+    tags: ["Security", "Prompt Injection", "DLP"],
+    providerKey: "Neuraltrust",
+  },
   {
     id: "noma",
     name: "Noma Security",

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx

@@ -211,6 +211,20 @@ describe("guardrail_info_helpers", () => {
       expect(result.logo).toContain("noma_security.png");
     });
 
+    it("should resolve NeuralTrust logo and display name", () => {
+      populateGuardrailProviders({
+        neuraltrust: { ui_friendly_name: "NeuralTrust" },
+      });
+      populateGuardrailProviderMap({
+        neuraltrust: { ui_friendly_name: "NeuralTrust" },
+      });
+
+      const result = getGuardrailLogoAndName("neuraltrust");
+
+      expect(result.displayName).toBe("NeuralTrust");
+      expect(result.logo).toContain("neuraltrust.svg");
+    });
+
     it("should resolve RepelloAI Argus logo and display name", () => {
       populateGuardrailProviders({
         repelloai: { ui_friendly_name: "RepelloAI Argus" },

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx

@@ -17,6 +17,7 @@ import lakeraAiLogo from "../../../../../public/assets/logos/lakeraai.jpeg";
 import lassoLogo from "../../../../../public/assets/logos/lasso.png";
 import litellmLogo from "../../../../../public/assets/logos/litellm_monogram.svg";
 import microsoftAzureLogo from "../../../../../public/assets/logos/microsoft_azure.svg";
+import neuraltrustLogo from "../../../../../public/assets/logos/neuraltrust.svg";
 import nomaSecurityLogo from "../../../../../public/assets/logos/noma_security.png";
 import openaiSmallLogo from "../../../../../public/assets/logos/openai_small.svg";
 import paloAltoNetworksLogo from "../../../../../public/assets/logos/palo_alto_networks.jpeg";
@@ -299,6 +300,7 @@ export const guardrailLogoMap = {
   "Aporia AI": aporiaLogo.src,
   "PANW Prisma AIRS": paloAltoNetworksLogo.src,
   "Cisco AI Defense": ciscoLogo.src,
+  NeuralTrust: neuraltrustLogo.src,
   "Noma Security": nomaSecurityLogo.src,
   "Javelin Guardrails": javelinLogo.src,
   "Pillar Guardrail": pillarLogo.src,

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -27259,6 +27259,11 @@ export interface components {
             stream_scope?: ("streaming" | "non_streaming" | "both") | {
                 [key: string]: "streaming" | "non_streaming" | "both";
             } | null;
+            /**
+             * Streaming Transform Mode
+             * @description Whether a guardrail's text rewrite reaches a streaming client. Implemented by guardrail='prompt_security' and 'neuraltrust'; generic_guardrail_api takes the same setting under optional_params. 'block_only' (default) streams the raw model chunks, so a block still ends the stream but rewrites are dropped. 'incremental_diff' withholds those chunks and streams the guardrail's rewritten text instead. OpenAI chat completions streaming only.
+             */
+            streaming_transform_mode?: ("block_only" | "incremental_diff") | null;
             /**
              * Template Id
              * @description The ID of your Model Armor template
@@ -27271,7 +27276,7 @@ export interface components {
             timeout?: number | null;
             /**
              * Unreachable Fallback
-             * @description Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. Also implemented by guardrail='decision_model'.
+             * @description Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', 'typesafe', and 'neuraltrust'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. Also implemented by guardrail='decision_model'.
              * @default fail_closed
              * @enum {string}
              */
@@ -36572,6 +36577,11 @@ export interface components {
              * @description Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable.
              */
             client_secret?: string | null;
+            /**
+             * Collector Key
+             * @description TrustGuard collector key (tgcol_...). Optional when the API key is bound to a collector. Env: TRUSTGUARD_COLLECTOR_KEY.
+             */
+            collector_key?: string | null;
             /**
              * Confidence Threshold
              * @description Only block or mask when detection confidence >= this value; below threshold, allow or log_only.
@@ -37165,6 +37175,11 @@ export interface components {
             stream_scope?: ("streaming" | "non_streaming" | "both") | {
                 [key: string]: "streaming" | "non_streaming" | "both";
             } | null;
+            /**
+             * Streaming Transform Mode
+             * @description Whether a guardrail's text rewrite reaches a streaming client. Implemented by guardrail='prompt_security' and 'neuraltrust'; generic_guardrail_api takes the same setting under optional_params. 'block_only' (default) streams the raw model chunks, so a block still ends the stream but rewrites are dropped. 'incremental_diff' withholds those chunks and streams the guardrail's rewritten text instead. OpenAI chat completions streaming only.
+             */
+            streaming_transform_mode?: ("block_only" | "incremental_diff") | null;
             /**
              * Template Id
              * @description The ID of your Model Armor template