Added test coverage for nemoguardrails/rails/llm/config.py

Author: tgasser-nv

nemoguardrails/rails/llm/options.py

@@ -434,4 +434,6 @@ class GenerationResponse(BaseModel):
 
 
 if __name__ == "__main__":
-    print(GenerationOptions(rails=GenerationRailsOptions(input=False)))
+    print(
+        GenerationOptions(rails=GenerationRailsOptions(input=False))
+    )  # pragma: no cover (Can't run as script for test coverage)

tests/rails/llm/test_config.py

@@ -13,19 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 from unittest.mock import MagicMock
 
 import pytest
 from langchain.llms.base import BaseLLM
 from pydantic import ValidationError
 
-from nemoguardrails.rails.llm.config import (
-    Document,
-    Instruction,
-    Model,
-    RailsConfig,
-    TaskPrompt,
-)
+from nemoguardrails.rails.llm.config import Model, RailsConfig, TaskPrompt
 from nemoguardrails.rails.llm.llmrails import LLMRails
 
 
@@ -342,3 +337,45 @@ def test_llm_rails_configure_streaming_without_attr(caplog):
     rails._configure_main_llm_streaming(mock_llm)
 
     assert caplog.messages[-1] == "Provided main LLM does not support streaming."
+
+
+def test_rails_config_streaming_supported_no_output_flows():
+    """Check `streaming_supported` property doesn't depend on RailsConfig.streaming with no output flows"""
+
+    config = RailsConfig(
+        models=[],
+        streaming=False,
+    )
+    assert config.streaming_supported
+
+
+def test_rails_config_flows_streaming_supported_true():
+    """Create RailsConfig and check the `streaming_supported Check LLM has the streaming attribute set if RailsConfig has it"""
+
+    rails = {
+        "output": {
+            "flows": ["content_safety_check_output"],
+            "streaming": {"enabled": True},
+        }
+    }
+    prompts = [{"task": "content safety check output", "content": "..."}]
+    rails_config = RailsConfig.model_validate(
+        {"models": [], "rails": rails, "prompts": prompts}
+    )
+    assert rails_config.streaming_supported
+
+
+def test_rails_config_flows_streaming_supported_false():
+    """Create RailsConfig and check the `streaming_supported Check LLM has the streaming attribute set if RailsConfig has it"""
+
+    rails = {
+        "output": {
+            "flows": ["content_safety_check_output"],
+            "streaming": {"enabled": False},
+        }
+    }
+    prompts = [{"task": "content safety check output", "content": "..."}]
+    rails_config = RailsConfig.model_validate(
+        {"models": [], "rails": rails, "prompts": prompts}
+    )
+    assert not rails_config.streaming_supported