feat(tool-rails): add support for tool output rails and validation

nemoguardrails/actions/llm/generation.py

@@ -582,7 +582,21 @@ async def generate_user_intent(
             if streaming_handler:
                 await streaming_handler.push_chunk(text)
 
-            output_events.append(new_event_dict("BotMessage", text=text))
+            if self.config.passthrough:
+                from nemoguardrails.actions.llm.utils import (
+                    get_and_clear_tool_calls_contextvar,
+                )
+
+                tool_calls = get_and_clear_tool_calls_contextvar()
+
+                if tool_calls:
+                    output_events.append(
+                        new_event_dict("BotToolCall", tool_calls=tool_calls)
+                    )
+                else:
+                    output_events.append(new_event_dict("BotMessage", text=text))
+            else:
+                output_events.append(new_event_dict("BotMessage", text=text))
 
             return ActionResult(events=output_events)
 

nemoguardrails/actions/llm/utils.py

@@ -673,6 +673,21 @@ def get_and_clear_tool_calls_contextvar() -> Optional[list]:
     return None
 
 
+def extract_tool_calls_from_events(events: list) -> Optional[list]:
+    """Extract tool_calls from BotToolCall events.
+
+    Args:
+        events: List of events to search through
+
+    Returns:
+        tool_calls if found in BotToolCall event, None otherwise
+    """
+    for event in events:
+        if event.get("type") == "BotToolCall":
+            return event.get("tool_calls")
+    return None
+
+
 def get_and_clear_response_metadata_contextvar() -> Optional[dict]:
     """Get the current response metadata and clear it from the context.
 

nemoguardrails/rails/llm/config.py

@@ -527,6 +527,40 @@ class ActionRails(BaseModel):
     )
 
 
+class ToolOutputRails(BaseModel):
+    """Configuration of tool output rails.
+
+    Tool output rails are applied to tool calls before they are executed.
+    They can validate tool names, parameters, and context to ensure safe tool usage.
+    """
+
+    flows: List[str] = Field(
+        default_factory=list,
+        description="The names of all the flows that implement tool output rails.",
+    )
+    parallel: Optional[bool] = Field(
+        default=False,
+        description="If True, the tool output rails are executed in parallel.",
+    )
+
+
+class ToolInputRails(BaseModel):
+    """Configuration of tool input rails.
+
+    Tool input rails are applied to tool results before they are processed.
+    They can validate, filter, or transform tool outputs for security and safety.
+    """
+
+    flows: List[str] = Field(
+        default_factory=list,
+        description="The names of all the flows that implement tool input rails.",
+    )
+    parallel: Optional[bool] = Field(
+        default=False,
+        description="If True, the tool input rails are executed in parallel.",
+    )
+
+
 class SingleCallConfig(BaseModel):
     """Configuration for the single LLM call option for topical rails."""
 
@@ -912,6 +946,14 @@ class Rails(BaseModel):
     actions: ActionRails = Field(
         default_factory=ActionRails, description="Configuration of action rails."
     )
+    tool_output: ToolOutputRails = Field(
+        default_factory=ToolOutputRails,
+        description="Configuration of tool output rails.",
+    )
+    tool_input: ToolInputRails = Field(
+        default_factory=ToolInputRails,
+        description="Configuration of tool input rails.",
+    )
 
 
 def merge_two_dicts(dict_1: dict, dict_2: dict, ignore_keys: Set[str]) -> None:

nemoguardrails/rails/llm/llm_flows.co

@@ -102,6 +102,34 @@ define parallel extension flow generate bot message
   execute generate_bot_message
 
 
+define parallel extension flow process bot tool call
+  """Processes tool calls from the bot."""
+  priority 100
+
+  event BotToolCall
+
+  $tool_calls = $event.tool_calls
+
+  # Run tool-specific output rails if configured (Phase 2)
+  if $config.rails.tool_output.flows
+    # If we have generation options, we make sure the tool output rails are enabled.
+    if $generation_options is None or $generation_options.rails.tool_output:
+      # Create a marker event.
+      create event StartToolOutputRails
+      event StartToolOutputRails
+
+      # Run all the tool output rails
+      # This can potentially alter or block the tool calls
+      do run tool output rails
+
+      # Create a marker event.
+      create event ToolOutputRailsFinished
+      event ToolOutputRailsFinished
+
+  # Create the action event for tool execution
+  create event StartToolCallBotAction(tool_calls=$tool_calls)
+
+
 define parallel extension flow process bot message
   """Runs the output rails on a bot message."""
   priority 100
@@ -164,3 +192,25 @@ define subflow run retrieval rails
   while $i < len($retrieval_flows)
     do $retrieval_flows[$i]
     $i = $i + 1
+
+
+define subflow run tool output rails
+  """Runs all the tool output rails in a sequential order."""
+  $tool_output_flows = $config.rails.tool_output.flows
+
+  $i = 0
+  while $i < len($tool_output_flows)
+    # We set the current rail as being triggered.
+    $triggered_tool_output_rail = $tool_output_flows[$i]
+
+    create event StartToolOutputRail(flow_id=$triggered_tool_output_rail)
+    event StartToolOutputRail
+
+    do $tool_output_flows[$i]
+    $i = $i + 1
+
+    create event ToolOutputRailFinished(flow_id=$triggered_tool_output_rail)
+    event ToolOutputRailFinished
+
+    # If all went smooth, we remove it.
+    $triggered_tool_output_rail = None

nemoguardrails/rails/llm/llmrails.py

@@ -32,9 +32,9 @@
 
 from nemoguardrails.actions.llm.generation import LLMGenerationActions
 from nemoguardrails.actions.llm.utils import (
+    extract_tool_calls_from_events,
     get_and_clear_reasoning_trace_contextvar,
     get_and_clear_response_metadata_contextvar,
-    get_and_clear_tool_calls_contextvar,
     get_colang_history,
 )
 from nemoguardrails.actions.output_mapping import is_output_blocked
@@ -1086,7 +1086,7 @@ async def generate_async(
                 options.log.llm_calls = True
                 options.log.internal_events = True
 
-        tool_calls = get_and_clear_tool_calls_contextvar()
+        tool_calls = extract_tool_calls_from_events(new_events)
         llm_metadata = get_and_clear_response_metadata_contextvar()
 
         # If we have generation options, we prepare a GenerationResponse instance.

nemoguardrails/rails/llm/options.py

@@ -127,6 +127,16 @@ class GenerationRailsOptions(BaseModel):
         default=True,
         description="Whether the dialog rails are enabled or not.",
     )
+    tool_output: Union[bool, List[str]] = Field(
+        default=True,
+        description="Whether the tool output rails are enabled or not. "
+        "If a list of names is specified, then only the specified tool output rails will be applied.",
+    )
+    tool_input: Union[bool, List[str]] = Field(
+        default=True,
+        description="Whether the tool input rails are enabled or not. "
+        "If a list of names is specified, then only the specified tool input rails will be applied.",
+    )
 
 
 class GenerationOptions(BaseModel):
@@ -177,6 +187,8 @@ def check_fields(cls, values):
                 "dialog": False,
                 "retrieval": False,
                 "output": False,
+                "tool_output": False,
+                "tool_input": False,
             }
             for rail_type in values["rails"]:
                 _rails[rail_type] = True