create explain_info context variable in stream_async to keep trac…

docs/user-guides/advanced/streaming.md

@@ -1,9 +1,11 @@
 # Streaming
 
-To use a guardrails configuration in streaming mode, the following must be met:
+If the application LLM supports streaming, you can configure NeMo Guardrails to stream tokens as well.
 
-1. The main LLM must support streaming.
-2. There are no output rails.
+For information about configuring streaming with output guardrails, refer to the following:
+
+- For configuration, refer to [streaming output configuration](../../user-guides/configuration-guide.md#streaming-output-configuration).
+- For sample Python client code, refer to [streaming output](../../getting-started/5-output-rails/README.md#streaming-output).
 
 ## Configuration
 
@@ -26,6 +28,7 @@ nemoguardrails chat --config=examples/configs/streaming --streaming
 ### Python API
 
 You can use the streaming directly from the python API in two ways:
+
 1. Simple: receive just the chunks (tokens).
 2. Full: receive both the chunks as they are generated and the full response at the end.
 
@@ -73,9 +76,11 @@ For the complete working example, check out this [demo script](https://github.co
 ### Server API
 
 To make a call to the NeMo Guardrails Server in streaming mode, you have to set the `stream` parameter to `True` inside the JSON body. For example, to get the completion for a chat session using the `/v1/chat/completions` endpoint:
+
 ```
 POST /v1/chat/completions
 ```
+
 ```json
 {
     "config_id": "some_config_id",

docs/user-guides/configuration-guide.md

@@ -75,23 +75,29 @@ models:
 
 The meaning of the attributes is as follows:
 
-- `type`: is set to "main" indicating the main LLM model.
-- `engine`: the LLM provider, e.g., `openai`, `huggingface_endpoint`, `self_hosted`, etc.
-- `model`: the name of the model, e.g., `gpt-3.5-turbo-instruct`.
-- `parameters`: any additional parameters, e.g., `temperature`, `top_k`, etc.
+- `type`: is set to _main_ to indicate the model is the application LLM.
+- `engine`: the LLM provider, such as `openai`, `huggingface_endpoint`, `self_hosted`, and so on.
+- `model`: the name of the model, such as `gpt-3.5-turbo-instruct`.
+- `parameters`: arguments to pass to the LangChain class used by the LLM provider.
+  For example, when `engine` is set to `openai`, the toolkit loads the `ChatOpenAI` class.
+  The [ChatOpenAI class](https://python.langchain.com/api_reference/openai/chat_models/langchain_openai.chat_models.base.ChatOpenAI.html)
+  supports `temperature`, `max_tokens`, and other class-specific arguments.
 
 #### Supported LLM Providers
 
-You can use any LLM provider that is supported by LangChain, e.g., `ai21`, `aleph_alpha`, `anthropic`, `anyscale`, `azure`, `cohere`, `huggingface_endpoint`, `huggingface_hub`, `openai`, `self_hosted`, `self_hosted_hugging_face`. Check out the LangChain official documentation for the full list.
+You can use any LLM provider that is supported by LangChain, such as `ai21`, `aleph_alpha`, `anthropic`, `anyscale`, `azure`, `cohere`, `huggingface_endpoint`, `huggingface_hub`, `openai`, `self_hosted`, `self_hosted_hugging_face`. Check out the LangChain official documentation for the full list.
 
-In addition to the above LangChain providers, connecting to [Nvidia NIMs](https://docs.nvidia.com/nim/index.html) is supported using the engine `nvidia_ai_endpoints` or synonymously `nim`, for both Nvidia hosted NIMs (accessible through an Nvidia AI Enterprise license) and for locally downloaded and elf-hosted NIM containers.
+In addition to the above LangChain providers, connecting to [NVIDIA NIM microservices](https://docs.nvidia.com/nim/index.html) is supported using the `nim` engine.
+The `nvidia_ai_endpoints` engine is an alias for the `nim` engine.
+The engine provides access to locally-deployed NIM microservices or NVIDIA hosted models that you can view from <https://build.nvidia.com/models>.
 
-```{note}
-To use any of the providers, you must install additional packages; when you first try to use a configuration with a new provider, you typically receive an error from LangChain that instructs which packages you should install.
-```
+To use any of the LLM providers, you must install the LangChain package for the provider.
+When you first try to use a configuration with a new provider, you typically receive an error from LangChain that instructs which packages you should install.
 
 ```{important}
-Although you can instantiate any of the previously mentioned LLM providers, depending on the capabilities of the model, the NeMo Guardrails toolkit works better with some providers than others. The toolkit includes prompts that have been optimized for certain types of models, such as models provided by`openai` or `llama3` models. For others, you can optimize the prompts yourself following the information in the [LLM Prompts](#llm-prompts) section.
+Although you can instantiate any of the previously mentioned LLM providers, depending on the capabilities of the model, the NeMo Guardrails toolkit works better with some providers than others.
+The toolkit includes prompts that have been optimized for certain types of models, such as models provided by `openai` or `llama3` models.
+For others, you can optimize the prompts yourself following the information in the [LLM Prompts](#llm-prompts) section.
 ```
 
 #### Exploring Available Providers
@@ -797,7 +803,7 @@ rails:
 
 On a typical RAG (Retrieval Augmented Generation) scenario, using this option brings a 3x improvement in terms of latency and uses 37% fewer tokens.
 
-**IMPORTANT**: currently, the *Single Call Mode* can only predict bot messages as next steps. This means that if you want the LLM to generalize and decide to execute an action on a dynamically generated user canonical form message, it will not work.
+**IMPORTANT**: currently, the _Single Call Mode_ can only predict bot messages as next steps. This means that if you want the LLM to generalize and decide to execute an action on a dynamically generated user canonical form message, it will not work.
 
 #### Embeddings Only
 

nemoguardrails/library/injection_detection/actions.py

@@ -32,7 +32,7 @@
 import re
 from functools import lru_cache
 from pathlib import Path
-from typing import Dict, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, TypedDict, Union
 
 yara = None
 try:
@@ -49,6 +49,12 @@
 log = logging.getLogger(__name__)
 
 
+class InjectionDetectionResult(TypedDict):
+    is_injection: bool
+    text: str
+    detections: List[str]
+
+
 def _check_yara_available():
     if yara is None:
         raise ImportError(
@@ -197,13 +203,13 @@ def _load_rules(
             }
             rules = yara.compile(filepaths=rules_to_load)
     except yara.SyntaxError as e:
-        msg = f"Encountered SyntaxError: {e}"
+        msg = f"Failed to initialize injection detection due to configuration or YARA rule error: YARA compilation failed: {e}"
         log.error(msg)
-        raise e
+        return None
     return rules
 
 
-def _omit_injection(text: str, matches: list["yara.Match"]) -> str:
+def _omit_injection(text: str, matches: list["yara.Match"]) -> Tuple[bool, str]:
     """
     Attempts to strip the offending injection attempts from the provided text.
 
@@ -216,14 +222,18 @@ def _omit_injection(text: str, matches: list["yara.Match"]) -> str:
         matches (list['yara.Match']): A list of YARA rule matches.
 
     Returns:
-        str: The text with the detected injections stripped out.
+        Tuple[bool, str]: A tuple containing:
+            - bool: True if injection was detected and modified,
+                    False if the text is safe (i.e., not modified).
+            - str: The text, with detected injections stripped out if modified.
 
     Raises:
         ImportError: If the yara module is not installed.
     """
 
-    # Copy the text to a placeholder variable
+    original_text = text
     modified_text = text
+    is_injection = False
     for match in matches:
         if match.strings:
             for match_string in match.strings:
@@ -234,10 +244,16 @@ def _omit_injection(text: str, matches: list["yara.Match"]) -> str:
                             modified_text = modified_text.replace(plaintext, "")
                     except (AttributeError, UnicodeDecodeError) as e:
                         log.warning(f"Error processing match: {e}")
-    return modified_text
+
+    if modified_text != original_text:
+        is_injection = True
+        return is_injection, modified_text
+    else:
+        is_injection = False
+        return is_injection, original_text
 
 
-def _sanitize_injection(text: str, matches: list["yara.Match"]) -> str:
+def _sanitize_injection(text: str, matches: list["yara.Match"]) -> Tuple[bool, str]:
     """
     Attempts to sanitize the offending injection attempts in the provided text.
     This is done by 'de-fanging' the offending content, transforming it into a state that will not execute
@@ -253,19 +269,27 @@ def _sanitize_injection(text: str, matches: list["yara.Match"]) -> str:
         matches (list['yara.Match']): A list of YARA rule matches.
 
     Returns:
-        str: The text with the detected injections sanitized.
+        Tuple[bool, str]: A tuple containing:
+            - bool: True if injection was detected, False otherwise.
+            - str: The sanitized text, or original text depending on sanitization outcome.
+                   Currently, this function will always raise NotImplementedError.
 
     Raises:
         NotImplementedError: If the sanitization logic is not implemented.
         ImportError: If the yara module is not installed.
     """
-
     raise NotImplementedError(
         "Injection sanitization is not yet implemented. Please use 'reject' or 'omit'"
     )
+    # Hypothetical logic if implemented, to match existing behavior in injection_detection:
+    # sanitized_text_attempt = "..." # result of sanitization
+    # if sanitized_text_attempt != text:
+    #     return True, text  # Original text returned, marked as injection detected
+    # else:
+    #     return False, sanitized_text_attempt
 
 
-def _reject_injection(text: str, rules: "yara.Rules") -> Tuple[bool, str]:
+def _reject_injection(text: str, rules: "yara.Rules") -> Tuple[bool, List[str]]:
     """
     Detects whether the provided text contains potential injection attempts.
 
@@ -277,8 +301,9 @@ def _reject_injection(text: str, rules: "yara.Rules") -> Tuple[bool, str]:
         rules ('yara.Rules'): The loaded YARA rules.
 
     Returns:
-        bool: True if attempted exploitation is detected, False otherwise.
-        str: list of matches as a string
+        Tuple[bool, List[str]]: A tuple containing:
+            - bool: True if attempted exploitation is detected, False otherwise.
+            - List[str]: List of matched rule names.
 
     Raises:
         ValueError: If the `action` parameter in the configuration is invalid.
@@ -289,18 +314,20 @@ def _reject_injection(text: str, rules: "yara.Rules") -> Tuple[bool, str]:
         log.warning(
             "reject_injection guardrail was invoked but no rules were specified in the InjectionDetection config."
         )
-        return False, ""
+        return False, []
     matches = rules.match(data=text)
     if matches:
-        matches_string = ", ".join([match_name.rule for match_name in matches])
-        log.info(f"Input matched on rule {matches_string}.")
-        return True, matches_string
+        matched_rules = [match_name.rule for match_name in matches]
+        log.info(f"Input matched on rule {', '.join(matched_rules)}.")
+        return True, matched_rules
     else:
-        return False, ""
+        return False, []
 
 
 @action()
-async def injection_detection(text: str, config: RailsConfig) -> str:
+async def injection_detection(
+    text: str, config: RailsConfig
+) -> InjectionDetectionResult:
     """
     Detects and mitigates potential injection attempts in the provided text.
 
@@ -310,45 +337,68 @@ async def injection_detection(text: str, config: RailsConfig) -> str:
 
     Args:
         text (str): The text to check for command injection.
+
         config (RailsConfig): The Rails configuration object containing injection detection settings.
 
     Returns:
-        str: The sanitized or original text, depending on the action specified in the configuration.
+        InjectionDetectionResult: A TypedDict containing:
+            - is_injection (bool): Whether an injection was detected. True if any injection is detected,
+                            False if no injection is detected.
+            - text (str): The sanitized or original text
+            - detections (List[str]): List of matched rule names if any injection is detected
 
     Raises:
         ValueError: If the `action` parameter in the configuration is invalid.
         NotImplementedError: If an unsupported action is encountered.
+        ImportError: If the yara module is not installed.
     """
     _check_yara_available()
 
     _validate_injection_config(config)
+
     action_option, yara_path, rule_names, yara_rules = _extract_injection_config(config)
 
     rules = _load_rules(yara_path, rule_names, yara_rules)
 
-    if action_option == "reject":
-        verdict, detections = _reject_injection(text, rules)
-        if verdict:
-            return f"I'm sorry, the desired output triggered rule(s) designed to mitigate exploitation of {detections}."
-        else:
-            return text
     if rules is None:
         log.warning(
             "injection detection guardrail was invoked but no rules were specified in the InjectionDetection config."
         )
-        return text
-    matches = rules.match(data=text)
-    if matches:
-        matches_string = ", ".join([match_name.rule for match_name in matches])
-        log.info(f"Input matched on rule {matches_string}.")
-        if action_option == "omit":
-            return _omit_injection(text, matches)
-        elif action_option == "sanitize":
-            return _sanitize_injection(text, matches)
+        return InjectionDetectionResult(is_injection=False, text=text, detections=[])
+
+    if action_option == "reject":
+        is_injection, detected_rules = _reject_injection(text, rules)
+        return InjectionDetectionResult(
+            is_injection=is_injection, text=text, detections=detected_rules
+        )
+    else:
+        matches = rules.match(data=text)
+        if matches:
+            detected_rules_list = [match_name.rule for match_name in matches]
+            log.info(f"Input matched on rule {', '.join(detected_rules_list)}.")
+
+            if action_option == "omit":
+                is_injection, result_text = _omit_injection(text, matches)
+                return InjectionDetectionResult(
+                    is_injection=is_injection,
+                    text=result_text,
+                    detections=detected_rules_list,
+                )
+            elif action_option == "sanitize":
+                # _sanitize_injection will raise NotImplementedError before returning a tuple.
+                # the assignment below is for structural consistency if it were implemented.
+                is_injection, result_text = _sanitize_injection(text, matches)
+                return InjectionDetectionResult(
+                    is_injection=is_injection,
+                    text=result_text,
+                    detections=detected_rules_list,
+                )
+            else:
+                raise NotImplementedError(
+                    f"Expected `action` parameter to be 'reject', 'omit', or 'sanitize' but got {action_option} instead."
+                )
+        # no matches found
         else:
-            # We should never ever hit this since we inspect the action option above, but putting an error here anyway.
-            raise NotImplementedError(
-                f"Expected `action` parameter to be 'omit' or 'sanitize' but got {action_option} instead."
+            return InjectionDetectionResult(
+                is_injection=False, text=text, detections=[]
             )
-    else:
-        return text

nemoguardrails/library/injection_detection/flows.co

@@ -1,7 +1,19 @@
-# OUTPUT RAILS
-
 flow injection detection
   """
   Reject, omit, or sanitize injection attempts from the bot.
+  This rail operates on the $bot_message.
   """
-  $bot_message = await InjectionDetectionAction(text=$bot_message)
+  response = await InjectionDetectionAction(text=$bot_message)
+  join_separator = ", "
+  injection_detection_action = $config.rails.config.injection_detection.action
+
+  if response["is_injection"]
+    if $config.enable_rails_exceptions
+      send InjectionDetectionRailException(message="Output not allowed. The output was blocked by the 'injection detection' flow.")
+    else if injection_detection_action == "reject"
+      bot "I'm sorry, the desired output triggered rule(s) designed to mitigate exploitation of {{ response.detections | join(join_separator) }}."
+      abort
+    else if injection_detection_action == "omit" or injection_detection_action == "sanitize"
+      $bot_message = response["text"]
+  else
+    $bot_message = response["text"]

nemoguardrails/library/injection_detection/flows.v1.co

@@ -1,5 +1,19 @@
-define subflow injection detection
+
+define flow injection detection
   """
   Reject, omit, or sanitize injection attempts from the bot.
   """
-  $bot_message = execute injection_detection(text=$bot_message)
+  $response = execute injection_detection(text=$bot_message)
+  $join_separator = ", "
+  $injection_detection_action = $config.rails.config.injection_detection.action
+  if $response["is_injection"]
+    if $config.enable_rails_exceptions
+      create event InjectionDetectionRailException(message="Output not allowed. The output was blocked by the 'injection detection' flow.")
+      stop
+    else if $config.rails.config.injection_detection.action == "reject"
+      bot say "I'm sorry, the desired output triggered rule(s) designed to mitigate exploitation of {{ response.detections | join(join_separator) }}."
+      stop
+    else if $injection_detection_action == "omit" or $injection_detection_action == "sanitize"
+      $bot_message = $response["text"]
+  else
+    $bot_message = $response["text"]