test(self-check): add openai config smoke coverage

Author: Pouyanpi

tests/test_self_check_openai_smoke.py

@@ -0,0 +1,361 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from unittest.mock import patch
+
+import pytest
+
+from nemoguardrails import LLMRails, RailsConfig
+from nemoguardrails.imports import check_optional_dependency
+from nemoguardrails.testing.fake_model import FakeLLMModel
+
+COLANG = """
+define user express greeting
+    "hello"
+    "hi"
+
+define bot express greeting
+    "Hey!"
+
+define flow greeting
+    user express greeting
+    bot express greeting
+"""
+
+OPENAI_MODEL = """
+models:
+  - type: main
+    engine: openai
+    model: gpt-4o-mini
+"""
+
+LIVE_TEST_MODE = os.environ.get("LIVE_TEST_MODE") or os.environ.get("TEST_LIVE_MODE") or os.environ.get("LIVE_TEST")
+
+
+def _response_content(response):
+    return response["content"] if isinstance(response, dict) else response.response[0]["content"]
+
+
+@patch("nemoguardrails.rails.llm.llmrails.init_llm_model")
+def test_llmrails_openai_default_self_check_uses_main_model(mock_init_llm_model):
+    main_llm = FakeLLMModel(responses=["Yes"])
+    mock_init_llm_model.return_value = main_llm
+    config = RailsConfig.from_content(
+        colang_content=COLANG,
+        yaml_content=OPENAI_MODEL
+        + """
+rails:
+  input:
+    flows:
+      - self check input
+prompts:
+  - task: self_check_input
+    content: |
+      User message: "{{ user_input }}"
+      Answer No.
+enable_rails_exceptions: True
+""",
+    )
+
+    rails = LLMRails(config, verbose=False)
+    response = rails.generate(messages=[{"role": "user", "content": "hello"}])
+
+    assert response["role"] == "exception"
+    assert response["content"]["type"] == "InputRailException"
+    assert main_llm.inference_count == 1
+    assert mock_init_llm_model.call_args.kwargs["provider_name"] == "openai"
+
+
+@patch("nemoguardrails.rails.llm.llmrails.init_llm_model")
+def test_llmrails_openai_custom_self_check_models_route_by_task(mock_init_llm_model):
+    main_llm = FakeLLMModel(responses=["  ask question", "  Here is the answer."])
+    check_harmful_llm = FakeLLMModel(responses=["No"])
+    check_data_leakage_llm = FakeLLMModel(responses=["No"])
+    mock_init_llm_model.side_effect = [main_llm, check_harmful_llm, check_data_leakage_llm]
+    config = RailsConfig.from_content(
+        colang_content="""
+define user ask question
+    "tell me something"
+
+define flow
+    user ask question
+    bot respond
+""",
+        yaml_content="""
+models:
+  - type: main
+    engine: openai
+    model: gpt-4o-mini
+  - type: check_harmful
+    engine: openai
+    model: gpt-4o-mini
+  - type: check_data_leakage
+    engine: openai
+    model: gpt-4o-mini
+rails:
+  input:
+    flows:
+      - self check input $task=check_harmful
+  output:
+    flows:
+      - self check output $task=check_data_leakage
+prompts:
+  - task: check_harmful
+    content: |
+      User message: "{{ user_input }}"
+      Answer No.
+  - task: check_data_leakage
+    content: |
+      Bot response: "{{ bot_response }}"
+      Answer No.
+enable_rails_exceptions: True
+""",
+    )
+
+    rails = LLMRails(config, verbose=False)
+    response = rails.generate(messages=[{"role": "user", "content": "tell me something"}])
+
+    assert _response_content(response) == "Here is the answer."
+    assert check_harmful_llm.inference_count == 1
+    assert check_data_leakage_llm.inference_count == 1
+    assert [call.kwargs["provider_name"] for call in mock_init_llm_model.call_args_list] == [
+        "openai",
+        "openai",
+        "openai",
+    ]
+
+
+@patch("nemoguardrails.rails.llm.llmrails.init_llm_model")
+def test_llmrails_openai_custom_task_falls_back_to_default_self_check_model(mock_init_llm_model):
+    main_llm = FakeLLMModel(responses=[])
+    default_self_check_llm = FakeLLMModel(responses=["Yes"])
+    mock_init_llm_model.side_effect = [main_llm, default_self_check_llm]
+    config = RailsConfig.from_content(
+        colang_content=COLANG,
+        yaml_content="""
+models:
+  - type: main
+    engine: openai
+    model: gpt-4o-mini
+  - type: self_check_input
+    engine: openai
+    model: gpt-4o-mini
+rails:
+  input:
+    flows:
+      - self check input $task=check_harmful
+prompts:
+  - task: check_harmful
+    content: |
+      User message: "{{ user_input }}"
+      Answer No.
+enable_rails_exceptions: True
+""",
+    )
+
+    rails = LLMRails(config, verbose=False)
+    response = rails.generate(messages=[{"role": "user", "content": "hello"}])
+
+    assert response["role"] == "exception"
+    assert response["content"]["type"] == "InputRailException"
+    assert main_llm.inference_count == 0
+    assert default_self_check_llm.inference_count == 1
+
+
+@patch("nemoguardrails.rails.llm.llmrails.init_llm_model")
+def test_llmrails_openai_parallel_input_model_routes_by_task(mock_init_llm_model):
+    main_llm = FakeLLMModel(responses=["  express greeting", '  "Hey!"'])
+    check_harmful_llm = FakeLLMModel(responses=["No"])
+    default_input_llm = FakeLLMModel(responses=["Yes"])
+    mock_init_llm_model.side_effect = [main_llm, check_harmful_llm, default_input_llm]
+    config = RailsConfig.from_content(
+        colang_content=COLANG,
+        yaml_content="""
+models:
+  - type: main
+    engine: openai
+    model: gpt-4o-mini
+  - type: check_harmful
+    engine: openai
+    model: gpt-4o-mini
+  - type: self_check_input
+    engine: openai
+    model: gpt-4o-mini
+rails:
+  input:
+    parallel: true
+    flows:
+      - self check input $task=check_harmful
+prompts:
+  - task: check_harmful
+    content: |
+      User message: "{{ user_input }}"
+      Answer No.
+  - task: self_check_input
+    content: |
+      User message: "{{ user_input }}"
+      Answer Yes.
+enable_rails_exceptions: True
+""",
+    )
+
+    rails = LLMRails(config, verbose=False)
+    response = rails.generate(messages=[{"role": "user", "content": "hello"}])
+
+    assert _response_content(response) == "Hey!"
+    assert check_harmful_llm.inference_count == 1
+    assert default_input_llm.inference_count == 0
+
+
+@patch("nemoguardrails.rails.llm.llmrails.init_llm_model")
+def test_llmrails_openai_parallel_output_model_routes_by_task(mock_init_llm_model):
+    main_llm = FakeLLMModel(responses=["  ask question", "  Here is the answer."])
+    check_data_leakage_llm = FakeLLMModel(responses=["No"])
+    default_output_llm = FakeLLMModel(responses=["Yes"])
+    mock_init_llm_model.side_effect = [main_llm, check_data_leakage_llm, default_output_llm]
+    config = RailsConfig.from_content(
+        colang_content="""
+define user ask question
+    "tell me something"
+
+define flow
+    user ask question
+    bot respond
+""",
+        yaml_content="""
+models:
+  - type: main
+    engine: openai
+    model: gpt-4o-mini
+  - type: check_data_leakage
+    engine: openai
+    model: gpt-4o-mini
+  - type: self_check_output
+    engine: openai
+    model: gpt-4o-mini
+rails:
+  output:
+    parallel: true
+    flows:
+      - self check output $task=check_data_leakage
+prompts:
+  - task: check_data_leakage
+    content: |
+      Bot response: "{{ bot_response }}"
+      Answer No.
+  - task: self_check_output
+    content: |
+      Bot response: "{{ bot_response }}"
+      Answer Yes.
+enable_rails_exceptions: True
+""",
+    )
+
+    rails = LLMRails(config, verbose=False)
+    response = rails.generate(messages=[{"role": "user", "content": "tell me something"}])
+
+    assert _response_content(response) == "Here is the answer."
+    assert check_data_leakage_llm.inference_count == 1
+    assert default_output_llm.inference_count == 0
+
+
+@pytest.mark.asyncio
+@patch("nemoguardrails.rails.llm.llmrails.init_llm_model")
+async def test_llmrails_openai_streaming_output_rail_routes_custom_task(mock_init_llm_model):
+    main_llm = FakeLLMModel(
+        responses=[
+            '  express greeting\nbot express greeting\n  "Hi, how are you doing?"',
+            "This response should stream safely.",
+        ]
+    )
+    check_data_leakage_llm = FakeLLMModel(responses=["No", "No"])
+    default_output_llm = FakeLLMModel(responses=["Yes"])
+    mock_init_llm_model.side_effect = [main_llm, check_data_leakage_llm, default_output_llm]
+    config = RailsConfig.from_content(
+        colang_content="""
+define user express greeting
+    "hi"
+
+define flow
+    user express greeting
+    bot tell joke
+""",
+        yaml_content="""
+models:
+  - type: main
+    engine: openai
+    model: gpt-4o-mini
+  - type: check_data_leakage
+    engine: openai
+    model: gpt-4o-mini
+  - type: self_check_output
+    engine: openai
+    model: gpt-4o-mini
+rails:
+  output:
+    flows:
+      - self check output $task=check_data_leakage
+    streaming:
+      enabled: true
+      chunk_size: 4
+      context_size: 2
+      stream_first: false
+prompts:
+  - task: check_data_leakage
+    content: |
+      Bot response: "{{ bot_response }}"
+      Answer No.
+  - task: self_check_output
+    content: |
+      Bot response: "{{ bot_response }}"
+      Answer Yes.
+""",
+    )
+
+    rails = LLMRails(config, verbose=False)
+    chunks = []
+    async for chunk in rails.stream_async(messages=[{"role": "user", "content": "hi"}]):
+        chunks.append(chunk)
+
+    assert "This response should stream safely." in "".join(chunks)
+    assert check_data_leakage_llm.inference_count > 0
+    assert default_output_llm.inference_count == 0
+
+
+@pytest.mark.skipif(
+    not (LIVE_TEST_MODE and os.environ.get("OPENAI_API_KEY") and check_optional_dependency("langchain_openai")),
+    reason="Requires LIVE_TEST_MODE or TEST_LIVE_MODE, OPENAI_API_KEY, and langchain_openai",
+)
+def test_llmrails_live_openai_default_self_check_smoke():
+    config = RailsConfig.from_content(
+        colang_content=COLANG,
+        yaml_content=OPENAI_MODEL
+        + """
+rails:
+  input:
+    flows:
+      - self check input
+prompts:
+  - task: self_check_input
+    content: |
+      Always answer exactly No.
+""",
+    )
+
+    rails = LLMRails(config, verbose=False)
+    response = rails.generate(messages=[{"role": "user", "content": "hello"}])
+
+    assert response["role"] == "assistant"