Feat: Support models with reasoning traces (#996)

* Fix jailbreak models test for cases when JAILBREAK_SETUP_PRESENT env variable is not defined

* Add support from removing reasoning traces for reasoning LLMs (e.g. Deepseek-R1)

* Add default prompts for Deepseek models.

* Changed how output parsers are used by the self check rails.

* Added test for a config with LLM reasoning for LLM rails.

* Added DeepSeek-R1 example config and updated documentation on reasoning models.

* Added support for deepseek and google_genai as providers

* Refactored model name config for langchain providers

Author: trebedea

docs/user-guides/configuration-guide.md

@@ -94,6 +94,29 @@ To use any of the providers, you must install additional packages; when you firs
 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 `openai` and `nemollm`. For others, you can optimize the prompts yourself following the information in the [LLM Prompts](#llm-prompts) section.
 ```
 
+#### Using LLMs with Reasoning Traces
+
+To use an LLM that outputs the reasoning traces as part of the response (e.g. [DeepSeek-R1](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d)), the following model config should be used:
+
+```yaml
+models:
+  - type: main
+    engine: deepseek
+    model: deepseek-reasoner
+    reasoning_config:
+      remove_thinking_traces: True
+      start_token: "<think>"
+      end_token: "</think>"
+```
+
+The `reasoning_config` attribute for a model contains all the required configuration for a reasoning model that outputs reasoning traces.
+In most of the cases, the reasoning traces need to be removed and the guardrails runtime will only process the actual responses from the LLM.
+
+The attributes that can be configured for a reasoning model are:
+- `remove_thinking_traces`: if the reasoning traces should be ignored (defaults to `True`.
+- `start_token`: the start token for the reasoning process (e.g. `<think>` for DeepSeek-R1).
+- `end_token`: the end token for the reasoning process (e.g. `</think>` for DeepSeek-R1).
+
 #### NIM for LLMs
 
 [NVIDIA NIM](https://docs.nvidia.com/nim/index.html) is a set of easy-to-use microservices designed to accelerate the deployment of generative AI models across the cloud, data center, and workstations.

examples/configs/llm/deepseek-r1/config.yml

@@ -0,0 +1,8 @@
+models:
+  - type: main
+    engine: deepseek
+    model: deepseek-reasoner
+    reasoning_config:
+      remove_thinking_traces: True
+      start_token: "<think>"
+      end_token: "</think>"

nemoguardrails/library/self_check/facts/actions.py

@@ -78,7 +78,9 @@ async def self_check_facts(
     if llm_task_manager.has_output_parser(task):
         result = llm_task_manager.parse_task_output(task, output=response)
     else:
-        result = llm_task_manager.output_parsers["is_content_safe"](response)
+        result = llm_task_manager.parse_task_output(
+            task, output=response, forced_output_parser="is_content_safe"
+        )
 
     is_not_safe, _ = result
 

nemoguardrails/library/self_check/input_check/actions.py

@@ -79,7 +79,9 @@ async def self_check_input(
             result = llm_task_manager.parse_task_output(task, output=response)
 
         else:
-            result = llm_task_manager.output_parsers["is_content_safe"](response)
+            result = llm_task_manager.parse_task_output(
+                task, output=response, forced_output_parser="is_content_safe"
+            )
 
         is_safe, _ = result
 

nemoguardrails/library/self_check/output_check/actions.py

@@ -83,7 +83,9 @@ async def self_check_output(
         if llm_task_manager.has_output_parser(task):
             result = llm_task_manager.parse_task_output(task, output=response)
         else:
-            result = llm_task_manager.output_parsers["is_content_safe"](response)
+            result = llm_task_manager.parse_task_output(
+                task, output=response, forced_output_parser="is_content_safe"
+            )
 
         is_safe, _ = result
 

nemoguardrails/llm/filters.py

@@ -482,3 +482,26 @@ def conversation_to_events(conversation: List) -> List[dict]:
             )
 
     return events
+
+
+def remove_reasoning_traces(response: str, start_token: str, end_token: str) -> str:
+    """Removes the text between the first occurrence of the start token and the
+    last occurrence of the last token, if these tokens exist in the response.
+
+    This utility function is useful to strip reasoning traces from reasoning LLMs
+    that encode the reasoning traces between specific tokens.
+    """
+    if start_token and end_token:
+        start_index = response.find(start_token)
+        # If the start index is missing, this is probably a continuation of a bot message
+        # started in the prompt.
+        if start_index == -1:
+            start_index = 0
+        end_index = response.rfind(end_token)
+        if end_index == -1:
+            return response
+
+        if start_index != -1 and end_index != -1 and start_index < end_index:
+            return response[:start_index] + response[end_index + len(end_token) :]
+
+    return response

nemoguardrails/llm/prompts.py

@@ -20,7 +20,7 @@
 import yaml
 
 from nemoguardrails.llm.types import Task
-from nemoguardrails.rails.llm.config import RailsConfig, TaskPrompt
+from nemoguardrails.rails.llm.config import Model, RailsConfig, TaskPrompt
 
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 
@@ -116,21 +116,33 @@ def _get_prompt(
     raise ValueError(f"Could not find prompt for task {task_name} and model {model}")
 
 
-def get_prompt(config: RailsConfig, task: Union[str, Task]) -> TaskPrompt:
-    """Return the prompt for the given task."""
-
+def get_task_model(config: RailsConfig, task: Union[str, Task]) -> Model:
+    """Return the model for the given task in the current config."""
     # Fetch current task parameters like name, models to use, and the prompting mode
     task_name = str(task.value) if isinstance(task, Task) else task
 
-    task_model = "unknown"
     if config.models:
         _models = [model for model in config.models if model.type == task_name]
         if not _models:
             _models = [model for model in config.models if model.type == "main"]
 
-        task_model = _models[0].engine
-        if _models[0].model:
-            task_model += "/" + _models[0].model
+        return _models[0]
+
+    return None
+
+
+def get_prompt(config: RailsConfig, task: Union[str, Task]) -> TaskPrompt:
+    """Return the prompt for the given task."""
+
+    # Fetch current task parameters like name, models to use, and the prompting mode
+    task_name = str(task.value) if isinstance(task, Task) else task
+
+    task_model = "unknown"
+    _model = get_task_model(config, task)
+    if _model:
+        task_model = _model.engine
+        if _model.model:
+            task_model += "/" + _model.model
 
     task_prompting_mode = "standard"
     if config.prompting_mode: