@@ -812,9 +812,9 @@ def test_main_llm_supports_streaming_flag_disabled_when_no_streaming():
812812 fake_llm = FakeLLM (responses = ["test" ], streaming = False )
813813 rails = LLMRails (config , llm = fake_llm )
814814
815- assert (
816- rails . main_llm_supports_streaming is False
817- ), "main_llm_supports_streaming should be False when streaming is disabled"
815+ assert rails . main_llm_supports_streaming is False , (
816+ " main_llm_supports_streaming should be False when streaming is disabled"
817+ )
818818
819819
820820def test_main_llm_supports_streaming_with_multiple_model_types (
@@ -826,12 +826,12 @@ def test_main_llm_supports_streaming_with_multiple_model_types(
826826 "models" : [
827827 {
828828 "type" : "main" ,
829- "engine" : "custom_streaming_1 " ,
829+ "engine" : "custom_streaming " ,
830830 "model" : "test-model" ,
831831 },
832832 {
833833 "type" : "content_safety" ,
834- "engine" : "custom_streaming_2 " ,
834+ "engine" : "custom_streaming " ,
835835 "model" : "safety-model" ,
836836 },
837837 ],
@@ -846,17 +846,17 @@ def test_main_llm_supports_streaming_with_multiple_model_types(
846846 "and config has multiple model types including a streaming-capable main LLM"
847847 )
848848 # Verify the main LLM's streaming attribute was set
849- assert (
850- hasattr ( rails . llm , "streaming" ) and rails . llm . streaming is True
851- ), "Main LLM's streaming attribute should be set to True"
849+ assert hasattr ( rails . llm , "streaming" ) and rails . llm . streaming is True , (
850+ "Main LLM's streaming attribute should be set to True"
851+ )
852852
853853
854854def test_main_llm_supports_streaming_with_specialized_models_only (
855855 custom_streaming_providers ,
856856):
857857 """Test streaming config when only specialized models are defined (no main)."""
858858 config = RailsConfig .from_content (
859- config = {g
859+ config = {
860860 "models" : [
861861 {
862862 "type" : "content_safety" ,
@@ -871,6 +871,6 @@ def test_main_llm_supports_streaming_with_specialized_models_only(
871871 rails = LLMRails (config )
872872
873873 # Verify that main_llm_supports_streaming is False when no main LLM is configured
874- assert (
875- rails . main_llm_supports_streaming is False
876- ), "main_llm_supports_streaming should be False when no main LLM is configured"
874+ assert rails . main_llm_supports_streaming is False , (
875+ " main_llm_supports_streaming should be False when no main LLM is configured"
876+ )
0 commit comments