refactor(library): make context bloat verdict flow-independent (#2190)

Return RailOutcome from context-bloat checks so execution backends can derive the same block, transform, or allow verdict from the configured source and input. Update both Colang dialects and cover the action and flow contracts. Follow-up on #2150


Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com>

Author: Pouyanpi

nemoguardrails/library/context_bloat_detection/actions.py

@@ -33,10 +33,11 @@
 import logging
 import math
 from collections import Counter
-from typing import List, Optional, TypedDict
+from typing import List, Literal, Optional, TypedDict
 
 from nemoguardrails import RailsConfig
 from nemoguardrails.actions import action
+from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget
 
 log = logging.getLogger(__name__)
 
@@ -54,6 +55,9 @@ class ContextBloatResult(TypedDict):
     metrics: dict
 
 
+ContextBloatSource = Literal["input", "retrieval"]
+
+
 def _stratified_sample(text: str, sample_chars: int) -> str:
     third = sample_chars // 3
     mid = len(text) // 2
@@ -162,21 +166,48 @@ def _check_repetition(text: str, cfg, detections: List[str], metrics: dict) -> O
     return None
 
 
+def _context_bloat_outcome(source: ContextBloatSource, original_text: str, result: ContextBloatResult) -> RailOutcome:
+    targets = {
+        "input": TransformTarget.USER_MESSAGE,
+        "retrieval": TransformTarget.RELEVANT_CHUNKS,
+    }
+
+    metadata = {
+        "is_bloat": result["is_bloat"],
+        "action": result["action"],
+        "detections": result["detections"],
+        "metrics": result["metrics"],
+    }
+    if result["is_bloat"] and result["action"] == "reject":
+        return RailOutcome.block(reason=result["reason"], metadata=metadata)
+    if result["is_bloat"] and result["action"] == "truncate" and result["text"] != original_text:
+        return RailOutcome.transform(
+            [(targets[source], result["text"])],
+            reason=result["reason"],
+            metadata=metadata,
+        )
+    return RailOutcome.allow(reason=result["reason"], metadata=metadata)
+
+
 @action()
-async def context_bloat_detection(text: str, config: RailsConfig) -> ContextBloatResult:
+async def context_bloat_detection(source: ContextBloatSource, text: str, config: RailsConfig) -> RailOutcome:
     """Detect context-bloat / context-manipulation attacks.
     Check order is cheapest first to enable early-exit.
 
     Args:
+        source: Whether the text came from input or retrieved chunks.
         text: The text to inspect (joined chunks or user message).
         config: RailsConfig with rails.config.context_bloat_detection settings.
 
     Returns:
-        ContextBloatResult with is_bloat flag, processed text, reason, metrics.
+        The allow, block, or transform verdict with detection evidence.
     """
+    if source not in ("input", "retrieval"):
+        raise ValueError("source must be either 'input' or 'retrieval'")
     _validate_config(config)
     cfg = config.rails.config.context_bloat_detection
 
+    original_text = text
     char_count = len(text) if text else 0
     detections: List[str] = []
     metrics: dict = {"chars": char_count}
@@ -185,13 +216,17 @@ async def context_bloat_detection(text: str, config: RailsConfig) -> ContextBloa
         detections.append("size_cap_exceeded")
         log.info(f"context bloat detected: size_cap_exceeded | chars={char_count}")
         if cfg.action == "reject":
-            return ContextBloatResult(
-                is_bloat=True,
-                action=cfg.action,
-                text=text,
-                reason="size_cap_exceeded",
-                detections=detections,
-                metrics=metrics,
+            return _context_bloat_outcome(
+                source,
+                original_text,
+                ContextBloatResult(
+                    is_bloat=True,
+                    action=cfg.action,
+                    text=text,
+                    reason="size_cap_exceeded",
+                    detections=detections,
+                    metrics=metrics,
+                ),
             )
         if cfg.action == "truncate":
             text = text[: cfg.max_chars]
@@ -200,17 +235,21 @@ async def context_bloat_detection(text: str, config: RailsConfig) -> ContextBloa
         for check in (_check_entropy, _check_longest_run, _check_repetition):
             result = check(text, cfg, detections, metrics)
             if result is not None:
-                return result
+                return _context_bloat_outcome(source, original_text, result)
 
     is_bloat = bool(detections)
     reason = ", ".join(detections) if detections else None
     if is_bloat:
         log.info(f"context bloat detected: {reason} | metrics={metrics}")
-    return ContextBloatResult(
-        is_bloat=is_bloat,
-        action=cfg.action,
-        text=text,
-        reason=reason,
-        detections=detections,
-        metrics=metrics,
+    return _context_bloat_outcome(
+        source,
+        original_text=original_text,
+        result=ContextBloatResult(
+            is_bloat=is_bloat,
+            action=cfg.action,
+            text=text,
+            reason=reason,
+            detections=detections,
+            metrics=metrics,
+        ),
     )

nemoguardrails/library/context_bloat_detection/flows.co

@@ -4,12 +4,12 @@
 flow context bloat detection on retrieval
   """Detect bloated or padded content in retrieved RAG chunks."""
   global $relevant_chunks
-  $bloat_result = await ContextBloatDetectionAction(text=$relevant_chunks)
-  if $bloat_result.is_bloat and $bloat_result.action == "reject"
+  $bloat_result = await ContextBloatDetectionAction(source="retrieval", text=$relevant_chunks)
+  if $bloat_result.is_blocked
     bot inform retrieval bloated
     abort
-  if $bloat_result.is_bloat and $bloat_result.action == "truncate"
-    $relevant_chunks = $bloat_result.text
+  if $bloat_result.is_transform
+    $relevant_chunks = $bloat_result.transform_text["relevant_chunks"]
 
 flow bot inform retrieval bloated
   bot say "The retrieved sources for this question appeared oversized or padded. I'm not using them to avoid context manipulation."
@@ -19,12 +19,12 @@ flow bot inform retrieval bloated
 flow context bloat detection on input
   """Detect bloated or padded user-supplied content."""
   global $user_message
-  $bloat_result = await ContextBloatDetectionAction(text=$user_message)
-  if $bloat_result.is_bloat and $bloat_result.action == "reject"
+  $bloat_result = await ContextBloatDetectionAction(source="input", text=$user_message)
+  if $bloat_result.is_blocked
     bot refuse bloated input
     abort
-  if $bloat_result.is_bloat and $bloat_result.action == "truncate"
-    $user_message = $bloat_result.text
+  if $bloat_result.is_transform
+    $user_message = $bloat_result.transform_text["user_message"]
 
 flow bot refuse bloated input
   bot say "Your message appears oversized or padded with repetitive content. Please send a shorter message focused on your question."

nemoguardrails/library/context_bloat_detection/flows.v1.co

@@ -3,12 +3,12 @@
 
 define subflow context bloat detection on retrieval
   """Detect bloated or padded content in retrieved RAG chunks."""
-  $bloat_result = execute context_bloat_detection(text=$relevant_chunks)
-  if $bloat_result.is_bloat and $bloat_result.action == "reject"
+  $bloat_result = execute context_bloat_detection(source="retrieval", text=$relevant_chunks)
+  if $bloat_result.is_blocked
     bot inform retrieval bloated
     stop
-  if $bloat_result.is_bloat and $bloat_result.action == "truncate"
-    $relevant_chunks = $bloat_result.text
+  if $bloat_result.is_transform
+    $relevant_chunks = $bloat_result.transform_text["relevant_chunks"]
 
 define bot inform retrieval bloated
   "The retrieved sources for this question appeared oversized or padded. I'm not using them to avoid context manipulation."
@@ -17,12 +17,12 @@ define bot inform retrieval bloated
 
 define subflow context bloat detection on input
   """Detect bloated or padded user-supplied content."""
-  $bloat_result = execute context_bloat_detection(text=$user_message)
-  if $bloat_result.is_bloat and $bloat_result.action == "reject"
+  $bloat_result = execute context_bloat_detection(source="input", text=$user_message)
+  if $bloat_result.is_blocked
     bot refuse bloated input
     stop
-  if $bloat_result.is_bloat and $bloat_result.action == "truncate"
-    $user_message = $bloat_result.text
+  if $bloat_result.is_transform
+    $user_message = $bloat_result.transform_text["user_message"]
 
 define bot refuse bloated input
   "Your message appears oversized or padded with repetitive content. Please send a shorter message focused on your question."